Diffusers welcomes Stable Diffusion 3
dn6, YiYiXu, sayakpaul, OzzyGT, kashif, multimodalart
[Stable Diffusion 3](https://stability.ai/news/stable-diffusion-3-research-paper) (SD3), Stability AI’s latest iteration of the Stable Diffusion family of models, is now available on the Hugging Face Hub and can be used with 🧨 Diffusers.

The model released today is Stable Diffusion 3 Medium, with 2B parameters.

As part of this release, we have provided:

1. Models on the Hub
2. Diffusers Integration
3. SD3 Dreambooth and LoRA training scripts

## Table Of Contents

- [What’s new with SD3](#whats-new-with-sd3)
- [Using SD3 with Diffusers](#using-sd3-with-diffusers)
- [Memory optimizations to enable running SD3 on a variety of hardware](#memory-optimizations-for-sd3)
- [Performance optimizations to speed things up](#performance-optimizations-for-sd3)
- [Finetuning and creating LoRAs for SD3](#dreambooth-and-lora-fine-tuning)

## What’s New With SD3?

### Model

SD3 is a latent diffusion model that consists of three different text encoders ([CLIP L/14](https://huggingface.co/openai/clip-vit-large-patch14), [OpenCLIP bigG/14](https://huggingface.co/laion/CLIP-ViT-bigG-14-laion2B-39B-b160k), and [T5-v1.1-XXL](https://huggingface.co/google/t5-v1_1-xxl)), a novel Multimodal Diffusion Transformer (MMDiT) model, and a 16 channel AutoEncoder model that is similar to the one used in [Stable Diffusion XL](https://arxiv.org/abs/2307.01952).

SD3 processes text inputs and pixel latents as a sequence of embeddings. Positional encodings are added to 2x2 patches of the latents which are then flattened into a patch encoding sequence. This sequence, along with the text encoding sequence are fed into the MMDiT blocks, where they are embedded to a common dimensionality, concatenated, and passed through a sequence of modulated attentions and MLPs.

In order to account for the differences between the two modalities, the MMDiT blocks use two separate sets of weights to embed the text and image sequences to a common dimensionality.  These sequences are joined before the attention operation, which allows both representations to work in their own space while taking the other one into account during the attention operation. This two-way flow of information between text and image data differs from previous approaches for text-to-image synthesis, where text information is incorporated into the latent via cross-attention with a fixed text representation.

SD3 also makes use of the pooled text embeddings from both its CLIP models as part of its timestep conditioning. These embeddings are first concatenated and added to the timestep embedding before being passed to each of the MMDiT blocks.

### Training with Rectified Flow Matching

In addition to architectural changes, SD3 applies a [conditional flow-matching objective to train the model](https://arxiv.org/html/2403.03206v1#S2). In this approach, the forward noising process is defined as a [rectified flow](https://arxiv.org/html/2403.03206v1#S3) that connects the data and noise distributions on a straight line.

The rectified flow-matching sampling process is simpler and performs well when reducing the number of sampling steps. To support inference with SD3, we have introduced a new scheduler (`FlowMatchEulerDiscreteScheduler`) with a rectified flow-matching formulation and Euler method steps. It also implements resolution-dependent shifting of the timestep schedule via a `shift` parameter. Increasing the `shift` value handles noise scaling better for higher resolutions. It is recommended to use `shift=3.0` for the 2B model.

To quickly try out SD3, refer to the application below:

<script type="module" src="https://gradio.s3-us-west-2.amazonaws.com/4.36.1/gradio.js"> </script>
<gradio-app theme_mode="light" space="stabilityai/stable-diffusion-3-medium"></gradio-app>

## Using SD3 with Diffusers

To use SD3 with Diffusers, make sure to upgrade to the latest Diffusers release.

```sh
pip install --upgrade diffusers
```

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 log in so that your system knows you’ve accepted the gate. Use the command below to log in:

```bash
huggingface-cli login
```

The following snippet will download the 2B parameter version of SD3 in `fp16` precision. This is the format used in the original checkpoint published by Stability AI, and is the recommended way to run inference.

### Text-To-Image

```python
import torch
from diffusers import StableDiffusion3Pipeline

pipe = StableDiffusion3Pipeline.from_pretrained(
    "stabilityai/stable-diffusion-3-medium-diffusers", torch_dtype=torch.float16
).to("cuda")

image = pipe(
	"A cat holding a sign that says hello world",
	negative_prompt="",
    num_inference_steps=28,
    guidance_scale=7.0,
).images[0]
image
```

![hello_world_cat](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/sd3/hello_world_cat.png)

### Image-To-Image

```python
import torch
from diffusers import StableDiffusion3Img2ImgPipeline
from diffusers.utils import load_image

pipe = StableDiffusion3Img2ImgPipeline.from_pretrained(
    "stabilityai/stable-diffusion-3-medium-diffusers", torch_dtype=torch.float16
).to("cuda")

init_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cat.png")
prompt = "cat wizard, gandalf, lord of the rings, detailed, fantasy, cute, adorable, Pixar, Disney, 8k"
image = pipe(prompt, image=init_image).images[0]
image
```

![wizard_cat](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/sd3/wizard_cat.png)

You can check out the SD3 documentation [here](https://huggingface.co/docs/diffusers/main/en/api/pipelines/stable_diffusion/stable_diffusion_3).

## Memory Optimizations for SD3

SD3 uses three text encoders, one of which is the very large [T5-XXL model](https://huggingface.co/google/t5-v1_1-xxl). This makes running the model on GPUs with less than 24GB of VRAM challenging, even when using `fp16` precision.

To account for this, the Diffusers integration ships with memory optimizations that allow SD3 to be run on a wider range of devices.

### Running Inference with Model Offloading

The most basic memory optimization available in Diffusers allows you to offload the components of the model to the 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.

```python
import torch
from diffusers import StableDiffusion3Pipeline

pipe = StableDiffusion3Pipeline.from_pretrained(
    "stabilityai/stable-diffusion-3-medium-diffusers", torch_dtype=torch.float16
)
pipe.enable_model_cpu_offload()

prompt = "smiling cartoon dog sits at a table, coffee mug on hand, as a room goes up in flames. “This is fine,” the dog assures himself."
image = pipe(prompt).images[0]
```

### Dropping the T5 Text Encoder during Inference

[Removing the memory-intensive 4.7B parameter T5-XXL text encoder during inference](https://arxiv.org/html/2403.03206v1#S5.F9) can significantly decrease the memory requirements for SD3 with only a slight loss in performance.

```python
import torch
from diffusers import StableDiffusion3Pipeline

pipe = StableDiffusion3Pipeline.from_pretrained(
    "stabilityai/stable-diffusion-3-medium-diffusers",
    text_encoder_3=None, 
    tokenizer_3=None, 
    torch_dtype=torch.float16
).to("cuda")

prompt = "smiling cartoon dog sits at a table, coffee mug on hand, as a room goes up in flames. “This is fine,” the dog assures himself."
image = pipe(prompt).images[0]
```

## Using A Quantized Version of the T5-XXL Model

You can load the T5-XXL model in 8 bits using the `bitsandbytes` library to reduce the memory requirements further.

```python
import torch
from diffusers import StableDiffusion3Pipeline
from transformers import T5EncoderModel, BitsAndBytesConfig

# Make sure you have `bitsandbytes` installed.
quantization_config = BitsAndBytesConfig(load_in_8bit=True)

model_id = "stabilityai/stable-diffusion-3-medium-diffusers"
text_encoder = T5EncoderModel.from_pretrained(
    model_id,
    subfolder="text_encoder_3",
    quantization_config=quantization_config,
)
pipe = StableDiffusion3Pipeline.from_pretrained(
    model_id,
    text_encoder_3=text_encoder,
    device_map="balanced",
    torch_dtype=torch.float16
)
```

*You can find the full code snippet [here](https://gist.github.com/sayakpaul/82acb5976509851f2db1a83456e504f1).*

### Summary of Memory Optimizations

All benchmarking runs were conducted using the 2B version of the SD3 model on an A100 GPU with 80GB of VRAM using `fp16` precision and PyTorch 2.3.

For our memory benchmarks, we use 3 iterations of pipeline calls for warming up and report an average inference time of 10 iterations of pipeline calls. We use the default arguments of the [`StableDiffusion3Pipeline` `__call__()` method](https://github.com/huggingface/diffusers/blob/adc31940a9cedbbe2fca8142d09bb81db14a8a52/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3.py#L634).

| **Technique**         | **Inference Time (secs)** | **Memory (GB)** |
|: