定制屬于你自己的 Stable Diffusion

作者：二旺 2024-09-14 14:09:40

在本文中，我想分享如何使用現(xiàn)代技術(shù)LoRA逐步微調(diào)Stable Diffusion模型的指南。整個(gè)練習(xí)在Kaggle筆記本上進(jìn)行，這樣每個(gè)人都可以在沒有適當(dāng)設(shè)備的情況下完成這個(gè)練習(xí)。

我們很幸運(yùn)生活在AI時(shí)代。人工智能迅速發(fā)展，現(xiàn)在變得更加易于獲取。大多數(shù)開源模型都經(jīng)過了充分的訓(xùn)練，擁有來自廣泛類別的合理數(shù)量的數(shù)據(jù)，因此我們可以生成許多不同類別的圖像。如果模型從未學(xué)習(xí)過我們想要生成的特定主題怎么辦？例如，我想要生成我今天早上剛做的蛋糕的圖像。盡管模型可以生成無數(shù)種蛋糕圖像，但模型不知道如何生成我剛剛制作的蛋糕的圖像。

盡管從頭開始訓(xùn)練對(duì)個(gè)人來說幾乎是不可能的，但由于不斷發(fā)展的模型微調(diào)技術(shù)，將新主題納入模型可以輕松實(shí)現(xiàn)。自2022年以來，Stable Diffusion作為最新趨勢(shì)的一代模型，為圖像生成開啟了一個(gè)全新的時(shí)代。在本文中，我想分享如何使用現(xiàn)代技術(shù)LoRA逐步微調(diào)Stable Diffusion模型的指南。整個(gè)練習(xí)在Kaggle筆記本上進(jìn)行，這樣每個(gè)人都可以在沒有適當(dāng)設(shè)備的情況下完成這個(gè)練習(xí)。

在我們開始實(shí)現(xiàn)之前，讓我們快速了解一下Stable Diffusion的概念和微調(diào)技術(shù)LoRA。如果你已經(jīng)熟悉背后的知識(shí)，請(qǐng)隨時(shí)跳到實(shí)施部分。

Stable Diffusion

圖1. Stable Diffusion中的反向擴(kuò)散

Stable Diffusion是一種可以從文本和圖像生成逼真照片的生成性AI模型。擴(kuò)散模型在潛在空間而不是圖像空間中施展魔法，這使得圖像生成對(duì)公眾來說更加可行，因?yàn)樗璧挠?jì)算能力更低。人們甚至可以在CPU機(jī)器上的小憩期間生成一個(gè)不存在的圖像。

有許多在線的穩(wěn)定擴(kuò)散平臺(tái)供公眾和開發(fā)者使用，如HuggingFace1和Stable Diffusion Online2。我們可以在沒有機(jī)器學(xué)習(xí)知識(shí)的情況下生成圖像。

LoRA（低秩適應(yīng)）

圖2. Stable Diffusion網(wǎng)絡(luò)

LoRA是一種微調(diào)方法，它在穩(wěn)定擴(kuò)散網(wǎng)絡(luò)的交叉注意力層中添加了額外的權(quán)重。交叉注意力層在圖2中由黃色塊表示。交叉注意力層融合了圖像ZT和文本τθ的中間信息。向交叉注意力層添加權(quán)重?cái)U(kuò)展了文本提示和圖像之間的相關(guān)性。這就是LoRA如何向穩(wěn)定擴(kuò)散模型引入新知識(shí)。

修改后的權(quán)重被分解為更小的（低秩）矩陣。這些較小的矩陣攜帶的參數(shù)更少，并單獨(dú)存儲(chǔ)。這就是為什么LoRA權(quán)重文件的尺寸可控，但必須與底層的Stable Diffusion模型一起使用。

LoRA的權(quán)重大小通常在2到200 MB之間。這是共享和管理模型的優(yōu)勢(shì)之一。與其以幾GB的整個(gè)模型共享，更有效的方法是以更小的尺寸共享LoRA權(quán)重。

環(huán)境

這個(gè)練習(xí)是在Kaggle筆記本中進(jìn)行的，環(huán)境設(shè)置如下：

NVIDIA Tesla P100 (RAM=16G)
diffusers-0.30.0.dev0
Pytorch 2.1.2
Python 3.10

Kaggle筆記本已經(jīng)安裝了Python和PyTorch。我們需要安裝的唯一庫(kù)是來自HuggingFace的Diffusers。我們可以使用以下命令從源安裝Diffusers。從源安裝確保了庫(kù)的最新版本。路徑/kaggle/working是Kaggle運(yùn)行時(shí)的默認(rèn)工作空間。人們可以調(diào)整路徑以適應(yīng)自己的開發(fā)環(huán)境。

%cd /kaggle/working/
!pip install accelerate
!git clone https://github.com/huggingface/diffusers%cd /kaggle/working/diffusers
!pip install /kaggle/working/diffusers/. -q
%cd /kaggle/working/diffusers/examples/text_to_image
!pip install -r /kaggle/working/diffusers/examples/text_to_image/requirements.txt -q
!pip install --upgrade peft transformers xformers bitsandbytes -q

實(shí)驗(yàn)在這個(gè)練習(xí)中，我想向模型介紹蠟筆小新雕像。讓我們?cè)谝韵赂鞴?jié)中經(jīng)歷步驟。

第1步：數(shù)據(jù)收集

圖3. 訓(xùn)練樣本

首先，我們需要收集一些主題的圖像。在這個(gè)練習(xí)中，我們收集了25張蠟筆小新雕像的照片作為我們的訓(xùn)練數(shù)據(jù)。一些樣本可以在圖3中找到。以下是我遵循的數(shù)據(jù)收集推薦提示：

統(tǒng)一圖像大小：例如，512x512。
高圖像質(zhì)量：圖片以12M分辨率拍攝。
保持主題在圖像中心。

此外，訓(xùn)練集中盡可能多地覆蓋多樣性也很重要，以便模型學(xué)習(xí)一般特征，如角色的眼睛、發(fā)型和身材。換句話說，我們希望模型學(xué)習(xí)什么是蠟筆小新雕像，而不是特定著裝風(fēng)格（例如紅襯衫和黃色短褲）和姿勢(shì)（例如站立）。

我已經(jīng)在5、15、25的培訓(xùn)規(guī)模上進(jìn)行了實(shí)驗(yàn)。盡管每一輪之間只有10張圖像的差異，但模型的泛化程度有了顯著提高。

第2步：基線評(píng)估

from diffusers import AutoPipelineForText2Image
pipeline = AutoPipelineForText2Image.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16).to("cuda")
prompt = "A crayon shin-chan figurine."
image = pipeline(prompt, num_inference_steps=100).images[0]
image.save("baseline.png")

在開始訓(xùn)練之前，我們想檢查一下預(yù)訓(xùn)練模型是否已經(jīng)知道我們想要微調(diào)的特定雕像，使用上述Python腳本。

AutoPipelineForText2Image類用于加載Stable Diffusion預(yù)訓(xùn)練模型stable-diffusion-v1–5。
接下來，我們將文本提示設(shè)置為“A crayon shin-chan figurine”。我保持文本提示簡(jiǎn)潔，以便文本信息對(duì)模型來說直接了當(dāng)。
將提示和推理次數(shù)傳遞給AutoPipelineForText2Image的實(shí)例并進(jìn)行生成過程。

圖4. 基線評(píng)估結(jié)果

腳本多次運(yùn)行，圖4中展示了三個(gè)生成的圖像。結(jié)果顯示，預(yù)訓(xùn)練模型知道可能是一個(gè)人角色的雕像，但不知道前綴‘crayon shin-chan’具體指的是哪個(gè)角色。生成圖像中的限制驗(yàn)證了我們微調(diào)工作的努力，因?yàn)槲覀兊某h是用一個(gè)未見過的主體來微調(diào)模型。

第3步：微調(diào)

圖5. 蠟筆

定義一個(gè)新的標(biāo)識(shí)符

除了收集一些訓(xùn)練圖像外，我們還需要為這個(gè)主題選擇一個(gè)新的名稱。每當(dāng)新名稱出現(xiàn)在文本提示中時(shí)，修改后的模型就會(huì)理解它指的是新的主題，并相應(yīng)地生成內(nèi)容。

我們?yōu)槭裁葱枰粋€(gè)新的名稱？原始名稱Crayon Shin-Chan不是更具代表性嗎？我們這樣做的想法是，我們需要一個(gè)獨(dú)特的名稱，不要與已經(jīng)訓(xùn)練的其他現(xiàn)有名稱混在一起。在我們的例子中，內(nèi)部的蠟筆暗示了與我們用于繪畫的蠟筆的潛在聯(lián)系。如果關(guān)聯(lián)的文本提示包含Crayon Shin-Chan，生成的圖像可能包括與蠟筆相關(guān)的內(nèi)容。

建議使用隨機(jī)名稱作為新標(biāo)識(shí)符。我為這個(gè)練習(xí)選擇了aawxbc這個(gè)名字，它幾乎不會(huì)與任何現(xiàn)有名稱重疊。其他人也可以為新主題使用任何其他獨(dú)特的名稱。

訓(xùn)練

圖6. 模型生成的類別圖像

在Kaggle筆記本中，除了工作區(qū)：/kaggle/working/，另一個(gè)常用的目錄是/kaggle/input/，我們?cè)谶@里存儲(chǔ)工作的輸入，如檢查點(diǎn)和訓(xùn)練圖像。以下是基于腳本train_dreambooth_lora.py的完整訓(xùn)練命令。

!accelerate config default
!mkdir /kaggle/working/figurine

!accelerate launch --mixed_precision="fp16"  /kaggle/working/diffusers/examples/dreambooth/train_dreambooth_lora.py 
--mixed_precision="fp16" 
--pretrained_model_name_or_path="runwayml/stable-diffusion-v1-5" 
--instance_data_dir="/kaggle/input/aawxbc" 
--instance_prompt="a photo of aawxbc figurine" 
--class_data_dir="/kaggle/working/figurine" 
--class_prompt="a photo of figurine" 
--with_prior_preservation 
--prior_loss_weight=1.0 
--use_8bit_adam 
--lr_scheduler="constant" 
--gradient_checkpointing 
--enable_xformers_memory_efficient_attention 
--num_train_epochs=50 
--output_dir="/kaggle/working/dreambooth-lora"

讓我們經(jīng)歷一些關(guān)鍵配置：

mixed_precision [可選]：選擇平臺(tái)支持的精度。選項(xiàng)有fp32、fp16和bf16。
class_prompt [可選]：存儲(chǔ)類別數(shù)據(jù)的提示和文件夾。在這個(gè)實(shí)驗(yàn)中，類別名稱是figurine，因?yàn)槲覀兿Ｍ麛U(kuò)散模型學(xué)習(xí)一種新的雕像類型。
class_data_dir [可選]：存儲(chǔ)模型生成的類別圖像的目錄。在我們的例子中，模型生成了一般的雕像圖像并將其保存在這個(gè)目錄下。
instance_prompt [必須]：我們?yōu)樾轮黝}定義了一個(gè)文本提示。格式是{新名稱} {類別名稱}，例如aawxbc figurine、uvhhhl cat等。
instance_data_dir [必須]：存儲(chǔ)訓(xùn)練樣本的目錄。在這個(gè)練習(xí)中，我們將前面提到的25個(gè)樣本上傳到這個(gè)目錄。
num_train_epochs [可選]：訓(xùn)練的輪數(shù)。

通過遵循上述步驟，包括數(shù)據(jù)收集、安裝庫(kù)和完成訓(xùn)練，我們現(xiàn)在可以在由output_dir定義的目錄中找到LoRA權(quán)重pytorch_lora_weights.safetensors。權(quán)重文件的大小為3.2 MB，相當(dāng)容易管理。

第4步：推理—生成新圖像

from diffusers import AutoPipelineForText2Image
pipeline = AutoPipelineForText2Image.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16).to("cuda")
pipeline.load_lora_weights("/kaggle/working/dreambooth-lora", weight_name="pytorch_lora_weights.safetensors")
prompt = "A photo of aawxbc figurine."
image = pipeline(prompt, num_inference_steps=180).images[0]
image.save("aawxbc.png")

在成功導(dǎo)出LoRA權(quán)重后，我們可以進(jìn)入最后也是最令人興奮的部分：生成新圖像。請(qǐng)注意，LoRA權(quán)重不能獨(dú)立工作，因此我們需要同時(shí)加載原始的穩(wěn)定擴(kuò)散和LoRA權(quán)重文件。

在Python腳本中，除了我們?cè)诨€評(píng)估中使用的步驟外，我們還有一個(gè)額外的步驟，使用.load_lora_weights()函數(shù)加載LoRA模型，傳入目錄和LoRA權(quán)重文件的名稱。