[Stable Diffusion] LoRAの学習手順

LoRAの学習をやってみたので、その手順を掲載。

関連記事

目次

学習環境構築

LoRAの学習には、Kohya氏によるPythonプログラム「sd-scripts」を利用する。
sd-scriptsを更に扱いやすくするための補助ツールも幾つか公開されている。

今回の手順では、補助ツールは使わずにsd-scriptsを直接使うことにする。

sd-scriptsのインストール

適当なフォルダーを作り、sd-scripts公式のインストールガイドに従ってその中に環境を構築する。その手順を以下でも解説する。

PowerShellのセキュリティ設定の変更

PowerShellはデフォルトではスクリプトを実行することができないので設定を変更する必要がある。

PowerShellを管理者として実行。
「Set-ExecutionPolicy Unrestricted」を実行してYと入力する。

sd-scriptsのインストール

以下を1行ずつ実行。

(PowerShellの場合)

git clone https://github.com/kohya-ss/sd-scripts.git
cd sd-scripts

python -m venv venv
.\venv\Scripts\activate

pip install torch==1.12.1+cu116 torchvision==0.13.1+cu116 --extra-index-url https://download.pytorch.org/whl/cu116
pip install --upgrade -r requirements.txt
pip install -U -I --no-deps https://github.com/C43H66N12O12S2/stable-diffusion-webui/releases/download/f/xformers-0.0.14.dev0-cp310-cp310-win_amd64.whl

cp .\bitsandbytes_windows\*.dll .\venv\Lib\site-packages\bitsandbytes\
cp .\bitsandbytes_windows\cextension.py .\venv\Lib\site-packages\bitsandbytes\cextension.py
cp .\bitsandbytes_windows\main.py .\venv\Lib\site-packages\bitsandbytes\cuda_setup\main.py

accelerate config

accelerate configの設定ではそれぞれ以下のように入力。

  • This machine (そのままEnterキーで決定)
  • No distributed training (そのままEnterキーで決定)
  • NO (「NO」と入力)
  • NO (「NO」と入力)
  • NO (「NO」と入力)
  • all (「all」と入力)
  • fp16 (「1」を入力)

Prodigyのインストール

今回の学習に利用するので予めインストールしておく。

ProdigyはSamsungとMetaの技術者が開発したSGD(確率的勾配降下法)のOptimizer(=最適化アルゴリズム)。学習率を自動調整する機能を持つ。同じ研究者が過去に考案した「D-Adaptation」というOptimizerを更に改良したもの。
sd-scriptsでは2023/06/15から使用できるようになった。

sd-scriptsのフォルダーで以下を実行してvenvにインストールする。

.\venv\Scripts\activate
pip install prodigyopt

学習

今回は「DB / キャプション」方式の設定方法で学習する。(これが最も主流な学習方式と言われている)
また、今回は正則化画像は用意しない。したがって、原理としてはDBではなくfine-tuneを行っている(ものと思われる)。

題材

今回はSDのモデル「Realistic Vision」に「埴輪」(はにわ)の概念を覚えさせる例を題材にして、LoRAの学習を解説する。(埴輪にも色々あるが、特に「踊る人々」と呼ばれる有名な埴輪のみを扱う)

Realistic Visionで「haniwa」あるいは「haniwa, japan, clay figure」などとPromptに入力しても、よく知られる埴輪は出て来ない(以下の図のように、こけしのようなものが生成される)。そこで、埴輪の画像を「haniwa」という単語で学習させ、Realistic Visionでも埴輪の画像を出力できるようにしてみる。


haniwaの出力例

画像の用意

学習させたい画像を少なくとも20枚程度用意する。(多ければ多い方が良い)

画像はできれば学習時に設定する解像度(この例では512)よりも大きくした方が良い。必要ならば、AUTOMATIC1111のExtrasタブで拡大することができる。
また、LoRAの学習がうまく行かなかった場合は、背景を白色化するなどの加工を行って試行錯誤する場合もある。

今回は画像をWikimedia Commonsから埴輪の画像を24枚収集した。
1枚の画像に2つの埴輪が写っている場合は、1体ずつに分離して別ファイルとした。


学習に利用する画像
正則化画像を用意する場合

正則化画像を用意する場合は、学習させたい画像のクラスに相当するPromptで予めStable Diffusionから画像を100枚ほど生成する。
この例で正則化画像を使う場合は、例えば「clay figure」などのPromptで正則化画像を生成しておく。(その場合、学習時には教師画像にも必ず「clay figure」を含むようなPromptを対応させる)

学習時には、以下で説明する「config_dataset.toml」に教師画像と同じように[[datasets.subsets]]を追加し、is_reg = trueを設定することで利用できる。

Prompt(キャプション)の用意

学習に使用するPromptには、自然文形式のキャプションと、単語のリスト形式のタグがある。
sd-scriptsの学習にはキャプションかタグを使用する。キャプションとタグを併用する機能もあり、その場合学習に使用されるpromptは「キャプション, タグ1, タグ2, ...」という形になる。

Promptの自動生成

Promptは手動で入力しても良いが、AIで自動生成することもできる。自動生成には以下の方法がある。

  • BLIP - キャプション自動生成器。AUTOMATIC1111ではデフォルトで利用可能。
  • DeepDanbooru (=DeepBooru) - タグ自動生成器。イラストのデータセットに向いている。AUTOMATIC1111ではデフォルトで利用可能。
  • WD14Tagger - タグ自動生成器。イラストのデータセットに向いている。AUTOMATIC1111の拡張機能として別途インストールすることで利用できる。

Trainタブ→Preprocess imagesタブの「Use BLIP for caption」や「Use deepbooru for caption」を利用することで、キャプションを一括で自動生成することができる。


キャプション自動生成の実行画面

Promptの修正

自動生成されたキャプション・タグを手作業で修正する。(手間ならばAIが生成したPromptをそのまま使っても良い)
修正内容は主に以下の3つの内容。

  • トリガーワードを追加
  • トリガーワードに集約させたい概念を削除
  • AIが生成した誤った内容の要素を削除

LoRA学習の定番の方法では、学習したい概念を表す1単語をPromptの先頭に設定する。このような単語をトリガーワードと言う。(DreamBoothの論文に忠実に従って、「class identifier」の2単語を設定しても良い)

また、全ての教師画像に共通する概念はキャプションから除外する。こうすることで、学習によってそれらの共通の要素がトリガーワードに集約されるような効果がもたらされる。その結果、LoRAの利用時にトリガーワードを入れるだけで集約した全ての概念を自動で再現するようになる(と考えられている)。

今回は「haniwa」という単語に埴輪の特徴を集約させる。例えば埴輪には材質が粘土という共通の特徴があるので、学習時のPromptから「clay」という言葉を削除する。
BLIPとDeepBooruで処理した結果、一枚の埴輪の画像に対して次のような文章が生成された。

a clay statue of a person with a hat on his head and a nose on his head, with a small hole in the middle of the body, animal, gradient, gradient background, grey background, no humans, simple background」

赤い要素は「埴輪の特徴」なので、全て削除して先頭に「haniwa」という単語を追加する。
青い要素は誤りなので削除。緑の要素は重複や効果の薄い説明なのでこれらも削除した。
また、撮影した方向が画像のファイル名に書かれていたので、その情報を追加した。

haniwa, with a small hole in the middle of the body, gradient background, grey background, simple background, shot from the front right

このような修正を全ての自動生成された文章に対して手作業で行う。

Promptは画像と同じフォルダーに.txtファイル(あるいは.caption)で保存する。自動生成した場合は、生成されたtxtファイルを直接編集してそのまま学習に利用しても良い。
画像とPromptが入ったフォルダーの名前に「画像を学習させる回数」と「トリガーワード」(あるいは「Class Identifier」)を設定する方法もあるが、今回は「dataset_config」の機能を使うのでフォルダー名は適当で構わない。

学習

まずは、データセットに関する設定ファイル「config_dataset.toml」と、それ以外の設定ファイル「config.toml」を作成する。(ファイル名は自由に変更して構わない)

config.toml

データセット以外に関する設定を記述する。この例では次のように設定した。
オプションの詳しい意味は、venv内で「python train_network.py -h」を実行することで確認できる。
また、train_network.py実行時にオプション--output_configを与えると、起動引数を全てtomlファイルに出力することもできる。

pretrained_model_name_or_path = 'D:\sd\repository\stable-diffusion-webui\models\Stable-diffusion\realisticVisionV50_v50VAE.safetensors'
output_dir = 'D:\experiment\haniwa\save'
network_module = 'networks.lora'
network_dim = 128
#network_alpha = 1
clip_skip = 2
mixed_precision = 'fp16'
xformers = true
optimizer_type = 'prodigy'
learning_rate = 1
optimizer_args = ["d0=1e-5", "d_coef=0.4",]
#optimizer_type = 'AdamW8bit'
#learning_rate = 1e-4
lr_scheduler = 'cosine_with_restarts'
lr_scheduler_num_cycles = 4
lr_warmup_steps = 500
offset_noise = 0.1
max_train_epochs = 10
persistent_data_loader_workers = true
cache_latents = true
seed = 42
#save_model_as = 'safetensors'
save_every_n_epochs = 1

特に重要なもの・調整の余地があるものを解説する。

  • pretrained_model_name_or_path - 学習対象モデルのパスを指定。今回はrealisticVisionV50_v50VAE.safetensorsを学習する。
  • output_dir - 学習したLoRAのパラメーターファイルの保存先パス。
  • network_module - LoRAの場合は'networks.lora'を指定。今回の例では使っていないが、更にnetwork_argsにも値も設定することでLoConを学習することもできる。
  • network_dim - LoRAのランク。実写系は高いランクが良いようなので、今回は128という高い値を設定した。イラスト系は更に低ランクでも良いらしい。
  • clip_skip - 2を使用。
  • optimizer_type - 'prodigy'を指定。
  • learning_rate - Prodigyを使う場合は1を指定。
  • optimizer_args - Prodigyに与えるオプション。
  • offset_noise - 全体的に明るい画像や暗い画像が生成されにくくなる問題を解消するためのパラメーター。
  • max_train_epochs - 学習にデータセットを繰り返し使用する回数。
  • cache_latents - trueを設定するとVAEの結果をキャッシュし、VRAMを節約できる。

config_dataset.toml

データセットに関する設定を記述する。この例では次のように設定した。

[general]
resolution = 512  # [512, 512]
batch_size = 2
enable_bucket = true
#min_bucket_reso = 320
#max_bucket_reso = 1024
bucket_no_upscale = true
caption_dropout_rate = 0.05

[[datasets]]
caption_extension = '.txt'
shuffle_caption = true
keep_tokens = 1
#color_aug = true

  [[datasets.subsets]]
  image_dir = 'D:\experiment\haniwa\data\3_preprocessed'
  num_repeats = 10
#  class_tokens = 'clay figure'
  • enable_bucket - Aspect Ratio Bucketingを有効化するオプション。Aspect Ratio Bucketingは、学習データに異なる縦横比の画像が混在しているときに、同じ縦横比の画像をまとめて1バッチにする仕組み。
  • caption_extension - キャプションファイルの拡張子。デフォルトは.captionなので、ここで.txtと指定する必要がある。
  • shuffle_caption - trueを指定すると、学習時にキャプションをシャッフルする。
  • keep_tokens - Prompt(キャプション)をカンマ区切りで分割した各要素の内、先頭からここで指定した数の要素は学習時にシャッフルされない。
  • num_repeats - 1エポックで画像が学習に使われる回数。画像フォルダー毎に指定することができる。
  • image_dir - 画像とキャプションファイルが入ったフォルダーのパスを入力。フォルダーが複数ある場合は、[[datasets.subsets]]を追加で記載し、それぞれのimage_dirに各フォルダーのパスを指定すれば良い。
  • class_tokens - 画像に対応するキャプションファイルがない場合に使われるPromptを設定できる。DBの元の理論のように、「Identifier Class」形式のPromptを指定したり、正則化画像の場合は「Class」のPromptを指定したりする。すべての画像にキャプションファイルがある場合は設定不要。
  • is_reg - 正則化画像を格納するsubsetsにはtrueを指定する。

[[datasets]]や[[datasets.subsets]]は.tomlファイル内に複数記述することができ、それぞれ異なる設定値で学習させることができる。

学習

以下を.batファイルに保存。
(1行目と4行目の各パス部分は実行環境に合わせて書き換えが必要)

cd /d D:\sd\train\sd-scripts
call venv\Scripts\activate.bat
 
accelerate launch --num_cpu_threads_per_process 4 train_network.py --config_file="D:\experiment\haniwa\config.toml" --dataset_config="D:\experiment\haniwa\config_dataset.toml"
pause

作成したbatファイルを実行すると学習が始まる。
(VRAMを消費するので、AUTOMATIC1111は予め停止しておく)

この例の場合、画像枚数31、num_repeats=10、epoch数10、batch_size=2なので、step数は31×10×10÷2=1550。
学習に要した時間は12分38秒だった。

学習結果は設定ファイルで指定したパスに出力される。

LoRAの利用

学習したLoRAをstable-diffusion-webui/models/Loraに置くことで、AUTOMATIC1111で利用できるようになる。
今回は各チェックポイントをhaniwa-{epoch数}.safetensorsという名前に変更してLoRAフォルダーにコピーした。


LoRAを所定フォルダーへコピー

AUTOMATIC1111を起動。
LoRAタブを開くと、フォルダーにコピーしたLoRAが認識されていることを確認できる。
(1.6.0より前のバージョンでは、Generateボタンの下にある🎴ボタンを押すことでLoRAタブを開くことができるようになる。)


LoRAの選択画面

🛈ボタンを押すと、学習設定を確認することもできる。(2023/07/25機能追加)


LoRAの情報

LoRAをクリックするとPromptに反映され、LoRAを適用することができる。Prompt欄に直接書いても良い。
通常のPromptと同じように、Ctrl+↑or↓でLoRAの重みを調整することもできる。


LoRAの利用

出力例

以下のPromptで画像を出力。Realistic Visionで埴輪を出力することに成功した。

<lora:haniwa-10:1>haniwa

LoRAを使って出力した埴輪の画像

Promptを追加することで、「屋外」・「跳ねる埴輪」など学習元画像にない特徴も追加することができる。
下の出力例のように、各チェックポイントの出力結果を比較して、最も品質の高いチェックポイントを採用すると良い。


各epoch・様々なPromptでの比較

参考

学習に使用した画像の著作権表示

  1. https://commons.wikimedia.org/?curid=19820046, ©sailko, 2012, CC BY-SA 3.0.
  2. https://commons.wikimedia.org/?curid=31437479, Public Domain.
  3. https://commons.wikimedia.org/?curid=31596322, ©Sailko, 2014, CC BY-SA 3.0.
  4. https://commons.wikimedia.org/?curid=39961357, ©x768, 2014, CC BY 2.0.
  5. https://commons.wikimedia.org/?curid=75119158, Public Domain.
  6. https://commons.wikimedia.org/?curid=83164448, ©Jean-Pierre Dalbéra, 2018, CC BY 2.0.
  7. https://commons.wikimedia.org/?curid=118792744, ©Ka23 13, 2019, CC BY-SA 4.0.
  8. https://commons.wikimedia.org/?curid=88870485, ©国立文化財機構, 2019, CC BY 4.0.
  9. https://commons.wikimedia.org/?curid=88870475, ©国立文化財機構, 2019, CC BY 4.0.
  10. https://commons.wikimedia.org/?curid=131844858, Public Domain.
  11. https://commons.wikimedia.org/?curid=131465164, ©国立文化財機構, 2023, CC BY 4.0.
  12. https://commons.wikimedia.org/?curid=131465160, ©国立文化財機構, 2023, CC BY 4.0.
  13. https://commons.wikimedia.org/?curid=131465165, ©国立文化財機構, 2023, CC BY 4.0.
  14. https://commons.wikimedia.org/?curid=131465159, ©国立文化財機構, 2023, CC BY 4.0.
  15. https://commons.wikimedia.org/?curid=131465154, ©国立文化財機構, 2023, CC BY 4.0.
  16. https://commons.wikimedia.org/?curid=131465153, ©国立文化財機構, 2023, CC BY 4.0.
  17. https://commons.wikimedia.org/?curid=131465158, ©国立文化財機構, 2023, CC BY 4.0.
  18. https://commons.wikimedia.org/?curid=131465155, ©国立文化財機構, 2023, CC BY 4.0.
  19. https://commons.wikimedia.org/?curid=131464852, ©国立文化財機構, 2023, CC BY 4.0.
  20. https://commons.wikimedia.org/?curid=131464851, ©国立文化財機構, 2023, CC BY 4.0.
  21. https://commons.wikimedia.org/?curid=131464849, ©国立文化財機構, 2023, CC BY 4.0.
  22. https://commons.wikimedia.org/?curid=131464845, ©国立文化財機構, 2023, CC BY 4.0.
  23. https://commons.wikimedia.org/?curid=131464848, ©国立文化財機構, 2023, CC BY 4.0.
  24. https://commons.wikimedia.org/?curid=131464844, ©国立文化財機構, 2023, CC BY 4.0.