Diffusers 支持 Stable Diffusion 3 Medium，含内存与性能优化

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

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

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

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

目录

• SD3 的新特性
• 将 SD3 与 Diffusers 一起使用
• 内存优化，使 SD3 能在各种硬件上运行
• 性能优化，加快速度
• 微调和为 SD3 创建 LoRA

SD3 的新特性？

模型

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

SD3 将文本输入和像素潜变量作为嵌入序列进行处理。位置编码被添加到潜变量的 2x2 补丁中，然后将这些补丁展平为补丁编码序列。该序列与文本编码序列一起被送入 MMDiT 块，在其中被嵌入到公共维度，连接，并通过一系列调制注意力和 MLP 传递。

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

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

使用矫正流匹配进行训练

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

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

要快速试用 SD3，请参考以下应用程序：

将 SD3 与 Diffusers 一起使用

要将 SD3 与 Diffusers 一起使用，请确保升级到最新的 Diffusers 版本。

pip install --upgrade diffusers

由于该模型受到门控限制，在使用 diffusers 之前，您首先需要前往 Stable Diffusion 3 Medium HF Mirror 页面，填写表格并接受门控。进入后，您需要登录，以便系统知道您已接受门控。使用以下命令进行登录：

huggingface-cli login

以下代码片段将下载 fp16 精度的 2B 参数版 SD3。这是 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://hf-mirror.com/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 模型。这使得在显存小于 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
)

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

内存优化总结

所有基准测试均使用 2B 版 SD3 模型，在具有 80GB 显存的 A100 GPU 上，使用 fp16 精度和 PyTorch 2.3 进行。

在我们的内存基准测试中，我们使用 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，对 torch.compile() 在 SD3 上的性能进行了基准测试。我们运行了 10 次流水线推理调用，包含 20 个扩散步骤。我们发现，使用编译版模型的平均推理时间为0.585 秒，比即时执行快 4 倍。

Dreambooth 和 LoRA 微调

此外，我们为 SD3 提供了一个利用 LoRA 的 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 帮助我们制作博客文章缩略图。