【Stable Diffusion】LoRAによる追加学習

【Stable Diffusion】LoRAによる追加学習 機械学習

「LoRAによる追加学習を試したい」
「WindowsでLoRAによる学習が実行できない・・・」

このような場合には、この記事の内容が参考になります。
この記事では、LoRAによる追加学習について解説しています。

本記事の内容

  • LoRAによる追加学習
  • LoRAによる追加学習の効果
  • WindowsでLoRAによる追加学習ができない!?

それでは、上記に沿って解説していきます。

LoRAによる追加学習

LoRAについては、次の記事で解説しています。
インストールして動作確認までは、実施しておいてください。

LoRAによる追加学習は、基本的にはDreamBoothと同じです。
そのため、不明点がある場合は次の記事を参考にしてください。

では、LoRAによる追加学習を行っていきます。
まずは、LoRAのリポジトリルートまで移動します。

そこでは、次のファイルを確認できます。

「train_lora_dreambooth.py」を用いて、追加学習を実行できます。
まずは、ヘルプを確認しましょう。

> accelerate launch train_lora_dreambooth.py -h
usage: train_lora_dreambooth.py [-h] --pretrained_model_name_or_path PRETRAINED_MODEL_NAME_OR_PATH [--pretrained_vae_name_or_path PRETRAINED_VAE_NAME_OR_PATH]
                                [--revision REVISION] [--tokenizer_name TOKENIZER_NAME] --instance_data_dir INSTANCE_DATA_DIR [--class_data_dir CLASS_DATA_DIR]
                                --instance_prompt INSTANCE_PROMPT [--class_prompt CLASS_PROMPT] [--with_prior_preservation] [--prior_loss_weight PRIOR_LOSS_WEIGHT]
                                [--num_class_images NUM_CLASS_IMAGES] [--output_dir OUTPUT_DIR] [--seed SEED] [--resolution RESOLUTION] [--center_crop]
                                [--color_jitter] [--train_text_encoder] [--train_batch_size TRAIN_BATCH_SIZE] [--sample_batch_size SAMPLE_BATCH_SIZE]
                                [--num_train_epochs NUM_TRAIN_EPOCHS] [--max_train_steps MAX_TRAIN_STEPS] [--save_steps SAVE_STEPS]
                                [--gradient_accumulation_steps GRADIENT_ACCUMULATION_STEPS] [--gradient_checkpointing] [--lora_rank LORA_RANK]
                                [--learning_rate LEARNING_RATE] [--learning_rate_text LEARNING_RATE_TEXT] [--scale_lr] [--lr_scheduler LR_SCHEDULER]
                                [--lr_warmup_steps LR_WARMUP_STEPS] [--use_8bit_adam] [--adam_beta1 ADAM_BETA1] [--adam_beta2 ADAM_BETA2]
                                [--adam_weight_decay ADAM_WEIGHT_DECAY] [--adam_epsilon ADAM_EPSILON] [--max_grad_norm MAX_GRAD_NORM] [--push_to_hub]
                                [--hub_token HUB_TOKEN] [--logging_dir LOGGING_DIR] [--mixed_precision {no,fp16,bf16}] [--local_rank LOCAL_RANK]
                                [--resume_unet RESUME_UNET] [--resume_text_encoder RESUME_TEXT_ENCODER]

Simple example of a training script.

options:
  -h, --help            show this help message and exit
  --pretrained_model_name_or_path PRETRAINED_MODEL_NAME_OR_PATH
                        Path to pretrained model or model identifier from huggingface.co/models.
  --pretrained_vae_name_or_path PRETRAINED_VAE_NAME_OR_PATH
                        Path to pretrained vae or vae identifier from huggingface.co/models.
  --revision REVISION   Revision of pretrained model identifier from huggingface.co/models.
  --tokenizer_name TOKENIZER_NAME
                        Pretrained tokenizer name or path if not the same as model_name
  --instance_data_dir INSTANCE_DATA_DIR
                        A folder containing the training data of instance images.
  --class_data_dir CLASS_DATA_DIR
                        A folder containing the training data of class images.
  --instance_prompt INSTANCE_PROMPT
                        The prompt with identifier specifying the instance
  --class_prompt CLASS_PROMPT
                        The prompt to specify images in the same class as provided instance images.
  --with_prior_preservation
                        Flag to add prior preservation loss.
  --prior_loss_weight PRIOR_LOSS_WEIGHT
                        The weight of prior preservation loss.
  --num_class_images NUM_CLASS_IMAGES
                        Minimal class images for prior preservation loss. If not have enough images, additional images will be sampled with class_prompt.
  --output_dir OUTPUT_DIR
                        The output directory where the model predictions and checkpoints will be written.
  --seed SEED           A seed for reproducible training.
  --resolution RESOLUTION
                        The resolution for input images, all the images in the train/validation dataset will be resized to this resolution
  --center_crop         Whether to center crop images before resizing to resolution
  --color_jitter        Whether to apply color jitter to images
  --train_text_encoder  Whether to train the text encoder
  --train_batch_size TRAIN_BATCH_SIZE
                        Batch size (per device) for the training dataloader.
  --sample_batch_size SAMPLE_BATCH_SIZE
                        Batch size (per device) for sampling images.
  --num_train_epochs NUM_TRAIN_EPOCHS
  --max_train_steps MAX_TRAIN_STEPS
                        Total number of training steps to perform. If provided, overrides num_train_epochs.
  --save_steps SAVE_STEPS
                        Save checkpoint every X updates steps.
  --gradient_accumulation_steps GRADIENT_ACCUMULATION_STEPS
                        Number of updates steps to accumulate before performing a backward/update pass.
  --gradient_checkpointing
                        Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.
  --lora_rank LORA_RANK
                        Rank of LoRA approximation.
  --learning_rate LEARNING_RATE
                        Initial learning rate (after the potential warmup period) to use.
  --learning_rate_text LEARNING_RATE_TEXT
                        Initial learning rate for text encoder (after the potential warmup period) to use.
  --scale_lr            Scale the learning rate by the number of GPUs, gradient accumulation steps, and batch size.
  --lr_scheduler LR_SCHEDULER
                        The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial", "constant", "constant_with_warmup"]     
  --lr_warmup_steps LR_WARMUP_STEPS
                        Number of steps for the warmup in the lr scheduler.
  --use_8bit_adam       Whether or not to use 8-bit Adam from bitsandbytes.
  --adam_beta1 ADAM_BETA1
                        The beta1 parameter for the Adam optimizer.
  --adam_beta2 ADAM_BETA2
                        The beta2 parameter for the Adam optimizer.
  --adam_weight_decay ADAM_WEIGHT_DECAY
                        Weight decay to use.
  --adam_epsilon ADAM_EPSILON
                        Epsilon value for the Adam optimizer
  --max_grad_norm MAX_GRAD_NORM
                        Max gradient norm.
  --push_to_hub         Whether or not to push the model to the Hub.
  --hub_token HUB_TOKEN
                        The token to use to push to the Model Hub.
  --logging_dir LOGGING_DIR
                        [TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***.        
  --mixed_precision {no,fp16,bf16}
                        Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10.and an Nvidia Ampere GPU. Default to  
                        the value of accelerate config of the current system or the flag passed with the `accelerate.launch` command. Use this argument to override  
                        the accelerate config.
  --local_rank LOCAL_RANK
                        For distributed training: local_rank
  --resume_unet RESUME_UNET
                        File path for unet lora to resume training.
  --resume_text_encoder RESUME_TEXT_ENCODER
                        File path for text encoder lora to resume training.

ほぼDreamBoothと同じですね。
おそらく、DreamBoothのプログラムをベースにしているのでしょう。

そのため、実行コマンドもほぼ変わりません。
そして、利用例となるコマンドは以下のスクリプトに記載されています。

これらを参考にして、次のコマンドを利用します。

accelerate launch train_lora_dreambooth.py `
  --pretrained_model_name_or_path="stabilityai/stable-diffusion-2-1-base"  `
  --instance_data_dir="./training_images" `
  --output_dir="./model" `
  --instance_prompt="lora_irasutoya" `
  --train_text_encoder `
  --resolution=512 --color_jitter `
  --train_batch_size=1 `
  --gradient_accumulation_steps=1 `
  --learning_rate=1e-4 `
  --learning_rate_text=5e-5 `
  --lr_scheduler="constant" `
  --lr_warmup_steps=0 `
  --max_train_steps=30000

正則化画像については、今回は用意しません。
とりあえず、学習画像だけでどれくらいの効果があるのかを確認します。

学習画像は、DreamBoothで用いたモノを流用しています。

DreamBoothとの大きな違いは、以下のオプションです。

--train_text_encoder

これを追加すれば、CLIP用のLoRAモデルも生成されることになります。
追加しないと、Unet用のLoRAモデルだけが生成されます。

あと、学習ステップに注意が必要です。

  --max_train_steps=30000

この「30000」は、参考スクリプトの数字をそのまま利用しています。
実際、「30000」を試したところ、2時間20分ほどかかりました。

この時間は、DreamBoothの「8000」と同じぐらいの時間になります。
もちろん、この時間は各自のマシンスペックにより異なることに注意です。

単純に、テスト稼働だけしたい場合は「300」ぐらいのステップ数でも問題ありません。

実際に追加学習した結果は、指定したディレクトリ(今回は「./model」)に保存されます。

下線のあるモノが、最終成果物(約3MB)です。
500ステップ毎に、それぞれの成果物が保存されています。

以上、LoRAによる追加学習を説明しました。
次は、LoRAによる追加学習の効果を説明します。

LoRAによる追加学習の効果

追加学習の成果を利用して、その効果を確認していきます。
今回であれば、次の2ファイルですね。

  • lora_weight.pt
  • lora_weight.text_encoder.pt

これらの利用方法は、すでに上げた記事内で説明しています。

上記記事内の方法をベースに、次のコードを用意しました。
これで、LoRAによる追加学習の効果を確認していきます。

import torch
from lora_diffusion import monkeypatch_lora, tune_lora_scale
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler

model_id = "stabilityai/stable-diffusion-2-1-base"
unet_pt = "./model/lora_weight.pt"
text_encoder_pt = "./model/lora_weight.text_encoder.pt"

pipe = StableDiffusionPipeline.from_pretrained(model_id)
pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
pipe = pipe.to("cuda")

monkeypatch_lora(pipe.unet, torch.load(unet_pt))
monkeypatch_lora(pipe.text_encoder, torch.load(text_encoder_pt), target_replace_module=["CLIPAttention"])

tune_lora_scale(pipe.unet, 1.0)
tune_lora_scale(pipe.text_encoder, 1.0)

prompt = "a boy, lora_irasutoya"

image = pipe(
    prompt,
    num_inference_steps=25,
    guidance_scale=7
).images[0]

image.save("test.png")

次の箇所で指定している数字は、結果を見ながら適当に調整してください。

tune_lora_scale(pipe.unet, 1.0) 
tune_lora_scale(pipe.text_encoder, 1.0)

その際に参考となるのは、次のページです。
https://github.com/cloneofsimo/lora/discussions/37

あと、プロンプトには各自で設定した「–instance_prompt」を含める必要があります。

prompt = "a boy, lora_irasutoya"

上記コードを実行した結果、次のような画像が生成されました。

test.png

いらすとや風の画像です。
では、動きを付けてみましょう。

a boy swimming in the river, lora_irasutoya

ここから確認可能。

泳いでいるようには見えませんが、川の中にいるのでOKとしましょう。
次は、いらすとや風のスタイルでキャラを描いてみましょう。

Iron Man, lora_irasutoya

想像以上にアイアンマンですね。
最後は、人間ではないモノを描画してみましょう。

Totoro in the forest, lora_irasutoya

うーん、微妙ですね。
そもそも、学習画像は人の画像だけでした。

それもたった33枚です。
それを考えれば、ここまで表現できるのはスゴイのかもしれません。
何と言っても、合計でたった6MBのファイルを読み込んだだけですからね。

以上、LoRAによる追加学習の効果を説明しました。
次は、WindowsでLoRAによる追加学習ができない!?を説明します。

WindowsでLoRAによる追加学習ができない!?

Windows環境では、LoRAの追加学習ができない場合があります。
実際、私はできませんでした。

その事象については、次のページで記載されています。
解決方法も提示してあります。

Doesn’t seem to work on Windows
https://github.com/cloneofsimo/lora/issues/4

「num_workers」の記述が、Windowsではアウトのようです。
ダウンロードしたリポジトリ以下では、次の2ファイルで記載されています。

train_lora_dreambooth.pyにありますね。
これを無効にします。

train_lora_dreambooth.py

コメントにするだけです。

これで、WindowsでもLoRAによる追加学習が可能になります。
以上、WindowsでLoRAによる追加学習ができない!? を説明しました。

タイトルとURLをコピーしました