介紹 Würstchen：用於圖像生成的快速擴散

TL;DR

Hugging Face 已經推出 Würstchen，一種為極致效率而設計的文本到圖像擴散模型。透過利用新穎的兩階段壓縮架構，Würstchen 實現了 42x 空間壓縮，使其能夠以更快的速度生成高解析度圖像，且所需的記憶體和訓練計算資源顯著低於先前的模型，如 Stable Diffusion。

高壓縮架構

Würstchen 透過在高度壓縮的潛在空間中運行來降低訓練和推理的計算成本。雖然典型的擴散模型使用 4x 到 8x 範圍的空間壓縮，但 Würstchen 採用 42x 空間壓縮，以在不犧牲圖像細節的情況下，在更小的潛在表示上進行訓練。

此架構由三個主要階段組成：

  • 解碼器（階段 A 和 B）: 此組件將壓縮圖像解碼回像素空間。它由階段 A（一個 VQGAN）和階段 B（一個擴散自編碼器）組成。
  • 先驅（階段 C）: 此模型在高度壓縮的潛在空間中學習。它處理圖像生成過程中的文本條件部分。

效能與效率提升

與 Stable Diffusion XL (SDXL) 等模型相比，Würstchen 在推理速度和記憶體使用方面提供了顯著的改進。它專門設計為無需高端硬體（如 A100 GPU）的使用者也能輕鬆使用。

訓練計算減少

該模型展示了訓練所需 GPU 小時的顯著減少：

  • Würstchen v1 (512x512): 需要 9,000 GPU 小時，與 Stable Diffusion 1.4 所花費的 150,000 GPU 小時相比，成本降低了 16x。
  • Würstchen v2 (最高 1536 解析度): 需要 24,602 GPU 小時，儘管在更高解析度下訓練，但相比 SD1.4 大約便宜 6x。

技術實作與使用

Würstchen 完全整合到 diffusers 庫中，可透過 AutoPipelineForText2Image 介面使用。該模型在 1024x1024 到 1536x1536 的解析度間進行訓練，儘管它可以在 1024x2048 產出高品質輸出，並且可以以低成本微調至 2048x2048 解析度。

整合優化

透過其 diffusers 整合，Würstchen 支援多種 out-of-the-box 優化：

  • 記憶體管理: 模型卸載和順序 CPU 卸載可最小化 VRAM 使用。
  • 硬體支援: 支援 Apple Silicon Mac 上的 mps 裝置。
  • 提示詞: 透過 Compel 庫進行提示詞加權。
  • 可重複性: 使用生成器以獲得一致的結果。

進階效能調整

使用者可以透過兩種主要技術進一步加速 Würstchen：

  1. 快速注意力: 模型自動利用 PyTorch 2.0 的 scaled_dot_product_attention (SDPA) 來進行記憶體高效的注意力。使用 PyTorch 1.x 的使用者可以透過 pipeline.enable_xformers_memory_efficient_attention() 使用 xFormers 庫。
  2. Torch 編譯: 對先驅和解碼器模型應用 torch.compile（使用 mode="reduce-overhead" 和 fullgraph=True）在初始編譯期間（最多兩分鐘）後提供額外的效能提升。

Sources