Google Colab で はじめる Stable Diffusion v1.4 (3) - Textual Inversionによるファインチューニング

「Textual Inversion」による「Stable Diffusion」のファインチューニングに挑戦したので、その記録をまとめました。

RAMとGPUのメモリが必要なため「Colab (無料版)」では動きません。
「Colab Pro / Pro+」が必要です。

・Stable Diffusion v1.4
・rinongal/textual_inversion

前回

1. Textual Inversionによるファインチューニング

「Textual Inversion」は、3～5枚の画像を使って「Stable Diffusion」のファインチューニングを行うことで、オリジナルの画風やキャラクターを追加学習させることができる手法です。

An Image is Worth One Word: Personalizing Text-to-Image Generation using Textual Inversion

[Tutorial] "Fine Tuning" Stable Diffusion using only 5 Images Using Textual Inversion.

2. 学習用画像の準備

ファインチューニングには5枚ほどの学習用画像が必要です。

・画像サイズ は512×512。
・Exifで回転してない画像。
・最良の結果を得るには約 3～5個の画像。
　使用する画像の数が多すぎると、モデルが収束しない場合がある。
・画像には共通のコンテキスト情報が含まれている必要がある。
　スタイルベースでは、配色とアートスタイルが一貫していること。
　オブジェクトベースでは、同じオブジェクトを異なる視点から写すこと。

3. ファインチューニングの実行

Colabでファインチューニングを行う手順は、次のとおりです。

(1) 新規のColabノートブックを開き、「GPU」と「ハイメモリ」を選択。

# GPUの確認 
!nvidia-smi

(3) Googleドライブのマウントと作業フォルダへの移動。
あらかじめGoogleドライブのマイドライブ直下に「work」フォルダを作成しておきます。

# Googleドライブのマウントと作業フォルダへの移動
from google.colab import drive
drive.mount('/content/drive')
%cd '/content/drive/My Drive/work'

(4) パッケージのインストール。

# パッケージのインストール
!pip install -e git+https://github.com/CompVis/taming-transformers.git@master#egg=taming-transformers
!pip install pytorch_lightning tensorboard==2.8 omegaconf einops taming-transformers==0.0.1 clip transformers kornia test-tube

(5) 「Textual Inversionの」インストール。

# Textual Inversionのインストール
!git clone https://github.com/rinongal/textual_inversion.git
%cd textual_inversion
!pip install -e .
%mkdir checkpoint inputs outputs

(7) フォント「DejaVuSans.ttf」を「textual_inversion/data/」に配置。
以下のサイトで入手できます。

DejaVu Sans フォント

(8) Stable Diffusionのcheckpoint「sd-v1-4-full-ema.ckpt」を「textual_inversion/checkpoints/」に配置。
以下のサイトで入手できます。

CompVis/stable-diffusion-v-1-4-original · Hugging Face

(9) 「textual_inversion/configs/stable-diffusion/v1-finetune.yaml」の編集。
ファインチューニングの各種パラメータを調整します。

確認用画像の生成に時間がかかるので、batch_frequency(何ステップ毎に生成するか)を増やし、max_iamges(最大画像数)を減らして調整します。

・batch_frequency : 500 → 1000 (例)
・max_images : 8 → 1 (例)

GPU 20GB以下の場合は、そのままではメモリが足りないのでbatch_sizeとnum_workersを減らします。

・batch_size：2 → 1
・num_workers：2 → 1

画像サイズを512から256に減らすことで、学習精度は落ちると思いますが、さらにメモリと学習時間を抑えることができます。

・resolution : 512 → 256
・train: size : 512 → 256
・validation : size: 512 → 256

(10) main.pyで学習実行。
今回は、「Epoch 24」まで学習して停止しました。

# ファインチューニング
!python main.py \
    --no-test \
    --base ./configs/stable-diffusion/v1-finetune.yaml \
    -t \
    --actual_resume ./checkpoint/sd-v1-4-full-ema.ckpt \
    -n FINETUNE \
    --gpus 0, \
    --data_root ./inputs

主なパラメータは、次のとおりです。

・--train, -t : 学習を行う
・--no-test : テストを行わない
・--base : コンフィグファイルのパス (v1-finetune.yaml)
・--data_root : 入力画像フォルダのパス
・--actual_resume : モデルのパス (sd-v1-4-full-ema.ckpt)
・--name, -n : 学習の実行名 (logsフォルダ名に利用)
・--gpus : GPU

結果は、「logs/inputs<DATE>_<NAME/checkpoints」フォルダに出力されます。ckptファイルは学習再開、ptファイルは推論に利用します。

・embeddings_gs-XXXX.pt : XXXXステップの特徴ベクトルファイル
・epoch=XXXXX.ckpt : XXXXXエポックでの学習チェックポイント
・embeddeds.pt : 最終ステップ特徴ベクトルファイル
・last.ckpt : 最終エポックの学習チェックポイント

4. 推論の実行

Colabで推論を行う手順は、次のとおりです。

(1) stable_text2img.pyで推論実行。
embedding_pathには自分の環境のフォルダ名を指定してください。

!python scripts/stable_txt2img.py \
    --n_samples 1 \
    --n_iter 5 \
    --scale 10.0 \
    --ddim_eta 0.0 \
    --ddim_steps 50 \
    --embedding_path ./logs/inputs2022-09-01T11-06-38_sorami/checkpoints/embeddings.pt \
    --config ./logs/inputs2022-09-01T11-06-38_sorami/configs/*project.yaml \
    --ckpt  ./checkpoint/sd-v1-4-full-ema.ckpt \
    --prompt "super hero of * style"   

主なパラメータは、次のとおりです。

・--n_samples : バッチサイズ
・--n_iter : 出力画像数
・--scale : プロンプトに従う度合い (7〜11が良い)
・--ddim_eta : ddim eta (eta=0.0 は決定論的サンプリング)
・--ddim_steps : ddim サンプリング ステップ数
・--embedding_path : 特徴ベクトルファイルのパス
・--config : コンフィグファイルのパス (v1-inference.yaml)
・--ckpt : 事前学習モデルのパス (sd-v1-4-full-ema.ckpt)
・--prompt : プロンプト
・--seed : 乱数シード

outputフォルダに推論結果が出力されています。キャラクター（そらみとへにへに）を学習させたつもりが、画風を学習してた模様。

・super hero of * style

・cat of * style

・girl of * style

・robot of * style

5. その他のコンフィグファイルのパラメータ

その他のコンフィグファイル (v1-finetune.yaml) のパラメータで、重要そうなものをリストアップしてます。いろいろ変更して試すのはこれから。

◎ initializer_words
学習しているものを説明する単語のリストを指定します。学習開始時の特徴ベクトルとなり、学習によって特定の概念に最適化されます。

# 彫刻 (デフォルト)
initializer_words: ["sculpture"]

# 特定のスタイルの車
initializer_words: ["car","style","artistic"]

1つの単語を使用したい場合は、main.pyの引数でも指定できます。

--init_word [SINGLE_WORD]

◎ placeholder_strings
学習したものを表す新しい単語を指定します。デフォルトは"*"です。

placeholder_strings: ["*"]

既存の単語と衝突すると実行時にエラーになります。<hoge>のように<>を付加するのが推奨されています。

AssertionError: String 'sirachan' maps to more than a single token. Please use another string

6. 参考

birdManさんの記事をかなり参考にさせてもらいました。

Stable Diffusionをいらすとやでファインチューニングする

Stable Diffusionファインチューニング用のColab

以下の記事もかなり参考にさせてもらってます。

How to Fine-tune Stable Diffusion using Textual Inversion

関連

BERT/GPT-3/DALL-E 自然言語処理・画像処理・音声処理 人工知能プログラミング実践入門