使用 diffusers 訓練你的 ControlNet 🧨

簡介

ControlNet 是一種神經網路結構，它透過新增額外的條件來實現對擴散模型的細粒度控制。這項技術首次出現在論文 《為文字到影像擴散模型新增條件控制》(Adding Conditional Control to Text-to-Image Diffusion Models) 中，並迅速在開源擴散社群中流行起來，作者釋出了 8 種不同的條件來控制 Stable Diffusion v1-5，包括姿態估計、深度圖、Canny 邊緣、草圖，以及更多。

在這篇博文中，我們將詳細介紹我們是如何訓練 詭異人臉模型 (Uncanny Faces model) 的每一步——這是一個基於 3D 合成人臉的面部姿勢模型 (詭異人臉實際上是一個意想不到的結果，請繼續關注它是如何產生的)。

開始為 Stable Diffusion 訓練你的 ControlNet

訓練你自己的 ControlNet 需要 3 個步驟

1. 規劃你的條件: ControlNet 的靈活性足以讓 Stable Diffusion 適應許多工。預訓練模型展示了廣泛的條件，社群也構建了其他條件，例如基於畫素化調色盤的條件。

2. 構建你的資料集: 一旦確定了條件，就該構建資料集了。為此，你可以從頭開始構建資料集，也可以使用現有資料集的子集。你需要資料集中的三列來訓練模型：一個基準真相 image、一個 conditioning_image 和一個 prompt。

3. 訓練模型: 資料集準備好後，就可以訓練模型了。感謝 diffusers 訓練指令碼，這是最簡單的部分。你需要一個至少有 8GB VRAM 的 GPU。

1. 規劃你的條件

為了規劃你的條件，思考以下兩個問題會很有幫助

1. 我想使用什麼樣的條件？
2. 是否已經有現成的模型可以將“常規”影像轉換為我的條件？

在我們的例子中，我們考慮使用面部關鍵點作為條件。我們的理由是：1. 基於通用關鍵點條件的 ControlNet 效果很好。2. 面部關鍵點是一種足夠普及的技術，並且有多種模型可以在常規圖片上計算面部關鍵點。3. 讓 Stable Diffusion 遵循特定的面部關鍵點或模仿你自己的面部表情可能會很有趣。

2. 構建你的資料集

好的！我們決定使用面部關鍵點作為 Stable Diffusion 的條件。所以，為了準備資料集，我們需要

• 基準真相 image：在這裡是人臉影像
• conditioning_image：在這裡是可視化了面部關鍵點的影像
• caption：描述所用影像的標題

對於這個專案，我們決定使用微軟的 FaceSynthetics 資料集：它是一個包含 10 萬張合成人臉的資料集。其他包含真實人臉的面部研究資料集，如 Celeb-A HQ、FFHQ - 但我們決定在這個專案中使用合成人臉。

FaceSynthetics 資料集聽起來是一個很好的起點：它包含了人臉的基準真相影像，以及以 iBUG 68 點面部關鍵點格式標註的關鍵點，還有人臉的分割影像。

完美。是嗎？不幸的是，並非如此。還記得在“規劃你的條件”步驟中的第二個問題嗎——我們應該有模型能將常規影像轉換為條件影像？結果發現，沒有已知的模型可以將人臉轉換成該資料集的標註關鍵點格式。

所以我們決定走另一條路

• 使用 FaceSynthetics 資料集中的人臉基準真相 image
• 使用一個已知的模型，該模型可以將任何一張人臉影像轉換為 iBUG 的 68 點面部關鍵點格式 (在我們的例子中，我們使用了 SOTA 模型 SPIGA)
• 使用自定義程式碼將面部關鍵點轉換為一個美觀的插圖蒙版，用作 conditioning_image
• 將其儲存為 Hugging Face 資料集

你可以在這裡找到用於將 FaceSynthetics 資料集中的基準真相影像轉換為插圖蒙版並將其儲存為 Hugging Face 資料集的程式碼。

現在，資料集中有了基準真相 image 和 conditioning_image，我們還差一步：為每張圖片新增標題。這一步是強烈推薦的，但你也可以嘗試使用空提示詞，並反饋你的結果。由於我們沒有 FaceSynthetics 資料集的標題，我們透過 BLIP 影像描述模型 為其生成了標題。你可以在這裡檢視為所有影像生成標題的程式碼。

這樣，我們得到了最終的資料集！帶標題的 Face Synthetics SPIGA (Face Synthetics SPIGA with captions) 資料集包含了 FaceSynthetics 資料集中 10 萬張影像的基準真相影像、分割圖和標題。我們準備好訓練模型了！

3. 訓練模型

我們的資料集準備好了，是時候訓練模型了！雖然這本應是整個過程中最難的部分，但有了 diffusers 訓練指令碼，它反而成了最簡單的部分。我們使用了一臺在 LambdaLabs 上租用的 A100，每小時 1.10 美元。

我們的訓練經驗

我們對模型進行了 3 個 epoch 的訓練 (這意味著 10 萬張圖片的批次被模型看了 3 次)，批處理大小為 4 (每一步向模型展示 4 張圖片)。結果證明這樣做太多了並且導致了過擬合 (所以它忘記了那些與真實人臉稍有不同的概念，例如在提示詞中輸入“shrek”或“a cat”不會生成史萊克或貓，而是一個人，並且也開始忽略風格)。

僅用 1 個 epoch (即模型“看過”10 萬張圖片後)，它就已經收斂到能夠遵循姿勢並且沒有過擬合。所以它成功了，但是...由於我們使用了 FaceSynthetics 資料集，模型最終學會了生成詭異的 3D 臉，而不是逼真的人臉。考慮到我們使用的是合成人臉資料集而不是真實人臉，這是合理的，並且可以用於有趣/梗圖的目的。這是 uncannyfaces_25K 模型。

在這個互動式表格中，你可以透過下面的撥盤來檢視模型經歷了多少訓練步驟，以及這對訓練過程有何影響。在大約 1.5 萬步時，它已經開始學習姿勢。在大約 2.5 萬步時，它已經成熟。

我們是如何進行訓練的

我們所要做的就是，安裝依賴項

pip install git+https://github.com/huggingface/diffusers.git transformers accelerate xformers==0.0.16 wandb
huggingface-cli login
wandb login 

然後執行 train_controlnet.py 程式碼

!accelerate launch train_controlnet.py \
 --pretrained_model_name_or_path="stabilityai/stable-diffusion-2-1-base" \
 --output_dir="model_out" \
 --dataset_name=multimodalart/facesyntheticsspigacaptioned \
 --conditioning_image_column=spiga_seg \
 --image_column=image \
 --caption_column=image_caption \
 --resolution=512 \
 --learning_rate=1e-5 \
 --validation_image "./face_landmarks1.jpeg" "./face_landmarks2.jpeg" "./face_landmarks3.jpeg" \
 --validation_prompt "High-quality close-up dslr photo of man wearing a hat with trees in the background" "Girl smiling, professional dslr photograph, dark background, studio lights, high quality" "Portrait of a clown face, oil on canvas, bittersweet expression" \
 --train_batch_size=4 \
 --num_train_epochs=3 \
 --tracker_project_name="controlnet" \
 --enable_xformers_memory_efficient_attention \
 --checkpointing_steps=5000 \
 --validation_steps=5000 \
 --report_to wandb \
 --push_to_hub

讓我們來分析一些設定，並介紹一些最佳化技巧，以便在低至 8GB VRAM 的顯示卡上進行訓練。

• pretrained_model_name_or_path: 你想使用的 Stable Diffusion 基礎模型 (我們在這裡選擇了 v2-1，因為它能更好地渲染人臉)
• output_dir: 你希望模型儲存的目錄
• dataset_name: 用於訓練的資料集。在我們的例子中是 帶標題的 Face Synthetics SPIGA
• conditioning_image_column: 資料集中包含條件影像的列名 (在我們的例子中是 spiga_seg)
• image_column: 資料集中包含基準真相影像的列名 (在我們的例子中是 image)
• caption_column: 資料集中包含該影像標題的列名 (在我們的例子中是 image_caption)
• resolution: 條件影像和基準真相影像的解析度 (在我們的例子中是 512x512)
• learning_rate: 學習率。我們發現 1e-5 在這些例子中效果很好，但你可以嘗試不同的值，例如在 1e-4 和 2e-6 之間。
• validation_image: 這可以讓你在訓練過程中進行預覽！驗證影像會根據 validation_steps 的設定定期執行，這樣你就可以看到訓練進展如何。在這裡插入任意數量的條件影像的本地路徑。
• validation_prompt: 與你的驗證影像一起執行的提示詞。可以是任何可以測試模型訓練效果的內容。
• train_batch_size: 這是為了適應 GPU 而設定的訓練批次大小。由於我們有 A100，所以可以負擔得起 4，但如果你的 GPU VRAM 較低，我們建議將此值降低到 1。
• num_train_epochs: 每個 epoch 對應訓練集中的影像被模型“看到”的次數。我們實驗了 3 個 epoch，但結果表明最好的結果只需要略多於 1 個 epoch，3 個 epoch 會導致我們的模型過擬合。
• checkpointing_steps: 每 x 步儲存一箇中間檢查點 (在我們的例子中是 5000)。每 5000 步，就會儲存一箇中間檢查點。
• validation_steps: 每 x 步執行一次 validation_prompt 和 validation_image。
• report_to: 向哪裡報告你的訓練情況。這裡我們使用了 Weights and Biases，它給了我們這份漂亮的報告。但是，將 train_batch_size 從 4 減少到 1 可能不足以讓訓練在小型 GPU 上執行，這裡是針對不同 GPU VRAM 大小需要新增的一些額外引數。
• push_to_hub: 一個將最終訓練好的模型推送到 Hugging Face Hub 的引數。

在 16GB VRAM GPU 上進行適配

pip install bitsandbytes

--train_batch_size=1 \
--gradient_accumulation_steps=4 \
--gradient_checkpointing \
--use_8bit_adam

批處理大小為 1 和 4 個梯度累積步驟的組合，等效於我們示例中使用的原始批處理大小 4。此外，我們啟用了梯度檢查點和 8 位 Adam 最佳化器以節省更多記憶體。

在 12GB VRAM GPU 上進行適配

--gradient_accumulation_steps=4 \
--gradient_checkpointing \
--use_8bit_adam
--set_grads_to_none

在 8GB VRAM GPU 上進行適配

請遵循我們這裡的指南

4. 總結！

這次訓練 ControlNet 的經歷非常有趣。我們成功訓練了一個可以遵循真實人臉姿勢的模型——然而，它學會了生成詭異的 3D 人臉，而不是逼真的人臉，因為這是它訓練時使用的資料集，這也有它自己的魅力和風格。

試試我們的 Hugging Face Space

至於我們的下一步——為了建立看起來逼真的人臉，同時仍然不使用真實人臉資料集，一個想法是執行整個 FaceSynthetics 資料集透過 Stable Diffusion Image2Image，將 3D 臉轉換為逼真的臉，然後訓練另一個 ControlNet。

請繼續關注，因為我們很快將舉辦一個 ControlNet 訓練活動！在 Twitter 上關注 Hugging Face 或加入我們的 Discord 以獲取最新訊息。