DiffusersベースのDreamBooth(SD2.0対応、Windows対応、必要VRAM 12GB~)

(12/20)新たにリポジトリを作成しましたので最新版はそちらから取得してください。解説記事は以下になります。

(12/17)スクリプトをv18に更新しました。目次から「学習用スクリプトのダウンロード」に飛んでください。

はじめに

先日、Diffusers版DreamBoothText Encoderまで学習するようにして精度を上げる記事を書きましたが、依然としてXavierXiao氏のStableDiffusion版よりも若干精度が落ちるように感じていました。調査の結果、その原因らしきものが分かりましたので、それを解消した学習スクリプトを作成しました。
手元では精度が向上しているように思われますが、追試、確認いただけると幸いです。

Pythonで仮想環境を構築できるくらいの方を対象にしています。また細かいところは省略していますのでご容赦ください。

(10/28)VRAM 12GBの動作報告をいただきました!(ただしv2ではVRAM 12GBは厳しいようです)

同じスクリプトを利用したaugmentation編fine tuning編の記事も公開しておりますので興味があればあわせてお読みください。

※使用に当たっては自己責任でお願いいたします。

概要

  • Diffusers版とStableDiffusion版のDreamBoothを同じ学習データ、設定で学習しても結果が異なる。StableDiffusion版のほうが個人的には良い結果を得られている。

  • 調査したところ、Diffusers版DreamBoothはStableDiffusion版と異なり、Tokenizerの出力をpaddingせずにText Encoderに渡している。

    • その後のupdateでDiffUsers版もpaddingするようになりました(追記)

  • 学習スクリプトを改造し77トークンまでpaddingするようにして学習したところ、学習結果がかなり近くなった。

  • スクリプトにはあわせて以下の機能追加を行った。

    • 8bit Adam optimizerおよびlatentのキャッシュによる省メモリ化(ShivamShrirao氏版と同様)。

    • xformersによる省メモリ化。

    • 512*512だけではなく任意サイズでの学習。

    • augmentationによる品質の向上。

    • DreamBoothだけではなくText Encoder+U-Netのfine tuningに対応。

    • StableDiffusion形式でのモデルの読み書き。

    • Aspect Ratio Bucketing。

  • Stable Diffusion v2.0対応(11/27)。

  • 導入を容易にするよう、既存仮想環境へのソースコードの変更等は最低限とした。

以下、環境構築から簡単に解説します。

環境構築

※Windowsでxformersを使う場合、Python 3.10をお勧めします(3.8、3.9の場合、自前でxformersをビルドする必要があります)。

適当なディレクトリに仮想環境を構築します。以下はvenvを使った例です。

E:\Workへ作りました
python -m venv --system-site-packages venv_diffusers
venv_diffusers\Scripts\activate
python.exe -m pip install --upgrade pip

pipを更新しておきます。

PyTorchとTorchvisionを入れる

Diffusersを入れる前にPyTorch 1.12.1とTorchvision 0.13.1を入れます。以前のバージョンよりもわずかにメモリ使用量が少ないようです(DiffusersをインストールするとPyTorchもインストールされますがCPU版になるようです)。お使いのCUDAのバージョンにあわせたバージョンを入れてください。

※12/15現在、1.13.0でのaccelerateのエラーは解消されたようですが、後述のxformersがエラーになるようです。

以下はCUDA 11.6の例です。

pip install torch==1.12.1+cu116 torchvision==0.13.1+cu116 --extra-index-url https://download.pytorch.org/whl/cu116

11.3ならこちらになります。
※後述のxformersを入れる場合は1.12.1をお勧めします。

pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 --extra-index-url https://download.pytorch.org/whl/cu113

Diffusersと依存ライブラリを入れる

Diffusersのページに従いインストールします。0.10.0以降が必要となりますが0.10.0には問題があるようですので0.10.2を入れてください。私は0.10.0で動作確認しています(0.10.2で問題があればご連絡ください)。0.10.2を入れるには以下の通りです。

pip install diffusers[torch]==0.10.2

次にDreamBoothの学習スクリプトに必要な以下のライブラリを入れます。(12/10・スクリプト側でaccelerateの0.15.0へ対応しましたが、0.14.0指定を当面入れておきます。)

accelerate==0.14.0
transformers>=4.21.0
ftfy
albumentations
opencv-python
einops
safetensors

以下のrequirements.txtを作業ディレクトリにダウンロードしてインストールしてください。

pip install -r requirements.txt

またStableDiffusion版のcheckpointを読み書きする場合、以下のライブラリをインストールします。

pip install pytorch_lightning

※11/8:OmegaConfはたぶん不要になりました。

xformersのインストール(オプション)

CrossAttentionを省メモリ化、高速化するためxformersを入れます。Linuxの場合、xformersのページに従ってインストールしてください。

WindowsでPython 3.10、GeForce 10XX(Pascal)以降のGPUの場合、Automatic1111氏のStable Diffusion web UIのWikiを参考に、以下のコマンドでインストールできます。

pip install -U -I --no-deps https://github.com/C43H66N12O12S2/stable-diffusion-webui/releases/download/f/xformers-0.0.14.dev0-cp310-cp310-win_amd64.whl

それ以外のPythonバージョン、GPUの場合、Wikiを参考に自前でビルドしてください(かなりの時間が掛かります)。

bitsandbytes (8-bit Adam optimizer)のインストール(オプション)

省メモリ化のため8-bit Adam optimizerを入れます。Linuxの場合、bitsandbytesのページに従ってインストールしてください。

Windowsの場合は以下の記事に従ってインストールしてください。

accelerateを設定する

Diffusers版Dreamboothのページにあるようにaccelerate configを実行します。以下はバージョン0.13での例です。

accelerate config

NOTE: Redirects are currently not supported in Windows or MacOs.
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]:
Do you want to use DeepSpeed? [yes/NO]:
What GPU(s) (by id) should be used for training on this machine as a comma-seperated list? [all]:
Do you wish to use FP16 or BF16 (mixed precision)? [NO/fp16/bf16]: bf16

GPU 1枚で学習する場合、0、0、そのままエンター、そのままエンター、そのままエンターでOKでしょう。
mixed precision設定は最も省メモリなのはbf16またはfp16になります(bf16でも動作しました)。使用VRAMを極限まで減らしたい場合、速度を上げたい場合、xformersでエラーが出る場合にはbf16またはfp16を、それ以外は精度が高いnoを指定するとよいでしょう。

学習用スクリプトのダウンロード

こちらのzipファイルをダウンロード、展開して作業ディレクトリに配置してください。二つのスクリプトが含まれていますので、同じディレクトリに置いてください(train_dreambooth.pyを参考に書いたものです)。

※12/17 (v18)の更新情報:
読み込んだモデル形式に関係なく任意の形式で保存できるようになりました。save_model_asオプションに、「ckpt、safetensors、diffusers、diffusers_safetensors」のいずれかを指定してください。

過去の更新履歴は末尾に移しました。ファイルがバグっていた場合に備えて旧バージョンを記事の末尾に置いておきます。

前述のようにTokenizerの出力を77トークンまでpaddingしています。また単純化のためめったに使わないようなオプションは省略しました(ソースコードを直接、修正してください)。またU-Netのみの学習はできません。

※省メモリ化のため、ステップ当たりの学習回数がtrain_dreambooth.pyの半分になっています(対象の画像と正則化画像を同一のバッチではなく別のバッチに分割して学習するため)。元のDiffusers版やXavierXiao氏のStableDiffusion版とほぼ同じ学習を行うには、ステップ数を倍にしてください
(shuffle=Trueのため厳密にはデータの順番が変わってしまいますが、学習には大きな影響はないと思います。)

以上で環境構築は完了です。

学習を行う

StableDiffusion版のcheckpointをDiffusers版に変換する方法については、以前の記事をご参照ください。

学習用画像の準備

学習用画像を格納するディレクトを作成します。さらにその中に、以下の名前でディレクトリを作成します。

<繰り返し回数>_<プロンプト>

繰り返し回数は、正則化画像と枚数を合わせるために指定します(後述します)。
たとえば「sls frog」というプロンプトで、データを20回繰り返す場合、「20_sls frog」となります。以下のようになります。

ダ鳥獣戯画さんのカエル画像15枚を用意した

正則化画像の準備

正則化画像を格納するディレクトを作成します。学習用画像と同様に「<繰り返し回数>_<プロンプト>」という名前でディレクトリを作成します。たとえば「frog」というプロンプトで、データを繰り返さない(1回だけ)場合、以下のようになります。

StableDiffusionで生成した

繰り返し回数は「学習用画像の繰り返し回数×学習用画像の枚数≧正則化画像の繰り返し回数×正則化画像の枚数」となるように指定してください。
(1 epochのデータ数が「学習用画像の繰り返し回数×学習用画像の枚数」となります。正則化画像の枚数がそれより多いと、余った部分の正則化画像は使用されません。)

学習の実行

スクリプトを実行します。最大限、メモリを節約したコマンドは以下のようになります(実際には1行で入力します)。

accelerate launch --num_cpu_threads_per_process 8 train_db_fixed.py 
    --pretrained_model_name_or_path=<Diffusers版モデルのディレクトリ> 
    --train_data_dir=<学習用データのディレクトリ> 
    --reg_data_dir=<正則化画像のディレクトリ> 
    --output_dir=<学習したモデルの出力先ディレクトリ> 
    --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

num_cpu_threads_per_processにはCPUコア数を指定するとよいようです。

pretrained_model_name_or_pathに追加学習を行う元となるモデルを指定します。Stable Diffusionのcheckpointファイル、Diffusersのローカルディスクにあるモデルディレクトリ、DiffusersのモデルID("stabilityai/stable-diffusion-2"など)が指定できます。学習後のモデルの保存形式は元のモデルと同じになります。

prior_loss_weightは正則化画像のlossの重みです。通常は1.0を指定します。
resolutionは画像のサイズ(解像度、幅と高さ)になります。学習用画像、正則化画像はこのサイズとしてください。
train_batch_sizeは学習時のバッチサイズです。また前述のとおりmax_train_stepsを1600とします。
学習率learning_rateは、diffusers版では5e-6ですがStableDiffusion版は1e-6ですのでここでは1e-6を指定しています。

省メモリ化のためmixed_precision="bf16"(エラーが出る場合は"fp16"としてください)、およびgradient_checkpointing を指定します。

xformersオプションを指定し、xformersのCrossAttentionを用います。xformersをインストールしていない場合、エラーとなる場合(mixed_precisionなしの場合、私の環境ではエラーとなりました)、代わりにmem_eff_attnオプションを指定すると省メモリ版CrossAttentionを使用します(速度は遅くなります)。

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

ある程度メモリがある場合はたとえば以下のように指定します。

accelerate launch --num_cpu_threads_per_process 8 train_db_fixed.py 
    --pretrained_model_name_or_path=<Diffusers版モデルのディレクトリ> 
    --train_data_dir=<学習用データのディレクトリ> 
    --reg_data_dir=<正則化画像のディレクトリ> 
    --output_dir=<学習したモデルの出力先ディレクトリ> 
    --prior_loss_weight=1.0 
    --resolution=512 
    --train_batch_size=4 
    --learning_rate=1e-6 
    --max_train_steps=400 
    --use_8bit_adam 
    --xformers 
    --cache_latents

gradient_checkpointingを外し高速化します(メモリ使用量は増えます)。またmixed_precisionを外し精度向上を図ります。
バッチサイズを増やし、高速化と精度向上を図ります。

このあたりのパラメータの関係については前回の記事の末尾をご参照ください。また学習後のDiffusers版モデルをStableDiffusion版checkpointに変換する方法も前の記事にあります。

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

学習が終わると指定したフォルダにlast.ckptという名前でcheckpointが出力されます(DiffUsers版モデルを学習した場合はlastフォルダになります)。

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

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

hasui0sd2-2-epoch-000001.ckptとlast.ckptを置いた例

各yamlファイルはStability AIのSD2.0のリポジトリにありますが、以下にも添付しておきます(MITライセンス)。

その他の学習オプション

Stable Diffusion 2.0対応  --v2 / --v_parameterization

今のところ(11/27時点では)DiffUsersのDreamBoothはSD2.0(特にv-parameterization)に対応していないため、独自に実装したものです。そのため不具合等あるかもしれませんのでご容赦ください。

Hugging Faceのstable-diffusion-2-baseを使う場合はv2オプションを、stable-diffusion-2または768-v-ema.ckptを使う場合はv2とv_parameterizationの両方のオプションを指定してください。

なおSD 2.0の学習はText Encoderが大きくなっているためVRAM 12GBでは厳しいようです。

Stable Diffusion 2.0では大きく以下の点が変わっています。

  1. 使用するTokenizer

  2. 使用するText Encoderおよび使用する出力層(2.0は最後から二番目の層を使う)

  3. Text Encoderの出力次元数(768->1024)

  4. U-Netの構造(CrossAttentionのhead数など)

  5. v-parameterization(サンプリング方法が変更されているらしい)

このうちbaseでは1~4が、baseのつかない方(768-v)では1~5が採用されています。1~4を有効にするのがv2オプション、5を有効にするのがv_parameterizationオプションです。

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

stop_text_encoder_trainingオプションに数値を指定すると、そのステップ数以降はText Encoderの学習を行わずU-Netだけ学習します。場合によっては精度の向上が期待できるかもしれません。
(恐らくText Encoderだけ先に過学習することがあり、それを防げるのではないかと推測していますが、詳細な影響は不明です。)

VAEを別途読み込んで学習する --vae

vaeオプションにStable Diffusionのcheckpoint、VAEのcheckpointファイル、DiffusesのモデルまたはVAE(ともにローカルまたはHugging FaceのモデルIDが指定できます)のいずれかを指定すると、そのVAEを使って学習します(latentsのキャッシュ時または学習中のlatents取得時)。
保存されるモデルはこのVAEを組み込んだものになります。

学習途中での保存 --save_every_n_epochs / --save_state / --resume

save_every_n_epochsオプションに数値を指定すると、そのエポックごとに学習途中のモデルを保存します。
save_stateオプションを同時に指定すると、optimizer等の状態も含めた学習状態を合わせて保存します(checkpointから学習再開するのに比べて、精度の向上、学習時間の短縮が期待できます)。学習状態は保存先フォルダに"epoch-??????-state"(??????はエポック数)という名前のフォルダで出力されます。長時間にわたる学習時にご利用ください。

保存された学習状態から学習を再開するにはresumeオプションを使います。学習状態のフォルダを指定してください。

なおAcceleratorの仕様により(?)、エポック数、global stepは保存されておらず、resumeしたときにも1からになりますがご容赦ください。

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

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

任意サイズの画像での学習 --resolution

正方形以外で学習できます。resolutionに「448,640」のように「幅,高さ」で指定してください。幅と高さは64で割り切れる必要があります。学習用画像、正則化画像のサイズを合わせてください。

個人的には縦長の画像を生成することが多いため、最近は「448,640」で学習しています。

Aspect Ratio Bucketing --enable_bucket / --min_bucket_reso / --max_bucket_reso

enable_bucketオプションを指定すると有効になります。Stable Diffusionは512*512で学習されていますが、それに加えて256*768や384*640といった解像度でも学習します。

このオプションを指定した場合は、学習用画像、正則化画像を特定の解像度に統一する必要はありません。いくつかの解像度(アスペクト比)から最適なものを選び、その解像度で学習します。
解像度は64ピクセル単位のため、元画像とアスペクト比が完全に一致しない場合がありますが、その場合は、はみ出した部分がわずかにトリミングされます。

解像度の最小サイズをmin_bucket_resoオプションで、最大サイズをmax_bucket_resoで指定できます。デフォルトはそれぞれ256、1024です。たとえば最小サイズに384を指定すると、256*1024や320*768などの解像度は使わなくなります。
解像度を768*768のように大きくした場合、最大サイズに1280などを指定しても良いかもしれません。

なおAspect Ratio Bucketingを有効にするときには、正則化画像についても、学習用画像と似た傾向の様々な解像度を用意した方がいいかもしれません。
(ひとつのバッチ内の画像が学習用画像、正則化画像に偏らなくなるため。そこまで大きな影響はないと思いますが……。)

保存時のデータ精度の指定 --save_precision

save_precisionオプションにfloat、fp16、bf16のいずれかを指定すると、その形式でcheckpointを保存します(Stable Diffusion形式で保存する場合のみ)。checkpointのサイズを削減したい場合などにお使いください。

任意の形式で保存する --save_model_as

モデルの保存形式を指定します。ckpt、safetensors、diffusers、diffusers_safetensorsのいずれかを指定してください。
Stable Diffusion形式(ckptまたはsafetensors)を読み込み、Diffusers形式で保存する場合、不足する情報はHugging Faceからv1.5またはv2.1の情報を落としてきて補完します。

safetensors形式で保存する --use_safetensors

このオプションを指定するとsafetensors形式でcheckpointを保存します。保存形式はデフォルト(読み込んだ形式と同じ)になります。

学習ログの保存 --logging_dir / --log_prefix

logging_dirオプションにログ保存先フォルダを指定してください。TensorBoard形式のログが保存されます。

たとえば--logging_dir=logsと指定すると、作業フォルダにlogsフォルダが作成され、その中の日時フォルダにログが保存されます。
また--log_prefixオプションを指定すると、日時の前に指定した文字列が追加されます。「--logging_dir=logs --log_prefix=db_style1_」などとして識別用にお使いください。

TensorBoardでログを確認するには、別のコマンドプロンプトを開き、作業フォルダで以下のように入力します(tensorboardはDiffusersのインストール時にあわせてインストールされると思いますが、もし入っていないならpip install tensorboardで入れてください)。

tensorboard --logdir=logs

その後ブラウザを開き、http://localhost:6006/へアクセスすると表示されます。

学習率のスケジューラ関連の指定 --lr_scheduler / --lr_warmup_steps

lr_schedulerオプションで学習率のスケジューラをlinear, cosine, cosine_with_restarts, polynomial, constant, constant_with_warmupから選べます。デフォルトはconstantです。lr_warmup_stepsでスケジューラのウォームアップ(だんだん学習率を変えていく)ステップ数を指定できます。詳細については各自お調べください。

勾配をfp16とした学習(実験的機能) --full_fp16

full_fp16オプションを指定すると勾配を通常のfloat32からfloat16(fp16)に変更して学習します(mixed precisionではなく完全なfp16学習になるようです)。これによりSD1.xの512*512サイズでは8GB未満、SD2.xの512*512サイズで12GB未満のVRAM使用量で学習できるようです。

あらかじめaccelerate configでfp16を指定し、オプションでmixed_precision="fp16"としてください(bf16では動作しません)。

メモリ使用量を最小化するためには、xformers、use_8bit_adam、cache_latents、gradient_checkpointingの各オプションを指定し、train_batch_sizeを1としてください。
(余裕があるようならtrain_batch_sizeを段階的に増やすと若干精度が上がるはずです。)

PyTorchのソースにパッチを当てて無理やり実現しています(PyTorch 1.12.1と1.13.0で確認)。精度はかなり落ちますし、途中で学習失敗する確率も高くなります。学習率やステップ数の設定もシビアなようです。それらを認識したうえで自己責任でお使いください。

その他の学習方法

マルチクラス、マルチサブジェクト学習

(ここでクラスは「dog」「1girl」など学習対象の分類、サブジェクトは特定のキャラクタなどを指します。)

方法は単純で、学習用画像、正則化画像のディレクトリ内に「繰り返し回数_プロンプト」のディレクトリを複数用意してください。

たとえば「sls frog」と「cpc rabbit」を同時に学習する場合、以下のようになります。

クラスがひとつでサブジェクトが複数の場合、正則化画像ディレクトリはひとつで構いません。たとえば1girlにキャラAとキャラBがいる場合は次のようにします。

  • train_girls

    • 10_sls 1girl

    • 10_cpc 1girl

  • reg_girls

    • 1_1girl

いずれかのサブジェクト、クラスが過学習する場合、枚数や繰り返し回数を減らすと改善されるかもしれません。

DreamBoothでキャプションを使う

学習用画像、正則化画像のフォルダに、画像と同じファイル名で、拡張子.caption(オプションで変えられます)のファイルを置くと、そのファイルからキャプションを読み込み、フォルダ名に指定したプロンプトに追加して学習します。たとえばフォルダ名が「10_cpc frog」で、キャプションが「, lying, white background」(先頭にカンマがあります)の場合、学習時のプロンプト(キャプション)は「cpc frog, lying, white background」になります。

各画像にキャプションを付けることで(BLIP等を使っても良いでしょう)、学習したい属性をより明確にできるかもしれません。

キャプションファイルの拡張子はデフォルトで.captionです。caption_extensionで変更できます。shuffle_captionオプションで学習時のキャプションについて、カンマ区切りの各部分をシャッフルしながら学習します。

事前に顔検出した画像での顔を中心としたランダム切り出しaugmentation、色合いのaugmentation、左右反転augmentation、モデルのfine tuning

記事が長くなりすぎたので別記事としました。

おわりに

今回のスクリプトでDiffusers版DreamBoothの微妙な点が改善されたように思います。少しでも多くの方がより高品質なDreamBoothを利用できるようになれば幸いです。

ラーメン店で箸でラーメンを食べる鳥獣戯画のカエル
ラーメンではなく野菜炒めだな

スクリプトの更新履歴

※12/13 (v17)の更新情報:勾配までfp16にする学習に対応しました(実験的機能)。SD1.xでVRAM 8GBで学習可能となります。full_fp16オプションを指定してください。詳細は目次から「勾配をfp16とした学習(実験的機能)」をご覧ください。

※12/10 (v16)の更新情報:DiffUsers 0.10.2が必要となります(0.10.0以降で動きますが0.10.0には問題があるようですので0.10.2を使ってください)。仮想環境内で「pip install -U diffusers[torch]==0.10.2」として更新してください。Diffusers 0.10に対応しました(v-parameterizationの学習にDiffusers内のコードを使います。またDiffusersでもsafetensorsに対応します)。accelerate 0.15.0に対応しました。

※12/5 (v15)の更新情報:一部のモデルでsafetensor形式の保存時にエラーになるとのことで暫定対応しました。v14でエラーになる場合はv15をお試しください。動作確認をしっかり行ったらv14は消します。

※12/5 (v14)の更新情報:スクリプトが2つに分かれたため.zip形式としました。展開してください。.safetensors形式に対応しました。「pip install safetensors」としてsafetensorsをインストールしてください。読み込み時は拡張子で自動判定します。保存時はuse_safetensorsオプションを指定してください。VAEを別途読み込むvaeオプションを追加しました。日時のログディレクトリ名の前に任意の文字列を追加するlog_prefixオプションを追加しました。

※11/30 (v13)の更新情報:stop_text_encoder_trainingオプションが正しく動作していない(全体の学習が止まる)ようなので修正しました。申し訳ありません。
※float32での学習がエラーになる場合、mem_eff_attnオプションを追加してください。

※11/29 (v12)の更新情報:Text Encoderの学習を指定ステップで止める--stop_text_encoder_trainingオプションを追加しました。

※11/27 (v11)の更新情報:
DiffUsers 0.9.0が必要となります。仮想環境内で「pip install -U diffusers[torch]==0.9.0」として更新してください。またその他のエラーが出るときには「pip install --upgrade -r requirements.txt」として依存ライブラリの更新を行ってください。

DreamBooth時のキャプションの扱いが変わりました。キャプションファイルが存在したとき、v10まではフォルダのキャプションにファイルのキャプションが追加されていましたが、v11からはファイルのキャプションのみになります。ご注意ください。
prior_loss_weightが学習画像のほうに適用されていたバグを修正。ご不便をお掛けしました。
Stable Diffusion v2.0に対応。--v2オプションを追加してください。stable-diffusion-v2-baseではなく768-v-ema.ckptまたはstable-diffusion-2を使う場合は--v_parameterizationも合わせて追加してください。詳しくはその他のオプションをご覧ください。「学習したモデルで画像生成する」もあわせてごらんください。
学習率のスケジューラに関するオプションを追加。
DiffUsersのモデルをHugging Faceから直接ダウンロードして学習できるように。またDiffUsersのモデルでも学習途中の保存が可能に。

※11/21 (v10):Aspect Ratio Bucketing使用時の最小/最大解像度指定を追加(min_bucket_reso/max_bucket_resoオプション)。
キャプションファイルの拡張子指定を追加(caption_extention)。(旧バージョンまでは.captionと.txtを自動的に読み込んでいましたが当バージョンでは.captionのみとなります。.txtをお使いの場合、お手数ですが当オプションで--caption_extention=.txtと指定してください。)(思い切りスペルミスしてるので次バージョンで直します……。)
拡張子.webpの画像に対応。学習用画像・正則化画像にキャプションを付けられる機能を追加。

※11/19 (v9):Aspect Ratio Bucketingに対応しました(enable_bucketオプション)。checkpoint保存時のデータ形式の選択(fp16/bf16/float)、学習状態の保存、ロギングに対応しました。詳細は「その他の学習オプション」をご覧ください。

※11/8 (v8):Diffusers 0.7.2に対応しました。

※11/7 (v7):Text Encoderの格納形式が異なるcheckpointファイルに対応(読み込み時に変換しますので書き出しは通常形式になります)。epochのlossの平均値を画面に出力するよう変更。SD形式のcheckpointにepochとglobal stepを保存するよう機能追加(既存のデータがある場合は値を加算)。fine tuning時にreg_data_dirオプションが有効に(正則化画像を混ぜながらfine tuningします)。fine tuning時に有効なdataset_repeatsオプションを追加(教師画像枚数が少なくepochが極端に短くなる時に指定)。

※10/30 (v6):debug_datasetがバグっていたのを修正しました。たびたびすみません。

※10/30 (v5):DreamBoothでlatentsをキャッシュする場合、ほぼ変わりませんのでv4のままで構いません。DataLoaderをマルチプロセスで動かすようにしました(latentsをキャッシュしない場合に速度向上)。latentsのキャッシュの仕組みを修正し、fine tuning時にshuffle_captionが効くようにしました。clip_skipオプションを追加しました。checkpointが最後のepochでも保存されるのを修正しました。

※10/29 (v4):paddingしないモード(no_token_paddingオプション、後述)を追加しました。またText Encoderがtrainモードになっていないバグ(参照したtrain_dreambooth.pyのバグ)を修正しました。

※10/29 (v3):StableDiffusion形式のモデルの読み書きに対応しました。オプションは変わりません(ファイルを指定するとStableDiffusion形式で、ディレクトリを指定するとDiffusers形式で読み込みます)。書き出しは読み込みと同形式になります。またsave_every_n_epochsオプション(後述)を追加しました。

※10/28:VRAM使用量を削減するよう修正しました(optimizerのzero_gradにset_to_none=Trueを追加)。

以前のバージョンのスクリプト

この記事が気に入ったらサポートをしてみませんか?