Stable Diffusionを手塚治虫のキャラクターでファインチューニングしてみる（Textual Inversion）

この記事では、Textual Inversionによって、数枚のオリジナル画像でStable Diffusionをファインチューニングし、新たに好きなキャラクターや画風を学習させる方法を紹介します。

Stable Diffusionを用いたtext-to-image、image-to-image、image inpaintingの方法は以下の記事をご覧ください。

Textual Inversionとは
デモ
まとめ
参考文献

Textual Inversionとは

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

Textual Inversionは、好きなコンセプトを表す3～5枚の画像を用いて、新たにそのコンセプトをモデルに学習させる手法です。

それらの画像によってStable Diffusionのようなtext-to-imageモデルの単語埋め込み層をファインチューニングし、画像の画風やモチーフを新しいワードに圧縮します。そして、そのワードを使って画像生成のための指示を与えれば、圧縮された情報を再現することができるようになります。

以下の例では、左端の画像の画風が「S*」という単語に圧縮されており、「S*風の○○」といった指示を与えるだけで似た画風の画像が生成されています。

デモ

それでは、Google Colaboratoryを使って実際に実行していきます。

なお、記事内で紹介したコードをすべて含むノートブックは、以下のリンクから直接参照することができます。

事前準備

まず、モデルが公開されているHugging Faceのアカウントを作成します。

Hugging Faceのトップページから右上の「Sign Up」に進みます。アドレスとパスワードを入力すると、登録したアドレスに認証メールが届くので、メール内のリンクをクリックして作成は完了です。

次に、Stable Diffusionのリポジトリを開き「Access repository」をクリックします。

最後に、学習済みモデルをダウンロードするためのアクセストークンを取得します。

トップページから先ほど作成したアカウントでログイン後、右上のアカウントのアイコンをクリックし、「Settings」、「Access Tokens」と進みます。画面中央の「New Token」をクリックしてアクセストークンを発行し、コピーしておきます。

これで事前準備は完了です。

環境設定

ここからは、Google Colaboratory上で作業を行います。

はじめに、画面上部のメニューから、「ランタイム」、「ランタイムのタイプを変更」と進み、「ハードウェアアクセラレータ」を「GPU」に変更しておきます。これにより、推論時にGPUを利用することができます。

続いて、必要となるライブラリのインストール、インポートを行います。

# ライブラリのインストール
!pip install -qq diffusers['training']==0.3.0 transformers accelerate ftfy

# ライブラリのインポート
import os
import re
import glob
from matplotlib import pyplot as plt
from PIL import Image
import torch
from diffusers import StableDiffusionPipeline
from torch import autocast

先ほど取得したアクセストークンを用いて、Hugging Faceにログインします。以下の内容を記載したセルを実行後、表示される指示にしたがってアクセストークンを入力してください。

from huggingface_hub import notebook_login
notebook_login()

次に、HuggingFace Accelerateを初期化します。セル実行後、表示される質問に対して、colabの環境に対応するように「0, 0, NO, NO, NO」と順に入力します。

!accelerate config

最後に、diffusersリポジトリをクローンして環境設定は完了です。

!git clone https://github.com/huggingface/diffusers.git
%cd diffusers

ファインチューニング

新しいコンセプトの画像をcolab上にアップロードします。

アップロードの方法は下記の記事をご覧ください。

今回は、こちらからお借りした、手塚治虫先生のキャラクター画像5枚を利用します（※ 利用規約を遵守しています。再配布時にはご留意ください）。

手塚治虫 TEZUKA OSAMU OFFICIAL

続いて、学習用の設定を行います。

まず、画像に含まれる「対象」を学習させるのか、「画風」を学習させるのか選択します。今回は「画風」を学習させたいので「style」を選択します。

what_to_teach = 'style'  # 'object' or 'style'

次に、学習させるコンセプトを表すカテゴリや、要約した単語を指定します。この値は学習の初期値として利用されます。ここでは、漫画絵というカテゴリを意識して「comic」と指定します。

initializer_token = 'comic'

以下のコードで学習を実行します。これにより、「new_concept」というワードに画像群の画風が圧縮されます。

!accelerate launch examples/textual_inversion/textual_inversion.py \
  --pretrained_model_name_or_path="CompVis/stable-diffusion-v1-4" \
  --use_auth_token \
  --train_data_dir="画像を保存したディレクトリ" \
  --learnable_property=$what_to_teach \
  --placeholder_token="<new_concept>" \
  --initializer_token=$initializer_token \
  --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="result"

学習結果は、カレントディレクトリのresultに保存されます。

推論

学習したコンセプトを用いて画像を生成してみます。

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

# プロンプトの設定
prompt = 'a man of <new_concept> style'

# 画像の生成
with autocast('cuda'):
    image = pipe(prompt)['images'][0]
    image.save('test.png')

9枚生成した結果がこちらです。