Würstchen
Wuerstchen: 大规模文本到图像扩散模型的高效架构 由 Pablo Pernias、Dominic Rampas、Mats L. Richter 和 Christopher Pal 以及 Marc Aubreville 共同撰写。
论文的摘要如下:
我们介绍了 Würstchen,这是一种用于文本到图像合成的新型架构,结合了竞争性的性能和前所未有的成本效益,适用于大规模文本到图像扩散模型。我们工作的关键贡献之一是开发了一种潜在扩散技术,在这种技术中,我们学习了一种详细但极其紧凑的语义图像表示,用于指导扩散过程。与语言的潜在表示相比,这种高度压缩的图像表示提供了更详细的指导,这显著减少了实现最先进结果所需的计算需求。我们的方法还基于用户偏好研究提高了文本条件图像生成的质量。我们的方法的训练需求包括 24,602 个 A100-GPU 小时——相比之下,Stable Diffusion 2.1 需要 200,000 个 GPU 小时。我们的方法还减少了实现这些结果所需的训练数据。此外,我们紧凑的潜在表示使我们能够以两倍的速度进行推理,显著降低了通常的成本和最先进(SOTA)扩散模型的碳足迹,而不会影响最终性能。在与 SOTA 模型的更广泛比较中,我们的方法在效率上显著更高,并且在图像质量方面表现出色。我们相信这项工作强调了在性能和计算可访问性之间优先考虑的重要性。
Würstchen 概述
Würstchen 是一种扩散模型,其文本条件模型在高度压缩的图像潜在空间中工作。为什么这很重要?压缩数据可以大幅降低训练和推理的计算成本。在 1024x1024 图像上训练比在 32x32 图像上训练要昂贵得多。通常,其他工作使用相对较小的压缩,空间压缩范围在 4x - 8x 之间。Würstchen 将这一极限推向了极致。通过其新颖的设计,我们实现了 42x 的空间压缩。这在以前是未曾见过的,因为常见的方法在 16x 空间压缩后无法忠实地重建详细图像。Würstchen 采用两阶段压缩,我们称之为阶段 A 和阶段 B。阶段 A 是一个 VQGAN,阶段 B 是一个扩散自编码器(更多细节可以在论文中找到)。在高度压缩的潜在空间中学习了第三个模型,阶段 C。这种训练所需的计算量仅为当前顶级模型的一小部分,同时允许更便宜和更快的推理。
Würstchen v2 来到 Diffusers
在最初的论文发布后,我们在架构、训练和采样方面进行了多项改进,使 Würstchen 在许多方面与当前最先进的模型具有竞争力。我们很高兴与 Diffusers 一起发布这个新版本。以下是改进的列表。
- 更高的分辨率(1024x1024 至 2048x2048)
- 更快的推理
- 多方面分辨率采样
- 更好的质量
我们为文本条件图像生成模型(阶段 C)发布了 3 个检查点。这些是:
- v2-base
- v2-aesthetic
- (默认) v2-interpolated(v2-base 和 v2-aesthetic 之间的 50% 插值)
我们推荐使用 v2-interpolated,因为它在写实和美学之间取得了很好的平衡。使用 v2-base 进行微调,因为它没有风格偏见,并使用 v2-aesthetic 进行非常艺术化的生成。比较可以在这里看到:
文本到图像生成
为了便于使用,Würstchen 可以使用单一管道。该管道可以按如下方式使用:
import torch
from diffusers import AutoPipelineForText2Image
from diffusers.pipelines.wuerstchen import DEFAULT_STAGE_C_TIMESTEPS
pipe = AutoPipelineForText2Image.from_pretrained("warp-ai/wuerstchen", torch_dtype=torch.float16).to("cuda")
caption = "Anthropomorphic cat dressed as a fire fighter"
images = pipe(
caption,
width=1024,
height=1536,
prior_timesteps=DEFAULT_STAGE_C_TIMESTEPS,
prior_guidance_scale=4.0,
num_images_per_prompt=2,
).images
为了解释的目的,我们也可以分别初始化 Würstchen 的两个主要管道。Würstchen 由三个阶段组成:阶段 C、阶段 B、阶段 A。它们都有不同的任务,并且只有在一起工作时才能发挥作用。在生成文本条件图像时,阶段 C 首先会在一个非常压缩的潜在空间中生成潜在变量。这就是在 prior_pipeline
中发生的事情。之后,生成的潜在变量将被传递到阶段 B,阶段 B 将潜在变量解压缩到一个更大的 VQGAN 潜在空间中。然后,这些潜在变量可以被阶段 A(一个 VQGAN)解码到像素空间中。阶段 B 和阶段 A 都被封装在 decoder_pipeline
中。更多详情,请参阅 论文。
import torch
from diffusers import WuerstchenDecoderPipeline, WuerstchenPriorPipeline
from diffusers.pipelines.wuerstchen import DEFAULT_STAGE_C_TIMESTEPS
device = "cuda"
dtype = torch.float16
num_images_per_prompt = 2
prior_pipeline = WuerstchenPriorPipeline.from_pretrained(
"warp-ai/wuerstchen-prior", torch_dtype=dtype
).to(device)
decoder_pipeline = WuerstchenDecoderPipeline.from_pretrained(
"warp-ai/wuerstchen", torch_dtype=dtype
).to(device)
caption = "Anthropomorphic cat dressed as a fire fighter"
negative_prompt = ""
prior_output = prior_pipeline(
prompt=caption,
height=1024,
width=1536,
timesteps=DEFAULT_STAGE_C_TIMESTEPS,
negative_prompt=negative_prompt,
guidance_scale=4.0,
num_images_per_prompt=num_images_per_prompt,
)
decoder_output = decoder_pipeline(
image_embeddings=prior_output.image_embeddings,
prompt=caption,
negative_prompt=negative_prompt,
guidance_scale=0.0,
output_type="pil",
).images[0]
decoder_output
加速推理
你可以利用 torch.compile
函数,获得大约 2-3 倍的加速:
prior_pipeline.prior = torch.compile(prior_pipeline.prior, mode="reduce-overhead", fullgraph=True)
decoder_pipeline.decoder = torch.compile(decoder_pipeline.decoder, mode="reduce-overhead", fullgraph=True)
限制
- 由于Würstchen采用了高压缩技术,生成的图像可能会缺乏大量细节。对于人类的眼睛来说,这在面部、手部等部位尤为明显。
- 图像只能以128像素的步长生成,例如,在1024x1024之后,下一个更高的分辨率是1152x1152。
- 该模型缺乏在图像中正确渲染文本的能力。
- 该模型通常无法实现照片级真实感。
- 复杂的组合提示对模型来说较为困难。
原始代码库以及实验性想法可以在dome272/Wuerstchen找到。
WuerstchenCombinedPipeline
[[autodoc]] WuerstchenCombinedPipeline - all - call
WuerstchenPriorPipeline
[[autodoc]] WuerstchenPriorPipeline - all - call
WuerstchenPriorPipelineOutput
[[autodoc]] pipelines.wuerstchen.pipeline_wuerstchen_prior.WuerstchenPriorPipelineOutput
WuerstchenDecoderPipeline
[[autodoc]] WuerstchenDecoderPipeline - all - call
引用
@misc{pernias2023wuerstchen,
title={Wuerstchen: An Efficient Architecture for Large-Scale Text-to-Image Diffusion Models},
author={Pablo Pernias and Dominic Rampas and Mats L. Richter and Christopher J. Pal and Marc Aubreville},
year={2023},
eprint={2306.00637},
archivePrefix={arXiv},
primaryClass={cs.CV}
}