aMUSEd：高效文字轉圖像生成

aMUSEd：高效文字轉圖像生成

Hugging Face 已發布 aMUSEd，一款基於遮罩影像建模（MIM）的高效非擴散文字轉圖像模型，並且是 Google MUSE 的開源再現。

aMUSEd：高效文字轉圖像生成

Hugging Face 推出了 aMUSEd，一款高效的非擴散文字轉圖像模型。作為 Google MUSE 的開源再現，aMUSEd 利用遮罩影像建模（MIM）提供比傳統潛在擴散模型更快速且更具可解釋性的替代方案。

遮罩影像建模架構

aMUSEd 採用遮罩影像建模（MIM）方法，與潛在擴散相比，可減少推理步驟並提升可解釋性。模型架構由三個主要組件構成：用於分詞的 VQGAN、用於提示嵌入的 CLIP-L/14 文字編碼器，以及用於預測遮罩區塊的 U‑ViT 模型。

訓練流程

在訓練過程中，模型遵循以下步驟：

  1. Tokenization：使用 VQGAN 將輸入圖像轉換為圖像 token。
  2. Masking：根據餘弦遮罩排程對圖像 token 進行遮罩。
  3. Prediction：將帶有來自 CLIP-L/14 文字編碼器的提示嵌入條件的遮罩 token 傳入 U‑ViT 模型，以預測遮罩區塊。

推理流程

在推理過程中，模型透過迭代方式生成圖像：

  1. Embedding：使用 CLIP-L/14 文字編碼器對輸入提示進行嵌入。
  2. Iterative Prediction：過程從隨機遮罩的 token 開始。U‑ViT 模型預測遮罩 token，僅保留最有信心的預測的一部分（由遮罩排程與步數 N 決定）。其餘 token 重新遮罩並再次傳入 U‑ViT 模型。
  3. Decoding：最終輸出經過 VQGAN 解碼器產生最終圖像。

與 MUSE 的主要差異

aMUSEd 與原始 MUSE 模型在三個關鍵方面有所不同：

  • Single-Stage Prediction：不使用兩階段的最終遮罩區塊預測方法。
  • Text Conditioning：使用 CLIP L/14 取代 T5 來計算文字嵌入。
  • Micro-conditioning：遵循 SDXL，加入圖像尺寸與裁剪等額外條件。

能力與效能

aMUSEd 旨在提升效率與速度。模型約有 800M 參數（包括文字編碼器與 VQ‑GAN），具備極佳的效能。

推理速度與延遲

在 A100 GPU 上的基準測試顯示，aMUSEd 的推理延遲顯著低於其他模型。此效率使其成為適合於裝置端應用的候選模型。

零樣本圖像修補

由於其預訓練目標，aMUSEd 能夠執行零樣本圖像修補，這是如 SDXL 等模型所不具備的功能。

風格轉換

繼承 MUSE 的能力，aMUSEd 展示了使用單張圖像進行風格轉換的能力，為個人化與特定風格的圖像生成提供了潛力。

微調與可及性

aMUSEd 以 OpenRAIL 授權釋出，具商業友好性。Hugging Face 提供了用於自訂資料集微調的訓練腳本。

微調的硬體需求

  • Standard Fine-Tuning：使用 8 位元 Adam 優化器與 float16 精度，微調所需的 GPU 記憶體略低於 11GB。
  • LoRA：使用低秩適應（LoRA）時，記憶體需求進一步降低至 7GB VRAM。

限制

雖然在原始圖像品質方面 aMUSEd 並非最先進，但其發布旨在鼓勵社群探索 MIM 在圖像生成上的潛力。作者強調了 MIM 框架的多項優勢，包括推理效率、適合裝置端使用的較小模型尺寸，以及在不需昂貴微調的情況下執行任務轉移的能力。

Sources