【悪用厳禁】Google GolabでDreamboothを使う。Stable Diffusionを再学習(ファインチューニング)

Stable Diffusionで画像生成していると「あの有名人の画像を生成したい」といったことを思うかもしれません。

Stable Diffusionで有名人の画像を学習させて生成することは可能です。Stable Diffusionを0から学習させるとなると数億円のお金が必要ですが、転移学習(ファインチューニング)という元あるモデルを追加で学習させる場合だと、クラウドを使えば月額1000円出せば、学習させ放題でほとんどお金はかかりません。

ただし、実際に有名人の画像を使って記事を書いてしまうと肖像権を侵害してしまいます。

そこで実際に有名人の画像を使わずに解説するのであれば大丈夫だろうということでStable Diffusionを転移学習(ファインチューニング)することで、特定の人物を画像生成の呪文(プロンプト)から画像に登場させる方法について解説します。

この記事ではDreamboothの手法とGoogle Colabを使用したファインチューニングの方法を紹介します。

ちなみにdreambooth関連でStable DiffusionのLoRAを学習させるkohya lora dreamboothの使い方については以下を参考にしてみてください。

ファインチューニングの手法の1つDreambooth

Stable Diffusionモデルの学習には次の3つの方法があります。

  1. Stable Diffusionを追加データを使って普通に学習させる
  2. Stable DiffusionをDreamboothで学習させる
  3. Stable DiffusionのLoRAを学習させる。

その1つの手法がDreamboothです。Stable Diffusionのファインチューニングの手法の1つでAWSなどでもDreamboothをサービスの一つとして提供されています。メリットとしては次のようなものが挙げられます。

  • 少ない学習量やデータでも十分な結果が得られる。
  • ファインチューニングのためのスクリプトが多数用意されている。
  • 既存の呪文(プロンプト)に影響を与えずに学習させることができる。

個人的には、通常のファインチューニングよりも、調整しやすくて良い結果が得られやすい印象があります。本記事では、Dreamboothの具体的な理論説明については省きますが、もし理論的な背景に興味があれば論文を読むことをおすすめします。

DreamBooth

LoRAによるファインチューニングも人気の方法なので、以下の記事も参考にしてみてください。

Google ColabでDreamboothによるファインチューニングを行う

Dreamboothによるファインチューニングの手順ですが大きく分けて、次の6つの工程を行います。

  1. ノートブックを取得する
  2. 環境の設定
  3. ベースとなるモデルのダウンロード
  4. 学習データの設定/アップロード
  5. トレーニングスクリプトの引数を指定して実行
  6. Stable Diffusion Web UI(Auto1111) から使えるようにモデルを変換

ノートブックを取得する(1/6)

diffusers/examples/dreambooth at main · ShivamShrirao/diffusers
🤗 Diffusers: State-of-the-art diffusion models for image and audio generation in PyTorch - ShivamShrirao/diffusers

へアクセスして”Open in Colab”をクリックしてください。

Google Colabの画面が表示されます。

画面上部の「ファイル」→「ドライブにコピーを保存」をクリックし、「このノートブックを自分のGoogle Driveにコピー」を選択してコピーを作成します。

環境の設定(2/6)

次に、ランタイムタイプを設定します。「ランタイム」→「ランタイムのタイプを変更」をクリックし、GPUのタイプをT4に、ランタイムの仕様を「標準」に設定します。

それではノートブックを実行していきましょう。一番上のセルから順次実行していきます。

ベースとなるモデルのダウンロードと設定(3/6)

Login to HuggingFace 🤗のセルでは、自身のHUGGINFACE_TOKENを指定する必要があります。無料でHUGGINFACEのアカウントは作成できるのでHUGIINGFACEに登録してAccess Tokenを取得してください。

取得したAccess Tokenを以下のように入力してください(画像上では1111…)。セルを実行してください。(Shift+Enterを押すとセルを実行)

次のセルでは以下のような設定を行いセルを実行してください。

save_to_drive: チェックを入れます。チェックをいれることでGoogle Driveにモデルが保存されます。

MODEL_NAME: runwayml/stable-diffusion-v1-5をベースモデルとしてしていします。ベースモデルですがHUGGGIN FACEのリポジトリがあるものについてはこのようにパスを指定すればダウンロードがされます。

OUPUT_DIR: モデルが保存される場所です。ここではstable_diffusion_weights/finetune-sdv15としますが、自由に指定してください。Driveへのアクセス権を尋ねられるので、許可してください。

次の訓練の目安の説明です。これはただの説明なので実行される処理はありません。

学習データの設定(4/6)

次のセルでは、concepts_listを指定します。ここで学習させたい、人物の呪文(プロンプト)やファイルを指定します。

instance_prompt: 学習させたい人物の名前などを入れたオリジナルな呪文(プロンプト)です。人物の名前ですが、今回はNanasiさんとしておきます。

class_prompt: 学習させたい対象を含むクラスです。女性を学習させたい場合はphoto of womanとしておきます。

instance_data_dir: 学習させたい人物の画像が保存されたフォルダです。次のステップで10枚程度画像をアップロードするので、ここではアップロード予定のディレクトリを指定しておきます。アップロード先はGoogle Drive直下のにしますが、パスは/content/drive/MyDrive/から始める必要があります。今回は/content/drive/MyDrive/Nanasiを指定します。

class_data_dir: こちらはdreamboothの特徴である、事前学習モデルからclass_promptに対応する画像を生成して訓練に使うのでそれの保存先のフォルダです。こちらはアップロードする必要はないので、どこに保存しても構いません。

# You can also add multiple concepts here. Try tweaking `--max_train_steps` accordingly.

concepts_list = [
    {
        "instance_prompt":      "photo of a nanasi woman",
        "class_prompt":         "photo of a woman",
        "instance_data_dir":    "/content/drive/MyDrive/nanasi",
        "class_data_dir":       "/content/drive/MyDrive/woman"
    },
]

# `class_data_dir` contains regularization images
import json
import os
for c in concepts_list:
    os.makedirs(c["instance_data_dir"], exist_ok=True)

with open("concepts_list.json", "w") as f:
    json.dump(concepts_list, f, indent=4)

次のセルでは画像をドライブにアップロードします。このセルに学習させたい対象画像(nanasiの画像を10枚程度)をドラッグアンドドロップするとアップロードされます。Google Driveから直接アップロードしても構いません。直接アップロードする場合は、content/drive/MyDrive/nanasiを指定しえいる場合は、nanasiというフォルダを作成してその下に画像を10枚程度アップロードしてください。

トレーニングの引数を設定(5/6)

それでは訓練を開始しますがここでコマンドの引数を調整します。

  • revision=”f16”を削除
  • max_train_stepsを2000に変更(10枚ぐらいだとこのぐらい)
  • save_sample_promptを”photo of a nanasi woman”に変更
  • save_intervalを500に設定(Option)
!python3 train_dreambooth.py \\
  --pretrained_model_name_or_path=$MODEL_NAME \\
  --pretrained_vae_name_or_path="stabilityai/sd-vae-ft-mse" \\
  --output_dir=$OUTPUT_DIR \\
  --with_prior_preservation --prior_loss_weight=1.0 \\
  --seed=1337 \\
  --resolution=512 \\
  --train_batch_size=1 \\
  --train_text_encoder \\
  --mixed_precision="fp16" \\
  --use_8bit_adam \\
  --gradient_accumulation_steps=1 \\
  --learning_rate=1e-6 \\
  --lr_scheduler="constant" \\
  --lr_warmup_steps=0 \\
  --num_class_images=50 \\
  --sample_batch_size=4 \\
  --max_train_steps=2000 \\
  --save_interval=500 \\
  --save_sample_prompt="photo of a nanasi woman" \\
  --concepts_list="concepts_list.json"

それではセルを実行してみましょう。1時間ほどすると訓練が完了します。

Stable Diffusion Web UI(Auto1111) から使えるようにモデルを変換(6/6)

訓練が完了したらモデルファイルを使いやすい形式に変換します。まずはWEIGHT_DIRを設定します。空欄にすると訓練結果のモデルが指定されるので空欄のまま実行してください。

次のセルを実行するとckptファイルが保存されます。このckptファイルが重要となります。

また、最近はckptよりもsafetensorsのほうが一般的なので、その保存については新しくセルを作って、以下のコードを実行するとsafetensors形式でも保存されます。

#@markdown Run conversion.
ckpt_path = WEIGHTS_DIR + "/model.safetensors"

half_arg = ""
#@markdown  Whether to convert to fp16, takes half the space (2GB).
fp16 = True #@param {type: "boolean"}
if fp16:
    half_arg = "--half"

!python convert_diffusers_to_original_stable_diffusion.py --model_path $WEIGHTS_DIR  --checkpoint_path $ckpt_path $half_arg --use_safetensors
print(f"[*] Converted ckpt saved at {ckpt_path}")

使い方

モデルのチェックポイント(ckptまたはsafetensors)をダウンロードしてStable Diffusion web UIでチェックポイントが保存されているディレクトリにチェックポイントを配置して、text2imgのチェックポイントの更新をクリックすることで今回作成したチェックポイントが利用できます。

promptに”photo of a nanasi woman”を入力すると、学習させた人物が生成されます。

またinpaintで、顔を置き換えしたい写真の顔にマスクをして、promptにphoto of a nanasi womanを入力し、inpaintを実行すると、顔の置き換えができます。体全体が写っている画像だと精度が悪化しがちなので、胴より上の画像を置き換え元として選ぶのをおすすめします。

そのほかコツについて

  • 画像は大量に用意するよりも少ない画像で良いので高画質であることを重視
  • 画像は大量(100~200枚)に用意すると逆に品質が悪化する
  • 複数の別々のインスタンス呪文(プロンプト)を同時に学習させたりできるが、個別に学習させてモデルを分けたほうが品質は高い
  • 横や縦に引き伸ばしたりするとinpaintなどで利用しやすい
  • うつむき、横向きなど様々な画角を用意するのが良い
  • 胸から上の画像を学習させるのがおすすめ
  • 全身画像を生成するとなると精度が落ちる。
  • Beautiful Realistic AsianやChilloutmixなどを高品質な事前学習済みモデルを利用すると学習が上手く行かないので。、SD 1.5がおすすめです。(SD 2系も学習の進捗が遅いです。)
  • 体全体の学習は難易度がたかく、inpaintを使う場合も精度が悪化するので、胴より上を切り抜いて入力画像にして、顔にマスクをかけて、inpaintしてから、もとの切り抜いた位置に画像を合成すると高品質の画像が生成できます。

LoRAもおすすめ

ファインチューニングの方法にはLoRAという方法もあり、こちらのほうが主流になりつつあるため、LoRAについて確認しておきましょう!

コメント

  1. Convert weights to ckpt to use in web UIs like AUTOMATIC1111.で

    FileNotFoundError: [Errno 2] No such file or directory: ‘/content/drive/MyDrive/stable_diffusion_weights/finetune-sdv15/2000/unet/diffusion_pytorch_model.bin’

    と出ます

  2. エラーメッセージを見る限りはWEIGHT_DIRは/content/drive/MyDrive/stable_diffusion_weights/finetune-sdv15/2000/に設定されています。何らかの原因でファイルの保存に失敗しているか、途中で処理が停止してフォルダだけは作成されて中身が保存されてないなどの可能性があります。