在stable-diffusion-webui-directml中优化VAE显存占用的技术方案

问题背景

在使用stable-diffusion-webui-directml项目进行图像生成时，许多用户特别是使用低显存显卡（如4GB或更低）的用户会遇到一个常见问题：在生成过程的最后阶段，当VAE（变分自编码器）进行解码时，系统会因显存不足而崩溃。这个问题尤其令人沮丧，因为生成过程已经完成了大部分计算，却在最后一步功亏一篑。

技术分析

VAE在Stable Diffusion模型中负责将潜在空间表示解码为最终的图像像素。这一过程虽然计算量不大，但需要加载整个VAE模型到显存中，对于低显存设备来说是一个挑战。通过分析，我们发现：

1. 在生成过程中，显存占用通常维持在50%左右
2. 当进行VAE解码时，显存需求突然增加
3. 4GB显存设备在512x768分辨率下就可能出现OOM（内存不足）错误

解决方案

方案一：使用低显存模式

项目本身提供了--lowvram参数来优化显存使用。这是最推荐的解决方案，因为它经过了官方优化和测试。使用方式是在启动参数中添加：

--use-directml --lowvram --opt-sub-quad-attention --opt-split-attention --no-half-vae --upcast-sampling

方案二：分块VAE处理

对于更复杂的情况，可以使用分块VAE处理技术：

1. 安装分块VAE扩展
2. 仅启用"Tiled VAE"选项
3. 根据模型类型设置合适的编码器/解码器分辨率：
   • 对于1.5基础模型：编码器1024，解码器128
   • 对于SDXL模型：编码器1280，解码器128

方案三：VAE CPU运行（高级方案）

对于极端情况，可以将VAE强制运行在CPU上。这需要修改源代码：

1. 在processing.py中添加代码将潜在表示转移到CPU
2. 在lowvram.py中注释相关代码行
3. 注意：此方法会导致解码速度显著下降

最佳实践建议

1. 对于4GB显存设备：

   • 基础分辨率不超过512x512
   • 避免使用Hires Fix功能
   • 在img2img中使用SD-Upscale脚本而非"Resize to"

2. 对于2GB显存设备：

   • 考虑使用TAESD（轻量级VAE）
   • 分辨率控制在512x512以下
   • 编码器/解码器分辨率设置为512/64

3. 通用优化：

   • 保持项目版本最新
   • 使用精简版模型（2GB左右的1.5基础模型）
   • 避免复杂的提示词和过多LoRA

性能权衡

每种解决方案都有其优缺点：

• 低显存模式：平衡了性能和显存使用，是最佳折中方案
• 分块VAE：可以处理更高分辨率，但可能出现边缘瑕疵
• CPU运行VAE：完全避免显存问题，但速度最慢

结论

通过合理配置stable-diffusion-webui-directml的参数和采用适当的优化技术，即使是低显存设备也能获得可用的图像生成体验。对于大多数用户，推荐优先尝试官方提供的低显存模式，只有在特殊情况下才考虑更高级的定制方案。随着项目的持续更新，未来可能会有更多针对低显存设备的优化方案出现。