在AMD GPU版Stable Diffusion WebUI中优化VAE显存占用的技术方案

问题背景

AMD GPU版Stable Diffusion WebUI（基于DirectML）在低显存显卡（如4GB或更低）上运行时，经常在VAE解码阶段出现显存不足的问题。虽然生成过程可能顺利完成，但最后的VAE解码步骤却会导致"Out of Memory"错误，使得整个生成过程功亏一篑。

现有解决方案分析

1. 低显存模式(--lowvram)

这是最基础的解决方案，通过添加--lowvram参数可以显著降低显存使用量。但根据用户反馈，在某些情况下，即使启用了低显存模式，VAE解码阶段仍然会耗尽显存。

2. 平铺VAE(Tiled VAE)

平铺VAE是一种将大图像分割成小块分别处理的技术，可以有效降低显存需求。使用时需要注意：

• 对于SD1.5模型，建议设置编码器分辨率为1024，解码器为128
• 对于SDXL模型，建议设置"Upscale by"为1.5，分辨率为768x1024，编码器为1280，解码器为128

但平铺VAE在AMD显卡上可能会出现左下角灰色区块的渲染问题，这是DirectML驱动的一个已知bug。

进阶解决方案：强制VAE在CPU运行

对于极端低显存环境（如2-4GB显卡），可以考虑强制VAE在CPU上运行。这种方法虽然会显著降低处理速度，但能彻底解决显存不足问题。

实施步骤

1. 确保使用最新版WebUI
2. 在启动参数中添加--lowvram --no-half-vae
3. 修改源代码：
   • 在modules/processing.py中，在约1082行后添加：

     samples_ddim = samples_ddim.to(devices.dtype_vae).to("cpu")
     x_samples_ddim = decode_first_stage(p.sd_model, samples_ddim)
     

   • 在modules/lowvram.py中，注释掉第73行

注意事项

• 此方法会显著降低处理速度，仅建议在无法使用其他方法时采用
• 仅测试过SD1.5模型，SDXL模型可能需要更多测试
• 修改源代码前请做好备份

其他优化建议

1. 使用轻量级VAE：如TAESD，可在设置中启用
2. 合理设置分辨率：4GB显存建议最大768x768分辨率
3. 避免使用Hires Fix：低显存环境下建议使用SD-Upscale脚本替代
4. 优化参数组合：--use-directml --lowvram --opt-sub-quad-attention --opt-split-attention --no-half-vae --upcast-sampling

总结

针对AMD GPU版Stable Diffusion WebUI在低显存环境下的VAE显存问题，开发者可以根据自身硬件条件和需求选择适合的解决方案。从简单的启动参数优化到复杂的源代码修改，各种方法各有优劣。理解这些技术方案的原理和实施细节，能够帮助用户在有限的硬件资源下获得最佳的AI绘画体验。