Stable Diffusion 転移学習

本記事では、Stable Diffusionの転移学習の進め方を紹介します！
GPU環境をお持ちの方は試してみてください🌟

初期設定

環境

• Linux 20.0.4
• GPU
  • NVIDIA GeForce
  • 20GB以上 推奨
• CUDA 11.8
• NVIDIA Driver 520.61.05
• Anaconda 4.12.0
  • Python

Anaconda・NVIDIAの導入は今回は扱いませんので
以下の記事を参考にしてください。

導入

textual_inversionのレポジトリからクローンします。

ターミナル

git clone https://github.com/rinongal/textual_inversion.git

次に、プロジェクトに移動後、Anacondaで仮想環境を作成します。

ターミナル

conda env create -f environment.yaml
conda activate ldm
pip install -e .

最後に、使用するモデルをhugging-faceからダウンロードしておきます。

↑のサイトでsd-v1-4.ckptをクリックすればダウンロードが始まります！
（ログインしていないとダウンロードできない可能性があります）
そして、クローンしたtextual_inversionの直下にモデルを配置します。

これで導入はできました！

学習

コマンドで実行

ターミナル

python main.py --base configs/stable-diffusion/v1-finetune.yaml
               -t 
               --actual_resume sd-v-1-4.ckpt 
               -n st_train
               --gpus 1, 
               --data_root 

ここから学習が実行されるので、所要時間はとても長くなります。
筆者の場合、

• 学習させるデータ数74
• --gpus 3
• パラメータの設定は下の項目の内容

で実行したところ、大体45分くらいで終わりました！

ランタイムエラー

RuntimeError: CUDA out of memory
GPUのメモリ容量を超えていることで起きるエラーです。
解決方法は主に以下の３点です

1. GPUのスケールアップ
2. パラメータの変更
3. 処理する画像サイズなど学習させるデータの変更

今回は2番、パラメータの変更で対応します。

v1-finetune.yaml

# 変更点がないため、上部は省略
data:
  target: main.DataModuleFromConfig
  params:
-    batch_size: 2
+    batch_size: 1
-    num_workers: 2
+    num_workers: 1
    wrap: false
    train:
      target: ldm.data.personalized.PersonalizedBase
      params:
-        size: 512
+        size: 256
        set: train
        per_image_tokens: false
        repeats: 100
    validation:
      target: ldm.data.personalized.PersonalizedBase
      params:
-        size: 512
+        size: 256
        set: val
        per_image_tokens: false
        repeats: 10

# 変更点がないため、下部は省略

Pytorchのエラー No test_dataloader()

pytorch_lightning.utilities.exceptions.MisconfigurationException: No test_dataloader() method defined to run Trainer.test

814行目付近をコメントアウトする

main.py

#        if not opt.no_test and not trainer.interrupted:
#            trainer.test(model, data)

実行してみる

ターミナル

python scripts/stable_txt2img.py 
    --ddim_eta 0.0 --n_samples 1 --n_iter 2 --scale 10.0 --ddim_steps 50 
    --embedding_path ./logs/最新のもの/checkpoints/embeddings.pt 
    --ckpt ./sd-v1-4.ckpt 
    --config ./logs/最新のもの/configs/最新のもの-project.yaml 
    --prompt "任意の文字列"

このコマンドで実行できます！
実行すると、outputsディレクトリ内に画像が生成されます！
これで転移学習自体は完了です🙌

精度に関して

redditの記載

redditでは512x512のサイズの画像を5枚程度で学習に用いると精度が上がるそうです。

Then, it's preferred to get 5 images of your subject at 512x512 resolution. From the paper, 5 images are the optimal amount for textual inversion. On a single V100, training should take about two hours give or take. More images will increase training time, and may or may not improve results. You are free to test this and let us know how it goes!

ただ、確実に上がるかどうかは保証してくれないので
どのくらいの枚数が良いか、は試すしかなさそうです、、、😅
試してみて、いい結果が出たら共有します🙇

参考文献