M1 Pro设备上Stable Diffusion黑图及PyTorch内存问题求助

• 为适配stabilityai/stable-diffusion-x4-upscaler模型升级PyTorch后，除超分模型外，其余Stable Diffusion模型生成的图像均被标记为NSFW，最终显示为黑色
• 升级Mac OS到13.3.1后，测试多组PyTorch、Diffusers、Transformers版本组合，均无法解决NSFW误判问题
• 使用PyTorch 2.0时NSFW误判消失，但无法在pipeline中设置torch_dtype=torch.float16，且触发MPS内存不足报错：

RuntimeError: MPS backend out of memory (MPS allocated: 18.04 GB, other allocations: 94.99 MB, max allowed: 18.13 GB). Tried to allocate 4.00 KB on private pool. Use PYTORCH_MPS_HIGH_WATERMARK_RATIO=0.0 to disable upper limit for memory allocations (may cause system failure). INFO:     Stopping reloader process [15702]

使用的模型与代码 #

main_model_id = "runwayml/stable-diffusion-v1-5"
inpainting_model_id = "runwayml/stable-diffusion-inpainting"
upscaler_model_id = "stabilityai/stable-diffusion-x4-upscaler"

text2imgPipe = StableDiffusionPipeline.from_pretrained(main_model_id, torch_dtype=torch.float16).to(
    device)
text2imgPipe.enable_attention_slicing()

img2imgPipe = StableDiffusionImg2ImgPipeline.from_pretrained(main_model_id, torch_dtype=torch.float16).to(
    device)
img2imgPipe.enable_attention_slicing()

inpaintingPipe = StableDiffusionInpaintPipeline.from_pretrained(inpainting_model_id, torch_dtype=torch.float16).to(
    device)
inpaintingPipe.enable_attention_slicing()

upscalerPipe = StableDiffusionUpscalePipeline.from_pretrained(upscaler_model_id, torch_dtype=torch.float16).to(
    device)
upscalerPipe.enable_attention_slicing()

当前环境版本 #

transformers              4.26.0
torchaudio                2.0.0
torchvision               0.14.1
pytorch                   2.0.0
diffusers                 0.12.0
huggingface-hub           0.13.4
python                    3.10.10

一、解决NSFW误判导致图像变黑的问题 #

1. 强制禁用NSFW检测器
   在加载pipeline后添加以下代码，直接跳过NSFW检测逻辑：

   text2imgPipe.safety_checker = lambda images, clip_input: (images, False)
   img2imgPipe.safety_checker = lambda images, clip_input: (images, False)
   inpaintingPipe.safety_checker = lambda images, clip_input: (images, False)
   

   注：超分模型本身不带NSFW检测器，无需处理。

2. 回退到兼容的版本组合
   尝试使用经过验证的稳定版本组合，避免版本冲突：

   • PyTorch: 1.13.1
   • Diffusers: 0.11.1
   • Transformers: 4.25.1
     执行安装命令：

   pip install torch==1.13.1 torchvision==0.14.1 torchaudio==0.13.1 diffusers==0.11.1 transformers==4.25.1
   

二、解决PyTorch 2.0下的float16与MPS内存问题 #

1. 修复float16无法使用的问题
   PyTorch 2.0对MPS的float16支持存在兼容性问题，可通过手动转换模型参数解决：

   # 加载时先不指定dtype，加载后再转换各模块精度
   text2imgPipe = StableDiffusionPipeline.from_pretrained(main_model_id).to(device)
   text2imgPipe.unet.to(torch.float16)
   text2imgPipe.vae.to(torch.float16)
   text2imgPipe.text_encoder.to(torch.float32)  # 文本编码器在MPS上用float32更稳定
   text2imgPipe.enable_attention_slicing()
   

2. 缓解MPS内存不足问题

   • 启用更高效的内存优化：

     text2imgPipe.enable_attention_slicing("max")  # 最大化注意力切片，减少内存占用
     text2imgPipe.enable_model_cpu_offload()  # 将模型部分层卸载到CPU，仅在需要时加载到MPS
     

   • 降低生成图像分辨率：将默认的512x512改为448x448或更小尺寸
   • 放宽内存限制（谨慎使用，可能导致系统卡顿）：
     在运行脚本前执行以下命令：

     export PYTORCH_MPS_HIGH_WATERMARK_RATIO=0.0
     

内容的提问来源于stack exchange，提问作者itsDanial