wuerstchen.md 8.82 KB
Newer Older
Aryan's avatar
Aryan committed
1
<!--Copyright 2025 The HuggingFace Team. All rights reserved.
2
3
4
5
6
7
8
9
10
11
12

Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
the License. You may obtain a copy of the License at

http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
specific language governing permissions and limitations under the License.
-->

Kashif Rasul's avatar
Kashif Rasul committed
13
14
# Würstchen

15
16
17
> [!WARNING]
> This pipeline is deprecated but it can still be used. However, we won't test the pipeline anymore and won't accept any changes to it. If you run into any issues, reinstall the last Diffusers version that supported this model.

Steven Liu's avatar
Steven Liu committed
18
19
20
21
<div class="flex flex-wrap space-x-1">
  <img alt="LoRA" src="https://img.shields.io/badge/LoRA-d8b4fe?style=flat"/>
</div>

Kashif Rasul's avatar
Kashif Rasul committed
22
23
<img src="https://github.com/dome272/Wuerstchen/assets/61938694/0617c863-165a-43ee-9303-2a17299a0cf9">

24
[Wuerstchen: An Efficient Architecture for Large-Scale Text-to-Image Diffusion Models](https://huggingface.co/papers/2306.00637) is by Pablo Pernias, Dominic Rampas, Mats L. Richter and Christopher Pal and Marc Aubreville.
Kashif Rasul's avatar
Kashif Rasul committed
25
26
27

The abstract from the paper is:

28
*We introduce Würstchen, a novel architecture for text-to-image synthesis that combines competitive performance with unprecedented cost-effectiveness for large-scale text-to-image diffusion models. A key contribution of our work is to develop a latent diffusion technique in which we learn a detailed but extremely compact semantic image representation used to guide the diffusion process. This highly compressed representation of an image provides much more detailed guidance compared to latent representations of language and this significantly reduces the computational requirements to achieve state-of-the-art results. Our approach also improves the quality of text-conditioned image generation based on our user preference study. The training requirements of our approach consists of 24,602 A100-GPU hours - compared to Stable Diffusion 2.1's 200,000 GPU hours. Our approach also requires less training data to achieve these results. Furthermore, our compact latent representations allows us to perform inference over twice as fast, slashing the usual costs and carbon footprint of a state-of-the-art (SOTA) diffusion model significantly, without compromising the end performance. In a broader comparison against SOTA models our approach is substantially more efficient and compares favorably in terms of image quality. We believe that this work motivates more emphasis on the prioritization of both performance and computational accessibility.*
Kashif Rasul's avatar
Kashif Rasul committed
29

30
## Würstchen Overview
31
Würstchen is a diffusion model, whose text-conditional model works in a highly compressed latent space of images. Why is this important? Compressing data can reduce computational costs for both training and inference by magnitudes. Training on 1024x1024 images is way more expensive than training on 32x32. Usually, other works make use of a relatively small compression, in the range of 4x - 8x spatial compression. Würstchen takes this to an extreme. Through its novel design, we achieve a 42x spatial compression. This was unseen before because common methods fail to faithfully reconstruct detailed images after 16x spatial compression. Würstchen employs a two-stage compression, what we call Stage A and Stage B. Stage A is a VQGAN, and Stage B is a Diffusion Autoencoder (more details can be found in the [paper](https://huggingface.co/papers/2306.00637)). A third model, Stage C, is learned in that highly compressed latent space. This training requires fractions of the compute used for current top-performing models, while also allowing cheaper and faster inference.
32

Kashif Rasul's avatar
Kashif Rasul committed
33
34
## Würstchen v2 comes to Diffusers

35
After the initial paper release, we have improved numerous things in the architecture, training and sampling, making Würstchen competitive to current state-of-the-art models in many ways. We are excited to release this new version together with Diffusers. Here is a list of the improvements.
Kashif Rasul's avatar
Kashif Rasul committed
36
37
38
39
40
41

- Higher resolution (1024x1024 up to 2048x2048)
- Faster inference
- Multi Aspect Resolution Sampling
- Better quality

Kashif Rasul's avatar
Kashif Rasul committed
42

43
We are releasing 3 checkpoints for the text-conditional image generation model (Stage C). Those are:
44

Kashif Rasul's avatar
Kashif Rasul committed
45
46
- v2-base
- v2-aesthetic
47
- **(default)** v2-interpolated (50% interpolation between v2-base and v2-aesthetic)
Kashif Rasul's avatar
Kashif Rasul committed
48

49
We recommend using v2-interpolated, as it has a nice touch of both photorealism and aesthetics. Use v2-base for finetunings as it does not have a style bias and use v2-aesthetic for very artistic generations.
Kashif Rasul's avatar
Kashif Rasul committed
50
A comparison can be seen here:
Kashif Rasul's avatar
Kashif Rasul committed
51
52
53
54
55

<img src="https://github.com/dome272/Wuerstchen/assets/61938694/2914830f-cbd3-461c-be64-d50734f4b49d" width=500>

## Text-to-Image Generation

56
For the sake of usability, Würstchen can be used with a single pipeline. This pipeline can be used as follows:
Kashif Rasul's avatar
Kashif Rasul committed
57
58
59
60

```python
import torch
from diffusers import AutoPipelineForText2Image
Kashif Rasul's avatar
Kashif Rasul committed
61
from diffusers.pipelines.wuerstchen import DEFAULT_STAGE_C_TIMESTEPS
Kashif Rasul's avatar
Kashif Rasul committed
62

Kashif Rasul's avatar
Kashif Rasul committed
63
pipe = AutoPipelineForText2Image.from_pretrained("warp-ai/wuerstchen", torch_dtype=torch.float16).to("cuda")
Kashif Rasul's avatar
Kashif Rasul committed
64
65

caption = "Anthropomorphic cat dressed as a fire fighter"
Kashif Rasul's avatar
Kashif Rasul committed
66
images = pipe(
67
    caption,
Kashif Rasul's avatar
Kashif Rasul committed
68
    width=1024,
Kashif Rasul's avatar
Kashif Rasul committed
69
70
    height=1536,
    prior_timesteps=DEFAULT_STAGE_C_TIMESTEPS,
Kashif Rasul's avatar
Kashif Rasul committed
71
    prior_guidance_scale=4.0,
Kashif Rasul's avatar
Kashif Rasul committed
72
    num_images_per_prompt=2,
Kashif Rasul's avatar
Kashif Rasul committed
73
74
75
).images
```

76
For explanation purposes, we can also initialize the two main pipelines of Würstchen individually. Würstchen consists of 3 stages: Stage C, Stage B, Stage A. They all have different jobs and work only together. When generating text-conditional images, Stage C will first generate the latents in a very compressed latent space. This is what happens in the `prior_pipeline`. Afterwards, the generated latents will be passed to Stage B, which decompresses the latents into a bigger latent space of a VQGAN. These latents can then be decoded by Stage A, which is a VQGAN, into the pixel-space. Stage B & Stage A are both encapsulated in the `decoder_pipeline`. For more details, take a look at the [paper](https://huggingface.co/papers/2306.00637).
Kashif Rasul's avatar
Kashif Rasul committed
77
78
79
80

```python
import torch
from diffusers import WuerstchenDecoderPipeline, WuerstchenPriorPipeline
Kashif Rasul's avatar
Kashif Rasul committed
81
from diffusers.pipelines.wuerstchen import DEFAULT_STAGE_C_TIMESTEPS
Kashif Rasul's avatar
Kashif Rasul committed
82
83
84
85
86
87

device = "cuda"
dtype = torch.float16
num_images_per_prompt = 2

prior_pipeline = WuerstchenPriorPipeline.from_pretrained(
Kashif Rasul's avatar
Kashif Rasul committed
88
    "warp-ai/wuerstchen-prior", torch_dtype=dtype
Kashif Rasul's avatar
Kashif Rasul committed
89
90
).to(device)
decoder_pipeline = WuerstchenDecoderPipeline.from_pretrained(
Kashif Rasul's avatar
Kashif Rasul committed
91
    "warp-ai/wuerstchen", torch_dtype=dtype
Kashif Rasul's avatar
Kashif Rasul committed
92
93
).to(device)

Kashif Rasul's avatar
Kashif Rasul committed
94
caption = "Anthropomorphic cat dressed as a fire fighter"
Kashif Rasul's avatar
Kashif Rasul committed
95
96
97
98
99
negative_prompt = ""

prior_output = prior_pipeline(
    prompt=caption,
    height=1024,
Kashif Rasul's avatar
Kashif Rasul committed
100
101
    width=1536,
    timesteps=DEFAULT_STAGE_C_TIMESTEPS,
Kashif Rasul's avatar
Kashif Rasul committed
102
    negative_prompt=negative_prompt,
Kashif Rasul's avatar
Kashif Rasul committed
103
    guidance_scale=4.0,
Kashif Rasul's avatar
Kashif Rasul committed
104
105
106
107
108
109
110
111
    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",
112
113
).images[0]
decoder_output
Kashif Rasul's avatar
Kashif Rasul committed
114
115
116
```

## Speed-Up Inference
117
You can make use of `torch.compile` function and gain a speed-up of about 2-3x:
Kashif Rasul's avatar
Kashif Rasul committed
118
119

```python
120
121
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)
Kashif Rasul's avatar
Kashif Rasul committed
122
123
124
```

## Limitations
125

Kashif Rasul's avatar
Kashif Rasul committed
126
127
- Due to the high compression employed by Würstchen, generations can lack a good amount
of detail. To our human eye, this is especially noticeable in faces, hands etc.
Kashif Rasul's avatar
Kashif Rasul committed
128
- **Images can only be generated in 128-pixel steps**, e.g. the next higher resolution
Kashif Rasul's avatar
Kashif Rasul committed
129
130
131
132
133
134
135
after 1024x1024 is 1152x1152
- The model lacks the ability to render correct text in images
- The model often does not achieve photorealism
- Difficult compositional prompts are hard for the model

The original codebase, as well as experimental ideas, can be found at [dome272/Wuerstchen](https://github.com/dome272/Wuerstchen).

136

137
## WuerstchenCombinedPipeline
Kashif Rasul's avatar
Kashif Rasul committed
138
139
140
141
142
143
144

[[autodoc]] WuerstchenCombinedPipeline
	- all
	- __call__

## WuerstchenPriorPipeline

145
[[autodoc]] WuerstchenPriorPipeline
Kashif Rasul's avatar
Kashif Rasul committed
146
147
148
149
150
151
152
153
154
155
156
157
	- all
	- __call__

## WuerstchenPriorPipelineOutput

[[autodoc]] pipelines.wuerstchen.pipeline_wuerstchen_prior.WuerstchenPriorPipelineOutput

## WuerstchenDecoderPipeline

[[autodoc]] WuerstchenDecoderPipeline
	- all
	- __call__
MLRichter's avatar
MLRichter committed
158
159

## Citation
160
161

```bibtex
MLRichter's avatar
MLRichter committed
162
      @misc{pernias2023wuerstchen,
163
164
            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},
MLRichter's avatar
MLRichter committed
165
166
167
168
169
            year={2023},
            eprint={2306.00637},
            archivePrefix={arXiv},
            primaryClass={cs.CV}
      }
Patrick von Platen's avatar
Patrick von Platen committed
170
```