Stable Diffusion 2.1 用 LoRA 訓練 (on Google Colab) (ブログ)

作成 : Masashi Okumura (@ClassCat)
作成日時 : 10/26/2023

* サンプルコードの動作確認はしておりますが、動作環境によりコードの追加変更が必要な場合はあるかもしれません。
* ご自由にリンクを張って頂いてかまいませんが、sales-info@classcat.com までご一報いただけると嬉しいです。

 

 

Stable Diffusion 2.1 用 LoRA 訓練 (on Google Colab)

Stable Diffusion のポピュラーな訓練スクリプト・ツール kohya-ss/sd-scripts による Stable Diffusion 2.1 用 LoRA の訓練についても具体的な手順を簡単に記述しておきます。

 

1. 概要

SD 1.5 で kohya-ss/sd-scripts による LoRA 訓練の経験があれば、基本的な流れに違いはありません。スクリプト “train_network.py” を使用すれば良く、引数フラグを 2 つ追加するだけです。

注意点としては、解像度 768×768 の訓練画像を中心に用意することです。

 

2. kohya-ss/sd-scripts

最初に kohya-ss/sd-scripts を “git clone” して依存関係をインストールしましょう :

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

%cd /content/sd-scripts
!pip install --upgrade -q -r requirements.txt

!pip install -q xformers==0.0.22 triton

!pip install -q bitsandbytes

 

3. データセット

データセットは通常通りに (画像ファイル、キャプションファイル) のペアを用意します。またデータセットについて記述した toml 設定ファイルも用意します。

ここでは以下のようなディレクトリ構成を想定します :

/content/LoRA
├── config
│   └── dataset_config.toml
└── train_data
    ├── 0001.png
    ├── 0001.txt

    ├── 0xxx.png
    └── 0xxx.txt

toml 設定ファイルの内容は以下のようなものです :

[general]
shuffle_caption = true
enable_bucket = true
caption_extension = ".txt"

[[datasets]]
resolution = 768
min_bucket_reso = 512
max_bucket_reso = 1024

[[datasets.subsets]]
image_dir = "/content/LoRA/train_data"
class_tokens = "mysample"
num_repeats = 1

 

Hugging Face Accelerate

Accelerate 用の設定ファイルも /content/LoRA/config に配置することにします :

import os

accelerate_config = "/content/LoRA/config/accelerate_config.yaml"

from accelerate.utils import write_basic_config

if not os.path.exists(accelerate_config):
    write_basic_config(save_location=accelerate_config)

/content/LoRA
├── config
│   ├── accelerate_config.yaml
│   └── dataset_config.toml
└── train_data
    ├── 0001.png
    ├── 0001.txt
    ├── 0002.png
    ├── 0002.txt
...

 

Stable Diffusion 2.1 の取得配置

Stable Diffusion 2.1 は /content/pretrained_model に配置しました :

!mkdir -p /content/pretrained_model
!wget https://huggingface.co/stabilityai/stable-diffusion-2-1/resolve/main/v2-1_768-ema-pruned.safetensors \
   -O /content/pretrained_model/v2-1_768-ema-pruned.safetensors

 

訓練の実行

Stable Diffusion 1.x 用 LoRA 訓練との違いは次の 2 つの引数フラグです : “–v2” 及び “–v_parameterization”。

必要最小限の引数とともに訓練を実行します。出力先は /content/LoRA/output です。
バッチサイズは 1 にしてありますが、GPU メモリに応じて増やすことができます。

%%time
%cd /content/sd-scripts

!accelerate launch \
    --config_file=/content/LoRA/config/accelerate_config.yaml \
    train_network.py \
    --v2 \
    --v_parameterization \
    --dataset_config=/content/LoRA/config/dataset_config.toml \
    --pretrained_model_name_or_path=/content/pretrained_model/v2-1_768-ema-pruned.safetensors \
    --network_module=networks.lora \
    --mixed_precision="fp16" \
    --network_dim=8 --network_alpha=4 \
    --learning_rate=1e-4 --optimizer_type="AdamW8bit" \
    --output_dir=/content/LoRA/output --output_name="mysample" \
    --save_precision="fp16" --save_every_n_epochs=1 --save_model_as=safetensors \
    --max_train_epochs=5 --train_batch_size=1 \
    --xformers --no_half_vae

最終的なディレクトリ構成は以下のようなものです :

/content/LoRA
├── config
│   ├── accelerate_config.yaml
│   └── dataset_config.toml
├── output
│   ├── mysample-000001.safetensors
│   ├── mysample-000002.safetensors
│   ├── mysample-000003.safetensors
│   ├── mysample-000004.safetensors
│   └── mysample.safetensors
└── train_data
    ├── 0001.png
    ├── 0001.txt
    ├── 0002.png
    ├── 0002.txt

    ├── 0100.png
    └── 0100.txt
/content/pretrained_model
└── v2-1_768-ema-pruned.safetensors
/content/sd-scripts

以上