ホーム / stable-diffusion /

LoRAで追加学習を試す(google colaboratory, stable diffusion, fine tuning)

2023年3月18日土曜日

stable-diffusion

t f B! P L

目的

  • Stable Diffusionの画像生成で使えるLoRAを作成する
  • LoRAを使った追加学習の流れを知る(sd-scriptsを用いたfine tuning方式)
  • タグの整理、正則化などはやらずにとにかく流れを把握することだけを目標にする
  • Google Colaboratoryで動かす

参考

https://github.com/kohya-ss/sd-scripts/blob/main/README-ja.mdを参考にしました。最近書いてくれた方ではなく、設定ファイルを用いずコマンドラインから指定する方法でやります。

必要なもの、環境

  • 学習対象の画像
  • google colabolatory
  • Googleドライブ
  • モデル(AbyssOrangeMixなど)

モデルと学習対象の画像はGoogleドライブに入れておくとやりやすいです。学習結果の保存にGoogleドライブを使うので1Gbほど空き容量が必要になります。

今回使う画像

例としてユニティちゃんを学習させることにしました。公式からダウンロードできるイラスト5枚(© UTJ/UCL)でやってみます。これらをzipにしてGoogleドライブに保存しておきます。

python3.10のインストール

colabにはpython3.9が入っているが、sd-scriptsで推奨されているpythonのバージョンが3.10なので入れておく。

!sudo apt autoremove
!sudo apt-get -y update
!sudo apt-get -y install python3.10

!sudo cp `which python3.10` /usr/local/bin/python

!wget https://bootstrap.pypa.io/get-pip.py
!sudo python get-pip.py

!python --version
!pip --version

sd-scriptsと必要なライブラリをインストール

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

!pip install torch==1.13.1+cu117 torchvision==0.14.1+cu117 --extra-index-url https://download.pytorch.org/whl/cu117
!pip install --upgrade -r requirements.txt
!pip install -U --pre triton
!pip install xformers==0.0.16rc425

colabの最初にいるディレクトリが/contentなので/content/sd-scriptsというディレクトリに展開されること前提で進めます。

accelerateの設定

!accelerate config
(以下accelerate configの質問)
In which compute environment are you running? ([0] This machine, [1] AWS (Amazon SageMaker)): 0
Which type of machine are you using? ([0] No distributed training, [1] multi-CPU, [2] multi-GPU, [3] TPU [4] MPS): 0
Do you want to run your training on CPU only (even if a GPU is available)? [yes/NO]:NO
Do you want to use DeepSpeed? [yes/NO]: NO
Do you wish to use FP16 or BF16 (mixed precision)? [NO/fp16/bf16]: fp16

accelerate configの質問に答えるとき、ブラウザ上のUIだと上手く入力できなくて動かないことがあるのでそのときはdefaultを指定します。

!accelerate config default --mixed_precision fp16

学習対象の画像をcolabに送る

この後ディレクトリの指定が必要になるので置いた場所は覚えておきます。今回はsd-scriptsの中に展開。/content/sd-scripts/unitychanの中に画像を置いています。ディレクトリは保存場所に応じて読み替えてください。

%cd /content/sd-scripts
from google.colab import drive
drive.mount('/content/drive')
!cp /content/drive/MyDrive/your/path/to/unitychan.zip ./
!unzip unitychan.zip

キャプションとタグを生成する

キャプション

!python finetune/make_captions.py --batch_size 8 ./unitychan

タグ

!python finetune/tag_images_by_wd14_tagger.py --batch_size 4 ./unitychan

どのようなキャプション、タグが生成されたか覗いてみます。

そこそこの精度。

キャプションとタグ情報の前処理

参考サイトの通りにキャプションとタグをjsonにまとめ、クリーニング(girlとwomanをgirlに統一するなど)します。

!python finetune/merge_captions_to_metadata.py ./unitychan meta_cap.json
!python finetune/merge_dd_tags_to_metadata.py ./unitychan --in_json meta_cap.json meta_cap_dd.json
!python finetune/clean_captions_and_tags.py meta_cap_dd.json meta_clean.json

latentsの事前取得

この例ではAOM3A1をモデルに利用します。

!python finetune/prepare_buckets_latents.py ./unitychan meta_clean.json meta_lat.json /content/drive/MyDrive/path/to/AOM3A1.safetensors --batch_size 4 --max_resolution 512,512 --mixed_precision no

学習実行

!accelerate launch  --num_cpu_threads_per_process 1 train_network.py --pretrained_model_name_or_path=/content/drive/MyDrive/path/to/AOM3A1.safetensors --in_json meta_lat.json --train_data_dir=unitychan --output_dir=unitychanft_s400 --shuffle_caption --train_batch_size=1 --learning_rate=1e-4 --max_train_steps=400 --use_8bit_adam --xformers --gradient_checkpointing --mixed_precision=fp16 --save_every_n_epochs=4 --save_precision=fp16 --save_model_as=safetensors  --network_module=networks.lora

一部引数のメモ

  • --pretrained_model_name_or_path=/content/drive/MyDrive/path/to/AOM3A1.safetensors: 学習済みモデルへのパス
  • --train_data_dir=unitychan: 学習対象の画像があるディレクトリ
  • --output_dir=unitychanft_s400: 学習結果が出力されるディレクトリ。好きに決める
  • --max_train_steps=400: 学習回数。上手く動くか400ぐらいで試してみるといいかも
  • --network_module=networks.lora: 忘れると1.99Gbのファイルが生成され続けcolabだとストレージが足りなくなる

colabでこの設定だと学習は数分で終わるはず。

保存

途中経過含めて全部zipにしてGoogleドライブに送って保存。last.safetensorsが最終出力なのでこれを使用します。

!cp unitychanft_s400/last.safetensors /content/drive/MyDrive/path/to/loradir/unitychan_s400.safetensors
!zip unitychanft_s400.zip unitychanft_s400/*
!cp unitychanft_s400.zip /content/drive/MyDrive/path/to/yourdir/unitychanft_s400.zip

学習できているか試す

タグ生成したワードで学習されているのでそのまま使ってみます。タグは.txtファイルに保存されているので見てみましょう。

!cat /content/sd-scripts/unitychan/*.txt
(出力)
1girl, solo, long_hair, blonde_hair, thighhighs, twintails, very_long_hair, green_eyes, shoes, striped, striped_thighhighs, low-tied_long_hair
1girl, solo, long_hair, looking_at_viewer, blush, smile, blonde_hair, very_long_hair, green_eyes, one_eye_closed, headband
1girl, solo, long_hair, smile, open_mouth, blonde_hair, thighhighs, green_eyes, striped, striped_thighhighs
1girl, solo, long_hair, looking_at_viewer, smile, blonde_hair, thighhighs, very_long_hair, green_eyes, detached_sleeves, striped, hood, hand_on_hip, hoodie, striped_thighhighs, low-tied_long_hair, salute, sleeveless_hoodie
1girl, solo, long_hair, open_mouth, skirt, blonde_hair, thighhighs, very_long_hair, green_eyes, weapon, striped, striped_thighhighs, flag, orange_thighhighs

この中の一つをそのまま使ってみます。

1girl, solo, long_hair, looking_at_viewer, smile, blonde_hair, thighhighs, very_long_hair, green_eyes, detached_sleeves, striped, hood, hand_on_hip, hoodie, striped_thighhighs, low-tied_long_hair, salute, sleeveless_hoodie

で生成してみる。

今回の学習の場合、clipskipが1のほうがユニティちゃんに近い画像が生成されます。

違うけど近い。weightを2にあげてみる。

近くなった。同一seedでLoRAなしも見てみます。

ぜんぜん違うのでLoRAが効いていることがわかります。

ちなみに学習するとき--max_train_steps=2000にするとこうなります。その場合、学習するのにcolabだと1時間ほど時間がかかります。

学習させるイラストの数が少ないので元画像と近い構図になりがちですが、より近い絵柄になりました。LoRAのweightが1だと強すぎたので0.7に下げています。

おわり

便利なスクリプトのおかげで簡単に追加学習ができました。学習させる画像を増やしたり、step数など各種パラメータを変えたり、タグを微調整したりするとより精度の高い画像生成ができるようになります。

QooQ