歡迎 Stable Diffusion 3 加入 ?? Diffusers
共 11273字，需瀏覽 23分鐘
·
2024-06-21 09:43
作為 Stability AI 的 Stable Diffusion 家族最新的模型，
-
Stable Diffusion 3 https://stability.ai/news/stable-diffusion-3-research-paper
當前放出的模型版本是 Stable Diffusion 3 Medium，有二十億 (2B) 的參數(shù)量。
針對當前發(fā)布版本，我們提供了:
-
Hub 上可供下載的模型 -
Diffusers 的代碼集成 -
SD3 的 Dreambooth 和 LoRA 訓練腳本
SD3 新特性
模型
作為一個隱變量擴散模型，SD3 包含了三個不同的文本編碼器 (
-
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 的各個模塊中去。兩種特征序列被轉(zhuǎn)化成相同特征維度，拼接在一起，然后送入一系列注意力機制模塊和多層感知機 (MLP) 里。
為應(yīng)對兩種模態(tài)間的差異，MMDiT 模塊使用兩組不同的權(quán)重去轉(zhuǎn)換文本和圖像序列的特征維度。兩個序列之后會在注意力操作之前被合并在一起。這種設(shè)計使得兩種表征能在自己的特征空間里工作，同時也使得它們之間可以通過注意力機制 [1] 從對方的特征中提取有用的信息。這種文本和圖像間雙向的信息流動有別于以前的文生圖模型，后者的文本信息是通過 cross-attention 送入模型的，且不同層輸入的文本特征均是文本編碼器的輸出，不隨深度的變化而改變。
此外，SD3 還在時間步 (timestep) 這一條件信息上加入了匯合過的文本特征，這些文本特征來自使用的兩個 CLIP 模型。這些匯合過的文本特征被拼接在一起，然后加到時間步的 Embedding 上，再送入每個 MMDiT 模塊。
使用 Rectified Flow Matching 訓練
除了結(jié)構(gòu)上的創(chuàng)新，SD3 也使用了
-
conditional flow-matching https://arxiv.org/html/2403.03206v1#S2 -
rectified flow https://arxiv.org/html/2403.03206v1#S3
采樣過程也變得更簡單了，當采樣步數(shù)減少的時候，模型性能也很穩(wěn)定。為此，我們也引入了新的 scheduler ( FlowMatchEulerDiscreteScheduler )，集成了 rectified flow-matching 的運算公式以及歐拉方法 (Euler Method) 的采樣步驟。同時還提出了一個與生成分辨率相關(guān)的 shift 參數(shù)。對于高分辨率，增大 shift 的值可以更好地處理 noise scaling。針對 2B 模型，我們建議設(shè)置 shift=3.0 。
如想快速嘗試 SD3，可以使用下面的一個基于 Gradio 的應(yīng)用:
stabilityai/stable-diffusion-3-medium
在 Diffusers 中使用 SD3
如想在 diffusers 中使用 SD3，首先請確保安裝的 diffusers 是最新版本:
pip install --upgrade diffusers
使用模型前，你需要先到
-
Stable Diffusion 3 Medium 在 Hugging Face 的頁面 https://hf.co/stabilityai/stable-diffusion-3-medium-diffusers
huggingface-cli login
下面程序?qū)螺d SD3 的 2B 參數(shù)模型，并使用 fp16 精度。Stability AI 原本發(fā)布的模型精度就是 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
圖生圖
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("")
prompt = "cat wizard, gandalf, lord of the rings, detailed, fantasy, cute, adorable, Pixar, Disney, 8k"
image = pipe(prompt, image=init_image).images[0]
image
相關(guān)的 SD3 文檔可在
-
完整代碼示例 https://hf.co/docs/diffusers/main/en/api/pipelines/stable_diffusion/stable_diffusion_3
對 SD3 進行內(nèi)存優(yōu)化
SD3 使用了三個文本編碼器，其中一個是fp16 精度。
-
T5-XXL model https://hf.co/google/t5-v1_1-xxl
對此，diffusers 集成了一些內(nèi)存優(yōu)化手段，來讓 SD3 能在更多的 GPU 上跑起來。
使用 Model Offloading 推理
Diffusers 上一個最常用的內(nèi)存優(yōu)化手段就是 model offloading。它使得你可以在推理時，把一些當前不需要的模型組件卸載到 CPU 上，以此節(jié)省 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 參數(shù)量的 T5-XXL 文本編碼器 https://arxiv.org/html/2403.03206v1#S5.F9
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
)
完整代碼在
-
完整代碼示例 https://gist.github.com/sayakpaul/82acb5976509851f2db1a83456e504f1
顯存優(yōu)化小結(jié)
所有的基準測試都用了 2B 參數(shù)量的 SD3 模型，測試在一個 A100-80G 上進行，使用 fp16 精度推理，PyTorch 版本為 2.3。
我們對每個推理調(diào)用跑十次，記錄平均峰值顯存用量和 20 步采樣的平均時長。
SD3 性能優(yōu)化
為加速推理，我們可以使用 torch.compile() 來獲取優(yōu)化過的 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")
完整代碼可參考
-
完整代碼示例 https://gist.github.com/sayakpaul/508d89d7aad4f454900813da5d42ca97
我們測量了使用過 torch.compile() 的 SD3 的推理速度 (在 A100-80G 上，使用 fp16 推理，PyTorch 版本為 2.3)。我們針對每個生成任務(wù)跑 10 遍，每次推理使用 20 步采樣。平均推理耗時是 0.585 秒， 這比 eager execution 模式下快了四倍 。
使用 DreamBooth 和 LoRA 進行微調(diào)
最后，我們還提供了使用
-
LoRA https://hf.co/blog/lora -
DreamBooth https://dreambooth.github.io/ -
minRF https://github.com/cloneofsimo/minRF/
如果需要使用該程序，首先需要確保各項設(shè)置都已完成，同時準備好一個數(shù)據(jù)集 (比如peft 和 bitsandbytes ，然后再開始運行訓練程序:
-
這個 https://hf.co/datasets/diffusers/dog-example
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 團隊開發(fā)并開源了 Stable Diffusion 3 并讓我們提早體驗，也感謝
-
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