大家好,今天我要给大家介绍一款名为 Würstchen 的新型图像生成模型。这款模型由 Hugging Face 公司开发,它是一种扩散模型,能够在高度压缩的图像潜在空间中进行文本条件操作。这项技术的优势在于能大大降低训练和推理的计算成本。以前我们在处理 1024×1024 的图像时,需要花费大量的计算资源,而现在,通过 Würstchen,我们可以用相当于 32×32 图像的资源来完成这项工作,这无疑是一个巨大的突破。
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 在 1024x1024 到 1536x1536 之间的图像分辨率上进行训练,同时我们也注意到,在 1024x2048 这样的分辨率上,Würstchen 也能产生很好的结果。我们还发现,Prior(第三阶段模型)能够非常快速地适应新的分辨率,因此在 2048x2048 的分辨率上进行微调应该是计算成本非常低的。
此外,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 的旅程才刚刚开始。这个新模型将在图片生成领域带来革命性影响,无论是在计算成本、效率还是在图像质量上,都有显著的提升。我们期待它在未来能够带给更多研究人员和开发者便利和启发。感谢大家的阅读,我们下次再见!