SD3 新特性
模型
作为一个隐变量扩散模型,SD3 包含了三个不同的文本编码器 (CLIP L/14、OpenCLIP bigG/14 和 T5-v1.1-XXL) 、一个新提出的多模态 Diffusion Transformer (MMDiT) 模型,以及一个 16 通道的 AutoEncoder 模型 (与 Stable Diffusion XL 中的类似)。
SD3 以序列 Embedding 的形式处理文本输入和视觉隐空间特征。位置编码 (Positional Encoding) 是施加在隐空间特征的 2×2 patch 上的,随后被展开成 patch 的 Enbedding 序列。这一序列和文本的特征序列一起,被送入 MMDiT 的各个模块中去。两种特征序列被转化成相同特征维度,拼接在一起,然后送入一系列注意力机制模块和多层感知机 (MLP) 里。
为应对两种模态间的差异,MMDiT 模块使用两组不同的权重去转换文本和图像序列的特征维度。两个序列之后会在注意力操作之前被合并在一起。这种设计使得两种表征能在自己的特征空间里工作,同时也使得它们之间可以通过注意力机制 [1] 从对方的特征中提取有用的信息。这种文本和图像间双向的信息流动有别于以前的文生图模型,后者的文本信息是通过 cross-attention 送入模型的,且不同层输入的文本特征均是文本编码器的输出,不随深度的变化而改变。
此外,SD3 还在时间步 (timestep) 这一条件信息上加入了汇合过的文本特征,这些文本特征来自使用的两个 CLIP 模型。这些汇合过的文本特征被拼接在一起,然后加到时间步的 Embedding 上,再送入每个 MMDiT 模块。
使用 Rectified Flow Matching 训练
除了结构上的创新,SD3 也使用了 conditional flow-matching 作为训练目标函数来训练模型。这一方法中,前向加噪过程被定义为一个 rectified flow,以一条直线连接数据分布和噪声分布。
采样过程也变得更简单了,当采样步数减少的时候,模型性能也很稳定。为此,我们也引入了新的 scheduler ( FlowMatchEulerDiscreteScheduler
),集成了 rectified flow-matching 的运算公式以及欧拉方法 (Euler Method) 的采样步骤。同时还提出了一个与生成分辨率相关的 shift
参数。对于高分辨率,增大 shift
的值可以更好地处理 noise scaling。针对 2B 模型,我们建议设置 shift=3.0
。
如想快速尝试 SD3,可以使用下面的一个基于 Gradio 的应用:
![](https://pic4.zhimg.com/80/v2-b8b494b6192739d4f449d0175a9169ab_1440w.webp)
在 Diffusers 中使用 SD3
如想在 diffusers 中使用 SD3,首先请确保安装的 diffusers 是最新版本:
pip install --upgrade diffusers
使用模型前,你需要先到 Stable Diffusion 3 Medium 在 Hugging Face 的页面,填写表格并同意相关内容。一切就绪后,你需要登录你的 huggingface 账号:
huggingface-cli login
下面程序将会下载 SD3 的 2B 参数模型,并使用 fp16
精度。Stability AI 原本发布的模型精度就是 fp16
,这也是推荐的模型推理精度。
文生图
import torch
from diffusers import StableDiffusion3Pipeline
pipe = StableDiffusion3Pipeline.from_pretrained("stabilityai/stable-diffusion-3-medium-diffusers", torch_dtype=torch.float16)
pipe = pipe.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
![](https://pic1.zhimg.com/80/v2-2a833b18db87861369887522f7bf6500_1440w.webp)
图生图
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)
pipe = pipe.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
![](https://pic2.zhimg.com/80/v2-95836571ceeedad483ab1fd45a8618f5_1440w.webp)
相关的 SD3 文档可在 这里 查看。
对 SD3 进行内存优化
SD3 使用了三个文本编码器,其中一个是 T5-XXL model,是一个很大的模型。这使得在显存小于 24GB 的 GPU 上跑模型非常困难,即使使用的是 fp16
精度。
对此,diffusers 集成了一些内存优化手段,来让 SD3 能在更多的 GPU 上跑起来。
使用 Model Offloading 推理
Diffusers 上一个最常用的内存优化手段就是 model offloading。它使得你可以在推理时,把一些当前不需要的模型组件卸载到 CPU 上,以此节省 GPU 显存。但这会引入少量的推理时长增长。在推理时,model offloading 只会将模型当前需要参与计算的部分放在 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 文本编码器 可以很大程度地减少内存需求,带来的性能损失却很小。
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)
pipe = pipe.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("").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 模型,测试在一个 A100-80G 上进行,使用 fp16
精度推理,PyTorch 版本为 2.3。
我们对每个推理调用跑十次,记录平均峰值显存用量和 20 步采样的平均时长。
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")
完整代码可参考 这里。
我们测量了使用过 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。
如果需要使用该程序,首先需要确保各项设置都已完成,同时准备好一个数据集 (比如 这个)。你需要安装 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 对撰写此文的帮助。
原文链接: https://hf.co/blog/sd3
原文作者: Dhruv Nair, YiYi Xu, Sayak Paul, Alvaro Somoza, Kashif Rasul, Apolinário from multimodal AI art
译者: hugging-hoi2022
发表回复