优化 Stable Diffusion XL (SDXL) 的推理速度与内存占用
Hugging Face 详细介绍了一系列简单的优化方案，旨在使 Stable Diffusion XL (SDXL) 在推理时更加实用，尤其是在消费级 GPU 上。通过实施低精度、内存高效的注意力机制以及各种卸载技术，内存占用可以从 28GB 降低到 11.47GB，推理延迟可以从 72.2 秒降低到大约 10.3 秒。
推理速度优化
在扩散模型中，针对速度进行优化至关重要，因为图像生成的迭代性质通常需要多次运行才能达到理想结果。以下技术侧重于降低延迟：
低精度 (fp16)
使用 float16 (fp16) 代替标准的 float32 (fp32) 可以将内存占用减少一半并提高计算速度，因为 fp16 捕获的浮点数范围较窄，且得到现代 GPU 硬件更好的支持。在测试的 SDXL 流水线中，切换到 fp16 将内存占用减少到 21.7GB，并将推理时间从 72.2 秒降低到 14.8 秒。
内存高效注意力机制 (SDPA)
Transformer 中的注意力块可能会造成内存瓶颈，因为内存需求随输入序列长度呈平方级增长。PyTorch 2.0 引入了 Scaled Dot Product Attention (SDPA)，它提供了 Flash Attention 和内存高效注意力 (xFormers) 的融合实现。在 🤗Diffusers 中，对于 PyTorch ≥ 2.0，SDPA 默认启用，在保持 21.7GB 内存占用的同时，进一步将推理时间缩短至 11.4 秒。
使用 torch.compile 进行 JIT 编译
PyTorch 2.0 的 torch.compile API 允许将 PyTorch 代码即时 (JIT) 编译为优化的内核。使用 torch.compile（使用 mode="reduce-overhead"）对 SDXL UNet 进行封装，可进一步将推理时间提升至 10.2 秒。请注意，第一次编译运行速度较慢，但随后的调用速度会显著加快。
减少模型内存占用
由于 SDXL 的规模大约是之前 Stable Diffusion 模型的 3 倍（拥有 3.5B 参数的 UNet），将其装入 VRAM 是主要挑战。有几种技术可以减少内存占用：
CPU 卸载 (CPU Offloading)
模型卸载会在流水线的组件不需要在 GPU 上活动时将其移动到 CPU：
- Model CPU Offloading：将 UNet 加载到 GPU 内存中，同时将文本编码器和 VAE 保留在 CPU 上。这将内存占用降低到 20.2GB。
- Sequential CPU Offloading：将单个 UNet 子模块的权重卸载到 CPU，仅在进行前向传播前立即将其加载到 GPU。这将内存降低到 19.9GB，但显著增加了延迟至 67 秒。
VAE 切片 (VAE Slicing)
变分自编码器 (VAE) 将潜变量解码为图像，这是一个内存占用随 Batch Size 规模化的过程。VAE 切片将输入张量拆分为较小的切片，并分多个步骤进行解码。这种优化可以在不显著影响延迟的情况下，将内存占用降低到 15.4GB。
缓存计算结果
SDXL 利用两个文本编码器从提示词中计算嵌入 (embeddings)。由于这些嵌入在整个反向扩散过程中保持不变，因此可以预先计算并缓存。一旦生成了嵌入，文本编码器和分词器就可以从 GPU 内存中移除，结合 SDPA 和 fp16 后，内存占用为 21.9GB。
微型自编码器 (TAESD)
使用蒸馏版本替换标准 VAE（例如由 madebyollin 开发的 Tiny Autoencoder，约 10MB），可将内存占用降低到 15.6GB 并减少推理延迟。然而，Tiny Autoencoder 可能会忽略细微细节，主要推荐用于图像预览。
性能总结
在 A100 GPU (40 GB) 上针对每个提示词生成 4 张图像进行的测试显示了内存与延迟之间的权衡：
| 技术 | 内存 (GB) | 推理延迟 (ms) |
|---|---|---|
| 未优化流水线 | 28.09 | 72200.5 |
| fp16 | 21.72 | 14800.9 |
| fp16 + SDPA (默认) | 21.72 | 11413.0 |
默认 + torch.compile |
21.73 | 10296.7 |
| 默认 + model CPU offload | 20.21 | 16082.2 |
| 默认 + sequential CPU offload | 19.91 | 67034.0 |
| 默认 + VAE slicing | 15.40 | 11232.2 |
| 默认 + VAE slicing + sequential CPU offload | 11.47 | 66869.2 |
| 默认 + 预计算文本嵌入 | 21.85 | 11909.0 |
| 默认 + Tiny Autoencoder | 15.48 | 10449.7 |