Merge pull request #278 from kohya-ss/dev

Dev
restore comment
2026-04-06 21:52:27 +00:00 · 2023-03-10 22:05:36 +09:00 · 2023-03-10 22:02:17 +09:00 · 2023-03-10 21:59:25 +09:00 · 2023-03-10 21:43:32 +09:00 · 2023-03-10 21:31:59 +09:00
18 changed files with 1312 additions and 967 deletions
--- a/README-ja.md
+++ b/README-ja.md
@@ -16,9 +16,10 @@ GUIやPowerShellスクリプトなど、より使いやすくする機能が[bma

 当リポジトリ内およびnote.comに記事がありますのでそちらをご覧ください（将来的にはすべてこちらへ移すかもしれません）。

+* [学習について、共通編](./train_README-ja.md) : データ整備やオプションなど
+    * [データセット設定](./config_README-ja.md)
 * [DreamBoothの学習について](./train_db_README-ja.md)
 * [fine-tuningのガイド](./fine_tune_README_ja.md):
-BLIPによるキャプショニングと、DeepDanbooruまたはWD14 taggerによるタグ付けを含みます
 * [LoRAの学習について](./train_network_README-ja.md)
 * [Textual Inversionの学習について](./train_ti_README-ja.md)
 * note.com [画像生成スクリプト](https://note.com/kohya_ss/n/n2693183a798e)
@@ -131,6 +132,8 @@ pip install --use-pep517 --upgrade -r requirements.txt

 LoRAの実装は[cloneofsimo氏のリポジトリ](https://github.com/cloneofsimo/lora)を基にしたものです。感謝申し上げます。

+Conv2d 3x3への拡大は [cloneofsimo氏](https://github.com/cloneofsimo/lora) が最初にリリースし、KohakuBlueleaf氏が [LoCon](https://github.com/KohakuBlueleaf/LoCon) でその有効性を明らかにしたものです。KohakuBlueleaf氏に深く感謝します。
+
 ## ライセンス

 スクリプトのライセンスはASL 2.0ですが（Diffusersおよびcloneofsimo氏のリポジトリ由来のものも同様）、一部他のライセンスのコードを含みます。
--- a/README.md
+++ b/README.md
@@ -28,9 +28,10 @@ The scripts are tested with PyTorch 1.12.1 and 1.13.0, Diffusers 0.10.2.

 All documents are in Japanese currently.

+* [Training guide - common](./train_README-ja.md) : data preparation, options etc...
+    * [Dataset config](./config_README-ja.md)
 * [DreamBooth training guide](./train_db_README-ja.md)
 * [Step by Step fine-tuning guide](./fine_tune_README_ja.md):
-Including BLIP captioning and tagging by DeepDanbooru or WD14 tagger
 * [training LoRA](./train_network_README-ja.md)
 * [training Textual Inversion](./train_ti_README-ja.md)
 * note.com [Image generation](https://note.com/kohya_ss/n/n2693183a798e)
@@ -110,11 +111,13 @@ Once the commands have completed successfully you should be ready to use the new

 ## Credits

-The implementation for LoRA is based on [cloneofsimo's repo](https://github.com/cloneofsimo/lora). Thank you for great work!!!
+The implementation for LoRA is based on [cloneofsimo's repo](https://github.com/cloneofsimo/lora). Thank you for great work!
+
+The LoRA expansion to Conv2d 3x3 was initially released by cloneofsimo and its effectiveness was demonstrated at [LoCon](https://github.com/KohakuBlueleaf/LoCon) by KohakuBlueleaf. Thank you so much KohakuBlueleaf!

 ## License

-The majority of scripts is licensed under ASL 2.0 (including codes from Diffusers, cloneofsimo's), however portions of the project are available under separate license terms:
+The majority of scripts is licensed under ASL 2.0 (including codes from Diffusers, cloneofsimo's and LoCon), however portions of the project are available under separate license terms:

 [Memory Efficient Attention Pytorch](https://github.com/lucidrains/memory-efficient-attention-pytorch): MIT

@@ -124,22 +127,57 @@ The majority of scripts is licensed under ASL 2.0 (including codes from Diffuser

 ## Change History

- 2 Mar. 2023, 2023/3/2:
+- 10 Mar. 2023, 2023/3/10: release v0.5.1
+  - Fix to LoRA modules in the model are same to the previous (before 0.5.0) if Conv2d-3x3 is disabled (no `conv_dim` arg, default).
+    - Conv2D with kernel size 1x1 in ResNet modules were accidentally included in v0.5.0.
+    - Trained models with v0.5.0 will work with Web UI's built-in LoRA and Additional Networks extension.
+  - Fix an issue that dim (rank) of LoRA module is limited to the in/out dimensions of the target Linear/Conv2d (in case of the dim > 320).
+  - `resize_lora.py` now have a feature to `dynamic resizing` which means each LoRA module can have different ranks (dims). Thanks to mgz-dev for this great work!
+    - The appropriate rank is selected based on the complexity of each module with an algorithm specified in the command line arguments. For details: https://github.com/kohya-ss/sd-scripts/pull/243
+  - Multiple GPUs training is finally supported in `train_network.py`. Thanks to ddPn08 to solve this long running issue!
+  - Dataset with fine-tuning method (with metadata json) now works without images if `.npz` files exist. Thanks to rvhfxb!
+  - `train_network.py` can work if the current directory is not the directory where the script is in. Thanks to mio2333!
+  - Fix `extract_lora_from_models.py` and `svd_merge_lora.py` doesn't work with higher rank (>320).
+
+  - LoRAのConv2d-3x3拡張を行わない場合（`conv_dim` を指定しない場合）、以前（v0.5.0）と同じ構成になるよう修正しました。
+    - ResNetのカーネルサイズ1x1のConv2dが誤って対象になっていました。
+    - ただv0.5.0で学習したモデルは Additional Networks 拡張、およびWeb UIのLoRA機能で問題なく使えると思われます。
+  - LoRAモジュールの dim (rank) が、対象モジュールの次元数以下に制限される不具合を修正しました（320より大きい dim を指定した場合）。
+  - `resize_lora.py` に `dynamic resizing` （リサイズ後の各LoRAモジュールが異なるrank (dim) を持てる機能）を追加しました。mgz-dev 氏の貢献に感謝します。
+    - 適切なランクがコマンドライン引数で指定したアルゴリズムにより自動的に選択されます。詳細はこちらをご覧ください: https://github.com/kohya-ss/sd-scripts/pull/243
+  - `train_network.py` でマルチGPU学習をサポートしました。長年の懸案を解決された ddPn08 氏に感謝します。
+  - fine-tuning方式のデータセット（メタデータ.jsonファイルを使うデータセット）で `.npz` が存在するときには画像がなくても動作するようになりました。rvhfxb 氏に感謝します。
+  - 他のディレクトリから `train_network.py` を呼び出しても動作するよう変更しました。 mio2333 氏に感謝します。
+  - `extract_lora_from_models.py` および `svd_merge_lora.py` が320より大きいrankを指定すると動かない不具合を修正しました。
+  
+- 9 Mar. 2023, 2023/3/9: release v0.5.0
  - There may be problems due to major changes. If you cannot revert back to the previous version when problems occur, please do not update for a while.
-  - Dependencies are updated, Please [upgrade](#upgrade) the repo.
-  - Add detail dataset config feature by extra config file. Thanks to fur0ut0 for this great contribution!
-    - Documentation is [here](./config_README-ja.md) (only in Japanese currently.)
-    - Specify ``.toml`` file with ``--dataset_config`` option.
-    - The previous options for dataset can be used as is.
-    - There might be a bug due to the large scale of update, please report any problems if you find.
-  - Add feature to generate sample images in the middle of training for each training scripts.
-    - ``--sample_every_n_steps`` and ``--sample_every_n_epochs`` options: frequency to generate.
-    - ``--sample_prompts`` option: the file contains prompts (each line generates one image.)
-      - The prompt is subset of ``gen_img_diffusers.py``. The prompt options ``w, h, d, l, s, n`` are supported.
-    - ``--sample_sampler`` option: sampler (scheduler) for generating, such as ddim or k_euler. See help for useable samplers.
-  - Add ``--tokenizer_cache_dir`` to each training and generation scripts to cache Tokenizer locally from Diffusers.
-    - Scripts will support offline training/generation after caching.
-  - Support letents upscaling for highres. fix, and VAE batch size in ``gen_img_diffusers.py`` (no documentation yet.)
+  - Minimum metadata (module name, dim, alpha and network_args) is recorded even with `--no_metadata`, issue https://github.com/kohya-ss/sd-scripts/issues/254
+  - `train_network.py` supports LoRA for Conv2d-3x3 (extended to conv2d with a kernel size not 1x1).
+    - Same as a current version of [LoCon](https://github.com/KohakuBlueleaf/LoCon). __Thank you very much KohakuBlueleaf for your help!__
+      - LoCon will be enhanced in the future. Compatibility for future versions is not guaranteed.
+    - Specify `--network_args` option like: `--network_args "conv_dim=4" "conv_alpha=1"`
+    - [Additional Networks extension](https://github.com/kohya-ss/sd-webui-additional-networks) version 0.5.0 or later is required to use 'LoRA for Conv2d-3x3' in Stable Diffusion web UI.
+    - __Stable Diffusion web UI built-in LoRA does not support 'LoRA for Conv2d-3x3' now. Consider carefully whether or not to use it.__
+  - Merging/extracting scripts also support LoRA for Conv2d-3x3.
+  - Free CUDA memory after sample generation to reduce VRAM usage, issue https://github.com/kohya-ss/sd-scripts/issues/260 
+  - Empty caption doesn't cause error now, issue https://github.com/kohya-ss/sd-scripts/issues/258
+  - Fix sample generation is crashing in Textual Inversion training when using templates, or if height/width is not divisible by 8.
+  - Update documents (Japanese only).
+
+  - 大きく変更したため不具合があるかもしれません。問題が起きた時にスクリプトを前のバージョンに戻せない場合は、しばらく更新を控えてください。
+  - 最低限のメタデータ（module name, dim, alpha および network_args）が `--no_metadata` オプション指定時にも記録されます。issue https://github.com/kohya-ss/sd-scripts/issues/254
+  - `train_network.py` で LoRAの Conv2d-3x3 拡張に対応しました（カーネルサイズ1x1以外のConv2dにも対象範囲を拡大します）。
+    - 現在のバージョンの [LoCon](https://github.com/KohakuBlueleaf/LoCon) と同一の仕様です。__KohakuBlueleaf氏のご支援に深く感謝します。__
+      - LoCon が将来的に拡張された場合、それらのバージョンでの互換性は保証できません。
+    - `--network_args` オプションを `--network_args "conv_dim=4" "conv_alpha=1"` のように指定してください。
+    - Stable Diffusion web UI での使用には [Additional Networks extension](https://github.com/kohya-ss/sd-webui-additional-networks) のversion 0.5.0 以降が必要です。
+    - __Stable Diffusion web UI の LoRA 機能は LoRAの Conv2d-3x3 拡張に対応していないようです。使用するか否か慎重にご検討ください。__
+  - マージ、抽出のスクリプトについても LoRA の Conv2d-3x3 拡張に対応しました.
+  - サンプル画像生成後にCUDAメモリを解放しVRAM使用量を削減しました。 issue https://github.com/kohya-ss/sd-scripts/issues/260 
+  - 空のキャプションが使えるようになりました。 issue https://github.com/kohya-ss/sd-scripts/issues/258
+  - Textual Inversion 学習でテンプレートを使ったとき、height/width が 8 で割り切れなかったときにサンプル画像生成がクラッシュするのを修正しました。
+  - ドキュメント類を更新しました。

  - Sample image generation:
    A prompt file might look like this, for example
@@ -163,22 +201,6 @@ The majority of scripts is licensed under ASL 2.0 (including codes from Diffuser

    The prompt weighting such as `( )` and `[ ]` are not working.

-  - 大きく変更したため不具合があるかもしれません。問題が起きた時にスクリプトを前のバージョンに戻せない場合は、しばらく更新を控えてください。
-  - ライブラリを更新しました。[アップグレード](https://github.com/kohya-ss/sd-scripts/blob/main/README-ja.md#%E3%82%A2%E3%83%83%E3%83%97%E3%82%B0%E3%83%AC%E3%83%BC%E3%83%89)に従って更新してください。
-  - 設定ファイルによるデータセット定義機能を追加しました。素晴らしいPRを提供していただいた fur0ut0 氏に感謝します。
-    - ドキュメントは[こちら](./config_README-ja.md)。
-    - ``--dataset_config`` オプションで ``.toml`` ファイルを指定してください。
-    - 今までのオプションはそのまま使えます。
-    - 大規模なアップデートのため、もし不具合がありましたらご報告ください。
-  - 学習の途中でサンプル画像を生成する機能を各学習スクリプトに追加しました。
-    - ``--sample_every_n_steps`` と ``--sample_every_n_epochs`` オプション：生成頻度を指定
-    - ``--sample_prompts`` オプション：プロンプトを記述したファイルを指定（1行ごとに1枚の画像を生成）
-      - プロンプトには ``gen_img_diffusers.py`` のプロンプトオプションの一部、 ``w, h, d, l, s, n`` が使えます。
-    - ``--sample_sampler`` オプション：ddim や k_euler などの sampler (scheduler) を指定します。使用できる sampler についてはヘルプをご覧ください。
-  - ``--tokenizer_cache_dir`` オプションを各学習スクリプトおよび生成スクリプトに追加しました。Diffusers から Tokenizer を取得してきてろーかるに保存します。
-    - 一度キャッシュしておくことでオフライン学習、生成ができるかもしれません。
-  - ``gen_img_diffusers.py`` で highres. fix での letents upscaling と VAE のバッチサイズ指定に対応しました。
-
  - サンプル画像生成：
    プロンプトファイルは例えば以下のようになります。

--- a/fine_tune_README_ja.md
+++ b/fine_tune_README_ja.md
@@ -1,6 +1,9 @@
-NovelAIの提案した学習手法、自動キャプションニング、タグ付け、Windows＋VRAM 12GB（v1.4/1.5の場合）環境等に対応したfine tuningです。
+NovelAIの提案した学習手法、自動キャプションニング、タグ付け、Windows＋VRAM 12GB（SD v1.xの場合）環境等に対応したfine tuningです。ここでfine tuningとは、モデルを画像とキャプションで学習することを指します（LoRAやTextual Inversion、Hypernetworksは含みません）
+
+[学習についての共通ドキュメント](./train_README-ja.md) もあわせてご覧ください。
+
+# 概要

-## 概要
 Diffusersを用いてStable DiffusionのU-Netのfine tuningを行います。NovelAIの記事にある以下の改善に対応しています（Aspect Ratio BucketingについてはNovelAIのコードを参考にしましたが、最終的なコードはすべてオリジナルです）。

 * CLIP（Text Encoder）の最後の層ではなく最後から二番目の層の出力を用いる。
@@ -13,19 +16,24 @@ Diffusersを用いてStable DiffusionのU-Netのfine tuningを行います。Nov

 デフォルトではText Encoderの学習は行いません。モデル全体のfine tuningではU-Netだけを学習するのが一般的なようです（NovelAIもそのようです）。オプション指定でText Encoderも学習対象とできます。

-## 追加機能について
-### CLIPの出力の変更
+# 追加機能について
+
+## CLIPの出力の変更
+
 プロンプトを画像に反映するため、テキストの特徴量への変換を行うのがCLIP（Text Encoder）です。Stable DiffusionではCLIPの最後の層の出力を用いていますが、それを最後から二番目の層の出力を用いるよう変更できます。NovelAIによると、これによりより正確にプロンプトが反映されるようになるとのことです。
 元のまま、最後の層の出力を用いることも可能です。
+
 ※Stable Diffusion 2.0では最後から二番目の層をデフォルトで使います。clip_skipオプションを指定しないでください。

-### 正方形以外の解像度での学習
+## 正方形以外の解像度での学習
+
 Stable Diffusionは512\*512で学習されていますが、それに加えて256\*1024や384\*640といった解像度でも学習します。これによりトリミングされる部分が減り、より正しくプロンプトと画像の関係が学習されることが期待されます。
 学習解像度はパラメータとして与えられた解像度の面積（＝メモリ使用量）を超えない範囲で、64ピクセル単位で縦横に調整、作成されます。

 機械学習では入力サイズをすべて統一するのが一般的ですが、特に制約があるわけではなく、実際は同一のバッチ内で統一されていれば大丈夫です。NovelAIの言うbucketingは、あらかじめ教師データを、アスペクト比に応じた学習解像度ごとに分類しておくことを指しているようです。そしてバッチを各bucket内の画像で作成することで、バッチの画像サイズを統一します。

-### トークン長の75から225への拡張
+## トークン長の75から225への拡張
+
 Stable Diffusionでは最大75トークン（開始・終了を含むと77トークン）ですが、それを225トークンまで拡張します。
 ただしCLIPが受け付ける最大長は75トークンですので、225トークンの場合、単純に三分割してCLIPを呼び出してから結果を連結しています。

@@ -33,296 +41,67 @@ Stable Diffusionでは最大75トークン（開始・終了を含むと77トー

 ※Automatic1111氏のWeb UIではカンマを意識して分割、といったこともしているようですが、私の場合はそこまでしておらず単純な分割です。

-## 環境整備
+# 学習の手順

-このリポジトリの[README](./README-ja.md)を参照してください。
+あらかじめこのリポジトリのREADMEを参照し、環境整備を行ってください。

-## 教師データの用意
-
-学習させたい画像データを用意し、任意のフォルダに入れてください。リサイズ等の事前の準備は必要ありません。
-ただし学習解像度よりもサイズが小さい画像については、超解像などで品質を保ったまま拡大しておくことをお勧めします。
-
-複数の教師データフォルダにも対応しています。前処理をそれぞれのフォルダに対して実行する形となります。
-
-たとえば以下のように画像を格納します。
-
-![教師データフォルダのスクショ](https://user-images.githubusercontent.com/52813779/208907739-8e89d5fa-6ca8-4b60-8927-f484d2a9ae04.png)
-
-## 自動キャプショニング
-キャプションを使わずタグだけで学習する場合はスキップしてください。
-
-また手動でキャプションを用意する場合、キャプションは教師データ画像と同じディレクトリに、同じファイル名、拡張子.caption等で用意してください。各ファイルは1行のみのテキストファイルとします。
-
-### BLIPによるキャプショニング
-
-最新版ではBLIPのダウンロード、重みのダウンロード、仮想環境の追加は不要になりました。そのままで動作します。
-
-finetuneフォルダ内のmake_captions.pyを実行します。
-
-```
-python finetune\make_captions.py --batch_size <バッチサイズ> <教師データフォルダ>
-```
-
-バッチサイズ8、教師データを親フォルダのtrain_dataに置いた場合、以下のようになります。
-
-```
-python finetune\make_captions.py --batch_size 8 ..\train_data
-```
-
-キャプションファイルが教師データ画像と同じディレクトリに、同じファイル名、拡張子.captionで作成されます。
-
-batch_sizeはGPUのVRAM容量に応じて増減してください。大きいほうが速くなります（VRAM 12GBでももう少し増やせると思います）。
-max_lengthオプションでキャプションの最大長を指定できます。デフォルトは75です。モデルをトークン長225で学習する場合には長くしても良いかもしれません。
-caption_extensionオプションでキャプションの拡張子を変更できます。デフォルトは.captionです（.txtにすると後述のDeepDanbooruと競合します）。
-
-複数の教師データフォルダがある場合には、それぞれのフォルダに対して実行してください。
-
-なお、推論にランダム性があるため、実行するたびに結果が変わります。固定する場合には--seedオプションで「--seed 42」のように乱数seedを指定してください。
-
-その他のオプションは--helpでヘルプをご参照ください（パラメータの意味についてはドキュメントがまとまっていないようで、ソースを見るしかないようです）。
-
-デフォルトでは拡張子.captionでキャプションファイルが生成されます。
-
-![captionが生成されたフォルダ](https://user-images.githubusercontent.com/52813779/208908845-48a9d36c-f6ee-4dae-af71-9ab462d1459e.png)
-
-たとえば以下のようなキャプションが付きます。
-
-![キャプションと画像](https://user-images.githubusercontent.com/52813779/208908947-af936957-5d73-4339-b6c8-945a52857373.png)
-
-## DeepDanbooruによるタグ付け
-danbooruタグのタグ付け自体を行わない場合は「キャプションとタグ情報の前処理」に進んでください。
-
-タグ付けはDeepDanbooruまたはWD14Taggerで行います。WD14Taggerのほうが精度が良いようです。WD14Taggerでタグ付けする場合は、次の章へ進んでください。
-
-### 環境整備
-DeepDanbooru https://github.com/KichangKim/DeepDanbooru  を作業フォルダにcloneしてくるか、zipをダウンロードして展開します。私はzipで展開しました。
-またDeepDanbooruのReleasesのページ https://github.com/KichangKim/DeepDanbooru/releases  の「DeepDanbooru Pretrained Model v3-20211112-sgd-e28」のAssetsから、deepdanbooru-v3-20211112-sgd-e28.zipをダウンロードしてきてDeepDanbooruのフォルダに展開します。
-
-以下からダウンロードします。Assetsをクリックして開き、そこからダウンロードします。
-
-![DeepDanbooruダウンロードページ](https://user-images.githubusercontent.com/52813779/208909417-10e597df-7085-41ee-bd06-3e856a1339df.png)
-
-以下のようなこういうディレクトリ構造にしてください
-
-![DeepDanbooruのディレクトリ構造](https://user-images.githubusercontent.com/52813779/208909486-38935d8b-8dc6-43f1-84d3-fef99bc471aa.png)
-
-Diffusersの環境に必要なライブラリをインストールします。DeepDanbooruのフォルダに移動してインストールします（実質的にはtensorflow-ioが追加されるだけだと思います）。
-
-```
-pip install -r requirements.txt
-```
-
-続いてDeepDanbooru自体をインストールします。
-
-```
-pip install .
-```
-
-以上でタグ付けの環境整備は完了です。
-
-### タグ付けの実施
-DeepDanbooruのフォルダに移動し、deepdanbooruを実行してタグ付けを行います。
-
-```
-deepdanbooru evaluate <教師データフォルダ> --project-path deepdanbooru-v3-20211112-sgd-e28 --allow-folder --save-txt
-```
-
-教師データを親フォルダのtrain_dataに置いた場合、以下のようになります。
-
-```
-deepdanbooru evaluate ../train_data --project-path deepdanbooru-v3-20211112-sgd-e28 --allow-folder --save-txt
-```
-
-タグファイルが教師データ画像と同じディレクトリに、同じファイル名、拡張子.txtで作成されます。1件ずつ処理されるためわりと遅いです。
-
-複数の教師データフォルダがある場合には、それぞれのフォルダに対して実行してください。
-
-以下のように生成されます。
-
-![DeepDanbooruの生成ファイル](https://user-images.githubusercontent.com/52813779/208909855-d21b9c98-f2d3-4283-8238-5b0e5aad6691.png)
-
-こんな感じにタグが付きます（すごい情報量……）。
-
-![DeepDanbooruタグと画像](https://user-images.githubusercontent.com/52813779/208909908-a7920174-266e-48d5-aaef-940aba709519.png)
-
-## WD14Taggerによるタグ付け
-DeepDanbooruの代わりにWD14Taggerを用いる手順です。
-
-Automatic1111氏のWebUIで使用しているtaggerを利用します。こちらのgithubページ（https://github.com/toriato/stable-diffusion-webui-wd14-tagger#mrsmilingwolfs-model-aka-waifu-diffusion-14-tagger ）の情報を参考にさせていただきました。
-
-最初の環境整備で必要なモジュールはインストール済みです。また重みはHugging Faceから自動的にダウンロードしてきます。
-
-### タグ付けの実施
-スクリプトを実行してタグ付けを行います。
-```
-python tag_images_by_wd14_tagger.py --batch_size <バッチサイズ> <教師データフォルダ>
-```
-
-教師データを親フォルダのtrain_dataに置いた場合、以下のようになります。
-```
-python tag_images_by_wd14_tagger.py --batch_size 4 ..\train_data
-```
-
-初回起動時にはモデルファイルがwd14_tagger_modelフォルダに自動的にダウンロードされます（フォルダはオプションで変えられます）。以下のようになります。
-
-![ダウンロードされたファイル](https://user-images.githubusercontent.com/52813779/208910447-f7eb0582-90d6-49d3-a666-2b508c7d1842.png)
-
-タグファイルが教師データ画像と同じディレクトリに、同じファイル名、拡張子.txtで作成されます。
-
-![生成されたタグファイル](https://user-images.githubusercontent.com/52813779/208910534-ea514373-1185-4b7d-9ae3-61eb50bc294e.png)
-
-![タグと画像](https://user-images.githubusercontent.com/52813779/208910599-29070c15-7639-474f-b3e4-06bd5a3df29e.png)
-
-threshオプションで、判定されたタグのconfidence（確信度）がいくつ以上でタグをつけるかが指定できます。デフォルトはWD14Taggerのサンプルと同じ0.35です。値を下げるとより多くのタグが付与されますが、精度は下がります。
-batch_sizeはGPUのVRAM容量に応じて増減してください。大きいほうが速くなります（VRAM 12GBでももう少し増やせると思います）。caption_extensionオプションでタグファイルの拡張子を変更できます。デフォルトは.txtです。
-model_dirオプションでモデルの保存先フォルダを指定できます。
-またforce_downloadオプションを指定すると保存先フォルダがあってもモデルを再ダウンロードします。
-
-複数の教師データフォルダがある場合には、それぞれのフォルダに対して実行してください。
-
-## キャプションとタグ情報の前処理
-
-スクリプトから処理しやすいようにキャプションとタグをメタデータとしてひとつのファイルにまとめます。
-
-### キャプションの前処理
-
-キャプションをメタデータに入れるには、作業フォルダ内で以下を実行してください（キャプションを学習に使わない場合は実行不要です）（実際は1行で記述します、以下同様）。
-
-```
-python merge_captions_to_metadata.py <教師データフォルダ>
-　  --in_json <読み込むメタデータファイル名> 
-    <メタデータファイル名>
-```
-
-メタデータファイル名は任意の名前です。
-教師データがtrain_data、読み込むメタデータファイルなし、メタデータファイルがmeta_cap.jsonの場合、以下のようになります。
-
-```
-python merge_captions_to_metadata.py train_data meta_cap.json
-```
-
-caption_extensionオプションでキャプションの拡張子を指定できます。
-
-複数の教師データフォルダがある場合には、full_path引数を指定してください（メタデータにフルパスで情報を持つようになります）。そして、それぞれのフォルダに対して実行してください。
-
-```
-python merge_captions_to_metadata.py --full_path 
-    train_data1 meta_cap1.json
-python merge_captions_to_metadata.py --full_path --in_json meta_cap1.json 
-    train_data2 meta_cap2.json
-```
-
-in_jsonを省略すると書き込み先メタデータファイルがあるとそこから読み込み、そこに上書きします。
-
-__※in_jsonオプションと書き込み先を都度書き換えて、別のメタデータファイルへ書き出すようにすると安全です。__
-
-### タグの前処理
-
-同様にタグもメタデータにまとめます（タグを学習に使わない場合は実行不要です）。
-```
-python merge_dd_tags_to_metadata.py <教師データフォルダ> 
-    --in_json <読み込むメタデータファイル名>
-    <書き込むメタデータファイル名>
-```
-
-先と同じディレクトリ構成で、meta_cap.jsonを読み、meta_cap_dd.jsonに書きだす場合、以下となります。
-```
-python merge_dd_tags_to_metadata.py train_data --in_json meta_cap.json meta_cap_dd.json
-```
-
-複数の教師データフォルダがある場合には、full_path引数を指定してください。そして、それぞれのフォルダに対して実行してください。
-
-```
-python merge_dd_tags_to_metadata.py --full_path --in_json meta_cap2.json
-    train_data1 meta_cap_dd1.json
-python merge_dd_tags_to_metadata.py --full_path --in_json meta_cap_dd1.json 
-    train_data2 meta_cap_dd2.json
-```
-
-in_jsonを省略すると書き込み先メタデータファイルがあるとそこから読み込み、そこに上書きします。
-
-__※in_jsonオプションと書き込み先を都度書き換えて、別のメタデータファイルへ書き出すようにすると安全です。__
-
-### キャプションとタグのクリーニング
-ここまででメタデータファイルにキャプションとDeepDanbooruのタグがまとめられています。ただ自動キャプショニングにしたキャプションは表記ゆれなどがあり微妙（※）ですし、タグにはアンダースコアが含まれていたりratingが付いていたりしますので（DeepDanbooruの場合）、エディタの置換機能などを用いてキャプションとタグのクリーニングをしたほうがいいでしょう。
-
-※たとえばアニメ絵の少女を学習する場合、キャプションにはgirl/girls/woman/womenなどのばらつきがあります。また「anime girl」なども単に「girl」としたほうが適切かもしれません。
-
-クリーニング用のスクリプトが用意してありますので、スクリプトの内容を状況に応じて編集してお使いください。
-
-（教師データフォルダの指定は不要になりました。メタデータ内の全データをクリーニングします。）
-
-```
-python clean_captions_and_tags.py <読み込むメタデータファイル名> <書き込むメタデータファイル名>
-```
-
--in_jsonは付きませんのでご注意ください。たとえば次のようになります。
-
-```
-python clean_captions_and_tags.py meta_cap_dd.json meta_clean.json
-```
-
-以上でキャプションとタグの前処理は完了です。
-
-## latentsの事前取得
-
-学習を高速に進めるためあらかじめ画像の潜在表現を取得しディスクに保存しておきます。あわせてbucketing（教師データをアスペクト比に応じて分類する）を行います。
-
-作業フォルダで以下のように入力してください。
-```
-python prepare_buckets_latents.py <教師データフォルダ>  
-    <読み込むメタデータファイル名> <書き込むメタデータファイル名> 
-    <fine tuningするモデル名またはcheckpoint> 
-    --batch_size <バッチサイズ> 
-    --max_resolution <解像度 幅,高さ> 
-    --mixed_precision <精度>
-```
-
-モデルがmodel.ckpt、バッチサイズ4、学習解像度は512\*512、精度no（float32）で、meta_clean.jsonからメタデータを読み込み、meta_lat.jsonに書き込む場合、以下のようになります。
-
-```
-python prepare_buckets_latents.py 
-    train_data meta_clean.json meta_lat.json model.ckpt 
-    --batch_size 4 --max_resolution 512,512 --mixed_precision no
-```
-
-教師データフォルダにnumpyのnpz形式でlatentsが保存されます。
-
-Stable Diffusion 2.0のモデルを読み込む場合は--v2オプションを指定してください（--v_parameterizationは不要です）。
-
-解像度の最小サイズを--min_bucket_resoオプションで、最大サイズを--max_bucket_resoで指定できます。デフォルトはそれぞれ256、1024です。たとえば最小サイズに384を指定すると、256\*1024や320\*768などの解像度は使わなくなります。
-解像度を768\*768のように大きくした場合、最大サイズに1280などを指定すると良いでしょう。
-
--flip_augオプションを指定すると左右反転のaugmentation（データ拡張）を行います。疑似的にデータ量を二倍に増やすことができますが、データが左右対称でない場合に指定すると（例えばキャラクタの外見、髪型など）学習がうまく行かなくなります。
-（反転した画像についてもlatentsを取得し、\*\_flip.npzファイルを保存する単純な実装です。fline_tune.pyには特にオプション指定は必要ありません。\_flip付きのファイルがある場合、flip付き・なしのファイルを、ランダムに読み込みます。）
-
-バッチサイズはVRAM 12GBでももう少し増やせるかもしれません。
-解像度は64で割り切れる数字で、"幅,高さ"で指定します。解像度はfine tuning時のメモリサイズに直結します。VRAM 12GBでは512,512が限界と思われます（※）。16GBなら512,704や512,768まで上げられるかもしれません。なお256,256等にしてもVRAM 8GBでは厳しいようです（パラメータやoptimizerなどは解像度に関係せず一定のメモリが必要なため）。
-
-※batch size 1の学習で12GB VRAM、640,640で動いたとの報告もありました。
-
-以下のようにbucketingの結果が表示されます。
-
-![bucketingの結果](https://user-images.githubusercontent.com/52813779/208911419-71c00fbb-2ce6-49d5-89b5-b78d7715e441.png)
-
-複数の教師データフォルダがある場合には、full_path引数を指定してください。そして、それぞれのフォルダに対して実行してください。
-```
-python prepare_buckets_latents.py --full_path  
-    train_data1 meta_clean.json meta_lat1.json model.ckpt 
-    --batch_size 4 --max_resolution 512,512 --mixed_precision no
-
-python prepare_buckets_latents.py --full_path 
-    train_data2 meta_lat1.json meta_lat2.json model.ckpt 
-    --batch_size 4 --max_resolution 512,512 --mixed_precision no
-
-```
-読み込み元と書き込み先を同じにすることも可能ですが別々の方が安全です。
-
-__※引数を都度書き換えて、別のメタデータファイルに書き込むと安全です。__
+## データの準備

+[学習データの準備について](./train_README-ja.md) を参照してください。fine tuningではメタデータを用いるfine tuning方式のみ対応しています。

 ## 学習の実行
-たとえば以下のように実行します。以下は省メモリ化のための設定です。
+たとえば以下のように実行します。以下は省メモリ化のための設定です。それぞれの行を必要に応じて書き換えてください。
+
+```
+accelerate launch --num_cpu_threads_per_process 1 fine_tune.py 
+    --pretrained_model_name_or_path=<.ckptまたは.safetensordまたはDiffusers版モデルのディレクトリ> 
+    --output_dir=<学習したモデルの出力先フォルダ>  
+    --output_name=<学習したモデル出力時のファイル名> 
+    --dataset_config=<データ準備で作成した.tomlファイル> 
+    --save_model_as=safetensors 
+    --learning_rate=5e-6 --max_train_steps=10000 
+    --use_8bit_adam --xformers --gradient_checkpointing
+    --mixed_precision=fp16
+```
+
+`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` としてください。
+
+学習させるステップ数 `max_train_steps` を10000とします。学習率 `learning_rate` はここでは5e-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を使用します（速度は遅くなります）。
+
+ある程度メモリがある場合は、`.toml` ファイルを編集してバッチサイズをたとえば `4` くらいに増やしてください（高速化と精度向上の可能性があります）。
+
+### よく使われるオプションについて
+
+以下の場合にはオプションに関するドキュメントを参照してください。
+
+- Stable Diffusion 2.xまたはそこからの派生モデルを学習する
+- clip skipを2以上を前提としたモデルを学習する
+- 75トークンを超えたキャプションで学習する
+
+### バッチサイズについて
+
+モデル全体を学習するためLoRA等の学習に比べるとメモリ消費量は多くなります（DreamBoothと同じ）。
+
+### 学習率について
+
+1e-6から5e-6程度が一般的なようです。他のfine tuningの例なども参照してみてください。
+
+### 以前の形式のデータセット指定をした場合のコマンドライン
+
+解像度やバッチサイズをオプションで指定します。コマンドラインの例は以下の通りです。
+
 ```
 accelerate launch --num_cpu_threads_per_process 1 fine_tune.py 
    --pretrained_model_name_or_path=model.ckpt 
@@ -336,76 +115,7 @@ accelerate launch --num_cpu_threads_per_process 1 fine_tune.py
    --save_every_n_epochs=4
 ```

-accelerateのnum_cpu_threads_per_processには通常は1を指定するとよいようです。
-
-pretrained_model_name_or_pathに学習対象のモデルを指定します（Stable DiffusionのcheckpointかDiffusersのモデル）。Stable Diffusionのcheckpointは.ckptと.safetensorsに対応しています（拡張子で自動判定）。
-
-in_jsonにlatentをキャッシュしたときのメタデータファイルを指定します。
-
-train_data_dirに教師データのフォルダを、output_dirに学習後のモデルの出力先フォルダを指定します。
-
-shuffle_captionを指定すると、キャプション、タグをカンマ区切りされた単位でシャッフルして学習します（Waifu Diffusion v1.3で行っている手法です）。
-（先頭のトークンのいくつかをシャッフルせずに固定できます。その他のオプションのkeep_tokensをご覧ください。）
-
-train_batch_sizeにバッチサイズを指定します。VRAM 12GBでは1か2程度を指定してください。解像度によっても指定可能な数は変わってきます。
-学習に使用される実際のデータ量は「バッチサイズ×ステップ数」です。バッチサイズを増やした時には、それに応じてステップ数を下げることが可能です。
-
-learning_rateに学習率を指定します。たとえばWaifu Diffusion v1.3は5e-6のようです。
-max_train_stepsにステップ数を指定します。
-
-use_8bit_adamを指定すると8-bit Adam Optimizerを使用します。省メモリ化、高速化されますが精度は下がる可能性があります。
-
-xformersを指定するとCrossAttentionを置換して省メモリ化、高速化します。
-※11/9時点ではfloat32の学習ではxformersがエラーになるため、bf16/fp16を使うか、代わりにmem_eff_attnを指定して省メモリ版CrossAttentionを使ってください（速度はxformersに劣ります）。
-
-gradient_checkpointingで勾配の途中保存を有効にします。速度は遅くなりますが使用メモリ量が減ります。
-
-mixed_precisionで混合精度を使うか否かを指定します。"fp16"または"bf16"を指定すると省メモリになりますが精度は劣ります。
-"fp16"と"bf16"は使用メモリ量はほぼ同じで、bf16の方が学習結果は良くなるとの話もあります（試した範囲ではあまり違いは感じられませんでした）。
-"no"を指定すると使用しません（float32になります）。
-
-※bf16で学習したcheckpointをAUTOMATIC1111氏のWeb UIで読み込むとエラーになるようです。これはデータ型のbfloat16がWeb UIのモデルsafety checkerでエラーとなるためのようです。save_precisionオプションを指定してfp16またはfloat32形式で保存してください。またはsafetensors形式で保管しても良さそうです。
-
-save_every_n_epochsを指定するとそのエポックだけ経過するたびに学習中のモデルを保存します。
-
-### Stable Diffusion 2.0対応
-Hugging Faceのstable-diffusion-2-baseを使う場合は--v2オプションを、stable-diffusion-2または768-v-ema.ckptを使う場合は--v2と--v_parameterizationの両方のオプションを指定してください。
-
-### メモリに余裕がある場合に精度や速度を上げる
-まずgradient_checkpointingを外すと速度が上がります。ただし設定できるバッチサイズが減りますので、精度と速度のバランスを見ながら設定してください。
-
-バッチサイズを増やすと速度、精度が上がります。メモリが足りる範囲で、1データ当たりの速度を確認しながら増やしてください（メモリがぎりぎりになるとかえって速度が落ちることがあります）。
-
-### 使用するCLIP出力の変更
-clip_skipオプションに2を指定すると、後ろから二番目の層の出力を用います。1またはオプション省略時は最後の層を用います。
-学習したモデルはAutomatic1111氏のWeb UIで推論できるはずです。
-
-※SD2.0はデフォルトで後ろから二番目の層を使うため、SD2.0の学習では指定しないでください。
-
-学習対象のモデルがもともと二番目の層を使うように学習されている場合は、2を指定するとよいでしょう。
-
-そうではなく最後の層を使用していた場合はモデル全体がそれを前提に学習されています。そのため改めて二番目の層を使用して学習すると、望ましい学習結果を得るにはある程度の枚数の教師データ、長めの学習が必要になるかもしれません。
-
-### トークン長の拡張
-max_token_lengthに150または225を指定することでトークン長を拡張して学習できます。
-学習したモデルはAutomatic1111氏のWeb UIで推論できるはずです。
-
-clip_skipと同様に、モデルの学習状態と異なる長さで学習するには、ある程度の教師データ枚数、長めの学習時間が必要になると思われます。
-
-### 学習ログの保存
-logging_dirオプションにログ保存先フォルダを指定してください。TensorBoard形式のログが保存されます。
-
-たとえば--logging_dir=logsと指定すると、作業フォルダにlogsフォルダが作成され、その中の日時フォルダにログが保存されます。
-また--log_prefixオプションを指定すると、日時の前に指定した文字列が追加されます。「--logging_dir=logs --log_prefix=fine_tune_style1」などとして識別用にお使いください。
-
-TensorBoardでログを確認するには、別のコマンドプロンプトを開き、作業フォルダで以下のように入力します（tensorboardはDiffusersのインストール時にあわせてインストールされると思いますが、もし入っていないならpip install tensorboardで入れてください）。
-```
-tensorboard --logdir=logs
-```
-
-### Hypernetworkの学習
-別の記事で解説予定です。
-
+<!-- 
 ### 勾配をfp16とした学習（実験的機能）
 full_fp16オプションを指定すると勾配を通常のfloat32からfloat16（fp16）に変更して学習します（mixed precisionではなく完全なfp16学習になるようです）。これによりSD1.xの512*512サイズでは8GB未満、SD2.xの512*512サイズで12GB未満のVRAM使用量で学習できるようです。

@@ -415,51 +125,16 @@ full_fp16オプションを指定すると勾配を通常のfloat32からfloat16
 （余裕があるようならtrain_batch_sizeを段階的に増やすと若干精度が上がるはずです。）

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

-### その他のオプション
+# fine tuning特有のその他の主なオプション

-#### keep_tokens
-数値を指定するとキャプションの先頭から、指定した数だけのトークン（カンマ区切りの文字列）をシャッフルせず固定します。
+すべてのオプションについては別文書を参照してください。

-キャプションとタグが両方ある場合、学習時のプロンプトは「キャプション,タグ1,タグ2……」のように連結されますので、「--keep_tokens=1」とすれば、学習時にキャプションが必ず先頭に来るようになります。
-
-#### dataset_repeats
-データセットの枚数が極端に少ない場合、epochがすぐに終わってしまうため（epochの区切りで少し時間が掛かります）、数値を指定してデータを何倍かしてepochを長めにしてください。
-
-#### train_text_encoder
+## `train_text_encoder`
 Text Encoderも学習対象とします。メモリ使用量が若干増加します。

 通常のfine tuningではText Encoderは学習対象としませんが（恐らくText Encoderの出力に従うようにU-Netを学習するため）、学習データ数が少ない場合には、DreamBoothのようにText Encoder側に学習させるのも有効的なようです。

-#### save_precision
-checkpoint保存時のデータ形式をfloat、fp16、bf16から指定できます（未指定時は学習中のデータ形式と同じ）。ディスク容量が節約できますがモデルによる生成結果は変わってきます。またfloatやfp16を指定すると、1111氏のWeb UIでも読めるようになるはずです。
-
-※VAEについては元のcheckpointのデータ形式のままになりますので、fp16でもモデルサイズが2GB強まで小さくならない場合があります。
-
-#### save_model_as
-モデルの保存形式を指定します。ckpt、safetensors、diffusers、diffusers_safetensorsのいずれかを指定してください。
-
-Stable Diffusion形式（ckptまたはsafetensors）を読み込み、Diffusers形式で保存する場合、不足する情報はHugging Faceからv1.5またはv2.1の情報を落としてきて補完します。
-
-#### use_safetensors
-このオプションを指定するとsafetensors形式でcheckpointを保存します。保存形式はデフォルト（読み込んだ形式と同じ）になります。
-
-#### save_stateとresume
-save_stateオプションで、途中保存時および最終保存時に、checkpointに加えてoptimizer等の学習状態をフォルダに保存します。これにより中断してから学習再開したときの精度低下が避けられます（optimizerは状態を持ちながら最適化をしていくため、その状態がリセットされると再び初期状態から最適化を行わなくてはなりません）。なお、Accelerateの仕様でステップ数は保存されません。
-
-スクリプト起動時、resumeオプションで状態の保存されたフォルダを指定すると再開できます。
-
-学習状態は一回の保存あたり5GB程度になりますのでディスク容量にご注意ください。
-
-#### gradient_accumulation_steps
-指定したステップ数だけまとめて勾配を更新します。バッチサイズを増やすのと同様の効果がありますが、メモリを若干消費します。
-
-※Accelerateの仕様で学習モデルが複数の場合には対応していないとのことですので、Text Encoderを学習対象にして、このオプションに2以上の値を指定するとエラーになるかもしれません。
-
-#### lr_scheduler / lr_warmup_steps
-lr_schedulerオプションで学習率のスケジューラをlinear, cosine, cosine_with_restarts, polynomial, constant, constant_with_warmupから選べます。デフォルトはconstantです。
-
-lr_warmup_stepsでスケジューラのウォームアップ（だんだん学習率を変えていく）ステップ数を指定できます。詳細については各自お調べください。
-
-#### diffusers_xformers
+## `diffusers_xformers`
 スクリプト独自のxformers置換機能ではなくDiffusersのxformers機能を利用します。Hypernetworkの学習はできなくなります。
--- a/gen_img_diffusers.py
+++ b/gen_img_diffusers.py
@@ -1649,10 +1649,11 @@ def get_unweighted_text_embeddings(
      if pad == eos:                        # v1
        text_input_chunk[:, -1] = text_input[0, -1]
      else:                                 # v2
-        if text_input_chunk[:, -1] != eos and text_input_chunk[:, -1] != pad:     # 最後に普通の文字がある
-          text_input_chunk[:, -1] = eos
-        if text_input_chunk[:, 1] == pad:                                         # BOSだけであとはPAD
-          text_input_chunk[:, 1] = eos
+        for j in range(len(text_input_chunk)):
+          if text_input_chunk[j, -1] != eos and text_input_chunk[j, -1] != pad:     # 最後に普通の文字がある
+            text_input_chunk[j, -1] = eos
+          if text_input_chunk[j, 1] == pad:                                         # BOSだけであとはPAD
+            text_input_chunk[j, 1] = eos

      if clip_skip is None or clip_skip == 1:
        text_embedding = pipe.text_encoder(text_input_chunk)[0]
@@ -2276,13 +2277,26 @@ def main(args):
      mask_images = l

  # 画像サイズにオプション指定があるときはリサイズする
-  if init_images is not None and args.W is not None and args.H is not None:
-    print(f"resize img2img source images to {args.W}*{args.H}")
-    init_images = resize_images(init_images, (args.W, args.H))
+  if args.W is not None and args.H is not None:
+    if init_images is not None:
+      print(f"resize img2img source images to {args.W}*{args.H}")
+      init_images = resize_images(init_images, (args.W, args.H))
    if mask_images is not None:
      print(f"resize img2img mask images to {args.W}*{args.H}")
      mask_images = resize_images(mask_images, (args.W, args.H))

+  if networks and mask_images:
+    # mask を領域情報として流用する、現在は1枚だけ対応
+    # TODO 複数のnetwork classの混在時の考慮
+    print("use mask as region")
+    # import cv2
+    # for i in range(3):
+    #   cv2.imshow("msk", np.array(mask_images[0])[:,:,i])
+    #   cv2.waitKey()
+    #   cv2.destroyAllWindows()
+    networks[0].__class__.set_regions(networks, np.array(mask_images[0]))
+    mask_images = None
+
  prev_image = None               # for VGG16 guided
  if args.guide_image_path is not None:
    print(f"load image for CLIP/VGG16/ControlNet guidance: {args.guide_image_path}")
--- a/library/model_util.py
+++ b/library/model_util.py
@@ -4,7 +4,7 @@
 import math
 import os
 import torch
-from transformers import CLIPTextModel, CLIPTokenizer, CLIPTextConfig
+from transformers import CLIPTextModel, CLIPTokenizer, CLIPTextConfig, logging
 from diffusers import AutoencoderKL, DDIMScheduler, StableDiffusionPipeline, UNet2DConditionModel
 from safetensors.torch import load_file, save_file

@@ -916,7 +916,11 @@ def load_models_from_stable_diffusion_checkpoint(v2, ckpt_path, dtype=None):
    info = text_model.load_state_dict(converted_text_encoder_checkpoint)
  else:
    converted_text_encoder_checkpoint = convert_ldm_clip_checkpoint_v1(state_dict)
+
+    logging.set_verbosity_error()                                                       # don't show annoying warning
    text_model = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14")
+    logging.set_verbosity_warning()
+
    info = text_model.load_state_dict(converted_text_encoder_checkpoint)
  print("loading text encoder:", info)

--- a/library/train_util.py
+++ b/library/train_util.py
@@ -7,13 +7,13 @@ import re
 import shutil
 import time
 from typing import (
-  Dict,
-  List,
-  NamedTuple,
-  Optional,
-  Sequence,
-  Tuple,
-  Union,
+    Dict,
+    List,
+    NamedTuple,
+    Optional,
+    Sequence,
+    Tuple,
+    Union,
 )
 from accelerate import Accelerator
 import glob
@@ -214,24 +214,24 @@ class AugHelper:
  def __init__(self):
    # prepare all possible augmentators
    color_aug_method = albu.OneOf([
-      albu.HueSaturationValue(8, 0, 0, p=.5),
-      albu.RandomGamma((95, 105), p=.5),
+        albu.HueSaturationValue(8, 0, 0, p=.5),
+        albu.RandomGamma((95, 105), p=.5),
    ], p=.33)
    flip_aug_method = albu.HorizontalFlip(p=0.5)

    # key: (use_color_aug, use_flip_aug)
    self.augmentors = {
-      (True, True): albu.Compose([
-          color_aug_method,
-          flip_aug_method,
-      ], p=1.),
-      (True, False): albu.Compose([
-          color_aug_method,
-      ], p=1.),
-      (False, True): albu.Compose([
-          flip_aug_method,
-      ], p=1.),
-      (False, False): None
+        (True, True): albu.Compose([
+            color_aug_method,
+            flip_aug_method,
+        ], p=1.),
+        (True, False): albu.Compose([
+            color_aug_method,
+        ], p=1.),
+        (False, True): albu.Compose([
+            flip_aug_method,
+        ], p=1.),
+        (False, False): None
    }

  def get_augmentor(self, use_color_aug: bool, use_flip_aug: bool) -> Optional[albu.Compose]:
@@ -260,7 +260,7 @@ class DreamBoothSubset(BaseSubset):
    assert image_dir is not None, "image_dir must be specified / image_dirは指定が必須です"

    super().__init__(image_dir, num_repeats, shuffle_caption, keep_tokens, color_aug, flip_aug,
-          face_crop_aug_range, random_crop, caption_dropout_rate, caption_dropout_every_n_epochs, caption_tag_dropout_rate)
+                     face_crop_aug_range, random_crop, caption_dropout_rate, caption_dropout_every_n_epochs, caption_tag_dropout_rate)

    self.is_reg = is_reg
    self.class_tokens = class_tokens
@@ -271,12 +271,13 @@ class DreamBoothSubset(BaseSubset):
      return NotImplemented
    return self.image_dir == other.image_dir

+
 class FineTuningSubset(BaseSubset):
  def __init__(self, image_dir, metadata_file: str, num_repeats, shuffle_caption, keep_tokens, color_aug, flip_aug, face_crop_aug_range, random_crop, caption_dropout_rate, caption_dropout_every_n_epochs, caption_tag_dropout_rate) -> None:
    assert metadata_file is not None, "metadata_file must be specified / metadata_fileは指定が必須です"

    super().__init__(image_dir, num_repeats, shuffle_caption, keep_tokens, color_aug, flip_aug,
-          face_crop_aug_range, random_crop, caption_dropout_rate, caption_dropout_every_n_epochs, caption_tag_dropout_rate)
+                     face_crop_aug_range, random_crop, caption_dropout_rate, caption_dropout_every_n_epochs, caption_tag_dropout_rate)

    self.metadata_file = metadata_file

@@ -285,6 +286,7 @@ class FineTuningSubset(BaseSubset):
      return NotImplemented
    return self.metadata_file == other.metadata_file

+
 class BaseDataset(torch.utils.data.Dataset):
  def __init__(self, tokenizer: CLIPTokenizer, max_token_length: int, resolution: Optional[Tuple[int, int]], debug_dataset: bool) -> None:
    super().__init__()
@@ -804,7 +806,7 @@ class DreamBoothDataset(BaseDataset):
          captions.append("")
        else:
          captions.append(subset.class_tokens if cap_for_img is None else cap_for_img)
-          
+
      self.set_tag_frequency(os.path.basename(subset.image_dir), captions)         # タグ頻度を記録

      return img_paths, captions
@@ -815,11 +817,13 @@ class DreamBoothDataset(BaseDataset):
    reg_infos: List[ImageInfo] = []
    for subset in subsets:
      if subset.num_repeats < 1:
-        print(f"ignore subset with image_dir='{subset.image_dir}': num_repeats is less than 1 / num_repeatsが1を下回っているためサブセットを無視します: {subset.num_repeats}")
+        print(
+            f"ignore subset with image_dir='{subset.image_dir}': num_repeats is less than 1 / num_repeatsが1を下回っているためサブセットを無視します: {subset.num_repeats}")
        continue

      if subset in self.subsets:
-        print(f"ignore duplicated subset with image_dir='{subset.image_dir}': use the first one / 既にサブセットが登録されているため、重複した後発のサブセットを無視します")
+        print(
+            f"ignore duplicated subset with image_dir='{subset.image_dir}': use the first one / 既にサブセットが登録されているため、重複した後発のサブセットを無視します")
        continue

      img_paths, captions = load_dreambooth_dir(subset)
@@ -881,11 +885,13 @@ class FineTuningDataset(BaseDataset):

    for subset in subsets:
      if subset.num_repeats < 1:
-        print(f"ignore subset with metadata_file='{subset.metadata_file}': num_repeats is less than 1 / num_repeatsが1を下回っているためサブセットを無視します: {subset.num_repeats}")
+        print(
+            f"ignore subset with metadata_file='{subset.metadata_file}': num_repeats is less than 1 / num_repeatsが1を下回っているためサブセットを無視します: {subset.num_repeats}")
        continue

      if subset in self.subsets:
-        print(f"ignore duplicated subset with metadata_file='{subset.metadata_file}': use the first one / 既にサブセットが登録されているため、重複した後発のサブセットを無視します")
+        print(
+            f"ignore duplicated subset with metadata_file='{subset.metadata_file}': use the first one / 既にサブセットが登録されているため、重複した後発のサブセットを無視します")
        continue

      # メタデータを読み込む
@@ -906,10 +912,14 @@ class FineTuningDataset(BaseDataset):
        if os.path.exists(image_key):
          abs_path = image_key
        else:
-          # わりといい加減だがいい方法が思いつかん
-          abs_path = glob_images(subset.image_dir, image_key)
-          assert len(abs_path) >= 1, f"no image / 画像がありません: {image_key}"
-          abs_path = abs_path[0]
+          npz_path = os.path.join(subset.image_dir, image_key + ".npz")
+          if os.path.exists(npz_path):
+            abs_path = npz_path
+          else:
+            # わりといい加減だがいい方法が思いつかん
+            abs_path = glob_images(subset.image_dir, image_key)
+            assert len(abs_path) >= 1, f"no image / 画像がありません: {image_key}"
+            abs_path = abs_path[0]

        caption = img_md.get('caption')
        tags = img_md.get('tags')
@@ -918,7 +928,9 @@ class FineTuningDataset(BaseDataset):
        elif tags is not None and len(tags) > 0:
          caption = caption + ', ' + tags
          tags_list.append(tags)
-        assert caption is not None and len(caption) > 0, f"caption or tag is required / キャプションまたはタグは必須です:{abs_path}"
+
+        if caption is None:
+          caption = ""

        image_info = ImageInfo(image_key, subset.num_repeats, caption, False, abs_path)
        image_info.image_size = img_md.get('train_resolution')
@@ -937,7 +949,7 @@ class FineTuningDataset(BaseDataset):
      self.subsets.append(subset)

    # check existence of all npz files
-    use_npz_latents = all([not(subset.color_aug or subset.random_crop) for subset in self.subsets])
+    use_npz_latents = all([not (subset.color_aug or subset.random_crop) for subset in self.subsets])
    if use_npz_latents:
      flip_aug_in_subset = False
      npz_any = False
@@ -1749,15 +1761,22 @@ def get_optimizer(args, trainable_params):
      raise ImportError("No dadaptation / dadaptation がインストールされていないようです")
    print(f"use D-Adaptation Adam optimizer | {optimizer_kwargs}")

-    min_lr = lr
+    actual_lr = lr
+    lr_count = 1
    if type(trainable_params) == list and type(trainable_params[0]) == dict:
+      lrs = set()
+      actual_lr = trainable_params[0].get("lr", actual_lr)
      for group in trainable_params:
-        min_lr = min(min_lr, group.get("lr", lr))
+        lrs.add(group.get("lr", actual_lr))
+      lr_count = len(lrs)

-    if min_lr <= 0.1:
+    if actual_lr <= 0.1:
      print(
-          f'learning rate is too low. If using dadaptation, set learning rate around 1.0 / 学習率が低すぎるようです。1.0前後の値を指定してください: {min_lr}')
+          f'learning rate is too low. If using dadaptation, set learning rate around 1.0 / 学習率が低すぎるようです。1.0前後の値を指定してください: lr={actual_lr}')
      print('recommend option: lr=1.0 / 推奨は1.0です')
+    if lr_count > 1:
+      print(
+          f"when multiple learning rates are specified with dadaptation (e.g. for Text Encoder and U-Net), only the first one will take effect / D-Adaptationで複数の学習率を指定した場合（Text EncoderとU-Netなど）、最初の学習率のみが有効になります: lr={actual_lr}")

    optimizer_class = dadaptation.DAdaptAdam
    optimizer = optimizer_class(trainable_params, lr=lr, **optimizer_kwargs)
@@ -2201,7 +2220,7 @@ def sample_images(accelerator, args: argparse.Namespace, epoch, steps, device, v
    if epoch is None or epoch % args.sample_every_n_epochs != 0:
      return
  else:
-    if steps % args.sample_every_n_steps != 0:
+    if steps % args.sample_every_n_steps != 0 or epoch is not None:       # steps is not divisible or end of epoch
      return

  print(f"generating sample images at step / サンプル画像生成 ステップ: {steps}")
@@ -2209,8 +2228,6 @@ def sample_images(accelerator, args: argparse.Namespace, epoch, steps, device, v
    print(f"No prompt file / プロンプトファイルがありません: {args.sample_prompts}")
    return

-  # ここでCUDAのキャッシュクリアとかしたほうがいいのか……
-
  org_vae_device = vae.device                           # CPUにいるはず
  vae.to(device)

@@ -2290,6 +2307,8 @@ def sample_images(accelerator, args: argparse.Namespace, epoch, steps, device, v
  with torch.no_grad():
    with accelerator.autocast():
      for i, prompt in enumerate(prompts):
+        if not accelerator.is_main_process:
+          continue
        prompt = prompt.strip()
        if len(prompt) == 0 or prompt[0] == '#':
          continue
@@ -2346,7 +2365,15 @@ def sample_images(accelerator, args: argparse.Namespace, epoch, steps, device, v
          prompt = prompt.replace(prompt_replacement[0], prompt_replacement[1])
          if negative_prompt is not None:
            negative_prompt = negative_prompt.replace(prompt_replacement[0], prompt_replacement[1])
-            
+
+        height = max(64, height - height % 8)                 # round to divisible by 8
+        width = max(64, width - width % 8)                 # round to divisible by 8
+        print(f"prompt: {prompt}")
+        print(f"negative_prompt: {negative_prompt}")
+        print(f"height: {height}")
+        print(f"width: {width}")
+        print(f"sample_steps: {sample_steps}")
+        print(f"scale: {scale}")
        image = pipeline(prompt, height, width, sample_steps, scale, negative_prompt).images[0]

        ts_str = time.strftime('%Y%m%d%H%M%S', time.localtime())
@@ -2356,6 +2383,10 @@ def sample_images(accelerator, args: argparse.Namespace, epoch, steps, device, v

        image.save(os.path.join(save_dir, img_filename))

+  # clear pipeline and cache to reduce vram usage
+  del pipeline
+  torch.cuda.empty_cache()
+
  torch.set_rng_state(rng_state)
  torch.cuda.set_rng_state(cuda_rng_state)
  vae.to(org_vae_device)
--- a/networks/check_lora_weights.py
+++ b/networks/check_lora_weights.py
@@ -21,7 +21,7 @@ def main(file):

  for key, value in values:
    value = value.to(torch.float32)
-    print(f"{key},{torch.mean(torch.abs(value))},{torch.min(torch.abs(value))}")
+    print(f"{key},{str(tuple(value.size())).replace(', ', '-')},{torch.mean(torch.abs(value))},{torch.min(torch.abs(value))}")


 if __name__ == '__main__':
--- a/networks/extract_lora_from_models.py
+++ b/networks/extract_lora_from_models.py
@@ -45,8 +45,13 @@ def svd(args):
  text_encoder_t, _, unet_t = model_util.load_models_from_stable_diffusion_checkpoint(args.v2, args.model_tuned)

  # create LoRA network to extract weights: Use dim (rank) as alpha
-  lora_network_o = lora.create_network(1.0, args.dim, args.dim, None, text_encoder_o, unet_o)
-  lora_network_t = lora.create_network(1.0, args.dim, args.dim, None, text_encoder_t, unet_t)
+  if args.conv_dim is None:
+    kwargs = {}
+  else:
+    kwargs = {"conv_dim": args.conv_dim, "conv_alpha": args.conv_dim}
+
+  lora_network_o = lora.create_network(1.0, args.dim, args.dim, None, text_encoder_o, unet_o, **kwargs)
+  lora_network_t = lora.create_network(1.0, args.dim, args.dim, None, text_encoder_t, unet_t, **kwargs)
  assert len(lora_network_o.text_encoder_loras) == len(
      lora_network_t.text_encoder_loras), f"model version is different (SD1.x vs SD2.x) / それぞれのモデルのバージョンが違います（SD1.xベースとSD2.xベース） "

@@ -85,13 +90,28 @@ def svd(args):

  # make LoRA with svd
  print("calculating by svd")
-  rank = args.dim
  lora_weights = {}
  with torch.no_grad():
    for lora_name, mat in tqdm(list(diffs.items())):
+      # if args.conv_dim is None, diffs do not include LoRAs for conv2d-3x3
      conv2d = (len(mat.size()) == 4)
+      kernel_size = None if not conv2d else mat.size()[2:4]
+      conv2d_3x3 = conv2d and kernel_size != (1, 1)
+
+      rank = args.dim if not conv2d_3x3 or args.conv_dim is None else args.conv_dim
+      out_dim, in_dim = mat.size()[0:2]
+
+      if args.device:
+        mat = mat.to(args.device)
+
+      # print(lora_name, mat.size(), mat.device, rank, in_dim, out_dim)
+      rank = min(rank, in_dim, out_dim)                           # LoRA rank cannot exceed the original dim
+
      if conv2d:
-        mat = mat.squeeze()
+        if conv2d_3x3:
+          mat = mat.flatten(start_dim=1)
+        else:
+          mat = mat.squeeze()

      U, S, Vh = torch.linalg.svd(mat)

@@ -108,30 +128,27 @@ def svd(args):
      U = U.clamp(low_val, hi_val)
      Vh = Vh.clamp(low_val, hi_val)

+      if conv2d:
+        U = U.reshape(out_dim, rank, 1, 1)
+        Vh = Vh.reshape(rank, in_dim, kernel_size[0], kernel_size[1])
+
+      U = U.to("cpu").contiguous()
+      Vh = Vh.to("cpu").contiguous()
+
      lora_weights[lora_name] = (U, Vh)

  # make state dict for LoRA
-  lora_network_o.apply_to(text_encoder_o, unet_o, text_encoder_different, True)   # to make state dict
-  lora_sd = lora_network_o.state_dict()
-  print(f"LoRA has {len(lora_sd)} weights.")
-
-  for key in list(lora_sd.keys()):
-    if "alpha" in key:
-      continue
-
-    lora_name = key.split('.')[0]
-    i = 0 if "lora_up" in key else 1
-
-    weights = lora_weights[lora_name][i]
-    # print(key, i, weights.size(), lora_sd[key].size())
-    if len(lora_sd[key].size()) == 4:
-      weights = weights.unsqueeze(2).unsqueeze(3)
-
-    assert weights.size() == lora_sd[key].size(), f"size unmatch: {key}"
-    lora_sd[key] = weights
+  lora_sd = {}
+  for lora_name, (up_weight, down_weight) in lora_weights.items():
+    lora_sd[lora_name + '.lora_up.weight'] = up_weight
+    lora_sd[lora_name + '.lora_down.weight'] = down_weight
+    lora_sd[lora_name + '.alpha'] = torch.tensor(down_weight.size()[0])

  # load state dict to LoRA and save it
-  info = lora_network_o.load_state_dict(lora_sd)
+  lora_network_save = lora.create_network_from_weights(1.0, None, None, text_encoder_o, unet_o, weights_sd=lora_sd)
+  lora_network_save.apply_to(text_encoder_o, unet_o)        # create internal module references for state_dict
+
+  info = lora_network_save.load_state_dict(lora_sd)
  print(f"Loading extracted LoRA weights: {info}")

  dir_name = os.path.dirname(args.save_to)
@@ -139,9 +156,9 @@ def svd(args):
    os.makedirs(dir_name, exist_ok=True)

  # minimum metadata
-  metadata = {"ss_network_dim": str(args.dim), "ss_network_alpha": str(args.dim)}
+  metadata = {"ss_network_module": "networks.lora", "ss_network_dim": str(args.dim), "ss_network_alpha": str(args.dim)}

-  lora_network_o.save_weights(args.save_to, save_dtype, metadata)
+  lora_network_save.save_weights(args.save_to, save_dtype, metadata)
  print(f"LoRA weights are saved to: {args.save_to}")


@@ -158,6 +175,8 @@ if __name__ == '__main__':
  parser.add_argument("--save_to", type=str, default=None,
                      help="destination file name: ckpt or safetensors file / 保存先のファイル名、ckptまたはsafetensors")
  parser.add_argument("--dim", type=int, default=4, help="dimension (rank) of LoRA (default 4) / LoRAの次元数（rank）（デフォルト4）")
+  parser.add_argument("--conv_dim", type=int, default=None,
+                      help="dimension (rank) of LoRA for Conv2d-3x3 (default None, disabled) / LoRAのConv2d-3x3の次元数（rank）（デフォルトNone、適用なし）")
  parser.add_argument("--device", type=str, default=None, help="device to use, cuda for GPU / 計算を行うデバイス、cuda でGPUを使う")

  args = parser.parse_args()
--- a/networks/lora.py
+++ b/networks/lora.py
@@ -6,6 +6,7 @@
 import math
 import os
 from typing import List
+import numpy as np
 import torch

 from library import train_util
@@ -20,22 +21,34 @@ class LoRAModule(torch.nn.Module):
    """ if alpha == 0 or None, alpha is rank (no scaling). """
    super().__init__()
    self.lora_name = lora_name
-    self.lora_dim = lora_dim

    if org_module.__class__.__name__ == 'Conv2d':
      in_dim = org_module.in_channels
      out_dim = org_module.out_channels
-      self.lora_down = torch.nn.Conv2d(in_dim, lora_dim, (1, 1), bias=False)
-      self.lora_up = torch.nn.Conv2d(lora_dim, out_dim, (1, 1), bias=False)
    else:
      in_dim = org_module.in_features
      out_dim = org_module.out_features
-      self.lora_down = torch.nn.Linear(in_dim, lora_dim, bias=False)
-      self.lora_up = torch.nn.Linear(lora_dim, out_dim, bias=False)
+
+    # if limit_rank:
+    #   self.lora_dim = min(lora_dim, in_dim, out_dim)
+    #   if self.lora_dim != lora_dim:
+    #     print(f"{lora_name} dim (rank) is changed to: {self.lora_dim}")
+    # else:
+    self.lora_dim = lora_dim
+
+    if org_module.__class__.__name__ == 'Conv2d':
+      kernel_size = org_module.kernel_size
+      stride = org_module.stride
+      padding = org_module.padding
+      self.lora_down = torch.nn.Conv2d(in_dim, self.lora_dim, kernel_size, stride, padding, bias=False)
+      self.lora_up = torch.nn.Conv2d(self.lora_dim, out_dim, (1, 1), (1, 1), bias=False)
+    else:
+      self.lora_down = torch.nn.Linear(in_dim, self.lora_dim, bias=False)
+      self.lora_up = torch.nn.Linear(self.lora_dim, out_dim, bias=False)

    if type(alpha) == torch.Tensor:
      alpha = alpha.detach().float().numpy()                              # without casting, bf16 causes error
-    alpha = lora_dim if alpha is None or alpha == 0 else alpha
+    alpha = self.lora_dim if alpha is None or alpha == 0 else alpha
    self.scale = alpha / self.lora_dim
    self.register_buffer('alpha', torch.tensor(alpha))                    # 定数として扱える

@@ -45,69 +58,192 @@ class LoRAModule(torch.nn.Module):

    self.multiplier = multiplier
    self.org_module = org_module                  # remove in applying
+    self.region = None
+    self.region_mask = None

  def apply_to(self):
    self.org_forward = self.org_module.forward
    self.org_module.forward = self.forward
    del self.org_module

+  def set_region(self, region):
+    self.region = region
+    self.region_mask = None
+
  def forward(self, x):
-    return self.org_forward(x) + self.lora_up(self.lora_down(x)) * self.multiplier * self.scale
+    if self.region is None:
+      return self.org_forward(x) + self.lora_up(self.lora_down(x)) * self.multiplier * self.scale
+
+    # regional LoRA   FIXME same as additional-network extension
+    if x.size()[1] % 77 == 0:
+      # print(f"LoRA for context: {self.lora_name}")
+      self.region = None
+      return self.org_forward(x) + self.lora_up(self.lora_down(x)) * self.multiplier * self.scale
+
+    # calculate region mask first time
+    if self.region_mask is None:
+      if len(x.size()) == 4:
+        h, w = x.size()[2:4]
+      else:
+        seq_len = x.size()[1]
+        ratio = math.sqrt((self.region.size()[0] * self.region.size()[1]) / seq_len)
+        h = int(self.region.size()[0] / ratio + .5)
+        w = seq_len // h
+
+      r = self.region.to(x.device)
+      if r.dtype == torch.bfloat16:
+        r = r.to(torch.float)
+      r = r.unsqueeze(0).unsqueeze(1)
+      # print(self.lora_name, self.region.size(), x.size(), r.size(), h, w)
+      r = torch.nn.functional.interpolate(r, (h, w), mode='bilinear')
+      r = r.to(x.dtype)
+
+      if len(x.size()) == 3:
+        r = torch.reshape(r, (1, x.size()[1], -1))
+
+      self.region_mask = r
+
+    return self.org_forward(x) + self.lora_up(self.lora_down(x)) * self.multiplier * self.scale * self.region_mask


 def create_network(multiplier, network_dim, network_alpha, vae, text_encoder, unet, **kwargs):
  if network_dim is None:
    network_dim = 4                     # default
-  network = LoRANetwork(text_encoder, unet, multiplier=multiplier, lora_dim=network_dim, alpha=network_alpha)
+
+  # extract dim/alpha for conv2d, and block dim
+  conv_dim = kwargs.get('conv_dim', None)
+  conv_alpha = kwargs.get('conv_alpha', None)
+  if conv_dim is not None:
+    conv_dim = int(conv_dim)
+    if conv_alpha is None:
+      conv_alpha = 1.0
+    else:
+      conv_alpha = float(conv_alpha)
+
+  """
+  block_dims = kwargs.get("block_dims")
+  block_alphas = None
+
+  if block_dims is not None:
+    block_dims = [int(d) for d in block_dims.split(',')]
+    assert len(block_dims) == NUM_BLOCKS, f"Number of block dimensions is not same to {NUM_BLOCKS}"
+    block_alphas = kwargs.get("block_alphas")
+    if block_alphas is None:
+      block_alphas = [1] * len(block_dims)
+    else:
+      block_alphas = [int(a) for a in block_alphas(',')]
+    assert len(block_alphas) == NUM_BLOCKS, f"Number of block alphas is not same to {NUM_BLOCKS}"
+
+  conv_block_dims = kwargs.get("conv_block_dims")
+  conv_block_alphas = None
+
+  if conv_block_dims is not None:
+    conv_block_dims = [int(d) for d in conv_block_dims.split(',')]
+    assert len(conv_block_dims) == NUM_BLOCKS, f"Number of block dimensions is not same to {NUM_BLOCKS}"
+    conv_block_alphas = kwargs.get("conv_block_alphas")
+    if conv_block_alphas is None:
+      conv_block_alphas = [1] * len(conv_block_dims)
+    else:
+      conv_block_alphas = [int(a) for a in conv_block_alphas(',')]
+    assert len(conv_block_alphas) == NUM_BLOCKS, f"Number of block alphas is not same to {NUM_BLOCKS}"
+  """
+
+  network = LoRANetwork(text_encoder, unet, multiplier=multiplier, lora_dim=network_dim,
+                        alpha=network_alpha, conv_lora_dim=conv_dim, conv_alpha=conv_alpha)
  return network


-def create_network_from_weights(multiplier, file, vae, text_encoder, unet, **kwargs):
-  if os.path.splitext(file)[1] == '.safetensors':
-    from safetensors.torch import load_file, safe_open
-    weights_sd = load_file(file)
-  else:
-    weights_sd = torch.load(file, map_location='cpu')
+def create_network_from_weights(multiplier, file, vae, text_encoder, unet, weights_sd=None, **kwargs):
+  if weights_sd is None:
+    if os.path.splitext(file)[1] == '.safetensors':
+      from safetensors.torch import load_file, safe_open
+      weights_sd = load_file(file)
+    else:
+      weights_sd = torch.load(file, map_location='cpu')

-  # get dim (rank)
-  network_alpha = None
-  network_dim = None
+  # get dim/alpha mapping
+  modules_dim = {}
+  modules_alpha = {}
  for key, value in weights_sd.items():
-    if network_alpha is None and 'alpha' in key:
-      network_alpha = value
-    if network_dim is None and 'lora_down' in key and len(value.size()) == 2:
-      network_dim = value.size()[0]
+    if '.' not in key:
+      continue

-  if network_alpha is None:
-    network_alpha = network_dim
+    lora_name = key.split('.')[0]
+    if 'alpha' in key:
+      modules_alpha[lora_name] = value
+    elif 'lora_down' in key:
+      dim = value.size()[0]
+      modules_dim[lora_name] = dim
+      # print(lora_name, value.size(), dim)

-  network = LoRANetwork(text_encoder, unet, multiplier=multiplier, lora_dim=network_dim, alpha=network_alpha)
+  # support old LoRA without alpha
+  for key in modules_dim.keys():
+    if key not in modules_alpha:
+      modules_alpha = modules_dim[key]
+
+  network = LoRANetwork(text_encoder, unet, multiplier=multiplier, modules_dim=modules_dim, modules_alpha=modules_alpha)
  network.weights_sd = weights_sd
  return network


 class LoRANetwork(torch.nn.Module):
+  # is it possible to apply conv_in and conv_out?
  UNET_TARGET_REPLACE_MODULE = ["Transformer2DModel", "Attention"]
+  UNET_TARGET_REPLACE_MODULE_CONV2D_3X3 = ["ResnetBlock2D", "Downsample2D", "Upsample2D"]
  TEXT_ENCODER_TARGET_REPLACE_MODULE = ["CLIPAttention", "CLIPMLP"]
  LORA_PREFIX_UNET = 'lora_unet'
  LORA_PREFIX_TEXT_ENCODER = 'lora_te'

-  def __init__(self, text_encoder, unet, multiplier=1.0, lora_dim=4, alpha=1) -> None:
+  def __init__(self, text_encoder, unet, multiplier=1.0, lora_dim=4, alpha=1, conv_lora_dim=None, conv_alpha=None, modules_dim=None, modules_alpha=None) -> None:
    super().__init__()
    self.multiplier = multiplier
+
    self.lora_dim = lora_dim
    self.alpha = alpha
+    self.conv_lora_dim = conv_lora_dim
+    self.conv_alpha = conv_alpha
+
+    if modules_dim is not None:
+      print(f"create LoRA network from weights")
+    else:
+      print(f"create LoRA network. base dim (rank): {lora_dim}, alpha: {alpha}")
+
+    self.apply_to_conv2d_3x3 = self.conv_lora_dim is not None
+    if self.apply_to_conv2d_3x3:
+      if self.conv_alpha is None:
+        self.conv_alpha = self.alpha
+      print(f"apply LoRA to Conv2d with kernel size (3,3). dim (rank): {self.conv_lora_dim}, alpha: {self.conv_alpha}")

    # create module instances
    def create_modules(prefix, root_module: torch.nn.Module, target_replace_modules) -> List[LoRAModule]:
      loras = []
      for name, module in root_module.named_modules():
        if module.__class__.__name__ in target_replace_modules:
+          # TODO get block index here
          for child_name, child_module in module.named_modules():
-            if child_module.__class__.__name__ == "Linear" or (child_module.__class__.__name__ == "Conv2d" and child_module.kernel_size == (1, 1)):
+            is_linear = child_module.__class__.__name__ == "Linear"
+            is_conv2d = child_module.__class__.__name__ == "Conv2d"
+            is_conv2d_1x1 = is_conv2d and child_module.kernel_size == (1, 1)
+            if is_linear or is_conv2d:
              lora_name = prefix + '.' + name + '.' + child_name
              lora_name = lora_name.replace('.', '_')
-              lora = LoRAModule(lora_name, child_module, self.multiplier, self.lora_dim, self.alpha)
+
+              if modules_dim is not None:
+                if lora_name not in modules_dim:
+                  continue                                      # no LoRA module in this weights file
+                dim = modules_dim[lora_name]
+                alpha = modules_alpha[lora_name]
+              else:
+                if is_linear or is_conv2d_1x1:
+                  dim = self.lora_dim
+                  alpha = self.alpha
+                elif self.apply_to_conv2d_3x3:
+                  dim = self.conv_lora_dim
+                  alpha = self.conv_alpha
+                else:
+                  continue
+
+              lora = LoRAModule(lora_name, child_module, self.multiplier, dim, alpha)
              loras.append(lora)
      return loras

@@ -115,7 +251,12 @@ class LoRANetwork(torch.nn.Module):
                                             text_encoder, LoRANetwork.TEXT_ENCODER_TARGET_REPLACE_MODULE)
    print(f"create LoRA for Text Encoder: {len(self.text_encoder_loras)} modules.")

-    self.unet_loras = create_modules(LoRANetwork.LORA_PREFIX_UNET, unet, LoRANetwork.UNET_TARGET_REPLACE_MODULE)
+    # extend U-Net target modules if conv2d 3x3 is enabled, or load from weights
+    target_modules = LoRANetwork.UNET_TARGET_REPLACE_MODULE
+    if modules_dim is not None or self.conv_lora_dim is not None:
+      target_modules += LoRANetwork.UNET_TARGET_REPLACE_MODULE_CONV2D_3X3
+
+    self.unet_loras = create_modules(LoRANetwork.LORA_PREFIX_UNET, unet, target_modules)
    print(f"create LoRA for U-Net: {len(self.unet_loras)} modules.")

    self.weights_sd = None
@@ -130,7 +271,7 @@ class LoRANetwork(torch.nn.Module):
    self.multiplier = multiplier
    for lora in self.text_encoder_loras + self.unet_loras:
      lora.multiplier = self.multiplier
-      
+
  def load_weights(self, file):
    if os.path.splitext(file)[1] == '.safetensors':
      from safetensors.torch import load_file, safe_open
@@ -240,3 +381,18 @@ class LoRANetwork(torch.nn.Module):
      save_file(state_dict, file, metadata)
    else:
      torch.save(state_dict, file)
+
+  @ staticmethod
+  def set_regions(networks, image):
+    image = image.astype(np.float32) / 255.0
+    for i, network in enumerate(networks[:3]):
+      # NOTE: consider averaging overwrapping area
+      region = image[:, :, i]
+      if region.max() == 0:
+        continue
+      region = torch.tensor(region)
+      network.set_region(region)
+
+  def set_region(self, region):
+    for lora in self.unet_loras:
+      lora.set_region(region)
--- a/networks/merge_lora.py
+++ b/networks/merge_lora.py
@@ -48,7 +48,7 @@ def merge_to_sd_model(text_encoder, unet, models, ratios, merge_dtype):
    for name, module in root_module.named_modules():
      if module.__class__.__name__ in target_replace_modules:
        for child_name, child_module in module.named_modules():
-          if child_module.__class__.__name__ == "Linear" or (child_module.__class__.__name__ == "Conv2d" and child_module.kernel_size == (1, 1)):
+          if child_module.__class__.__name__ == "Linear" or child_module.__class__.__name__ == "Conv2d":
            lora_name = prefix + '.' + name + '.' + child_name
            lora_name = lora_name.replace('.', '_')
            name_to_module[lora_name] = child_module
@@ -80,13 +80,19 @@ def merge_to_sd_model(text_encoder, unet, models, ratios, merge_dtype):

        # W <- W + U * D
        weight = module.weight
+        # print(module_name, down_weight.size(), up_weight.size())
        if len(weight.size()) == 2:
          # linear
          weight = weight + ratio * (up_weight @ down_weight) * scale
-        else:
-          # conv2d
+        elif down_weight.size()[2:4] == (1, 1):
+          # conv2d 1x1
          weight = weight + ratio * (up_weight.squeeze(3).squeeze(2) @ down_weight.squeeze(3).squeeze(2)
                                     ).unsqueeze(2).unsqueeze(3) * scale
+        else:
+          # conv2d 3x3
+          conved = torch.nn.functional.conv2d(down_weight.permute(1, 0, 2, 3), up_weight).permute(1, 0, 2, 3)
+          # print(conved.size(), weight.size(), module.stride, module.padding)
+          weight = weight + ratio * conved * scale

        module.weight = torch.nn.Parameter(weight)

@@ -123,7 +129,7 @@ def merge_lora_models(models, ratios, merge_dtype):
        alphas[lora_module_name] = alpha
        if lora_module_name not in base_alphas:
          base_alphas[lora_module_name] = alpha
-    
+
    print(f"dim: {list(set(dims.values()))}, alpha: {list(set(alphas.values()))}")

    # merge
@@ -145,7 +151,7 @@ def merge_lora_models(models, ratios, merge_dtype):
        merged_sd[key] = merged_sd[key] + lora_sd[key] * scale
      else:
        merged_sd[key] = lora_sd[key] * scale
-  
+
  # set alpha to sd
  for lora_module_name, alpha in base_alphas.items():
    key = lora_module_name + ".alpha"
--- a/networks/resize_lora.py
+++ b/networks/resize_lora.py
@@ -1,14 +1,15 @@
 # Convert LoRA to different rank approximation (should only be used to go to lower rank)
 # This code is based off the extract_lora_from_models.py file which is based on https://github.com/cloneofsimo/lora/blob/develop/lora_diffusion/cli_svd.py
-# Thanks to cloneofsimo and kohya
+# Thanks to cloneofsimo

 import argparse
-import os
 import torch
 from safetensors.torch import load_file, save_file, safe_open
 from tqdm import tqdm
 from library import train_util, model_util
+import numpy as np

+MIN_SV = 1e-6

 def load_state_dict(file_name, dtype):
  if model_util.is_safetensors(file_name):
@@ -38,12 +39,149 @@ def save_to_file(file_name, model, state_dict, dtype, metadata):
    torch.save(model, file_name)


-def resize_lora_model(lora_sd, new_rank, save_dtype, device, verbose):
+def index_sv_cumulative(S, target):
+  original_sum = float(torch.sum(S))
+  cumulative_sums = torch.cumsum(S, dim=0)/original_sum
+  index = int(torch.searchsorted(cumulative_sums, target)) + 1
+  if index >= len(S):
+    index = len(S) - 1
+
+  return index
+
+
+def index_sv_fro(S, target):
+  S_squared = S.pow(2)
+  s_fro_sq = float(torch.sum(S_squared))
+  sum_S_squared = torch.cumsum(S_squared, dim=0)/s_fro_sq
+  index = int(torch.searchsorted(sum_S_squared, target**2)) + 1
+  if index >= len(S):
+    index = len(S) - 1
+
+  return index
+
+
+# Modified from Kohaku-blueleaf's extract/merge functions
+def extract_conv(weight, lora_rank, dynamic_method, dynamic_param, device, scale=1):
+    out_size, in_size, kernel_size, _ = weight.size()
+    U, S, Vh = torch.linalg.svd(weight.reshape(out_size, -1).to(device))
+    
+    param_dict = rank_resize(S, lora_rank, dynamic_method, dynamic_param, scale)
+    lora_rank = param_dict["new_rank"]
+
+    U = U[:, :lora_rank]
+    S = S[:lora_rank]
+    U = U @ torch.diag(S)
+    Vh = Vh[:lora_rank, :]
+
+    param_dict["lora_down"] = Vh.reshape(lora_rank, in_size, kernel_size, kernel_size).cpu()
+    param_dict["lora_up"] = U.reshape(out_size, lora_rank, 1, 1).cpu()
+    del U, S, Vh, weight
+    return param_dict
+
+
+def extract_linear(weight, lora_rank, dynamic_method, dynamic_param, device, scale=1):
+    out_size, in_size = weight.size()
+    
+    U, S, Vh = torch.linalg.svd(weight.to(device))
+    
+    param_dict = rank_resize(S, lora_rank, dynamic_method, dynamic_param, scale)
+    lora_rank = param_dict["new_rank"]
+    
+    U = U[:, :lora_rank]
+    S = S[:lora_rank]
+    U = U @ torch.diag(S)
+    Vh = Vh[:lora_rank, :]
+    
+    param_dict["lora_down"] = Vh.reshape(lora_rank, in_size).cpu()
+    param_dict["lora_up"] = U.reshape(out_size, lora_rank).cpu()
+    del U, S, Vh, weight
+    return param_dict
+
+
+def merge_conv(lora_down, lora_up, device):
+    in_rank, in_size, kernel_size, k_ = lora_down.shape
+    out_size, out_rank, _, _ = lora_up.shape
+    assert in_rank == out_rank and kernel_size == k_, f"rank {in_rank} {out_rank} or kernel {kernel_size} {k_} mismatch"
+    
+    lora_down = lora_down.to(device)
+    lora_up = lora_up.to(device)
+
+    merged = lora_up.reshape(out_size, -1) @ lora_down.reshape(in_rank, -1)
+    weight = merged.reshape(out_size, in_size, kernel_size, kernel_size)
+    del lora_up, lora_down
+    return weight
+
+
+def merge_linear(lora_down, lora_up, device):
+    in_rank, in_size = lora_down.shape
+    out_size, out_rank = lora_up.shape
+    assert in_rank == out_rank, f"rank {in_rank} {out_rank} mismatch"
+    
+    lora_down = lora_down.to(device)
+    lora_up = lora_up.to(device)
+    
+    weight = lora_up @ lora_down
+    del lora_up, lora_down
+    return weight
+  
+
+def rank_resize(S, rank, dynamic_method, dynamic_param, scale=1):
+    param_dict = {}
+
+    if dynamic_method=="sv_ratio":
+        # Calculate new dim and alpha based off ratio
+        max_sv = S[0]
+        min_sv = max_sv/dynamic_param
+        new_rank = max(torch.sum(S > min_sv).item(),1)
+        new_alpha = float(scale*new_rank)
+
+    elif dynamic_method=="sv_cumulative":
+        # Calculate new dim and alpha based off cumulative sum
+        new_rank = index_sv_cumulative(S, dynamic_param)
+        new_rank = max(new_rank, 1)
+        new_alpha = float(scale*new_rank)
+
+    elif dynamic_method=="sv_fro":
+        # Calculate new dim and alpha based off sqrt sum of squares
+        new_rank = index_sv_fro(S, dynamic_param)
+        new_rank = min(max(new_rank, 1), len(S)-1)
+        new_alpha = float(scale*new_rank)
+    else:
+        new_rank = rank
+        new_alpha = float(scale*new_rank)
+
+    
+    if S[0] <= MIN_SV: # Zero matrix, set dim to 1
+        new_rank = 1
+        new_alpha = float(scale*new_rank)
+    elif new_rank > rank: # cap max rank at rank
+        new_rank = rank
+        new_alpha = float(scale*new_rank)
+
+
+    # Calculate resize info
+    s_sum = torch.sum(torch.abs(S))
+    s_rank = torch.sum(torch.abs(S[:new_rank]))
+    
+    S_squared = S.pow(2)
+    s_fro = torch.sqrt(torch.sum(S_squared))
+    s_red_fro = torch.sqrt(torch.sum(S_squared[:new_rank]))
+    fro_percent = float(s_red_fro/s_fro)
+
+    param_dict["new_rank"] = new_rank
+    param_dict["new_alpha"] = new_alpha
+    param_dict["sum_retained"] = (s_rank)/s_sum
+    param_dict["fro_retained"] = fro_percent
+    param_dict["max_ratio"] = S[0]/S[new_rank]
+
+    return param_dict
+
+
+def resize_lora_model(lora_sd, new_rank, save_dtype, device, dynamic_method, dynamic_param, verbose):
  network_alpha = None
  network_dim = None
  verbose_str = "\n"
-
-  CLAMP_QUANTILE = 0.99
+  fro_list = []

  # Extract loaded lora dim and alpha
  for key, value in lora_sd.items():
@@ -57,9 +195,9 @@ def resize_lora_model(lora_sd, new_rank, save_dtype, device, verbose):
      network_alpha = network_dim

  scale = network_alpha/network_dim
-  new_alpha = float(scale*new_rank)  # calculate new alpha from scale

-  print(f"old dimension: {network_dim}, old alpha: {network_alpha}, new alpha: {new_alpha}")
+  if dynamic_method:
+    print(f"Dynamically determining new alphas and dims based off {dynamic_method}: {dynamic_param}, max rank is {new_rank}")

  lora_down_weight = None
  lora_up_weight = None
@@ -68,7 +206,6 @@ def resize_lora_model(lora_sd, new_rank, save_dtype, device, verbose):
  block_down_name = None
  block_up_name = None

-  print("resizing lora...")
  with torch.no_grad():
    for key, value in tqdm(lora_sd.items()):
      if 'lora_down' in key:
@@ -85,57 +222,43 @@ def resize_lora_model(lora_sd, new_rank, save_dtype, device, verbose):
        conv2d = (len(lora_down_weight.size()) == 4)

        if conv2d:
-          lora_down_weight = lora_down_weight.squeeze()
-          lora_up_weight = lora_up_weight.squeeze()
-
-        if device:
-          org_device = lora_up_weight.device
-          lora_up_weight = lora_up_weight.to(args.device)
-          lora_down_weight = lora_down_weight.to(args.device)
-
-        full_weight_matrix = torch.matmul(lora_up_weight, lora_down_weight)
-
-        U, S, Vh = torch.linalg.svd(full_weight_matrix)
+          full_weight_matrix = merge_conv(lora_down_weight, lora_up_weight, device)
+          param_dict = extract_conv(full_weight_matrix, new_rank, dynamic_method, dynamic_param, device, scale)
+        else:
+          full_weight_matrix = merge_linear(lora_down_weight, lora_up_weight, device)
+          param_dict = extract_linear(full_weight_matrix, new_rank, dynamic_method, dynamic_param, device, scale)

        if verbose:
-          s_sum = torch.sum(torch.abs(S))
-          s_rank = torch.sum(torch.abs(S[:new_rank]))
-          verbose_str+=f"{block_down_name:76} | "
-          verbose_str+=f"sum(S) retained: {(s_rank)/s_sum:.1%}, max(S) ratio: {S[0]/S[new_rank]:0.1f}\n"
+          max_ratio = param_dict['max_ratio']
+          sum_retained = param_dict['sum_retained']
+          fro_retained = param_dict['fro_retained']
+          if not np.isnan(fro_retained):
+            fro_list.append(float(fro_retained))

-        U = U[:, :new_rank]
-        S = S[:new_rank]
-        U = U @ torch.diag(S)
+          verbose_str+=f"{block_down_name:75} | "
+          verbose_str+=f"sum(S) retained: {sum_retained:.1%}, fro retained: {fro_retained:.1%}, max(S) ratio: {max_ratio:0.1f}"

-        Vh = Vh[:new_rank, :]
+        if verbose and dynamic_method:
+          verbose_str+=f", dynamic | dim: {param_dict['new_rank']}, alpha: {param_dict['new_alpha']}\n"
+        else:
+          verbose_str+=f"\n"

-        dist = torch.cat([U.flatten(), Vh.flatten()])
-        hi_val = torch.quantile(dist, CLAMP_QUANTILE)
-        low_val = -hi_val
-
-        U = U.clamp(low_val, hi_val)
-        Vh = Vh.clamp(low_val, hi_val)
-
-        if conv2d:
-          U = U.unsqueeze(2).unsqueeze(3)
-          Vh = Vh.unsqueeze(2).unsqueeze(3)
-
-        if device:
-          U = U.to(org_device)
-          Vh = Vh.to(org_device)
-
-        o_lora_sd[block_down_name + "." + "lora_down.weight"] = Vh.to(save_dtype).contiguous()
-        o_lora_sd[block_up_name + "." + "lora_up.weight"] = U.to(save_dtype).contiguous()
-        o_lora_sd[block_up_name + "." "alpha"] = torch.tensor(new_alpha).to(save_dtype)
+        new_alpha = param_dict['new_alpha']
+        o_lora_sd[block_down_name + "." + "lora_down.weight"] = param_dict["lora_down"].to(save_dtype).contiguous()
+        o_lora_sd[block_up_name + "." + "lora_up.weight"] = param_dict["lora_up"].to(save_dtype).contiguous()
+        o_lora_sd[block_up_name + "." "alpha"] = torch.tensor(param_dict['new_alpha']).to(save_dtype)

        block_down_name = None
        block_up_name = None
        lora_down_weight = None
        lora_up_weight = None
        weights_loaded = False
+        del param_dict

  if verbose:
    print(verbose_str)
+
+    print(f"Average Frobenius norm retention: {np.mean(fro_list):.2%} | std: {np.std(fro_list):0.3f}")
  print("resizing complete")
  return o_lora_sd, network_dim, new_alpha

@@ -151,6 +274,9 @@ def resize(args):
      return torch.bfloat16
    return None

+  if args.dynamic_method and not args.dynamic_param:
+    raise Exception("If using dynamic_method, then dynamic_param is required")
+
  merge_dtype = str_to_dtype('float')  # matmul method above only seems to work in float32
  save_dtype = str_to_dtype(args.save_precision)
  if save_dtype is None:
@@ -159,17 +285,23 @@ def resize(args):
  print("loading Model...")
  lora_sd, metadata = load_state_dict(args.model, merge_dtype)

-  print("resizing rank...")
-  state_dict, old_dim, new_alpha = resize_lora_model(lora_sd, args.new_rank, save_dtype, args.device, args.verbose)
+  print("Resizing Lora...")
+  state_dict, old_dim, new_alpha = resize_lora_model(lora_sd, args.new_rank, save_dtype, args.device, args.dynamic_method, args.dynamic_param, args.verbose)

  # update metadata
  if metadata is None:
    metadata = {}

  comment = metadata.get("ss_training_comment", "")
-  metadata["ss_training_comment"] = f"dimension is resized from {old_dim} to {args.new_rank}; {comment}"
-  metadata["ss_network_dim"] = str(args.new_rank)
-  metadata["ss_network_alpha"] = str(new_alpha)
+
+  if not args.dynamic_method:
+    metadata["ss_training_comment"] = f"dimension is resized from {old_dim} to {args.new_rank}; {comment}"
+    metadata["ss_network_dim"] = str(args.new_rank)
+    metadata["ss_network_alpha"] = str(new_alpha)
+  else:
+    metadata["ss_training_comment"] = f"Dynamic resize with {args.dynamic_method}: {args.dynamic_param} from {old_dim}; {comment}"
+    metadata["ss_network_dim"] = 'Dynamic'
+    metadata["ss_network_alpha"] = 'Dynamic'

  model_hash, legacy_hash = train_util.precalculate_safetensors_hashes(state_dict, metadata)
  metadata["sshs_model_hash"] = model_hash
@@ -193,6 +325,11 @@ if __name__ == '__main__':
  parser.add_argument("--device", type=str, default=None, help="device to use, cuda for GPU / 計算を行うデバイス、cuda でGPUを使う")
  parser.add_argument("--verbose", action="store_true", 
                      help="Display verbose resizing information / rank変更時の詳細情報を出力する")
+  parser.add_argument("--dynamic_method", type=str, default=None, choices=[None, "sv_ratio", "sv_fro", "sv_cumulative"],
+                      help="Specify dynamic resizing method, --new_rank is used as a hard limit for max rank")
+  parser.add_argument("--dynamic_param", type=float, default=None,
+                      help="Specify target for dynamic reduction")
+                                           

  args = parser.parse_args()
  resize(args)
--- a/networks/svd_merge_lora.py
+++ b/networks/svd_merge_lora.py
@@ -23,19 +23,20 @@ def load_state_dict(file_name, dtype):
  return sd


-def save_to_file(file_name, model, state_dict, dtype):
+def save_to_file(file_name, state_dict, dtype):
  if dtype is not None:
    for key in list(state_dict.keys()):
      if type(state_dict[key]) == torch.Tensor:
        state_dict[key] = state_dict[key].to(dtype)

  if os.path.splitext(file_name)[1] == '.safetensors':
-    save_file(model, file_name)
+    save_file(state_dict, file_name)
  else:
-    torch.save(model, file_name)
+    torch.save(state_dict, file_name)


-def merge_lora_models(models, ratios, new_rank, device,  merge_dtype):
+def merge_lora_models(models, ratios, new_rank, new_conv_rank, device, merge_dtype):
+  print(f"new rank: {new_rank}, new conv rank: {new_conv_rank}")
  merged_sd = {}
  for model, ratio in zip(models, ratios):
    print(f"loading: {model}")
@@ -58,11 +59,12 @@ def merge_lora_models(models, ratios, new_rank, device,  merge_dtype):
      in_dim = down_weight.size()[1]
      out_dim = up_weight.size()[0]
      conv2d = len(down_weight.size()) == 4
-      print(lora_module_name, network_dim, alpha, in_dim, out_dim)
+      kernel_size = None if not conv2d else down_weight.size()[2:4]
+      # print(lora_module_name, network_dim, alpha, in_dim, out_dim, kernel_size)

      # make original weight if not exist
      if lora_module_name not in merged_sd:
-        weight = torch.zeros((out_dim, in_dim, 1, 1) if conv2d else (out_dim, in_dim), dtype=merge_dtype)
+        weight = torch.zeros((out_dim, in_dim, *kernel_size) if conv2d else (out_dim, in_dim), dtype=merge_dtype)
        if device:
          weight = weight.to(device)
      else:
@@ -77,9 +79,12 @@ def merge_lora_models(models, ratios, new_rank, device,  merge_dtype):
      scale = (alpha / network_dim)
      if not conv2d:        # linear
        weight = weight + ratio * (up_weight @ down_weight) * scale
-      else:
+      elif kernel_size == (1, 1):
        weight = weight + ratio * (up_weight.squeeze(3).squeeze(2) @ down_weight.squeeze(3).squeeze(2)
                                   ).unsqueeze(2).unsqueeze(3) * scale
+      else:
+        conved = torch.nn.functional.conv2d(down_weight.permute(1, 0, 2, 3), up_weight).permute(1, 0, 2, 3)
+        weight = weight + ratio * conved * scale

      merged_sd[lora_module_name] = weight

@@ -89,16 +94,26 @@ def merge_lora_models(models, ratios, new_rank, device,  merge_dtype):
  with torch.no_grad():
    for lora_module_name, mat in tqdm(list(merged_sd.items())):
      conv2d = (len(mat.size()) == 4)
+      kernel_size = None if not conv2d else mat.size()[2:4]
+      conv2d_3x3 = conv2d and kernel_size != (1, 1)
+      out_dim, in_dim = mat.size()[0:2]
+
      if conv2d:
-        mat = mat.squeeze()
+        if conv2d_3x3:
+          mat = mat.flatten(start_dim=1)
+        else:
+          mat = mat.squeeze()
+
+      module_new_rank = new_conv_rank if conv2d_3x3 else new_rank
+      module_new_rank = min(module_new_rank, in_dim, out_dim)                           # LoRA rank cannot exceed the original dim

      U, S, Vh = torch.linalg.svd(mat)

-      U = U[:, :new_rank]
-      S = S[:new_rank]
+      U = U[:, :module_new_rank]
+      S = S[:module_new_rank]
      U = U @ torch.diag(S)

-      Vh = Vh[:new_rank, :]
+      Vh = Vh[:module_new_rank, :]

      dist = torch.cat([U.flatten(), Vh.flatten()])
      hi_val = torch.quantile(dist, CLAMP_QUANTILE)
@@ -107,16 +122,16 @@ def merge_lora_models(models, ratios, new_rank, device,  merge_dtype):
      U = U.clamp(low_val, hi_val)
      Vh = Vh.clamp(low_val, hi_val)

+      if conv2d:
+        U = U.reshape(out_dim, module_new_rank, 1, 1)
+        Vh = Vh.reshape(module_new_rank, in_dim, kernel_size[0], kernel_size[1])
+
      up_weight = U
      down_weight = Vh

-      if conv2d:
-        up_weight = up_weight.unsqueeze(2).unsqueeze(3)
-        down_weight = down_weight.unsqueeze(2).unsqueeze(3)
-
      merged_lora_sd[lora_module_name + '.lora_up.weight'] = up_weight.to("cpu").contiguous()
      merged_lora_sd[lora_module_name + '.lora_down.weight'] = down_weight.to("cpu").contiguous()
-      merged_lora_sd[lora_module_name + '.alpha'] = torch.tensor(new_rank)
+      merged_lora_sd[lora_module_name + '.alpha'] = torch.tensor(module_new_rank)

  return merged_lora_sd

@@ -138,10 +153,11 @@ def merge(args):
  if save_dtype is None:
    save_dtype = merge_dtype

-  state_dict = merge_lora_models(args.models, args.ratios, args.new_rank, args.device, merge_dtype)
+  new_conv_rank = args.new_conv_rank if args.new_conv_rank is not None else args.new_rank
+  state_dict = merge_lora_models(args.models, args.ratios, args.new_rank, new_conv_rank, args.device, merge_dtype)

  print(f"saving model to: {args.save_to}")
-  save_to_file(args.save_to, state_dict, state_dict, save_dtype)
+  save_to_file(args.save_to, state_dict, save_dtype)


 if __name__ == '__main__':
@@ -158,6 +174,8 @@ if __name__ == '__main__':
                      help="ratios for each model / それぞれのLoRAモデルの比率")
  parser.add_argument("--new_rank", type=int, default=4,
                      help="Specify rank of output LoRA / 出力するLoRAのrank (dim)")
+  parser.add_argument("--new_conv_rank", type=int, default=None,
+                      help="Specify rank of output LoRA for Conv2d 3x3, None for same as new_rank / 出力するConv2D 3x3 LoRAのrank (dim)、Noneでnew_rankと同じ")
  parser.add_argument("--device", type=str, default=None, help="device to use, cuda for GPU / 計算を行うデバイス、cuda でGPUを使う")

  args = parser.parse_args()
--- a/train_README-ja.md
+++ b/train_README-ja.md
@@ -1,4 +1,8 @@
-当リポジトリではモデルのfine tuning、DreamBooth、およびLoRAとTextual Inversionの学習をサポートします。この文書ではそれらに共通する、学習データの準備方法やスクリプトオプションについて説明します。
+__ドキュメント更新中のため記述に誤りがあるかもしれません。__
+
+# 学習について、共通編
+
+当リポジトリではモデルのfine tuning、DreamBooth、およびLoRAとTextual Inversionの学習をサポートします。この文書ではそれらに共通する、学習データの準備方法やオプション等について説明します。

 # 概要

@@ -8,15 +12,14 @@
 以下について説明します。

 1. 学習データの準備について（設定ファイルを用いる新形式）
-1. Aspect Ratio Bucketingについて
+1. 学習で使われる用語のごく簡単な解説
 1. 以前の指定形式（設定ファイルを用いずコマンドラインから指定）
+1. 学習途中のサンプル画像生成
+1. 各スクリプトで共通の、よく使われるオプション
 1. fine tuning 方式のメタデータ準備：キャプションニングなど

 1.だけ実行すればとりあえず学習は可能です（学習については各スクリプトのドキュメントを参照）。2.以降は必要に応じて参照してください。

-<!--
-1. 各スクリプトで共通のオプション
-->

 # 学習データの準備について

@@ -36,7 +39,7 @@

 1. fine tuning方式（正則化画像使用不可）

-    あらかじめキャプションをメタデータファイルにまとめます。タグとキャプションを分けて管理したり、学習を高速化するためlatentsを事前キャッシュしたりなどの機能をサポートします（いずれも別文書で説明しています）。
+    あらかじめキャプションをメタデータファイルにまとめます。タグとキャプションを分けて管理したり、学習を高速化するためlatentsを事前キャッシュしたりなどの機能をサポートします（いずれも別文書で説明しています）。（fine tuning方式という名前ですが fine tuning 以外でも使えます。）

 学習したいものと使用できる指定方法の組み合わせは以下の通りです。

@@ -124,7 +127,7 @@ batch_size = 4                              # バッチサイズ
  num_repeats = 1                           # 正則化画像の繰り返し回数、基本的には1でよい
 ```

-基本的には以下を場所のみ書き換えれば学習できます。
+基本的には以下の場所のみ書き換えれば学習できます。

 1. 学習解像度

@@ -132,7 +135,7 @@ batch_size = 4                              # バッチサイズ

 1. バッチサイズ

-    同時に何件のデータを学習するかを指定します。GPUのVRAMサイズ、学習解像度によって変わってきます。またfine tuning/DreamBooth/LoRA等でも変わってきますので、詳しくは各スクリプトの説明をご覧ください。
+    同時に何件のデータを学習するかを指定します。GPUのVRAMサイズ、学習解像度によって変わってきます。詳しくは後述します。またfine tuning/DreamBooth/LoRA等でも変わってきますので各スクリプトの説明もご覧ください。

 1. フォルダ指定

@@ -248,7 +251,45 @@ batch_size = 4                                      # バッチサイズ

 それぞれのドキュメントを参考に学習を行ってください。

-# Aspect Ratio Bucketing について
+# 学習で使われる用語のごく簡単な解説
+
+細かいことは省略していますし私も完全には理解していないため、詳しくは各自お調べください。
+
+## fine tuning（ファインチューニング）
+
+モデルを学習して微調整することを指します。使われ方によって意味が異なってきますが、狭義のfine tuningはStable Diffusionの場合、モデルを画像とキャプションで学習することです。DreamBoothは狭義のfine tuningのひとつの特殊なやり方と言えます。広義のfine tuningは、LoRAやTextual Inversion、Hypernetworksなどを含み、モデルを学習することすべてを含みます。
+
+## ステップ
+
+ざっくりいうと学習データで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 +301,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 +339,7 @@ classがひとつで対象が複数の場合、正則化画像フォルダはひ
 - reg_girls
  - 1_1girl

-### DreamBoothでキャプションを使う
-
-学習用画像、正則化画像のフォルダに、画像と同じファイル名で、拡張子.caption（オプションで変えられます）のファイルを置くと、そのファイルからキャプションを読み込みプロンプトとして学習します。
-
-※それらの画像の学習に、フォルダ名（identifier class）は使用されなくなります。
-
-キャプションファイルの拡張子はデフォルトで.captionです。学習スクリプトの `--caption_extension` オプションで変更できます。`--shuffle_caption` オプションで学習時のキャプションについて、カンマ区切りの各部分をシャッフルしながら学習します。
-
-## step 2. 正則化画像の準備
+### step 2. 正則化画像の準備

 正則化画像を使う場合の手順です。

@@ -313,16 +350,296 @@ 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、またはそこからのfine tuningモデルを使う場合（推論時に `v2-inference.yaml` を使うように指示されているモデルの場合）は `--v2` オプションを、stable-diffusion-2や768-v-ema.ckpt、およびそれらのfine tuningモデルを使う場合（推論時に `v2-inference-v.yaml` を使うモデルの場合）は `--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と同様に、モデルの学習状態と異なる長さで学習するには、ある程度の教師データ枚数、長めの学習時間が必要になると思われます。
+
+- `--persistent_data_loader_workers`
+
+    Windows環境で指定するとエポック間の待ち時間が大幅に短縮されます。
+
+- `--max_data_loader_n_workers`
+
+    データ読み込みのプロセス数を指定します。プロセス数が多いとデータ読み込みが速くなりGPUを効率的に利用できますが、メインメモリを消費します。デフォルトは「`8` または `CPU同時実行スレッド数-1` の小さいほう」なので、メインメモリに余裕がない場合や、GPU使用率が90%程度以上なら、それらの数値を見ながら `2` または `1` 程度まで下げてください。
+
+- `--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 : [torch.optim.AdamW](https://pytorch.org/docs/stable/generated/torch.optim.AdamW.html)
+    - 過去のバージョンのオプション未指定時と同じ
+    - AdamW8bit : 引数は同上
+    - 過去のバージョンの--use_8bit_adam指定時と同じ
+    - Lion : https://github.com/lucidrains/lion-pytorch
+    - 過去のバージョンの--use_lion_optimizer指定時と同じ
+    - SGDNesterov : [torch.optim.SGD](https://pytorch.org/docs/stable/generated/torch.optim.SGD.html), nesterov=True
+    - SGDNesterov8bit : 引数は同上
+    - DAdaptation : https://github.com/facebookresearch/dadaptation
+    - AdaFactor : [Transformers AdaFactor](https://huggingface.co/docs/transformers/main_classes/optimizer_schedules)
+    - 任意のオプティマイザ
+
+- `--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 です。
+
+    詳細については各自お調べください。
+
+### オプティマイザの指定について
+
+オプティマイザのオプション引数は--optimizer_argsオプションで指定してください。key=valueの形式で、複数の値が指定できます。また、valueはカンマ区切りで複数の値が指定できます。たとえばAdamWオプティマイザに引数を指定する場合は、``--optimizer_args weight_decay=0.01 betas=.9,.999``のようになります。
+
+オプション引数を指定する場合は、それぞれのオプティマイザの仕様をご確認ください。
+
+一部のオプティマイザでは必須の引数があり、省略すると自動的に追加されます（SGDNesterovのmomentumなど）。コンソールの出力を確認してください。
+
+D-Adaptationオプティマイザは学習率を自動調整します。学習率のオプションに指定した値は学習率そのものではなくD-Adaptationが決定した学習率の適用率になりますので、通常は1.0を指定してください。Text EncoderにU-Netの半分の学習率を指定したい場合は、``--text_encoder_lr=0.5 --unet_lr=1.0``と指定します。
+
+AdaFactorオプティマイザはrelative_step=Trueを指定すると学習率を自動調整できます（省略時はデフォルトで追加されます）。自動調整する場合は学習率のスケジューラにはadafactor_schedulerが強制的に使用されます。またscale_parameterとwarmup_initを指定するとよいようです。
+
+自動調整する場合のオプション指定はたとえば ``--optimizer_args "relative_step=True" "scale_parameter=True" "warmup_init=True"`` のようになります。
+
+学習率を自動調整しない場合はオプション引数 ``relative_step=False`` を追加してください。その場合、学習率のスケジューラにはconstant_with_warmupが、また勾配のclip normをしないことが推奨されているようです。そのため引数は ``--optimizer_type=adafactor --optimizer_args "relative_step=False" --lr_scheduler="constant_with_warmup" --max_grad_norm=0.0`` のようになります。
+
+### 任意のオプティマイザを使う
+
+``torch.optim`` のオプティマイザを使う場合にはクラス名のみを（``--optimizer_type=RMSprop``など）、他のモジュールのオプティマイザを使う時は「モジュール名.クラス名」を指定してください（``--optimizer_type=bitsandbytes.optim.lamb.LAMB``など）。
+
+（内部でimportlibしているだけで動作は未確認です。必要ならパッケージをインストールしてください。）
+
+
+<!-- 
+## 任意サイズの画像での学習 --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で確認）。精度はかなり落ちますし、途中で学習失敗する確率も高くなります。
+学習率やステップ数の設定もシビアなようです。それらを認識したうえで自己責任でお使いください。
+
 -->

 # メタデータファイルの作成
--- a/train_db_README-ja.md
+++ b/train_db_README-ja.md
@@ -1,75 +1,104 @@
-DreamBoothのガイドです。LoRA等の追加ネットワークの学習にも同じ手順を使います。
+DreamBoothのガイドです。
+
+[学習についての共通ドキュメント](./train_README-ja.md) もあわせてご覧ください。

 # 概要

+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を参照し、環境整備を行ってください。

-このリポジトリの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のガイドを参照してください。
+### よく使われるオプションについて
+
+以下の場合には [学習の共通ドキュメント](./train_README-ja.md) の「よく使われるオプション」を参照してください。
+
+- 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 
@@ -77,6 +106,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 +119,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 +163,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オプションで学習時のキャプションについて、カンマ区切りの各部分をシャッフルしながら学習します。

+-->
--- a/train_network.py
+++ b/train_network.py
@@ -106,6 +106,7 @@ def train(args):
  # acceleratorを準備する
  print("prepare accelerator")
  accelerator, unwrap_model = train_util.prepare_accelerator(args)
+  is_main_process = accelerator.is_main_process

  # mixed precisionに対応した型を用意しておき適宜castする
  weight_dtype, save_dtype = train_util.prepare_dtype(args)
@@ -134,6 +135,8 @@ def train(args):
    gc.collect()

  # prepare network
+  import sys
+  sys.path.append(os.path.dirname(__file__))
  print("import network module:", args.network_module)
  network_module = importlib.import_module(args.network_module)

@@ -175,12 +178,13 @@ def train(args):

  # 学習ステップ数を計算する
  if args.max_train_epochs is not None:
-    args.max_train_steps = args.max_train_epochs * len(train_dataloader)
-    print(f"override steps. steps for {args.max_train_epochs} epochs is / 指定エポックまでのステップ数: {args.max_train_steps}")
+    args.max_train_steps = args.max_train_epochs * math.ceil(len(train_dataloader) / accelerator.num_processes)
+    if is_main_process:
+      print(f"override steps. steps for {args.max_train_epochs} epochs is / 指定エポックまでのステップ数: {args.max_train_steps}")

  # lr schedulerを用意する
  lr_scheduler = train_util.get_scheduler_fix(args.lr_scheduler, optimizer, num_warmup_steps=args.lr_warmup_steps,
-                                              num_training_steps=args.max_train_steps * args.gradient_accumulation_steps,
+                                              num_training_steps=args.max_train_steps * accelerator.num_processes * args.gradient_accumulation_steps,
                                              num_cycles=args.lr_scheduler_num_cycles, power=args.lr_scheduler_power)

  # 実験的機能：勾配も含めたfp16学習を行う　モデル全体をfp16にする
@@ -251,15 +255,17 @@ def train(args):
  # 学習する
  # TODO: find a way to handle total batch size when there are multiple datasets
  total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
-  print("running training / 学習開始")
-  print(f"  num train images * repeats / 学習画像の数×繰り返し回数: {train_dataset_group.num_train_images}")
-  print(f"  num reg images / 正則化画像の数: {train_dataset_group.num_reg_images}")
-  print(f"  num batches per epoch / 1epochのバッチ数: {len(train_dataloader)}")
-  print(f"  num epochs / epoch数: {num_train_epochs}")
-  print(f"  batch size per device / バッチサイズ: {', '.join([str(d.batch_size) for d in train_dataset_group.datasets])}")
-  # print(f"  total train batch size (with parallel & distributed & accumulation) / 総バッチサイズ（並列学習、勾配合計含む）: {total_batch_size}")
-  print(f"  gradient accumulation steps / 勾配を合計するステップ数 = {args.gradient_accumulation_steps}")
-  print(f"  total optimization steps / 学習ステップ数: {args.max_train_steps}")
+  
+  if is_main_process:
+    print("running training / 学習開始")
+    print(f"  num train images * repeats / 学習画像の数×繰り返し回数: {train_dataset_group.num_train_images}")
+    print(f"  num reg images / 正則化画像の数: {train_dataset_group.num_reg_images}")
+    print(f"  num batches per epoch / 1epochのバッチ数: {len(train_dataloader)}")
+    print(f"  num epochs / epoch数: {num_train_epochs}")
+    print(f"  batch size per device / バッチサイズ: {', '.join([str(d.batch_size) for d in train_dataset_group.datasets])}")
+    # print(f"  total train batch size (with parallel & distributed & accumulation) / 総バッチサイズ（並列学習、勾配合計含む）: {total_batch_size}")
+    print(f"  gradient accumulation steps / 勾配を合計するステップ数 = {args.gradient_accumulation_steps}")
+    print(f"  total optimization steps / 学習ステップ数: {args.max_train_steps}")

  # TODO refactor metadata creation and move to util
  metadata = {
@@ -427,10 +433,13 @@ def train(args):
        "ss_bucket_info": json.dumps(dataset.bucket_info),
    })

-  # uncomment if another network is added
-  # for key, value in net_kwargs.items():
-  #   metadata["ss_arg_" + key] = value
+  # add extra args
+  if args.network_args:
+    metadata["ss_network_args"] = json.dumps(net_kwargs)
+    # for key, value in net_kwargs.items():
+    #   metadata["ss_arg_" + key] = value

+  # model name and hash
  if args.pretrained_model_name_or_path is not None:
    sd_model_name = args.pretrained_model_name_or_path
    if os.path.exists(sd_model_name):
@@ -449,6 +458,13 @@ def train(args):

  metadata = {k: str(v) for k, v in metadata.items()}

+  # make minimum metadata for filtering
+  minimum_keys = ["ss_network_module", "ss_network_dim", "ss_network_alpha", "ss_network_args"]
+  minimum_metadata = {}
+  for key in minimum_keys:
+    if key in metadata:
+      minimum_metadata[key] = metadata[key]
+
  progress_bar = tqdm(range(args.max_train_steps), smoothing=0, disable=not accelerator.is_local_main_process, desc="steps")
  global_step = 0

@@ -461,7 +477,8 @@ def train(args):
  loss_list = []
  loss_total = 0.0
  for epoch in range(num_train_epochs):
-    print(f"epoch {epoch+1}/{num_train_epochs}")
+    if is_main_process:
+      print(f"epoch {epoch+1}/{num_train_epochs}")
    train_dataset_group.set_current_epoch(epoch + 1)

    metadata["ss_epoch"] = str(epoch+1)
@@ -564,7 +581,7 @@ def train(args):
        ckpt_file = os.path.join(args.output_dir, ckpt_name)
        metadata["ss_training_finished_at"] = str(time.time())
        print(f"saving checkpoint: {ckpt_file}")
-        unwrap_model(network).save_weights(ckpt_file, save_dtype, None if args.no_metadata else metadata)
+        unwrap_model(network).save_weights(ckpt_file, save_dtype, minimum_metadata if args.no_metadata else metadata)

      def remove_old_func(old_epoch_no):
        old_ckpt_name = train_util.EPOCH_FILE_NAME.format(model_name, old_epoch_no) + '.' + args.save_model_as
@@ -573,9 +590,10 @@ def train(args):
          print(f"removing old checkpoint: {old_ckpt_file}")
          os.remove(old_ckpt_file)

-      saving = train_util.save_on_epoch_end(args, save_func, remove_old_func, epoch + 1, num_train_epochs)
-      if saving and args.save_state:
-        train_util.save_state_on_epoch_end(args, accelerator, model_name, epoch + 1)
+      if is_main_process:
+        saving = train_util.save_on_epoch_end(args, save_func, remove_old_func, epoch + 1, num_train_epochs)
+        if saving and args.save_state:
+          train_util.save_state_on_epoch_end(args, accelerator, model_name, epoch + 1)

    train_util.sample_images(accelerator, args, epoch + 1, global_step, accelerator.device, vae, tokenizer, text_encoder, unet)

@@ -584,7 +602,6 @@ def train(args):
  metadata["ss_epoch"] = str(num_train_epochs)
  metadata["ss_training_finished_at"] = str(time.time())

-  is_main_process = accelerator.is_main_process
  if is_main_process:
    network = unwrap_model(network)

@@ -603,7 +620,7 @@ def train(args):
    ckpt_file = os.path.join(args.output_dir, ckpt_name)

    print(f"save trained model to {ckpt_file}")
-    network.save_weights(ckpt_file, save_dtype, None if args.no_metadata else metadata)
+    network.save_weights(ckpt_file, save_dtype, minimum_metadata if args.no_metadata else metadata)
    print("model saved.")


--- a/train_network_README-ja.md
+++ b/train_network_README-ja.md
@@ -1,118 +1,103 @@
-## LoRAの学習について
+# LoRAの学習について

 [LoRA: Low-Rank Adaptation of Large Language Models](https://arxiv.org/abs/2106.09685)（arxiv）、[LoRA](https://github.com/microsoft/LoRA)（github）をStable Diffusionに適用したものです。

 [cloneofsimo氏のリポジトリ](https://github.com/cloneofsimo/lora)を大いに参考にさせていただきました。ありがとうございます。

+通常のLoRAは Linear およぴカーネルサイズ 1x1 の Conv2d にのみ適用されますが、カーネルサイズ 3x3 のConv2dに適用を拡大することもできます。
+
+Conv2d 3x3への拡大は [cloneofsimo氏](https://github.com/cloneofsimo/lora) が最初にリリースし、KohakuBlueleaf氏が [LoCon](https://github.com/KohakuBlueleaf/LoCon) でその有効性を明らかにしたものです。KohakuBlueleaf氏に深く感謝します。
+
 8GB VRAMでもぎりぎり動作するようです。

+[学習についての共通ドキュメント](./train_README-ja.md) もあわせてご覧ください。
+
 ## 学習したモデルに関する注意

 cloneofsimo氏のリポジトリ、およびd8ahazard氏の[Dreambooth Extension for Stable-Diffusion-WebUI](https://github.com/d8ahazard/sd_dreambooth_extension)とは、現時点では互換性がありません。いくつかの機能拡張を行っているためです（後述）。

 WebUI等で画像生成する場合には、学習したLoRAのモデルを学習元のStable Diffusionのモデルにこのリポジトリ内のスクリプトであらかじめマージしておくか、こちらの[WebUI用extension](https://github.com/kohya-ss/sd-webui-additional-networks)を使ってください。

-## 学習方法
+# 学習の手順

-train_network.pyを用います。
+あらかじめこのリポジトリのREADMEを参照し、環境整備を行ってください。

-DreamBoothの手法（identifier（sksなど）とclass、オプションで正則化画像を用いる）と、キャプションを用いるfine tuningの手法の両方で学習できます。
+## データの準備

-どちらの方法も既存のスクリプトとほぼ同じ方法で学習できます。異なる点については後述します。
+[学習データの準備について](./train_README-ja.md) を参照してください。

-### DreamBoothの手法を用いる場合

-[DreamBoothのガイド](./train_db_README-ja.md) を参照してデータを用意してください。
+## 学習の実行

-学習するとき、train_db.pyの代わりにtrain_network.pyを指定してください。そして「LoRAの学習のためのオプション」にあるようにLoRA関連のオプション（``network_dim``や``network_alpha``など）を追加してください。
+`train_network.py`を用います。

-ほぼすべてのオプション（Stable Diffusionのモデル保存関係を除く）が使えますが、stop_text_encoder_trainingはサポートしていません。
-
-### キャプションを用いる場合
-
-[fine-tuningのガイド](./fine_tune_README_ja.md) を参照し、各手順を実行してください。
-
-学習するとき、fine_tune.pyの代わりにtrain_network.pyを指定してください。ほぼすべてのオプション（モデル保存関係を除く）がそのまま使えます。そして「LoRAの学習のためのオプション」にあるようにLoRA関連のオプション（``network_dim``や``network_alpha``など）を追加してください。
-
-なお「latentsの事前取得」は行わなくても動作します。VAEから学習時（またはキャッシュ時）にlatentを取得するため学習速度は遅くなりますが、代わりにcolor_augが使えるようになります。
-
-### LoRAの学習のためのオプション
-
-train_network.pyでは--network_moduleオプションに、学習対象のモジュール名を指定します。LoRAに対応するのはnetwork.loraとなりますので、それを指定してください。
+`train_network.py`では `--network_module` オプションに、学習対象のモジュール名を指定します。LoRAに対応するのはnetwork.loraとなりますので、それを指定してください。

 なお学習率は通常のDreamBoothやfine tuningよりも高めの、1e-4程度を指定するとよいようです。

-以下はコマンドラインの例です（DreamBooth手法）。
+以下はコマンドラインの例です。

 ```
 accelerate launch --num_cpu_threads_per_process 1 train_network.py 
-    --pretrained_model_name_or_path=..\models\model.ckpt 
-    --train_data_dir=..\data\db\char1 --output_dir=..\lora_train1 
-    --reg_data_dir=..\data\db\reg1 --prior_loss_weight=1.0 
-    --resolution=448,640 --train_batch_size=1 --learning_rate=1e-4 
-    --max_train_steps=400 --optimizer_type=AdamW8bit --xformers --mixed_precision=fp16 
-    --save_every_n_epochs=1 --save_model_as=safetensors --clip_skip=2 --seed=42 --color_aug 
+    --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=400 
+    --learning_rate=1e-4 
+    --optimizer_type="AdamW8bit" 
+    --xformers 
+    --mixed_precision="fp16" 
+    --cache_latents 
+    --gradient_checkpointing
+    --save_every_n_epochs=1 
    --network_module=networks.lora
 ```

-（2023/2/22:オプティマイザの指定方法が変わりました。[こちら](#オプティマイザの指定について）をご覧ください。）
-
--output_dirオプションで指定したフォルダに、LoRAのモデルが保存されます。
+`--output_dir` オプションで指定したフォルダに、LoRAのモデルが保存されます。他のオプション、オプティマイザ等については [学習の共通ドキュメント](./train_README-ja.md) の「よく使われるオプション」も参照してください。

 その他、以下のオプションが指定できます。

-* --network_dim
+* `--network_dim`
  * LoRAのRANKを指定します（``--networkdim=4``など）。省略時は4になります。数が多いほど表現力は増しますが、学習に必要なメモリ、時間は増えます。また闇雲に増やしても良くないようです。
-* --network_alpha
+* `--network_alpha`
  *  アンダーフローを防ぎ安定して学習するための ``alpha`` 値を指定します。デフォルトは1です。``network_dim``と同じ値を指定すると以前のバージョンと同じ動作になります。
-* --network_weights
+* `--persistent_data_loader_workers`
+  * Windows環境で指定するとエポック間の待ち時間が大幅に短縮されます。
+* `--max_data_loader_n_workers`
+  * データ読み込みのプロセス数を指定します。プロセス数が多いとデータ読み込みが速くなりGPUを効率的に利用できますが、メインメモリを消費します。デフォルトは「`8` または `CPU同時実行スレッド数-1` の小さいほう」なので、メインメモリに余裕がない場合や、GPU使用率が90%程度以上なら、それらの数値を見ながら `2` または `1` 程度まで下げてください。
+* `--network_weights`
  * 学習前に学習済みのLoRAの重みを読み込み、そこから追加で学習します。
-* --network_train_unet_only
+* `--network_train_unet_only`
  * U-Netに関連するLoRAモジュールのみ有効とします。fine tuning的な学習で指定するとよいかもしれません。
-* --network_train_text_encoder_only
+* `--network_train_text_encoder_only`
  * Text Encoderに関連するLoRAモジュールのみ有効とします。Textual Inversion的な効果が期待できるかもしれません。
-* --unet_lr
+* `--unet_lr`
  * U-Netに関連するLoRAモジュールに、通常の学習率（--learning_rateオプションで指定）とは異なる学習率を使う時に指定します。
-* --text_encoder_lr
+* `--text_encoder_lr`
  * Text Encoderに関連するLoRAモジュールに、通常の学習率（--learning_rateオプションで指定）とは異なる学習率を使う時に指定します。Text Encoderのほうを若干低めの学習率（5e-5など）にしたほうが良い、という話もあるようです。
+* `--network_args`
+  * 複数の引数を指定できます。後述します。

--network_train_unet_onlyと--network_train_text_encoder_onlyの両方とも未指定時（デフォルト）はText EncoderとU-Netの両方のLoRAモジュールを有効にします。
+`--network_train_unet_only` と `--network_train_text_encoder_only` の両方とも未指定時（デフォルト）はText EncoderとU-Netの両方のLoRAモジュールを有効にします。

-## オプティマイザの指定について
+## LoRA を Conv2d に拡大して適用する

--optimizer_type オプションでオプティマイザの種類を指定します。以下が指定できます。
+通常のLoRAは Linear およぴカーネルサイズ 1x1 の Conv2d にのみ適用されますが、カーネルサイズ 3x3 のConv2dに適用を拡大することもできます。

- AdamW : [torch.optim.AdamW](https://pytorch.org/docs/stable/generated/torch.optim.AdamW.html)
-  - 過去のバージョンのオプション未指定時と同じ
- AdamW8bit : 引数は同上
-  - 過去のバージョンの--use_8bit_adam指定時と同じ
- Lion : https://github.com/lucidrains/lion-pytorch
-  - 過去のバージョンの--use_lion_optimizer指定時と同じ
- SGDNesterov : [torch.optim.SGD](https://pytorch.org/docs/stable/generated/torch.optim.SGD.html), nesterov=True
- SGDNesterov8bit : 引数は同上
- DAdaptation : https://github.com/facebookresearch/dadaptation
- AdaFactor : [Transformers AdaFactor](https://huggingface.co/docs/transformers/main_classes/optimizer_schedules)
- 任意のオプティマイザ
+`--network_args` に以下のように指定してください。`conv_dim` で Conv2d (3x3) の rank を、`conv_alpha` で alpha を指定してください。

-オプティマイザのオプション引数は--optimizer_argsオプションで指定してください。key=valueの形式で、複数の値が指定できます。また、valueはカンマ区切りで複数の値が指定できます。たとえばAdamWオプティマイザに引数を指定する場合は、``--optimizer_args weight_decay=0.01 betas=.9,.999``のようになります。
+```
+--network_args "conv_dim=1" "conv_alpha=1"
+```

-オプション引数を指定する場合は、それぞれのオプティマイザの仕様をご確認ください。
+以下のように alpha 省略時は1になります。

-一部のオプティマイザでは必須の引数があり、省略すると自動的に追加されます（SGDNesterovのmomentumなど）。コンソールの出力を確認してください。
-
-D-Adaptationオプティマイザは学習率を自動調整します。学習率のオプションに指定した値は学習率そのものではなくD-Adaptationが決定した学習率の適用率になりますので、通常は1.0を指定してください。Text EncoderにU-Netの半分の学習率を指定したい場合は、``--text_encoder_lr=0.5 --unet_lr=1.0``と指定します。
-
-AdaFactorオプティマイザはrelative_step=Trueを指定すると学習率を自動調整できます（省略時はデフォルトで追加されます）。自動調整する場合は学習率のスケジューラにはadafactor_schedulerが強制的に使用されます。またscale_parameterとwarmup_initを指定するとよいようです。
-
-自動調整する場合のオプション指定はたとえば ``--optimizer_args "relative_step=True" "scale_parameter=True" "warmup_init=True"`` のようになります。
-
-学習率を自動調整しない場合はオプション引数 ``relative_step=False`` を追加してください。その場合、学習率のスケジューラにはconstant_with_warmupが、また勾配のclip normをしないことが推奨されているようです。そのため引数は ``--optimizer_type=adafactor --optimizer_args "relative_step=False" --lr_scheduler="constant_with_warmup" --max_grad_norm=0.0`` のようになります。
-
-### 任意のオプティマイザを使う
-
-``torch.optim`` のオプティマイザを使う場合にはクラス名のみを（``--optimizer_type=RMSprop``など）、他のモジュールのオプティマイザを使う時は「モジュール名.クラス名」を指定してください（``--optimizer_type=bitsandbytes.optim.lamb.LAMB``など）。
-
-（内部でimportlibしているだけで動作は未確認です。必要ならパッケージをインストールしてください。）
+```
+--network_args "conv_dim=1"
+```

 ## マージスクリプトについて

@@ -176,6 +161,27 @@ v1で学習したLoRAとv2で学習したLoRA、rank（次元数）や``alpha``
 * save_precision
  * モデル保存時の精度をfloat、fp16、bf16から指定できます。省略時はprecisionと同じ精度になります。

+
+## 複数のrankが異なるLoRAのモデルをマージする
+
+複数のLoRAをひとつのLoRAで近似します（完全な再現はできません）。`svd_merge_lora.py`を用います。たとえば以下のようなコマンドラインになります。
+
+```
+python networks\svd_merge_lora.py 
+    --save_to ..\lora_train1\model-char1-style1-merged.safetensors 
+    --models ..\lora_train1\last.safetensors ..\lora_train2\last.safetensors 
+    --ratios 0.6 0.4 --new_rank 32 --device cuda
+```
+
+`merge_lora.py` と主なオプションは同一です。以下のオプションが追加されています。
+
+- `--new_rank`
+  - 作成するLoRAのrankを指定します。
+- `--new_conv_rank`
+  - 作成する Conv2d 3x3 LoRA の rank を指定します。省略時は `new_rank` と同じになります。
+- `--device`
+  - `--device cuda`としてcudaを指定すると計算をGPU上で行います。処理が速くなります。
+
 ## 当リポジトリ内の画像生成スクリプトで生成する

 gen_img_diffusers.pyに、--network_module、--network_weightsの各オプションを追加してください。意味は学習時と同様です。
@@ -209,12 +215,14 @@ Text Encoderが二つのモデルで同じ場合にはLoRAはU-NetのみのLoRA

 ### その他のオプション

- --v2
+- `--v2`
  - v2.xのStable Diffusionモデルを使う場合に指定してください。
- --device
+- `--device`
  - ``--device cuda``としてcudaを指定すると計算をGPU上で行います。処理が速くなります（CPUでもそこまで遅くないため、せいぜい倍～数倍程度のようです）。
- --save_precision
+- `--save_precision`
  - LoRAの保存形式を"float", "fp16", "bf16"から指定します。省略時はfloatになります。
+- `--conv_dim`
+  - 指定するとLoRAの適用範囲を Conv2d 3x3 へ拡大します。Conv2d 3x3 の rank を指定します。

 ## 画像リサイズスクリプト

@@ -252,7 +260,7 @@ python tools\resize_images_to_resolution.py --max_resolution 512x512,384x384,256

 ### cloneofsimo氏のリポジトリとの違い

-12/25時点では、当リポジトリはLoRAの適用個所をText EncoderのMLP、U-NetのFFN、Transformerのin/out projectionに拡大し、表現力が増しています。ただその代わりメモリ使用量は増え、8GBぎりぎりになりました。
+2022/12/25時点では、当リポジトリはLoRAの適用個所をText EncoderのMLP、U-NetのFFN、Transformerのin/out projectionに拡大し、表現力が増しています。ただその代わりメモリ使用量は増え、8GBぎりぎりになりました。

 またモジュール入れ替え機構は全く異なります。

--- a/train_textual_inversion.py
+++ b/train_textual_inversion.py
@@ -181,6 +181,11 @@ def train(args):
    for tmpl in templates:
      captions.append(tmpl.format(replace_to))
    train_dataset_group.add_replacement("", captions)
+
+    if args.num_vectors_per_token > 1:
+      prompt_replacement = (args.token_string, replace_to)
+    else:
+      prompt_replacement = None
  else:
    if args.num_vectors_per_token > 1:
      replace_to = " ".join(token_strings)
--- a/train_ti_README-ja.md
+++ b/train_ti_README-ja.md
@@ -1,32 +1,41 @@
-## Textual Inversionの学習について
+[Textual Inversion](https://textual-inversion.github.io/) の学習についての説明です。

-[Textual Inversion](https://textual-inversion.github.io/)です。実装に当たっては https://github.com/huggingface/diffusers/tree/main/examples/textual_inversion を大いに参考にしました。
+[学習についての共通ドキュメント](./train_README-ja.md) もあわせてご覧ください。

-学習したモデルはWeb UIでもそのまま使えます。
+実装に当たっては https://github.com/huggingface/diffusers/tree/main/examples/textual_inversion を大いに参考にしました。

-なお恐らくSD2.xにも対応していますが現時点では未テストです。
+学習したモデルはWeb UIでもそのまま使えます。なお恐らくSD2.xにも対応していますが現時点では未テストです。

-## 学習方法
+# 学習の手順

-``train_textual_inversion.py`` を用います。
+あらかじめこのリポジトリのREADMEを参照し、環境整備を行ってください。

-データの準備については ``train_network.py`` と全く同じですので、[そちらのドキュメント](./train_network_README-ja.md)を参照してください。
+## データの準備

-## オプション
+[学習データの準備について](./train_README-ja.md) を参照してください。

-以下はコマンドラインの例です（DreamBooth手法）。
+## 学習の実行
+
+``train_textual_inversion.py`` を用います。以下はコマンドラインの例です（DreamBooth手法）。

 ```
 accelerate launch --num_cpu_threads_per_process 1 train_textual_inversion.py 
-    --pretrained_model_name_or_path=..\models\model.ckpt 
-    --train_data_dir=..\data\db\char1 --output_dir=..\ti_train1 
-    --resolution=448,640 --train_batch_size=1 --learning_rate=1e-4 
-    --max_train_steps=400 --use_8bit_adam --xformers --mixed_precision=fp16 
-    --save_every_n_epochs=1 --save_model_as=safetensors --clip_skip=2 --seed=42 --color_aug 
+    --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
    --token_string=mychar4 --init_word=cute --num_vectors_per_token=4
 ```

-``--token_string`` に学習時のトークン文字列を指定します。__学習時のプロンプトは、この文字列を含むようにしてください（token_stringがmychar4なら、``mychar4 1girl`` など）__。プロンプトのこの文字列の部分が、Textual Inversionの新しいtokenに置換されて学習されます。
+``--token_string`` に学習時のトークン文字列を指定します。__学習時のプロンプトは、この文字列を含むようにしてください（token_stringがmychar4なら、``mychar4 1girl`` など）__。プロンプトのこの文字列の部分が、Textual Inversionの新しいtokenに置換されて学習されます。DreamBooth, class+identifier形式のデータセットとして、`token_string` をトークン文字列にするのが最も簡単で確実です。

 プロンプトにトークン文字列が含まれているかどうかは、``--debug_dataset`` で置換後のtoken idが表示されますので、以下のように ``49408`` 以降のtokenが存在するかどうかで確認できます。

@@ -47,14 +56,47 @@ tokenizerがすでに持っている単語（一般的な単語）は使用で

 ``--num_vectors_per_token`` にいくつのトークンをこの学習で使うかを指定します。多いほうが表現力が増しますが、その分多くのトークンを消費します。たとえばnum_vectors_per_token=8の場合、指定したトークン文字列は（一般的なプロンプトの77トークン制限のうち）8トークンを消費します。

+以上がTextual Inversionのための主なオプションです。以降は他の学習スクリプトと同様です。

-その他、以下のオプションが指定できます。
+`num_cpu_threads_per_process` には通常は1を指定するとよいようです。

-* --weights
+`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` としてください。
+
+学習させるステップ数 `max_train_steps` を10000とします。学習率 `learning_rate` はここでは5e-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を使用します（速度は遅くなります）。
+
+ある程度メモリがある場合は、`.toml` ファイルを編集してバッチサイズをたとえば `8` くらいに増やしてください（高速化と精度向上の可能性があります）。
+
+### よく使われるオプションについて
+
+以下の場合にはオプションに関するドキュメントを参照してください。
+
+- Stable Diffusion 2.xまたはそこからの派生モデルを学習する
+- clip skipを2以上を前提としたモデルを学習する
+- 75トークンを超えたキャプションで学習する
+
+### Textual Inversionでのバッチサイズについて
+
+モデル全体を学習するDreamBoothやfine tuningに比べてメモリ使用量が少ないため、バッチサイズは大きめにできます。
+
+# Textual Inversionのその他の主なオプション
+
+すべてのオプションについては別文書を参照してください。
+
+* `--weights`
  * 学習前に学習済みのembeddingsを読み込み、そこから追加で学習します。
-* --use_object_template
+* `--use_object_template`
  * キャプションではなく既定の物体用テンプレート文字列（``a photo of a {}``など）で学習します。公式実装と同じになります。キャプションは無視されます。
-* --use_style_template
+* `--use_style_template`
  * キャプションではなく既定のスタイル用テンプレート文字列で学習します（``a painting in the style of {}``など）。公式実装と同じになります。キャプションは無視されます。

 ## 当リポジトリ内の画像生成スクリプトで生成する
Author	SHA1	Message	Date
Kohya S	b80431de30	Merge pull request #278 from kohya-ss/dev Dev	2023-03-10 22:05:36 +09:00
Kohya S	b177460807	restore comment	2023-03-10 22:02:17 +09:00
Kohya S	c78c51c78f	update documents	2023-03-10 21:59:25 +09:00
Kohya S	2652c9a66c	Merge pull request #276 from mio2333/main Append sys path for import_module	2023-03-10 21:43:32 +09:00
Kohya S	618592c52b	npz check to use subset, add dadap warn close #274	2023-03-10 21:31:59 +09:00
Kohya S	75d1883da6	fix LoRA rank is limited to target dim	2023-03-10 21:12:15 +09:00
Kohya S	4ad8e75291	fix to work with dim>320	2023-03-10 21:10:22 +09:00
Kohya S	e355b5e1d3	Merge pull request #269 from rvhfxb/patch-2 Allow to delete images after getting latents	2023-03-10 20:56:11 +09:00
mio	68cd874bb6	Append sys path for import_module This will be better if we run the scripts we do not run the training script from the current directory. This is reasonable as some other projects will use this as a subfolder, such as https://github.com/ddPn08/kohya-sd-scripts-webui. I can not run the script without adding this.	2023-03-10 18:29:34 +08:00
Kohya S	00a9d734d9	Merge pull request #247 from ddPn08/dev fix for multi gpu training	2023-03-10 13:01:52 +09:00
Kohya S	458173da5e	Merge branch 'dev' into dev	2023-03-10 13:00:49 +09:00
Kohya S	1932c31c66	Merge pull request #243 from mgz-dev/dynamic-dim-lora-resize Enable ability to resize lora dim based off sv ratios	2023-03-10 12:59:39 +09:00
Kohya S	dd05d99efd	Merge pull request #272 from kohya-ss/dev support conv2d-3x3, update documents etc	2023-03-09 21:54:41 +09:00
Kohya S	cf2bc437ec	update readme	2023-03-09 21:51:22 +09:00
Kohya S	aa317d4f57	Merge branch 'main' into dev	2023-03-09 20:56:54 +09:00
Kohya S	51249b1ba0	support conv2d 3x3 LoRA	2023-03-09 20:56:33 +09:00
Kohya S	e7051d427c	fix default conv alpha to 1	2023-03-09 20:26:14 +09:00
Kohya S	b885c6f9d2	disable annoying warning in CLIP loading	2023-03-09 20:25:21 +09:00
Kohya S	ad443e172a	fix samle gen failed if use templates	2023-03-09 20:24:53 +09:00
Kohya S	c4b4d1cb40	fix LoRA always expanded to Conv2d-3x3	2023-03-09 08:47:13 +09:00
rvhfxb	82aac26469	Update train_util.py	2023-03-08 22:42:41 +09:00
Kohya S	3ce846525b	set minimum metadata even with no_metadata	2023-03-08 21:19:12 +09:00
Kohya S	8929bf31d9	sample gen h/w to div by 8, fix in steps=epoch	2023-03-08 21:18:28 +09:00
ddPn08	87846c043f	fix for multi gpu training	2023-03-08 09:46:37 +09:00
Kohya S	225c533279	accept empty caption #258	2023-03-07 08:23:34 +09:00
Kohya S	8d5ba29363	free pipe and cache after sample gen #260	2023-03-07 08:06:36 +09:00
Kohya S	19386df6e9	expand LoRA to all Conv2d	2023-03-06 22:03:09 +09:00
Kohya S	5bb571ccc0	Merge branch 'main' into dev	2023-03-06 17:49:47 +09:00
Kohya S	573aa8b5e7	Merge branch 'dev' of https://github.com/kohya-ss/sd-scripts into dev	2023-03-06 17:48:27 +09:00
Kohya S	c2a8290965	Merge pull request #255 from Isotr0py/dev Add network_args to metadata if using another network	2023-03-06 17:46:45 +09:00
Kohya S	1c00764d01	updating documents	2023-03-05 22:32:26 +09:00
Kohya S	2ba6d74af8	Merge branch 'main' into dev	2023-03-05 21:29:40 +09:00
Kohya S	db8c79c463	Update documentation	2023-03-05 19:51:05 +09:00
Kohya S	2b6e9d83fa	Merge branch 'main' into dev	2023-03-05 09:55:41 +09:00
mgz-dev	4a4450d6b6	make new_rank limit max rank, fix zero matrices -new_rank arg changed to limit the max rank of any layer. -added logic to make sure zero-ed layers do not create large lora dim	2023-03-04 03:10:04 -06:00
Kohya S	fe4f4446f1	Add region control for LoRA	2023-03-04 18:03:11 +09:00
mgz-dev	214ed092f2	add support to extract lora with resnet and 2d blocks Modified resize script so support different types of LoRA networks (refer to Kohaku-Blueleaf module implementation for structure).	2023-03-04 02:01:10 -06:00
Isotr0py	4396350271	Add network_args to meta if using another network	2023-03-04 13:59:22 +08:00
mgz-dev	80be6fa130	refactor and bug fix for too large sv_ratio - code refactor to be able to re-use same function for dynamic extract lora - remove clamp - fix issue where if sv_ratio is too high index goes out of bounds	2023-03-03 23:32:46 -06:00
mgz-dev	52ca6c515c	add options to resize based off frobenius norm or cumulative sum	2023-03-01 13:35:24 -06:00
mgz-dev	efe4c98341	Enable ability to resize lora dim based off ratios	2023-02-28 14:55:15 -06:00