Diffusers 支持 Stable Diffusion 3.5 Large：含 8B 模型与量化、LoRA 训练

Stable Diffusion 3.5 是其前身 Stable Diffusion 3 的改进版本。 截至今日，这些模型已在 HF Mirror Hub 上提供，并可与 🧨 Diffusers 一起使用。

此次发布包含两个检查点：

• 一个大型（8B）模型
• 一个大型（8B）时间步蒸馏模型，支持少步推理

在本文中，我们将重点介绍如何配合 Diffusers 使用 Stable Diffusion 3.5（SD3.5），涵盖推理和训练两方面。

目录

• 架构变更
• 使用 SD3.5 与 Diffusers
• 使用量化进行推理
• 使用量化训练 LoRA
• 使用单文件加载
• 重要链接

架构变更

SD3.5（大型）的 Transformer 架构与 SD3（中型）非常相似，但有如下变化：

• QK 归一化：对于训练大型 Transformer 模型，QK 归一化已成为标准做法，SD3.5 Large 也不例外。
• 双重注意力层：在 MMDiT 块中，SD3.5 不再为每种模态流使用单层注意力，而是使用双重注意力层。

文本编码器、VAE 和噪声调度器方面的其余细节与 SD3 Medium 完全相同。有关 SD3 的更多信息，我们建议查看原始论文。

使用 SD3.5 与 Diffusers

请确保安装最新版本的 diffusers：

pip install -U diffusers

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

huggingface-cli login

以下代码片段将下载 torch.bfloat16 精度的 8B 参数版 SD3.5。 这是 Stability AI 发布的原始检查点所使用的格式，也是推荐用于运行推理的方式。

import torch
from diffusers import StableDiffusion3Pipeline

pipe = StableDiffusion3Pipeline.from_pretrained(
    "stabilityai/stable-diffusion-3.5-large", torch_dtype=torch.bfloat16
).to("cuda")

image = pipe(
    prompt="a photo of a cat holding a sign that says hello world",
    negative_prompt="",
    num_inference_steps=40,
    height=1024,
    width=1024,
    guidance_scale=4.5,
).images[0]

image.save("sd3_hello_world.png")

该版本还包含一个“时间步蒸馏”模型，它消除了无分类器引导，让我们能以更少的步数生成图像（通常为 4-8 步）。

import torch
from diffusers import StableDiffusion3Pipeline

pipe = StableDiffusion3Pipeline.from_pretrained(
    "stabilityai/stable-diffusion-3.5-large-turbo", torch_dtype=torch.bfloat16
).to("cuda")

image = pipe(
    prompt="a photo of a cat holding a sign that says hello world",
    num_inference_steps=4,
    height=1024,
    width=1024,
    guidance_scale=1.0,
).images[0]

image.save("sd3_hello_world.png")

我们在 SD3 博客文章和官方 Diffusers 文档中展示的所有示例应该都已适用于 SD3.5。 特别是，这两份资源都深入探讨了如何优化运行推理时的内存需求。 由于 SD3.5 Large 比 SD3 Medium 大得多，内存优化对于在消费级设备上运行推理变得至关重要。

使用量化运行推理

Diffusers 原生支持 bitsandbytes 量化，可进一步优化内存。

首先，确保安装所有必要的库：

pip install -Uq git+https://github.com/huggingface/transformers@main
pip install -Uq bitsandbytes

然后以 “NF4”精度加载 Transformer：

from diffusers import BitsAndBytesConfig, SD3Transformer2DModel
import torch

model_id = "stabilityai/stable-diffusion-3.5-large"
nf4_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16
)
model_nf4 = SD3Transformer2DModel.from_pretrained(
    model_id,
    subfolder="transformer",
    quantization_config=nf4_config,
    torch_dtype=torch.bfloat16
)

现在，我们已准备好运行推理：

from diffusers import StableDiffusion3Pipeline

pipeline = StableDiffusion3Pipeline.from_pretrained(
    model_id, 
    transformer=model_nf4,
    torch_dtype=torch.bfloat16
)
pipeline.enable_model_cpu_offload()

prompt = "A whimsical and creative image depicting a hybrid creature that is a mix of a waffle and a hippopotamus, basking in a river of melted butter amidst a breakfast-themed landscape. It features the distinctive, bulky body shape of a hippo. However, instead of the usual grey skin, the creature's body resembles a golden-brown, crispy waffle fresh off the griddle. The skin is textured with the familiar grid pattern of a waffle, each square filled with a glistening sheen of syrup. The environment combines the natural habitat of a hippo with elements of a breakfast table setting, a river of warm, melted butter, with oversized utensils or plates peeking out from the lush, pancake-like foliage in the background, a towering pepper mill standing in for a tree.  As the sun rises in this fantastical world, it casts a warm, buttery glow over the scene. The creature, content in its butter river, lets out a yawn. Nearby, a flock of birds take flight"
image = pipeline(
    prompt=prompt,
    negative_prompt="",
    num_inference_steps=28,
    guidance_scale=4.5,
    max_sequence_length=512,
).images[0]
image.save("whimsical.png")

您可以在 BitsAndBytesConfig 中控制其他参数。详情请参阅文档。

也可以直接加载使用与上述相同 nf4_config 量化的模型。 这对于 RAM 较低的机器特别有帮助。有关端到端示例，请参阅此 Colab Notebook。

使用量化训练 SD3.5 Large 的 LoRA

得益于像 bitsandbytes 和 peft 这样的库，在拥有 24GB 显存的消费级显卡上微调像 SD3.5 Large 这样的大模型成为可能。我们已经可以利用现有的 SD3 训练脚本 来训练 LoRA 了。 下面的训练命令已经可以正常工作：

accelerate launch train_dreambooth_lora_sd3.py \
  --pretrained_model_name_or_path="stabilityai/stable-diffusion-3.5-large"  \
  --dataset_name="Norod78/Yarn-art-style" \
  --output_dir="yart_art_sd3-5_lora" \
  --mixed_precision="bf16" \
  --instance_prompt="Frog, yarn art style" \
  --caption_column="text"\
  --resolution=768 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=1 \
  --learning_rate=4e-4 \
  --report_to="wandb" \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --max_train_steps=700 \
  --rank=16 \
  --seed="0" \
  --push_to_hub

然而，要使其与量化配合工作，我们需要调整几个参数。下面，我们提供如何做到这一点的指引：

• 我们使用量化配置初始化 transformer，或者直接加载量化检查点。
• 然后，我们使用 peft 中的 prepare_model_for_kbit_training() 来准备它。
• 得益于 peft 对 bitsandbytes 的强力支持，其余流程保持不变！

更完整的示例请参考这个示例脚本。

使用单文件加载 Stable Diffusion 3.5 Transformer

你可以使用 Stability AI 发布的原始检查点文件，通过 from_single_file 方法加载 Stable Diffusion 3.5 Transformer 模型：

import torch
from diffusers import SD3Transformer2DModel, StableDiffusion3Pipeline

transformer = SD3Transformer2DModel.from_single_file(
    "https://hf-mirror.com/stabilityai/stable-diffusion-3.5-large-turbo/blob/main/sd3.5_large.safetensors",
    torch_dtype=torch.bfloat16,
)
pipe = StableDiffusion3Pipeline.from_pretrained(
    "stabilityai/stable-diffusion-3.5-large",
    transformer=transformer,
    torch_dtype=torch.bfloat16,
)
pipe.enable_model_cpu_offload()
image = pipe("a cat holding a sign that says hello world").images[0]
image.save("sd35.png") 

重要链接

• Hub 上的 Stable Diffusion 3.5 Large 合集
• Stable Diffusion 3.5 的官方 Diffusers 文档
• Colab Notebook 用于运行量化推理
• 训练 LoRA
• Stable Diffusion 3 论文
• Stable Diffusion 3 博客文章

致谢：感谢 Daniel Frank 提供本博客文章缩略图所用的背景照片。感谢 Pedro Cuenca 和 Tom Aarsen 对文章草稿的审阅。