From db8c79c46332c2e3e4b25f84aa98078b73609859 Mon Sep 17 00:00:00 2001
From: Kohya S <ykumeykume@gmail.com>
Date: Sun, 5 Mar 2023 19:51:05 +0900
Subject: [PATCH] Update documentation

---
 train_README-ja.md    | 321 ++++++++++++++++++++++++++++++++++++++----
 train_db_README-ja.md | 291 ++++++++++----------------------------
 2 files changed, 371 insertions(+), 241 deletions(-)

diff --git a/train_README-ja.md b/train_README-ja.md
index bf0d9f9c..5e9a6282 100644
--- a/train_README-ja.md
+++ b/train_README-ja.md
@@ -1,4 +1,6 @@
-当リポジトリではモデルのfine tuning、DreamBooth、およびLoRAとTextual Inversionの学習をサポートします。この文書ではそれらに共通する、学習データの準備方法やスクリプトオプションについて説明します。
+# 学習データの準備について
+
+当リポジトリではモデルのfine tuning、DreamBooth、およびLoRAとTextual Inversionの学習をサポートします。この文書ではそれらに共通する、学習データの準備方法について説明します。
 
 # 概要
 
@@ -8,15 +10,14 @@
 以下について説明します。
 
 1. 学習データの準備について（設定ファイルを用いる新形式）
-1. Aspect Ratio Bucketingについて
+1. 学習で使われる用語のごく簡単な解説
 1. 以前の指定形式（設定ファイルを用いずコマンドラインから指定）
+1. 学習途中のサンプル画像生成
+1. 各スクリプトで共通の、よく使われるオプション
 1. fine tuning 方式のメタデータ準備：キャプションニングなど
 
 1.だけ実行すればとりあえず学習は可能です（学習については各スクリプトのドキュメントを参照）。2.以降は必要に応じて参照してください。
 
-<!--
-1. 各スクリプトで共通のオプション
--->
 
 # 学習データの準備について
 
@@ -36,7 +37,7 @@
 
 1. fine tuning方式（正則化画像使用不可）
 
-    あらかじめキャプションをメタデータファイルにまとめます。タグとキャプションを分けて管理したり、学習を高速化するためlatentsを事前キャッシュしたりなどの機能をサポートします（いずれも別文書で説明しています）。
+    あらかじめキャプションをメタデータファイルにまとめます。タグとキャプションを分けて管理したり、学習を高速化するためlatentsを事前キャッシュしたりなどの機能をサポートします（いずれも別文書で説明しています）。（fine tuning方式という名前ですが fine tuning 以外でも使えます。）
 
 学習したいものと使用できる指定方法の組み合わせは以下の通りです。
 
@@ -124,7 +125,7 @@ batch_size = 4                              # バッチサイズ
   num_repeats = 1                           # 正則化画像の繰り返し回数、基本的には1でよい
 ```
 
-基本的には以下を場所のみ書き換えれば学習できます。
+基本的には以下の場所のみ書き換えれば学習できます。
 
 1. 学習解像度
 
@@ -132,7 +133,7 @@ batch_size = 4                              # バッチサイズ
 
 1. バッチサイズ
 
-    同時に何件のデータを学習するかを指定します。GPUのVRAMサイズ、学習解像度によって変わってきます。またfine tuning/DreamBooth/LoRA等でも変わってきますので、詳しくは各スクリプトの説明をご覧ください。
+    同時に何件のデータを学習するかを指定します。GPUのVRAMサイズ、学習解像度によって変わってきます。詳しくは後述します。またfine tuning/DreamBooth/LoRA等でも変わってきますので各スクリプトの説明もご覧ください。
 
 1. フォルダ指定
 
@@ -248,7 +249,41 @@ batch_size = 4                                      # バッチサイズ
 
 それぞれのドキュメントを参考に学習を行ってください。
 
-# Aspect Ratio Bucketing について
+# 学習で使われる用語のごく簡単な解説
+
+細かいことは省略していますし私も完全には理解していないため、詳しくは各自お調べください。
+
+## ステップ
+
+ざっくりいうと学習データで1回計算すると1ステップです。「学習データのキャプションを今のモデルに流してみて、出てくる画像を学習データの画像と比較し、学習データに近づくようにモデルをわずかに変更する」のが1ステップです。
+
+## バッチサイズ
+
+バッチサイズは1ステップで何件のデータをまとめて計算するかを指定する値です。まとめて計算するため速度は相対的に向上します。また一般的には精度も高くなるといわれています。
+
+`バッチサイズ×ステップ数` が学習に使われるデータの件数になります。そのため、バッチサイズを増やした分だけステップ数を減らすとよいでしょう。
+
+（ただし、たとえば「バッチサイズ1で1600ステップ」と「バッチサイズ4で400ステップ」は同じ結果にはなりません。同じ学習率の場合、一般的には後者のほうが学習不足になります。学習率を多少大きくするか（たとえば `2e-6` など）、ステップ数をたとえば500ステップにするなどして工夫してください。）
+
+バッチサイズを大きくするとその分だけGPUメモリを消費します。メモリが足りなくなるとエラーになりますし、エラーにならないギリギリでは学習速度が低下します。タスクマネージャーや `nvidia-smi` コマンドで使用メモリ量を確認しながら調整するとよいでしょう。
+
+なお、バッチは「一塊のデータ」位の意味です。
+
+## 学習率
+
+ざっくりいうと1ステップごとにどのくらい変化させるかを表します。大きな値を指定するとそれだけ速く学習が進みますが、変化しすぎてモデルが壊れたり、最適な状態にまで至れない場合があります。小さい値を指定すると学習速度は遅くなり、また最適な状態にやはり至れない場合があります。
+
+fine tuning、DreamBoooth、LoRAそれぞれで大きく異なり、また学習データや学習させたいモデル、バッチサイズやステップ数によっても変わってきます。一般的な値から初めて学習状態を見ながら増減してください。
+
+デフォルトでは学習全体を通して学習率は固定です。スケジューラの指定で学習率をどう変化させるか決められますので、それらによっても結果は変わってきます。
+
+## エポック（epoch）
+
+学習データが一通り学習されると（データが一周すると）1 epochです。繰り返し回数を指定した場合は、その繰り返し後のデータが一周すると1 epochです。
+
+1 epochのステップ数は、基本的には `データ件数÷バッチサイズ` ですが、Aspect Ratio Bucketing を使うと微妙に増えます（異なるbucketのデータは同じバッチにできないため、ステップ数が増えます）。
+
+## Aspect Ratio Bucketing
 
 Stable Diffusion のv1は512\*512で学習されていますが、それに加えて256\*1024や384\*640といった解像度でも学習します。これによりトリミングされる部分が減り、より正しくキャプションと画像の関係が学習されることが期待されます。
 
@@ -260,11 +295,15 @@ Stable Diffusion のv1は512\*512で学習されていますが、それに加
 
 機械学習では入力サイズをすべて統一するのが一般的ですが、特に制約があるわけではなく、実際は同一のバッチ内で統一されていれば大丈夫です。NovelAIの言うbucketingは、あらかじめ教師データを、アスペクト比に応じた学習解像度ごとに分類しておくことを指しているようです。そしてバッチを各bucket内の画像で作成することで、バッチの画像サイズを統一します。
 
-# 以前のデータ指定方法
+# 以前の指定形式（設定ファイルを用いずコマンドラインから指定）
 
-フォルダ名で繰り返し回数を指定する方法です。
+`.toml` ファイルを指定せずコマンドラインオプションで指定する方法です。DreamBooth class+identifier方式、DreamBooth キャプション方式、fine tuning方式があります。
 
-## step 1. 学習用画像の準備
+## DreamBooth、class+identifier方式
+
+フォルダ名で繰り返し回数を指定します。また `train_data_dir` オプションと `reg_data_dir` オプションを用います。
+
+### step 1. 学習用画像の準備
 
 学習用画像を格納するフォルダを作成します。 __さらにその中に__ 、以下の名前でディレクトリを作成します。
 
@@ -294,15 +333,7 @@ classがひとつで対象が複数の場合、正則化画像フォルダはひ
 - reg_girls
   - 1_1girl
 
-### DreamBoothでキャプションを使う
-
-学習用画像、正則化画像のフォルダに、画像と同じファイル名で、拡張子.caption（オプションで変えられます）のファイルを置くと、そのファイルからキャプションを読み込みプロンプトとして学習します。
-
-※それらの画像の学習に、フォルダ名（identifier class）は使用されなくなります。
-
-キャプションファイルの拡張子はデフォルトで.captionです。学習スクリプトの `--caption_extension` オプションで変更できます。`--shuffle_caption` オプションで学習時のキャプションについて、カンマ区切りの各部分をシャッフルしながら学習します。
-
-## step 2. 正則化画像の準備
+### step 2. 正則化画像の準備
 
 正則化画像を使う場合の手順です。
 
@@ -313,16 +344,256 @@ classがひとつで対象が複数の場合、正則化画像フォルダはひ
 ![image](https://user-images.githubusercontent.com/52813779/210770897-329758e5-3675-49f1-b345-c135f1725832.png)
 
 
-## step 3. 学習の実行
+### step 3. 学習の実行
 
 各学習スクリプトを実行します。 `--train_data_dir` オプションで前述の学習用データのフォルダを（__画像を含むフォルダではなく、その親フォルダ__）、`--reg_data_dir` オプションで正則化画像のフォルダ（__画像を含むフォルダではなく、その親フォルダ__）を指定してください。
 
-<!-- 
-# 学習スクリプト共通のオプション
+## DreamBooth、キャプション方式
+
+学習用画像、正則化画像のフォルダに、画像と同じファイル名で、拡張子.caption（オプションで変えられます）のファイルを置くと、そのファイルからキャプションを読み込みプロンプトとして学習します。
+
+※それらの画像の学習に、フォルダ名（identifier class）は使用されなくなります。
+
+キャプションファイルの拡張子はデフォルトで.captionです。学習スクリプトの `--caption_extension` オプションで変更できます。`--shuffle_caption` オプションで学習時のキャプションについて、カンマ区切りの各部分をシャッフルしながら学習します。
+
+## fine tuning 方式
+
+メタデータを作るところまでは設定ファイルを使う場合と同様です。`in_json` オプションでメタデータファイルを指定します。
+
+# 学習途中でのサンプル出力
+
+学習中のモデルで試しに画像生成することで学習の進み方を確認できます。学習スクリプトに以下のオプションを指定します。
+
+- `--sample_every_n_steps` / `--sample_every_n_epochs`
+    
+    サンプル出力するステップ数またはエポック数を指定します。この数ごとにサンプル出力します。両方指定するとエポック数が優先されます。
+
+- `--sample_prompts`
+
+    サンプル出力用プロンプトのファイルを指定します。
+
+- `--sample_sampler`
+
+    サンプル出力に使うサンプラーを指定します。
+    `'ddim', 'pndm', 'heun', 'dpmsolver', 'dpmsolver++', 'dpmsingle', 'k_lms', 'k_euler', 'k_euler_a', 'k_dpm_2', 'k_dpm_2_a'`が選べます。
+
+サンプル出力を行うにはあらかじめプロンプトを記述したテキストファイルを用意しておく必要があります。1行につき1プロンプトで記述します。
+
+たとえば以下のようになります。
+
+```txt
+# prompt 1
+masterpiece, best quality, 1girl, in white shirts, upper body, looking at viewer, simple background --n low quality, worst quality, bad anatomy,bad composition, poor, low effort --w 768 --h 768 --d 1 --l 7.5 --s 28
+
+# prompt 2
+masterpiece, best quality, 1boy, in business suit, standing at street, looking back --n low quality, worst quality, bad anatomy,bad composition, poor, low effort --w 576 --h 832 --d 2 --l 5.5 --s 40
+```
+
+先頭が `#` の行はコメントになります。`--n` のように 「`--` + 英小文字」で生成画像へのオプションを指定できます。以下が使えます。
+
+- `--n` 次のオプションまでをネガティブプロンプトとします。
+- `--w` 生成画像の横幅を指定します。
+- `--h` 生成画像の高さを指定します。
+- `--d` 生成画像のseedを指定します。
+- `--l` 生成画像のCFG scaleを指定します。
+- `--s` 生成時のステップ数を指定します。
+
+
+# 各スクリプトで共通の、よく使われるオプション
 
 スクリプトの更新後、ドキュメントの更新が追い付いていない場合があります。その場合は `--help` オプションで使用できるオプションを確認してください。
 
-## TODO 書きます
+## 学習に使うモデル指定
+
+- `--v2` / `--v_parameterization`
+    
+    Hugging Faceのstable-diffusion-2-baseを使う場合はv2オプションを、stable-diffusion-2または768-v-ema.ckptを使う場合はv2とv_parameterizationの両方のオプションを指定してください。
+
+    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オプションです。
+
+- `--pretrained_model_name_or_path` 
+    
+    追加学習を行う元となるモデルを指定します。Stable Diffusionのcheckpointファイル（.ckptまたは.safetensors）、Diffusersのローカルディスクにあるモデルディレクトリ、DiffusersのモデルID（"stabilityai/stable-diffusion-2"など）が指定できます。
+
+## 学習に関する設定
+
+- `--output_dir` 
+
+    学習後のモデルを保存するフォルダを指定します。
+    
+- `--output_name` 
+    
+    モデルのファイル名を拡張子を除いて指定します。
+    
+- `--dataset_config` 
+
+    データセットの設定を記述した `.toml` ファイルを指定します。
+
+- `--max_train_steps` / `--max_train_epochs`
+
+    学習するステップ数やエポック数を指定します。両方指定するとエポック数のほうが優先されます。
+
+- `--mixed_precision`
+
+    省メモリ化のため mixed precision （混合精度）で学習します。`--mixed_precision="fp16"` のように指定します。mixed precision なし（デフォルト）と比べて精度が低くなる可能性がありますが、学習に必要なGPUメモリ量が大きく減ります。
+    
+    （RTX30 シリーズ以降では `bf16` も指定できます。環境整備時にaccelerateに行った設定と合わせてください）。
+    
+- `--gradient_checkpointing`
+
+    学習時の重みの計算をまとめて行うのではなく少しずつ行うことで、学習に必要なGPUメモリ量を減らします。オンオフは精度には影響しませんが、オンにするとバッチサイズを大きくできるため、そちらでの影響はあります。
+    
+    また一般的にはオンにすると速度は低下しますが、バッチサイズを大きくできるので、トータルでの学習時間はむしろ速くなるかもしれません。
+
+- `--xformers` / `--mem_eff_attn`
+
+    xformersオプションを指定するとxformersのCrossAttentionを用います。xformersをインストールしていない場合やエラーとなる場合（環境にもよりますが `mixed_precision="no"` の場合など）、代わりに `mem_eff_attn` オプションを指定すると省メモリ版CrossAttentionを使用します（xformersよりも速度は遅くなります）。
+
+- `--save_precision`
+
+    保存時のデータ精度を指定します。save_precisionオプションにfloat、fp16、bf16のいずれかを指定すると、その形式でモデルを保存します（DreamBooth、fine tuningでDiffusers形式でモデルを保存する場合は無効です）。モデルのサイズを削減したい場合などにお使いください。
+
+- `--save_every_n_epochs` / `--save_state` / `--resume`
+    save_every_n_epochsオプションに数値を指定すると、そのエポックごとに学習途中のモデルを保存します。
+
+    save_stateオプションを同時に指定すると、optimizer等の状態も含めた学習状態を合わせて保存します（保存したモデルからも学習再開できますが、それに比べると精度の向上、学習時間の短縮が期待できます）。保存先はフォルダになります。
+    
+    学習状態は保存先フォルダに `<output_name>-??????-state`（??????はエポック数）という名前のフォルダで出力されます。長時間にわたる学習時にご利用ください。
+
+    保存された学習状態から学習を再開するにはresumeオプションを使います。学習状態のフォルダ（`output_dir` ではなくその中のstateのフォルダ）を指定してください。
+
+    なおAcceleratorの仕様により、エポック数、global stepは保存されておらず、resumeしたときにも1からになりますがご容赦ください。
+
+- `--save_model_as` （DreamBooth, fine tuning のみ）
+
+    モデルの保存形式を`ckpt, safetensors, diffusers, diffusers_safetensors` から選べます。
+    
+    `--save_model_as=safetensors` のように指定します。Stable Diffusion形式（ckptまたはsafetensors）を読み込み、Diffusers形式で保存する場合、不足する情報はHugging Faceからv1.5またはv2.1の情報を落としてきて補完します。
+    
+- `--clip_skip`
+    
+    `2` を指定すると、Text Encoder (CLIP) の後ろから二番目の層の出力を用います。1またはオプション省略時は最後の層を用います。
+
+    ※SD2.0はデフォルトで後ろから二番目の層を使うため、SD2.0の学習では指定しないでください。
+
+    学習対象のモデルがもともと二番目の層を使うように学習されている場合は、2を指定するとよいでしょう。
+
+    そうではなく最後の層を使用していた場合はモデル全体がそれを前提に学習されています。そのため改めて二番目の層を使用して学習すると、望ましい学習結果を得るにはある程度の枚数の教師データ、長めの学習が必要になるかもしれません。
+
+- `--max_token_length`
+
+    デフォルトは75です。`150` または `225` を指定することでトークン長を拡張して学習できます。長いキャプションで学習する場合に指定してください。
+    
+    ただし学習時のトークン拡張の仕様は Automatic1111 氏のWeb UIとは微妙に異なるため（分割の仕様など）、必要なければ75で学習することをお勧めします。
+
+    clip_skipと同様に、モデルの学習状態と異なる長さで学習するには、ある程度の教師データ枚数、長めの学習時間が必要になると思われます。
+
+- `--logging_dir` / `--log_prefix`
+
+    学習ログの保存に関するオプションです。logging_dirオプションにログ保存先フォルダを指定してください。TensorBoard形式のログが保存されます。
+
+    たとえば--logging_dir=logsと指定すると、作業フォルダにlogsフォルダが作成され、その中の日時フォルダにログが保存されます。
+    また--log_prefixオプションを指定すると、日時の前に指定した文字列が追加されます。「--logging_dir=logs --log_prefix=db_style1_」などとして識別用にお使いください。
+
+    TensorBoardでログを確認するには、別のコマンドプロンプトを開き、作業フォルダで以下のように入力します。
+
+    ```
+    tensorboard --logdir=logs
+    ```
+
+    （tensorboardは環境整備時にあわせてインストールされると思いますが、もし入っていないなら `pip install tensorboard` で入れてください。）
+
+    その後ブラウザを開き、http://localhost:6006/ へアクセスすると表示されます。
+
+- `--noise_offset`
+
+    こちらの記事の実装になります: https://www.crosslabs.org//blog/diffusion-with-offset-noise
+    
+    全体的に暗い、明るい画像の生成結果が良くなる可能性があるようです。LoRA学習でも有効なようです。`0.1` 程度の値を指定するとよいようです。
+
+- `--debug_dataset`
+
+    このオプションを付けることで学習を行う前に事前にどのような画像データ、キャプションで学習されるかを確認できます。Escキーを押すと終了してコマンドラインに戻ります。
+
+    ※Linux環境（Colabを含む）では画像は表示されません。
+
+- `--vae`
+
+    vaeオプションにStable Diffusionのcheckpoint、VAEのcheckpointファイル、DiffusesのモデルまたはVAE（ともにローカルまたはHugging FaceのモデルIDが指定できます）のいずれかを指定すると、そのVAEを使って学習します（latentsのキャッシュ時または学習中のlatents取得時）。
+
+    DreamBoothおよびfine tuningでは、保存されるモデルはこのVAEを組み込んだものになります。
+
+
+## オプティマイザ関係
+
+- `--optimizer_type`
+
+    オプティマイザの種類を次から指定します: AdamW (default), AdamW8bit, Lion, SGDNesterov, SGDNesterov8bit, DAdaptation, AdaFactor
+
+- `--learning_rate`
+
+    学習率を指定します。適切な学習率は学習スクリプトにより異なりますので、それぞれの説明を参照してください。
+
+- `--lr_scheduler` / `--lr_warmup_steps` / `--lr_scheduler_num_cycles` / `--lr_scheduler_power`
+  
+    学習率のスケジューラ関連の指定です。
+
+    lr_schedulerオプションで学習率のスケジューラをlinear, cosine, cosine_with_restarts, polynomial, constant, constant_with_warmupから選べます。デフォルトはconstantです。
+    
+    lr_warmup_stepsでスケジューラのウォームアップ（だんだん学習率を変えていく）ステップ数を指定できます。
+    
+    lr_scheduler_num_cycles は cosine with restartsスケジューラでのリスタート回数、lr_scheduler_power は polynomialスケジューラでのpolynomial power です。
+
+    詳細については各自お調べください。
+
+
+<!-- 
+## 任意サイズの画像での学習 --resolution
+正方形以外で学習できます。resolutionに「448,640」のように「幅,高さ」で指定してください。幅と高さは64で割り切れる必要があります。学習用画像、正則化画像のサイズを合わせてください。
+
+個人的には縦長の画像を生成することが多いため「448,640」などで学習することもあります。
+
+## Aspect Ratio Bucketing --enable_bucket / --min_bucket_reso / --max_bucket_reso
+enable_bucketオプションを指定すると有効になります。Stable Diffusionは512x512で学習されていますが、それに加えて256x768や384x640といった解像度でも学習します。
+
+このオプションを指定した場合は、学習用画像、正則化画像を特定の解像度に統一する必要はありません。いくつかの解像度（アスペクト比）から最適なものを選び、その解像度で学習します。
+解像度は64ピクセル単位のため、元画像とアスペクト比が完全に一致しない場合がありますが、その場合は、はみ出した部分がわずかにトリミングされます。
+
+解像度の最小サイズをmin_bucket_resoオプションで、最大サイズをmax_bucket_resoで指定できます。デフォルトはそれぞれ256、1024です。
+たとえば最小サイズに384を指定すると、256x1024や320x768などの解像度は使わなくなります。
+解像度を768x768のように大きくした場合、最大サイズに1280などを指定しても良いかもしれません。
+
+なおAspect Ratio Bucketingを有効にするときには、正則化画像についても、学習用画像と似た傾向の様々な解像度を用意した方がいいかもしれません。
+
+（ひとつのバッチ内の画像が学習用画像、正則化画像に偏らなくなるため。そこまで大きな影響はないと思いますが……。）
+
+## augmentation --color_aug / --flip_aug
+augmentationは学習時に動的にデータを変化させることで、モデルの性能を上げる手法です。color_augで色合いを微妙に変えつつ、flip_augで左右反転をしつつ、学習します。
+
+動的にデータを変化させるため、cache_latentsオプションと同時に指定できません。
+
+
+## 勾配をfp16とした学習（実験的機能） --full_fp16
+full_fp16オプションを指定すると勾配を通常のfloat32からfloat16（fp16）に変更して学習します（mixed precisionではなく完全なfp16学習になるようです）。
+これによりSD1.xの512x512サイズでは8GB未満、SD2.xの512x512サイズで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で確認）。精度はかなり落ちますし、途中で学習失敗する確率も高くなります。
+学習率やステップ数の設定もシビアなようです。それらを認識したうえで自己責任でお使いください。
+
 -->
 
 # メタデータファイルの作成
diff --git a/train_db_README-ja.md b/train_db_README-ja.md
index 85ae35aa..903f214c 100644
--- a/train_db_README-ja.md
+++ b/train_db_README-ja.md
@@ -1,75 +1,92 @@
-DreamBoothのガイドです。LoRA等の追加ネットワークの学習にも同じ手順を使います。
+DreamBoothのガイドです。
 
 # 概要
 
+DreamBoothとは、画像生成モデルに特定の主題を追加学習し、それを特定の識別子で生成する技術です。[論文はこちら](https://arxiv.org/abs/2208.12242)。
+
+具体的には、Stable Diffusionのモデルにキャラや画風などを学ばせ、それを `shs` のような特定の単語で呼び出せる（生成画像に出現させる）ことができます。
+
+スクリプトは[DiffusersのDreamBooth](https://github.com/huggingface/diffusers/tree/main/examples/dreambooth)を元にしていますが、以下のような機能追加を行っています（いくつかの機能は元のスクリプト側もその後対応しています）。
+
 スクリプトの主な機能は以下の通りです。
 
-- 8bit Adam optimizerおよびlatentのキャッシュによる省メモリ化（ShivamShrirao氏版と同様）。
+- 8bit Adam optimizerおよびlatentのキャッシュによる省メモリ化（[Shivam Shrirao氏版](https://github.com/ShivamShrirao/diffusers/tree/main/examples/dreambooth)と同様）。
 - xformersによる省メモリ化。
 - 512x512だけではなく任意サイズでの学習。
 - augmentationによる品質の向上。
 - DreamBoothだけではなくText Encoder+U-Netのfine tuningに対応。
-- StableDiffusion形式でのモデルの読み書き。
+- Stable Diffusion形式でのモデルの読み書き。
 - Aspect Ratio Bucketing。
 - Stable Diffusion v2.0対応。
 
 # 学習の手順
 
-## step 1. 環境整備
+## データの準備
 
-このリポジトリのREADMEを参照してください。
+[学習データの準備について](./train_README-ja.md) を参照してください。
 
+## 学習の実行
 
-## step 2. identifierとclassを決める
-
-学ばせたい対象を結びつける単語identifierと、対象の属するclassを決めます。
-
-（instanceなどいろいろな呼び方がありますが、とりあえず元の論文に合わせます。）
-
-以下ごく簡単に説明します（詳しくは調べてください）。
-
-classは学習対象の一般的な種別です。たとえば特定の犬種を学ばせる場合には、classはdogになります。アニメキャラならモデルによりboyやgirl、1boyや1girlになるでしょう。
-
-identifierは学習対象を識別して学習するためのものです。任意の単語で構いませんが、元論文によると「tokinizerで1トークンになる3文字以下でレアな単語」が良いとのことです。
-
-identifierとclassを使い、たとえば「shs dog」などでモデルを学習することで、学習させたい対象をclassから識別して学習できます。
-
-画像生成時には「shs dog」とすれば学ばせた犬種の画像が生成されます。
-
-（identifierとして私が最近使っているものを参考までに挙げると、``shs sts scs cpc coc cic msm usu ici lvl cic dii muk ori hru rik koo yos wny`` などです。）
-
-## step 3. 学習用画像の準備
-学習用画像を格納するフォルダを作成します。 __さらにその中に__ 、以下の名前でディレクトリを作成します。
+スクリプトを実行します。最大限、メモリを節約したコマンドは以下のようになります（実際には1行で入力します）。`< >` 内を書き換えてください。12GB程度のVRAMで動作するようです。
 
 ```
-<繰り返し回数>_<identifier> <class>
+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"など）が指定できます。
 
-たとえば「sls frog」というプロンプトで、データを20回繰り返す場合、「20_sls frog」となります。以下のようになります。
+`output_dir` に学習後のモデルを保存するフォルダを指定します。`output_name` にモデルのファイル名を拡張子を除いて指定します。`save_model_as` でsafetensors形式での保存を指定しています。
 
-![image](https://user-images.githubusercontent.com/52813779/210770636-1c851377-5936-4c15-90b7-8ac8ad6c2074.png)
+`dataset_config` に `.toml` ファイルを指定します。ファイル内でのバッチサイズ指定は、メモリ消費を抑えるために `1` としてください。
 
-## step 4. 正則化画像の準備
-正則化画像を使う場合の手順です。使わずに学習することもできます（正則化画像を使わないと区別ができなくなるので対象class全体が影響を受けます）。
+`prior_loss_weight` は正則化画像のlossの重みです。通常は1.0を指定します。
 
-正則化画像を格納するフォルダを作成します。 __さらにその中に__  ``<繰り返し回数>_<class>`` という名前でディレクトリを作成します。
+学習させるステップ数 `max_train_steps` を1600とします。学習率 `learning_rate` はここでは1e-6を指定しています。
 
-たとえば「frog」というプロンプトで、データを繰り返さない（1回だけ）場合、以下のようになります。
+省メモリ化のため `mixed_precision="fp16"` を指定します（RTX30 シリーズ以降では `bf16` も指定できます。環境整備時にaccelerateに行った設定と合わせてください）。また `gradient_checkpointing` を指定します。
 
-![image](https://user-images.githubusercontent.com/52813779/210770897-329758e5-3675-49f1-b345-c135f1725832.png)
+オプティマイザ（モデルを学習データにあうように最適化＝学習させるクラス）にメモリ消費の少ない 8bit AdamW を使うため、 `optimizer_type="AdamW8bit"` を指定します。
 
-繰り返し回数は「 __学習用画像の繰り返し回数×学習用画像の枚数≧正則化画像の繰り返し回数×正則化画像の枚数__ 」となるように指定してください。
+xformersオプションを指定し、xformersのCrossAttentionを用います。xformersをインストールしていない場合やエラーとなる場合（環境にもよりますが `mixed_precision="no"` の場合など）、代わりに `mem_eff_attn` オプションを指定すると省メモリ版CrossAttentionを使用します（速度は遅くなります）。
 
-（1 epochのデータ数が「学習用画像の繰り返し回数×学習用画像の枚数」となります。正則化画像の枚数がそれより多いと、余った部分の正則化画像は使用されません。）
+省メモリ化のため `cache_latents` オプションを指定してVAEの出力をキャッシュします。
 
-## step 5. 学習の実行
-スクリプトを実行します。最大限、メモリを節約したコマンドは以下のようになります（実際には1行で入力します）。
+ある程度メモリがある場合は、`.toml` ファイルを編集してバッチサイズをたとえば `4` くらいに増やしてください（高速化と精度向上の可能性があります）。また `cache_latents` を外すことで augmentation が可能になります。
 
-※LoRA等の追加ネットワークを学習する場合のコマンドは ``train_db.py`` ではなく ``train_network.py`` となります。また追加でnetwork_\*オプションが必要となりますので、LoRAのガイドを参照してください。
+### 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 
@@ -77,6 +94,7 @@ accelerate launch --num_cpu_threads_per_process 1 train_db.py
     --train_data_dir=<学習用データのディレクトリ> 
     --reg_data_dir=<正則化画像のディレクトリ> 
     --output_dir=<学習したモデルの出力先ディレクトリ> 
+    --output_name=<学習したモデル出力時のファイル名> 
     --prior_loss_weight=1.0 
     --resolution=512 
     --train_batch_size=1 
@@ -89,43 +107,33 @@ accelerate launch --num_cpu_threads_per_process 1 train_db.py
     --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"など）が指定できます。学習後のモデルの保存形式はデフォルトでは元のモデルと同じになります（save_model_asオプションで変更できます）。
+学習が終わると指定したフォルダに指定した名前でsafetensorsファイルが出力されます。
 
-prior_loss_weightは正則化画像のlossの重みです。通常は1.0を指定します。
+v1.4/1.5およびその他の派生モデルの場合、このモデルでAutomatic1111氏のWebUIなどで推論できます。models\Stable-diffusionフォルダに置いてください。
 
-resolutionは画像のサイズ（解像度、幅と高さ）になります。bucketing（後述）を用いない場合、学習用画像、正則化画像はこのサイズとしてください。
+v2.xモデルでWebUIで画像生成する場合、モデルの仕様が記述された.yamlファイルが別途必要になります。v2.x baseの場合はv2-inference.yamlを、768/vの場合はv2-inference-v.yamlを、同じフォルダに置き、拡張子の前の部分をモデルと同じ名前にしてください。
 
-train_batch_sizeは学習時のバッチサイズです。max_train_stepsを1600とします。学習率learning_rateは、diffusers版では5e-6ですがStableDiffusion版は1e-6ですのでここでは1e-6を指定しています。
+![image](https://user-images.githubusercontent.com/52813779/210776915-061d79c3-6582-42c2-8884-8b91d2f07313.png)
 
-省メモリ化のためmixed_precision="bf16"（または"fp16"）、およびgradient_checkpointing を指定します。
+各yamlファイルは[Stability AIのSD2.0のリポジトリ](https://github.com/Stability-AI/stablediffusion/tree/main/configs/stable-diffusion)にあります。
 
-xformersオプションを指定し、xformersのCrossAttentionを用います。xformersをインストールしていない場合、エラーとなる場合（mixed_precisionなしの場合、私の環境ではエラーとなりました）、代わりにmem_eff_attnオプションを指定すると省メモリ版CrossAttentionを使用します（速度は遅くなります）。
+# DreamBooth特有のその他の主なオプション
 
-省メモリ化のためcache_latentsオプションを指定してVAEの出力をキャッシュします。
+すべてのオプションについては別文書を参照してください。
 
-ある程度メモリがある場合はたとえば以下のように指定します。
+## Text Encoderの学習を途中から行わない --stop_text_encoder_training
 
-```
-accelerate launch --num_cpu_threads_per_process 8 train_db.py 
-    --pretrained_model_name_or_path=<.ckptまたは.safetensordまたは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 
-    --mixed_precision="bf16" 
-    --cache_latents
-```
+stop_text_encoder_trainingオプションに数値を指定すると、そのステップ数以降はText Encoderの学習を行わずU-Netだけ学習します。場合によっては精度の向上が期待できるかもしれません。
 
-gradient_checkpointingを外し高速化します（メモリ使用量は増えます）。バッチサイズを増やし、高速化と精度向上を図ります。
+（恐らくText Encoderだけ先に過学習することがあり、それを防げるのではないかと推測していますが、詳細な影響は不明です。）
 
+## Tokenizerのパディングをしない --no_token_padding
+no_token_paddingオプションを指定するとTokenizerの出力をpaddingしません（Diffusers版の旧DreamBoothと同じ動きになります）。
+
+
+<!-- 
 bucketing（後述）を利用しかつaugmentation（後述）を使う場合の例は以下のようになります。
 
 ```
@@ -143,154 +151,5 @@ accelerate launch --num_cpu_threads_per_process 8 train_db.py
     --color_aug --flip_aug --gradient_checkpointing --seed 42
 ```
 
-### ステップ数について
-省メモリ化のため、ステップ当たりの学習回数がtrain_dreambooth.pyの半分になっています（対象の画像と正則化画像を同一のバッチではなく別のバッチに分割して学習するため）。
-元のDiffusers版やXavierXiao氏のStableDiffusion版とほぼ同じ学習を行うには、ステップ数を倍にしてください。
-
-（shuffle=Trueのため厳密にはデータの順番が変わってしまいますが、学習には大きな影響はないと思います。）
-
-## 学習したモデルで画像生成する
-
-学習が終わると指定したフォルダにlast.ckptという名前でcheckpointが出力されます（DiffUsers版モデルを学習した場合はlastフォルダになります）。
-
-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を、同じフォルダに置き、拡張子の前の部分をモデルと同じ名前にしてください。
-
-![image](https://user-images.githubusercontent.com/52813779/210776915-061d79c3-6582-42c2-8884-8b91d2f07313.png)
-
-各yamlファイルは[Stability AIのSD2.0のリポジトリ](https://github.com/Stability-AI/stablediffusion/tree/main/configs/stable-diffusion)にあります。
-
-# その他の学習オプション
-
-## Stable Diffusion 2.0対応 --v2 / --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オプションです。
-
-## 学習データの確認 --debug_dataset
-このオプションを付けることで学習を行う前に事前にどのような画像データ、キャプションで学習されるかを確認できます。Escキーを押すと終了してコマンドラインに戻ります。
-
-※Colabなど画面が存在しない環境で実行するとハングするようですのでご注意ください。
-
-## 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は512x512で学習されていますが、それに加えて256x768や384x640といった解像度でも学習します。
-
-このオプションを指定した場合は、学習用画像、正則化画像を特定の解像度に統一する必要はありません。いくつかの解像度（アスペクト比）から最適なものを選び、その解像度で学習します。
-解像度は64ピクセル単位のため、元画像とアスペクト比が完全に一致しない場合がありますが、その場合は、はみ出した部分がわずかにトリミングされます。
-
-解像度の最小サイズをmin_bucket_resoオプションで、最大サイズをmax_bucket_resoで指定できます。デフォルトはそれぞれ256、1024です。
-たとえば最小サイズに384を指定すると、256x1024や320x768などの解像度は使わなくなります。
-解像度を768x768のように大きくした場合、最大サイズに1280などを指定しても良いかもしれません。
-
-なおAspect Ratio Bucketingを有効にするときには、正則化画像についても、学習用画像と似た傾向の様々な解像度を用意した方がいいかもしれません。
-
-（ひとつのバッチ内の画像が学習用画像、正則化画像に偏らなくなるため。そこまで大きな影響はないと思いますが……。）
-
-## augmentation --color_aug / --flip_aug
-augmentationは学習時に動的にデータを変化させることで、モデルの性能を上げる手法です。color_augで色合いを微妙に変えつつ、flip_augで左右反転をしつつ、学習します。
-
-動的にデータを変化させるため、cache_latentsオプションと同時に指定できません。
-
-## 保存時のデータ精度の指定 --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の情報を落としてきて補完します。
-
-## 学習ログの保存 --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の512x512サイズでは8GB未満、SD2.xの512x512サイズで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で確認）。精度はかなり落ちますし、途中で学習失敗する確率も高くなります。
-学習率やステップ数の設定もシビアなようです。それらを認識したうえで自己責任でお使いください。
-
-# その他の学習方法
-
-## 複数class、複数対象（identifier）の学習
-方法は単純で、学習用画像のフォルダ内に ``繰り返し回数_<identifier> <class>`` のフォルダを複数、正則化画像フォルダにも同様に ``繰り返し回数_<class>`` のフォルダを複数、用意してください。
-
-たとえば「sls frog」と「cpc rabbit」を同時に学習する場合、以下のようになります。
-
-![image](https://user-images.githubusercontent.com/52813779/210777933-a22229db-b219-4cd8-83ca-e87320fc4192.png)
-
-classがひとつで対象が複数の場合、正則化画像フォルダはひとつで構いません。たとえば1girlにキャラAとキャラBがいる場合は次のようにします。
-
-- train_girls
-  - 10_sls 1girl
-  - 10_cpc 1girl
-- reg_girls
-  - 1_1girl
-
-データ数にばらつきがある場合、繰り返し回数を調整してclass、identifierごとの枚数を統一すると良い結果が得られることがあるようです。
-
-## DreamBoothでキャプションを使う
-学習用画像、正則化画像のフォルダに、画像と同じファイル名で、拡張子.caption（オプションで変えられます）のファイルを置くと、そのファイルからキャプションを読み込みプロンプトとして学習します。
-
-※それらの画像の学習に、フォルダ名（identifier class）は使用されなくなります。
-
-各画像にキャプションを付けることで（BLIP等を使っても良いでしょう）、学習したい属性をより明確にできるかもしれません。
-
-キャプションファイルの拡張子はデフォルトで.captionです。--caption_extensionで変更できます。--shuffle_captionオプションで学習時のキャプションについて、カンマ区切りの各部分をシャッフルしながら学習します。
 
+-->