Wuerstchen

🚀 Würstchen - 概述

Würstchen是一个扩散模型，其文本条件模型在高度压缩的图像潜在空间中工作。为什么这很重要呢？数据压缩可以将训练和推理的计算成本大幅降低。在1024x1024的图像上进行训练，要比在32x32的图像上训练昂贵得多。通常，其他工作使用的压缩率相对较小，空间压缩范围在4倍 - 8倍。而Würstchen将这一点发挥到了极致。通过其新颖的设计，我们实现了42倍的空间压缩。这在以前是从未见过的，因为常见的方法在16倍空间压缩后无法忠实地重建出细节丰富的图像。Würstchen采用了两阶段压缩，我们称之为A阶段和B阶段。A阶段是一个VQGAN，B阶段是一个扩散自编码器（更多细节可在论文中找到）。第三个模型，即C阶段，是在这个高度压缩的潜在空间中学习的。这种训练所需的计算量只是当前顶级模型的一小部分，同时也使得推理更便宜、更快速。

✨ 主要特性

Würstchen - 解码器

解码器指的是“A阶段”和“B阶段”。解码器接收图像嵌入（可以是由先验模型（C阶段）生成的，也可以是从真实图像中提取的），并将这些潜在特征解码回像素空间。具体来说，B阶段首先将图像嵌入解码到VQGAN空间，然后A阶段（即VQGAN）将潜在特征解码到像素空间。它们共同实现了42倍的空间压缩。

注意：这种重建是有损的，会丢失图像的部分信息。目前的B阶段在重建时往往缺乏细节，当我们观察人脸、手部等部位时，这些细节缺失尤其明显。我们正在努力在未来改进这些重建效果！

图像尺寸

Würstchen是在1024x1024到1536x1536的图像分辨率上进行训练的。我们有时也会在1024x2048等分辨率上观察到不错的输出效果。你可以随意尝试。我们还发现，先验模型（C阶段）能非常快速地适应新的分辨率。因此，在2048x2048的分辨率上对其进行微调，计算成本应该较低。

📦 安装指南

此管道应与先验模型https://huggingface.co/warp-ai/wuerstchen-prior 一起运行。

import torch
from diffusers import AutoPipelineForText2Image

device = "cuda"
dtype = torch.float16

pipeline =  AutoPipelineForText2Image.from_pretrained(
    "warp-diffusion/wuerstchen", torch_dtype=dtype
).to(device)

caption = "Anthropomorphic cat dressed as a fire fighter"

output = pipeline(
    prompt=caption,
    height=1024,
    width=1024,
    prior_guidance_scale=4.0,
    decoder_guidance_scale=0.0,
).images

图像采样时间

该图展示了在Würstchen上不同批量大小（num_images_per_prompt）的推理时间（在A100上），并与Stable Diffusion XL（无精炼器）进行了比较。左图显示了推理时间（使用torch > 2.0），而右图则提前对两个管道应用了torch.compile。

📚 详细文档

模型详情

@inproceedings{
      pernias2024wrstchen,
      title={W"urstchen: An Efficient Architecture for Large - Scale Text - to - Image Diffusion Models},
      author={Pablo Pernias and Dominic Rampas and Mats Leon Richter and Christopher Pal and Marc Aubreville},
      booktitle={The Twelfth International Conference on Learning Representations},
      year={2024},
      url={https://openreview.net/forum?id=gU58d5QeGv}
}

环境影响

Würstchen v2 估计排放量 基于这些信息，我们使用Lacoste等人（2019）提出的机器学习影响计算器来估计以下二氧化碳排放量。我们利用硬件、运行时间、云服务提供商和计算区域来估计碳影响。

• 硬件类型：A100 PCIe 40GB
• 使用时长：24602小时
• 云服务提供商：AWS
• 计算区域：美国东部
• 碳排放（功耗 x 时间 x 基于电网位置产生的碳排放量）：2275.68千克二氧化碳当量。

📄 许可证

本项目采用MIT许可证。