Merge pull request #208 from space-nuko/add-optimizer-to-metadata

Add optimizer to metadata

.github/workflows/typos.yml

@@ -0,0 +1,21 @@
+---
+# yamllint disable rule:line-length
+name: Typos
+
+on:  # yamllint disable-line rule:truthy
+  push:
+  pull_request:
+    types:
+      - opened
+      - synchronize
+      - reopened
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+
+    steps:
+      - uses: actions/checkout@v3
+
+      - name: typos-action
+        uses: crate-ci/typos@v1.13.10

README-ja.md

@@ -64,6 +64,12 @@ cp .\bitsandbytes_windows\main.py .\venv\Lib\site-packages\bitsandbytes\cuda_set
 accelerate config
 ```
 
+<!-- 
+pip install torch==1.13.1+cu117 torchvision==0.14.1+cu117 --extra-index-url https://download.pytorch.org/whl/cu117
+pip install --use-pep517 --upgrade -r requirements.txt
+pip install -U -I --no-deps xformers==0.0.16
+-->
+
 コマンドプロンプトでは以下になります。
 
 
@@ -116,7 +122,7 @@ accelerate configの質問には以下のように答えてください。（bf1
 cd sd-scripts
 git pull
 .\venv\Scripts\activate
-pip install --upgrade -r requirements.txt
+pip install --use-pep517 --upgrade -r requirements.txt
 ```
 
 コマンドが成功すれば新しいバージョンが使用できます。

README.md

@@ -1,64 +1,7 @@
 This repository contains training, generation and utility scripts for Stable Diffusion.
 
-## Updates
-
-__Stable Diffusion web UI now seems to support LoRA trained by ``sd-scripts``.__ Thank you for great work!!! 
-
-Note: The LoRA models for SD 2.x is not supported too in Web UI.
-
-- 4 Feb. 2023, 2023/2/4
-  - ``--persistent_data_loader_workers`` option is added to ``fine_tune.py``, ``train_db.py`` and ``train_network.py``. This option may significantly reduce the waiting time between epochs. Thanks to hitomi!
-  - ``--debug_dataset`` option is now working on non-Windows environment. Thanks to tsukimiya!
-  - ``networks/resize_lora.py`` script is added. This can approximate the higher-rank (dim) LoRA model by a lower-rank LoRA model, e.g. 128 by 4. Thanks to mgz-dev!
-    - ``--help`` option shows usage.
-    - Currently the metadata is not copied. This will be fixed in the near future.
-  -  ``--persistent_data_loader_workers``オプションが ``fine_tune.py``、 ``train_db.py``、``train_network.py``の各スクリプトに追加されました。このオプションを指定するとエポック間の待ち時間が大幅に短縮される可能性があります。hitomi氏に感謝します。
-  - ``--debug_dataset``オプションがWindows環境以外でも動くようになりました。tsukimiya氏に感謝します。
-  - ``networks/resize_lora.py``スクリプトを追加しました。高rankのLoRAモデルを低rankのLoRAモデルで近似します（つまり128 rank (dim)のLoRAに似た、4 rank (dim)のLoRAを作ることができます）。mgz-dev氏に感謝します。
-    - 使い方は``--help``オプションを指定して参照してください。
-    - 現時点ではメタデータはコピーされません。近日中に対応予定です。
-- 3 Feb. 2023, 2023/2/3
-  - Update finetune preprocessing scripts.
-    - ``.bmp`` and ``.jpeg`` are supported. Thanks to breakcore2 and p1atdev!
-    - The default weights of ``tag_images_by_wd14_tagger.py`` is now ``SmilingWolf/wd-v1-4-convnext-tagger-v2``. You can specify another model id from ``SmilingWolf`` by ``--repo_id`` option. Thanks to SmilingWolf for the great work.
-      - To change the weight, remove ``wd14_tagger_model`` folder, and run the script again.
-    - ``--max_data_loader_n_workers`` option is added to each script. This option uses the DataLoader for data loading to speed up loading, 20%~30% faster.
-      - Please specify 2 or 4, depends on the number of CPU cores.
-    - ``--recursive`` option is added to ``merge_dd_tags_to_metadata.py`` and ``merge_captions_to_metadata.py``, only works with ``--full_path``.
-    - ``make_captions_by_git.py`` is added. It uses [GIT microsoft/git-large-textcaps](https://huggingface.co/microsoft/git-large-textcaps) for captioning. 
-      - ``requirements.txt`` is updated. If you use this script, [please update the libraries](https://github.com/kohya-ss/sd-scripts#upgrade).
-      - Usage is almost the same as ``make_captions.py``, but batch size should be smaller.
-      - ``--remove_words`` option removes as much text as possible (such as ``the word "XXXX" on it``).
-    - ``--skip_existing`` option is added to ``prepare_buckets_latents.py``. Images with existing npz files are ignored by this option.
-    - ``clean_captions_and_tags.py`` is updated to remove duplicated or conflicting tags, e.g. ``shirt`` is removed when ``white shirt`` exists. if ``black hair`` is with ``red hair``, both are removed.
-  - Tag frequency is added to the metadata in ``train_network.py``. Thanks to space-nuko!
-    - __All tags and number of occurrences of the tag are recorded.__ If you do not want it, disable metadata storing with ``--no_metadata`` option.
-  
-  - fine tuning用の前処理スクリプト群を更新しました。
-    - 拡張子 ``.bmp`` と ``.jpeg`` をサポートしました。breakcore2氏およびp1atdev氏に感謝します。
-    - ``tag_images_by_wd14_tagger.py`` のデフォルトの重みを ``SmilingWolf/wd-v1-4-convnext-tagger-v2`` に更新しました。他の ``SmilingWolf`` 氏の重みも ``--repo_id`` オプションで指定可能です。SmilingWolf氏に感謝します。
-      - 重みを変更するときには ``wd14_tagger_model`` フォルダを削除してからスクリプトを再実行してください。
-    - ``--max_data_loader_n_workers`` オプションが各スクリプトに追加されました。DataLoaderを用いることで読み込み処理を並列化し、処理を20~30%程度高速化します。
-      - CPUのコア数に応じて2~4程度の値を指定してください。
-    - ``--recursive`` オプションを ``merge_dd_tags_to_metadata.py`` と ``merge_captions_to_metadata.py`` に追加しました。``--full_path`` を指定したときのみ使用可能です。
-    - ``make_captions_by_git.py`` を追加しました。[GIT microsoft/git-large-textcaps](https://huggingface.co/microsoft/git-large-textcaps) を用いてキャプションニングを行います。
-      - ``requirements.txt`` が更新されていますので、[ライブラリをアップデート](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)してください。
-      - 使用法は ``make_captions.py``とほぼ同じですがバッチサイズは小さめにしてください。
-      - ``--remove_words`` オプションを指定するとテキスト読み取りを可能な限り削除します（``the word "XXXX" on it``のようなもの）。    
-    - ``--skip_existing`` を ``prepare_buckets_latents.py`` に追加しました。すでにnpzファイルがある画像の処理をスキップします。
-    - ``clean_captions_and_tags.py``を重複タグや矛盾するタグを削除するよう機能追加しました。例：``white shirt`` タグがある場合、 ``shirt`` タグは削除されます。また``black hair``と``red hair``の両方がある場合、両方とも削除されます。
-  - ``train_network.py``で使用されているタグと回数をメタデータに記録するようになりました。space-nuko氏に感謝します。
-    - __すべてのタグと回数がメタデータに記録されます__ 望まない場合には``--no_metadata option``オプションでメタデータの記録を停止してください。
-    
-
-Stable Diffusion web UI本体で当リポジトリで学習したLoRAモデルによる画像生成がサポートされたようです。
-
-注：SD2.x用のLoRAモデルはサポートされないようです。
-
-Please read [Releases](https://github.com/kohya-ss/sd-scripts/releases) for recent updates.
-最近の更新情報は [Release](https://github.com/kohya-ss/sd-scripts/releases) をご覧ください。
-
-##
+[__Change History__](#change-history) is moved to the bottom of the page.
+更新履歴は[ページ末尾](#change-history)に移しました。
 
 [日本語版README](./README-ja.md)
 
@@ -67,10 +10,13 @@ For easier use (GUI and PowerShell scripts etc...), please visit [the repository
 This repository contains the scripts for:
 
 * DreamBooth training, including U-Net and Text Encoder
-* fine-tuning (native training), including U-Net and Text Encoder
+* Fine-tuning (native training), including U-Net and Text Encoder
 * LoRA training
-* image generation
-* model conversion (supports 1.x and 2.x, Stable Diffision ckpt/safetensors and Diffusers)
+* Texutl Inversion training
+* Image generation
+* Model conversion (supports 1.x and 2.x, Stable Diffision ckpt/safetensors and Diffusers)
+
+__Stable Diffusion web UI now seems to support LoRA trained by ``sd-scripts``.__ (SD 1.x based only) Thank you for great work!!! 
 
 ## About requirements.txt
 
@@ -157,7 +103,7 @@ When a new release comes out you can upgrade your repo with the following comman
 cd sd-scripts
 git pull
 .\venv\Scripts\activate
-pip install --upgrade -r requirements.txt
+pip install --use-pep517 --upgrade -r requirements.txt
 ```
 
 Once the commands have completed successfully you should be ready to use the new version.
@@ -175,3 +121,115 @@ The majority of scripts is licensed under ASL 2.0 (including codes from Diffuser
 [bitsandbytes](https://github.com/TimDettmers/bitsandbytes): MIT
 
 [BLIP](https://github.com/salesforce/BLIP): BSD-3-Clause
+
+## Change History
+
+- 19 Feb. 2023, 2023/2/19:
+  - Add ``--use_lion_optimizer`` to each training script to use [Lion optimizer](https://github.com/lucidrains/lion-pytorch).
+    - Please install Lion optimizer with ``pip install lion-pytorch`` (it is not in ``requirements.txt`` currently.)
+  - Add ``--lowram`` option to ``train_network.py``. Load models to VRAM instead of VRAM (for machines which have bigger VRAM than RAM such as Colab and Kaggle). Thanks to Isotr0py!
+    - Default behavior (without lowram) has reverted to the same as before 14 Feb.
+  - Fixed git commit hash to be set correctly regardless of the working directory. Thanks to vladmandic!
+
+  - ``--use_lion_optimizer`` オプションを各学習スクリプトに追加しました。 [Lion optimizer](https://github.com/lucidrains/lion-pytorch) を使用できます。
+    - あらかじめ ``pip install lion-pytorch`` でインストールしてください（現在は ``requirements.txt`` に含まれていません）。
+  - ``--lowram`` オプションを ``train_network.py`` に追加しました。モデルをRAMではなくVRAMに読み込みます（ColabやKaggleなど、VRAMがRAMに比べて多い環境で有効です）。 Isotr0py 氏に感謝します。
+    - lowram オプションなしのデフォルト動作は2/14より前と同じに戻しました。
+  - git commit hash を現在のフォルダ位置に関わらず正しく取得するように修正しました。vladmandic 氏に感謝します。
+  
+- 16 Feb. 2023, 2023/2/16:
+  - Noise offset is recorded to the metadata. Thanks to space-nuko!
+  - Show the moving average loss to prevent loss jumping in ``train_network.py`` and ``train_db.py``. Thanks to shirayu!
+  - Noise offsetがメタデータに記録されるようになりました。space-nuko氏に感謝します。
+  - ``train_network.py``と``train_db.py``で学習中に表示されるlossの値が移動平均になりました。epochの先頭で表示されるlossが大きく変動する事象を解決します。shirayu氏に感謝します。
+- 14 Feb. 2023, 2023/2/14:
+  - Add support with multi-gpu trainining for ``train_network.py``. Thanks to Isotr0py!
+  - Add ``--verbose`` option for ``resize_lora.py``. For details, see [this PR](https://github.com/kohya-ss/sd-scripts/pull/179). Thanks to mgz-dev!
+  - Git commit hash is added to the metadata for LoRA. Thanks to space-nuko!
+  - Add ``--noise_offset`` option for each training scripts. 
+    - Implementation of https://www.crosslabs.org//blog/diffusion-with-offset-noise 
+    - This option may improve ability to generate darker/lighter images. May work with LoRA.
+  - ``train_network.py``でマルチGPU学習をサポートしました。Isotr0py氏に感謝します。
+  - ``--verbose``オプションを ``resize_lora.py`` に追加しました。表示される情報の詳細は [こちらのPR](https://github.com/kohya-ss/sd-scripts/pull/179) をご参照ください。mgz-dev氏に感謝します。
+  - LoRAのメタデータにgitのcommit hashを追加しました。space-nuko氏に感謝します。
+  - ``--noise_offset`` オプションを各学習スクリプトに追加しました。
+    - こちらの記事の実装になります: https://www.crosslabs.org//blog/diffusion-with-offset-noise 
+    - 全体的に暗い、明るい画像の生成結果が良くなる可能性があるようです。LoRA学習でも有効なようです。
+
+- 11 Feb. 2023, 2023/2/11:
+  - ``lora_interrogator.py`` is added in ``networks`` folder. See ``python networks\lora_interrogator.py -h`` for usage.
+    - For LoRAs where the activation word is unknown, this script compares the output of Text Encoder after applying LoRA to that of unapplied to find out which token is affected by LoRA. Hopefully you can figure out the activation word. LoRA trained with captions does not seem to be able to interrogate.
+    - Batch size can be large (like 64 or 128).
+  - ``train_textual_inversion.py`` now supports multiple init words.
+  - Following feature is reverted to be the same as before. Sorry for confusion:
+    > Now the number of data in each batch is limited to the number of actual images (not duplicated). Because a certain bucket may contain smaller number of actual images, so the batch may contain same (duplicated) images.
+  
+  - ``lora_interrogator.py`` を ``network``フォルダに追加しました。使用法は ``python networks\lora_interrogator.py -h`` でご確認ください。
+    - このスクリプトは、起動promptがわからないLoRAについて、LoRA適用前後のText Encoderの出力を比較することで、どのtokenの出力が変化しているかを調べます。運が良ければ起動用の単語が分かります。キャプション付きで学習されたLoRAは影響が広範囲に及ぶため、調査は難しいようです。
+    - バッチサイズはわりと大きくできます（64や128など）。
+  - ``train_textual_inversion.py`` で複数のinit_word指定が可能になりました。
+  - 次の機能を削除し元に戻しました。混乱を招き申し訳ありません。
+    >  これらのオプションによりbucketが細分化され、ひとつのバッチ内に同一画像が重複して存在することが増えたため、バッチサイズを``そのbucketの画像種類数``までに制限する機能を追加しました。
+
+- 10 Feb. 2023, 2023/2/10:
+  - Updated ``requirements.txt`` to prevent upgrading with pip taking a long time or failure to upgrade.
+  - ``resize_lora.py`` keeps the metadata of the model. ``dimension is resized from ...`` is added to the top of  ``ss_training_comment``.
+  - ``merge_lora.py`` supports models with different ``alpha``s. If there is a problem, old version is ``merge_lora_old.py``.
+  - ``svd_merge_lora.py`` is added. This script merges LoRA models with any rank (dim) and alpha, and approximate a new LoRA with svd for a specified rank (dim).
+  - Note: merging scripts erase the metadata currently.
+  - ``resize_images_to_resolution.py`` supports multibyte characters in filenames.
+  - pipでの更新が長時間掛かったり、更新に失敗したりするのを防ぐため、``requirements.txt``を更新しました。
+  - ``resize_lora.py``がメタデータを保持するようになりました。 ``dimension is resized from ...`` という文字列が ``ss_training_comment`` の先頭に追加されます。
+  - ``merge_lora.py``がalphaが異なるモデルをサポートしました。 何か問題がありましたら旧バージョン ``merge_lora_old.py`` をお使いください。
+  - ``svd_merge_lora.py`` を追加しました。 複数の任意のdim (rank)、alphaのLoRAモデルをマージし、svdで任意dim(rank)のLoRAで近似します。 
+  - 注：マージ系のスクリプトは現時点ではメタデータを消去しますのでご注意ください。
+  - ``resize_images_to_resolution.py``が日本語ファイル名をサポートしました。
+  
+- 9 Feb. 2023, 2023/2/9:
+  - Caption dropout is supported in ``train_db.py``, ``fine_tune.py`` and ``train_network.py``. Thanks to forestsource!
+    - ``--caption_dropout_rate`` option specifies the dropout rate for captions (0~1.0, 0.1 means 10% chance for dropout). If dropout occurs, the image is trained with the empty caption. Default is 0 (no dropout).
+    - ``--caption_dropout_every_n_epochs`` option specifies how many epochs to drop captions. If ``3`` is specified, in epoch 3, 6, 9 ..., images are trained with all captions empty. Default is None (no dropout).
+    - ``--caption_tag_dropout_rate`` option specified the dropout rate for tags (comma separated tokens) (0~1.0, 0.1 means 10% chance for dropout). If dropout occurs, the tag is removed from the caption. If ``--keep_tokens`` option is set, these tokens (tags) are not dropped. Default is 0 (no droupout).
+    - The bulk image downsampling script is added. Documentation is [here](https://github.com/kohya-ss/sd-scripts/blob/main/train_network_README-ja.md#%E7%94%BB%E5%83%8F%E3%83%AA%E3%82%B5%E3%82%A4%E3%82%BA%E3%82%B9%E3%82%AF%E3%83%AA%E3%83%97%E3%83%88) (in Jpanaese). Thanks to bmaltais!
+    - Typo check is added. Thanks to shirayu!
+  - キャプションのドロップアウトを``train_db.py``、``fine_tune.py``、``train_network.py``の各スクリプトに追加しました。forestsource氏に感謝します。
+    - ``--caption_dropout_rate``オプションでキャプションのドロップアウト率を指定します（0~1.0、 0.1を指定すると10%の確率でドロップアウト）。ドロップアウトされた場合、画像は空のキャプションで学習されます。デフォルトは 0 （ドロップアウトなし）です。
+    - ``--caption_dropout_every_n_epochs`` オプションで何エポックごとにキャプションを完全にドロップアウトするか指定します。たとえば``3``を指定すると、エポック3、6、9……で、すべての画像がキャプションなしで学習されます。デフォルトは None （ドロップアウトなし）です。
+    - ``--caption_tag_dropout_rate`` オプションで各タグ（カンマ区切りの各部分）のドロップアウト率を指定します（0~1.0、 0.1を指定すると10%の確率でドロップアウト）。ドロップアウトが起きるとそのタグはそのときだけキャプションから取り除かれて学習されます。``--keep_tokens`` オプションを指定していると、シャッフルされない部分のタグはドロップアウトされません。デフォルトは 0 （ドロップアウトなし）です。
+    - 画像の一括縮小スクリプトを追加しました。ドキュメントは [こちら](https://github.com/kohya-ss/sd-scripts/blob/main/train_network_README-ja.md#%E7%94%BB%E5%83%8F%E3%83%AA%E3%82%B5%E3%82%A4%E3%82%BA%E3%82%B9%E3%82%AF%E3%83%AA%E3%83%97%E3%83%88) です。bmaltais氏に感謝します。
+    - 誤字チェッカが追加されました。shirayu氏に感謝します。
+    
+- 6 Feb. 2023, 2023/2/6：
+  - ``--bucket_reso_steps`` and ``--bucket_no_upscale`` options are added to training scripts (fine tuning, DreamBooth, LoRA and Textual Inversion) and ``prepare_buckets_latents.py``.
+  - ``--bucket_reso_steps`` takes the steps for buckets in aspect ratio bucketing. Default is 64, same as before.
+    - Any value greater than or equal to 1 can be specified; 64 is highly recommended and a value divisible by 8 is recommended.
+    - If less than 64 is specified, padding will occur within U-Net. The result is unknown.
+    - If you specify a value that is not divisible by 8, it will be truncated to divisible by 8 inside VAE, because the size of the latent is 1/8 of the image size.
+  - If ``--bucket_no_upscale`` option is specified, images smaller than the bucket size will be processed without upscaling.
+    - Internally, a bucket smaller than the image size is created (for example, if the image is 300x300 and ``bucket_reso_steps=64``, the bucket is 256x256). The image will be trimmed.
+    - Implementation of [#130](https://github.com/kohya-ss/sd-scripts/issues/130).
+    - Images with an area larger than the maximum size specified by ``--resolution`` are downsampled to the max bucket size.
+  - Now the number of data in each batch is limited to the number of actual images (not duplicated). Because a certain bucket may contain smaller number of actual images, so the batch may contain same (duplicated) images.
+  - ``--random_crop`` now also works with buckets enabled.
+    - Instead of always cropping the center of the image, the image is shifted left, right, up, and down to be used as the training data. This is expected to train to the edges of the image.
+    - Implementation of discussion [#34](https://github.com/kohya-ss/sd-scripts/discussions/34).
+
+  - ``--bucket_reso_steps``および``--bucket_no_upscale``オプションを、学習スクリプトおよび``prepare_buckets_latents.py``に追加しました。
+  - ``--bucket_reso_steps``オプションでは、bucketの解像度の単位を指定できます。デフォルトは64で、今までと同じ動作です。
+    - 1以上の任意の値を指定できます。基本的には64を推奨します。64以外の値では、8で割り切れる値を推奨します。
+    - 64未満を指定するとU-Netの内部でpaddingが発生します。どのような結果になるかは未知数です。
+    - 8で割り切れない値を指定すると余りはVAE内部で切り捨てられます。
+  - ``--bucket_no_upscale``オプションを指定すると、bucketサイズよりも小さい画像は拡大せずそのまま処理します。
+    - 内部的には画像サイズ以下のサイズのbucketを作成します（たとえば画像が300x300で``bucket_reso_steps=64``の場合、256x256のbucket）。余りは都度trimmingされます。
+    - [#130](https://github.com/kohya-ss/sd-scripts/issues/130) を実装したものです。
+    - ``--resolution``で指定した最大サイズよりも面積が大きい画像は、最大サイズと同じ面積になるようアスペクト比を維持したまま縮小され、そのサイズを元にbucketが作られます。
+  - これらのオプションによりbucketが細分化され、ひとつのバッチ内に同一画像が重複して存在することが増えたため、バッチサイズを``そのbucketの画像種類数``までに制限する機能を追加しました。
+    - たとえば繰り返し回数10で、あるbucketに1枚しか画像がなく、バッチサイズが10以上のとき、今まではepoch内で、同一画像を10枚含むバッチが1回だけ使用されていました。
+    - 機能追加後はepoch内にサイズ1のバッチが10回、使用されます。
+  - ``--random_crop``がbucketを有効にした場合にも機能するようになりました。
+    - 常に画像の中央を切り取るのではなく、左右、上下にずらして教師データにします。これにより画像端まで学習されることが期待されます。
+    - discussionの[#34](https://github.com/kohya-ss/sd-scripts/discussions/34)を実装したものです。
+    
+
+Please read [Releases](https://github.com/kohya-ss/sd-scripts/releases) for recent updates.
+最近の更新情報は [Release](https://github.com/kohya-ss/sd-scripts/releases) をご覧ください。

_typos.toml

@@ -0,0 +1,15 @@
+# Files for typos
+# Instruction:  https://github.com/marketplace/actions/typos-action#getting-started
+
+[default.extend-identifiers]
+
+[default.extend-words]
+NIN="NIN"
+parms="parms"
+nin="nin"
+extention="extention" # Intentionally left
+nd="nd"
+
+
+[files]
+extend-exclude = ["_typos.toml"]

fine_tune.py

@@ -33,8 +33,13 @@ def train(args):
   train_dataset = train_util.FineTuningDataset(args.in_json, args.train_batch_size, args.train_data_dir,
                                                tokenizer, args.max_token_length, args.shuffle_caption, args.keep_tokens,
                                                args.resolution, args.enable_bucket, args.min_bucket_reso, args.max_bucket_reso,
+                                               args.bucket_reso_steps, args.bucket_no_upscale,
                                                args.flip_aug, args.color_aug, args.face_crop_aug_range, args.random_crop,
                                                args.dataset_repeats, args.debug_dataset)
+
+  # 学習データのdropout率を設定する
+  train_dataset.set_caption_dropout(args.caption_dropout_rate, args.caption_dropout_every_n_epochs, args.caption_tag_dropout_rate)
+
   train_dataset.make_buckets()
 
   if args.debug_dataset:
@@ -153,6 +158,13 @@ def train(args):
       raise ImportError("No bitsand bytes / bitsandbytesがインストールされていないようです")
     print("use 8-bit Adam optimizer")
     optimizer_class = bnb.optim.AdamW8bit
+  elif args.use_lion_optimizer:
+    try:
+      import lion_pytorch
+    except ImportError:
+      raise ImportError("No lion_pytorch / lion_pytorch がインストールされていないようです")
+    print("use Lion optimizer")
+    optimizer_class = lion_pytorch.Lion
   else:
     optimizer_class = torch.optim.AdamW
 
@@ -225,6 +237,8 @@ def train(args):
 
   for epoch in range(num_train_epochs):
     print(f"epoch {epoch+1}/{num_train_epochs}")
+    train_dataset.set_current_epoch(epoch + 1)
+
     for m in training_models:
       m.train()
 
@@ -248,6 +262,9 @@ def train(args):
 
         # Sample noise that we'll add to the latents
         noise = torch.randn_like(latents, device=latents.device)
+        if args.noise_offset:
+          # https://www.crosslabs.org//blog/diffusion-with-offset-noise
+          noise += args.noise_offset * torch.randn((latents.shape[0], latents.shape[1], 1, 1), device=latents.device)
 
         # Sample a random timestep for each image
         timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (b_size,), device=latents.device)
@@ -331,7 +348,7 @@ if __name__ == '__main__':
   parser = argparse.ArgumentParser()
 
   train_util.add_sd_models_arguments(parser)
-  train_util.add_dataset_arguments(parser, False, True)
+  train_util.add_dataset_arguments(parser, False, True, True)
   train_util.add_training_arguments(parser, False)
   train_util.add_sd_saving_arguments(parser)
 

finetune/prepare_buckets_latents.py

@@ -52,6 +52,10 @@ def get_npz_filename_wo_ext(data_dir, image_key, is_full_path, flip):
 
 
 def main(args):
+  # assert args.bucket_reso_steps % 8 == 0, f"bucket_reso_steps must be divisible by 8 / bucket_reso_stepは8で割り切れる必要があります"
+  if args.bucket_reso_steps % 8 > 0:
+    print(f"resolution of buckets in training time is a multiple of 8 / 学習時の各bucketの解像度は8単位になります")
+
   image_paths = train_util.glob_images(args.train_data_dir)
   print(f"found {len(image_paths)} images.")
 
@@ -77,32 +81,41 @@ def main(args):
   max_reso = tuple([int(t) for t in args.max_resolution.split(',')])
   assert len(max_reso) == 2, f"illegal resolution (not 'width,height') / 画像サイズに誤りがあります。'幅,高さ'で指定してください: {args.max_resolution}"
 
-  bucket_resos, bucket_aspect_ratios = model_util.make_bucket_resolutions(
-      max_reso, args.min_bucket_reso, args.max_bucket_reso)
+  bucket_manager = train_util.BucketManager(args.bucket_no_upscale, max_reso,
+                                            args.min_bucket_reso, args.max_bucket_reso, args.bucket_reso_steps)
+  if not args.bucket_no_upscale:
+    bucket_manager.make_buckets()
+  else:
+    print("min_bucket_reso and max_bucket_reso are ignored if bucket_no_upscale is set, because bucket reso is defined by image size automatically / bucket_no_upscaleが指定された場合は、bucketの解像度は画像サイズから自動計算されるため、min_bucket_resoとmax_bucket_resoは無視されます")
 
   # 画像をひとつずつ適切なbucketに割り当てながらlatentを計算する
-  bucket_aspect_ratios = np.array(bucket_aspect_ratios)
-  buckets_imgs = [[] for _ in range(len(bucket_resos))]
-  bucket_counts = [0 for _ in range(len(bucket_resos))]
   img_ar_errors = []
 
   def process_batch(is_last):
-    for j in range(len(buckets_imgs)):
-      bucket = buckets_imgs[j]
+    for bucket in bucket_manager.buckets:
       if (is_last and len(bucket) > 0) or len(bucket) >= args.batch_size:
-        latents = get_latents(vae, [img for _, _, img in bucket], weight_dtype)
+        latents = get_latents(vae, [img for _, img in bucket], weight_dtype)
+        assert latents.shape[2] == bucket[0][1].shape[0] // 8 and latents.shape[3] == bucket[0][1].shape[1] // 8, \
+            f"latent shape {latents.shape}, {bucket[0][1].shape}"
 
-        for (image_key, _, _), latent in zip(bucket, latents):
+        for (image_key, _), latent in zip(bucket, latents):
           npz_file_name = get_npz_filename_wo_ext(args.train_data_dir, image_key, args.full_path, False)
           np.savez(npz_file_name, latent)
 
         # flip
         if args.flip_aug:
-          latents = get_latents(vae, [img[:, ::-1].copy() for _, _, img in bucket], weight_dtype)   # copyがないとTensor変換できない
+          latents = get_latents(vae, [img[:, ::-1].copy() for _, img in bucket], weight_dtype)   # copyがないとTensor変換できない
 
-          for (image_key, _, _), latent in zip(bucket, latents):
+          for (image_key, _), latent in zip(bucket, latents):
             npz_file_name = get_npz_filename_wo_ext(args.train_data_dir, image_key, args.full_path, True)
             np.savez(npz_file_name, latent)
+        else:
+          # remove existing flipped npz
+          for image_key, _ in bucket:
+            npz_file_name = get_npz_filename_wo_ext(args.train_data_dir, image_key, args.full_path, True) + ".npz"
+            if os.path.isfile(npz_file_name):
+              print(f"remove existing flipped npz / 既存のflipされたnpzファイルを削除します: {npz_file_name}")
+              os.remove(npz_file_name)
 
         bucket.clear()
 
@@ -114,6 +127,7 @@ def main(args):
   else:
     data = [[(None, ip)] for ip in image_paths]
 
+  bucket_counts = {}
   for data_entry in tqdm(data, smoothing=0.0):
     if data_entry[0] is None:
       continue
@@ -134,29 +148,24 @@ def main(args):
     if image_key not in metadata:
       metadata[image_key] = {}
 
-    # 本当はこの部分もDataSetに持っていけば高速化できるがいろいろ大変
-    aspect_ratio = image.width / image.height
-    ar_errors = bucket_aspect_ratios - aspect_ratio
-    bucket_id = np.abs(ar_errors).argmin()
-    reso = bucket_resos[bucket_id]
-    ar_error = ar_errors[bucket_id]
+    # 本当はこのあとの部分もDataSetに持っていけば高速化できるがいろいろ大変
+
+    reso, resized_size, ar_error = bucket_manager.select_bucket(image.width, image.height)
     img_ar_errors.append(abs(ar_error))
+    bucket_counts[reso] = bucket_counts.get(reso, 0) + 1
 
-    # どのサイズにリサイズするか→トリミングする方向で
-    if ar_error <= 0:                   # 横が長い→縦を合わせる
-      scale = reso[1] / image.height
-    else:
-      scale = reso[0] / image.width
+    # メタデータに記録する解像度はlatent単位とするので、8単位で切り捨て
+    metadata[image_key]['train_resolution'] = (reso[0] - reso[0] % 8, reso[1] - reso[1] % 8)
 
-    resized_size = (int(image.width * scale + .5), int(image.height * scale + .5))
+    if not args.bucket_no_upscale:
+      # upscaleを行わないときには、resize後のサイズは、bucketのサイズと、縦横どちらかが同じであることを確認する
+      assert resized_size[0] == reso[0] or resized_size[1] == reso[
+          1], f"internal error, resized size not match: {reso}, {resized_size}, {image.width}, {image.height}"
+      assert resized_size[0] >= reso[0] and resized_size[1] >= reso[
+          1], f"internal error, resized size too small: {reso}, {resized_size}, {image.width}, {image.height}"
 
-    # print(image.width, image.height, bucket_id, bucket_resos[bucket_id], ar_errors[bucket_id], resized_size,
-    #       bucket_resos[bucket_id][0] - resized_size[0], bucket_resos[bucket_id][1] - resized_size[1])
-
-    assert resized_size[0] == reso[0] or resized_size[1] == reso[
-        1], f"internal error, resized size not match: {reso}, {resized_size}, {image.width}, {image.height}"
     assert resized_size[0] >= reso[0] and resized_size[1] >= reso[
-        1], f"internal error, resized size too small: {reso}, {resized_size}, {image.width}, {image.height}"
+        1], f"internal error resized size is small: {resized_size}, {reso}"
 
     # 既に存在するファイルがあればshapeを確認して同じならskipする
     if args.skip_existing:
@@ -180,22 +189,24 @@ def main(args):
     # 画像をリサイズしてトリミングする
     # PILにinter_areaがないのでcv2で……
     image = np.array(image)
-    image = cv2.resize(image, resized_size, interpolation=cv2.INTER_AREA)
+    if resized_size[0] != image.shape[1] or resized_size[1] != image.shape[0]:            # リサイズ処理が必要？
+      image = cv2.resize(image, resized_size, interpolation=cv2.INTER_AREA)
+
     if resized_size[0] > reso[0]:
       trim_size = resized_size[0] - reso[0]
       image = image[:, trim_size//2:trim_size//2 + reso[0]]
-    elif resized_size[1] > reso[1]:
+
+    if resized_size[1] > reso[1]:
       trim_size = resized_size[1] - reso[1]
       image = image[trim_size//2:trim_size//2 + reso[1]]
+
     assert image.shape[0] == reso[1] and image.shape[1] == reso[0], f"internal error, illegal trimmed size: {image.shape}, {reso}"
 
     # # debug
-    # cv2.imwrite(f"r:\\test\\img_{i:05d}.jpg", image[:, :, ::-1])
+    # cv2.imwrite(f"r:\\test\\img_{len(img_ar_errors)}.jpg", image[:, :, ::-1])
 
     # バッチへ追加
-    buckets_imgs[bucket_id].append((image_key, reso, image))
-    bucket_counts[bucket_id] += 1
-    metadata[image_key]['train_resolution'] = reso
+    bucket_manager.add_image(reso, (image_key, image))
 
     # バッチを推論するか判定して推論する
     process_batch(False)
@@ -203,8 +214,11 @@ def main(args):
   # 残りを処理する
   process_batch(True)
 
-  for i, (reso, count) in enumerate(zip(bucket_resos, bucket_counts)):
-    print(f"bucket {i} {reso}: {count}")
+  bucket_manager.sort()
+  for i, reso in enumerate(bucket_manager.resos):
+    count = bucket_counts.get(reso, 0)
+    if count > 0:
+      print(f"bucket {i} {reso}: {count}")
   img_ar_errors = np.array(img_ar_errors)
   print(f"mean ar error: {np.mean(img_ar_errors)}")
 
@@ -230,6 +244,10 @@ if __name__ == '__main__':
                       help="max resolution in fine tuning (width,height) / fine tuning時の最大画像サイズ 「幅,高さ」（使用メモリ量に関係します）")
   parser.add_argument("--min_bucket_reso", type=int, default=256, help="minimum resolution for buckets / bucketの最小解像度")
   parser.add_argument("--max_bucket_reso", type=int, default=1024, help="maximum resolution for buckets / bucketの最小解像度")
+  parser.add_argument("--bucket_reso_steps", type=int, default=64,
+                      help="steps of resolution for buckets, divisible by 8 is recommended / bucketの解像度の単位、8で割り切れる値を推奨します")
+  parser.add_argument("--bucket_no_upscale", action="store_true",
+                      help="make bucket for each image without upscaling / 画像を拡大せずbucketを作成します")
   parser.add_argument("--mixed_precision", type=str, default="no",
                       choices=["no", "fp16", "bf16"], help="use mixed precision / 混合精度を使う場合、その精度")
   parser.add_argument("--full_path", action="store_true",

library/model_util.py

@@ -1163,15 +1163,14 @@ def make_bucket_resolutions(max_reso, min_size=256, max_size=1024, divisible=64)
 
   resos = list(resos)
   resos.sort()
-
-  aspect_ratios = [w / h for w, h in resos]
-  return resos, aspect_ratios
+  return resos
 
 
 if __name__ == '__main__':
-  resos, aspect_ratios = make_bucket_resolutions((512, 768))
+  resos = make_bucket_resolutions((512, 768))
   print(len(resos))
   print(resos)
+  aspect_ratios = [w / h for w, h in resos]
   print(aspect_ratios)
 
   ars = set()

library/train_util.py

@@ -4,7 +4,7 @@ import argparse
 import json
 import shutil
 import time
-from typing import NamedTuple
+from typing import Dict, List, NamedTuple, Tuple
 from accelerate import Accelerator
 from torch.autograd.function import Function
 import glob
@@ -12,6 +12,7 @@ import math
 import os
 import random
 import hashlib
+import subprocess
 from io import BytesIO
 
 from tqdm import tqdm
@@ -55,16 +56,142 @@ class ImageInfo():
     self.caption: str = caption
     self.is_reg: bool = is_reg
     self.absolute_path: str = absolute_path
-    self.image_size: tuple[int, int] = None
-    self.bucket_reso: tuple[int, int] = None
+    self.image_size: Tuple[int, int] = None
+    self.resized_size: Tuple[int, int] = None
+    self.bucket_reso: Tuple[int, int] = None
     self.latents: torch.Tensor = None
     self.latents_flipped: torch.Tensor = None
     self.latents_npz: str = None
     self.latents_npz_flipped: str = None
 
 
+class BucketManager():
+  def __init__(self, no_upscale, max_reso, min_size, max_size, reso_steps) -> None:
+    self.no_upscale = no_upscale
+    if max_reso is None:
+      self.max_reso = None
+      self.max_area = None
+    else:
+      self.max_reso = max_reso
+      self.max_area = max_reso[0] * max_reso[1]
+    self.min_size = min_size
+    self.max_size = max_size
+    self.reso_steps = reso_steps
+
+    self.resos = []
+    self.reso_to_id = {}
+    self.buckets = []                     # 前処理時は (image_key, image)、学習時は image_key
+
+  def add_image(self, reso, image):
+    bucket_id = self.reso_to_id[reso]
+    self.buckets[bucket_id].append(image)
+
+  def shuffle(self):
+    for bucket in self.buckets:
+      random.shuffle(bucket)
+
+  def sort(self):
+    # 解像度順にソートする（表示時、メタデータ格納時の見栄えをよくするためだけ）。bucketsも入れ替えてreso_to_idも振り直す
+    sorted_resos = self.resos.copy()
+    sorted_resos.sort()
+
+    sorted_buckets = []
+    sorted_reso_to_id = {}
+    for i, reso in enumerate(sorted_resos):
+      bucket_id = self.reso_to_id[reso]
+      sorted_buckets.append(self.buckets[bucket_id])
+      sorted_reso_to_id[reso] = i
+
+    self.resos = sorted_resos
+    self.buckets = sorted_buckets
+    self.reso_to_id = sorted_reso_to_id
+
+  def make_buckets(self):
+    resos = model_util.make_bucket_resolutions(self.max_reso, self.min_size, self.max_size, self.reso_steps)
+    self.set_predefined_resos(resos)
+
+  def set_predefined_resos(self, resos):
+    # 規定サイズから選ぶ場合の解像度、aspect ratioの情報を格納しておく
+    self.predefined_resos = resos.copy()
+    self.predefined_resos_set = set(resos)
+    self.predefined_aspect_ratios = np.array([w / h for w, h in resos])
+
+  def add_if_new_reso(self, reso):
+    if reso not in self.reso_to_id:
+      bucket_id = len(self.resos)
+      self.reso_to_id[reso] = bucket_id
+      self.resos.append(reso)
+      self.buckets.append([])
+      # print(reso, bucket_id, len(self.buckets))
+
+  def round_to_steps(self, x):
+    x = int(x + .5)
+    return x - x % self.reso_steps
+
+  def select_bucket(self, image_width, image_height):
+    aspect_ratio = image_width / image_height
+    if not self.no_upscale:
+      # 同じaspect ratioがあるかもしれないので（fine tuningで、no_upscale=Trueで前処理した場合）、解像度が同じものを優先する
+      reso = (image_width, image_height)
+      if reso in self.predefined_resos_set:
+        pass
+      else:
+        ar_errors = self.predefined_aspect_ratios - aspect_ratio
+        predefined_bucket_id = np.abs(ar_errors).argmin()          # 当該解像度以外でaspect ratio errorが最も少ないもの
+        reso = self.predefined_resos[predefined_bucket_id]
+
+      ar_reso = reso[0] / reso[1]
+      if aspect_ratio > ar_reso:                   # 横が長い→縦を合わせる
+        scale = reso[1] / image_height
+      else:
+        scale = reso[0] / image_width
+
+      resized_size = (int(image_width * scale + .5), int(image_height * scale + .5))
+      # print("use predef", image_width, image_height, reso, resized_size)
+    else:
+      if image_width * image_height > self.max_area:
+        # 画像が大きすぎるのでアスペクト比を保ったまま縮小することを前提にbucketを決める
+        resized_width = math.sqrt(self.max_area * aspect_ratio)
+        resized_height = self.max_area / resized_width
+        assert abs(resized_width / resized_height - aspect_ratio) < 1e-2, "aspect is illegal"
+
+        # リサイズ後の短辺または長辺をreso_steps単位にする：aspect ratioの差が少ないほうを選ぶ
+        # 元のbucketingと同じロジック
+        b_width_rounded = self.round_to_steps(resized_width)
+        b_height_in_wr = self.round_to_steps(b_width_rounded / aspect_ratio)
+        ar_width_rounded = b_width_rounded / b_height_in_wr
+
+        b_height_rounded = self.round_to_steps(resized_height)
+        b_width_in_hr = self.round_to_steps(b_height_rounded * aspect_ratio)
+        ar_height_rounded = b_width_in_hr / b_height_rounded
+
+        # print(b_width_rounded, b_height_in_wr, ar_width_rounded)
+        # print(b_width_in_hr, b_height_rounded, ar_height_rounded)
+
+        if abs(ar_width_rounded - aspect_ratio) < abs(ar_height_rounded - aspect_ratio):
+          resized_size = (b_width_rounded, int(b_width_rounded / aspect_ratio + .5))
+        else:
+          resized_size = (int(b_height_rounded * aspect_ratio + .5), b_height_rounded)
+        # print(resized_size)
+      else:
+        resized_size = (image_width, image_height)              # リサイズは不要
+
+      # 画像のサイズ未満をbucketのサイズとする（paddingせずにcroppingする）
+      bucket_width = resized_size[0] - resized_size[0] % self.reso_steps
+      bucket_height = resized_size[1] - resized_size[1] % self.reso_steps
+      # print("use arbitrary", image_width, image_height, resized_size, bucket_width, bucket_height)
+
+      reso = (bucket_width, bucket_height)
+
+    self.add_if_new_reso(reso)
+
+    ar_error = (reso[0] / reso[1]) - aspect_ratio
+    return reso, resized_size, ar_error
+
+
 class BucketBatchIndex(NamedTuple):
   bucket_index: int
+  bucket_batch_size: int
   batch_index: int
 
 
@@ -85,14 +212,23 @@ class BaseDataset(torch.utils.data.Dataset):
     self.token_padding_disabled = False
     self.dataset_dirs_info = {}
     self.reg_dataset_dirs_info = {}
+    self.tag_frequency = {}
+
     self.enable_bucket = False
+    self.bucket_manager: BucketManager = None                         # not initialized
     self.min_bucket_reso = None
     self.max_bucket_reso = None
-    self.tag_frequency = {}
-    self.bucket_info = None
+    self.bucket_reso_steps = None
+    self.bucket_no_upscale = None
+    self.bucket_info = None                                           # for metadata
 
     self.tokenizer_max_length = self.tokenizer.model_max_length if max_token_length is None else max_token_length + 2
 
+    self.current_epoch: int = 0            # インスタンスがepochごとに新しく作られるようなので外側から渡さないとダメ
+    self.dropout_rate: float = 0
+    self.dropout_every_n_epochs: int = None
+    self.tag_dropout_rate: float = 0
+
     # augmentation
     flip_p = 0.5 if flip_aug else 0.0
     if color_aug:
@@ -113,10 +249,19 @@ class BaseDataset(torch.utils.data.Dataset):
 
     self.image_transforms = transforms.Compose([transforms.ToTensor(), transforms.Normalize([0.5], [0.5]), ])
 
-    self.image_data: dict[str, ImageInfo] = {}
+    self.image_data: Dict[str, ImageInfo] = {}
 
     self.replacements = {}
 
+  def set_current_epoch(self, epoch):
+    self.current_epoch = epoch
+
+  def set_caption_dropout(self, dropout_rate, dropout_every_n_epochs, tag_dropout_rate):
+    # コンストラクタで渡さないのはTextual Inversionで意識したくないから（ということにしておく）
+    self.dropout_rate = dropout_rate
+    self.dropout_every_n_epochs = dropout_every_n_epochs
+    self.tag_dropout_rate = tag_dropout_rate
+
   def set_tag_frequency(self, dir_name, captions):
     frequency_for_dir = self.tag_frequency.get(dir_name, {})
     self.tag_frequency[dir_name] = frequency_for_dir
@@ -134,27 +279,52 @@ class BaseDataset(torch.utils.data.Dataset):
     self.replacements[str_from] = str_to
 
   def process_caption(self, caption):
-    if self.shuffle_caption:
-      tokens = caption.strip().split(",")
-      if self.shuffle_keep_tokens is None:
-        random.shuffle(tokens)
-      else:
-        if len(tokens) > self.shuffle_keep_tokens:
-          keep_tokens = tokens[:self.shuffle_keep_tokens]
-          tokens = tokens[self.shuffle_keep_tokens:]
-          random.shuffle(tokens)
-          tokens = keep_tokens + tokens
-      caption = ",".join(tokens).strip()
+    # dropoutの決定：tag dropがこのメソッド内にあるのでここで行うのが良い
+    is_drop_out = self.dropout_rate > 0 and random.random() < self.dropout_rate
+    is_drop_out = is_drop_out or self.dropout_every_n_epochs and self.current_epoch % self.dropout_every_n_epochs == 0
 
-    for str_from, str_to in self.replacements.items():
-      if str_from == "":
-        # replace all
-        if type(str_to) == list:
-          caption = random.choice(str_to)
+    if is_drop_out:
+      caption = ""
+    else:
+      if self.shuffle_caption or self.tag_dropout_rate > 0:
+        def dropout_tags(tokens):
+          if self.tag_dropout_rate <= 0:
+            return tokens
+          l = []
+          for token in tokens:
+            if random.random() >= self.tag_dropout_rate:
+              l.append(token)
+          return l
+
+        tokens = [t.strip() for t in caption.strip().split(",")]
+        if self.shuffle_keep_tokens is None:
+          if self.shuffle_caption:
+            random.shuffle(tokens)
+
+          tokens = dropout_tags(tokens)
         else:
-          caption = str_to
-      else:
-        caption = caption.replace(str_from, str_to)
+          if len(tokens) > self.shuffle_keep_tokens:
+            keep_tokens = tokens[:self.shuffle_keep_tokens]
+            tokens = tokens[self.shuffle_keep_tokens:]
+
+            if self.shuffle_caption:
+              random.shuffle(tokens)
+
+            tokens = dropout_tags(tokens)
+
+            tokens = keep_tokens + tokens
+        caption = ", ".join(tokens)
+
+      # textual inversion対応
+      for str_from, str_to in self.replacements.items():
+        if str_from == "":
+          # replace all
+          if type(str_to) == list:
+            caption = random.choice(str_to)
+          else:
+            caption = str_to
+        else:
+          caption = caption.replace(str_from, str_to)
 
     return caption
 
@@ -215,64 +385,80 @@ class BaseDataset(torch.utils.data.Dataset):
     else:
       print("prepare dataset")
 
-    bucket_resos = self.bucket_resos
-    bucket_aspect_ratios = np.array(self.bucket_aspect_ratios)
-
-    # bucketを作成する
+    # bucketを作成し、画像をbucketに振り分ける
     if self.enable_bucket:
+      if self.bucket_manager is None:                         # fine tuningの場合でmetadataに定義がある場合は、すでに初期化済み
+        self.bucket_manager = BucketManager(self.bucket_no_upscale, (self.width, self.height),
+                                            self.min_bucket_reso, self.max_bucket_reso, self.bucket_reso_steps)
+        if not self.bucket_no_upscale:
+          self.bucket_manager.make_buckets()
+        else:
+          print("min_bucket_reso and max_bucket_reso are ignored if bucket_no_upscale is set, because bucket reso is defined by image size automatically / bucket_no_upscaleが指定された場合は、bucketの解像度は画像サイズから自動計算されるため、min_bucket_resoとmax_bucket_resoは無視されます")
+
       img_ar_errors = []
       for image_info in self.image_data.values():
-        # bucketを決める
         image_width, image_height = image_info.image_size
-        aspect_ratio = image_width / image_height
-        ar_errors = bucket_aspect_ratios - aspect_ratio
+        image_info.bucket_reso, image_info.resized_size, ar_error = self.bucket_manager.select_bucket(image_width, image_height)
 
-        bucket_id = np.abs(ar_errors).argmin()
-        image_info.bucket_reso = bucket_resos[bucket_id]
+        # print(image_info.image_key, image_info.bucket_reso)
+        img_ar_errors.append(abs(ar_error))
 
-        ar_error = ar_errors[bucket_id]
-        img_ar_errors.append(ar_error)
+      self.bucket_manager.sort()
     else:
+      self.bucket_manager = BucketManager(False, (self.width, self.height), None, None, None)
+      self.bucket_manager.set_predefined_resos([(self.width, self.height)])  # ひとつの固定サイズbucketのみ
       for image_info in self.image_data.values():
-        image_info.bucket_reso = bucket_resos[0]              # bucket_resos contains (width, height) only
-
-    # 画像をbucketに分割する
-    self.buckets: list[str] = [[] for _ in range(len(bucket_resos))]
-    reso_to_index = {}
-    for i, reso in enumerate(bucket_resos):
-      reso_to_index[reso] = i
+        image_width, image_height = image_info.image_size
+        image_info.bucket_reso, image_info.resized_size, _ = self.bucket_manager.select_bucket(image_width, image_height)
 
     for image_info in self.image_data.values():
-      bucket_index = reso_to_index[image_info.bucket_reso]
       for _ in range(image_info.num_repeats):
-        self.buckets[bucket_index].append(image_info.image_key)
+        self.bucket_manager.add_image(image_info.bucket_reso, image_info.image_key)
 
+    # bucket情報を表示、格納する
     if self.enable_bucket:
       self.bucket_info = {"buckets": {}}
       print("number of images (including repeats) / 各bucketの画像枚数（繰り返し回数を含む）")
-      for i, (reso, img_keys) in enumerate(zip(bucket_resos, self.buckets)):
-        self.bucket_info["buckets"][i] = {"resolution": reso, "count": len(img_keys)}
-        print(f"bucket {i}: resolution {reso}, count: {len(img_keys)}")
+      for i, (reso, bucket) in enumerate(zip(self.bucket_manager.resos, self.bucket_manager.buckets)):
+        count = len(bucket)
+        if count > 0:
+          self.bucket_info["buckets"][i] = {"resolution": reso, "count": len(bucket)}
+          print(f"bucket {i}: resolution {reso}, count: {len(bucket)}")
 
       img_ar_errors = np.array(img_ar_errors)
       mean_img_ar_error = np.mean(np.abs(img_ar_errors))
       self.bucket_info["mean_img_ar_error"] = mean_img_ar_error
       print(f"mean ar error (without repeats): {mean_img_ar_error}")
 
-    # 参照用indexを作る
-    self.buckets_indices: list(BucketBatchIndex) = []
-    for bucket_index, bucket in enumerate(self.buckets):
+    # データ参照用indexを作る。このindexはdatasetのshuffleに用いられる
+    self.buckets_indices: List(BucketBatchIndex) = []
+    for bucket_index, bucket in enumerate(self.bucket_manager.buckets):
       batch_count = int(math.ceil(len(bucket) / self.batch_size))
       for batch_index in range(batch_count):
-        self.buckets_indices.append(BucketBatchIndex(bucket_index, batch_index))
+        self.buckets_indices.append(BucketBatchIndex(bucket_index, self.batch_size, batch_index))
+
+      # ↓以下はbucketごとのbatch件数があまりにも増えて混乱を招くので元に戻す
+      # 　学習時はステップ数がランダムなので、同一画像が同一batch内にあってもそれほど悪影響はないであろう、と考えられる
+      #
+      # # bucketが細分化されることにより、ひとつのbucketに一種類の画像のみというケースが増え、つまりそれは
+      # # ひとつのbatchが同じ画像で占められることになるので、さすがに良くないであろう
+      # # そのためバッチサイズを画像種類までに制限する
+      # # ただそれでも同一画像が同一バッチに含まれる可能性はあるので、繰り返し回数が少ないほうがshuffleの品質は良くなることは間違いない？
+      # # TO DO 正則化画像をepochまたがりで利用する仕組み
+      # num_of_image_types = len(set(bucket))
+      # bucket_batch_size = min(self.batch_size, num_of_image_types)
+      # batch_count = int(math.ceil(len(bucket) / bucket_batch_size))
+      # # print(bucket_index, num_of_image_types, bucket_batch_size, batch_count)
+      # for batch_index in range(batch_count):
+      #   self.buckets_indices.append(BucketBatchIndex(bucket_index, bucket_batch_size, batch_index))
+      # ↑ここまで
 
     self.shuffle_buckets()
     self._length = len(self.buckets_indices)
 
   def shuffle_buckets(self):
     random.shuffle(self.buckets_indices)
-    for bucket in self.buckets:
-      random.shuffle(bucket)
+    self.bucket_manager.shuffle()
 
   def load_image(self, image_path):
     image = Image.open(image_path)
@@ -281,28 +467,30 @@ class BaseDataset(torch.utils.data.Dataset):
     img = np.array(image, np.uint8)
     return img
 
-  def resize_and_trim(self, image, reso):
+  def trim_and_resize_if_required(self, image, reso, resized_size):
     image_height, image_width = image.shape[0:2]
-    ar_img = image_width / image_height
-    ar_reso = reso[0] / reso[1]
-    if ar_img > ar_reso:                   # 横が長い→縦を合わせる
-      scale = reso[1] / image_height
-    else:
-      scale = reso[0] / image_width
-    resized_size = (int(image_width * scale + .5), int(image_height * scale + .5))
 
-    image = cv2.resize(image, resized_size, interpolation=cv2.INTER_AREA)       # INTER_AREAでやりたいのでcv2でリサイズ
-    if resized_size[0] > reso[0]:
-      trim_size = resized_size[0] - reso[0]
-      image = image[:, trim_size//2:trim_size//2 + reso[0]]
-    elif resized_size[1] > reso[1]:
-      trim_size = resized_size[1] - reso[1]
-      image = image[trim_size//2:trim_size//2 + reso[1]]
-    assert image.shape[0] == reso[1] and image.shape[1] == reso[0],  \
-        f"internal error, illegal trimmed size: {image.shape}, {reso}"
+    if image_width != resized_size[0] or image_height != resized_size[1]:
+      # リサイズする
+      image = cv2.resize(image, resized_size, interpolation=cv2.INTER_AREA)       # INTER_AREAでやりたいのでcv2でリサイズ
+
+    image_height, image_width = image.shape[0:2]
+    if image_width > reso[0]:
+      trim_size = image_width - reso[0]
+      p = trim_size // 2 if not self.random_crop else random.randint(0, trim_size)
+      # print("w", trim_size, p)
+      image = image[:, p:p + reso[0]]
+    if image_height > reso[1]:
+      trim_size = image_height - reso[1]
+      p = trim_size // 2 if not self.random_crop else random.randint(0, trim_size)
+      # print("h", trim_size, p)
+      image = image[p:p + reso[1]]
+
+    assert image.shape[0] == reso[1] and image.shape[1] == reso[0], f"internal error, illegal trimmed size: {image.shape}, {reso}"
     return image
 
   def cache_latents(self, vae):
+    # TODO ここを高速化したい
     print("caching latents.")
     for info in tqdm(self.image_data.values()):
       if info.latents_npz is not None:
@@ -314,7 +502,7 @@ class BaseDataset(torch.utils.data.Dataset):
         continue
 
       image = self.load_image(info.absolute_path)
-      image = self.resize_and_trim(image, info.bucket_reso)
+      image = self.trim_and_resize_if_required(image, info.bucket_reso, info.resized_size)
 
       img_tensor = self.image_transforms(image)
       img_tensor = img_tensor.unsqueeze(0).to(device=vae.device, dtype=vae.dtype)
@@ -404,8 +592,9 @@ class BaseDataset(torch.utils.data.Dataset):
     if index == 0:
       self.shuffle_buckets()
 
-    bucket = self.buckets[self.buckets_indices[index].bucket_index]
-    image_index = self.buckets_indices[index].batch_index * self.batch_size
+    bucket = self.bucket_manager.buckets[self.buckets_indices[index].bucket_index]
+    bucket_batch_size = self.buckets_indices[index].bucket_batch_size
+    image_index = self.buckets_indices[index].batch_index * bucket_batch_size
 
     loss_weights = []
     captions = []
@@ -413,7 +602,7 @@ class BaseDataset(torch.utils.data.Dataset):
     latents_list = []
     images = []
 
-    for image_key in bucket[image_index:image_index + self.batch_size]:
+    for image_key in bucket[image_index:image_index + bucket_batch_size]:
       image_info = self.image_data[image_key]
       loss_weights.append(self.prior_loss_weight if image_info.is_reg else 1.0)
 
@@ -431,7 +620,7 @@ class BaseDataset(torch.utils.data.Dataset):
         im_h, im_w = img.shape[0:2]
 
         if self.enable_bucket:
-          img = self.resize_and_trim(img, image_info.bucket_reso)
+          img = self.trim_and_resize_if_required(img, image_info.bucket_reso, image_info.resized_size)
         else:
           if face_cx > 0:                   # 顔位置情報あり
             img = self.crop_target(img, face_cx, face_cy, face_w, face_h)
@@ -488,7 +677,7 @@ class BaseDataset(torch.utils.data.Dataset):
 
 
 class DreamBoothDataset(BaseDataset):
-  def __init__(self, batch_size, train_data_dir, reg_data_dir, tokenizer, max_token_length, caption_extension, shuffle_caption, shuffle_keep_tokens, resolution, enable_bucket, min_bucket_reso, max_bucket_reso, prior_loss_weight, flip_aug, color_aug, face_crop_aug_range, random_crop, debug_dataset) -> None:
+  def __init__(self, batch_size, train_data_dir, reg_data_dir, tokenizer, max_token_length, caption_extension, shuffle_caption, shuffle_keep_tokens, resolution, enable_bucket, min_bucket_reso, max_bucket_reso, bucket_reso_steps, bucket_no_upscale, prior_loss_weight, flip_aug, color_aug, face_crop_aug_range, random_crop, debug_dataset) -> None:
     super().__init__(tokenizer, max_token_length, shuffle_caption, shuffle_keep_tokens,
                      resolution, flip_aug, color_aug, face_crop_aug_range, random_crop, debug_dataset)
 
@@ -503,13 +692,15 @@ class DreamBoothDataset(BaseDataset):
     if self.enable_bucket:
       assert min(resolution) >= min_bucket_reso, f"min_bucket_reso must be equal or less than resolution / min_bucket_resoは最小解像度より大きくできません。解像度を大きくするかmin_bucket_resoを小さくしてください"
       assert max(resolution) <= max_bucket_reso, f"max_bucket_reso must be equal or greater than resolution / max_bucket_resoは最大解像度より小さくできません。解像度を小さくするかmin_bucket_resoを大きくしてください"
-      self.bucket_resos, self.bucket_aspect_ratios = model_util.make_bucket_resolutions(
-          (self.width, self.height), min_bucket_reso, max_bucket_reso)
       self.min_bucket_reso = min_bucket_reso
       self.max_bucket_reso = max_bucket_reso
+      self.bucket_reso_steps = bucket_reso_steps
+      self.bucket_no_upscale = bucket_no_upscale
     else:
-      self.bucket_resos = [(self.width, self.height)]
-      self.bucket_aspect_ratios = [self.width / self.height]
+      self.min_bucket_reso = None
+      self.max_bucket_reso = None
+      self.bucket_reso_steps = None                              # この情報は使われない
+      self.bucket_no_upscale = False
 
     def read_caption(img_path):
       # captionの候補ファイル名を作る
@@ -580,7 +771,7 @@ class DreamBoothDataset(BaseDataset):
     num_reg_images = 0
     if reg_data_dir:
       print("prepare reg images.")
-      reg_infos: list[ImageInfo] = []
+      reg_infos: List[ImageInfo] = []
 
       reg_dirs = os.listdir(reg_data_dir)
       for dir in reg_dirs:
@@ -619,7 +810,7 @@ class DreamBoothDataset(BaseDataset):
 
 
 class FineTuningDataset(BaseDataset):
-  def __init__(self, json_file_name, batch_size, train_data_dir, tokenizer, max_token_length, shuffle_caption, shuffle_keep_tokens, resolution, enable_bucket, min_bucket_reso, max_bucket_reso, flip_aug, color_aug, face_crop_aug_range, random_crop, dataset_repeats, debug_dataset) -> None:
+  def __init__(self, json_file_name, batch_size, train_data_dir, tokenizer, max_token_length, shuffle_caption, shuffle_keep_tokens, resolution, enable_bucket, min_bucket_reso, max_bucket_reso, bucket_reso_steps, bucket_no_upscale, flip_aug, color_aug, face_crop_aug_range, random_crop, dataset_repeats, debug_dataset) -> None:
     super().__init__(tokenizer, max_token_length, shuffle_caption, shuffle_keep_tokens,
                      resolution, flip_aug, color_aug, face_crop_aug_range, random_crop, debug_dataset)
 
@@ -658,7 +849,7 @@ class FineTuningDataset(BaseDataset):
       image_info = ImageInfo(image_key, dataset_repeats, caption, False, abs_path)
       image_info.image_size = img_md.get('train_resolution')
 
-      if not self.color_aug:
+      if not self.color_aug and not self.random_crop:
         # if npz exists, use them
         image_info.latents_npz, image_info.latents_npz_flipped = self.image_key_to_npz_file(image_key)
 
@@ -666,11 +857,13 @@ class FineTuningDataset(BaseDataset):
     self.num_train_images = len(metadata) * dataset_repeats
     self.num_reg_images = 0
 
+    # TODO do not record tag freq when no tag
     self.set_tag_frequency(os.path.basename(json_file_name), tags_list)
     self.dataset_dirs_info[os.path.basename(json_file_name)] = {"n_repeats": dataset_repeats, "img_count": len(metadata)}
 
     # check existence of all npz files
-    if not self.color_aug:
+    use_npz_latents = not (self.color_aug or self.random_crop)
+    if use_npz_latents:
       npz_any = False
       npz_all = True
       for image_info in self.image_data.values():
@@ -685,13 +878,15 @@ class FineTuningDataset(BaseDataset):
           break
 
       if not npz_any:
-        print(f"npz file does not exist. make latents with VAE / npzファイルが見つからないためVAEを使ってlatentsを取得します")
+        use_npz_latents = False
+        print(f"npz file does not exist. ignore npz files / npzファイルが見つからないためnpzファイルを無視します")
       elif not npz_all:
+        use_npz_latents = False
         print(f"some of npz file does not exist. ignore npz files / いくつかのnpzファイルが見つからないためnpzファイルを無視します")
         if self.flip_aug:
           print("maybe no flipped files / 反転されたnpzファイルがないのかもしれません")
-        for image_info in self.image_data.values():
-          image_info.latents_npz = image_info.latents_npz_flipped = None
+    # else:
+    #   print("npz files are not used with color_aug and/or random_crop / color_augまたはrandom_cropが指定されているためnpzファイルは使用されません")
 
     # check min/max bucket size
     sizes = set()
@@ -705,30 +900,34 @@ class FineTuningDataset(BaseDataset):
       resos.add(tuple(image_info.image_size))
 
     if sizes is None:
+      if use_npz_latents:
+        use_npz_latents = False
+        print(f"npz files exist, but no bucket info in metadata. ignore npz files / メタデータにbucket情報がないためnpzファイルを無視します")
+
       assert resolution is not None, "if metadata doesn't have bucket info, resolution is required / メタデータにbucket情報がない場合はresolutionを指定してください"
 
       self.enable_bucket = enable_bucket
       if self.enable_bucket:
-        assert min(resolution) >= min_bucket_reso, f"min_bucket_reso must be equal or less than resolution / min_bucket_resoは最小解像度より大きくできません。解像度を大きくするかmin_bucket_resoを小さくしてください"
-        assert max(resolution) <= max_bucket_reso, f"max_bucket_reso must be equal or greater than resolution / max_bucket_resoは最大解像度より小さくできません。解像度を小さくするかmin_bucket_resoを大きくしてください"
-        self.bucket_resos, self.bucket_aspect_ratios = model_util.make_bucket_resolutions(
-            (self.width, self.height), min_bucket_reso, max_bucket_reso)
         self.min_bucket_reso = min_bucket_reso
         self.max_bucket_reso = max_bucket_reso
-      else:
-        self.bucket_resos = [(self.width, self.height)]
-        self.bucket_aspect_ratios = [self.width / self.height]
+        self.bucket_reso_steps = bucket_reso_steps
+        self.bucket_no_upscale = bucket_no_upscale
     else:
       if not enable_bucket:
         print("metadata has bucket info, enable bucketing / メタデータにbucket情報があるためbucketを有効にします")
       print("using bucket info in metadata / メタデータ内のbucket情報を使います")
       self.enable_bucket = True
-      self.bucket_resos = list(resos)
-      self.bucket_resos.sort()
-      self.bucket_aspect_ratios = [w / h for w, h in self.bucket_resos]
 
-      self.min_bucket_reso = min([min(reso) for reso in resos])
-      self.max_bucket_reso = max([max(reso) for reso in resos])
+      assert not bucket_no_upscale, "if metadata has bucket info, bucket reso is precalculated, so bucket_no_upscale cannot be used / メタデータ内にbucket情報がある場合はbucketの解像度は計算済みのため、bucket_no_upscaleは使えません"
+
+      # bucket情報を初期化しておく、make_bucketsで再作成しない
+      self.bucket_manager = BucketManager(False, None, None, None, None)
+      self.bucket_manager.set_predefined_resos(resos)
+
+    # npz情報をきれいにしておく
+    if not use_npz_latents:
+      for image_info in self.image_data.values():
+        image_info.latents_npz = image_info.latents_npz_flipped = None
 
   def image_key_to_npz_file(self, image_key):
     base_name = os.path.splitext(image_key)[0]
@@ -757,16 +956,19 @@ class FineTuningDataset(BaseDataset):
 def debug_dataset(train_dataset, show_input_ids=False):
   print(f"Total dataset length (steps) / データセットの長さ（ステップ数）: {len(train_dataset)}")
   print("Escape for exit. / Escキーで中断、終了します")
+
+  train_dataset.set_current_epoch(1)
   k = 0
-  for example in train_dataset:
+  for i, example in enumerate(train_dataset):
     if example['latents'] is not None:
-      print("sample has latents from npz file")
+      print(f"sample has latents from npz file: {example['latents'].size()}")
     for j, (ik, cap, lw, iid) in enumerate(zip(example['image_keys'], example['captions'], example['loss_weights'], example['input_ids'])):
-      print(f'{ik}, size: {train_dataset.image_data[ik].image_size}, caption: "{cap}", loss weight: {lw}')
+      print(f'{ik}, size: {train_dataset.image_data[ik].image_size}, loss weight: {lw}, caption: "{cap}"')
       if show_input_ids:
         print(f"input ids: {iid}")
       if example['images'] is not None:
         im = example['images'][j]
+        print(f"image size: {im.size()}")
         im = ((im.numpy() + 1.0) * 127.5).astype(np.uint8)
         im = np.transpose(im, (1, 2, 0))                # c,H,W -> H,W,c
         im = im[:, :, ::-1]                             # RGB -> BGR (OpenCV)
@@ -776,7 +978,7 @@ def debug_dataset(train_dataset, show_input_ids=False):
         cv2.destroyAllWindows()
         if k == 27:
           break
-    if k == 27 or example['images'] is None:
+    if k == 27 or (example['images'] is None and i >= 8):
       break
 
 
@@ -899,6 +1101,13 @@ def addnet_hash_safetensors(b):
   return hash_sha256.hexdigest()
 
 
+def get_git_revision_hash() -> str:
+  try:
+    return subprocess.check_output(['git', 'rev-parse', 'HEAD'], cwd=os.path.dirname(__file__)).decode('ascii').strip()
+  except:
+    return "(unknown)"
+
+
 # flash attention forwards and backwards
 
 # https://arxiv.org/abs/2205.14135
@@ -1180,6 +1389,8 @@ def add_training_arguments(parser: argparse.ArgumentParser, support_dreambooth:
                       help="max token length of text encoder (default for 75, 150 or 225) / text encoderのトークンの最大長（未指定で75、150または225が指定可）")
   parser.add_argument("--use_8bit_adam", action="store_true",
                       help="use 8bit Adam optimizer (requires bitsandbytes) / 8bit Adamオプティマイザを使う（bitsandbytesのインストールが必要）")
+  parser.add_argument("--use_lion_optimizer", action="store_true",
+                      help="use Lion optimizer (requires lion-pytorch) / Lionオプティマイザを使う（ lion-pytorch のインストールが必要）")
   parser.add_argument("--mem_eff_attn", action="store_true",
                       help="use memory efficient attention for CrossAttention / CrossAttentionに省メモリ版attentionを使う")
   parser.add_argument("--xformers", action="store_true",
@@ -1212,6 +1423,10 @@ def add_training_arguments(parser: argparse.ArgumentParser, support_dreambooth:
                       help="scheduler to use for learning rate / 学習率のスケジューラ: linear, cosine, cosine_with_restarts, polynomial, constant (default), constant_with_warmup")
   parser.add_argument("--lr_warmup_steps", type=int, default=0,
                       help="Number of steps for the warmup in the lr scheduler (default is 0) / 学習率のスケジューラをウォームアップするステップ数（デフォルト0）")
+  parser.add_argument("--noise_offset", type=float, default=None,
+                      help="enable noise offset with this value (if enabled, around 0.1 is recommended) / Noise offsetを有効にしてこの値を設定する（有効にする場合は0.1程度を推奨）")
+  parser.add_argument("--lowram", action="store_true",
+                      help="enable low RAM optimization. e.g. load models to VRAM instead of RAM (for machines which have bigger VRAM than RAM such as Colab and Kaggle) / メインメモリが少ない環境向け最適化を有効にする。たとえばVRAMにモデルを読み込むなど（ColabやKaggleなどRAMに比べてVRAMが多い環境向け）")
 
   if support_dreambooth:
     # DreamBooth training
@@ -1226,7 +1441,7 @@ def verify_training_args(args: argparse.Namespace):
     print("v2 with clip_skip will be unexpected / v2でclip_skipを使用することは想定されていません")
 
 
-def add_dataset_arguments(parser: argparse.ArgumentParser, support_dreambooth: bool, support_caption: bool):
+def add_dataset_arguments(parser: argparse.ArgumentParser, support_dreambooth: bool, support_caption: bool, support_caption_dropout: bool):
   # dataset common
   parser.add_argument("--train_data_dir", type=str, default=None, help="directory for train images / 学習画像データのディレクトリ")
   parser.add_argument("--shuffle_caption", action="store_true",
@@ -1252,6 +1467,20 @@ def add_dataset_arguments(parser: argparse.ArgumentParser, support_dreambooth: b
                       help="enable buckets for multi aspect ratio training / 複数解像度学習のためのbucketを有効にする")
   parser.add_argument("--min_bucket_reso", type=int, default=256, help="minimum resolution for buckets / bucketの最小解像度")
   parser.add_argument("--max_bucket_reso", type=int, default=1024, help="maximum resolution for buckets / bucketの最大解像度")
+  parser.add_argument("--bucket_reso_steps", type=int, default=64,
+                      help="steps of resolution for buckets, divisible by 8 is recommended / bucketの解像度の単位、8で割り切れる値を推奨します")
+  parser.add_argument("--bucket_no_upscale", action="store_true",
+                      help="make bucket for each image without upscaling / 画像を拡大せずbucketを作成します")
+
+  if support_caption_dropout:
+    # Textual Inversion はcaptionのdropoutをsupportしない
+    # いわゆるtensorのDropoutと紛らわしいのでprefixにcaptionを付けておく　every_n_epochsは他と平仄を合わせてdefault Noneに
+    parser.add_argument("--caption_dropout_rate", type=float, default=0,
+                        help="Rate out dropout caption(0.0~1.0) / captionをdropoutする割合")
+    parser.add_argument("--caption_dropout_every_n_epochs", type=int, default=None,
+                        help="Dropout all captions every N epochs / captionを指定エポックごとにdropoutする")
+    parser.add_argument("--caption_tag_dropout_rate", type=float, default=0,
+                        help="Rate out dropout comma separated tokens(0.0~1.0) / カンマ区切りのタグをdropoutする割合")
 
   if support_dreambooth:
     # DreamBooth dataset
@@ -1283,6 +1512,7 @@ def prepare_dataset_args(args: argparse.Namespace, support_metadata: bool):
 
   if args.cache_latents:
     assert not args.color_aug, "when caching latents, color_aug cannot be used / latentをキャッシュするときはcolor_augは使えません"
+    assert not args.random_crop, "when caching latents, random_crop cannot be used / latentをキャッシュするときはrandom_cropは使えません"
 
   # assert args.resolution is not None, f"resolution is required / resolution（解像度）を指定してください"
   if args.resolution is not None:
@@ -1294,14 +1524,14 @@ def prepare_dataset_args(args: argparse.Namespace, support_metadata: bool):
 
   if args.face_crop_aug_range is not None:
     args.face_crop_aug_range = tuple([float(r) for r in args.face_crop_aug_range.split(',')])
-    assert len(args.face_crop_aug_range) == 2, \
+    assert len(args.face_crop_aug_range) == 2 and args.face_crop_aug_range[0] <= args.face_crop_aug_range[1], \
         f"face_crop_aug_range must be two floats / face_crop_aug_rangeは'下限,上限'で指定してください: {args.face_crop_aug_range}"
   else:
     args.face_crop_aug_range = None
 
   if support_metadata:
-    if args.in_json is not None and args.color_aug:
-      print(f"latents in npz is ignored when color_aug is True / color_augを有効にした場合、npzファイルのlatentsは無視されます")
+    if args.in_json is not None and (args.color_aug or args.random_crop):
+      print(f"latents in npz is ignored when color_aug or random_crop is True / color_augまたはrandom_cropを有効にした場合、npzファイルのlatentsは無視されます")
 
 
 def load_tokenizer(args: argparse.Namespace):
@@ -1404,9 +1634,6 @@ def get_hidden_states(args: argparse.Namespace, input_ids, tokenizer, text_encod
   else:
     enc_out = text_encoder(input_ids, output_hidden_states=True, return_dict=True)
     encoder_hidden_states = enc_out['hidden_states'][-args.clip_skip]
-    if weight_dtype is not None:
-      # this is required for additional network training
-      encoder_hidden_states = encoder_hidden_states.to(weight_dtype)
     encoder_hidden_states = text_encoder.text_model.final_layer_norm(encoder_hidden_states)
 
   # bs*3, 77, 768 or 1024
@@ -1433,6 +1660,10 @@ def get_hidden_states(args: argparse.Namespace, input_ids, tokenizer, text_encod
       states_list.append(encoder_hidden_states[:, -1].unsqueeze(1))                         # <EOS>
       encoder_hidden_states = torch.cat(states_list, dim=1)
 
+  if weight_dtype is not None:
+    # this is required for additional network training
+    encoder_hidden_states = encoder_hidden_states.to(weight_dtype)
+
   return encoder_hidden_states
 
 

networks/lora.py

@@ -5,6 +5,7 @@
 
 import math
 import os
+from typing import List
 import torch
 
 from library import train_util
@@ -98,7 +99,7 @@ class LoRANetwork(torch.nn.Module):
     self.alpha = alpha
 
     # create module instances
-    def create_modules(prefix, root_module: torch.nn.Module, target_replace_modules) -> list[LoRAModule]:
+    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:

networks/lora_interrogator.py

@@ -0,0 +1,122 @@
+
+
+from tqdm import tqdm
+from library import model_util
+import argparse
+from transformers import CLIPTokenizer
+import torch
+
+import library.model_util as model_util
+import lora
+
+TOKENIZER_PATH = "openai/clip-vit-large-patch14"
+V2_STABLE_DIFFUSION_PATH = "stabilityai/stable-diffusion-2"     # ここからtokenizerだけ使う
+
+DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+
+def interrogate(args):
+  # いろいろ準備する
+  print(f"loading SD model: {args.sd_model}")
+  text_encoder, vae, unet = model_util.load_models_from_stable_diffusion_checkpoint(args.v2, args.sd_model)
+
+  print(f"loading LoRA: {args.model}")
+  network = lora.create_network_from_weights(1.0, args.model, vae, text_encoder, unet)
+
+  # text encoder向けの重みがあるかチェックする：本当はlora側でやるのがいい
+  has_te_weight = False
+  for key in network.weights_sd.keys():
+    if 'lora_te' in key:
+      has_te_weight = True
+      break
+  if not has_te_weight:
+    print("This LoRA does not have modules for Text Encoder, cannot interrogate / このLoRAはText Encoder向けのモジュールがないため調査できません")
+    return
+  del vae
+
+  print("loading tokenizer")
+  if args.v2:
+    tokenizer: CLIPTokenizer = CLIPTokenizer.from_pretrained(V2_STABLE_DIFFUSION_PATH, subfolder="tokenizer")
+  else:
+    tokenizer: CLIPTokenizer = CLIPTokenizer.from_pretrained(TOKENIZER_PATH)  # , model_max_length=max_token_length + 2)
+
+  text_encoder.to(DEVICE)
+  text_encoder.eval()
+  unet.to(DEVICE)
+  unet.eval()               # U-Netは呼び出さないので不要だけど
+
+  # トークンをひとつひとつ当たっていく
+  token_id_start = 0
+  token_id_end = max(tokenizer.all_special_ids)
+  print(f"interrogate tokens are: {token_id_start} to {token_id_end}")
+
+  def get_all_embeddings(text_encoder):
+    embs = []
+    with torch.no_grad():
+      for token_id in tqdm(range(token_id_start, token_id_end + 1, args.batch_size)):
+        batch = []
+        for tid in range(token_id, min(token_id_end + 1, token_id + args.batch_size)):
+          tokens = [tokenizer.bos_token_id, tid, tokenizer.eos_token_id]
+          # tokens = [tid]                                                    # こちらは結果がいまひとつ
+          batch.append(tokens)
+
+        # batch_embs = text_encoder(torch.tensor(batch).to(DEVICE))[0].to("cpu")  # bos/eosも含めたほうが差が出るようだ [:, 1]
+        # clip skip対応
+        batch = torch.tensor(batch).to(DEVICE)
+        if args.clip_skip is None:
+          encoder_hidden_states = text_encoder(batch)[0]
+        else:
+          enc_out = text_encoder(batch, output_hidden_states=True, return_dict=True)
+          encoder_hidden_states = enc_out['hidden_states'][-args.clip_skip]
+          encoder_hidden_states = text_encoder.text_model.final_layer_norm(encoder_hidden_states)
+        encoder_hidden_states = encoder_hidden_states.to("cpu")
+
+        embs.extend(encoder_hidden_states)
+    return torch.stack(embs)
+
+  print("get original text encoder embeddings.")
+  orig_embs = get_all_embeddings(text_encoder)
+
+  network.apply_to(text_encoder, unet, True, len(network.unet_loras) > 0)
+  network.to(DEVICE)
+  network.eval()
+
+  print("You can ignore warning messages start with '_IncompatibleKeys' (LoRA model does not have alpha because trained by older script) / '_IncompatibleKeys'の警告は無視して構いません（以前のスクリプトで学習されたLoRAモデルのためalphaの定義がありません）")
+  print("get text encoder embeddings with lora.")
+  lora_embs = get_all_embeddings(text_encoder)
+
+  # 比べる：とりあえず単純に差分の絶対値で
+  print("comparing...")
+  diffs = {}
+  for i, (orig_emb, lora_emb) in enumerate(zip(orig_embs, tqdm(lora_embs))):
+    diff = torch.mean(torch.abs(orig_emb - lora_emb))
+    # diff = torch.mean(torch.cosine_similarity(orig_emb, lora_emb, dim=1))       # うまく検出できない
+    diff = float(diff.detach().to('cpu').numpy())
+    diffs[token_id_start + i] = diff
+
+  diffs_sorted = sorted(diffs.items(), key=lambda x: -x[1])
+
+  # 結果を表示する
+  print("top 100:")
+  for i, (token, diff) in enumerate(diffs_sorted[:100]):
+    # if diff < 1e-6:
+    #   break
+    string = tokenizer.convert_tokens_to_string(tokenizer.convert_ids_to_tokens([token]))
+    print(f"[{i:3d}]: {token:5d} {string:<20s}: {diff:.5f}")
+
+
+if __name__ == '__main__':
+  parser = argparse.ArgumentParser()
+  parser.add_argument("--v2", action='store_true',
+                      help='load Stable Diffusion v2.x model / Stable Diffusion 2.xのモデルを読み込む')
+  parser.add_argument("--sd_model", type=str, default=None,
+                      help="Stable Diffusion model to load: ckpt or safetensors file / 読み込むSDのモデル、ckptまたはsafetensors")
+  parser.add_argument("--model", type=str, default=None,
+                      help="LoRA model to interrogate: ckpt or safetensors file / 調査するLoRAモデル、ckptまたはsafetensors")
+  parser.add_argument("--batch_size", type=int, default=16,
+                      help="batch size for processing with Text Encoder / Text Encoderで処理するときのバッチサイズ")
+  parser.add_argument("--clip_skip", type=int, default=None,
+                      help="use output of nth layer from back of text encoder (n>=1) / text encoderの後ろからn番目の層の出力を用いる（nは1以上）")
+
+  args = parser.parse_args()
+  interrogate(args)

networks/merge_lora.py

@@ -1,5 +1,5 @@
 
-
+import math
 import argparse
 import os
 import torch
@@ -85,43 +85,76 @@ def merge_to_sd_model(text_encoder, unet, models, ratios, merge_dtype):
           weight = weight + ratio * (up_weight @ down_weight) * scale
         else:
           # conv2d
-          weight = weight + ratio * (up_weight.squeeze(3).squeeze(2) @ down_weight.squeeze(3).squeeze(2)).unsqueeze(2).unsqueeze(3) * scale
+          weight = weight + ratio * (up_weight.squeeze(3).squeeze(2) @ down_weight.squeeze(3).squeeze(2)
+                                     ).unsqueeze(2).unsqueeze(3) * scale
 
         module.weight = torch.nn.Parameter(weight)
 
 
 def merge_lora_models(models, ratios, merge_dtype):
-  merged_sd = {}
+  base_alphas = {}                          # alpha for merged model
+  base_dims = {}
 
-  alpha = None
-  dim = None
+  merged_sd = {}
   for model, ratio in zip(models, ratios):
     print(f"loading: {model}")
     lora_sd = load_state_dict(model, merge_dtype)
 
+    # get alpha and dim
+    alphas = {}                             # alpha for current model
+    dims = {}                               # dims for current model
+    for key in lora_sd.keys():
+      if 'alpha' in key:
+        lora_module_name = key[:key.rfind(".alpha")]
+        alpha = float(lora_sd[key].detach().numpy())
+        alphas[lora_module_name] = alpha
+        if lora_module_name not in base_alphas:
+          base_alphas[lora_module_name] = alpha
+      elif "lora_down" in key:
+        lora_module_name = key[:key.rfind(".lora_down")]
+        dim = lora_sd[key].size()[0]
+        dims[lora_module_name] = dim
+        if lora_module_name not in base_dims:
+          base_dims[lora_module_name] = dim
+
+    for lora_module_name in dims.keys():
+      if lora_module_name not in alphas:
+        alpha = dims[lora_module_name]
+        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
     print(f"merging...")
     for key in lora_sd.keys():
       if 'alpha' in key:
-        if key in merged_sd:
-          assert merged_sd[key] == lora_sd[key], f"alpha mismatch / alphaが異なる場合、現時点ではマージできません"
-        else:
-          alpha = lora_sd[key].detach().numpy()
-          merged_sd[key] = lora_sd[key]
+        continue
+
+      lora_module_name = key[:key.rfind(".lora_")]
+
+      base_alpha = base_alphas[lora_module_name]
+      alpha = alphas[lora_module_name]
+
+      scale = math.sqrt(alpha / base_alpha) * ratio
+
+      if key in merged_sd:
+        assert merged_sd[key].size() == lora_sd[key].size(
+        ), f"weights shape mismatch merging v1 and v2, different dims? / 重みのサイズが合いません。v1とv2、または次元数の異なるモデルはマージできません"
+        merged_sd[key] = merged_sd[key] + lora_sd[key] * scale
       else:
-        if key in merged_sd:
-          assert merged_sd[key].size() == lora_sd[key].size(
-          ), f"weights shape mismatch merging v1 and v2, different dims? / 重みのサイズが合いません。v1とv2、または次元数の異なるモデルはマージできません"
-          merged_sd[key] = merged_sd[key] + lora_sd[key] * ratio
-        else:
-          if "lora_down" in key:
-            dim = lora_sd[key].size()[0]
-          merged_sd[key] = lora_sd[key] * ratio
+        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"
+    merged_sd[key] = torch.tensor(alpha)
 
-  print(f"dim (rank): {dim}, alpha: {alpha}")
-  if alpha is None:
-    alpha = dim
+  print("merged model")
+  print(f"dim: {list(set(base_dims.values()))}, alpha: {list(set(base_alphas.values()))}")
 
-  return merged_sd, dim, alpha
+  return merged_sd
 
 
 def merge(args):
@@ -152,7 +185,7 @@ def merge(args):
     model_util.save_stable_diffusion_checkpoint(args.v2, args.save_to, text_encoder, unet,
                                                 args.sd_model, 0, 0, save_dtype, vae)
   else:
-    state_dict, _, _ = merge_lora_models(args.models, args.ratios, merge_dtype)
+    state_dict = merge_lora_models(args.models, args.ratios, merge_dtype)
 
     print(f"saving model to: {args.save_to}")
     save_to_file(args.save_to, state_dict, state_dict, save_dtype)

networks/merge_lora_old.py

@@ -0,0 +1,179 @@
+
+
+import argparse
+import os
+import torch
+from safetensors.torch import load_file, save_file
+import library.model_util as model_util
+import lora
+
+
+def load_state_dict(file_name, dtype):
+  if os.path.splitext(file_name)[1] == '.safetensors':
+    sd = load_file(file_name)
+  else:
+    sd = torch.load(file_name, map_location='cpu')
+  for key in list(sd.keys()):
+    if type(sd[key]) == torch.Tensor:
+      sd[key] = sd[key].to(dtype)
+  return sd
+
+
+def save_to_file(file_name, model, 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)
+  else:
+    torch.save(model, file_name)
+
+
+def merge_to_sd_model(text_encoder, unet, models, ratios, merge_dtype):
+  text_encoder.to(merge_dtype)
+  unet.to(merge_dtype)
+
+  # create module map
+  name_to_module = {}
+  for i, root_module in enumerate([text_encoder, unet]):
+    if i == 0:
+      prefix = lora.LoRANetwork.LORA_PREFIX_TEXT_ENCODER
+      target_replace_modules = lora.LoRANetwork.TEXT_ENCODER_TARGET_REPLACE_MODULE
+    else:
+      prefix = lora.LoRANetwork.LORA_PREFIX_UNET
+      target_replace_modules = lora.LoRANetwork.UNET_TARGET_REPLACE_MODULE
+
+    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)):
+            lora_name = prefix + '.' + name + '.' + child_name
+            lora_name = lora_name.replace('.', '_')
+            name_to_module[lora_name] = child_module
+
+  for model, ratio in zip(models, ratios):
+    print(f"loading: {model}")
+    lora_sd = load_state_dict(model, merge_dtype)
+
+    print(f"merging...")
+    for key in lora_sd.keys():
+      if "lora_down" in key:
+        up_key = key.replace("lora_down", "lora_up")
+        alpha_key = key[:key.index("lora_down")] + 'alpha'
+
+        # find original module for this lora
+        module_name = '.'.join(key.split('.')[:-2])               # remove trailing ".lora_down.weight"
+        if module_name not in name_to_module:
+          print(f"no module found for LoRA weight: {key}")
+          continue
+        module = name_to_module[module_name]
+        # print(f"apply {key} to {module}")
+
+        down_weight = lora_sd[key]
+        up_weight = lora_sd[up_key]
+
+        dim = down_weight.size()[0]
+        alpha = lora_sd.get(alpha_key, dim)
+        scale = alpha / dim
+
+        # W <- W + U * D
+        weight = module.weight
+        if len(weight.size()) == 2:
+          # linear
+          weight = weight + ratio * (up_weight @ down_weight) * scale
+        else:
+          # conv2d
+          weight = weight + ratio * (up_weight.squeeze(3).squeeze(2) @ down_weight.squeeze(3).squeeze(2)).unsqueeze(2).unsqueeze(3) * scale
+
+        module.weight = torch.nn.Parameter(weight)
+
+
+def merge_lora_models(models, ratios, merge_dtype):
+  merged_sd = {}
+
+  alpha = None
+  dim = None
+  for model, ratio in zip(models, ratios):
+    print(f"loading: {model}")
+    lora_sd = load_state_dict(model, merge_dtype)
+
+    print(f"merging...")
+    for key in lora_sd.keys():
+      if 'alpha' in key:
+        if key in merged_sd:
+          assert merged_sd[key] == lora_sd[key], f"alpha mismatch / alphaが異なる場合、現時点ではマージできません"
+        else:
+          alpha = lora_sd[key].detach().numpy()
+          merged_sd[key] = lora_sd[key]
+      else:
+        if key in merged_sd:
+          assert merged_sd[key].size() == lora_sd[key].size(
+          ), f"weights shape mismatch merging v1 and v2, different dims? / 重みのサイズが合いません。v1とv2、または次元数の異なるモデルはマージできません"
+          merged_sd[key] = merged_sd[key] + lora_sd[key] * ratio
+        else:
+          if "lora_down" in key:
+            dim = lora_sd[key].size()[0]
+          merged_sd[key] = lora_sd[key] * ratio
+
+  print(f"dim (rank): {dim}, alpha: {alpha}")
+  if alpha is None:
+    alpha = dim
+
+  return merged_sd, dim, alpha
+
+
+def merge(args):
+  assert len(args.models) == len(args.ratios), f"number of models must be equal to number of ratios / モデルの数と重みの数は合わせてください"
+
+  def str_to_dtype(p):
+    if p == 'float':
+      return torch.float
+    if p == 'fp16':
+      return torch.float16
+    if p == 'bf16':
+      return torch.bfloat16
+    return None
+
+  merge_dtype = str_to_dtype(args.precision)
+  save_dtype = str_to_dtype(args.save_precision)
+  if save_dtype is None:
+    save_dtype = merge_dtype
+
+  if args.sd_model is not None:
+    print(f"loading SD model: {args.sd_model}")
+
+    text_encoder, vae, unet = model_util.load_models_from_stable_diffusion_checkpoint(args.v2, args.sd_model)
+
+    merge_to_sd_model(text_encoder, unet, args.models, args.ratios, merge_dtype)
+
+    print(f"saving SD model to: {args.save_to}")
+    model_util.save_stable_diffusion_checkpoint(args.v2, args.save_to, text_encoder, unet,
+                                                args.sd_model, 0, 0, save_dtype, vae)
+  else:
+    state_dict, _, _ = merge_lora_models(args.models, args.ratios, merge_dtype)
+
+    print(f"saving model to: {args.save_to}")
+    save_to_file(args.save_to, state_dict, state_dict, save_dtype)
+
+
+if __name__ == '__main__':
+  parser = argparse.ArgumentParser()
+  parser.add_argument("--v2", action='store_true',
+                      help='load Stable Diffusion v2.x model / Stable Diffusion 2.xのモデルを読み込む')
+  parser.add_argument("--save_precision", type=str, default=None,
+                      choices=[None, "float", "fp16", "bf16"], help="precision in saving, same to merging if omitted / 保存時に精度を変更して保存する、省略時はマージ時の精度と同じ")
+  parser.add_argument("--precision", type=str, default="float",
+                      choices=["float", "fp16", "bf16"], help="precision in merging (float is recommended) / マージの計算時の精度（floatを推奨）")
+  parser.add_argument("--sd_model", type=str, default=None,
+                      help="Stable Diffusion model to load: ckpt or safetensors file, merge LoRA models if omitted / 読み込むモデル、ckptまたはsafetensors。省略時はLoRAモデル同士をマージする")
+  parser.add_argument("--save_to", type=str, default=None,
+                      help="destination file name: ckpt or safetensors file / 保存先のファイル名、ckptまたはsafetensors")
+  parser.add_argument("--models", type=str, nargs='*',
+                      help="LoRA models to merge: ckpt or safetensors file / マージするLoRAモデル、ckptまたはsafetensors")
+  parser.add_argument("--ratios", type=float, nargs='*',
+                      help="ratios for each model / それぞれのLoRAモデルの比率")
+
+  args = parser.parse_args()
+  merge(args)

networks/resize_lora.py

@@ -5,148 +5,178 @@
 import argparse
 import os
 import torch
-from safetensors.torch import load_file, save_file
+from safetensors.torch import load_file, save_file, safe_open
 from tqdm import tqdm
+from library import train_util, model_util
+
 
 def load_state_dict(file_name, dtype):
-  if os.path.splitext(file_name)[1] == '.safetensors':
+  if model_util.is_safetensors(file_name):
     sd = load_file(file_name)
+    with safe_open(file_name, framework="pt") as f:
+      metadata = f.metadata()
   else:
     sd = torch.load(file_name, map_location='cpu')
+    metadata = None
+
   for key in list(sd.keys()):
     if type(sd[key]) == torch.Tensor:
       sd[key] = sd[key].to(dtype)
-  return sd
+
+  return sd, metadata
 
 
-def save_to_file(file_name, model, state_dict, dtype):
+def save_to_file(file_name, model, state_dict, dtype, metadata):
   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)
+  if model_util.is_safetensors(file_name):
+    save_file(model, file_name, metadata)
   else:
     torch.save(model, file_name)
-    
 
 
-def resize_lora_model(model, new_rank, merge_dtype, save_dtype):
-    print("Loading Model...")
-    lora_sd = load_state_dict(model, merge_dtype)
+def resize_lora_model(lora_sd, new_rank, save_dtype, device, verbose):
+  network_alpha = None
+  network_dim = None
+  verbose_str = "\n"
 
-    network_alpha = None
-    network_dim = None
+  CLAMP_QUANTILE = 0.99
 
-    CLAMP_QUANTILE = 0.99
+  # Extract loaded lora dim and alpha
+  for key, value in lora_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 network_alpha is not None and network_dim is not None:
+      break
+    if network_alpha is None:
+      network_alpha = network_dim
 
-    # Extract loaded lora dim and alpha
-    for key, value in lora_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 network_alpha is not None and network_dim is not None:
-            break
-        if network_alpha is None:
-            network_alpha = network_dim
+  scale = network_alpha/network_dim
+  new_alpha = float(scale*new_rank)  # calculate new alpha from scale
 
-    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}")
 
-    print(f"dimension: {network_dim}, alpha: {network_alpha}, new alpha: {new_alpha}")
+  lora_down_weight = None
+  lora_up_weight = None
 
-    lora_down_weight = None
-    lora_up_weight = None
+  o_lora_sd = lora_sd.copy()
+  block_down_name = None
+  block_up_name = None
 
-    o_lora_sd = lora_sd.copy()
-    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:
+        block_down_name = key.split(".")[0]
+        lora_down_weight = value
+      if 'lora_up' in key:
+        block_up_name = key.split(".")[0]
+        lora_up_weight = value
 
-    print("resizing lora...")
-    with torch.no_grad():
-        for key, value in tqdm(lora_sd.items()):
-            if 'lora_down' in key:
-                block_down_name = key.split(".")[0]
-                lora_down_weight = value
-            if 'lora_up' in key:
-                block_up_name = key.split(".")[0]
-                lora_up_weight = value
+      weights_loaded = (lora_down_weight is not None and lora_up_weight is not None)
 
-            weights_loaded = (lora_down_weight is not None and lora_up_weight is not None)
+      if (block_down_name == block_up_name) and weights_loaded:
 
-            if (block_down_name == block_up_name) and weights_loaded:
+        conv2d = (len(lora_down_weight.size()) == 4)
 
-                conv2d = (len(lora_down_weight.size()) == 4)
-                
-                if conv2d:
-                    lora_down_weight = lora_down_weight.squeeze()
-                    lora_up_weight = lora_up_weight.squeeze()
+        if conv2d:
+          lora_down_weight = lora_down_weight.squeeze()
+          lora_up_weight = lora_up_weight.squeeze()
 
-                if args.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)
+        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)
+        full_weight_matrix = torch.matmul(lora_up_weight, lora_down_weight)
 
-                U, S, Vh = torch.linalg.svd(full_weight_matrix)
+        U, S, Vh = torch.linalg.svd(full_weight_matrix)
 
-                U = U[:, :new_rank]
-                S = S[:new_rank]
-                U = U @ torch.diag(S)
+        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"
 
-                Vh = Vh[:new_rank, :]
+        U = U[:, :new_rank]
+        S = S[:new_rank]
+        U = U @ torch.diag(S)
 
-                dist = torch.cat([U.flatten(), Vh.flatten()])
-                hi_val = torch.quantile(dist, CLAMP_QUANTILE)
-                low_val = -hi_val
+        Vh = Vh[:new_rank, :]
 
-                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 args.device:
-                   U = U.to(org_device)
-                   Vh = Vh.to(org_device)
+        dist = torch.cat([U.flatten(), Vh.flatten()])
+        hi_val = torch.quantile(dist, CLAMP_QUANTILE)
+        low_val = -hi_val
 
-                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)
+        U = U.clamp(low_val, hi_val)
+        Vh = Vh.clamp(low_val, hi_val)
 
-                block_down_name = None
-                block_up_name = None
-                lora_down_weight = None
-                lora_up_weight = None
-                weights_loaded = False
+        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)
+
+        block_down_name = None
+        block_up_name = None
+        lora_down_weight = None
+        lora_up_weight = None
+        weights_loaded = False
+
+  if verbose:
+    print(verbose_str)
+  print("resizing complete")
+  return o_lora_sd, network_dim, new_alpha
 
-    print("resizing complete")
-    return o_lora_sd
 
 def resize(args):
 
-    def str_to_dtype(p):
-        if p == 'float':
-            return torch.float
-        if p == 'fp16':
-            return torch.float16
-        if p == 'bf16':
-            return torch.bfloat16
-        return None
+  def str_to_dtype(p):
+    if p == 'float':
+      return torch.float
+    if p == 'fp16':
+      return torch.float16
+    if p == 'bf16':
+      return torch.bfloat16
+    return None
 
-    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:
-        save_dtype = merge_dtype
+  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:
+    save_dtype = merge_dtype
 
-    state_dict =  resize_lora_model(args.model, args.new_rank, merge_dtype, save_dtype)
+  print("loading Model...")
+  lora_sd, metadata = load_state_dict(args.model, merge_dtype)
 
-    print(f"saving model to: {args.save_to}")
-    save_to_file(args.save_to, state_dict, state_dict, save_dtype)
+  print("resizing rank...")
+  state_dict, old_dim, new_alpha = resize_lora_model(lora_sd, args.new_rank, save_dtype, args.device, 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)
+
+  model_hash, legacy_hash = train_util.precalculate_safetensors_hashes(state_dict, metadata)
+  metadata["sshs_model_hash"] = model_hash
+  metadata["sshs_legacy_hash"] = legacy_hash
+
+  print(f"saving model to: {args.save_to}")
+  save_to_file(args.save_to, state_dict, state_dict, save_dtype, metadata)
 
 
 if __name__ == '__main__':
@@ -161,6 +191,8 @@ if __name__ == '__main__':
   parser.add_argument("--model", type=str, default=None,
                       help="LoRA model to resize at to new rank: ckpt or safetensors file / 読み込むLoRAモデル、ckptまたはsafetensors")
   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変更時の詳細情報を出力する")
 
   args = parser.parse_args()
   resize(args)

networks/svd_merge_lora.py

@@ -0,0 +1,164 @@
+
+import math
+import argparse
+import os
+import torch
+from safetensors.torch import load_file, save_file
+from tqdm import tqdm
+import library.model_util as model_util
+import lora
+
+
+CLAMP_QUANTILE = 0.99
+
+
+def load_state_dict(file_name, dtype):
+  if os.path.splitext(file_name)[1] == '.safetensors':
+    sd = load_file(file_name)
+  else:
+    sd = torch.load(file_name, map_location='cpu')
+  for key in list(sd.keys()):
+    if type(sd[key]) == torch.Tensor:
+      sd[key] = sd[key].to(dtype)
+  return sd
+
+
+def save_to_file(file_name, model, 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)
+  else:
+    torch.save(model, file_name)
+
+
+def merge_lora_models(models, ratios, new_rank, device,  merge_dtype):
+  merged_sd = {}
+  for model, ratio in zip(models, ratios):
+    print(f"loading: {model}")
+    lora_sd = load_state_dict(model, merge_dtype)
+
+    # merge
+    print(f"merging...")
+    for key in tqdm(list(lora_sd.keys())):
+      if 'lora_down' not in key:
+        continue
+
+      lora_module_name = key[:key.rfind(".lora_down")]
+
+      down_weight = lora_sd[key]
+      network_dim = down_weight.size()[0]
+
+      up_weight = lora_sd[lora_module_name + '.lora_up.weight']
+      alpha = lora_sd.get(lora_module_name + '.alpha', network_dim)
+
+      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)
+
+      # 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)
+        if device:
+          weight = weight.to(device)
+      else:
+        weight = merged_sd[lora_module_name]
+
+      # merge to weight
+      if device:
+        up_weight = up_weight.to(device)
+        down_weight = down_weight.to(device)
+
+      # W <- W + U * D
+      scale = (alpha / network_dim)
+      if not conv2d:        # linear
+        weight = weight + ratio * (up_weight @ down_weight) * scale
+      else:
+        weight = weight + ratio * (up_weight.squeeze(3).squeeze(2) @ down_weight.squeeze(3).squeeze(2)
+                                   ).unsqueeze(2).unsqueeze(3) * scale
+
+      merged_sd[lora_module_name] = weight
+
+  # extract from merged weights
+  print("extract new lora...")
+  merged_lora_sd = {}
+  with torch.no_grad():
+    for lora_module_name, mat in tqdm(list(merged_sd.items())):
+      conv2d = (len(mat.size()) == 4)
+      if conv2d:
+        mat = mat.squeeze()
+
+      U, S, Vh = torch.linalg.svd(mat)
+
+      U = U[:, :new_rank]
+      S = S[:new_rank]
+      U = U @ torch.diag(S)
+
+      Vh = Vh[:new_rank, :]
+
+      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)
+
+      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)
+
+  return merged_lora_sd
+
+
+def merge(args):
+  assert len(args.models) == len(args.ratios), f"number of models must be equal to number of ratios / モデルの数と重みの数は合わせてください"
+
+  def str_to_dtype(p):
+    if p == 'float':
+      return torch.float
+    if p == 'fp16':
+      return torch.float16
+    if p == 'bf16':
+      return torch.bfloat16
+    return None
+
+  merge_dtype = str_to_dtype(args.precision)
+  save_dtype = str_to_dtype(args.save_precision)
+  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)
+
+  print(f"saving model to: {args.save_to}")
+  save_to_file(args.save_to, state_dict, state_dict, save_dtype)
+
+
+if __name__ == '__main__':
+  parser = argparse.ArgumentParser()
+  parser.add_argument("--save_precision", type=str, default=None,
+                      choices=[None, "float", "fp16", "bf16"], help="precision in saving, same to merging if omitted / 保存時に精度を変更して保存する、省略時はマージ時の精度と同じ")
+  parser.add_argument("--precision", type=str, default="float",
+                      choices=["float", "fp16", "bf16"], help="precision in merging (float is recommended) / マージの計算時の精度（floatを推奨）")
+  parser.add_argument("--save_to", type=str, default=None,
+                      help="destination file name: ckpt or safetensors file / 保存先のファイル名、ckptまたはsafetensors")
+  parser.add_argument("--models", type=str, nargs='*',
+                      help="LoRA models to merge: ckpt or safetensors file / マージするLoRAモデル、ckptまたはsafetensors")
+  parser.add_argument("--ratios", type=float, nargs='*',
+                      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("--device", type=str, default=None, help="device to use, cuda for GPU / 計算を行うデバイス、cuda でGPUを使う")
+
+  args = parser.parse_args()
+  merge(args)

requirements.txt

@@ -1,23 +1,24 @@
 accelerate==0.15.0
 transformers==4.26.0
-ftfy
-albumentations
-opencv-python
-einops
+ftfy==6.1.1
+albumentations==1.3.0
+opencv-python==4.7.0.68
+einops==0.6.0
 diffusers[torch]==0.10.2
-pytorch_lightning
+pytorch-lightning==1.9.0
 bitsandbytes==0.35.0
-tensorboard
+tensorboard==2.10.1
 safetensors==0.2.6
-gradio
-altair
-easygui
+gradio==3.16.2
+altair==4.2.2
+easygui==0.98.3
 # for BLIP captioning
-requests
-timm==0.4.12
-fairscale==0.4.4
+requests==2.28.2
+timm==0.6.12
+fairscale==0.4.13
 # for WD14 captioning
-tensorflow<2.11
-huggingface-hub
+# tensorflow<2.11
+tensorflow==2.10.1
+huggingface-hub==0.12.0
 # for kohya_ss library
-.
+.

tools/resize_images_to_resolution.py

@@ -0,0 +1,122 @@
+import glob
+import os
+import cv2
+import argparse
+import shutil
+import math
+from PIL import Image
+import numpy as np
+
+
+def resize_images(src_img_folder, dst_img_folder, max_resolution="512x512", divisible_by=2, interpolation=None, save_as_png=False, copy_associated_files=False):
+  # Split the max_resolution string by "," and strip any whitespaces
+  max_resolutions = [res.strip() for res in max_resolution.split(',')]
+
+  # # Calculate max_pixels from max_resolution string
+  # max_pixels = int(max_resolution.split("x")[0]) * int(max_resolution.split("x")[1])
+
+  # Create destination folder if it does not exist
+  if not os.path.exists(dst_img_folder):
+    os.makedirs(dst_img_folder)
+
+  # Select interpolation method
+  if interpolation == 'lanczos4':
+    cv2_interpolation = cv2.INTER_LANCZOS4
+  elif interpolation == 'cubic':
+    cv2_interpolation = cv2.INTER_CUBIC
+  else:
+    cv2_interpolation = cv2.INTER_AREA
+
+  # Iterate through all files in src_img_folder
+  img_exts = (".png", ".jpg", ".jpeg", ".webp", ".bmp")                   # copy from train_util.py
+  for filename in os.listdir(src_img_folder):
+    # Check if the image is png, jpg or webp etc...
+    if not filename.endswith(img_exts):
+      # Copy the file to the destination folder if not png, jpg or webp etc (.txt or .caption or etc.)
+      shutil.copy(os.path.join(src_img_folder, filename), os.path.join(dst_img_folder, filename))
+      continue
+
+    # Load image
+    # img = cv2.imread(os.path.join(src_img_folder, filename))
+    image = Image.open(os.path.join(src_img_folder, filename))
+    if not image.mode == "RGB":
+      image = image.convert("RGB")
+    img = np.array(image, np.uint8)
+
+    base, _ = os.path.splitext(filename)
+    for max_resolution in max_resolutions:
+      # Calculate max_pixels from max_resolution string
+      max_pixels = int(max_resolution.split("x")[0]) * int(max_resolution.split("x")[1])
+
+      # Calculate current number of pixels
+      current_pixels = img.shape[0] * img.shape[1]
+
+      # Check if the image needs resizing
+      if current_pixels > max_pixels:
+        # Calculate scaling factor
+        scale_factor = max_pixels / current_pixels
+
+        # Calculate new dimensions
+        new_height = int(img.shape[0] * math.sqrt(scale_factor))
+        new_width = int(img.shape[1] * math.sqrt(scale_factor))
+
+        # Resize image
+        img = cv2.resize(img, (new_width, new_height), interpolation=cv2_interpolation)
+      else:
+        new_height, new_width = img.shape[0:2]
+
+      # Calculate the new height and width that are divisible by divisible_by (with/without resizing)
+      new_height = new_height if new_height % divisible_by == 0 else new_height - new_height % divisible_by
+      new_width = new_width if new_width % divisible_by == 0 else new_width - new_width % divisible_by
+
+      # Center crop the image to the calculated dimensions
+      y = int((img.shape[0] - new_height) / 2)
+      x = int((img.shape[1] - new_width) / 2)
+      img = img[y:y + new_height, x:x + new_width]
+
+      # Split filename into base and extension
+      new_filename = base + '+' + max_resolution + ('.png' if save_as_png else '.jpg')
+
+      # Save resized image in dst_img_folder
+      # cv2.imwrite(os.path.join(dst_img_folder, new_filename), img, [cv2.IMWRITE_JPEG_QUALITY, 100])
+      image = Image.fromarray(img)
+      image.save(os.path.join(dst_img_folder, new_filename), quality=100)
+
+      proc = "Resized" if current_pixels > max_pixels else "Saved"
+      print(f"{proc} image: {filename} with size {img.shape[0]}x{img.shape[1]} as {new_filename}")
+
+    # If other files with same basename, copy them with resolution suffix
+    if copy_associated_files:
+      asoc_files = glob.glob(os.path.join(src_img_folder, base + ".*"))
+      for asoc_file in asoc_files:
+        ext = os.path.splitext(asoc_file)[1]
+        if ext in img_exts:
+          continue
+        for max_resolution in max_resolutions:
+          new_asoc_file = base + '+' + max_resolution + ext
+          print(f"Copy {asoc_file} as {new_asoc_file}")
+          shutil.copy(os.path.join(src_img_folder, asoc_file), os.path.join(dst_img_folder, new_asoc_file))
+
+
+def main():
+  parser = argparse.ArgumentParser(
+      description='Resize images in a folder to a specified max resolution(s) / 指定されたフォルダ内の画像を指定した最大画像サイズ（面積）以下にアスペクト比を維持したままリサイズします')
+  parser.add_argument('src_img_folder', type=str, help='Source folder containing the images / 元画像のフォルダ')
+  parser.add_argument('dst_img_folder', type=str, help='Destination folder to save the resized images / リサイズ後の画像を保存するフォルダ')
+  parser.add_argument('--max_resolution', type=str,
+                      help='Maximum resolution(s) in the format "512x512,384x384, etc, etc" / 最大画像サイズをカンマ区切りで指定 ("512x512,384x384, etc, etc" など)', default="512x512,384x384,256x256,128x128")
+  parser.add_argument('--divisible_by', type=int,
+                      help='Ensure new dimensions are divisible by this value / リサイズ後の画像のサイズをこの値で割り切れるようにします', default=1)
+  parser.add_argument('--interpolation', type=str, choices=['area', 'cubic', 'lanczos4'],
+                      default='area', help='Interpolation method for resizing / リサイズ時の補完方法')
+  parser.add_argument('--save_as_png', action='store_true', help='Save as png format / png形式で保存')
+  parser.add_argument('--copy_associated_files', action='store_true',
+                      help='Copy files with same base name to images (captions etc) / 画像と同じファイル名（拡張子を除く）のファイルもコピーする')
+
+  args = parser.parse_args()
+  resize_images(args.src_img_folder, args.dst_img_folder, args.max_resolution,
+                args.divisible_by, args.interpolation, args.save_as_png, args.copy_associated_files)
+
+
+if __name__ == '__main__':
+  main()

train_db.py

@@ -35,10 +35,16 @@ def train(args):
 
   train_dataset = DreamBoothDataset(args.train_batch_size, args.train_data_dir, args.reg_data_dir,
                                     tokenizer, args.max_token_length, args.caption_extension, args.shuffle_caption, args.keep_tokens,
-                                    args.resolution, args.enable_bucket, args.min_bucket_reso, args.max_bucket_reso, args.prior_loss_weight,
-                                    args.flip_aug, args.color_aug, args.face_crop_aug_range, args.random_crop, args.debug_dataset)
+                                    args.resolution, args.enable_bucket, args.min_bucket_reso, args.max_bucket_reso,
+                                    args.bucket_reso_steps, args.bucket_no_upscale,
+                                    args.prior_loss_weight, args.flip_aug, args.color_aug, args.face_crop_aug_range, args.random_crop, args.debug_dataset)
+
   if args.no_token_padding:
     train_dataset.disable_token_padding()
+
+  # 学習データのdropout率を設定する
+  train_dataset.set_caption_dropout(args.caption_dropout_rate, args.caption_dropout_every_n_epochs, args.caption_tag_dropout_rate)
+
   train_dataset.make_buckets()
 
   if args.debug_dataset:
@@ -118,6 +124,13 @@ def train(args):
       raise ImportError("No bitsand bytes / bitsandbytesがインストールされていないようです")
     print("use 8-bit Adam optimizer")
     optimizer_class = bnb.optim.AdamW8bit
+  elif args.use_lion_optimizer:
+    try:
+      import lion_pytorch
+    except ImportError:
+      raise ImportError("No lion_pytorch / lion_pytorch がインストールされていないようです")
+    print("use Lion optimizer")
+    optimizer_class = lion_pytorch.Lion
   else:
     optimizer_class = torch.optim.AdamW
 
@@ -200,8 +213,11 @@ def train(args):
   if accelerator.is_main_process:
     accelerator.init_trackers("dreambooth")
 
+  loss_list = []
+  loss_total = 0.0
   for epoch in range(num_train_epochs):
     print(f"epoch {epoch+1}/{num_train_epochs}")
+    train_dataset.set_current_epoch(epoch + 1)
 
     # 指定したステップ数までText Encoderを学習する：epoch最初の状態
     unet.train()
@@ -209,7 +225,6 @@ def train(args):
     if args.gradient_checkpointing or global_step < args.stop_text_encoder_training:
       text_encoder.train()
 
-    loss_total = 0
     for step, batch in enumerate(train_dataloader):
       # 指定したステップ数でText Encoderの学習を止める
       if global_step == args.stop_text_encoder_training:
@@ -226,10 +241,13 @@ def train(args):
           else:
             latents = vae.encode(batch["images"].to(dtype=weight_dtype)).latent_dist.sample()
           latents = latents * 0.18215
+        b_size = latents.shape[0]
 
         # Sample noise that we'll add to the latents
         noise = torch.randn_like(latents, device=latents.device)
-        b_size = latents.shape[0]
+        if args.noise_offset:
+          # https://www.crosslabs.org//blog/diffusion-with-offset-noise
+          noise += args.noise_offset * torch.randn((latents.shape[0], latents.shape[1], 1, 1), device=latents.device)
 
         # Get the text embedding for conditioning
         with torch.set_grad_enabled(global_step < args.stop_text_encoder_training):
@@ -284,8 +302,13 @@ def train(args):
         logs = {"loss": current_loss, "lr": lr_scheduler.get_last_lr()[0]}
         accelerator.log(logs, step=global_step)
 
+      if epoch == 0:
+        loss_list.append(current_loss)
+      else:
+        loss_total -= loss_list[step]
+        loss_list[step] = current_loss
       loss_total += current_loss
-      avr_loss = loss_total / (step+1)
+      avr_loss = loss_total / len(loss_list)
       logs = {"loss": avr_loss}  # , "lr": lr_scheduler.get_last_lr()[0]}
       progress_bar.set_postfix(**logs)
 
@@ -293,7 +316,7 @@ def train(args):
         break
 
     if args.logging_dir is not None:
-      logs = {"epoch_loss": loss_total / len(train_dataloader)}
+      logs = {"loss/epoch": loss_total / len(loss_list)}
       accelerator.log(logs, step=epoch+1)
 
     accelerator.wait_for_everyone()
@@ -326,7 +349,7 @@ if __name__ == '__main__':
   parser = argparse.ArgumentParser()
 
   train_util.add_sd_models_arguments(parser)
-  train_util.add_dataset_arguments(parser, True, False)
+  train_util.add_dataset_arguments(parser, True, False, True)
   train_util.add_training_arguments(parser, True)
   train_util.add_sd_saving_arguments(parser)
 

train_network.py

@@ -1,5 +1,7 @@
 from diffusers.optimization import SchedulerType, TYPE_TO_SCHEDULER_FUNCTION
 from torch.optim import Optimizer
+from torch.cuda.amp import autocast
+from torch.nn.parallel import DistributedDataParallel as DDP
 from typing import Optional, Union
 import importlib
 import argparse
@@ -120,15 +122,22 @@ def train(args):
     print("Use DreamBooth method.")
     train_dataset = DreamBoothDataset(args.train_batch_size, args.train_data_dir, args.reg_data_dir,
                                       tokenizer, args.max_token_length, args.caption_extension, args.shuffle_caption, args.keep_tokens,
-                                      args.resolution, args.enable_bucket, args.min_bucket_reso, args.max_bucket_reso, args.prior_loss_weight,
-                                      args.flip_aug, args.color_aug, args.face_crop_aug_range, args.random_crop, args.debug_dataset)
+                                      args.resolution, args.enable_bucket, args.min_bucket_reso, args.max_bucket_reso,
+                                      args.bucket_reso_steps, args.bucket_no_upscale,
+                                      args.prior_loss_weight, args.flip_aug, args.color_aug, args.face_crop_aug_range,
+                                      args.random_crop, args.debug_dataset)
   else:
     print("Train with captions.")
     train_dataset = FineTuningDataset(args.in_json, args.train_batch_size, args.train_data_dir,
                                       tokenizer, args.max_token_length, args.shuffle_caption, args.keep_tokens,
                                       args.resolution, args.enable_bucket, args.min_bucket_reso, args.max_bucket_reso,
+                                      args.bucket_reso_steps, args.bucket_no_upscale,
                                       args.flip_aug, args.color_aug, args.face_crop_aug_range, args.random_crop,
                                       args.dataset_repeats, args.debug_dataset)
+
+  # 学習データのdropout率を設定する
+  train_dataset.set_caption_dropout(args.caption_dropout_rate, args.caption_dropout_every_n_epochs, args.caption_tag_dropout_rate)
+
   train_dataset.make_buckets()
 
   if args.debug_dataset:
@@ -148,6 +157,11 @@ def train(args):
   # モデルを読み込む
   text_encoder, vae, unet, _ = train_util.load_target_model(args, weight_dtype)
 
+  # work on low-ram device
+  if args.lowram:
+    text_encoder.to("cuda")
+    unet.to("cuda")
+  
   # モデルに xformers とか memory efficient attention を組み込む
   train_util.replace_unet_modules(unet, args.mem_eff_attn, args.xformers)
 
@@ -202,9 +216,18 @@ def train(args):
       raise ImportError("No bitsand bytes / bitsandbytesがインストールされていないようです")
     print("use 8-bit Adam optimizer")
     optimizer_class = bnb.optim.AdamW8bit
+  elif args.use_lion_optimizer:
+    try:
+      import lion_pytorch
+    except ImportError:
+      raise ImportError("No lion_pytorch / lion_pytorch がインストールされていないようです")
+    print("use Lion optimizer")
+    optimizer_class = lion_pytorch.Lion
   else:
     optimizer_class = torch.optim.AdamW
 
+  optimizer_name = optimizer_class.__module__ + "." + optimizer_class.__name__
+
   trainable_params = network.prepare_optimizer_params(args.text_encoder_lr, args.unet_lr)
 
   # betaやweight decayはdiffusers DreamBoothもDreamBooth SDもデフォルト値のようなのでオプションはとりあえず省略
@@ -251,17 +274,26 @@ def train(args):
   unet.requires_grad_(False)
   unet.to(accelerator.device, dtype=weight_dtype)
   text_encoder.requires_grad_(False)
-  text_encoder.to(accelerator.device, dtype=weight_dtype)
+  text_encoder.to(accelerator.device)
   if args.gradient_checkpointing:                       # according to TI example in Diffusers, train is required
     unet.train()
     text_encoder.train()
 
     # set top parameter requires_grad = True for gradient checkpointing works
-    text_encoder.text_model.embeddings.requires_grad_(True)
+    if type(text_encoder) == DDP:
+      text_encoder.module.text_model.embeddings.requires_grad_(True)
+    else:
+      text_encoder.text_model.embeddings.requires_grad_(True)
   else:
     unet.eval()
     text_encoder.eval()
 
+  # support DistributedDataParallel
+  if type(text_encoder) == DDP:
+    text_encoder = text_encoder.module
+    unet = unet.module
+    network = network.module
+
   network.prepare_grad_etc(text_encoder, unet)
 
   if not cache_latents:
@@ -333,11 +365,14 @@ def train(args):
       "ss_max_bucket_reso": train_dataset.max_bucket_reso,
       "ss_seed": args.seed,
       "ss_keep_tokens": args.keep_tokens,
+      "ss_noise_offset": args.noise_offset,
       "ss_dataset_dirs": json.dumps(train_dataset.dataset_dirs_info),
       "ss_reg_dataset_dirs": json.dumps(train_dataset.reg_dataset_dirs_info),
       "ss_tag_frequency": json.dumps(train_dataset.tag_frequency),
       "ss_bucket_info": json.dumps(train_dataset.bucket_info),
-      "ss_training_comment": args.training_comment        # will not be updated after training
+      "ss_training_comment": args.training_comment,       # will not be updated after training
+      "ss_sd_scripts_commit_hash": train_util.get_git_revision_hash(),
+      "ss_optimizer": optimizer_name
   }
 
   # uncomment if another network is added
@@ -371,13 +406,16 @@ def train(args):
   if accelerator.is_main_process:
     accelerator.init_trackers("network_train")
 
+  loss_list = []
+  loss_total = 0.0
   for epoch in range(num_train_epochs):
     print(f"epoch {epoch+1}/{num_train_epochs}")
+    train_dataset.set_current_epoch(epoch + 1)
+
     metadata["ss_epoch"] = str(epoch+1)
 
     network.on_epoch_start(text_encoder, unet)
 
-    loss_total = 0
     for step, batch in enumerate(train_dataloader):
       with accelerator.accumulate(network):
         with torch.no_grad():
@@ -396,6 +434,9 @@ def train(args):
 
         # Sample noise that we'll add to the latents
         noise = torch.randn_like(latents, device=latents.device)
+        if args.noise_offset:
+          # https://www.crosslabs.org//blog/diffusion-with-offset-noise
+          noise += args.noise_offset * torch.randn((latents.shape[0], latents.shape[1], 1, 1), device=latents.device)
 
         # Sample a random timestep for each image
         timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (b_size,), device=latents.device)
@@ -406,7 +447,8 @@ def train(args):
         noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)
 
         # Predict the noise residual
-        noise_pred = unet(noisy_latents, timesteps, encoder_hidden_states).sample
+        with autocast():
+          noise_pred = unet(noisy_latents, timesteps, encoder_hidden_states).sample
 
         if args.v_parameterization:
           # v-parameterization training
@@ -437,8 +479,13 @@ def train(args):
         global_step += 1
 
       current_loss = loss.detach().item()
+      if epoch == 0:
+        loss_list.append(current_loss)
+      else:
+        loss_total -= loss_list[step]
+        loss_list[step] = current_loss
       loss_total += current_loss
-      avr_loss = loss_total / (step+1)
+      avr_loss = loss_total / len(loss_list)
       logs = {"loss": avr_loss}  # , "lr": lr_scheduler.get_last_lr()[0]}
       progress_bar.set_postfix(**logs)
 
@@ -450,7 +497,7 @@ def train(args):
         break
 
     if args.logging_dir is not None:
-      logs = {"loss/epoch": loss_total / len(train_dataloader)}
+      logs = {"loss/epoch": loss_total / len(loss_list)}
       accelerator.log(logs, step=epoch+1)
 
     accelerator.wait_for_everyone()
@@ -506,7 +553,7 @@ if __name__ == '__main__':
   parser = argparse.ArgumentParser()
 
   train_util.add_sd_models_arguments(parser)
-  train_util.add_dataset_arguments(parser, True, True)
+  train_util.add_dataset_arguments(parser, True, True, True)
   train_util.add_training_arguments(parser, True)
 
   parser.add_argument("--no_metadata", action='store_true', help="do not save metadata in output model / メタデータを出力先モデルに保存しない")

train_network_README-ja.md

@@ -55,7 +55,7 @@ accelerate launch --num_cpu_threads_per_process 1 train_network.py
     --network_module=networks.lora
 ```
 
---output_dirオプションで指定したディレクトリに、LoRAのモデルが保存されます。
+--output_dirオプションで指定したフォルダに、LoRAのモデルが保存されます。
 
 その他、以下のオプションが指定できます。
 
@@ -178,6 +178,38 @@ Text Encoderが二つのモデルで同じ場合にはLoRAはU-NetのみのLoRA
 - --save_precision
   - LoRAの保存形式を"float", "fp16", "bf16"から指定します。省略時はfloatになります。
 
+## 画像リサイズスクリプト
+
+（のちほどドキュメントを整理しますがとりあえずここに説明を書いておきます。）
+
+Aspect Ratio Bucketingの機能拡張で、小さな画像については拡大しないでそのまま教師データとすることが可能になりました。元の教師画像を縮小した画像を、教師データに加えると精度が向上したという報告とともに前処理用のスクリプトをいただきましたので整備して追加しました。bmaltais氏に感謝します。
+
+### スクリプトの実行方法
+
+以下のように指定してください。元の画像そのまま、およびリサイズ後の画像が変換先フォルダに保存されます。リサイズ後の画像には、ファイル名に ``+512x512`` のようにリサイズ先の解像度が付け加えられます（画像サイズとは異なります）。リサイズ先の解像度より小さい画像は拡大されることはありません。
+
+```
+python tools\resize_images_to_resolution.py --max_resolution 512x512,384x384,256x256 --save_as_png 
+    --copy_associated_files 元画像フォルダ 変換先フォルダ
+```
+
+元画像フォルダ内の画像ファイルが、指定した解像度（複数指定可）と同じ面積になるようにリサイズされ、変換先フォルダに保存されます。画像以外のファイルはそのままコピーされます。
+
+``--max_resolution`` オプションにリサイズ先のサイズを例のように指定してください。面積がそのサイズになるようにリサイズします。複数指定すると、それぞれの解像度でリサイズされます。``512x512,384x384,256x256``なら、変換先フォルダの画像は、元サイズとリサイズ後サイズ×3の計4枚になります。
+
+``--save_as_png`` オプションを指定するとpng形式で保存します。省略するとjpeg形式（quality=100）で保存されます。
+
+``--copy_associated_files`` オプションを指定すると、拡張子を除き画像と同じファイル名（たとえばキャプションなど）のファイルが、リサイズ後の画像のファイル名と同じ名前でコピーされます。
+
+
+### その他のオプション
+
+- divisible_by
+  - リサイズ後の画像のサイズ（縦、横のそれぞれ）がこの値で割り切れるように、画像中心を切り出します。
+- interpolation
+  - 縮小時の補完方法を指定します。``area, cubic, lanczos4``から選択可能で、デフォルトは``area``です。
+
+
 ## 追加情報
 
 ### cloneofsimo氏のリポジトリとの違い

train_textual_inversion.py

@@ -98,12 +98,12 @@ def train(args):
 
   # Convert the init_word to token_id
   if args.init_word is not None:
-    init_token_id = tokenizer.encode(args.init_word, add_special_tokens=False)
-    assert len(
-        init_token_id) == 1, f"init word {args.init_word} is not converted to single token / 初期化単語が二つ以上のトークンに変換されます。別の単語を使ってください"
-    init_token_id = init_token_id[0]
+    init_token_ids = tokenizer.encode(args.init_word, add_special_tokens=False)
+    if len(init_token_ids) > 1 and len(init_token_ids) != args.num_vectors_per_token:
+      print(
+          f"token length for init words is not same to num_vectors_per_token, init words is repeated or truncated / 初期化単語のトークン長がnum_vectors_per_tokenと合わないため、繰り返しまたは切り捨てが発生します: length {len(init_token_ids)}")
   else:
-    init_token_id = None
+    init_token_ids = None
 
   # add new word to tokenizer, count is num_vectors_per_token
   token_strings = [args.token_string] + [f"{args.token_string}{i+1}" for i in range(args.num_vectors_per_token - 1)]
@@ -120,9 +120,9 @@ def train(args):
 
   # Initialise the newly added placeholder token with the embeddings of the initializer token
   token_embeds = text_encoder.get_input_embeddings().weight.data
-  if init_token_id is not None:
-    for token_id in token_ids:
-      token_embeds[token_id] = token_embeds[init_token_id]
+  if init_token_ids is not None:
+    for i, token_id in enumerate(token_ids):
+      token_embeds[token_id] = token_embeds[init_token_ids[i % len(init_token_ids)]]
       # print(token_id, token_embeds[token_id].mean(), token_embeds[token_id].min())
 
   # load weights
@@ -143,13 +143,15 @@ def train(args):
     print("Use DreamBooth method.")
     train_dataset = DreamBoothDataset(args.train_batch_size, args.train_data_dir, args.reg_data_dir,
                                       tokenizer, args.max_token_length, args.caption_extension, args.shuffle_caption, args.keep_tokens,
-                                      args.resolution, args.enable_bucket, args.min_bucket_reso, args.max_bucket_reso, args.prior_loss_weight,
-                                      args.flip_aug, args.color_aug, args.face_crop_aug_range, args.random_crop, args.debug_dataset)
+                                      args.resolution, args.enable_bucket, args.min_bucket_reso, args.max_bucket_reso,
+                                      args.bucket_reso_steps, args.bucket_no_upscale,
+                                      args.prior_loss_weight, args.flip_aug, args.color_aug, args.face_crop_aug_range, args.random_crop, args.debug_dataset)
   else:
     print("Train with captions.")
     train_dataset = FineTuningDataset(args.in_json, args.train_batch_size, args.train_data_dir,
                                       tokenizer, args.max_token_length, args.shuffle_caption, args.keep_tokens,
                                       args.resolution, args.enable_bucket, args.min_bucket_reso, args.max_bucket_reso,
+                                      args.bucket_reso_steps, args.bucket_no_upscale,
                                       args.flip_aug, args.color_aug, args.face_crop_aug_range, args.random_crop,
                                       args.dataset_repeats, args.debug_dataset)
 
@@ -205,6 +207,13 @@ def train(args):
       raise ImportError("No bitsand bytes / bitsandbytesがインストールされていないようです")
     print("use 8-bit Adam optimizer")
     optimizer_class = bnb.optim.AdamW8bit
+  elif args.use_lion_optimizer:
+    try:
+      import lion_pytorch
+    except ImportError:
+      raise ImportError("No lion_pytorch / lion_pytorch がインストールされていないようです")
+    print("use Lion optimizer")
+    optimizer_class = lion_pytorch.Lion
   else:
     optimizer_class = torch.optim.AdamW
 
@@ -233,7 +242,7 @@ def train(args):
       text_encoder, optimizer, train_dataloader, lr_scheduler)
 
   index_no_updates = torch.arange(len(tokenizer)) < token_ids[0]
-  print(len(index_no_updates), torch.sum(index_no_updates))
+  # print(len(index_no_updates), torch.sum(index_no_updates))
   orig_embeds_params = unwrap_model(text_encoder).get_input_embeddings().weight.data.detach().clone()
 
   # Freeze all parameters except for the token embeddings in text encoder
@@ -294,6 +303,7 @@ def train(args):
 
   for epoch in range(num_train_epochs):
     print(f"epoch {epoch+1}/{num_train_epochs}")
+    train_dataset.set_current_epoch(epoch + 1)
 
     text_encoder.train()
 
@@ -312,10 +322,14 @@ def train(args):
 
         # Get the text embedding for conditioning
         input_ids = batch["input_ids"].to(accelerator.device)
-        encoder_hidden_states = train_util.get_hidden_states(args, input_ids, tokenizer, text_encoder, torch.float) # weight_dtype) use float instead of fp16/bf16 because text encoder is float
+        # weight_dtype) use float instead of fp16/bf16 because text encoder is float
+        encoder_hidden_states = train_util.get_hidden_states(args, input_ids, tokenizer, text_encoder, torch.float)
 
         # Sample noise that we'll add to the latents
         noise = torch.randn_like(latents, device=latents.device)
+        if args.noise_offset:
+          # https://www.crosslabs.org//blog/diffusion-with-offset-noise
+          noise += args.noise_offset * torch.randn((latents.shape[0], latents.shape[1], 1, 1), device=latents.device)
 
         # Sample a random timestep for each image
         timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (b_size,), device=latents.device)
@@ -380,8 +394,8 @@ def train(args):
     accelerator.wait_for_everyone()
 
     updated_embs = unwrap_model(text_encoder).get_input_embeddings().weight[token_ids].data.detach().clone()
-    d = updated_embs - bef_epo_embs
-    print(bef_epo_embs.size(), updated_embs.size(), d.mean(), d.min())
+    # d = updated_embs - bef_epo_embs
+    # print(bef_epo_embs.size(), updated_embs.size(), d.mean(), d.min())
 
     if args.save_every_n_epochs is not None:
       model_name = train_util.DEFAULT_EPOCH_NAME if args.output_name is None else args.output_name
@@ -475,7 +489,7 @@ if __name__ == '__main__':
   parser = argparse.ArgumentParser()
 
   train_util.add_sd_models_arguments(parser)
-  train_util.add_dataset_arguments(parser, True, True)
+  train_util.add_dataset_arguments(parser, True, True, False)
   train_util.add_training_arguments(parser, True)
 
   parser.add_argument("--save_model_as", type=str, default="pt", choices=[None, "ckpt", "pt", "safetensors"],
@@ -488,7 +502,7 @@ if __name__ == '__main__':
   parser.add_argument("--token_string", type=str, default=None,
                       help="token string used in training, must not exist in tokenizer / 学習時に使用されるトークン文字列、tokenizerに存在しない文字であること")
   parser.add_argument("--init_word", type=str, default=None,
-                      help="word to initialize vector / ベクトルを初期化に使用する単語、tokenizerで一語になること")
+                      help="words to initialize vector / ベクトルを初期化に使用する単語、複数可")
   parser.add_argument("--use_object_template", action='store_true',
                       help="ignore caption and use default templates for object / キャプションは使わずデフォルトの物体用テンプレートで学習する")
   parser.add_argument("--use_style_template", action='store_true',