Würstchen：引领图像生成的速度革新

作者

大家好，今天我要给大家介绍一款名为 Würstchen 的新型图像生成模型。这款模型由 Hugging Face 公司开发，它是一种扩散模型，能够在高度压缩的图像潜在空间中进行文本条件操作。这项技术的优势在于能大大降低训练和推理的计算成本。以前我们在处理 1024×1024 的图像时，需要花费大量的计算资源，而现在，通过 Würstchen，我们可以用相当于 32×32 图像的资源来完成这项工作，这无疑是一个巨大的突破。

友情链接：ACEJoy

Würstchen 的设计新颖，实现了 42 倍的空间压缩，这在之前是无法想象的。它采用了两阶段压缩，我们称之为A阶段和B阶段。A阶段是一个 VQGAN，B阶段是一个扩散自编码器。A阶段和B阶段共同被称为解码器，因为它们将压缩后的图像解码回像素空间。还有一个第三阶段模型，称为 Prior，它在高度压缩的潜在空间中进行学习，这种训练需要的计算资源只是当前顶级模型的一小部分，同时也让推理变得更便宜、更快。

那么，为什么我们需要另一个文本到图像的模型呢？原因很简单，因为 Würstchen 非常快且高效。比起像 Stable Diffusion XL 这样的模型，Würstchen 可以更快地生成图像，同时使用的内存也更少。此外，Würstchen 的训练成本也大大降低，Würstchen v1 只需要 9,000 GPU 训练小时，而 Stable Diffusion 1.4 需要 150,000 GPU 训练小时。这意味着更多的组织可以训练这样的模型。

使用 Würstchen 也非常简单。您可以通过 Diffusers 库来使用它，下面是一个使用 AutoPipeline 进行推理的例子：

import torch
from diffusers import AutoPipelineForText2Image
from diffusers.pipelines.wuerstchen import DEFAULT_STAGE_C_TIMESTEPS

pipeline = AutoPipelineForText2Image.from_pretrained("warp-ai/wuerstchen", torch_dtype=torch.float16).to("cuda")

caption = "Anthropomorphic cat dressed as a firefighter"
images = pipeline(
    caption,
    height=1024,
    width=1536,
    prior_timesteps=DEFAULT_STAGE_C_TIMESTEPS,
    prior_guidance_scale=4.0,
    num_images_per_prompt=4,
).images

Würstchen 在 1024×1024 到 1536×1536 之间的图像分辨率上进行训练，同时我们也注意到，在 1024×2048 这样的分辨率上，Würstchen 也能产生很好的结果。我们还发现，Prior（第三阶段模型）能够非常快速地适应新的分辨率，因此在 2048×2048 的分辨率上进行微调应该是计算成本非常低的。

此外，Würstchen 还提供了一些优化技术，包括使用 PyTorch 2 SDPA 加速注意力机制，为 Apple Silicon Mac 提供的 mps 设备支持，以及使用生成器进行可复现性优化等。

在此，我们要感谢 Stability AI 提供的计算资源，使得我们能够训练出这样的模型，我们希望这项研究能够对更多的研究者和开发者带来帮助，提升他们在图像生成领域的工作效率。

Würstchen 的所有检查点都可以在 Huggingface Hub 上查看。在那里，您可以找到多个检查点，以及未来的演示和模型权重。目前，Prior有三个检查点，Decoder有一个检查点。请在文档中查看检查点的解释，以及不同的 Prior 模型的用途。

因为 Würstchen 完全集成在 Diffusers 中，所以它自带各种好处和优化。其中包括：

  • 自动使用 PyTorch 2 SDPA 加速的注意力机制。
  • 支持 xFormers 的 flash 注意力实现，如果你需要使用 PyTorch 1.x 而不是 2.x。
  • 模型卸载，将未使用的组件在不使用时移动到 CPU，这可以节省内存，而对性能影响微乎其微。
  • 顺序 CPU 卸载，适用于内存非常宝贵的情况。内存使用将被最小化，代价是推理速度较慢。
  • 使用 Comp 进行提示权重。

在这里，我们的介绍就要结束了，但 Würstchen 的旅程才刚刚开始。这个新模型将在图片生成领域带来革命性影响，无论是在计算成本、效率还是在图像质量上，都有显著的提升。我们期待它在未来能够带给更多研究人员和开发者便利和启发。感谢大家的阅读，我们下次再见！

发表评论