Google Colab ではじめる Textual Inversion

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

1. Textual Inversion

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

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

2. ライセンスの確認

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

3. 学習用画像の準備

ファインチューニングには、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」のパラメータは、次のとおりです。

7. 参考

・Textual InversionによるUgly Sonicの復活

・DreamBoothとTextual Inversionの比較

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

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

8. 関連

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