Google Colab ではじめる Textual Inversion

「Diffusers」の「textual_inversion.py」を使った「Textual Inversion」を試したのでまとめました。

・Stable Diffusion v1.4
・Diffusers v0.3.0

1. Textual Inversion

「Textual Inversion」は、3～5枚の画像を使ってファインチューニングを行う手法です。「Stable Diffusion」のモデルに、独自のオブジェクトや画風を覚えさせることができます。

An Image is Worth One Word: Personalizing Text-to-Image Generation using Textual Inversion

今回は、「Diffusers」の「textual_inversion.py」を使います。書くべきコード量も、GPUのメモリ消費量も少なくて良い感じです。

2. ライセンスの確認

以下のモデルカードにアクセスして、ライセンスを確認し、「Access Repository」を押し、「Hugging Face」にログインして（アカウントがない場合は作成）、同意します。

CompVis/stable-diffusion-v1-4 · Hugging Face

3. 学習用画像の準備

ファインチューニングには、3〜5枚の学習用画像が必要です。

・画像サイズ は512×512。
・Exifで回転してない画像。
・最良の結果を得るには約 3～5個の画像。
　使用する画像の枚数が多すぎると、モデルが収束しない場合がある。
・画像には共通のコンテキスト情報が含まれている必要がある。
　・オブジェクト : 同じオブジェクトを異なる視点から写す。
　・画風 : 配色とアートスタイルを一貫させる。

4. ファインチューニングの実行

Colabでファインチューニングを行う手順は、次のとおりです。

(1) 新規のColabノートブックを開き、「GPU」を選択 (ハイメモリでなくてOK)。
GPUのメモリが16GB以上 (以下は16280MiB) であることを確認します。

# GPUの確認 
!nvidia-smi

+-----------------------------------------------------------------------------+
| NVIDIA-SMI 460.32.03    Driver Version: 460.32.03    CUDA Version: 11.2     |
|-------------------------------+----------------------+----------------------+
| 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  Tesla P100-PCIE...  Off  | 00000000:00:04.0 Off |                    0 |
| N/A   33C    P0    27W / 250W |      0MiB / 16280MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+

(2) Googleドライブのマウントと作業フォルダへの移動。
あらかじめGoogleドライブのマイドライブ直下に「work」フォルダを作成しておきます。

# Googleドライブのマウントと作業フォルダへの移動
from google.colab import drive
drive.mount('/content/drive')
%cd '/content/drive/My Drive/work'

(3) パッケージのインストール。

# パッケージのインストール 
!pip install diffusers[training] accelerate transformers

(4) HuggingFaceにログイン。
リンク先でトークンをコピーして、テキストフィールドに入力し、Loginボタンを押します。

# HuggingFaceにログイン
from huggingface_hub import notebook_login
notebook_login()

(5) 「HuggingFace Accelerate」のセットアップ。
single GPUの設定 (0, 0, NO, NO, NO) を行います。

# HuggingFace Accelerateの初期化
!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]: NO

「HuggingFace Accelerate」については以下で簡単に解説してます。

(6) workフォルダ直下のinputsフォルダに学習用画像を配置。

(7) diffusersリポジトリのクローン。
diffusersリポジトリ内のexamplesを利用します。

# diffusersリポジトリのクローン
!git clone https://github.com/huggingface/diffusers.git

(8) ファインチューニングの実行。
今回は、「画風」を学習させてみます。フレースホルダートークンは<sorami>、初期トークンはcomicとします。P100で2時間ほどかかりました。

# 学習の実行
!accelerate launch ./diffusers/examples/textual_inversion/textual_inversion.py \
  --pretrained_model_name_or_path="CompVis/stable-diffusion-v1-4" \
  --use_auth_token \
  --train_data_dir="inputs" \
  --learnable_property="style" \
  --placeholder_token="<sorami>" \
  --initializer_token="comic" \
  --resolution=512 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=4 \
  --max_train_steps=3000 \
  --learning_rate=5.0e-04 \
  --scale_lr \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --output_dir="textual_inversion_sorami"

「--learnable_property」には、画風 (style) と オブジェクト (object) どちらを学習させるかを指定します。
「--placeholder_token」には、 学習した新しい概念を表す単語であるプレースホルダトークンを指定します。他の単語と衝突しないように<>を付加することが推奨されます。
「--initializer_token」には、学習開始時の学習する概念の初期値となる初期化トークンを指定します。

5. 推論の実行

Colabで推論を行う手順は、次のとおりです。

(1) StableDiffusionパイプラインの準備。

import torch
from diffusers import StableDiffusionPipeline

# StableDiffusionパイプラインの準備
pipe = StableDiffusionPipeline.from_pretrained(
    "./textual_inversion_sorami",
    torch_dtype=torch.float16
    ).to("cuda")

(2) 推論の実行。
今回は、猫耳メイド(cute cat ear maid of <sorami> style)で画像生成してみました。

from torch import autocast

# プロンプト
prompt = "cute cat ear maid of <sorami> style"

# 推論の実行
with autocast("cuda"):
    image = pipe(
        prompt, 
        num_inference_steps=50, 
        guidance_scale=7.5
        ).images[0]
image.save("output.png")

(3) 生成した画像の確認。
左端のフォルダアイコンでファイル一覧を表示し、output.pngをダブルクリックします。

6. APIリファレンス

「textual_inversion.py」のパラメータは、次のとおりです。

・-h, --help : ヘルプ
・--pretrained_model_name_or_path [PATH] : 事前学習モデルのパス
　・CompVis/stable-diffusion-v1-4
　・hakurei/waifu-diffusion
・--tokenizer_name [NAME] : トークナイザー名
・--train_data_dir [DIR] : 入力画像フォルダのパス
・--placeholder_token [PLACEHOLDER_TOKEN] : プレースホルダトークン。学習した新しい概念を表す単語。他の単語と衝突しないように<>を付加することが推奨される
・--initializer_token [INITIALIZER_TOKEN] : 初期トークン。学習開始時の学習する概念の初期値
・--learnable_property [LEARNABLE_PROPERTY] : オブジェクト(object) or 画風 (style)
・--repeats [REPEATS] : 学習データの繰り返し回数
・--output_dir [OUTPUT_DIR] : 出力フォルダのパス
・--seed [SEED] : 乱数シード
・--resolution [RESOLUTION] : 入力画像の解像度 (512)
・--center_crop : リサイズ時に画像を中央配置するかどうか
・--train_batch_size [TRAIN_BATCH_SIZE] : バッチサイズ
・--num_train_epochs [NUM_TRAIN_EPOCHS] : 学習エポック数。学習ステップ数指定時はそちらを優先
・--max_train_steps [MAX_TRAIN_STEPS] : 学習ステップ数 (3000〜7000)
・--gradient_accumulation_steps [GRADIENT_ACCUMULATION_STEPS] : backward/updateパスを実行する前に蓄積する更新ステップ数
・--learning_rate [LEARNING_RATE] : 学習率の初期値
・--scale_lr : GPUの数、勾配累積ステップ、バッチサイズによって学習率をスケーリング
・--lr_scheduler [LR_SCHEDULER] : scheduler種別
    ・cosine
    ・cosine_with_restarts
    ・polynomial
    ・constant
    ・constant_with_warmup
・--lr_warmup_steps [LR_WARMUP_STEPS] : lr schedulerのウォームアップステップ数
・--adam_beta1 [ADAM_BETA1] : Adam optimizerのbeta1
・--adam_beta2 [ADAM_BETA2] : Adam optimizerのbeta2
・--adam_weight_decay [ADAM_WEIGHT_DECAY] : Weight decay
・--adam_epsilon [ADAM_EPSILON] : Adam optimizerのepsilon
・--logging_dir [LOGGING] : ログフォルダのパス
・--use_auth_token : 「huggingface-cli login」で指定したトークンの利用
・--push_to_hub : ModelHubにパブリッシュするかどうか
・--hub_token HUB_TOKEN : ModelHubへのパブリッシュで利用するトークン
・--hub_model_id HUB_MODEL_ID : ModelHubへのパブリッシュで利用するID
・--mixed_precision {no,fp16,bf16} : 混合精度を使用するかどうか
・--local_rank [LOCAL_RANK] : 分散学習で利用するlocal_rank

7. 参考

・Textual InversionによるUgly Sonicの復活

I Resurrected "Ugly Sonic" with Stable Diffusion Textual Inversion | Max Woolf's Blog

・DreamBoothとTextual Inversionの比較

Comparison of DreamBooth and Textual Inversion

・22/09/18 Stable Diffusionの追加トレーニングについてのサーベイ

22/09/18 Stable Diffusionの追加トレーニングについてのサーベイ

・22/09/26 Stable Diffusion追加学習の記録

22/09/26 Stable Diffusion追加学習の記録

8. 関連

「Textual Inversion」リポジトリでの学習方法は、以下で紹介しています。

BERT/GPT-3/DALL-E 自然言語処理・画像処理・音声処理 人工知能プログラミング実践入門