Stable Diffusion 3 Medium 与 Diffusers 的集成

Hugging Face 已将 Stable Diffusion 3 Medium，一款来自 Stability AI 的 2B 参数模型，集成到 diffusers 库中。此版本使用户能够使用一种名为多模态扩散 Transformer（MMDiT）的新架构方法以及校正流匹配训练目标，进行高质量的文本到图像和图像到图像合成。

Stable Diffusion 3 的架构创新

Stable Diffusion 3（SD3）通过在文本和图像数据之间采用双向信息流，区别于以往的文本到图像架构。

多模态扩散 Transformer（MMDiT）

SD3 使用一种新颖的 MMDiT 模型，并配合三个文本编码器：CLIP L/14、OpenCLIP bigG/14 和 T5‑v1.1‑XXL。它还使用了类似于 Stable Diffusion XL 中的 16 通道 AutoEncoder。

与之前使用固定文本表示进行交叉注意力的模型不同，MMDiT 块将文本输入和像素潜在向量视为嵌入序列进行处理。这些序列使用各自的权重映射到相同维度后拼接，并通过调制注意力和 MLP 进行计算，使两种模态在注意力操作中相互影响。

校正流匹配

SD3 采用条件流匹配目标进行训练。该方法将前向噪声过程定义为一条将数据分布与噪声分布相连的“校正流”直线。

为支持此功能，Hugging Face 引入了 FlowMatchEulerDiscreteScheduler。该调度器实现了欧拉法步骤，并包含一个 shift 参数用于根据分辨率调整时间步调度。对于 2B 模型，推荐使用 shift=3.0 以更好地处理高分辨率下的噪声缩放。

内存优化技术

由于 SD3 使用了大型 T5‑XXL 文本编码器（4.7B 参数），在显存不足 24GB 的 GPU 上运行模型具有挑战性。Hugging Face 提供了多种优化策略以降低内存占用：

  • 模型卸载（Model Offloading）：使用 pipe.enable_model_cpu_offload() 在模型组件不使用时将其移动到 CPU，从而降低显存占用，但会增加延迟。
  • 去除 T5 编码器：在推理时将 T5‑XXL 编码器设为 None（text_encoder_3=None），可显著降低内存需求，代价是性能略有下降。
  • 8 位量化：通过 bitsandbytes 库以 8 位精度加载 T5‑XXL 模型，进一步降低内存消耗。

内存基准对比

在使用 fp16 精度和 PyTorch 2.3 的 A100 GPU（80GB 显存）上进行的基准测试显示了以下权衡：

技术 推理时间（秒） 内存（GB）
默认 4.762 18.765
卸载 32.765 12.0645
卸载 + 无 T5 19.110 4.266
8 位 T5 4.932 10.586

性能与微调

使用 torch.compile() 加速推理

通过使用 torch.compile() 对 VAE 和 Transformer 组件的计算图进行优化，用户可以实现显著的加速。在单块 80GB A100 上的基准测试中，平均推理时间降至 0.585 秒，相比 eager 执行提升约 4 倍。

DreamBooth 与 LoRA 训练

Hugging Face 已发布针对 SD3 的 DreamBooth 微调脚本，利用低秩适配（LoRA）技术。该脚本支持高效的模型定制，并可作为基于校正流的训练流水线实现参考。

Sources