Stable Diffusion 微調及推理優化

本文作者：劉遠，騰訊雲泛互聯網首席解決方案架構師。

概述

隨著 Stable Diffsuion 的迅速走紅，引發了 AI 繪圖的時代變革。然而對於大部分人來說，訓練擴散模型的門檻太高，對 Stable Diffusion 進行全量微調也很難入手。由此，社區催生了一系列針對 Stable Diffusion 的高效微調方案，在保留原模型泛化能力的同時，實現自定義風格的融合，最關鍵的是，操作簡單且資源消耗量低。

本文將介紹 Stable Diffsuion 微調方案選型，以及如何使用 Dreambooth 和 LoRA 進行微調實踐，最後，我們會使用騰訊雲 TACO 對微調後的 Dreambooth 和 LoRA 模型進行推理優化。

Stable Diffusion 微調

Stable Diffusion 微調的目標，是將新概念注入預訓練模型，利用新注入的概念以及模型的先驗知識，基於文本引導條件生成自定義圖片。目前主流訓練 Stable Diffusion 模型的方法有 Full FineTune、Dreambooth、Text Inversion 和 LoRA，不同方法的實現邏輯和使用場景不同，選型簡單對比如下：

訓練方法	方法	侷限性
Text Inversion	使用提供的一組圖片訓練一個新單詞的Embedding ，並將其與詞彙表中的已有單詞關聯起來，這個新單詞即為這組圖片概念的指代。	訓練過程只對應 Embedding，擴散模型沒有新知識輸入，所以也無法產生新的內容。
Full FineTune	最樸素的方式，使用圖片+ 標註的數據集，進行迭代訓練，數據集標註可以選擇BLIP來生成。訓練直接對原模型的所有權重進行調整。	容易過擬合，導致生成圖片的多樣性不夠，結果難以控制。模型體積大，不便於傳播。
Dreambooth	提供代表某個新概念（instance）對應的一組圖像，並使用罕見字符（identifier）進行概念Mapping，訓練過程充分考慮原有相關主題（class）生成，避免過擬合。訓練直接對原模型的所有權重進行調整。	訓練過程只針對新概念（instance），多樣性差。如果需要多概念生成，需要多次訓練。模型體積大，不便於傳播。
LoRA（w Dreambooth）	凍結預訓練模型參數，在每個Transformer塊插入可訓練層，不需要完整調整 UNet 模型的全部參數。訓練結果只保留新增的網絡層，模型體積小。	訓練效果不如Dreambooth

需要注意的是，LoRA 是一種加速訓練的方法，Stable Diffusion 從大語言模型微調中借鑑而來，可以搭配 Full FineTune 或 Dreambooth 使用。針對上述幾種訓練方法，我們在 A10-24G 機型上進行測試，5-10張訓練圖片，所需資源和時長對比如下：

	Dreambooth	Text Inversion	LoRA
模型體積	~ 4 GB	3-10 KB	3-10 MB
訓練顯存	~ 22GB	~ 5GB	~ 20GB
訓練時長	15分鐘	1小時	10分鐘（w Dreambooth） 4小時（w Full FineTune）

接下來，我們重點介紹如何使用 Dreambooth 和 Lora（w Dreambooth）對 Stable Diffusion 模型進行微調。

Dreambooth

Dreambooth 用一個罕見字符（identifier）來代表訓練圖片的概念，對 UNet 模型的所有權重進行調整。這裡選擇罕見字符（identifier），是希望原模型沒有該 identifier 的先驗知識，否則容易在模型先驗和新注入概念（instance）間產生混淆。

對比 Full FineTune，雖然都會調整原模型的所有權重，但 Dreambooth 的創新點在於，它會使用 Stable Diffusion 模型去生成一個已有相關主題（class）的先驗知識，並在訓練中充分考慮原 class 和新 instance 的 prior preservation loss，從而避免新 instance 圖片特徵滲透到其他生成裡。

另外，訓練中加入一個已有的相關主題（class）的描述，可以將 instance 和 class 進行綁定，這樣新 instance 也可以使用到 class 對應的先驗知識。

我們使用 Huggingface 提供的訓練代碼，準備5-10張圖片，在A10上使用以下腳本啟動訓練：

accelerate launch train_dreambooth.py \
  --pretrained_model_name_or_path=$MODEL_NAME  \
  --instance_data_dir=$INSTANCE_DIR \
  --class_data_dir=$CLASS_DIR \
  --output_dir=$OUTPUT_DIR \
  --with_prior_preservation --prior_loss_weight=1.0 \
  --mixed_precision=fp16 \
  --instance_prompt="a photo of az baby" \
  --class_prompt="a photo of baby" \
  --resolution=512 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=1 \
  --learning_rate=5e-6 \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --num_class_images=200 \
  --max_train_steps=800

其中 –instance_data_dir 為新 instance 的圖片目錄，在 –instance_prompt 參數裡設置對應的 identifier，在 –class_prompt 設置相關 class 描述。訓練代碼

訓練集圖片示例：

訓練完畢後，輸入“a photo of az baby”，可以看到生成的圖片具備訓練集人物特徵。

訓練好的模型，如果需要在 Stable Diffusion Web UI 上使用，先通過腳本進行轉換，輸出ckpt或者safetensors格式，再放入 $HOME/stable-diffusion-webui/models/Stable-diffusion 目錄。腳本鏈接

python  ../scripts/convert_diffusers_to_original_stable_diffusion.py --model_path ./dreambooth_baby --checkpoint_path dreambooth_baby.safetensors --use_safetensors

LoRA（w Dreambooth）

LoRA（Low-Rank Adaptation of Large Language Models ）是一種輕量級的微調方法，通過少量的圖片訓練出一個小模型，然後和基礎模型結合使用，並通過插層的方式影響模型結果。

LoRA 的一個創新點，是通過“矩陣分解”的方式，優化插入層的參數量。我們可以將一個權重矩陣分解為兩個矩陣進行存儲，如果W是d*d維矩陣，那麼A和B矩陣的尺寸可以減小到d*n，這樣n遠小於d，大幅度減少存儲空間。

訓練會凍結預訓練模型的參數，通過 W’ = W +△W 的方式來調整模型參數，這裡的△W= ABT，其中AB矩陣就是我們的訓練目標。如下圖所示:

LoRA 的優勢在於生成的模型較小，訓練速度快，但推理需要同時使用 LoRA 模型和基礎模型。LoRA 模型雖然會向原有模型中插入新的網絡層，但最終效果還是依賴基礎模型。

我們使用 Huggingface 提供的訓練代碼，準備好圖片後，在A10上使用以下腳本啟動訓練：

accelerate launch train_dreambooth_lora.py \
  --pretrained_model_name_or_path=$MODEL_NAME  \
  --instance_data_dir=$INSTANCE_DIR \
  --class_data_dir=$CLASS_DIR \
  --output_dir=$OUTPUT_DIR \
  --instance_prompt="a photo of az baby" \
  --class_prompt="a photo of baby" \
  --resolution=512 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=1 \
  --checkpointing_steps=100 \
  --learning_rate=1e-4 \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --max_train_steps=1000 \
  --validation_prompt="a photo of az baby" \
  --validation_epochs=50 \
  --seed="0"

因為我們採用 Dreambooth-LoRA 方式進行訓練，所以超參數基本與前述的 Dreambooth 一致。訓練代碼

LoRA 輸出默認為 Pytorch 文件格式，如果需要在 Stable Diffusion Web UI 裡使用，先將模型轉化為 safetensors 格式，然後放入 $HOME/stable-diffusion-webui/models/Lora 目錄使用。腳本鏈接

python diffusers-lora-to-safetensors.py --file pytorch_lora_weights.bin

Stable Diffusion 性能優化

與訓練階段側重於準確預測標籤和提高模型精度不同，推理階段更看重高效處理輸入並生成預測結果，同時減少資源消耗，在一些應用場景裡，還會採用量化技術，在精度和性能之間取得平衡。

Stable Diffusion 是一個多模型組成的擴散Pipeline，由三個部分組成：變分自編碼器 VAE、UNet 和文本編碼器 CLIP。模型的推理耗時主要集中在 UNet，我們選擇對這部分進行優化，提高推理性能和效率。

目前社區和硬件廠商提供了多種優化方案，但這些方案接口定義複雜，使用門檻高，使得難以被廣泛採用。騰訊雲 TACO 只需簡單操作，即可實現 Stable Diffusion 推理優化，輕鬆應用只被少數專家掌握的技術。

騰訊雲 TACO 使用自研的編譯後端，對 UNet 模型以靜態圖方式進行編譯優化，同時根據不同的底層硬件，動態選擇 Codegen 優化策略，輸出更高效的機器代碼，提升推理速度，減少資源佔用。

Dreambooth 優化

複用訓練使用的 A10 GPU 服務器，參考TACO Infer 優化 Stable Diffusion 模型，安裝 Docker runtime，並拉取預置優化環境的 sd_taco:v3 鏡像。因涉及編譯生成機器碼，最終部署的目標 GPU 型號，需要和優化時的 GPU 型號保持一致。

使用-v命令掛載微調後的 Dreambooth diffusers 模型目錄，交互式啟動容器。

docker run -it --gpus=all --network=host -v /[diffusers_model_directory]:/[custom_container_directory] sd_taco:v3 bash

在鏡像裡執行 python export_model.py，採用 TorchScript tracing 生成序列化的 UNet 模型文件。

script_model = torch.jit.trace(model, test_data, strict=False)
script_model.save("trace_module.pt")

在鏡像裡執行 python demo.py，對導出的 UNet Model 進行性能優化。這一步 TACO sdk 會對導出的 IR 進行編譯優化，包括計算圖結構優化、算子優化、以及其他針對代碼生成和執行的優化技術。

完成後，使用 jit 方式加載優化後的 UNet Model。對模型輸入 a. 圖像隱空間向量【batchsize，隱空間通道，圖片高度/8，圖片寬度/8】b. timesteps值【batchsize】c. 【batchsize，文本最大編碼長度，向量大小】，即可對優化結果進行測試。代碼參考如下：

import torch
import taco
import os

taco_path = os.path.dirname(taco.__file__)
torch.ops.load_library(os.path.join(taco_path, "torch_tensorrt/lib/libtorchtrt.so"))
optimized_model = torch.jit.load("optimized_recursive_script_module.pt")

pic = torch.rand(1, 4, 64, 64).cuda() // picture
timesteps = torch.tensor([1]*1) // timesteps
context = torch.randn(1, 77, 768) // text embedding

with torch.no_grad():
    output = optimized_model(pic, timesteps, context)
    print(output)

對比社區方案，TACO 優化後模型出圖速度提高50%，效果見下圖：

（20 steps，Euler a，512 * 512，torch 1.12，無xformers，1s出圖）

LoRA 優化

使用 LoRA合併腳本，將訓練得到的 LoRA 文件，和基礎模型進行合併。命令參考：

python networks/merge_lora.py --sd_model ../v1-5-pruned-emaonly.safetensors --save_to ../lora-v1-5-pruned-emaonly.safetensors --models <LoRA文件目錄> --ratios <LoRA權重>

參考上述 Dreambooth 的優化方法，對合並後的模型進行導出和優化。效果見下圖：

（20 steps，Euler a，512 * 512，anime-tarot-card，torch 1.12，無xformers，1s出圖）

ControlNet 優化

Dreambooth 及 LoRA 優化模型，依然適用於 ControlNet 使用場景，對比社區方案，TACO 優化後 ControlNet 的出圖速度可以提高30%以上，效果見下圖：

（20 steps，Euler a，512 * 512，ControlNet-canny，torch 1.12，無xformers，2s出圖）

經過 TACO 優化後的 UNet 模型，測試表明前向推理速度提高至開源方案的4倍。在實際應用中，512*512，20 steps 的配置下，Stable Diffusion Web UI 端到端的推理時間縮短 1 秒。以上優化詳細過程及環境獲取，參考 TACO Infer 優化 Stable Diffusion 系列模型。

總結

本文介紹了 Dreambooth 和 LoRA 在騰訊雲A10機型上的微調實踐，以及針對這兩種模型的 TACO 推理優化過程。感興趣的同學可以在文章的基礎上，嘗試訓練風格獨特的模型，輔以 TACO 推理優化能力，創造符合自身業務的雲上 Stable Diffusion。

概述