🧨 Diffusers 欢迎 Stable Diffusion 3

Stable Diffusion 3 (SD3)，Stability AI 最新一代的 Stable Diffusion 模型系列，现已在 Hugging Face Hub 上提供，并可与 🧨 Diffusers 一起使用。

今天发布的模型是 Stable Diffusion 3 Medium，拥有 2B 参数。

作为本次发布的一部分，我们提供了

1. Hub 上的模型
2. Diffusers 集成
3. SD3 Dreambooth 和 LoRA 训练脚本

目录

• SD3 有哪些新功能？
• 在 Diffusers 中使用 SD3
• 内存优化，使 SD3 能够在各种硬件上运行
• 性能优化以加速
• SD3 的微调和 LoRA 创建

SD3 有哪些新功能？

模型

SD3 是一个潜在扩散模型，由三个不同的文本编码器（CLIP L/14、OpenCLIP bigG/14 和 T5-v1.1-XXL）、一个新颖的多模态扩散 Transformer (MMDiT) 模型以及一个与 Stable Diffusion XL 中使用的模型类似的 16 通道自动编码器模型组成。

SD3 将文本输入和像素潜在编码处理为嵌入序列。位置编码被添加到潜在编码的 2x2 补丁中，然后将其展平为补丁编码序列。此序列与文本编码序列一起被输入到 MMDiT 块中，在那里它们被嵌入到共同的维度中，然后连接并通过一系列调制注意力层和 MLP。

为了解决两种模态之间的差异，MMDiT 块使用两组独立的权重将文本和图像序列嵌入到共同的维度。这些序列在注意力操作之前连接在一起，这使得两种表示可以在各自的空间中工作，同时在注意力操作期间考虑另一种表示。这种文本和图像数据之间的双向信息流与以前的文本到图像合成方法不同，在以前的方法中，文本信息通过与固定文本表示的交叉注意力合并到潜在空间中。

SD3 还使用其两个 CLIP 模型中的池化文本嵌入作为其时间步条件的一部分。这些嵌入首先被连接并添加到时间步嵌入中，然后传递给每个 MMDiT 块。

使用校正流匹配进行训练

除了架构更改之外，SD3 还应用了条件流匹配目标来训练模型。在这种方法中，前向噪声过程被定义为连接数据和噪声分布的直线流。

校正流匹配采样过程更简单，并且在减少采样步骤数时表现良好。为了支持 SD3 的推理，我们引入了一个新的调度器 (FlowMatchEulerDiscreteScheduler)，它具有校正流匹配公式和 Euler 方法步骤。它还通过 shift 参数实现了时间步调度器随分辨率的移动。增加 shift 值可以更好地处理更高分辨率的噪声缩放。建议对 2B 模型使用 shift=3.0。

要快速试用 SD3，请参考下面的应用程序

在 Diffusers 中使用 SD3

要在 Diffusers 中使用 SD3，请确保升级到最新的 Diffusers 版本。

pip install --upgrade diffusers

由于模型是受限的，在使用 diffusers 之前，您首先需要访问 Stable Diffusion 3 Medium Hugging Face 页面，填写表格并接受协议。一旦进入，您需要登录，以便您的系统知道您已接受协议。使用以下命令登录

huggingface-cli login

以下代码片段将以 fp16 精度下载 SD3 的 2B 参数版本。这是 Stability AI 发布原始检查点时使用的格式，也是运行推理的推荐方式。

文本到图像

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

图像到图像

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

您可以在此处查看 SD3 文档。

SD3 的内存优化

SD3 使用三个文本编码器，其中一个是庞大的T5-XXL 模型。这使得在 VRAM 小于 24GB 的 GPU 上运行模型具有挑战性，即使使用 fp16 精度也是如此。

为了解决这个问题，Diffusers 集成附带了内存优化，允许 SD3 在更广泛的设备上运行。

使用模型卸载运行推理

Diffusers 中最基本的内存优化允许您在推理期间将模型组件卸载到 CPU，以节省内存，同时推理延迟略有增加。模型卸载只会在需要执行时将模型组件移动到 GPU，同时将其余组件保留在 CPU 上。

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]

推理时丢弃 T5 文本编码器

在推理过程中移除内存密集型的 4.7B 参数 T5-XXL 文本编码器可以显著降低 SD3 的内存需求，而性能损失很小。

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]

使用 T5-XXL 模型的量化版本

您可以使用 bitsandbytes 库以 8 位加载 T5-XXL 模型，以进一步降低内存需求。

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
)

您可以在此处找到完整的代码片段。

内存优化总结

所有基准测试均使用 80GB VRAM 的 A100 GPU，采用 fp16 精度和 PyTorch 2.3，对 SD3 模型的 2B 版本进行。

对于我们的内存基准测试，我们使用 3 次管道调用进行预热，并报告 10 次管道调用的平均推理时间。我们使用 StableDiffusion3Pipeline 的 __call__() 方法的默认参数。

SD3 的性能优化

为了提高推理延迟，我们可以使用 torch.compile() 来获得 vae 和 transformer 组件的优化计算图。

import torch
from diffusers import StableDiffusion3Pipeline

torch.set_float32_matmul_precision("high")

torch._inductor.config.conv_1x1_as_mm = True
torch._inductor.config.coordinate_descent_tuning = True
torch._inductor.config.epilogue_fusion = False
torch._inductor.config.coordinate_descent_check_all_directions = True

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

pipe.transformer.to(memory_format=torch.channels_last)
pipe.vae.to(memory_format=torch.channels_last)

pipe.transformer = torch.compile(pipe.transformer, mode="max-autotune", fullgraph=True)
pipe.vae.decode = torch.compile(pipe.vae.decode, mode="max-autotune", fullgraph=True)

# Warm Up
prompt = "a photo of a cat holding a sign that says hello world",
for _ in range(3):
    _ = pipe(prompt=prompt, generator=torch.manual_seed(1))

# Run Inference
image = pipe(prompt=prompt, generator=torch.manual_seed(1)).images[0]
image.save("sd3_hello_world.png")

完整脚本请参阅此处。

我们在单台 80GB A100 机器上，使用 fp16 精度和 PyTorch 2.3，对 SD3 的 torch.compile() 性能进行了基准测试。我们运行了 10 次管道推理调用，包含 20 个扩散步骤。我们发现，编译后模型的平均推理时间为 0.585 秒，比急切执行快 4 倍。

Dreambooth 和 LoRA 微调

此外，我们还提供了利用 LoRA 的 SD3 DreamBooth 微调脚本。该脚本可用于高效微调 SD3，并可作为实现基于校正流的训练管道的参考。其他流行的校正流实现包括 minRF。

要开始使用该脚本，首先，请确保您有正确的设置和可用的演示数据集（例如此数据集）。详情请参阅此处。安装 peft 和 bitsandbytes 后，我们就可以开始了

export MODEL_NAME="stabilityai/stable-diffusion-3-medium-diffusers"
export INSTANCE_DIR="dog"
export OUTPUT_DIR="dreambooth-sd3-lora"

accelerate launch train_dreambooth_lora_sd3.py \
  --pretrained_model_name_or_path=${MODEL_NAME}  \
  --instance_data_dir=${INSTANCE_DIR} \
  --output_dir=/raid/.cache/${OUTPUT_DIR} \
  --mixed_precision="fp16" \
  --instance_prompt="a photo of sks dog" \
  --resolution=1024 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=4 \
  --learning_rate=1e-5 \
  --report_to="wandb" \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --max_train_steps=500 \
  --weighting_scheme="logit_normal" \
  --validation_prompt="A photo of sks dog in a bucket" \
  --validation_epochs=25 \
  --seed="0" \
  --push_to_hub

致谢

感谢 Stability AI 团队实现 Stable Diffusion 3 并提供早期访问权限。感谢 Linoy 帮助我们制作博客文章缩略图。