欢迎 Stable Diffusion 3 加入 🧨 Diffusers
作为 Stability AI 的 Stable Diffusion 家族最新的模型， Stable Diffusion 3 (SD3) 现已登陆 Hugging Face Hub，并且可用在 🧨 Diffusers 中使用了。
- Stable Diffusion 3 https://stability.ai/news/stable-diffusion-3-research-paper
当前放出的模型版本是 Stable Diffusion 3 Medium，有二十亿 (2B) 的参数量。
针对当前发布版本，我们提供了:
- Hub 上可供下载的模型
- Diffusers 的代码集成
- SD3 的 Dreambooth 和 LoRA 训练脚本
SD3 新特性
模型
作为一个隐变量扩散模型，SD3 包含了三个不同的文本编码器 ( CLIP L/14 、 OpenCLIP bigG/14 和 T5-v1.1-XXL ) 、一个新提出的多模态 Diffusion Transformer (MMDiT) 模型，以及一个 16 通道的 AutoEncoder 模型 (与 Stable Diffusion XL 中的类似)。
- CLIP L/14 https://hf.co/openai/clip-vit-large-patch14
- OpenCLIP bigG/14 https://hf.co/laion/CLIP-ViT-bigG-14-laion2B-39B-b160k
- T5-v1.1-XXL https://hf.co/google/t5-v1_1-xxl
- Stable Diffusion XL https://arxiv.org/abs/2307.01952
SD3 以序列 Embedding 的形式处理文本输入和视觉隐空间特征。位置编码 (Positional Encoding) 是施加在隐空间特征的 2x2 patch 上的，随后被展开成 patch 的 Enbedding 序列。这一序列和文本的特征序列一起，被送入 MMDiT 的各个模块中去。两种特征序列被转化成相同特征维度，拼接在一起，然后送入一系列注意力机制模块和多层感知机 (MLP) 里。
为应对两种模态间的差异，MMDiT 模块使用两组不同的权重去转换文本和图像序列的特征维度。两个序列之后会在注意力操作之前被合并在一起。这种设计使得两种表征能在自己的特征空间里工作，同时也使得它们之间可以通过注意力机制 [1] 从对方的特征中提取有用的信息。这种文本和图像间双向的信息流动有别于以前的文生图模型，后者的文本信息是通过 cross-attention 送入模型的，且不同层输入的文本特征均是文本编码器的输出，不随深度的变化而改变。
此外，SD3 还在时间步 (timestep) 这一条件信息上加入了汇合过的文本特征，这些文本特征来自使用的两个 CLIP 模型。这些汇合过的文本特征被拼接在一起，然后加到时间步的 Embedding 上，再送入每个 MMDiT 模块。
使用 Rectified Flow Matching 训练
除了结构上的创新，SD3 也使用了 conditional flow-matching 作为训练目标函数来训练模型。这一方法中，前向加噪过程被定义为一个 rectified flow ，以一条直线连接数据分布和噪声分布。
- conditional flow-matching https://arxiv.org/html/2403.03206v1#S2
- rectified flow https://arxiv.org/html/2403.03206v1#S3
采样过程也变得更简单了，当采样步数减少的时候，模型性能也很稳定。为此，我们也引入了新的 scheduler ( FlowMatchEulerDiscreteScheduler )，集成了 rectified flow-matching 的运算公式以及欧拉方法 (Euler Method) 的采样步骤。同时还提出了一个与生成分辨率相关的 shift 参数。对于高分辨率，增大 shift 的值可以更好地处理 noise scaling。针对 2B 模型，我们建议设置 shift=3.0 。
如想快速尝试 SD3，可以使用下面的一个基于 Gradio 的应用:
stabilityai/stable-diffusion-3-medium
在 Diffusers 中使用 SD3
如想在 diffusers 中使用 SD3，首先请确保安装的 diffusers 是最新版本:
pip install --upgrade diffusers<br>使用模型前，你需要先到 Stable Diffusion 3 Medium 在 Hugging Face 的页面 ，填写表格并同意相关内容。一切就绪后，你需要登录你的 huggingface 账号:
- Stable Diffusion 3 Medium 在 Hugging Face 的页面 https://hf.co/stabilityai/stable-diffusion-3-medium-diffusers
huggingface-cli login<br>下面程序将会下载 SD3 的 2B 参数模型，并使用 fp16 精度。Stability AI 原本发布的模型精度就是 fp16 ，这也是推荐的模型推理精度。
文生图
import torch<br>from diffusers import StableDiffusion3Pipeline<br><br>pipe = StableDiffusion3Pipeline.from_pretrained("stabilityai/stable-diffusion-3-medium-diffusers", torch_dtype=torch.float16)<br>pipe = pipe.to("cuda")<br><br>image = pipe(<br> "A cat holding a sign that says hello world",<br> negative_prompt="",<br> num_inference_steps=28,<br> guidance_scale=7.0,<br>).images[0]<br>image<br>
hello_world_cat 图生图
import torch<br>from diffusers import StableDiffusion3Img2ImgPipeline<br>from diffusers.utils import load_image<br><br>pipe = StableDiffusion3Img2ImgPipeline.from_pretrained("stabilityai/stable-diffusion-3-medium-diffusers", torch_dtype=torch.float16)<br>pipe = pipe.to("cuda")<br><br>init_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cat.png")<br>prompt = "cat wizard, gandalf, lord of the rings, detailed, fantasy, cute, adorable, Pixar, Disney, 8k"<br>image = pipe(prompt, image=init_image).images[0]<br>image<br>
wizard_cat 相关的 SD3 文档可在 这里 查看。
- 完整代码示例 https://hf.co/docs/diffusers/main/en/api/pipelines/stable_diffusion/stable_diffusion_3
对 SD3 进行内存优化
SD3 使用了三个文本编码器，其中一个是
T5-XXL model
，是一个很大的模型。这使得在显存小于 24GB 的 GPU 上跑模型非常困难，即使使用的是 fp16 精度。
- T5-XXL model https://hf.co/google/t5-v1_1-xxl
对此，diffusers 集成了一些内存优化手段，来让 SD3 能在更多的 GPU 上跑起来。
使用 Model Offloading 推理
Diffusers 上一个最常用的内存优化手段就是 model offloading。它使得你可以在推理时，把一些当前不需要的模型组件卸载到 CPU 上，以此节省 GPU 显存。但这会引入少量的推理时长增长。在推理时，model offloading 只会将模型当前需要参与计算的部分放在 GPU 上，而把剩余部分放在 CPU 上。
import torch<br>from diffusers import StableDiffusion3Pipeline<br><br>pipe = StableDiffusion3Pipeline.from_pretrained("stabilityai/stable-diffusion-3-medium-diffusers", torch_dtype=torch.float16)<br>pipe.enable_model_cpu_offload()<br><br>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."<br>image = pipe(prompt).images[0]<br>不使用 T5 模型进行推理
推理时移除掉 4.7B 参数量的 T5-XXL 文本编码器 可以很大程度地减少内存需求，带来的性能损失却很小。
- 推理时移除掉 4.7B 参数量的 T5-XXL 文本编码器 https://arxiv.org/html/2403.03206v1#S5.F9
import torch<br>from diffusers import StableDiffusion3Pipeline<br><br>pipe = StableDiffusion3Pipeline.from_pretrained("stabilityai/stable-diffusion-3-medium-diffusers", text_encoder_3=None, tokenizer_3=None, torch_dtype=torch.float16)<br>pipe = pipe.to("cuda")<br><br>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."<br>image = pipe("").images[0]<br>使用量化版的 T5-XXL 模型
使用 bitsandbytes 这个库，你也可以加载 8 比特量化版的 T5-XXL 模型，进一步减少显存需求。
import torch<br>from diffusers import StableDiffusion3Pipeline<br>from transformers import T5EncoderModel, BitsAndBytesConfig<br><br># Make sure you have `bitsandbytes` installed.<br>quantization_config = BitsAndBytesConfig(load_in_8bit=True)<br><br>model_id = "stabilityai/stable-diffusion-3-medium-diffusers"<br>text_encoder = T5EncoderModel.from_pretrained(<br> model_id,<br> subfolder="text_encoder_3",<br> quantization_config=quantization_config,<br>)<br>pipe = StableDiffusion3Pipeline.from_pretrained(<br> model_id,<br> text_encoder_3=text_encoder,<br> device_map="balanced",<br> torch_dtype=torch.float16<br>)<br>完整代码在 这里 。
- 完整代码示例 https://gist.github.com/sayakpaul/82acb5976509851f2db1a83456e504f1
显存优化小结
所有的基准测试都用了 2B 参数量的 SD3 模型，测试在一个 A100-80G 上进行，使用 fp16 精度推理，PyTorch 版本为 2.3。
我们对每个推理调用跑十次，记录平均峰值显存用量和 20 步采样的平均时长。
SD3 性能优化
为加速推理，我们可以使用 torch.compile() 来获取优化过的 vae 和 transformer 部分的计算图。
import torch<br>from diffusers import StableDiffusion3Pipeline<br><br>torch.set_float32_matmul_precision("high")<br><br>torch._inductor.config.conv_1x1_as_mm = True<br>torch._inductor.config.coordinate_descent_tuning = True<br>torch._inductor.config.epilogue_fusion = False<br>torch._inductor.config.coordinate_descent_check_all_directions = True<br><br>pipe = StableDiffusion3Pipeline.from_pretrained(<br> "stabilityai/stable-diffusion-3-medium-diffusers",<br> torch_dtype=torch.float16<br>).to("cuda")<br>pipe.set_progress_bar_config(disable=True)<br><br>pipe.transformer.to(memory_format=torch.channels_last)<br>pipe.vae.to(memory_format=torch.channels_last)<br><br>pipe.transformer = torch.compile(pipe.transformer, mode="max-autotune", fullgraph=True)<br>pipe.vae.decode = torch.compile(pipe.vae.decode, mode="max-autotune", fullgraph=True)<br><br># Warm Up<br>prompt = "a photo of a cat holding a sign that says hello world",<br>for _ in range(3):<br> _ = pipe(prompt=prompt, generator=torch.manual_seed(1))<br><br># Run Inference<br>image = pipe(prompt=prompt, generator=torch.manual_seed(1)).images[0]<br>image.save("sd3_hello_world.png")<br>完整代码可参考 这里 。
- 完整代码示例 https://gist.github.com/sayakpaul/508d89d7aad4f454900813da5d42ca97
我们测量了使用过 torch.compile() 的 SD3 的推理速度 (在 A100-80G 上，使用 fp16 推理，PyTorch 版本为 2.3)。我们针对每个生成任务跑 10 遍，每次推理使用 20 步采样。平均推理耗时是 0.585 秒， 这比 eager execution 模式下快了四倍 。
使用 DreamBooth 和 LoRA 进行微调
最后，我们还提供了使用 LoRA 的 DreamBooth 代码，用于微调 SD3。这一程序不仅能微调模型，还能作为一个参考，如果你想使用 rectified flow 来训练模型。当然，热门的 rectified flow 实现代码还有 minRF 。
- LoRA https://hf.co/blog/lora
- DreamBooth https://dreambooth.github.io/
- minRF https://github.com/cloneofsimo/minRF/
如果需要使用该程序，首先需要确保各项设置都已完成，同时准备好一个数据集 (比如
这个
)。你需要安装 peft 和 bitsandbytes ，然后再开始运行训练程序:
- 这个 https://hf.co/datasets/diffusers/dog-example
export MODEL_NAME="stabilityai/stable-diffusion-3-medium-diffusers"<br>export INSTANCE_DIR="dog"<br>export OUTPUT_DIR="dreambooth-sd3-lora"<br><br>accelerate launch train_dreambooth_lora_sd3.py \<br> --pretrained_model_name_or_path=${MODEL_NAME} \<br> --instance_data_dir=${INSTANCE_DIR} \<br> --output_dir=/raid/.cache/${OUTPUT_DIR} \<br> --mixed_precision="fp16" \<br> --instance_prompt="a photo of sks dog" \<br> --resolution=1024 \<br> --train_batch_size=1 \<br> --gradient_accumulation_steps=4 \<br> --learning_rate=1e-5 \<br> --report_to="wandb" \<br> --lr_scheduler="constant" \<br> --lr_warmup_steps=0 \<br> --max_train_steps=500 \<br> --weighting_scheme="logit_normal" \<br> --validation_prompt="A photo of sks dog in a bucket" \<br> --validation_epochs=25 \<br> --seed="0" \<br> --push_to_hub<br>声明
感谢 Stability AI 团队开发并开源了 Stable Diffusion 3 并让我们提早体验，也感谢 Linoy 对撰写此文的帮助。
- Linoy 个人主页 https://hf.co/linoyts
原文链接: https://hf.co/blog/sd3
原文作者: Dhruv Nair, YiYi Xu, Sayak Paul, Alvaro Somoza, Kashif Rasul, Apolinário from multimodal AI art
译者: hugging-hoi2022
本文分享自微信公众号 - Hugging Face（gh_504339124f0f）。
如有侵权，请联系 support@oschina.cn 删除。
本文参与“OSC源创计划”，欢迎正在阅读的你也加入，一起分享。
