|
| 1 | +<!--Copyright 2024 The HuggingFace Team. All rights reserved. |
| 2 | +
|
| 3 | +Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with |
| 4 | +the License. You may obtain a copy of the License at |
| 5 | +
|
| 6 | +http://www.apache.org/licenses/LICENSE-2.0 |
| 7 | +
|
| 8 | +Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on |
| 9 | +an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the |
| 10 | +specific language governing permissions and limitations under the License. |
| 11 | +--> |
| 12 | + |
| 13 | +# Stable Diffusion 3 |
| 14 | + |
| 15 | +Stable Diffusion 3 (SD3) was proposed in [Scaling Rectified Flow Transformers for High-Resolution Image Synthesis](https://arxiv.org/pdf/2403.03206.pdf) by Patrick Esser, Sumith Kulal, Andreas Blattmann, Rahim Entezari, Jonas Muller, Harry Saini, Yam Levi, Dominik Lorenz, Axel Sauer, Frederic Boesel, Dustin Podell, Tim Dockhorn, Zion English, Kyle Lacey, Alex Goodwin, Yannik Marek, and Robin Rombach. |
| 16 | + |
| 17 | +The abstract from the paper is: |
| 18 | + |
| 19 | +*Diffusion models create data from noise by inverting the forward paths of data towards noise and have emerged as a powerful generative modeling technique for high-dimensional, perceptual data such as images and videos. Rectified flow is a recent generative model formulation that connects data and noise in a straight line. Despite its better theoretical properties and conceptual simplicity, it is not yet decisively established as standard practice. In this work, we improve existing noise sampling techniques for training rectified flow models by biasing them towards perceptually relevant scales. Through a large-scale study, we demonstrate the superior performance of this approach compared to established diffusion formulations for high-resolution text-to-image synthesis. Additionally, we present a novel transformer-based architecture for text-to-image generation that uses separate weights for the two modalities and enables a bidirectional flow of information between image and text tokens, improving text comprehension typography, and human preference ratings. We demonstrate that this architecture follows predictable scaling trends and correlates lower validation loss to improved text-to-image synthesis as measured by various metrics and human evaluations.* |
| 20 | + |
| 21 | + |
| 22 | +## Usage Example |
| 23 | + |
| 24 | +_As the model is gated, before using it with diffusers you first need to go to the [Stable Diffusion 3 Medium Hugging Face page](https://huggingface.co/stabilityai/stable-diffusion-3-medium-diffusers), fill in the form and accept the gate. Once you are in, you need to login so that your system knows you’ve accepted the gate._ |
| 25 | + |
| 26 | +Use the command below to log in: |
| 27 | + |
| 28 | +```bash |
| 29 | +huggingface-cli login |
| 30 | +``` |
| 31 | + |
| 32 | +<Tip> |
| 33 | + |
| 34 | +The SD3 pipeline uses three text encoders to generate an image. Model offloading is necessary in order for it to run on most commodity hardware. Please use the `torch.float16` data type for additional memory savings. |
| 35 | + |
| 36 | +</Tip> |
| 37 | + |
| 38 | + |
| 39 | +```python |
| 40 | +import torch |
| 41 | +from diffusers import StableDiffusion3Pipeline |
| 42 | + |
| 43 | +pipe = StableDiffusion3Pipeline.from_pretrained("stabilityai/stable-diffusion-3-medium-diffusers", torch_dtype=torch.float16) |
| 44 | +pipe.to("cuda") |
| 45 | + |
| 46 | +image = pipe( |
| 47 | + prompt="a photo of a cat holding a sign that says hello world", |
| 48 | + negative_prompt="", |
| 49 | + num_inference_steps=28, |
| 50 | + height=1024, |
| 51 | + width=1024, |
| 52 | + guidance_scale=7.0, |
| 53 | +).images[0] |
| 54 | + |
| 55 | +image.save("sd3_hello_world.png") |
| 56 | +``` |
| 57 | + |
| 58 | +## Memory Optimisations for SD3 |
| 59 | + |
| 60 | +SD3 uses three text encoders, one if which is the very large T5-XXL model. This makes it challenging to run the model on GPUs with less than 24GB of VRAM, even when using `fp16` precision. The following section outlines a few memory optimizations in Diffusers that make it easier to run SD3 on low resource hardware. |
| 61 | + |
| 62 | +### Running Inference with Model Offloading |
| 63 | + |
| 64 | +The most basic memory optimization available in Diffusers allows you to offload the components of the model to CPU during inference in order to save memory, while seeing a slight increase in inference latency. Model offloading will only move a model component onto the GPU when it needs to be executed, while keeping the remaining components on the CPU. |
| 65 | + |
| 66 | +```python |
| 67 | +import torch |
| 68 | +from diffusers import StableDiffusion3Pipeline |
| 69 | + |
| 70 | +pipe = StableDiffusion3Pipeline.from_pretrained("stabilityai/stable-diffusion-3-medium-diffusers", torch_dtype=torch.float16) |
| 71 | +pipe.enable_model_cpu_offload() |
| 72 | + |
| 73 | +image = pipe( |
| 74 | + prompt="a photo of a cat holding a sign that says hello world", |
| 75 | + negative_prompt="", |
| 76 | + num_inference_steps=28, |
| 77 | + height=1024, |
| 78 | + width=1024, |
| 79 | + guidance_scale=7.0, |
| 80 | +).images[0] |
| 81 | + |
| 82 | +image.save("sd3_hello_world.png") |
| 83 | +``` |
| 84 | + |
| 85 | +### Dropping the T5 Text Encoder during Inference |
| 86 | + |
| 87 | +Removing the memory-intensive 4.7B parameter T5-XXL text encoder during inference can significantly decrease the memory requirements for SD3 with only a slight loss in performance. |
| 88 | + |
| 89 | +```python |
| 90 | +import torch |
| 91 | +from diffusers import StableDiffusion3Pipeline |
| 92 | + |
| 93 | +pipe = StableDiffusion3Pipeline.from_pretrained( |
| 94 | + "stabilityai/stable-diffusion-3-medium-diffusers", |
| 95 | + text_encoder_3=None, |
| 96 | + tokenizer_3=None, |
| 97 | + torch_dtype=torch.float16 |
| 98 | +) |
| 99 | +pipe.to("cuda") |
| 100 | + |
| 101 | +image = pipe( |
| 102 | + prompt="a photo of a cat holding a sign that says hello world", |
| 103 | + negative_prompt="", |
| 104 | + num_inference_steps=28, |
| 105 | + height=1024, |
| 106 | + width=1024, |
| 107 | + guidance_scale=7.0, |
| 108 | +).images[0] |
| 109 | + |
| 110 | +image.save("sd3_hello_world-no-T5.png") |
| 111 | +``` |
| 112 | + |
| 113 | +### Using a Quantized Version of the T5 Text Encoder |
| 114 | + |
| 115 | +We can leverage the `bitsandbytes` library to load and quantize the T5-XXL text encoder to 8-bit precision. This allows you to keep using all three text encoders while only slightly impacting performance. |
| 116 | + |
| 117 | +First install the `bitsandbytes` library. |
| 118 | + |
| 119 | +```shell |
| 120 | +pip install bitsandbytes |
| 121 | +``` |
| 122 | + |
| 123 | +Then load the T5-XXL model using the `BitsAndBytesConfig`. |
| 124 | + |
| 125 | +```python |
| 126 | +import torch |
| 127 | +from diffusers import StableDiffusion3Pipeline |
| 128 | +from transformers import T5EncoderModel, BitsAndBytesConfig |
| 129 | + |
| 130 | +quantization_config = BitsAndBytesConfig(load_in_8bit=True) |
| 131 | + |
| 132 | +model_id = "stabilityai/stable-diffusion-3-medium-diffusers" |
| 133 | +text_encoder = T5EncoderModel.from_pretrained( |
| 134 | + model_id, |
| 135 | + subfolder="text_encoder_3", |
| 136 | + quantization_config=quantization_config, |
| 137 | +) |
| 138 | +pipe = StableDiffusion3Pipeline.from_pretrained( |
| 139 | + model_id, |
| 140 | + text_encoder_3=text_encoder, |
| 141 | + device_map="balanced", |
| 142 | + torch_dtype=torch.float16 |
| 143 | +) |
| 144 | + |
| 145 | +image = pipe( |
| 146 | + prompt="a photo of a cat holding a sign that says hello world", |
| 147 | + negative_prompt="", |
| 148 | + num_inference_steps=28, |
| 149 | + height=1024, |
| 150 | + width=1024, |
| 151 | + guidance_scale=7.0, |
| 152 | +).images[0] |
| 153 | + |
| 154 | +image.save("sd3_hello_world-8bit-T5.png") |
| 155 | +``` |
| 156 | + |
| 157 | +You can find the end-to-end script [here](https://gist.github.com/sayakpaul/82acb5976509851f2db1a83456e504f1). |
| 158 | + |
| 159 | +## Performance Optimizations for SD3 |
| 160 | + |
| 161 | +### Using Torch Compile to Speed Up Inference |
| 162 | + |
| 163 | +Using compiled components in the SD3 pipeline can speed up inference by as much as 4X. The following code snippet demonstrates how to compile the Transformer and VAE components of the SD3 pipeline. |
| 164 | + |
| 165 | +```python |
| 166 | +import torch |
| 167 | +from diffusers import StableDiffusion3Pipeline |
| 168 | + |
| 169 | +torch.set_float32_matmul_precision("high") |
| 170 | + |
| 171 | +torch._inductor.config.conv_1x1_as_mm = True |
| 172 | +torch._inductor.config.coordinate_descent_tuning = True |
| 173 | +torch._inductor.config.epilogue_fusion = False |
| 174 | +torch._inductor.config.coordinate_descent_check_all_directions = True |
| 175 | + |
| 176 | +pipe = StableDiffusion3Pipeline.from_pretrained( |
| 177 | + "stabilityai/stable-diffusion-3-medium-diffusers", |
| 178 | + torch_dtype=torch.float16 |
| 179 | +).to("cuda") |
| 180 | +pipe.set_progress_bar_config(disable=True) |
| 181 | + |
| 182 | +pipe.transformer.to(memory_format=torch.channels_last) |
| 183 | +pipe.vae.to(memory_format=torch.channels_last) |
| 184 | + |
| 185 | +pipe.transformer = torch.compile(pipe.transformer, mode="max-autotune", fullgraph=True) |
| 186 | +pipe.vae.decode = torch.compile(pipe.vae.decode, mode="max-autotune", fullgraph=True) |
| 187 | + |
| 188 | +# Warm Up |
| 189 | +prompt = "a photo of a cat holding a sign that says hello world", |
| 190 | +for _ in range(3): |
| 191 | + _ = pipe(prompt=prompt, generator=torch.manual_seed(1)) |
| 192 | + |
| 193 | +# Run Inference |
| 194 | +image = pipe(prompt=prompt, generator=torch.manual_seed(1)).images[0] |
| 195 | +image.save("sd3_hello_world.png") |
| 196 | +``` |
| 197 | + |
| 198 | +Check out the full script [here](https://gist.github.com/sayakpaul/508d89d7aad4f454900813da5d42ca97). |
| 199 | + |
| 200 | +## Loading the original checkpoints via `from_single_file` |
| 201 | + |
| 202 | +The `SD3Transformer2DModel` and `StableDiffusion3Pipeline` classes support loading the original checkpoints via the `from_single_file` method. This method allows you to load the original checkpoint files that were used to train the models. |
| 203 | + |
| 204 | +## Loading the original checkpoints for the `SD3Transformer2DModel` |
| 205 | + |
| 206 | +```python |
| 207 | +from diffusers import SD3Transformer2DModel |
| 208 | + |
| 209 | +model = SD3Transformer2DModel.from_single_file("https://huggingface.co/stabilityai/stable-diffusion-3-medium/blob/main/sd3_medium.safetensors") |
| 210 | +``` |
| 211 | + |
| 212 | +## Loading the single checkpoint for the `StableDiffusion3Pipeline` |
| 213 | + |
| 214 | +```python |
| 215 | +from diffusers import StableDiffusion3Pipeline |
| 216 | +from transformers import T5EncoderModel |
| 217 | + |
| 218 | +text_encoder_3 = T5EncoderModel.from_pretrained("stabilityai/stable-diffusion-3-medium-diffusers", subfolder="text_encoder_3", torch_dtype=torch.float16) |
| 219 | +pipe = StableDiffusion3Pipeline.from_single_file("https://huggingface.co/stabilityai/stable-diffusion-3-medium/blob/main/sd3_medium_incl_clips.safetensors", torch_dtype=torch.float16, text_encoder_3=text_encoder_3) |
| 220 | +``` |
| 221 | + |
| 222 | +<Tip> |
| 223 | +`from_single_file` support for the `fp8` version of the checkpoints is coming soon. Watch this space. |
| 224 | +</Tip> |
| 225 | + |
| 226 | +## StableDiffusion3Pipeline |
| 227 | + |
| 228 | +[[autodoc]] StableDiffusion3Pipeline |
| 229 | + - all |
| 230 | + - __call__ |
0 commit comments