欢迎使用 aMUSEd:高效文本到图像生成

发布于 2024 年 1 月 4 日
在 GitHub 上更新

amused_grid

我们很高兴推出一个高效的非扩散文本到图像模型,名为 aMUSEd。之所以这样命名,是因为它是对 Google's MUSE 的开放复现。aMUSEd 的生成质量并非最佳,我们发布的是一个带有宽松许可的研究预览版。

与常用的潜在扩散方法 (Rombach 等人 (2022)) 相比,aMUSEd 采用了掩码图像模型 (MIM) 方法。正如 Chang 等人 (2023) 所指出的,这不仅需要更少的推理步骤,而且增强了模型的可解释性。

正如 MUSE 一样,aMUSEd 在使用单张图像进行风格迁移方面表现出卓越的能力,这一特性在 Sohn 等人 (2023) 的研究中进行了深入探讨。这方面可能会为个性化和特定风格的图像生成开辟新途径。

在这篇博文中,我们将向您介绍 aMUSEd 的一些内部机制,展示如何将其用于不同任务,包括文本到图像,并展示如何对其进行微调。在此过程中,我们将提供与 aMUSEd 相关的所有重要资源,包括其训练代码。让我们开始吧 🚀

目录

我们为读者制作了一个 aMUSEd 演示,您可以在 此空间 或嵌入在下方的游乐场中试用

它是如何工作的?

aMUSEd 基于 掩码图像建模 (Masked Image Modeling)。它为社区探索在图像生成背景下语言建模中已知有效组件提供了一个引人注目的用例。

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

amused_architecture

训练 期间

  • 输入图像使用 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

text2image_512.png

我们可以研究 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

image_grid_t2i_amused.png

至关重要的是,由于其体积小(仅约 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]

inpainting_grid_amused.png

aMUSEd 是 diffusers 中第一个非扩散系统。其迭代调度方法用于预测掩码补丁,使其成为 diffusers 的一个很好的候选。我们很高兴看到社区如何利用它。

我们鼓励您查阅技术报告,了解我们使用 aMUSEd 探索的所有任务。

微调 aMUSEd

我们提供了一个简单的训练脚本,用于在自定义数据集上微调 aMUSEd。使用 8 位 Adam 优化器和 float16 精度,可以在不到 11GB 的 GPU 显存下微调 aMUSEd。使用 LoRA,内存需求进一步降低到仅 7GB。

Fine-tuned result.
一个戴着红色方框眼镜的像素艺术角色

aMUSEd 带有 OpenRAIL 许可证,因此可以进行商业友好改编。有关微调的更多详细信息,请参阅此目录

局限性

aMUSEd 在图像质量方面并非最先进的图像生成模型。我们发布 aMUSEd 是为了鼓励社区探索非扩散框架,例如 MIM,用于图像生成。我们相信 MIM 的潜力尚未得到充分挖掘,考虑到它的优势:

  • 推理效率
  • 更小的体积,支持设备端应用
  • 无需昂贵的微调即可进行任务迁移
  • 语言建模领域成熟组件的优势

(请注意,MUSE 的原始作品是闭源的)

有关 aMUSEd 定量评估的详细描述,请参阅技术报告。

我们希望社区会发现这些资源有用,并有动力改进用于图像生成的 MIM 现状。

资源

论文:

代码 + 其他:

致谢

Suraj 负责训练。William 负责数据并支持训练。Patrick von Platen 支持训练和数据,并提供总体指导。Robin Rombach 进行了 VQGAN 训练并提供总体指导。Isamu Isozaki 提供了富有洞察力的讨论并贡献了代码。

感谢 Patrick von Platen 和 Pedro Cuenca 对博客文章草稿的审阅。

社区

注册登录 发表评论