概要

DreamBoothとは、画像生成モデルに特定の主題を追加学習し、それを特定の識別子で生成する技術です。論文はこちら。

具体的には、Stable Diffusionのモデルにキャラや画風などを学ばせ、それを shs のような特定の単語で呼び出せる（生成画像に出現させる）ことができます。

スクリプトはDiffusersのDreamBoothを元にしていますが、以下のような機能追加を行っています（いくつかの機能は元のスクリプト側もその後対応しています）。

スクリプトの主な機能は以下の通りです。

8bit Adam optimizerおよびlatentのキャッシュによる省メモリ化（Shivam Shrirao氏版と同様）。
xformersによる省メモリ化。
512x512だけではなく任意サイズでの学習。
augmentationによる品質の向上。
DreamBoothだけではなくText Encoder+U-Netのfine tuningに対応。
Stable Diffusion形式でのモデルの読み書き。
Aspect Ratio Bucketing。
Stable Diffusion v2.0対応。

学習の手順

あらかじめこのリポジトリのREADMEを参照し、環境整備を行ってください。

データの準備

学習データの準備についてを参照してください。

学習の実行

スクリプトを実行します。最大限、メモリを節約したコマンドは以下のようになります（実際には1行で入力します）。それぞれの行を必要に応じて書き換えてください。12GB程度のVRAMで動作するようです。

accelerate launch --num_cpu_threads_per_process 1 train_db.py 
    --pretrained_model_name_or_path=<.ckptまたは.safetensordまたはDiffusers版モデルのディレクトリ> 
    --dataset_config=<データ準備で作成した.tomlファイル> 
    --output_dir=<学習したモデルの出力先フォルダ>  
    --output_name=<学習したモデル出力時のファイル名> 
    --save_model_as=safetensors 
    --prior_loss_weight=1.0 
    --max_train_steps=1600 
    --learning_rate=1e-6 
    --optimizer_type="AdamW8bit" 
    --xformers 
    --mixed_precision="fp16" 
    --cache_latents 
    --gradient_checkpointing

num_cpu_threads_per_process には通常は1を指定するとよいようです。

pretrained_model_name_or_path に追加学習を行う元となるモデルを指定します。Stable Diffusionのcheckpointファイル（.ckptまたは.safetensors）、Diffusersのローカルディスクにあるモデルディレクトリ、DiffusersのモデルID（"stabilityai/stable-diffusion-2"など）が指定できます。

output_dir に学習後のモデルを保存するフォルダを指定します。output_name にモデルのファイル名を拡張子を除いて指定します。save_model_as でsafetensors形式での保存を指定しています。

dataset_config に .toml ファイルを指定します。ファイル内でのバッチサイズ指定は、当初はメモリ消費を抑えるために 1 としてください。

prior_loss_weight は正則化画像のlossの重みです。通常は1.0を指定します。

学習させるステップ数 max_train_steps を1600とします。学習率 learning_rate はここでは1e-6を指定しています。

省メモリ化のため mixed_precision="fp16" を指定します（RTX30 シリーズ以降では bf16 も指定できます。環境整備時にaccelerateに行った設定と合わせてください）。また gradient_checkpointing を指定します。

オプティマイザ（モデルを学習データにあうように最適化＝学習させるクラス）にメモリ消費の少ない 8bit AdamW を使うため、 optimizer_type="AdamW8bit" を指定します。

xformers オプションを指定し、xformersのCrossAttentionを用います。xformersをインストールしていない場合やエラーとなる場合（環境にもよりますが mixed_precision="no" の場合など）、代わりに mem_eff_attn オプションを指定すると省メモリ版CrossAttentionを使用します（速度は遅くなります）。

省メモリ化のため cache_latents オプションを指定してVAEの出力をキャッシュします。

ある程度メモリがある場合は、.toml ファイルを編集してバッチサイズをたとえば 4 くらいに増やしてください（高速化と精度向上の可能性があります）。また cache_latents を外すことで augmentation が可能になります。

よく使われるオプションについて

以下の場合には学習の共通ドキュメントの「よく使われるオプション」を参照してください。

Stable Diffusion 2.xまたはそこからの派生モデルを学習する
clip skipを2以上を前提としたモデルを学習する
75トークンを超えたキャプションで学習する

DreamBoothでのステップ数について

当スクリプトでは省メモリ化のため、ステップ当たりの学習回数が元のスクリプトの半分になっています（対象の画像と正則化画像を同一のバッチではなく別のバッチに分割して学習するため）。

元のDiffusers版やXavierXiao氏のStable Diffusion版とほぼ同じ学習を行うには、ステップ数を倍にしてください。

（学習画像と正則化画像をまとめてから shuffle するため厳密にはデータの順番が変わってしまいますが、学習には大きな影響はないと思います。）

DreamBoothでのバッチサイズについて

モデル全体を学習するためLoRA等の学習に比べるとメモリ消費量は多くなります（fine tuningと同じ）。

学習率について

Diffusers版では5e-6ですがStable Diffusion版は1e-6ですので、上のサンプルでは1e-6を指定しています。

以前の形式のデータセット指定をした場合のコマンドライン

解像度やバッチサイズをオプションで指定します。コマンドラインの例は以下の通りです。

accelerate launch --num_cpu_threads_per_process 1 train_db.py 
    --pretrained_model_name_or_path=<.ckptまたは.safetensordまたはDiffusers版モデルのディレクトリ> 
    --train_data_dir=<学習用データのディレクトリ> 
    --reg_data_dir=<正則化画像のディレクトリ> 
    --output_dir=<学習したモデルの出力先ディレクトリ> 
    --output_name=<学習したモデル出力時のファイル名> 
    --prior_loss_weight=1.0 
    --resolution=512 
    --train_batch_size=1 
    --learning_rate=1e-6 
    --max_train_steps=1600 
    --use_8bit_adam 
    --xformers 
    --mixed_precision="bf16" 
    --cache_latents
    --gradient_checkpointing

学習したモデルで画像生成する

学習が終わると指定したフォルダに指定した名前でsafetensorsファイルが出力されます。

v1.4/1.5およびその他の派生モデルの場合、このモデルでAutomatic1111氏のWebUIなどで推論できます。models\Stable-diffusionフォルダに置いてください。

v2.xモデルでWebUIで画像生成する場合、モデルの仕様が記述された.yamlファイルが別途必要になります。v2.x baseの場合はv2-inference.yamlを、768/vの場合はv2-inference-v.yamlを、同じフォルダに置き、拡張子の前の部分をモデルと同じ名前にしてください。

各yamlファイルはStability AIのSD2.0のリポジトリにあります。

DreamBooth特有のその他の主なオプション

すべてのオプションについては別文書を参照してください。

Text Encoderの学習を途中から行わない --stop_text_encoder_training

stop_text_encoder_trainingオプションに数値を指定すると、そのステップ数以降はText Encoderの学習を行わずU-Netだけ学習します。場合によっては精度の向上が期待できるかもしれません。

（恐らくText Encoderだけ先に過学習することがあり、それを防げるのではないかと推測していますが、詳細な影響は不明です。）

Tokenizerのパディングをしない --no_token_padding

no_token_paddingオプションを指定するとTokenizerの出力をpaddingしません（Diffusers版の旧DreamBoothと同じ動きになります）。