SDXL の LoRA トレーニングを kohya-ss/sd-scripts を使用して RunPod 上で行います。
SDXL : RunPod で LoRA トレーニング
作成 : Masashi Okumura (@ClassCat)
作成日時 : 10/04/2024
* サンプルコードの動作確認はしておりますが、必要な場合には適宜、追加改変しています。
* ご自由にリンクを張って頂いてかまいませんが、sales-info@classcat.com までご一報いただけると嬉しいです。
◆ お問合せ : 本件に関するお問合せは下記までお願いします。
- クラスキャット セールス・インフォメーション
- sales-info@classcat.com
- ClassCatJP
SDXL : RunPod で LoRA トレーニング
今回は SDXL の LoRA トレーニングを kohya-ss/sd-scripts を使用して RunPod 上で行います。
kohya-ss/sd-scripts による LoRA トレーニングについては、以前に Google Colab 上での基本的な手順は紹介しています ( SDXL 用 LoRA 訓練 (on Google Colab) ) ので、今回は具体例を示しながら RunPod 上で実行します。
RunPod
使用したポッド仕様は以下です :
- テンプレート: RunPod Pytorch 2.4.0 (runpod/pytorch:2.4.0-py3.11-cuda12.4.1-devel-ubuntu22.04)
- 1x A40 (48 GB VRAM)
- 9 vCPU 50 GB RAM
- ~50 GB Disk
- ~50 GB Pod Volume
- Start Jupyter Notebook
トレーニングデータ
トレーニングデータは、東北ずんこ・ずんだもんプロジェクト の AI画像モデル用学習データ を使用させていただきました。
記載されている使用規約は :
- 公式が配布する 東北3姉妹のイラストは非商用、もしくは東北企業であれば自由に使える そうです。
- また、AI画像モデル用学習データについては :
ずんずんPJ公式イラストを学習したデータのLoRAの配布はなどは大丈夫です（２次創作のデータが入る場合は２次創作者への配慮［２次創作者のガイドライン確認など］が別途必要になります）
実践
上記の仕様でポッドを起動し、Jupyter ノートブックでアクセスします。
ポッドの確認
最初にポッドの仕様を簡単に確認しましょう :
!nvidia-smi
!echo ""
!nvcc --version
!echo ""
!free -h
Fri Oct 4 03:12:16 2024
+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 550.90.07 Driver Version: 550.90.07 CUDA Version: 12.4 |
|-----------------------------------------+------------------------+----------------------+
| GPU Name Persistence-M | Bus-Id Disp.A | Volatile Uncorr. ECC |
| Fan Temp Perf Pwr:Usage/Cap | Memory-Usage | GPU-Util Compute M. |
| | | MIG M. |
|=========================================+========================+======================|
| 0 NVIDIA A40 On | 00000000:98:00.0 Off | 0 |
| 0% 30C P8 22W / 300W | 1MiB / 46068MiB | 0% Default |
| | | N/A |
+-----------------------------------------+------------------------+----------------------+
+-----------------------------------------------------------------------------------------+
| Processes: |
| GPU GI CI PID Type Process name GPU Memory |
| ID ID Usage |
|=========================================================================================|
| No running processes found |
+-----------------------------------------------------------------------------------------+
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2024 NVIDIA Corporation
Built on Thu_Mar_28_02:18:24_PDT_2024
Cuda compilation tools, release 12.4, V12.4.131
Build cuda_12.4.r12.4/compiler.34097967_0
total used free shared buff/cache available
Mem: 503Gi 115Gi 44Gi 7.9Gi 343Gi 376Gi
Swap: 0B 0B 0B
環境構築 (kohya-ss/sd-scripts)
レポジトリ kohya-ss/sd-scripts を複製して依存関係をインストールし、xformers と bitsandbytes をインストールするだけです。
※ 必要なら、SDXL 用 LoRA 訓練 (on Google Colab) を参照してください。
トレーニングデータの準備
次にトレーニングデータを準備します。前述の「東北ずんこ・ずんだもんプロジェクト」の AI 画像モデル用学習データから 東北イタコ のデータを使用しましたが、お好みでどうぞ。
zip ファイルをアップロードして解凍してから :
!unzip -q /workspace/itako.zip
以下のようなディレクトリ構成にします :
/content/LoRA
├── config
│ └── dataset_config.toml
└── train_data
├── itako(1).png
├── itako(1).txt
├── itako(9).png
└── itako(9).txt
toml 設定ファイルは少なくとも以下の内容があれば動作します :
[general]
shuffle_caption = true
enable_bucket = true
caption_extension = ".txt"
[[datasets]]
resolution = 1024
min_bucket_reso = 768
max_bucket_reso = 1280
[[datasets.subsets]]
image_dir = "/workspace/LoRA/train_data"
class_tokens = "itako "
num_repeats = 1
Hugging Face Accelerate
Accelerate 用の設定ファイルも /workspace/LoRA/config に配置することにします :
import os
accelerate_config = "/workspace/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)
その結果、以下のようなディレクトリ構成になります :
/workspace/LoRA
├── config
│ ├── accelerate_config.yaml
│ └── dataset_config.toml
└── train_data
├── itako(1).png
├── itako(1).txt
├── itako(9).png
└── itako(9).txt
SDXL モデルファイルの配置
モデルファイルは stabilityai/stable-diffusion-xl-base-1.0 から取得できます。
/content/pretrained_model に配置しました :
!mkdir -p /content/pretrained_model
!wget https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/resolve/main/sd_xl_base_1.0.safetensors \
-O /content/pretrained_model/sd_xl_base_1.0.safetensors
トレーニングの実行
必要最小限の引数とともに訓練を実行します。出力先は /workspace/LoRA/output です。
%cd /workspace/sd-scripts
!accelerate launch \
--config_file=/workspace/LoRA/config/accelerate_config.yaml \
sdxl_train_network.py \
--dataset_config=/workspace/LoRA/config/dataset_config.toml \
--pretrained_model_name_or_path=/workspace/pretrained_model/sd_xl_base_1.0.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="itako" \
--save_precision="fp16" \
--save_every_n_epochs=1 \
--save_model_as=safetensors \
--max_train_epochs=5 \
--train_batch_size=1 \
--xformers \
--no_half_vae \
--lowram
最終的なディレクトリ構成は以下のようなものです :
/workspace/LoRA
├── config
│ ├── accelerate_config.yaml
│ └── dataset_config.toml
├── output
│ ├── itako-000001.safetensors
│ └── itako.safetensors
└── train_data
├── itako(1).png
├── itako(1).txt
├── itako(9).png
└── itako(9).txt
/workspace/pretrained_model
└── sd_xl_base_1.0.safetensors
以下は、モデルファイル Lykon/AAM_XL_AnimeMix を使用した際のトレーニング過程です :
1 epoch
5 epoch
10 epoch
推論
トレーニングが完了したら、Hugging Face Diffusers を使用してサンプリングしてみます :
import torch
from diffusers import StableDiffusionXLPipeline
pipeline = StableDiffusionXLPipeline.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0",
#"Lykon/AAM_XL_AnimeMix",
torch_dtype=torch.float16,
)
pipeline.enable_model_cpu_offload()
pipeline.load_lora_weights("/workspace/LoRA/output/itako.safetensors")
prompt = "itako, 1girl, solo, kimono, hagoromo, tabi, animal ear fluff, hair ornament, dancing"
out = pipeline(
prompt=prompt,
guidance_scale=7.0,
width=1024,
height=1024,
num_inference_steps=20,
).images[0]
out.save("/workspace/image.png")
以上