Diffusers版Dreamboothを変更して精度を上げる（Text Encoderを学習対象とする）

10/29追記・Tokenizerの仕様を修正してさらに精度を上げる記事を書きましたので、そちらをご覧ください。

概要

Diffusers版のDreamboothは、Stable Diffusionを構成するU-Net (DAE)、Text Encoder、VAE (Autoencoder) のうち、U-Netしか学習していません。しかしながらXavierXiao氏のDreambooth on Stable DiffusionはU-Netに加えてText Encoderも学習しているようです。そのためDiffusers版のDreamboothはXavierXiao氏のバージョンより精度が低い（学習対象を学びにくい）傾向にあるようです。
そこでDiffusers版のDreamboothを変更してText Encoderも学習対象に加えました。

Colabでも動作するようですのでどうぞご利用ください。Colabでの環境整備などについては他の方の記事をご覧ください。
（学習用スクリプトとしてtrain_dreambooth.pyを使う形のnotebookが適しています。）

※無保証です。利用は自己責任でお願いします。
添付ファイルのライセンスは元ファイルに準拠します。ファイルを置き換えるときには念のため元ファイルを別名でコピーしておくことをお勧めします。

変更した学習用スクリプト

DiffusersのDreamboothのセットアップを行った後、以下のtrain_dreambooth_mod.pyをtrain_dreambooth.pyと同じディレクトリに置き、「accelerate launch train_dreambooth_mod.py ～」として学習してください。

train_dreambooth_mod.py

25.9 KB

ファイルダウンロードについて

ダウンロード

変更元は省メモリ化されたShivamShrirao氏のDreamboothです。10/3版を元に変更しました（10/4の変更も取り込んであります）。当該commitのバージョンはこちらです。
not_cache_latentsを指定すると動作しません。

その他のメモ

diffusers版の学習率は5e-6ですがXavierXiao版は1e-6です。
instance_promptは"sks XXX"、class_promptは"XXX"が学習結果が良いようです。
スクリプトの仕様でinstance imageとclass imageの数の多いほうがdatasetの長さになり、かつ短いほうは繰り返されて使用されるため、class imageの枚数はinstance imageの枚数で割り切れるようにすると良さそうです。
（instance imageが50枚、class imageが75枚の時、instance imageの1~25枚目はdatasetが一周するまでに二回使われ、26~50枚目は一回しか使われない。）