歡迎使用 aMUSEd：高效文字到影像生成

釋出於 2024 年 1 月 4 日

在 GitHub 上更新

贊

訪客

我們很高興推出一個高效的非擴散文字到影像模型，名為 aMUSEd。之所以這樣命名，是因為它是對 Google's MUSE 的開放復現。aMUSEd 的生成質量並非最佳，我們釋出的是一個帶有寬鬆許可的研究預覽版。

與常用的潛在擴散方法 (Rombach 等人 (2022)) 相比，aMUSEd 採用了掩碼影像模型 (MIM) 方法。正如 Chang 等人 (2023) 所指出的，這不僅需要更少的推理步驟，而且增強了模型的可解釋性。

正如 MUSE 一樣，aMUSEd 在使用單張影像進行風格遷移方面表現出卓越的能力，這一特性在 Sohn 等人 (2023) 的研究中進行了深入探討。這方面可能會為個性化和特定風格的影像生成開闢新途徑。

在這篇博文中，我們將向您介紹 aMUSEd 的一些內部機制，展示如何將其用於不同任務，包括文字到影像，並展示如何對其進行微調。在此過程中，我們將提供與 aMUSEd 相關的所有重要資源，包括其訓練程式碼。讓我們開始吧 🚀

它是如何工作的？

aMUSEd 基於 掩碼影像建模 (Masked Image Modeling)。它為社群探索在影像生成背景下語言建模中已知有效元件提供了一個引人注目的用例。

下圖展示了 aMUSEd 工作原理的示意圖。

在訓練期間

輸入影像使用 VQGAN 進行標記化，以獲得影像標記
然後根據餘弦掩碼排程對影像標記進行掩碼。
掩碼標記（以使用 CLIP-L/14 文字編碼器計算的提示嵌入為條件）被傳遞給 U-ViT 模型，該模型預測掩碼補丁

在推理期間

輸入提示使用 CLIP-L/14 文字編碼器進行嵌入。
迭代直到達到 N 步
- 從隨機掩碼標記開始，並將它們與提示嵌入一起傳遞給 U-ViT 模型
- 預測掩碼標記，並根據 N 和掩碼排程僅保留一定百分比的最自信預測。其餘的進行掩碼，並將其傳遞給 U-ViT 模型
將最終輸出傳遞給 VQGAN 解碼器以獲得最終影像

如前所述，aMUSEd 借鑑了 MUSE 的許多相似之處。但是，也有一些顯著差異

aMUSEd 不遵循預測最終掩碼補丁的兩階段方法。
Instead of using T5 for text conditioning, CLIP L/14 is used for computing the text embeddings. (此句已在原文字中被重複翻譯，故保留英文)
Instead of using T5 for text conditioning, CLIP L/14 is used for computing the text embeddings. (此句已在原文字中被重複翻譯，故保留英文)

要了解更多關於 aMUSEd 的資訊，我們建議閱讀此處的技術報告。

在 🧨 diffusers 中使用 aMUSEd

aMUSEd 已完全整合到 🧨 diffusers 中。要使用它，我們首先需要安裝庫

pip install -U diffusers accelerate transformers -q

讓我們從文字到影像生成開始

import torch
from diffusers import AmusedPipeline

pipe = AmusedPipeline.from_pretrained(
    "amused/amused-512", variant="fp16", torch_dtype=torch.float16
)
pipe = pipe.to("cuda")

prompt = "A mecha robot in a favela in expressionist style"
negative_prompt = "low quality, ugly"

image = pipe(prompt, negative_prompt=negative_prompt, generator=torch.manual_seed(0)).images[0]
image

我們可以研究 num_inference_steps 如何在固定種子下影響影像質量

from diffusers.utils import make_image_grid 

images = []
for step in [5, 10, 15]:
    image = pipe(prompt, negative_prompt=negative_prompt, num_inference_steps=step, generator=torch.manual_seed(0)).images[0]
    images.append(image)

grid = make_image_grid(images, rows=1, cols=3)
grid

至關重要的是，由於其體積小（僅約 800M 引數，包括文字編碼器和 VQ-GAN），aMUSEd 速度非常快。下圖提供了包括 aMUSEd 在內的不同模型的推理延遲比較研究

Speed Comparison — 除模型名稱外，元組的格式如下：（時間步長，解析度）。基準測試在 A100 上進行。更多詳情請參見技術報告。

作為其預訓練目標的直接副產品，aMUSEd 可以零樣本進行影像修復，這與 SDXL 等其他模型不同。

import torch
from diffusers import AmusedInpaintPipeline
from diffusers.utils import load_image
from PIL import Image

pipe = AmusedInpaintPipeline.from_pretrained(
    "amused/amused-512", variant="fp16", torch_dtype=torch.float16
)
pipe = pipe.to("cuda")

prompt = "a man with glasses"
input_image = (
    load_image(
        "https://huggingface.co/amused/amused-512/resolve/main/assets/inpainting_256_orig.png"
    )
    .resize((512, 512))
    .convert("RGB")
)
mask = (
    load_image(
        "https://huggingface.co/amused/amused-512/resolve/main/assets/inpainting_256_mask.png"
    )
    .resize((512, 512))
    .convert("L")
)   

image = pipe(prompt, input_image, mask, generator=torch.manual_seed(3)).images[0]