🧨 Diffusers 歡迎 Stable Diffusion 3

Stable Diffusion 3 (SD3)，Stability AI 最新一代的 Stable Diffusion 模型系列，現已在 Hugging Face Hub 上提供，並可與 🧨 Diffusers 一起使用。

今天釋出的模型是 Stable Diffusion 3 Medium，擁有 2B 引數。

作為本次釋出的一部分，我們提供了

1. Hub 上的模型
2. Diffusers 整合
3. SD3 Dreambooth 和 LoRA 訓練指令碼

目錄

• SD3 有哪些新功能？
• 在 Diffusers 中使用 SD3
• 記憶體最佳化，使 SD3 能夠在各種硬體上執行
• 效能最佳化以加速
• SD3 的微調和 LoRA 建立

SD3 有哪些新功能？

模型

SD3 是一個潛在擴散模型，由三個不同的文字編碼器（CLIP L/14、OpenCLIP bigG/14 和 T5-v1.1-XXL）、一個新穎的多模態擴散 Transformer (MMDiT) 模型以及一個與 Stable Diffusion XL 中使用的模型類似的 16 通道自動編碼器模型組成。

SD3 將文字輸入和畫素潛在編碼處理為嵌入序列。位置編碼被新增到潛在編碼的 2x2 補丁中，然後將其展平為補丁編碼序列。此序列與文字編碼序列一起被輸入到 MMDiT 塊中，在那裡它們被嵌入到共同的維度中，然後連線並透過一系列調製注意力層和 MLP。

為了解決兩種模態之間的差異，MMDiT 塊使用兩組獨立的權重將文字和影像序列嵌入到共同的維度。這些序列在注意力操作之前連線在一起，這使得兩種表示可以在各自的空間中工作，同時在注意力操作期間考慮另一種表示。這種文字和影像資料之間的雙向資訊流與以前的文字到影像合成方法不同，在以前的方法中，文字資訊透過與固定文字表示的交叉注意力合併到潛在空間中。

SD3 還使用其兩個 CLIP 模型中的池化文字嵌入作為其時間步條件的一部分。這些嵌入首先被連線並新增到時間步嵌入中，然後傳遞給每個 MMDiT 塊。

使用校正流匹配進行訓練

除了架構更改之外，SD3 還應用了條件流匹配目標來訓練模型。在這種方法中，前向噪聲過程被定義為連線資料和噪聲分佈的直線流。

校正流匹配取樣過程更簡單，並且在減少採樣步驟數時表現良好。為了支援 SD3 的推理，我們引入了一個新的排程器 (FlowMatchEulerDiscreteScheduler)，它具有校正流匹配公式和 Euler 方法步驟。它還透過 shift 引數實現了時間步排程器隨解析度的移動。增加 shift 值可以更好地處理更高解析度的噪聲縮放。建議對 2B 模型使用 shift=3.0。

要快速試用 SD3，請參考下面的應用程式

在 Diffusers 中使用 SD3

要在 Diffusers 中使用 SD3，請確保升級到最新的 Diffusers 版本。

pip install --upgrade diffusers

由於模型是受限的，在使用 diffusers 之前，您首先需要訪問 Stable Diffusion 3 Medium Hugging Face 頁面，填寫表格並接受協議。一旦進入，您需要登入，以便您的系統知道您已接受協議。使用以下命令登入

huggingface-cli login

以下程式碼片段將以 fp16 精度下載 SD3 的 2B 引數版本。這是 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://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

您可以在此處檢視 SD3 文件。

SD3 的記憶體最佳化

SD3 使用三個文字編碼器，其中一個是龐大的T5-XXL 模型。這使得在 VRAM 小於 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
)

您可以在此處找到完整的程式碼片段。

記憶體最佳化總結

所有基準測試均使用 80GB VRAM 的 A100 GPU，採用 fp16 精度和 PyTorch 2.3，對 SD3 模型的 2B 版本進行。

對於我們的記憶體基準測試，我們使用 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，對 SD3 的 torch.compile() 效能進行了基準測試。我們運行了 10 次管道推理呼叫，包含 20 個擴散步驟。我們發現，編譯後模型的平均推理時間為 0.585 秒，比急切執行快 4 倍。

Dreambooth 和 LoRA 微調

此外，我們還提供了利用 LoRA 的 SD3 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 幫助我們製作部落格文章縮圖。