該存儲庫提供了用於微調穩定擴散的代碼。它是通過擁抱臉來從這個腳本改編的。用於微調的預培訓模型來自Kerascv。要了解原始型號,請查看此文檔。
此存儲庫中提供的代碼僅用於研究目的。請查看本節,以了解有關潛在用例和限制的更多信息。
通過加載此模型,您可以在https://raw.githubusercontent.com/compvis/stable-diffusion/main/main/license上接受CreatiVeml Open Rail-M許可證。
如果您只是在尋找此存儲庫的隨附資源,則是以下鏈接:
目錄:
該存儲庫有一個姐妹存儲庫(Keras-SD服務),涵蓋了穩定擴散的各種部署模式。
2023年1月13日更新:該項目在Google組織的有史以來的首次Keras社區獎比賽中獲得了第二名。
按照擁抱面的原始腳本,此存儲庫還使用Pokemon數據集。但是它被重生以與tf.data更好地工作。數據集的再生版本在此處託管。查看該鏈接以獲取更多詳細信息。
finetune.py提供了微調代碼。在進行培訓之前,請確保您安裝了依賴關係(請參閱requirements.txt )。
您可以通過運行python finetune.py來啟動默認參數培訓。運行python finetune.py -h以了解支持的命令行參數。您可以通過傳遞--mp標誌來啟用混合精確培訓。
當您啟動訓練時,只有噹噹前損失低於前面的損失時,才會生成擴散模型檢查點。
為了避免OOM和更快的訓練,建議至少使用V100 GPU。我們使用了A100。
需要注意的一些重要細節:
培訓詳細信息:
我們對兩種不同的分辨率進行了微調:256x256和512x512。我們僅通過這兩種不同的分辨率改變了批處理大小和時代的數量。由於我們沒有使用梯度積累,因此我們使用此代碼段來得出時期的數量。
python finetune.py --batch_size 4 --num_epochs 577python finetune.py --img_height 512 --img_width 512 --batch_size 1 --num_epochs 72 --mp對於256x256分辨率,我們故意減少了時期的數量以節省計算時間。
微調重量:
您可以在此處找到微調的擴散模型權重。
此存儲庫中使用的默認寵物小精靈數據集附帶以下結構:
pokemon_dataset/
data.csv
image_24.png
image_3.png
image_550.png
image_700.png
... data.csv看起來像:
只要您的自定義數據集遵循此結構,您就無需更改當前代碼庫中的任何內容,除了dataset_archive 。
如果您的數據集每個圖像具有多個字幕,則可以在訓練期間從每個圖像的標題池中隨機選擇一個字幕。
根據數據集,您可能必須調整超參數。
import keras_cv
import matplotlib . pyplot as plt
from tensorflow import keras
IMG_HEIGHT = IMG_WIDTH = 512
def plot_images ( images , title ):
plt . figure ( figsize = ( 20 , 20 ))
for i in range ( len ( images )):
ax = plt . subplot ( 1 , len ( images ), i + 1 )
plt . title ( title )
plt . imshow ( images [ i ])
plt . axis ( "off" )
# We just have to load the fine-tuned weights into the diffusion model.
weights_path = keras . utils . get_file (
origin = "https://huggingface.co/sayakpaul/kerascv_sd_pokemon_finetuned/resolve/main/ckpt_epochs_72_res_512_mp_True.h5"
)
pokemon_model = keras_cv . models . StableDiffusion (
img_height = IMG_HEIGHT , img_width = IMG_WIDTH
)
pokemon_model . diffusion_model . load_weights ( weights_path )
# Generate images.
generated_images = pokemon_model . text_to_image ( "Yoda" , batch_size = 3 )
plot_images ( generated_images , "Fine-tuned on the Pokemon dataset" )您可以帶上weights_path (應該與diffusion_model兼容)並重複使用代碼片段。
查看此COLAB筆記本以播放推理代碼。
最初,我們以256x256的分辨率微調了該模型。以下是一些結果以及與原始模型的結果進行比較。
| 圖像 | 提示 |
|---|---|
| 尤達 | |
| 機器人貓與翅膀 | |
| Hello Kitty |
我們可以看到,微調模型比原始模型具有更穩定的輸出。即使在美學上可以改善結果,但可以看到微調效果。另外,我們從擁抱面孔的腳本進行256x256分辨率(除了時期和批處理大小)中遵循了相同的超參數。使用更好的超參數,結果可能會有所改善。
對於512x512分辨率,我們觀察到類似的東西。因此,我們嘗試了unconditional_guidance_scale參數,並註意到將其設置為40(在確定其他參數時)時,結果會更好。
| 圖像 | 提示 |
|---|---|
| 尤達 | |
| 機器人貓與翅膀 | |
| Hello Kitty |
注意:在512x512上進行微調仍在進行中。但是,在沒有分佈式訓練和梯度積累的情況下,完成一個時代需要大量時間。以上結果來自60個時期後得出的檢查點。
Lambda Labs有了類似的食譜(但經過培訓以獲得更優化的步驟),展示了令人驚嘆的結果。