介绍Würstchen:用于图像生成的快速扩散

Würstchen Rapid diffusion for image generation

Würstchen 是什么?

Würstchen 是一个扩散模型,其文本条件组件在高度压缩的图像潜空间中工作。为什么这很重要?压缩数据可以将训练和推理的计算成本降低数个数量级。与训练 1024×1024 图像相比,训练 32×32 图像的成本要高得多。通常,其他工作在 4x – 8x 空间压缩的范围内进行相对较小的压缩。Würstchen 将这一点推到了极致。通过其新颖的设计,它实现了 42x 的空间压缩!这在以前是从未见过的,因为常见方法在 16x 空间压缩后无法忠实地重建详细图像。Würstchen 使用了两阶段压缩,我们称之为阶段 A 和阶段 B。阶段 A 是一个 VQGAN,阶段 B 是一个扩散自编码器(更多细节可以在论文中找到)。阶段 A 和 B 一起被称为解码器,因为它们将压缩的图像解码回像素空间。在高度压缩的潜空间中还学习了第三个模型,我们称之为阶段 C。这种训练只需要当前性能最佳模型所使用计算的一小部分,同时还可以实现更便宜和更快速的推理。我们将阶段 C 称为先验模型。

为什么需要另一个文本到图像模型?

嗯,这个模型非常快速和高效。Würstchen 的最大优点在于它可以比 Stable Diffusion XL 等模型更快地生成图像,同时使用的内存要少得多!因此,对于我们这些没有 A100 的人来说,这非常方便。以下是与不同批次大小的 SDXL 进行比较的结果:

除此之外,Würstchen 的另一个非常重要的优点是降低了训练成本。Würstchen v1 在 512×512 的分辨率下仅需要 9,000 个 GPU 小时的训练。与 Stable Diffusion 1.4 花费的 150,000 个 GPU 小时相比,这意味着成本减少了 16 倍,不仅有利于研究人员进行新实验,而且也为更多组织培训这样的模型打开了大门。Würstchen v2 使用了 24,602 个 GPU 小时。随着分辨率提高到 1536,这仍然比 SD1.4 便宜 6 倍,后者仅在 512×512 上进行了训练。

您还可以在此处找到详细的解释视频:

如何使用 Würstchen?

您可以通过此处的演示尝试它:

否则,该模型可通过 Diffusers Library 获得,因此您可以使用您已熟悉的界面。例如,以下是如何使用 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 适用于哪些图像尺寸?

Würstchen 是在 1024×1024 和 1536×1536 之间的图像分辨率上进行训练的。我们有时还观察到在 1024×2048 等分辨率下有良好的输出。请随意尝试。我们还观察到先验模型(阶段 C)对新分辨率的适应非常快速。因此,在 2048×2048 上进行微调应该计算成本较低。

Hub 上的模型

所有检查点也可以在 Huggingface Hub 上查看。在那里可以找到多个检查点,以及未来的演示和模型权重。目前,Prior 有 3 个检查点可用,Decoder 有 1 个检查点可用。请查看文档,了解检查点的详细说明以及不同的 Prior 模型以及可以用于什么。

Diffusers 集成

因为 Würstchen 完全集成在 diffusers 中,所以它自动带有各种好处和优化。这些包括:

  • 自动使用 PyTorch 2 的 SDPA 加速注意力,如下所述。
  • 支持 xFormers 的闪电注意力实现,如果您需要使用 PyTorch 1.x 而不是 2。
  • 模型卸载,将未使用的组件移动到 CPU 上,以节省内存,并且性能几乎没有影响。
  • 顺序 CPU 卸载,用于内存非常宝贵的情况。内存使用将最小化,但推理速度较慢。
  • 使用 Compel 库进行提示加权。
  • 支持 Apple Silicon Mac 上的 mps 设备。
  • 使用生成器实现可重现性。
  • 推理的合理默认值,以在大多数情况下产生高质量的结果。当然,您可以随意调整所有参数!

优化技术 1:闪电注意力

从版本 2.0 开始,PyTorch 集成了一种高度优化和资源友好的注意力机制,称为 torch.nn.functional.scaled_dot_product_attention 或 SDPA。根据输入的性质,此函数利用多种底层优化。其性能和内存效率超过传统的注意力模型。值得注意的是,SDPA 函数反映了闪电注意力技术的特点,正如 Dao 和团队撰写的研究论文《具有 IO-Awareness 的快速和内存高效的精确注意力》中所强调的那样。

如果您正在使用 PyTorch 2.0 或更新版本的 Diffusers,并且可以访问 SDPA 函数,这些增强功能将自动应用。按照官方指南设置 torch 2.0 或更高版本,开始使用吧!

images = pipeline(caption, height=1024, width=1536, prior_timesteps=DEFAULT_STAGE_C_TIMESTEPS, prior_guidance_scale=4.0, num_images_per_prompt=4).images

要深入了解 diffusers 如何利用 SDPA,请查阅文档。

如果您使用的是早于 2.0 版本的 PyTorch,仍然可以使用 xFormers 库实现内存高效的注意力:

pipeline.enable_xformers_memory_efficient_attention()

优化技术 2:Torch 编译

如果您正在寻找额外的性能提升,可以使用 torch.compile。最好将其应用于 Prior 和 Decoder 的主模型,以获得最大的性能增益。

pipeline.prior_prior = torch.compile(pipeline.prior_prior , mode="reduce-overhead", fullgraph=True)
pipeline.decoder = torch.compile(pipeline.decoder, mode="reduce-overhead", fullgraph=True)

请注意,初始推理步骤将花费很长时间(最多 2 分钟),因为模型正在被编译。之后,您可以正常运行推理:

images = pipeline(caption, height=1024, width=1536, prior_timesteps=DEFAULT_STAGE_C_TIMESTEPS, prior_guidance_scale=4.0, num_images_per_prompt=4).images

好消息是,这种编译是一次性的执行。之后,对于相同的图像分辨率,您将始终体验到更快的推理速度。编译的初始时间投资很快就会被后续的速度优势所抵消。要详细了解 torch.compile 及其细微差别,请查阅官方文档。

资源

  • 有关此模型的更多信息,请参阅官方 diffusers 文档。
  • 所有检查点都可以在 Hub 上找到。
  • 您可以在此处尝试演示。
  • 如果您想讨论未来的项目,甚至为自己的想法做出贡献,请加入我们的 Discord!