Stable Diffusion 3 Medium 與 Diffusers 的整合

Hugging Face 已將 Stable Diffusion 3 Medium（Stability AI 的 2B 參數模型）整合至 diffusers 函式庫。此版本讓使用者能以全新架構方法——多模態擴散變換器（Multimodal Diffusion Transformer，MMDiT）以及校正流匹配訓練目標，執行高品質的文字到影像與影像到影像合成。

Stable Diffusion 3 的架構創新

Stable Diffusion 3（SD3）相較於以往的文字到影像架構，採用了文字與影像資料之間的雙向資訊流。

多模態擴散變換器（MMDiT）

SD3 使用新穎的 MMDiT 模型，並搭配三個文字編碼器：CLIP L/14、OpenCLIP bigG/14 與 T5‑v1.1‑XXL。它同時使用類似於 Stable Diffusion XL 的 16 通道 AutoEncoder。

與先前使用固定文字表示的交叉注意力模型不同，MMDiT 區塊將文字輸入與像素潛在向量視為嵌入序列處理。這些序列透過各自的權重嵌入至相同維度後串接，並通過調製的注意力與 MLP。此機制讓兩種模態在注意力運算中相互影響。

校正流匹配

SD3 以條件流匹配目標進行訓練。此方法將前向噪聲過程定義為「校正流」，即以直線連接資料與噪聲分佈。

為支援此方式，Hugging Face 引入了 FlowMatchEulerDiscreteScheduler。該排程器實作 Euler 方法步驟，並包含 shift 參數以因解析度調整時間步長的排程。對於 2B 模型，建議使用 shift=3.0 以更好地處理高解析度下的噪聲縮放。

記憶體最佳化技術

由於 SD3 使用大型 T5‑XXL 文字編碼器（4.7B 參數），在 VRAM 少於 24GB 的 GPU 上執行模型相當具挑戰性。Hugging Face 提供多項最佳化策略以降低記憶體占用：

  • 模型卸載（Model Offloading）：使用 pipe.enable_model_cpu_offload() 可在模型未使用時將其組件移至 CPU，降低 VRAM 使用量，但會增加延遲。
  • 移除 T5 編碼器：在推論時將 T5‑XXL 編碼器設為 None（text_encoder_3=None），可大幅減少記憶體需求，僅會有輕微的效能下降。
  • 8 位元量化：透過 bitsandbytes 套件以 8 位元精度載入 T5‑XXL 模型，進一步降低記憶體消耗。

記憶體基準比較

在搭載 80GB VRAM 的 A100 GPU 上，使用 fp16 精度與 PyTorch 2.3 進行測試，得到以下權衡結果：

Technique Inference Time (secs) Memory (GB)
Default 4.762 18.765
Offloading 32.765 12.0645
Offloading + no T5 19.110 4.266
8bit T5 4.932 10.586

效能與微調

使用 torch.compile() 加速推論

透過 torch.compile() 優化 VAE 與 transformer 元件的計算圖，使用者可獲得顯著的加速。在單張 80GB A100 上的基準測試中，平均推論時間縮短至 0.585 秒，較即時執行提升了 4 倍。

DreamBooth 與 LoRA 訓練

Hugging Face 已發布針對 SD3 的 DreamBooth 微調腳本，利用低秩適應（LoRA）。此腳本可高效地客製化模型，並作為基於校正流的訓練流程實作參考。

Sources