Upload 42 files

.gitattributes

@@ -1,34 +1,4 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.sh text eol=lf
+*.ps1 text eol=crlf
+*.bat text eol=crlf
+*.cmd text eol=crlf

.gitignore

@@ -0,0 +1,11 @@
+venv
+__pycache__
+cudnn_windows
+.vscode
+*.egg-info
+build
+wd14_tagger_model
+.DS_Store
+locon
+gui-user.bat
+gui-user.ps1

LICENSE.md

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [2022] [kohya-ss]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

README-ja.md

@@ -0,0 +1,147 @@
+## リポジトリについて
+Stable Diffusionの学習、画像生成、その他のスクリプトを入れたリポジトリです。
+
+[README in English](./README.md) ←更新情報はこちらにあります
+
+GUIやPowerShellスクリプトなど、より使いやすくする機能が[bmaltais氏のリポジトリ](https://github.com/bmaltais/kohya_ss)で提供されています（英語です）のであわせてご覧ください。bmaltais氏に感謝します。
+
+以下のスクリプトがあります。
+
+* DreamBooth、U-NetおよびText Encoderの学習をサポート
+* fine-tuning、同上
+* 画像生成
+* モデル変換（Stable Diffision ckpt/safetensorsとDiffusersの相互変換）
+
+## 使用法について
+
+当リポジトリ内およびnote.comに記事がありますのでそちらをご覧ください（将来的にはすべてこちらへ移すかもしれません）。
+
+* [学習について、共通編](./train_README-ja.md) : データ整備やオプションなど
+    * [データセット設定](./config_README-ja.md)
+* [DreamBoothの学習について](./train_db_README-ja.md)
+* [fine-tuningのガイド](./fine_tune_README_ja.md):
+* [LoRAの学習について](./train_network_README-ja.md)
+* [Textual Inversionの学習について](./train_ti_README-ja.md)
+* note.com [画像生成スクリプト](https://note.com/kohya_ss/n/n2693183a798e)
+* note.com [モデル変換スクリプト](https://note.com/kohya_ss/n/n374f316fe4ad)
+
+## Windowsでの動作に必要なプログラム
+
+Python 3.10.6およびGitが必要です。
+
+- Python 3.10.6: https://www.python.org/ftp/python/3.10.6/python-3.10.6-amd64.exe
+- git: https://git-scm.com/download/win
+
+PowerShellを使う場合、venvを使えるようにするためには以下の手順でセキュリティ設定を変更してください。
+（venvに限らずスクリプトの実行が可能になりますので注意してください。）
+
+- PowerShellを管理者として開きます。
+- 「Set-ExecutionPolicy Unrestricted」と入力し、Yと答えます。
+- 管理者のPowerShellを閉じます。
+
+## Windows環境でのインストール
+
+以下の例ではPyTorchは1.12.1／CUDA 11.6版をインストールします。CUDA 11.3版やPyTorch 1.13を使う場合は適宜書き換えください。
+
+（なお、python -m venv～の行で「python」とだけ表示された場合、py -m venv～のようにpythonをpyに変更してください。）
+
+通常の（管理者ではない）PowerShellを開き以下を順に実行します。
+
+```powershell
+git clone https://github.com/kohya-ss/sd-scripts.git
+cd sd-scripts
+
+python -m venv venv
+.\venv\Scripts\activate
+
+pip install torch==1.12.1+cu116 torchvision==0.13.1+cu116 --extra-index-url https://download.pytorch.org/whl/cu116
+pip install --upgrade -r requirements.txt
+pip install -U -I --no-deps https://github.com/C43H66N12O12S2/stable-diffusion-webui/releases/download/f/xformers-0.0.14.dev0-cp310-cp310-win_amd64.whl
+
+cp .\bitsandbytes_windows\*.dll .\venv\Lib\site-packages\bitsandbytes\
+cp .\bitsandbytes_windows\cextension.py .\venv\Lib\site-packages\bitsandbytes\cextension.py
+cp .\bitsandbytes_windows\main.py .\venv\Lib\site-packages\bitsandbytes\cuda_setup\main.py
+
+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
+-->
+
+コマンドプロンプトでは以下になります。
+
+
+```bat
+git clone https://github.com/kohya-ss/sd-scripts.git
+cd sd-scripts
+
+python -m venv venv
+.\venv\Scripts\activate
+
+pip install torch==1.12.1+cu116 torchvision==0.13.1+cu116 --extra-index-url https://download.pytorch.org/whl/cu116
+pip install --upgrade -r requirements.txt
+pip install -U -I --no-deps https://github.com/C43H66N12O12S2/stable-diffusion-webui/releases/download/f/xformers-0.0.14.dev0-cp310-cp310-win_amd64.whl
+
+copy /y .\bitsandbytes_windows\*.dll .\venv\Lib\site-packages\bitsandbytes\
+copy /y .\bitsandbytes_windows\cextension.py .\venv\Lib\site-packages\bitsandbytes\cextension.py
+copy /y .\bitsandbytes_windows\main.py .\venv\Lib\site-packages\bitsandbytes\cuda_setup\main.py
+
+accelerate config
+```
+
+（注:``python -m venv venv`` のほうが ``python -m venv --system-site-packages venv`` より安全そうなため書き換えました。globalなpythonにパッケージがインストールしてあると、後者だといろいろと問題が起きます。）
+
+accelerate configの質問には以下のように答えてください。（bf16で学習する場合、最後の質問にはbf16と答えてください。）
+
+※0.15.0から日本語環境では選択のためにカーソルキーを押すと落ちます（……）。数字キーの0、1、2……で選択できますので、そちらを使ってください。
+
+```txt
+- This machine
+- No distributed training
+- NO
+- NO
+- NO
+- all
+- fp16
+```
+
+※場合によって ``ValueError: fp16 mixed precision requires a GPU`` というエラーが出ることがあるようです。この場合、6番目の質問（
+``What GPU(s) (by id) should be used for training on this machine as a comma-separated list? [all]:``）に「0」と答えてください。（id `0`のGPUが使われます。）
+
+### PyTorchとxformersのバージョンについて
+
+他のバージョンでは学習がうまくいかない場合があるようです。特に他の理由がなければ指定のバージョンをお使いください。
+
+## アップグレード
+
+新しいリリースがあった場合、以下のコマンドで更新できます。
+
+```powershell
+cd sd-scripts
+git pull
+.\venv\Scripts\activate
+pip install --use-pep517 --upgrade -r requirements.txt
+```
+
+コマンドが成功すれば新しいバージョンが使用できます。
+
+## 謝意
+
+LoRAの実装は[cloneofsimo氏のリポジトリ](https://github.com/cloneofsimo/lora)を基にしたものです。感謝申し上げます。
+
+Conv2d 3x3への拡大は [cloneofsimo氏](https://github.com/cloneofsimo/lora) が最初にリリースし、KohakuBlueleaf氏が [LoCon](https://github.com/KohakuBlueleaf/LoCon) でその有効性を明らかにしたものです。KohakuBlueleaf氏に深く感謝します。
+
+## ライセンス
+
+スクリプトのライセンスはASL 2.0ですが（Diffusersおよびcloneofsimo氏のリポジトリ由来のものも同様）、一部他のライセンスのコードを含みます。
+
+[Memory Efficient Attention Pytorch](https://github.com/lucidrains/memory-efficient-attention-pytorch): MIT
+
+[bitsandbytes](https://github.com/TimDettmers/bitsandbytes): MIT
+
+[BLIP](https://github.com/salesforce/BLIP): BSD-3-Clause
+
+

README.md

@@ -0,0 +1,338 @@
+# Kohya's GUI
+
+This repository provides a Windows-focused Gradio GUI for [Kohya's Stable Diffusion trainers](https://github.com/kohya-ss/sd-scripts). The GUI allows you to set the training parameters and generate and run the required CLI commands to train the model.
+
+If you run on Linux and would like to use the GUI, there is now a port of it as a docker container. You can find the project [here](https://github.com/P2Enjoy/kohya_ss-docker).
+
+### Table of Contents
+
+- [Tutorials](#tutorials)
+- [Required Dependencies](#required-dependencies)
+  - [Linux/macOS](#linux-and-macos-dependencies)
+- [Installation](#installation)
+    - [Linux/macOS](#linux-and-macos)
+      - [Default Install Locations](#install-location)
+    - [Windows](#windows)
+    - [CUDNN 8.6](#optional--cudnn-86)
+- [Upgrading](#upgrading)
+  - [Windows](#windows-upgrade)
+  - [Linux/macOS](#linux-and-macos-upgrade)
+- [Launching the GUI](#starting-gui-service)
+  - [Windows](#launching-the-gui-on-windows)
+  - [Linux/macOS](#launching-the-gui-on-linux-and-macos)
+  - [Direct Launch via Python Script](#launching-the-gui-directly-using-kohyaguipy)
+- [Dreambooth](#dreambooth)
+- [Finetune](#finetune)
+- [Train Network](#train-network)
+- [LoRA](#lora)
+- [Troubleshooting](#troubleshooting)
+  - [Page File Limit](#page-file-limit)
+  - [No module called tkinter](#no-module-called-tkinter)
+  - [FileNotFoundError](#filenotfounderror)
+- [Change History](#change-history)
+
+## Tutorials
+
+[How to Create a LoRA Part 1: Dataset Preparation](https://www.youtube.com/watch?v=N4_-fB62Hwk):
+
+[![LoRA Part 1 Tutorial](https://img.youtube.com/vi/N4_-fB62Hwk/0.jpg)](https://www.youtube.com/watch?v=N4_-fB62Hwk)
+
+[How to Create a LoRA Part 2: Training the Model](https://www.youtube.com/watch?v=k5imq01uvUY):
+
+[![LoRA Part 2 Tutorial](https://img.youtube.com/vi/k5imq01uvUY/0.jpg)](https://www.youtube.com/watch?v=k5imq01uvUY)
+
+## Required Dependencies
+
+- Install [Python 3.10](https://www.python.org/ftp/python/3.10.9/python-3.10.9-amd64.exe) 
+  - make sure to tick the box to add Python to the 'PATH' environment variable
+- Install [Git](https://git-scm.com/download/win)
+- Install [Visual Studio 2015, 2017, 2019, and 2022 redistributable](https://aka.ms/vs/17/release/vc_redist.x64.exe)
+
+### Linux and macOS dependencies
+
+These dependencies are taken care of via `setup.sh` in the installation section. No additional steps should be needed unless the scripts inform you otherwise.
+
+## Installation
+
+### Runpod
+Follow the instructions found in this discussion: https://github.com/bmaltais/kohya_ss/discussions/379
+
+### Linux and macOS
+In the terminal, run
+
+```
+git clone https://github.com/bmaltais/kohya_ss.git
+cd kohya_ss
+# May need to chmod +x ./setup.sh if you're on a machine with stricter security.
+# There are additional options if needed for a runpod environment.
+# Call 'setup.sh -h' or 'setup.sh --help' for more information.
+./setup.sh
+```
+
+Setup.sh help included here:
+
+```bash
+Kohya_SS Installation Script for POSIX operating systems.
+
+The following options are useful in a runpod environment,
+but will not affect a local machine install.
+
+Usage:
+  setup.sh -b dev -d /workspace/kohya_ss -g https://mycustom.repo.tld/custom_fork.git
+  setup.sh --branch=dev --dir=/workspace/kohya_ss --git-repo=https://mycustom.repo.tld/custom_fork.git
+
+Options:
+  -b BRANCH, --branch=BRANCH    Select which branch of kohya to check out on new installs.
+  -d DIR, --dir=DIR             The full path you want kohya_ss installed to.
+  -g REPO, --git_repo=REPO      You can optionally provide a git repo to check out for runpod installation. Useful for custom forks.
+  -h, --help                    Show this screen.
+  -i, --interactive             Interactively configure accelerate instead of using default config file.
+  -n, --no-update               Do not update kohya_ss repo. No git pull or clone operations.
+  -p, --public                  Expose public URL in runpod mode. Won't have an effect in other modes.
+  -r, --runpod                  Forces a runpod installation. Useful if detection fails for any reason.
+  -s, --skip-space-check        Skip the 10Gb minimum storage space check.
+  -u, --no-gui                  Skips launching the GUI.
+  -v, --verbose                 Increase verbosity levels up to 3.
+```
+
+#### Install location
+
+The default install location for Linux is where the script is located if a previous installation is detected that location.
+Otherwise, it will fall to `/opt/kohya_ss`. If /opt is not writeable, the fallback is `$HOME/kohya_ss`. Lastly, if all else fails it will simply install to the current folder you are in (PWD).
+
+On macOS and other non-Linux machines, it will first try to detect an install where the script is run from and then run setup there if that's detected. 
+If a previous install isn't found at that location, then it will default install to `$HOME/kohya_ss` followed by where you're currently at if there's no access to $HOME.
+You can override this behavior by specifying an install directory with the -d option.
+
+If you are using the interactive mode, our default values for the accelerate config screen after running the script answer "This machine", "None", "No" for the remaining questions.
+These are the same answers as the Windows install.
+
+### Windows
+
+- Install [Python 3.10](https://www.python.org/ftp/python/3.10.9/python-3.10.9-amd64.exe) 
+  - make sure to tick the box to add Python to the 'PATH' environment variable
+- Install [Git](https://git-scm.com/download/win)
+- Install [Visual Studio 2015, 2017, 2019, and 2022 redistributable](https://aka.ms/vs/17/release/vc_redist.x64.exe)
+
+In the terminal, run:
+
+```
+git clone https://github.com/bmaltais/kohya_ss.git
+cd kohya_ss
+.\setup.bat
+```
+
+If this is a 1st install answer No when asked `Do you want to uninstall previous versions of torch and associated files before installing`.
+
+
+Then configure accelerate with the same answers as in the MacOS instructions when prompted.
+
+### Optional: CUDNN 8.6
+
+This step is optional but can improve the learning speed for NVIDIA 30X0/40X0 owners. It allows for larger training batch size and faster training speed.
+
+Due to the file size, I can't host the DLLs needed for CUDNN 8.6 on Github. I strongly advise you download them for a speed boost in sample generation (almost 50% on 4090 GPU) you can download them [here](https://b1.thefileditch.ch/mwxKTEtelILoIbMbruuM.zip).
+
+To install, simply unzip the directory and place the `cudnn_windows` folder in the root of the this repo.
+
+Run the following commands to install:
+
+```
+.\venv\Scripts\activate
+
+python .\tools\cudann_1.8_install.py
+```
+
+Once the commands have completed successfully you should be ready to use the new version. MacOS support is not tested and has been mostly taken from https://gist.github.com/jstayco/9f5733f05b9dc29de95c4056a023d645
+
+## Upgrading
+
+The following commands will work from the root directory of the project if you'd prefer to not run scripts.
+These commands will work on any OS.
+```bash
+git pull
+
+.\venv\Scripts\activate
+
+pip install --use-pep517 --upgrade -r requirements.txt
+```
+
+### Windows Upgrade
+When a new release comes out, you can upgrade your repo with the following commands in the root directory:
+
+```powershell
+upgrade.bat
+```
+
+### Linux and macOS Upgrade
+You can cd into the root directory and simply run
+
+```bash
+# Refresh and update everything
+./setup.sh
+
+# This will refresh everything, but NOT clone or pull the git repo.
+./setup.sh --no-git-update
+```
+
+Once the commands have completed successfully you should be ready to use the new version.
+
+# Starting GUI Service
+
+The following command line arguments can be passed to the scripts on any OS to configure the underlying service.
+```
+--listen: the IP address to listen on for connections to Gradio.
+--username: a username for authentication. 
+--password: a password for authentication. 
+--server_port: the port to run the server listener on. 
+--inbrowser: opens the Gradio UI in a web browser. 
+--share: shares the Gradio UI.
+```
+
+### Launching the GUI on Windows
+
+The two scripts to launch the GUI on Windows are gui.ps1 and gui.bat in the root directory.
+You can use whichever script you prefer.
+
+To launch the Gradio UI, run the script in a terminal with the desired command line arguments, for example:
+
+`gui.ps1 --listen 127.0.0.1 --server_port 7860 --inbrowser --share`
+
+or
+
+`gui.bat --listen 127.0.0.1 --server_port 7860 --inbrowser --share`
+
+## Launching the GUI on Linux and macOS
+
+Run the launcher script with the desired command line arguments similar to Windows.
+`gui.sh --listen 127.0.0.1 --server_port 7860 --inbrowser --share`
+
+## Launching the GUI directly using kohya_gui.py
+
+To run the GUI directly bypassing the wrapper scripts, simply use this command from the root project directory:
+
+```
+.\venv\Scripts\activate
+
+python .\kohya_gui.py
+```
+
+## Dreambooth
+
+You can find the dreambooth solution specific here: [Dreambooth README](train_db_README.md)
+
+## Finetune
+
+You can find the finetune solution specific here: [Finetune README](fine_tune_README.md)
+
+## Train Network
+
+You can find the train network solution specific here: [Train network README](train_network_README.md)
+
+## LoRA
+
+Training a LoRA currently uses the `train_network.py` code. You can create a LoRA network by using the all-in-one `gui.cmd` or by running the dedicated LoRA training GUI with:
+
+```
+.\venv\Scripts\activate
+
+python lora_gui.py
+```
+
+Once you have created the LoRA network, you can generate images via auto1111 by installing [this extension](https://github.com/kohya-ss/sd-webui-additional-networks).
+
+### Naming of LoRA
+
+The LoRA supported by `train_network.py` has been named to avoid confusion. The documentation has been updated. The following are the names of LoRA types in this repository.
+
+1. __LoRA-LierLa__ : (LoRA for __Li__ n __e__ a __r__  __La__ yers)
+
+    LoRA for Linear layers and Conv2d layers with 1x1 kernel
+
+2. __LoRA-C3Lier__ : (LoRA for __C__ olutional layers with __3__ x3 Kernel and  __Li__ n __e__ a __r__ layers)
+
+    In addition to 1., LoRA for Conv2d layers with 3x3 kernel 
+    
+LoRA-LierLa is the default LoRA type for `train_network.py` (without `conv_dim` network arg). LoRA-LierLa can be used with [our extension](https://github.com/kohya-ss/sd-webui-additional-networks) for AUTOMATIC1111's Web UI, or with the built-in LoRA feature of the Web UI.
+
+To use LoRA-C3Liar with Web UI, please use our extension.
+
+## Sample image generation during training
+A prompt file might look like this, for example
+
+```
+# prompt 1
+masterpiece, best quality, (1girl), in white shirts, upper body, looking at viewer, simple background --n low quality, worst quality, bad anatomy,bad composition, poor, low effort --w 768 --h 768 --d 1 --l 7.5 --s 28
+
+# prompt 2
+masterpiece, best quality, 1boy, in business suit, standing at street, looking back --n (low quality, worst quality), bad anatomy,bad composition, poor, low effort --w 576 --h 832 --d 2 --l 5.5 --s 40
+```
+
+  Lines beginning with `#` are comments. You can specify options for the generated image with options like `--n` after the prompt. The following can be used.
+
+  * `--n` Negative prompt up to the next option.
+  * `--w` Specifies the width of the generated image.
+  * `--h` Specifies the height of the generated image.
+  * `--d` Specifies the seed of the generated image.
+  * `--l` Specifies the CFG scale of the generated image.
+  * `--s` Specifies the number of steps in the generation.
+
+  The prompt weighting such as `( )` and `[ ]` are working.
+
+## Troubleshooting
+
+### Page File Limit
+
+- X error relating to `page file`: Increase the page file size limit in Windows.
+
+### No module called tkinter
+
+- Re-install [Python 3.10](https://www.python.org/ftp/python/3.10.9/python-3.10.9-amd64.exe) on your system.
+
+### FileNotFoundError
+
+This is usually related to an installation issue. Make sure you do not have any python modules installed locally that could conflict with the ones installed in the venv:
+
+1. Open a new powershell terminal and make sure no venv is active.
+2.  Run the following commands:
+
+```
+pip freeze > uninstall.txt
+pip uninstall -r uninstall.txt
+```
+
+This will store a backup file with your current locally installed pip packages and then uninstall them. Then, redo the installation instructions within the kohya_ss venv.
+
+## Change History
+
+* 2023/04/22 (v21.5.5)
+    - Update LoRA merge GUI to support SD checkpoint merge and up to 4 LoRA merging
+    - Fixed `lora_interrogator.py` not working. Please refer to [PR #392](https://github.com/kohya-ss/sd-scripts/pull/392) for details. Thank you A2va and heyalexchoi!
+    - Fixed the handling of tags containing `_` in `tag_images_by_wd14_tagger.py`.
+    - Add new Extract DyLoRA gui to the Utilities tab.
+    - Add new Merge LyCORIS models into checkpoint gui to the Utilities tab.
+    - Add new info on startup to help debug things
+* 2023/04/17 (v21.5.4)
+    - Fixed a bug that caused an error when loading DyLoRA with the `--network_weight` option in `train_network.py`.
+    - Added the `--recursive` option to each script in the `finetune` folder to process folders recursively. Please refer to [PR #400](https://github.com/kohya-ss/sd-scripts/pull/400/) for details. Thanks to Linaqruf!
+    - Upgrade Gradio to latest release
+    - Fix issue when Adafactor is used as optimizer and LR Warmup is not 0: https://github.com/bmaltais/kohya_ss/issues/617
+    - Added support for DyLoRA in `train_network.py`. Please refer to [here](./train_network_README-ja.md#dylora) for details (currently only in Japanese).
+    - Added support for caching latents to disk in each training script. Please specify __both__ `--cache_latents` and `--cache_latents_to_disk` options.
+        - The files are saved in the same folder as the images with the extension `.npz`. If you specify the `--flip_aug` option, the files with `_flip.npz` will also be saved.
+        - Multi-GPU training has not been tested.
+        - This feature is not tested with all combinations of datasets and training scripts, so there may be bugs.
+    - Added workaround for an error that occurs when training with `fp16` or `bf16` in `fine_tune.py`.
+    - Implemented DyLoRA GUI support. There will now be a new 'DyLoRA Unit` slider when the LoRA type is selected as `kohya DyLoRA` to specify the desired Unit value for DyLoRA training.
+    - Update gui.bat and gui.ps1 based on: https://github.com/bmaltais/kohya_ss/issues/188
+    - Update `setup.bat` to install torch 2.0.0 instead of 1.2.1. If you want to upgrade from 1.2.1 to 2.0.0 run setup.bat again, select 1 to uninstall the previous torch modules, then select 2 for torch 2.0.0
+
+* 2023/04/09 (v21.5.2)
+
+    - Added support for training with weighted captions. Thanks to AI-Casanova for the great contribution! 
+    - Please refer to the PR for details: [PR #336](https://github.com/kohya-ss/sd-scripts/pull/336)
+    - Specify the `--weighted_captions` option. It is available for all training scripts except Textual Inversion and XTI.
+    - This option is also applicable to token strings of the DreamBooth method.
+    - The syntax for weighted captions is almost the same as the Web UI, and you can use things like `(abc)`, `[abc]`, and `(abc:1.23)`. Nesting is also possible.
+    - If you include a comma in the parentheses, the parentheses will not be properly matched in the prompt shuffle/dropout, so do not include a comma in the parentheses.
+    - Run gui.sh from any place

XTI_hijack.py

@@ -0,0 +1,209 @@
+import torch
+from typing import Union, List, Optional, Dict, Any, Tuple
+from diffusers.models.unet_2d_condition import UNet2DConditionOutput
+
+def unet_forward_XTI(self,
+        sample: torch.FloatTensor,
+        timestep: Union[torch.Tensor, float, int],
+        encoder_hidden_states: torch.Tensor,
+        class_labels: Optional[torch.Tensor] = None,
+        return_dict: bool = True,
+    ) -> Union[UNet2DConditionOutput, Tuple]:
+        r"""
+        Args:
+            sample (`torch.FloatTensor`): (batch, channel, height, width) noisy inputs tensor
+            timestep (`torch.FloatTensor` or `float` or `int`): (batch) timesteps
+            encoder_hidden_states (`torch.FloatTensor`): (batch, sequence_length, feature_dim) encoder hidden states
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`models.unet_2d_condition.UNet2DConditionOutput`] instead of a plain tuple.
+
+        Returns:
+            [`~models.unet_2d_condition.UNet2DConditionOutput`] or `tuple`:
+            [`~models.unet_2d_condition.UNet2DConditionOutput`] if `return_dict` is True, otherwise a `tuple`. When
+            returning a tuple, the first element is the sample tensor.
+        """
+        # By default samples have to be AT least a multiple of the overall upsampling factor.
+        # The overall upsampling factor is equal to 2 ** (# num of upsampling layears).
+        # However, the upsampling interpolation output size can be forced to fit any upsampling size
+        # on the fly if necessary.
+        default_overall_up_factor = 2**self.num_upsamplers
+
+        # upsample size should be forwarded when sample is not a multiple of `default_overall_up_factor`
+        forward_upsample_size = False
+        upsample_size = None
+
+        if any(s % default_overall_up_factor != 0 for s in sample.shape[-2:]):
+            logger.info("Forward upsample size to force interpolation output size.")
+            forward_upsample_size = True
+
+        # 0. center input if necessary
+        if self.config.center_input_sample:
+            sample = 2 * sample - 1.0
+
+        # 1. time
+        timesteps = timestep
+        if not torch.is_tensor(timesteps):
+            # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can
+            # This would be a good case for the `match` statement (Python 3.10+)
+            is_mps = sample.device.type == "mps"
+            if isinstance(timestep, float):
+                dtype = torch.float32 if is_mps else torch.float64
+            else:
+                dtype = torch.int32 if is_mps else torch.int64
+            timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device)
+        elif len(timesteps.shape) == 0:
+            timesteps = timesteps[None].to(sample.device)
+
+        # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+        timesteps = timesteps.expand(sample.shape[0])
+
+        t_emb = self.time_proj(timesteps)
+
+        # timesteps does not contain any weights and will always return f32 tensors
+        # but time_embedding might actually be running in fp16. so we need to cast here.
+        # there might be better ways to encapsulate this.
+        t_emb = t_emb.to(dtype=self.dtype)
+        emb = self.time_embedding(t_emb)
+
+        if self.config.num_class_embeds is not None:
+            if class_labels is None:
+                raise ValueError("class_labels should be provided when num_class_embeds > 0")
+            class_emb = self.class_embedding(class_labels).to(dtype=self.dtype)
+            emb = emb + class_emb
+
+        # 2. pre-process
+        sample = self.conv_in(sample)
+
+        # 3. down
+        down_block_res_samples = (sample,)
+        down_i = 0
+        for downsample_block in self.down_blocks:
+            if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention:
+                sample, res_samples = downsample_block(
+                    hidden_states=sample,
+                    temb=emb,
+                    encoder_hidden_states=encoder_hidden_states[down_i:down_i+2],
+                )
+                down_i += 2
+            else:
+                sample, res_samples = downsample_block(hidden_states=sample, temb=emb)
+
+            down_block_res_samples += res_samples
+
+        # 4. mid
+        sample = self.mid_block(sample, emb, encoder_hidden_states=encoder_hidden_states[6])
+
+        # 5. up
+        up_i = 7
+        for i, upsample_block in enumerate(self.up_blocks):
+            is_final_block = i == len(self.up_blocks) - 1
+
+            res_samples = down_block_res_samples[-len(upsample_block.resnets) :]
+            down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)]
+
+            # if we have not reached the final block and need to forward the
+            # upsample size, we do it here
+            if not is_final_block and forward_upsample_size:
+                upsample_size = down_block_res_samples[-1].shape[2:]
+
+            if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention:
+                sample = upsample_block(
+                    hidden_states=sample,
+                    temb=emb,
+                    res_hidden_states_tuple=res_samples,
+                    encoder_hidden_states=encoder_hidden_states[up_i:up_i+3],
+                    upsample_size=upsample_size,
+                )
+                up_i += 3
+            else:
+                sample = upsample_block(
+                    hidden_states=sample, temb=emb, res_hidden_states_tuple=res_samples, upsample_size=upsample_size
+                )
+        # 6. post-process
+        sample = self.conv_norm_out(sample)
+        sample = self.conv_act(sample)
+        sample = self.conv_out(sample)
+
+        if not return_dict:
+            return (sample,)
+
+        return UNet2DConditionOutput(sample=sample)
+
+def downblock_forward_XTI(
+    self, hidden_states, temb=None, encoder_hidden_states=None, attention_mask=None, cross_attention_kwargs=None
+):
+    output_states = ()
+    i = 0
+
+    for resnet, attn in zip(self.resnets, self.attentions):
+        if self.training and self.gradient_checkpointing:
+
+            def create_custom_forward(module, return_dict=None):
+                def custom_forward(*inputs):
+                    if return_dict is not None:
+                        return module(*inputs, return_dict=return_dict)
+                    else:
+                        return module(*inputs)
+
+                return custom_forward
+
+            hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(resnet), hidden_states, temb)
+            hidden_states = torch.utils.checkpoint.checkpoint(
+                create_custom_forward(attn, return_dict=False), hidden_states, encoder_hidden_states[i]
+            )[0]
+        else:
+            hidden_states = resnet(hidden_states, temb)
+            hidden_states = attn(hidden_states, encoder_hidden_states=encoder_hidden_states[i]).sample
+
+        output_states += (hidden_states,)
+        i += 1
+
+    if self.downsamplers is not None:
+        for downsampler in self.downsamplers:
+            hidden_states = downsampler(hidden_states)
+
+        output_states += (hidden_states,)
+
+    return hidden_states, output_states
+
+def upblock_forward_XTI(
+    self,
+    hidden_states,
+    res_hidden_states_tuple,
+    temb=None,
+    encoder_hidden_states=None,
+    upsample_size=None,
+):
+    i = 0
+    for resnet, attn in zip(self.resnets, self.attentions):
+        # pop res hidden states
+        res_hidden_states = res_hidden_states_tuple[-1]
+        res_hidden_states_tuple = res_hidden_states_tuple[:-1]
+        hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)
+
+        if self.training and self.gradient_checkpointing:
+
+            def create_custom_forward(module, return_dict=None):
+                def custom_forward(*inputs):
+                    if return_dict is not None:
+                        return module(*inputs, return_dict=return_dict)
+                    else:
+                        return module(*inputs)
+
+                return custom_forward
+
+            hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(resnet), hidden_states, temb)
+            hidden_states = torch.utils.checkpoint.checkpoint(
+                create_custom_forward(attn, return_dict=False), hidden_states, encoder_hidden_states[i]
+            )[0]
+        else:
+            hidden_states = resnet(hidden_states, temb)
+            hidden_states = attn(hidden_states, encoder_hidden_states=encoder_hidden_states[i]).sample
+        
+        i += 1
+
+    if self.upsamplers is not None:
+        for upsampler in self.upsamplers:
+            hidden_states = upsampler(hidden_states, upsample_size)
+
+    return hidden_states

_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"]

activate.bat

@@ -0,0 +1,5 @@
+@echo off
+
+call venv\Scripts\activate.bat
+
+pause

activate.ps1

@@ -0,0 +1 @@
+.\venv\Scripts\activate

config_README-ja.md

@@ -0,0 +1,279 @@
+For non-Japanese speakers: this README is provided only in Japanese in the current state. Sorry for inconvenience. We will provide English version in the near future.
+
+`--dataset_config` で渡すことができる設定ファイルに関する説明です。
+
+## 概要
+
+設定ファイルを渡すことにより、ユーザが細かい設定を行えるようにします。
+
+* 複数のデータセットが設定可能になります
+    * 例えば `resolution` をデータセットごとに設定して、それらを混合して学習できます。
+    * DreamBooth の手法と fine tuning の手法の両方に対応している学習方法では、DreamBooth 方式と fine tuning 方式のデータセットを混合することが可能です。
+* サブセットごとに設定を変更することが可能になります
+    * データセットを画像ディレクトリ別またはメタデータ別に分割したものがサブセットです。いくつかのサブセットが集まってデータセットを構成します。
+    * `keep_tokens` や `flip_aug` 等のオプションはサブセットごとに設定可能です。一方、`resolution` や `batch_size` といったオプションはデータセットごとに設定可能で、同じデータセットに属するサブセットでは値が共通になります。詳しくは後述します。
+
+設定ファイルの形式は JSON か TOML を利用できます。記述のしやすさを考えると [TOML](https://toml.io/ja/v1.0.0-rc.2) を利用するのがオススメです。以下、TOML の利用を前提に説明します。
+
+TOML で記述した設定ファイルの例です。
+
+```toml
+[general]
+shuffle_caption = true
+caption_extension = '.txt'
+keep_tokens = 1
+
+# これは DreamBooth 方式のデータセット
+[[datasets]]
+resolution = 512
+batch_size = 4
+keep_tokens = 2
+
+  [[datasets.subsets]]
+  image_dir = 'C:\hoge'
+  class_tokens = 'hoge girl'
+  # このサブセットは keep_tokens = 2 （所属する datasets の値が使われる）
+
+  [[datasets.subsets]]
+  image_dir = 'C:\fuga'
+  class_tokens = 'fuga boy'
+  keep_tokens = 3
+
+  [[datasets.subsets]]
+  is_reg = true
+  image_dir = 'C:\reg'
+  class_tokens = 'human'
+  keep_tokens = 1
+
+# これは fine tuning 方式のデータセット
+[[datasets]]
+resolution = [768, 768]
+batch_size = 2
+
+  [[datasets.subsets]]
+  image_dir = 'C:\piyo'
+  metadata_file = 'C:\piyo\piyo_md.json'
+  # このサブセットは keep_tokens = 1 （general の値が使われる）
+```
+
+この例では、3 つのディレクトリを DreamBooth 方式のデータセットとして 512x512 (batch size 4) で学習させ、1 つのディレクトリを fine tuning 方式のデータセットとして 768x768 (batch size 2) で学習させることになります。
+
+## データセット・サブセットに関する設定
+
+データセット・サブセットに関する設定は、登録可能な箇所がいくつかに分かれています。
+
+* `[general]`
+    * 全データセットまたは全サブセットに適用されるオプションを指定する箇所です。
+    * データセットごとの設定及びサブセットごとの設定に同名のオプションが存在していた場合には、データセット・サブセットごとの設定が優先されます。
+* `[[datasets]]`
+    * `datasets` はデータセットに関する設定の登録箇所になります。各データセットに個別に適用されるオプションを指定する箇所です。
+    * サブセットごとの設定が存在していた場合には、サブセットごとの設定が優先されます。
+* `[[datasets.subsets]]`
+    * `datasets.subsets` はサブセットに関する設定の登録箇所になります。各サブセットに個別に適用されるオプションを指定する箇所です。
+
+先程の例における、画像ディレクトリと登録箇所の対応に関するイメージ図です。
+
+```
+C:\
+├─ hoge  ->  [[datasets.subsets]] No.1  ┐                        ┐
+├─ fuga  ->  [[datasets.subsets]] No.2  |->  [[datasets]] No.1   |->  [general]
+├─ reg   ->  [[datasets.subsets]] No.3  ┘                        |
+└─ piyo  ->  [[datasets.subsets]] No.4  -->  [[datasets]] No.2   ┘
+```
+
+画像ディレクトリがそれぞれ1つの `[[datasets.subsets]]` に対応しています。そして `[[datasets.subsets]]` が1つ以上組み合わさって1つの `[[datasets]]` を構成します。`[general]` には全ての `[[datasets]]`, `[[datasets.subsets]]` が属します。
+
+登録箇所ごとに指定可能なオプションは異なりますが、同名のオプションが指定された場合は下位の登録箇所にある値が優先されます。先程の例の `keep_tokens` オプションの扱われ方を確認してもらうと理解しやすいかと思います。
+
+加えて、学習方法が対応している手法によっても指定可能なオプションが変化します。
+
+* DreamBooth 方式専用のオプション
+* fine tuning 方式専用のオプション
+* caption dropout の手法が使える場合のオプション
+
+DreamBooth の手法と fine tuning の手法の両方とも利用可能な学習方法では、両者を併用することができます。
+併用する際の注意点として、DreamBooth 方式なのか fine tuning 方式なのかはデータセット単位で判別を行っているため、同じデータセット中に DreamBooth 方式のサブセットと fine tuning 方式のサブセットを混在させることはできません。
+つまり、これらを併用したい場合には異なる方式のサブセットが異なるデータセットに所属するように設定する必要があります。
+
+プログラムの挙動としては、後述する `metadata_file` オプションが存在していたら fine tuning 方式のサブセットだと判断します。
+そのため、同一のデータセットに所属するサブセットについて言うと、「全てが `metadata_file` オプションを持つ」か「全てが `metadata_file` オプションを持たない」かのどちらかになっていれば問題ありません。
+
+以下、利用可能なオプションを説明します。コマンドライン引数と名称が同一のオプションについては、基本的に説明を割愛します。他の README を参照してください。
+
+### 全学習方法で共通のオプション
+
+学習方法によらずに指定可能なオプションです。
+
+#### データセット向けオプション
+
+データセットの設定に関わるオプションです。`datasets.subsets` には記述できません。
+
+| オプション名 | 設定例 | `[general]` | `[[datasets]]` |
+| ---- | ---- | ---- | ---- |
+| `batch_size` | `1` | o | o |
+| `bucket_no_upscale` | `true` | o | o |
+| `bucket_reso_steps` | `64` | o | o |
+| `enable_bucket` | `true` | o | o |
+| `max_bucket_reso` | `1024` | o | o |
+| `min_bucket_reso` | `128` | o | o |
+| `resolution` | `256`, `[512, 512]` | o | o |
+
+* `batch_size`
+    * コマンドライン引数の `--train_batch_size` と同等です。
+
+これらの設定はデータセットごとに固定です。
+つまり、データセットに所属するサブセットはこれらの設定を共有することになります。
+例えば解像度が異なるデータセットを用意したい場合は、上に挙げた例のように別々のデータセットとして定義すれば別々の解像度を設定可能です。
+
+#### サブセット向けオプション
+
+サブセットの設定に関わるオプションです。
+
+| オプション名 | 設定例 | `[general]` | `[[datasets]]` | `[[dataset.subsets]]` |
+| ---- | ---- | ---- | ---- | ---- |
+| `color_aug` | `false` | o | o | o |
+| `face_crop_aug_range` | `[1.0, 3.0]` | o | o | o |
+| `flip_aug` | `true` | o | o | o |
+| `keep_tokens` | `2` | o | o | o |
+| `num_repeats` | `10` | o | o | o |
+| `random_crop` | `false` | o | o | o |
+| `shuffle_caption` | `true` | o | o | o |
+
+* `num_repeats`
+    * サブセットの画像の繰り返し回数を指定します。fine tuning における `--dataset_repeats` に相当しますが、`num_repeats` はどの学習方法でも指定可能です。
+
+### DreamBooth 方式専用のオプション
+
+DreamBooth 方式のオプションは、サブセット向けオプションのみ存在します。
+
+#### サブセット向けオプション
+
+DreamBooth 方式のサブセットの設定に関わるオプションです。
+
+| オプション名 | 設定例 | `[general]` | `[[datasets]]` | `[[dataset.subsets]]` |
+| ---- | ---- | ---- | ---- | ---- |
+| `image_dir` | `‘C:\hoge’` | - | - | o（必須） |
+| `caption_extension` | `".txt"` | o | o | o |
+| `class_tokens` | `“sks girl”` | - | - | o |
+| `is_reg` | `false` | - | - | o |
+
+まず注意点として、 `image_dir` には画像ファイルが直下に置かれているパスを指定する必要があります。従来の DreamBooth の手法ではサブディレクトリに画像を置く必要がありましたが、そちらとは仕様に互換性がありません。また、`5_cat` のようなフォルダ名にしても、画像の繰り返し回数とクラス名は反映されません。これらを個別に設定したい場合、`num_repeats` と `class_tokens` で明示的に指定する必要があることに注意してください。
+
+* `image_dir`
+    * 画像ディレクトリのパスを指定します。指定必須オプションです。
+    * 画像はディレクトリ直下に置かれている必要があります。
+* `class_tokens`
+    * クラストークンを設定します。
+    * 画像に対応する caption ファイルが存在しない場合にのみ学習時に利用されます。利用するかどうかの判定は画像ごとに行います。`class_tokens` を指定しなかった場合に caption ファイル���見つからなかった場合にはエラーになります。
+* `is_reg`
+    * サブセットの画像が正規化用かどうかを指定します。指定しなかった場合は `false` として、つまり正規化画像ではないとして扱います。
+
+### fine tuning 方式専用のオプション
+
+fine tuning 方式のオプションは、サブセット向けオプションのみ存在します。
+
+#### サブセット向けオプション
+
+fine tuning 方式のサブセットの設定に関わるオプションです。
+
+| オプション名 | 設定例 | `[general]` | `[[datasets]]` | `[[dataset.subsets]]` |
+| ---- | ---- | ---- | ---- | ---- |
+| `image_dir` | `‘C:\hoge’` | - | - | o |
+| `metadata_file` | `'C:\piyo\piyo_md.json'` | - | - | o（必須） |
+
+* `image_dir`
+    * 画像ディレクトリのパスを指定します。DreamBooth の手法の方とは異なり指定は必須ではありませんが、設定することを推奨します。
+        * 指定する必要がない状況としては、メタデータファイルの生成時に `--full_path` を付与して実行していた場合です。
+    * 画像はディレクトリ直下に置かれている必要があります。
+* `metadata_file`
+    * サブセットで利用されるメタデータファイルのパスを指定します。指定必須オプションです。
+        * コマンドライン引数の `--in_json` と同等です。
+    * サブセットごとにメタデータファイルを指定する必要がある仕様上、ディレクトリを跨いだメタデータを1つのメタデータファイルとして作成することは避けた方が良いでしょう。画像ディレクトリごとにメタデータファイルを用意し、それらを別々のサブセットとして登録することを強く推奨します。
+
+### caption dropout の手法が使える場合に指定可能なオプション
+
+caption dropout の手法が使える場合のオプションは、サブセット向けオプションのみ存在します。
+DreamBooth 方式か fine tuning 方式かに関わらず、caption dropout に対応している学習方法であれば指定可能です。
+
+#### サブセット向けオプション
+
+caption dropout が使えるサブセットの設定に関わるオプションです。
+
+| オプション名 | `[general]` | `[[datasets]]` | `[[dataset.subsets]]` |
+| ---- | ---- | ---- | ---- |
+| `caption_dropout_every_n_epochs` | o | o | o |
+| `caption_dropout_rate` | o | o | o |
+| `caption_tag_dropout_rate` | o | o | o |
+
+## 重複したサブセットが存在する時の挙動
+
+DreamBooth 方式のデータセットの場合、その中にある `image_dir` が同一のサブセットは重複していると見なされます。
+fine tuning 方式のデータセットの場合は、その中にある `metadata_file` が同一のサブセットは重複していると見なされます。
+データセット中に重複したサブセットが存在する場合、2個目以降は無視されます。
+
+一方、異なるデータセットに所属している場合は、重複しているとは見なされません。
+例えば、以下のように同一の `image_dir` を持つサブセットを別々のデータセットに入れた場合には、重複していないと見なします。
+これは、同じ画像でも異なる解像度で学習したい場合に役立ちます。
+
+```toml
+# 別々のデータセットに存在している場合は重複とは見なされず、両方とも学習に使われる
+
+[[datasets]]
+resolution = 512
+
+  [[datasets.subsets]]
+  image_dir = 'C:\hoge'
+
+[[datasets]]
+resolution = 768
+
+  [[datasets.subsets]]
+  image_dir = 'C:\hoge'
+```
+
+## コマンドライン引数との併用
+
+設定ファイルのオプションの中には、コマンドライン引数のオプションと役割が重複しているものがあります。
+
+以下に挙げるコマンドライン引数のオプションは、設定ファイルを渡した場合には無視されます。
+
+* `--train_data_dir`
+* `--reg_data_dir`
+* `--in_json`
+
+以下に挙げるコマンドライン引数のオプションは、コマンドライン引数と設定ファイルで同時に指定された場合、コマンドライン引数の値よりも設定ファイルの値が優先されます。特に断りがなければ同名のオプションとなります。
+
+| コマンドライン引数のオプション     | 優先される設定ファイルのオプション |
+| ---------------------------------- | ---------------------------------- |
+| `--bucket_no_upscale`              |                                    |
+| `--bucket_reso_steps`              |                                    |
+| `--caption_dropout_every_n_epochs` |                                    |
+| `--caption_dropout_rate`           |                                    |
+| `--caption_extension`              |                                    |
+| `--caption_tag_dropout_rate`       |                                    |
+| `--color_aug`                      |                                    |
+| `--dataset_repeats`                | `num_repeats`                      |
+| `--enable_bucket`                  |                                    |
+| `--face_crop_aug_range`            |                                    |
+| `--flip_aug`                       |                                    |
+| `--keep_tokens`                    |                                    |
+| `--min_bucket_reso`                |                                    |
+| `--random_crop`                    |                                    |
+| `--resolution`                     |                                    |
+| `--shuffle_caption`                |                                    |
+| `--train_batch_size`               | `batch_size`                       |
+
+## エラーの手引き
+
+現在、外部ライブラリを利用して設定ファイルの記述が正しいかどうかをチェックしているのですが、整備が行き届いておらずエラーメッセージがわかりづらいという問題があります。
+将来的にはこの問題の改善に取り組む予定です。
+
+次善策として、頻出のエラーとその対処法について載せておきます。
+正しいはずなのにエラーが出る場合、エラー内容がどうしても分からない場合は、バグかもしれないのでご連絡ください。
+
+* `voluptuous.error.MultipleInvalid: required key not provided @ ...`: 指定必須のオプションが指定されていないというエラーです。指定を忘れているか、オプション名を間違って記述している可能性が高いです。
+  * `...` の箇所にはエラーが発生した場所が載っています。例えば `voluptuous.error.MultipleInvalid: required key not provided @ data['datasets'][0]['subsets'][0]['image_dir']` のようなエラーが出たら、0 番目の `datasets` 中の 0 番目の `subsets` の設定に `image_dir` が存在しないということになります。
+* `voluptuous.error.MultipleInvalid: expected int for dictionary value @ ...`: 指定する値の形式が不正というエラーです。値の形式が間違っている可能性が高いです。`int` の部分は対象となるオプションによって変わります。この README に載っているオプションの「設定例」が役立つかもしれません。
+* `voluptuous.error.MultipleInvalid: extra keys not allowed @ ...`: 対応していないオプション名が存在している場合に発生するエラーです。オプション名を間違って記述しているか、誤って紛れ込んでいる可能性が高いです。
+
+

dreambooth_gui.py

@@ -0,0 +1,954 @@
+# v1: initial release
+# v2: add open and save folder icons
+# v3: Add new Utilities tab for Dreambooth folder preparation
+# v3.1: Adding captionning of images to utilities
+
+import gradio as gr
+import json
+import math
+import os
+import subprocess
+import pathlib
+import argparse
+from library.common_gui import (
+    get_folder_path,
+    remove_doublequote,
+    get_file_path,
+    get_any_file_path,
+    get_saveasfile_path,
+    color_aug_changed,
+    save_inference_file,
+    gradio_advanced_training,
+    run_cmd_advanced_training,
+    run_cmd_training,
+    gradio_training,
+    gradio_config,
+    gradio_source_model,
+    # set_legacy_8bitadam,
+    update_my_data,
+    check_if_model_exist,
+)
+from library.tensorboard_gui import (
+    gradio_tensorboard,
+    start_tensorboard,
+    stop_tensorboard,
+)
+from library.dreambooth_folder_creation_gui import (
+    gradio_dreambooth_folder_creation_tab,
+)
+from library.utilities import utilities_tab
+from library.sampler_gui import sample_gradio_config, run_cmd_sample
+from easygui import msgbox
+
+folder_symbol = '\U0001f4c2'  # 📂
+refresh_symbol = '\U0001f504'  # 🔄
+save_style_symbol = '\U0001f4be'  # 💾
+document_symbol = '\U0001F4C4'   # 📄
+
+
+def save_configuration(
+    save_as,
+    file_path,
+    pretrained_model_name_or_path,
+    v2,
+    v_parameterization,
+    logging_dir,
+    train_data_dir,
+    reg_data_dir,
+    output_dir,
+    max_resolution,
+    learning_rate,
+    lr_scheduler,
+    lr_warmup,
+    train_batch_size,
+    epoch,
+    save_every_n_epochs,
+    mixed_precision,
+    save_precision,
+    seed,
+    num_cpu_threads_per_process,
+    cache_latents,
+    caption_extension,
+    enable_bucket,
+    gradient_checkpointing,
+    full_fp16,
+    no_token_padding,
+    stop_text_encoder_training,
+    # use_8bit_adam,
+    xformers,
+    save_model_as,
+    shuffle_caption,
+    save_state,
+    resume,
+    prior_loss_weight,
+    color_aug,
+    flip_aug,
+    clip_skip,
+    vae,
+    output_name,
+    max_token_length,
+    max_train_epochs,
+    max_data_loader_n_workers,
+    mem_eff_attn,
+    gradient_accumulation_steps,
+    model_list,
+    keep_tokens,
+    persistent_data_loader_workers,
+    bucket_no_upscale,
+    random_crop,
+    bucket_reso_steps,
+    caption_dropout_every_n_epochs,
+    caption_dropout_rate,
+    optimizer,
+    optimizer_args,
+    noise_offset,
+    sample_every_n_steps,
+    sample_every_n_epochs,
+    sample_sampler,
+    sample_prompts,
+    additional_parameters,
+    vae_batch_size,
+    min_snr_gamma,weighted_captions,
+):
+    # Get list of function parameters and values
+    parameters = list(locals().items())
+
+    original_file_path = file_path
+
+    save_as_bool = True if save_as.get('label') == 'True' else False
+
+    if save_as_bool:
+        print('Save as...')
+        file_path = get_saveasfile_path(file_path)
+    else:
+        print('Save...')
+        if file_path == None or file_path == '':
+            file_path = get_saveasfile_path(file_path)
+
+    # print(file_path)
+
+    if file_path == None or file_path == '':
+        return original_file_path  # In case a file_path was provided and the user decide to cancel the open action
+
+    # Return the values of the variables as a dictionary
+    variables = {
+        name: value
+        for name, value in parameters  # locals().items()
+        if name
+        not in [
+            'file_path',
+            'save_as',
+        ]
+    }
+
+    # Extract the destination directory from the file path
+    destination_directory = os.path.dirname(file_path)
+
+    # Create the destination directory if it doesn't exist
+    if not os.path.exists(destination_directory):
+        os.makedirs(destination_directory)
+
+    # Save the data to the selected file
+    with open(file_path, 'w') as file:
+        json.dump(variables, file, indent=2)
+
+    return file_path
+
+
+def open_configuration(
+    ask_for_file,
+    file_path,
+    pretrained_model_name_or_path,
+    v2,
+    v_parameterization,
+    logging_dir,
+    train_data_dir,
+    reg_data_dir,
+    output_dir,
+    max_resolution,
+    learning_rate,
+    lr_scheduler,
+    lr_warmup,
+    train_batch_size,
+    epoch,
+    save_every_n_epochs,
+    mixed_precision,
+    save_precision,
+    seed,
+    num_cpu_threads_per_process,
+    cache_latents,
+    caption_extension,
+    enable_bucket,
+    gradient_checkpointing,
+    full_fp16,
+    no_token_padding,
+    stop_text_encoder_training,
+    # use_8bit_adam,
+    xformers,
+    save_model_as,
+    shuffle_caption,
+    save_state,
+    resume,
+    prior_loss_weight,
+    color_aug,
+    flip_aug,
+    clip_skip,
+    vae,
+    output_name,
+    max_token_length,
+    max_train_epochs,
+    max_data_loader_n_workers,
+    mem_eff_attn,
+    gradient_accumulation_steps,
+    model_list,
+    keep_tokens,
+    persistent_data_loader_workers,
+    bucket_no_upscale,
+    random_crop,
+    bucket_reso_steps,
+    caption_dropout_every_n_epochs,
+    caption_dropout_rate,
+    optimizer,
+    optimizer_args,
+    noise_offset,
+    sample_every_n_steps,
+    sample_every_n_epochs,
+    sample_sampler,
+    sample_prompts,
+    additional_parameters,
+    vae_batch_size,
+    min_snr_gamma,weighted_captions,
+):
+    # Get list of function parameters and values
+    parameters = list(locals().items())
+
+    ask_for_file = True if ask_for_file.get('label') == 'True' else False
+
+    original_file_path = file_path
+
+    if ask_for_file:
+        file_path = get_file_path(file_path)
+
+    if not file_path == '' and not file_path == None:
+        # load variables from JSON file
+        with open(file_path, 'r') as f:
+            my_data = json.load(f)
+            print('Loading config...')
+            # Update values to fix deprecated use_8bit_adam checkbox and set appropriate optimizer if it is set to True
+            my_data = update_my_data(my_data)
+    else:
+        file_path = original_file_path  # In case a file_path was provided and the user decide to cancel the open action
+        my_data = {}
+
+    values = [file_path]
+    for key, value in parameters:
+        # Set the value in the dictionary to the corresponding value in `my_data`, or the default value if not found
+        if not key in ['ask_for_file', 'file_path']:
+            values.append(my_data.get(key, value))
+    return tuple(values)
+
+
+def train_model(
+    pretrained_model_name_or_path,
+    v2,
+    v_parameterization,
+    logging_dir,
+    train_data_dir,
+    reg_data_dir,
+    output_dir,
+    max_resolution,
+    learning_rate,
+    lr_scheduler,
+    lr_warmup,
+    train_batch_size,
+    epoch,
+    save_every_n_epochs,
+    mixed_precision,
+    save_precision,
+    seed,
+    num_cpu_threads_per_process,
+    cache_latents,
+    caption_extension,
+    enable_bucket,
+    gradient_checkpointing,
+    full_fp16,
+    no_token_padding,
+    stop_text_encoder_training_pct,
+    # use_8bit_adam,
+    xformers,
+    save_model_as,
+    shuffle_caption,
+    save_state,
+    resume,
+    prior_loss_weight,
+    color_aug,
+    flip_aug,
+    clip_skip,
+    vae,
+    output_name,
+    max_token_length,
+    max_train_epochs,
+    max_data_loader_n_workers,
+    mem_eff_attn,
+    gradient_accumulation_steps,
+    model_list,  # Keep this. Yes, it is unused here but required given the common list used
+    keep_tokens,
+    persistent_data_loader_workers,
+    bucket_no_upscale,
+    random_crop,
+    bucket_reso_steps,
+    caption_dropout_every_n_epochs,
+    caption_dropout_rate,
+    optimizer,
+    optimizer_args,
+    noise_offset,
+    sample_every_n_steps,
+    sample_every_n_epochs,
+    sample_sampler,
+    sample_prompts,
+    additional_parameters,
+    vae_batch_size,
+    min_snr_gamma,weighted_captions,
+):
+    if pretrained_model_name_or_path == '':
+        msgbox('Source model information is missing')
+        return
+
+    if train_data_dir == '':
+        msgbox('Image folder path is missing')
+        return
+
+    if not os.path.exists(train_data_dir):
+        msgbox('Image folder does not exist')
+        return
+
+    if reg_data_dir != '':
+        if not os.path.exists(reg_data_dir):
+            msgbox('Regularisation folder does not exist')
+            return
+
+    if output_dir == '':
+        msgbox('Output folder path is missing')
+        return
+
+    if check_if_model_exist(output_name, output_dir, save_model_as):
+        return
+    
+    if optimizer == 'Adafactor' and lr_warmup != '0':
+        msgbox("Warning: lr_scheduler is set to 'Adafactor', so 'LR warmup (% of steps)' will be considered 0.", title="Warning")
+        lr_warmup = '0'
+
+    # Get a list of all subfolders in train_data_dir, excluding hidden folders
+    subfolders = [
+        f
+        for f in os.listdir(train_data_dir)
+        if os.path.isdir(os.path.join(train_data_dir, f))
+        and not f.startswith('.')
+    ]
+
+    # Check if subfolders are present. If not let the user know and return
+    if not subfolders:
+        print(
+            '\033[33mNo subfolders were found in',
+            train_data_dir,
+            " can't train\...033[0m",
+        )
+        return
+
+    total_steps = 0
+
+    # Loop through each subfolder and extract the number of repeats
+    for folder in subfolders:
+        # Extract the number of repeats from the folder name
+        try:
+            repeats = int(folder.split('_')[0])
+        except ValueError:
+            print(
+                '\033[33mSubfolder',
+                folder,
+                "does not have a proper repeat value, please correct the name or remove it... can't train...\033[0m",
+            )
+            continue
+
+        # Count the number of images in the folder
+        num_images = len(
+            [
+                f
+                for f, lower_f in (
+                    (file, file.lower())
+                    for file in os.listdir(
+                        os.path.join(train_data_dir, folder)
+                    )
+                )
+                if lower_f.endswith(('.jpg', '.jpeg', '.png', '.webp'))
+            ]
+        )
+
+        if num_images == 0:
+            print(f'{folder} folder contain no images, skipping...')
+        else:
+            # Calculate the total number of steps for this folder
+            steps = repeats * num_images
+            total_steps += steps
+
+            # Print the result
+            print('\033[33mFolder', folder, ':', steps, 'steps\033[0m')
+
+    if total_steps == 0:
+        print(
+            '\033[33mNo images were found in folder',
+            train_data_dir,
+            '... please rectify!\033[0m',
+        )
+        return
+
+    # Print the result
+    # print(f"{total_steps} total steps")
+
+    if reg_data_dir == '':
+        reg_factor = 1
+    else:
+        print(
+            '\033[94mRegularisation images are used... Will double the number of steps required...\033[0m'
+        )
+        reg_factor = 2
+
+    # calculate max_train_steps
+    max_train_steps = int(
+        math.ceil(
+            float(total_steps)
+            / int(train_batch_size)
+            * int(epoch)
+            * int(reg_factor)
+        )
+    )
+    print(f'max_train_steps = {max_train_steps}')
+
+    # calculate stop encoder training
+    if int(stop_text_encoder_training_pct) == -1:
+        stop_text_encoder_training = -1
+    elif stop_text_encoder_training_pct == None:
+        stop_text_encoder_training = 0
+    else:
+        stop_text_encoder_training = math.ceil(
+            float(max_train_steps) / 100 * int(stop_text_encoder_training_pct)
+        )
+    print(f'stop_text_encoder_training = {stop_text_encoder_training}')
+
+    lr_warmup_steps = round(float(int(lr_warmup) * int(max_train_steps) / 100))
+    print(f'lr_warmup_steps = {lr_warmup_steps}')
+
+    run_cmd = f'accelerate launch --num_cpu_threads_per_process={num_cpu_threads_per_process} "train_db.py"'
+    if v2:
+        run_cmd += ' --v2'
+    if v_parameterization:
+        run_cmd += ' --v_parameterization'
+    if enable_bucket:
+        run_cmd += ' --enable_bucket'
+    if no_token_padding:
+        run_cmd += ' --no_token_padding'
+    if weighted_captions:
+        run_cmd += ' --weighted_captions'
+    run_cmd += (
+        f' --pretrained_model_name_or_path="{pretrained_model_name_or_path}"'
+    )
+    run_cmd += f' --train_data_dir="{train_data_dir}"'
+    if len(reg_data_dir):
+        run_cmd += f' --reg_data_dir="{reg_data_dir}"'
+    run_cmd += f' --resolution={max_resolution}'
+    run_cmd += f' --output_dir="{output_dir}"'
+    run_cmd += f' --logging_dir="{logging_dir}"'
+    if not stop_text_encoder_training == 0:
+        run_cmd += (
+            f' --stop_text_encoder_training={stop_text_encoder_training}'
+        )
+    if not save_model_as == 'same as source model':
+        run_cmd += f' --save_model_as={save_model_as}'
+    # if not resume == '':
+    #     run_cmd += f' --resume={resume}'
+    if not float(prior_loss_weight) == 1.0:
+        run_cmd += f' --prior_loss_weight={prior_loss_weight}'
+    if not vae == '':
+        run_cmd += f' --vae="{vae}"'
+    if not output_name == '':
+        run_cmd += f' --output_name="{output_name}"'
+    if int(max_token_length) > 75:
+        run_cmd += f' --max_token_length={max_token_length}'
+    if not max_train_epochs == '':
+        run_cmd += f' --max_train_epochs="{max_train_epochs}"'
+    if not max_data_loader_n_workers == '':
+        run_cmd += (
+            f' --max_data_loader_n_workers="{max_data_loader_n_workers}"'
+        )
+    if int(gradient_accumulation_steps) > 1:
+        run_cmd += f' --gradient_accumulation_steps={int(gradient_accumulation_steps)}'
+
+    run_cmd += run_cmd_training(
+        learning_rate=learning_rate,
+        lr_scheduler=lr_scheduler,
+        lr_warmup_steps=lr_warmup_steps,
+        train_batch_size=train_batch_size,
+        max_train_steps=max_train_steps,
+        save_every_n_epochs=save_every_n_epochs,
+        mixed_precision=mixed_precision,
+        save_precision=save_precision,
+        seed=seed,
+        caption_extension=caption_extension,
+        cache_latents=cache_latents,
+        optimizer=optimizer,
+        optimizer_args=optimizer_args,
+    )
+
+    run_cmd += run_cmd_advanced_training(
+        max_train_epochs=max_train_epochs,
+        max_data_loader_n_workers=max_data_loader_n_workers,
+        max_token_length=max_token_length,
+        resume=resume,
+        save_state=save_state,
+        mem_eff_attn=mem_eff_attn,
+        clip_skip=clip_skip,
+        flip_aug=flip_aug,
+        color_aug=color_aug,
+        shuffle_caption=shuffle_caption,
+        gradient_checkpointing=gradient_checkpointing,
+        full_fp16=full_fp16,
+        xformers=xformers,
+        # use_8bit_adam=use_8bit_adam,
+        keep_tokens=keep_tokens,
+        persistent_data_loader_workers=persistent_data_loader_workers,
+        bucket_no_upscale=bucket_no_upscale,
+        random_crop=random_crop,
+        bucket_reso_steps=bucket_reso_steps,
+        caption_dropout_every_n_epochs=caption_dropout_every_n_epochs,
+        caption_dropout_rate=caption_dropout_rate,
+        noise_offset=noise_offset,
+        additional_parameters=additional_parameters,
+        vae_batch_size=vae_batch_size,
+        min_snr_gamma=min_snr_gamma,
+    )
+
+    run_cmd += run_cmd_sample(
+        sample_every_n_steps,
+        sample_every_n_epochs,
+        sample_sampler,
+        sample_prompts,
+        output_dir,
+    )
+
+    print(run_cmd)
+
+    # Run the command
+    if os.name == 'posix':
+        os.system(run_cmd)
+    else:
+        subprocess.run(run_cmd)
+
+    # check if output_dir/last is a folder... therefore it is a diffuser model
+    last_dir = pathlib.Path(f'{output_dir}/{output_name}')
+
+    if not last_dir.is_dir():
+        # Copy inference model for v2 if required
+        save_inference_file(output_dir, v2, v_parameterization, output_name)
+
+
+def dreambooth_tab(
+    train_data_dir=gr.Textbox(),
+    reg_data_dir=gr.Textbox(),
+    output_dir=gr.Textbox(),
+    logging_dir=gr.Textbox(),
+):
+    dummy_db_true = gr.Label(value=True, visible=False)
+    dummy_db_false = gr.Label(value=False, visible=False)
+    gr.Markdown('Train a custom model using kohya dreambooth python code...')
+    (
+        button_open_config,
+        button_save_config,
+        button_save_as_config,
+        config_file_name,
+        button_load_config,
+    ) = gradio_config()
+
+    (
+        pretrained_model_name_or_path,
+        v2,
+        v_parameterization,
+        save_model_as,
+        model_list,
+    ) = gradio_source_model()
+
+    with gr.Tab('Folders'):
+        with gr.Row():
+            train_data_dir = gr.Textbox(
+                label='Image folder',
+                placeholder='Folder where the training folders containing the images are located',
+            )
+            train_data_dir_input_folder = gr.Button(
+                '📂', elem_id='open_folder_small'
+            )
+            train_data_dir_input_folder.click(
+                get_folder_path,
+                outputs=train_data_dir,
+                show_progress=False,
+            )
+            reg_data_dir = gr.Textbox(
+                label='Regularisation folder',
+                placeholder='(Optional) Folder where where the regularization folders containing the images are located',
+            )
+            reg_data_dir_input_folder = gr.Button(
+                '📂', elem_id='open_folder_small'
+            )
+            reg_data_dir_input_folder.click(
+                get_folder_path,
+                outputs=reg_data_dir,
+                show_progress=False,
+            )
+        with gr.Row():
+            output_dir = gr.Textbox(
+                label='Model output folder',
+                placeholder='Folder to output trained model',
+            )
+            output_dir_input_folder = gr.Button(
+                '📂', elem_id='open_folder_small'
+            )
+            output_dir_input_folder.click(get_folder_path, outputs=output_dir)
+            logging_dir = gr.Textbox(
+                label='Logging folder',
+                placeholder='Optional: enable logging and output TensorBoard log to this folder',
+            )
+            logging_dir_input_folder = gr.Button(
+                '📂', elem_id='open_folder_small'
+            )
+            logging_dir_input_folder.click(
+                get_folder_path,
+                outputs=logging_dir,
+                show_progress=False,
+            )
+        with gr.Row():
+            output_name = gr.Textbox(
+                label='Model output name',
+                placeholder='Name of the model to output',
+                value='last',
+                interactive=True,
+            )
+        train_data_dir.change(
+            remove_doublequote,
+            inputs=[train_data_dir],
+            outputs=[train_data_dir],
+        )
+        reg_data_dir.change(
+            remove_doublequote,
+            inputs=[reg_data_dir],
+            outputs=[reg_data_dir],
+        )
+        output_dir.change(
+            remove_doublequote,
+            inputs=[output_dir],
+            outputs=[output_dir],
+        )
+        logging_dir.change(
+            remove_doublequote,
+            inputs=[logging_dir],
+            outputs=[logging_dir],
+        )
+    with gr.Tab('Training parameters'):
+        (
+            learning_rate,
+            lr_scheduler,
+            lr_warmup,
+            train_batch_size,
+            epoch,
+            save_every_n_epochs,
+            mixed_precision,
+            save_precision,
+            num_cpu_threads_per_process,
+            seed,
+            caption_extension,
+            cache_latents,
+            optimizer,
+            optimizer_args,
+        ) = gradio_training(
+            learning_rate_value='1e-5',
+            lr_scheduler_value='cosine',
+            lr_warmup_value='10',
+        )
+        with gr.Row():
+            max_resolution = gr.Textbox(
+                label='Max resolution',
+                value='512,512',
+                placeholder='512,512',
+            )
+            stop_text_encoder_training = gr.Slider(
+                minimum=-1,
+                maximum=100,
+                value=0,
+                step=1,
+                label='Stop text encoder training',
+            )
+            enable_bucket = gr.Checkbox(label='Enable buckets', value=True)
+        with gr.Accordion('Advanced Configuration', open=False):
+            with gr.Row():
+                no_token_padding = gr.Checkbox(
+                    label='No token padding', value=False
+                )
+                gradient_accumulation_steps = gr.Number(
+                    label='Gradient accumulate steps', value='1'
+                )
+                weighted_captions = gr.Checkbox(
+                    label='Weighted captions', value=False
+                )
+            with gr.Row():
+                prior_loss_weight = gr.Number(
+                    label='Prior loss weight', value=1.0
+                )
+                vae = gr.Textbox(
+                    label='VAE',
+                    placeholder='(Optiona) path to checkpoint of vae to replace for training',
+                )
+                vae_button = gr.Button('📂', elem_id='open_folder_small')
+                vae_button.click(
+                    get_any_file_path,
+                    outputs=vae,
+                    show_progress=False,
+                )
+            (
+                # use_8bit_adam,
+                xformers,
+                full_fp16,
+                gradient_checkpointing,
+                shuffle_caption,
+                color_aug,
+                flip_aug,
+                clip_skip,
+                mem_eff_attn,
+                save_state,
+                resume,
+                max_token_length,
+                max_train_epochs,
+                max_data_loader_n_workers,
+                keep_tokens,
+                persistent_data_loader_workers,
+                bucket_no_upscale,
+                random_crop,
+                bucket_reso_steps,
+                caption_dropout_every_n_epochs,
+                caption_dropout_rate,
+                noise_offset,
+                additional_parameters,
+                vae_batch_size,
+                min_snr_gamma,
+            ) = gradio_advanced_training()
+            color_aug.change(
+                color_aug_changed,
+                inputs=[color_aug],
+                outputs=[cache_latents],
+            )
+
+        (
+            sample_every_n_steps,
+            sample_every_n_epochs,
+            sample_sampler,
+            sample_prompts,
+        ) = sample_gradio_config()
+
+    with gr.Tab('Tools'):
+        gr.Markdown(
+            'This section provide Dreambooth tools to help setup your dataset...'
+        )
+        gradio_dreambooth_folder_creation_tab(
+            train_data_dir_input=train_data_dir,
+            reg_data_dir_input=reg_data_dir,
+            output_dir_input=output_dir,
+            logging_dir_input=logging_dir,
+        )
+
+    button_run = gr.Button('Train model', variant='primary')
+
+    # Setup gradio tensorboard buttons
+    button_start_tensorboard, button_stop_tensorboard = gradio_tensorboard()
+
+    button_start_tensorboard.click(
+        start_tensorboard,
+        inputs=logging_dir,
+        show_progress=False,
+    )
+
+    button_stop_tensorboard.click(
+        stop_tensorboard,
+        show_progress=False,
+    )
+
+    settings_list = [
+        pretrained_model_name_or_path,
+        v2,
+        v_parameterization,
+        logging_dir,
+        train_data_dir,
+        reg_data_dir,
+        output_dir,
+        max_resolution,
+        learning_rate,
+        lr_scheduler,
+        lr_warmup,
+        train_batch_size,
+        epoch,
+        save_every_n_epochs,
+        mixed_precision,
+        save_precision,
+        seed,
+        num_cpu_threads_per_process,
+        cache_latents,
+        caption_extension,
+        enable_bucket,
+        gradient_checkpointing,
+        full_fp16,
+        no_token_padding,
+        stop_text_encoder_training,
+        # use_8bit_adam,
+        xformers,
+        save_model_as,
+        shuffle_caption,
+        save_state,
+        resume,
+        prior_loss_weight,
+        color_aug,
+        flip_aug,
+        clip_skip,
+        vae,
+        output_name,
+        max_token_length,
+        max_train_epochs,
+        max_data_loader_n_workers,
+        mem_eff_attn,
+        gradient_accumulation_steps,
+        model_list,
+        keep_tokens,
+        persistent_data_loader_workers,
+        bucket_no_upscale,
+        random_crop,
+        bucket_reso_steps,
+        caption_dropout_every_n_epochs,
+        caption_dropout_rate,
+        optimizer,
+        optimizer_args,
+        noise_offset,
+        sample_every_n_steps,
+        sample_every_n_epochs,
+        sample_sampler,
+        sample_prompts,
+        additional_parameters,
+        vae_batch_size,
+        min_snr_gamma,
+        weighted_captions,
+    ]
+
+    button_open_config.click(
+        open_configuration,
+        inputs=[dummy_db_true, config_file_name] + settings_list,
+        outputs=[config_file_name] + settings_list,
+        show_progress=False,
+    )
+
+    button_load_config.click(
+        open_configuration,
+        inputs=[dummy_db_false, config_file_name] + settings_list,
+        outputs=[config_file_name] + settings_list,
+        show_progress=False,
+    )
+
+    button_save_config.click(
+        save_configuration,
+        inputs=[dummy_db_false, config_file_name] + settings_list,
+        outputs=[config_file_name],
+        show_progress=False,
+    )
+
+    button_save_as_config.click(
+        save_configuration,
+        inputs=[dummy_db_true, config_file_name] + settings_list,
+        outputs=[config_file_name],
+        show_progress=False,
+    )
+
+    button_run.click(
+        train_model,
+        inputs=settings_list,
+        show_progress=False,
+    )
+
+    return (
+        train_data_dir,
+        reg_data_dir,
+        output_dir,
+        logging_dir,
+    )
+
+
+def UI(**kwargs):
+    css = ''
+
+    if os.path.exists('./style.css'):
+        with open(os.path.join('./style.css'), 'r', encoding='utf8') as file:
+            print('Load CSS...')
+            css += file.read() + '\n'
+
+    interface = gr.Blocks(css=css)
+
+    with interface:
+        with gr.Tab('Dreambooth'):
+            (
+                train_data_dir_input,
+                reg_data_dir_input,
+                output_dir_input,
+                logging_dir_input,
+            ) = dreambooth_tab()
+        with gr.Tab('Utilities'):
+            utilities_tab(
+                train_data_dir_input=train_data_dir_input,
+                reg_data_dir_input=reg_data_dir_input,
+                output_dir_input=output_dir_input,
+                logging_dir_input=logging_dir_input,
+                enable_copy_info_button=True,
+            )
+
+    # Show the interface
+    launch_kwargs = {}
+    if not kwargs.get('username', None) == '':
+        launch_kwargs['auth'] = (
+            kwargs.get('username', None),
+            kwargs.get('password', None),
+        )
+    if kwargs.get('server_port', 0) > 0:
+        launch_kwargs['server_port'] = kwargs.get('server_port', 0)
+    if kwargs.get('inbrowser', False):
+        launch_kwargs['inbrowser'] = kwargs.get('inbrowser', False)
+    print(launch_kwargs)
+    interface.launch(**launch_kwargs)
+
+
+if __name__ == '__main__':
+    # torch.cuda.set_per_process_memory_fraction(0.48)
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        '--username', type=str, default='', help='Username for authentication'
+    )
+    parser.add_argument(
+        '--password', type=str, default='', help='Password for authentication'
+    )
+    parser.add_argument(
+        '--server_port',
+        type=int,
+        default=0,
+        help='Port to run the server listener on',
+    )
+    parser.add_argument(
+        '--inbrowser', action='store_true', help='Open in browser'
+    )
+
+    args = parser.parse_args()
+
+    UI(
+        username=args.username,
+        password=args.password,
+        inbrowser=args.inbrowser,
+        server_port=args.server_port,
+    )

fine_tune.py

@@ -0,0 +1,440 @@
+# training with captions
+# XXX dropped option: hypernetwork training
+
+import argparse
+import gc
+import math
+import os
+import toml
+from multiprocessing import Value
+
+from tqdm import tqdm
+import torch
+from accelerate.utils import set_seed
+import diffusers
+from diffusers import DDPMScheduler
+
+import library.train_util as train_util
+import library.config_util as config_util
+from library.config_util import (
+    ConfigSanitizer,
+    BlueprintGenerator,
+)
+import library.custom_train_functions as custom_train_functions
+from library.custom_train_functions import apply_snr_weight, get_weighted_text_embeddings
+
+
+def train(args):
+    train_util.verify_training_args(args)
+    train_util.prepare_dataset_args(args, True)
+
+    cache_latents = args.cache_latents
+
+    if args.seed is not None:
+        set_seed(args.seed)  # 乱数系列を初期化する
+
+    tokenizer = train_util.load_tokenizer(args)
+
+    blueprint_generator = BlueprintGenerator(ConfigSanitizer(False, True, True))
+    if args.dataset_config is not None:
+        print(f"Load dataset config from {args.dataset_config}")
+        user_config = config_util.load_user_config(args.dataset_config)
+        ignored = ["train_data_dir", "in_json"]
+        if any(getattr(args, attr) is not None for attr in ignored):
+            print(
+                "ignore following options because config file is found: {0} / 設定ファイルが利用されるため以下のオプションは無視されます: {0}".format(
+                    ", ".join(ignored)
+                )
+            )
+    else:
+        user_config = {
+            "datasets": [
+                {
+                    "subsets": [
+                        {
+                            "image_dir": args.train_data_dir,
+                            "metadata_file": args.in_json,
+                        }
+                    ]
+                }
+            ]
+        }
+
+    blueprint = blueprint_generator.generate(user_config, args, tokenizer=tokenizer)
+    train_dataset_group = config_util.generate_dataset_group_by_blueprint(blueprint.dataset_group)
+
+    current_epoch = Value("i", 0)
+    current_step = Value("i", 0)
+    ds_for_collater = train_dataset_group if args.max_data_loader_n_workers == 0 else None
+    collater = train_util.collater_class(current_epoch, current_step, ds_for_collater)
+
+    if args.debug_dataset:
+        train_util.debug_dataset(train_dataset_group)
+        return
+    if len(train_dataset_group) == 0:
+        print(
+            "No data found. Please verify the metadata file and train_data_dir option. / 画像がありません。メタデータおよびtrain_data_dirオプションを確認してください。"
+        )
+        return
+
+    if cache_latents:
+        assert (
+            train_dataset_group.is_latent_cacheable()
+        ), "when caching latents, either color_aug or random_crop cannot be used / latentをキャッシュするときはcolor_augとrandom_cropは使えません"
+
+    # acceleratorを準備する
+    print("prepare accelerator")
+    accelerator, unwrap_model = train_util.prepare_accelerator(args)
+
+    # mixed precisionに対応した型を用意しておき適宜castする
+    weight_dtype, save_dtype = train_util.prepare_dtype(args)
+
+    # モデルを読み込む
+    text_encoder, vae, unet, load_stable_diffusion_format = train_util.load_target_model(args, weight_dtype)
+
+    # verify load/save model formats
+    if load_stable_diffusion_format:
+        src_stable_diffusion_ckpt = args.pretrained_model_name_or_path
+        src_diffusers_model_path = None
+    else:
+        src_stable_diffusion_ckpt = None
+        src_diffusers_model_path = args.pretrained_model_name_or_path
+
+    if args.save_model_as is None:
+        save_stable_diffusion_format = load_stable_diffusion_format
+        use_safetensors = args.use_safetensors
+    else:
+        save_stable_diffusion_format = args.save_model_as.lower() == "ckpt" or args.save_model_as.lower() == "safetensors"
+        use_safetensors = args.use_safetensors or ("safetensors" in args.save_model_as.lower())
+
+    # Diffusers版のxformers使用フラグを設定する関数
+    def set_diffusers_xformers_flag(model, valid):
+        #   model.set_use_memory_efficient_attention_xformers(valid)            # 次のリリースでなくなりそう
+        # pipeが自動で再帰的にset_use_memory_efficient_attention_xformersを探すんだって(;´Д｀)
+        # U-Netだけ使う時にはどうすればいいのか……仕方ないからコピって使うか
+        # 0.10.2でなんか巻き戻って個別に指定するようになった(;^ω^)
+
+        # Recursively walk through all the children.
+        # Any children which exposes the set_use_memory_efficient_attention_xformers method
+        # gets the message
+        def fn_recursive_set_mem_eff(module: torch.nn.Module):
+            if hasattr(module, "set_use_memory_efficient_attention_xformers"):
+                module.set_use_memory_efficient_attention_xformers(valid)
+
+            for child in module.children():
+                fn_recursive_set_mem_eff(child)
+
+        fn_recursive_set_mem_eff(model)
+
+    # モデルに xformers とか memory efficient attention を組み込む
+    if args.diffusers_xformers:
+        print("Use xformers by Diffusers")
+        set_diffusers_xformers_flag(unet, True)
+    else:
+        # Windows版のxformersはfloatで学習できないのでxformersを使わない設定も可能にしておく必要がある
+        print("Disable Diffusers' xformers")
+        set_diffusers_xformers_flag(unet, False)
+        train_util.replace_unet_modules(unet, args.mem_eff_attn, args.xformers)
+
+    # 学習を準備する
+    if cache_latents:
+        vae.to(accelerator.device, dtype=weight_dtype)
+        vae.requires_grad_(False)
+        vae.eval()
+        with torch.no_grad():
+            train_dataset_group.cache_latents(vae, args.vae_batch_size, args.cache_latents_to_disk, accelerator.is_main_process)
+        vae.to("cpu")
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        gc.collect()
+
+        accelerator.wait_for_everyone()
+
+    # 学習を準備する：モデルを適切な状態にする
+    training_models = []
+    if args.gradient_checkpointing:
+        unet.enable_gradient_checkpointing()
+    training_models.append(unet)
+
+    if args.train_text_encoder:
+        print("enable text encoder training")
+        if args.gradient_checkpointing:
+            text_encoder.gradient_checkpointing_enable()
+        training_models.append(text_encoder)
+    else:
+        text_encoder.to(accelerator.device, dtype=weight_dtype)
+        text_encoder.requires_grad_(False)  # text encoderは学習しない
+        if args.gradient_checkpointing:
+            text_encoder.gradient_checkpointing_enable()
+            text_encoder.train()  # required for gradient_checkpointing
+        else:
+            text_encoder.eval()
+
+    if not cache_latents:
+        vae.requires_grad_(False)
+        vae.eval()
+        vae.to(accelerator.device, dtype=weight_dtype)
+
+    for m in training_models:
+        m.requires_grad_(True)
+    params = []
+    for m in training_models:
+        params.extend(m.parameters())
+    params_to_optimize = params
+
+    # 学習に必要なクラスを準備する
+    print("prepare optimizer, data loader etc.")
+    _, _, optimizer = train_util.get_optimizer(args, trainable_params=params_to_optimize)
+
+    # dataloaderを準備する
+    # DataLoaderのプロセス数：0はメインプロセスになる
+    n_workers = min(args.max_data_loader_n_workers, os.cpu_count() - 1)  # cpu_count-1 ただし最大で指定された数まで
+    train_dataloader = torch.utils.data.DataLoader(
+        train_dataset_group,
+        batch_size=1,
+        shuffle=True,
+        collate_fn=collater,
+        num_workers=n_workers,
+        persistent_workers=args.persistent_data_loader_workers,
+    )
+
+    # 学習ステップ数を計算する
+    if args.max_train_epochs is not None:
+        args.max_train_steps = args.max_train_epochs * math.ceil(
+            len(train_dataloader) / accelerator.num_processes / args.gradient_accumulation_steps
+        )
+        print(f"override steps. steps for {args.max_train_epochs} epochs is / 指定エポックまでのステップ数: {args.max_train_steps}")
+
+    # データセット側にも学習ステップを送信
+    train_dataset_group.set_max_train_steps(args.max_train_steps)
+
+    # lr schedulerを用意する
+    lr_scheduler = train_util.get_scheduler_fix(args, optimizer, accelerator.num_processes)
+
+    # 実験的機能：勾配も含めたfp16学習を行う　モデル全体をfp16にする
+    if args.full_fp16:
+        assert (
+            args.mixed_precision == "fp16"
+        ), "full_fp16 requires mixed precision='fp16' / full_fp16を使う場合はmixed_precision='fp16'を指定してください。"
+        print("enable full fp16 training.")
+        unet.to(weight_dtype)
+        text_encoder.to(weight_dtype)
+
+    # acceleratorがなんかよろしくやってくれるらしい
+    if args.train_text_encoder:
+        unet, text_encoder, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
+            unet, text_encoder, optimizer, train_dataloader, lr_scheduler
+        )
+    else:
+        unet, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(unet, optimizer, train_dataloader, lr_scheduler)
+
+    # 実験的機能：勾配も含めたfp16学習を行う　PyTorchにパッチを当ててfp16でのgrad scaleを有効にする
+    if args.full_fp16:
+        train_util.patch_accelerator_for_fp16_training(accelerator)
+
+    # resumeする
+    train_util.resume_from_local_or_hf_if_specified(accelerator, args)
+
+    # epoch数を計算する
+    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
+    num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)
+    if (args.save_n_epoch_ratio is not None) and (args.save_n_epoch_ratio > 0):
+        args.save_every_n_epochs = math.floor(num_train_epochs / args.save_n_epoch_ratio) or 1
+
+    # 学習する
+    total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
+    print("running training / 学習開始")
+    print(f"  num examples / サンプル数: {train_dataset_group.num_train_images}")
+    print(f"  num batches per epoch / 1epochのバッチ数: {len(train_dataloader)}")
+    print(f"  num epochs / epoch数: {num_train_epochs}")
+    print(f"  batch size per device / バッチサイズ: {args.train_batch_size}")
+    print(f"  total train batch size (with parallel & distributed & accumulation) / 総バッチサイズ（並列学習、勾配合計含む）: {total_batch_size}")
+    print(f"  gradient accumulation steps / 勾配を合計するステップ数 = {args.gradient_accumulation_steps}")
+    print(f"  total optimization steps / 学習ステップ数: {args.max_train_steps}")
+
+    progress_bar = tqdm(range(args.max_train_steps), smoothing=0, disable=not accelerator.is_local_main_process, desc="steps")
+    global_step = 0
+
+    noise_scheduler = DDPMScheduler(
+        beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000, clip_sample=False
+    )
+
+    if accelerator.is_main_process:
+        accelerator.init_trackers("finetuning")
+
+    for epoch in range(num_train_epochs):
+        print(f"epoch {epoch+1}/{num_train_epochs}")
+        current_epoch.value = epoch + 1
+
+        for m in training_models:
+            m.train()
+
+        loss_total = 0
+        for step, batch in enumerate(train_dataloader):
+            current_step.value = global_step
+            with accelerator.accumulate(training_models[0]):  # 複数モデルに対応していない模様だがとりあえずこうしておく
+                with torch.no_grad():
+                    if "latents" in batch and batch["latents"] is not None:
+                        latents = batch["latents"].to(accelerator.device) # .to(dtype=weight_dtype)
+                    else:
+                        # latentに変換
+                        latents = vae.encode(batch["images"].to(dtype=weight_dtype)).latent_dist.sample()
+                    latents = latents * 0.18215
+                b_size = latents.shape[0]
+
+                with torch.set_grad_enabled(args.train_text_encoder):
+                    # Get the text embedding for conditioning
+                    if args.weighted_captions:
+                      encoder_hidden_states = get_weighted_text_embeddings(tokenizer,
+                        text_encoder,
+                        batch["captions"],
+                        accelerator.device,
+                        args.max_token_length // 75 if args.max_token_length else 1,
+                        clip_skip=args.clip_skip,
+                        )
+                    else:
+                      input_ids = batch["input_ids"].to(accelerator.device)
+                      encoder_hidden_states = train_util.get_hidden_states(
+                          args, input_ids, tokenizer, text_encoder, None if not args.full_fp16 else weight_dtype
+                      )
+
+                # 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)
+                timesteps = timesteps.long()
+
+                # Add noise to the latents according to the noise magnitude at each timestep
+                # (this is the forward diffusion process)
+                noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)
+
+                # Predict the noise residual
+                with accelerator.autocast():
+                    noise_pred = unet(noisy_latents, timesteps, encoder_hidden_states).sample
+
+                if args.v_parameterization:
+                    # v-parameterization training
+                    target = noise_scheduler.get_velocity(latents, noise, timesteps)
+                else:
+                    target = noise
+
+                if args.min_snr_gamma:
+                    # do not mean over batch dimension for snr weight
+                    loss = torch.nn.functional.mse_loss(noise_pred.float(), target.float(), reduction="none")
+                    loss = loss.mean([1, 2, 3])
+                    loss = apply_snr_weight(loss, timesteps, noise_scheduler, args.min_snr_gamma)
+                    loss = loss.mean()  # mean over batch dimension
+                else:
+                    loss = torch.nn.functional.mse_loss(noise_pred.float(), target.float(), reduction="mean")
+
+                accelerator.backward(loss)
+                if accelerator.sync_gradients and args.max_grad_norm != 0.0:
+                    params_to_clip = []
+                    for m in training_models:
+                        params_to_clip.extend(m.parameters())
+                    accelerator.clip_grad_norm_(params_to_clip, args.max_grad_norm)
+
+                optimizer.step()
+                lr_scheduler.step()
+                optimizer.zero_grad(set_to_none=True)
+
+            # Checks if the accelerator has performed an optimization step behind the scenes
+            if accelerator.sync_gradients:
+                progress_bar.update(1)
+                global_step += 1
+
+                train_util.sample_images(
+                    accelerator, args, None, global_step, accelerator.device, vae, tokenizer, text_encoder, unet
+                )
+
+            current_loss = loss.detach().item()  # 平均なのでbatch sizeは関係ないはず
+            if args.logging_dir is not None:
+                logs = {"loss": current_loss, "lr": float(lr_scheduler.get_last_lr()[0])}
+                if args.optimizer_type.lower() == "DAdaptation".lower():  # tracking d*lr value
+                    logs["lr/d*lr"] = (
+                        lr_scheduler.optimizers[0].param_groups[0]["d"] * lr_scheduler.optimizers[0].param_groups[0]["lr"]
+                    )
+                accelerator.log(logs, step=global_step)
+
+            # TODO moving averageにする
+            loss_total += current_loss
+            avr_loss = loss_total / (step + 1)
+            logs = {"loss": avr_loss}  # , "lr": lr_scheduler.get_last_lr()[0]}
+            progress_bar.set_postfix(**logs)
+
+            if global_step >= args.max_train_steps:
+                break
+
+        if args.logging_dir is not None:
+            logs = {"loss/epoch": loss_total / len(train_dataloader)}
+            accelerator.log(logs, step=epoch + 1)
+
+        accelerator.wait_for_everyone()
+
+        if args.save_every_n_epochs is not None:
+            src_path = src_stable_diffusion_ckpt if save_stable_diffusion_format else src_diffusers_model_path
+            train_util.save_sd_model_on_epoch_end(
+                args,
+                accelerator,
+                src_path,
+                save_stable_diffusion_format,
+                use_safetensors,
+                save_dtype,
+                epoch,
+                num_train_epochs,
+                global_step,
+                unwrap_model(text_encoder),
+                unwrap_model(unet),
+                vae,
+            )
+
+        train_util.sample_images(accelerator, args, epoch + 1, global_step, accelerator.device, vae, tokenizer, text_encoder, unet)
+
+    is_main_process = accelerator.is_main_process
+    if is_main_process:
+        unet = unwrap_model(unet)
+        text_encoder = unwrap_model(text_encoder)
+
+    accelerator.end_training()
+
+    if args.save_state:
+        train_util.save_state_on_train_end(args, accelerator)
+
+    del accelerator  # この後メモリを使うのでこれは消す
+
+    if is_main_process:
+        src_path = src_stable_diffusion_ckpt if save_stable_diffusion_format else src_diffusers_model_path
+        train_util.save_sd_model_on_train_end(
+            args, src_path, save_stable_diffusion_format, use_safetensors, save_dtype, epoch, global_step, text_encoder, unet, vae
+        )
+        print("model saved.")
+
+
+def setup_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser()
+
+    train_util.add_sd_models_arguments(parser)
+    train_util.add_dataset_arguments(parser, False, True, True)
+    train_util.add_training_arguments(parser, False)
+    train_util.add_sd_saving_arguments(parser)
+    train_util.add_optimizer_arguments(parser)
+    config_util.add_config_arguments(parser)
+    custom_train_functions.add_custom_train_arguments(parser)
+
+    parser.add_argument("--diffusers_xformers", action="store_true", help="use xformers by diffusers / Diffusersでxformersを使用する")
+    parser.add_argument("--train_text_encoder", action="store_true", help="train text encoder / text encoderも学習する")
+
+    return parser
+
+
+if __name__ == "__main__":
+    parser = setup_parser()
+
+    args = parser.parse_args()
+    args = train_util.read_config_from_file(args, parser)
+
+    train(args)

fine_tune_README.md

@@ -0,0 +1,465 @@
+It is a fine tuning that corresponds to NovelAI's proposed learning method, automatic captioning, tagging, Windows + VRAM 12GB (for v1.4/1.5) environment, etc.
+
+## overview
+Fine tuning of U-Net of Stable Diffusion using Diffusers. It corresponds to the following improvements in NovelAI's article (For Aspect Ratio Bucketing, I referred to NovelAI's code, but the final code is all original).
+
+* Use the output of the penultimate layer instead of the last layer of CLIP (Text Encoder).
+* Learning at non-square resolutions (Aspect Ratio Bucketing).
+* Extend token length from 75 to 225.
+* Captioning with BLIP (automatic creation of captions), automatic tagging with DeepDanbooru or WD14Tagger.
+* Also supports Hypernetwork learning.
+* Supports Stable Diffusion v2.0 (base and 768/v).
+* By acquiring the output of VAE in advance and saving it to disk, we aim to save memory and speed up learning.
+
+Text Encoder is not trained by default. For fine tuning of the whole model, it seems common to learn only U-Net (NovelAI seems to be the same). Text Encoder can also be learned as an option.
+
+## Additional features
+### Change CLIP output
+CLIP (Text Encoder) converts the text into features in order to reflect the prompt in the image. Stable diffusion uses the output of the last layer of CLIP, but you can change it to use the output of the penultimate layer. According to NovelAI, this will reflect prompts more accurately.
+It is also possible to use the output of the last layer as is.
+*Stable Diffusion 2.0 uses the penultimate layer by default. Do not specify the clip_skip option.
+
+### Training in non-square resolutions
+Stable Diffusion is trained at 512\*512, but also at resolutions such as 256\*1024 and 384\*640. It is expected that this will reduce the cropped portion and learn the relationship between prompts and images more correctly.
+The learning resolution is adjusted vertically and horizontally in units of 64 pixels within a range that does not exceed the resolution area (= memory usage) given as a parameter.
+
+In machine learning, it is common to unify all input sizes, but there are no particular restrictions, and in fact it is okay as long as they are unified within the same batch. NovelAI's bucketing seems to refer to classifying training data in advance for each learning resolution according to the aspect ratio. And by creating a batch with the images in each bucket, the image size of the batch is unified.
+
+### Extending token length from 75 to 225
+Stable diffusion has a maximum of 75 tokens (77 tokens including the start and end), but we will extend it to 225 tokens.
+However, the maximum length that CLIP accepts is 75 tokens, so in the case of 225 tokens, we simply divide it into thirds, call CLIP, and then concatenate the results.
+
+*I'm not sure if this is the preferred implementation. It seems to be working for now. Especially in 2.0, there is no implementation that can be used as a reference, so I have implemented it independently.
+
+*Automatic1111's Web UI seems to divide the text with commas in mind, but in my case, it's a simple division.
+
+## Environmental arrangement
+
+See the [README](./README-en.md) in this repository.
+
+## Preparing teacher data
+
+Prepare the image data you want to learn and put it in any folder. No prior preparation such as resizing is required.
+However, for images that are smaller than the training resolution, it is recommended to enlarge them while maintaining the quality using super-resolution.
+
+It also supports multiple teacher data folders. Preprocessing will be executed for each folder.
+
+For example, store an image like this:
+
+![Teacher data folder screenshot](https://user-images.githubusercontent.com/52813779/208907739-8e89d5fa-6ca8-4b60-8927-f484d2a9ae04.png)
+
+## Automatic captioning
+Skip if you just want to learn tags without captions.
+
+Also, when preparing captions manually, prepare them in the same directory as the teacher data image, with the same file name, extension .caption, etc. Each file should be a text file with only one line.
+
+### Captioning with BLIP
+
+The latest version no longer requires BLIP downloads, weight downloads, and additional virtual environments. Works as-is.
+
+Run make_captions.py in the finetune folder.
+
+```
+python finetune\make_captions.py --batch_size <batch size> <teacher data folder>
+```
+
+If the batch size is 8 and the training data is placed in the parent folder train_data, it will be as follows.
+
+```
+python finetune\make_captions.py --batch_size 8 ..\train_data
+```
+
+A caption file is created in the same directory as the teacher data image with the same file name and extension .caption.
+
+Increase or decrease batch_size according to the VRAM capacity of the GPU. Bigger is faster (I think 12GB of VRAM can be a little more).
+You can specify the maximum length of the caption with the max_length option. Default is 75. It may be longer if the model is trained with a token length of 225.
+You can change the caption extension with the caption_extension option. Default is .caption (.txt conflicts with DeepDanbooru described later).
+
+If there are multiple teacher data folders, execute for each folder.
+
+Note that the inference is random, so the results will change each time you run it. If you want to fix it, specify a random number seed like "--seed 42" with the --seed option.
+
+For other options, please refer to the help with --help (there seems to be no documentation for the meaning of the parameters, so you have to look at the source).
+
+A caption file is generated with the extension .caption by default.
+
+![Folder where caption is generated](https://user-images.githubusercontent.com/52813779/208908845-48a9d36c-f6ee-4dae-af71-9ab462d1459e.png)
+
+For example, with captions like:
+
+![captions and images](https://user-images.githubusercontent.com/52813779/208908947-af936957-5d73-4339-b6c8-945a52857373.png)
+
+## Tagged by DeepDanbooru
+If you do not want to tag the danbooru tag itself, please proceed to "Preprocessing of caption and tag information".
+
+Tagging is done with DeepDanbooru or WD14Tagger. WD14Tagger seems to be more accurate. If you want to tag with WD14Tagger, skip to the next chapter.
+
+### Environmental arrangement
+Clone DeepDanbooru https://github.com/KichangKim/DeepDanbooru into your working folder, or download the zip and extract it. I unzipped it.
+Also, download deepdanbooru-v3-20211112-sgd-e28.zip from Assets of "DeepDanbooru Pretrained Model v3-20211112-sgd-e28" on the DeepDanbooru Releases page https://github.com/KichangKim/DeepDanbooru/releases and extract it to the DeepDanbooru folder.
+
+Download from below. Click to open Assets and download from there.
+
+![DeepDanbooru download page](https://user-images.githubusercontent.com/52813779/208909417-10e597df-7085-41ee-bd06-3e856a1339df.png)
+
+Make a directory structure like this
+
+![DeepDanbooru directory structure](https://user-images.githubusercontent.com/52813779/208909486-38935d8b-8dc6-43f1-84d3-fef99bc471aa.png)
+
+Install the necessary libraries for the Diffusers environment. Go to the DeepDanbooru folder and install it (I think it's actually just adding tensorflow-io).
+
+```
+pip install -r requirements.txt
+```
+
+Next, install DeepDanbooru itself.
+
+```
+pip install .
+```
+
+This completes the preparation of the environment for tagging.
+
+### Implementing tagging
+Go to DeepDanbooru's folder and run deepdanbooru to tag.
+
+```
+deepdanbooru evaluate <teacher data folder> --project-path deepdanbooru-v3-20211112-sgd-e28 --allow-folder --save-txt
+```
+
+If you put the training data in the parent folder train_data, it will be as follows.
+
+```
+deepdanbooru evaluate ../train_data --project-path deepdanbooru-v3-20211112-sgd-e28 --allow-folder --save-txt
+```
+
+A tag file is created in the same directory as the teacher data image with the same file name and extension .txt. It is slow because it is processed one by one.
+
+If there are multiple teacher data folders, execute for each folder.
+
+It is generated as follows.
+
+![DeepDanbooru generated files](https://user-images.githubusercontent.com/52813779/208909855-d21b9c98-f2d3-4283-8238-5b0e5aad6691.png)
+
+A tag is attached like this (great amount of information...).
+
+![Deep Danbooru tag and image](https://user-images.githubusercontent.com/52813779/208909908-a7920174-266e-48d5-aaef-940aba709519.png)
+
+## Tagging with WD14Tagger
+This procedure uses WD14Tagger instead of DeepDanbooru.
+
+Use the tagger used in Mr. Automatic1111's WebUI. I referred to the information on this github page (https://github.com/toriato/stable-diffusion-webui-wd14-tagger#mrsmilingwolfs-model-aka-waifu-diffusion-14-tagger).
+
+The modules required for the initial environment maintenance have already been installed. Weights are automatically downloaded from Hugging Face.
+
+### Implementing tagging
+Run the script to do the tagging.
+```
+python tag_images_by_wd14_tagger.py --batch_size <batch size> <teacher data folder>
+```
+
+If you put the training data in the parent folder train_data, it will be as follows.
+```
+python tag_images_by_wd14_tagger.py --batch_size 4 ..\train_data
+```
+
+The model file will be automatically downloaded to the wd14_tagger_model folder on first launch (folder can be changed in options). It will be as follows.
+
+![downloaded file](https://user-images.githubusercontent.com/52813779/208910447-f7eb0582-90d6-49d3-a666-2b508c7d1842.png)
+
+A tag file is created in the same directory as the teacher data image with the same file name and extension .txt.
+
+![generated tag file](https://user-images.githubusercontent.com/52813779/208910534-ea514373-1185-4b7d-9ae3-61eb50bc294e.png)
+
+![tags and images](https://user-images.githubusercontent.com/52813779/208910599-29070c15-7639-474f-b3e4-06bd5a3df29e.png)
+
+With the thresh option, you can specify the number of confidences of the determined tag to attach the tag. The default is 0.35, same as the WD14Tagger sample. Lower values give more tags, but less accuracy.
+Increase or decrease batch_size according to the VRAM capacity of the GPU. Bigger is faster (I think 12GB of VRAM can be a little more). You can change the tag file extension with the caption_extension option. Default is .txt.
+You can specify the folder where the model is saved with the model_dir option.
+Also, if you specify the force_download option, the model will be re-downloaded even if there is a save destination folder.
+
+If there are multiple teacher data folders, execute for each folder.
+
+## Preprocessing caption and tag information
+
+Combine captions and tags into a single file as metadata for easy processing from scripts.
+
+### Caption preprocessing
+
+To put captions into the metadata, run the following in your working folder (if you don't use captions for learning, you don't need to run this) (it's actually a single line, and so on).
+
+```
+python merge_captions_to_metadata.py <teacher data folder>
+--in_json <metadata file name to read>
+     <metadata file name>
+```
+
+The metadata file name is an arbitrary name.
+If the training data is train_data, there is no metadata file to read, and the metadata file is meta_cap.json, it will be as follows.
+
+```
+python merge_captions_to_metadata.py train_data meta_cap.json
+```
+
+You can specify the caption extension with the caption_extension option.
+
+If there are multiple teacher data folders, please specify the full_path argument (metadata will have full path information). Then run it for each folder.
+
+```
+python merge_captions_to_metadata.py --full_path
+     train_data1 meta_cap1.json
+python merge_captions_to_metadata.py --full_path --in_json meta_cap1.json
+     train_data2 meta_cap2.json
+```
+
+If in_json is omitted, if there is a write destination metadata file, it will be read from there and overwritten there.
+
+__*It is safe to rewrite the in_json option and the write destination each time and write to a separate metadata file. __
+
+### Tag preprocessing
+
+Similarly, tags are also collected in metadata (no need to do this if tags are not used for learning).
+```
+python merge_dd_tags_to_metadata.py <teacher data folder>
+     --in_json <metadata file name to load>
+     <metadata file name to write>
+```
+
+With the same directory structure as above, when reading meta_cap.json and writing to meta_cap_dd.json, it will be as follows.
+```
+python merge_dd_tags_to_metadata.py train_data --in_json meta_cap.json meta_cap_dd.json
+```
+
+If you have multiple teacher data folders, please specify the full_path argument. Then run it for each folder.
+
+```
+python merge_dd_tags_to_metadata.py --full_path --in_json meta_cap2.json
+     train_data1 meta_cap_dd1.json
+python merge_dd_tags_to_metadata.py --full_path --in_json meta_cap_dd1.json
+     train_data2 meta_cap_dd2.json
+```
+
+If in_json is omitted, if there is a write destination metadata file, it will be read from there and overwritten there.
+
+__*It is safe to rewrite the in_json option and the write destination each time and write to a separate metadata file. __
+
+### Cleaning captions and tags
+Up to this point, captions and DeepDanbooru tags have been put together in the metadata file. However, captions with automatic captioning are subtle due to spelling variations (*), and tags include underscores and ratings (in the case of DeepDanbooru), so the editor's replacement function etc. You should use it to clean your captions and tags.
+
+*For example, when learning a girl in an anime picture, there are variations in captions such as girl/girls/woman/women. Also, it may be more appropriate to simply use "girl" for things like "anime girl".
+
+A script for cleaning is provided, so please edit the contents of the script according to the situation and use it.
+
+(It is no longer necessary to specify the teacher data folder. All data in the metadata will be cleaned.)
+
+```
+python clean_captions_and_tags.py <metadata file name to read> <metadata file name to write>
+```
+
+Please note that --in_json is not included. For example:
+
+```
+python clean_captions_and_tags.py meta_cap_dd.json meta_clean.json
+```
+
+Preprocessing of captions and tags is now complete.
+
+## Get latents in advance
+
+In order to speed up the learning, we acquire the latent representation of the image in advance and save it to disk. At the same time, bucketing (classifying the training data according to the aspect ratio) is performed.
+
+In your working folder, type:
+```
+python prepare_buckets_latents.py <teacher data folder>
+     <metadata file name to read> <metadata file name to write>
+     <model name or checkpoint for fine tuning>
+     --batch_size <batch size>
+     --max_resolution <resolution width, height>
+     --mixed_precision <precision>
+```
+
+If the model is model.ckpt, batch size 4, training resolution is 512\*512, precision is no (float32), read metadata from meta_clean.json and write to meta_lat.json:
+
+```
+python prepare_buckets_latents.py
+     train_data meta_clean.json meta_lat.json model.ckpt
+     --batch_size 4 --max_resolution 512,512 --mixed_precision no
+```
+
+Latents are saved in numpy npz format in the teacher data folder.
+
+Specify the --v2 option when loading a Stable Diffusion 2.0 model (--v_parameterization is not required).
+
+You can specify the minimum resolution size with the --min_bucket_reso option and the maximum size with the --max_bucket_reso option. The defaults are 256 and 1024 respectively. For example, specifying a minimum size of 384 will not use resolutions such as 256\*1024 or 320\*768.
+If you increase the resolution to something like 768\*768, you should specify something like 1280 for the maximum size.
+
+If you specify the --flip_aug option, it will perform horizontal flip augmentation (data augmentation). You can artificially double the amount of data, but if you specify it when the data is not left-right symmetrical (for example, character appearance, hairstyle, etc.), learning will not go well.
+(This is a simple implementation that acquires the latents for the flipped image and saves the \*\_flip.npz file. No options are required for fline_tune.py. If there is a file with \_flip, Randomly load a file without
+
+The batch size may be increased a little more even with 12GB of VRAM.
+The resolution is a number divisible by 64, and is specified by "width, height". The resolution is directly linked to the memory size during fine tuning. 512,512 seems to be the limit with VRAM 12GB (*). 16GB may be raised to 512,704 or 512,768. Even with 256, 256, etc., it seems to be difficult with 8GB of VRAM (because parameters and optimizers require a certain amount of memory regardless of resolution).
+
+*There was also a report that learning batch size 1 worked with 12GB VRAM and 640,640.
+
+The result of bucketing is displayed as follows.
+
+![bucketing result](https://user-images.githubusercontent.com/52813779/208911419-71c00fbb-2ce6-49d5-89b5-b78d7715e441.png)
+
+If you have multiple teacher data folders, please specify the full_path argument. Then run it for each folder.
+```
+python prepare_buckets_latents.py --full_path
+     train_data1 meta_clean.json meta_lat1.json model.ckpt
+     --batch_size 4 --max_resolution 512,512 --mixed_precision no
+
+python prepare_buckets_latents.py --full_path
+     train_data2 meta_lat1.json meta_lat2.json model.ckpt
+     --batch_size 4 --max_resolution 512,512 --mixed_precision no
+
+```
+It is possible to make the read source and write destination the same, but separate is safer.
+
+__*It is safe to rewrite the argument each time and write it to a separate metadata file. __
+
+
+## Run training
+For example: Below are the settings for saving memory.
+```
+accelerate launch --num_cpu_threads_per_process 8 fine_tune.py
+     --pretrained_model_name_or_path=model.ckpt
+     --in_json meta_lat.json
+     --train_data_dir=train_data
+     --output_dir=fine_tuned
+     --shuffle_caption
+     --train_batch_size=1 --learning_rate=5e-6 --max_train_steps=10000
+     --use_8bit_adam --xformers --gradient_checkpointing
+     --mixed_precision=bf16
+     --save_every_n_epochs=4
+```
+
+It seems to be good to specify the number of CPU cores for num_cpu_threads_per_process of accelerate.
+
+Specify the model to be trained in pretrained_model_name_or_path (Stable Diffusion checkpoint or Diffusers model). Stable Diffusion checkpoint supports .ckpt and .safetensors (automatically determined by extension).
+
+Specifies the metadata file when caching latent to in_json.
+
+Specify the training data folder for train_data_dir and the output destination folder for the trained model for output_dir.
+
+If shuffle_caption is specified, captions and tags are shuffled and learned in units separated by commas (this is the method used in Waifu Diffusion v1.3).
+(You can keep some of the leading tokens fixed without shuffling. See keep_tokens for other options.)
+
+Specify the batch size in train_batch_size. Specify 1 or 2 for VRAM 12GB. The number that can be specified also changes depending on the resolution.
+The actual amount of data used for training is "batch size x number of steps". When increasing the batch size, the number of steps can be decreased accordingly.
+
+Specify the learning rate in learning_rate. For example Waifu Diffusion v1.3 seems to be 5e-6.
+Specify the number of steps in max_train_steps.
+
+Specify use_8bit_adam to use the 8-bit Adam Optimizer. It saves memory and speeds up, but accuracy may decrease.
+
+Specifying xformers replaces CrossAttention to save memory and speed up.
+* As of 11/9, xformers will cause an error in float32 learning, so please use bf16/fp16 or use memory-saving CrossAttention with mem_eff_attn instead (speed is inferior to xformers).
+
+Enable intermediate saving of gradients in gradient_checkpointing. It's slower, but uses less memory.
+
+Specifies whether to use mixed precision with mixed_precision. Specifying "fp16" or "bf16" saves memory, but accuracy is inferior.
+"fp16" and "bf16" use almost the same amount of memory, and it is said that bf16 has better learning results (I didn't feel much difference in the range I tried).
+If "no" is specified, it will not be used (it will be float32).
+
+* It seems that an error will occur when reading checkpoints learned with bf16 with Mr. AUTOMATIC1111's Web UI. This seems to be because the data type bfloat16 causes an error in the Web UI model safety checker. Save in fp16 or float32 format with the save_precision option. Or it seems to be good to store it in safetytensors format.
+
+Specifying save_every_n_epochs will save the model being trained every time that many epochs have passed.
+
+### Supports Stable Diffusion 2.0
+Specify the --v2 option when using Hugging Face's stable-diffusion-2-base, and specify both --v2 and --v_parameterization options when using stable-diffusion-2 or 768-v-ema.ckpt please.
+
+### Increase accuracy and speed when memory is available
+First, removing gradient_checkpointing will speed it up. However, the batch size that can be set is reduced, so please set while looking at the balance between accuracy and speed.
+
+Increasing the batch size increases speed and accuracy. Increase the speed while checking the speed per data within the range where the memory is sufficient (the speed may actually decrease when the memory is at the limit).
+
+### Change CLIP output used
+Specifying 2 for the clip_skip option uses the output of the next-to-last layer. If 1 or option is omitted, the last layer is used.
+The learned model should be able to be inferred by Automatic1111's web UI.
+
+*SD2.0 uses the second layer from the back by default, so please do not specify it when learning SD2.0.
+
+If the model being trained was originally trained to use the second layer, 2 is a good value.
+
+If you were using the last layer instead, the entire model would have been trained on that assumption. Therefore, if you train again using the second layer, you may need a certain number of teacher data and longer learning to obtain the desired learning result.
+
+### Extending Token Length
+You can learn by extending the token length by specifying 150 or 225 for max_token_length.
+The learned model should be able to be inferred by Automatic1111's web UI.
+
+As with clip_skip, learning with a length different from the learning state of the model may require a certain amount of teacher data and a longer learning time.
+
+### Save learning log
+Specify the log save destination folder in the logging_dir option. Logs in TensorBoard format are saved.
+
+For example, if you specify --logging_dir=logs, a logs folder will be created in your working folder, and logs will be saved in the date/time folder.
+Also, if you specify the --log_prefix option, the specified string will be added before the date and time. Use "--logging_dir=logs --log_prefix=fine_tune_style1" for identification.
+
+To check the log with TensorBoard, open another command prompt and enter the following in the working folder (I think tensorboard is installed when Diffusers is installed, but if it is not installed, pip install Please put it in tensorboard).
+```
+tensorboard --logdir=logs
+```
+
+### Learning Hypernetworks
+It will be explained in another article.
+
+### Learning with fp16 gradient (experimental feature)
+The full_fp16 option will change the gradient from normal float32 to float16 (fp16) and learn (it seems to be full fp16 learning instead of mixed precision). As a result, it seems that the SD1.x 512*512 size can be learned with a VRAM usage of less than 8GB, and the SD2.x 512*512 size can be learned with a VRAM usage of less than 12GB.
+
+Specify fp16 in advance in accelerate config and optionally set mixed_precision="fp16" (does not work with bf16).
+
+To minimize memory usage, use the xformers, use_8bit_adam, gradient_checkpointing options and set train_batch_size to 1.
+(If you can afford it, increasing the train_batch_size step by step should improve the accuracy a little.)
+
+It is realized by patching the PyTorch source (confirmed with PyTorch 1.12.1 and 1.13.0). The accuracy will drop considerably, and the probability of learning failure on the way will also increase. The setting of the learning rate and the number of steps seems to be severe. Please be aware of them and use them at your own risk.
+
+### Other Options
+
+#### keep_tokens
+If a number is specified, the specified number of tokens (comma-separated strings) from the beginning of the caption are fixed without being shuffled.
+
+If there are both captions and tags, the prompts during learning will be concatenated like "caption, tag 1, tag 2...", so if you set "--keep_tokens=1", the caption will always be at the beginning during learning. will come.
+
+#### dataset_repeats
+If the number of data sets is extremely small, the epoch will end soon (it will take some time at the epoch break), so please specify a numerical value and multiply the data by some to make the epoch longer.
+
+#### train_text_encoder
+Text Encoder is also a learning target. Slightly increased memory usage.
+
+In normal fine tuning, the Text Encoder is not targeted for training (probably because U-Net is trained to follow the output of the Text Encoder), but if the number of training data is small, the Text Encoder is trained like DreamBooth. also seems to be valid.
+
+#### save_precision
+The data format when saving checkpoints can be specified from float, fp16, and bf16 (if not specified, it is the same as the data format during learning). It saves disk space, but the model produces different results. Also, if you specify float or fp16, you should be able to read it on Mr. 1111's Web UI.
+
+*For VAE, the data format of the original checkpoint will remain, so the model size may not be reduced to a little over 2GB even with fp16.
+
+#### save_model_as
+Specify the save format of the model. Specify one of ckpt, safetensors, diffusers, diffusers_safetensors.
+
+When reading Stable Diffusion format (ckpt or safetensors) and saving in Diffusers format, missing information is supplemented by dropping v1.5 or v2.1 information from Hugging Face.
+
+#### use_safetensors
+This option saves checkpoints in safetyensors format. The save format will be the default (same format as loaded).
+
+#### save_state and resume
+The save_state option saves the learning state of the optimizer, etc. in addition to the checkpoint in the folder when saving midway and at the final save. This avoids a decrease in accuracy when learning is resumed after being interrupted (since the optimizer optimizes while having a state, if the state is reset, the optimization must be performed again from the initial state. not). Note that the number of steps is not saved due to Accelerate specifications.
+
+When starting the script, you can resume by specifying the folder where the state is saved with the resume option.
+
+Please note that the learning state will be about 5 GB per save, so please be careful of the disk capacity.
+
+#### gradient_accumulation_steps
+Updates the gradient in batches for the specified number of steps. Has a similar effect to increasing the batch size, but consumes slightly more memory.
+
+*The Accelerate specification does not support multiple learning models, so if you set Text Encoder as the learning target and specify a value of 2 or more for this option, an error may occur.
+
+#### lr_scheduler / lr_warmup_steps
+You can choose the learning rate scheduler from linear, cosine, cosine_with_restarts, polynomial, constant, constant_with_warmup with the lr_scheduler option. Default is constant.
+
+With lr_warmup_steps, you can specify the number of steps to warm up the scheduler (gradually changing the learning rate). Please do your own research for details.
+
+#### diffusers_xformers
+Uses Diffusers' xformers feature rather than the script's own xformers replacement feature. Hypernetwork learning is no longer possible.

fine_tune_README_ja.md

@@ -0,0 +1,140 @@
+NovelAIの提案した学習手法、自動キャプションニング、タグ付け、Windows＋VRAM 12GB（SD v1.xの場合）環境等に対応したfine tuningです。ここでfine tuningとは、モデルを画像とキャプションで学習することを指します（LoRAやTextual Inversion、Hypernetworksは含みません）
+
+[学習についての共通ドキュメント](./train_README-ja.md) もあわせてご覧ください。
+
+# 概要
+
+Diffusersを用いてStable DiffusionのU-Netのfine tuningを行います。NovelAIの記事にある以下の改善に対応しています（Aspect Ratio BucketingについてはNovelAIのコードを参考にしましたが、最終的なコードはすべてオリジナルです）。
+
+* CLIP（Text Encoder）の最後の層ではなく最後から二番目の層の出力を用いる。
+* 正方形以外の解像度での学習（Aspect Ratio Bucketing） 。
+* トークン長を75から225に拡張する。
+* BLIPによるキャプショニング（キャプションの自動作成）、DeepDanbooruまたはWD14Taggerによる自動タグ付けを行う。
+* Hypernetworkの学習にも対応する。
+* Stable Diffusion v2.0（baseおよび768/v）に対応。
+* VAEの出力をあらかじめ取得しディスクに保存しておくことで、学習の省メモリ化、高速化を図る。
+
+デフォルトではText Encoderの学習は行いません。モデル全体のfine tuningではU-Netだけを学習するのが一般的なようです（NovelAIもそのようです）。オプション指定でText Encoderも学習対象とできます。
+
+# 追加機能について
+
+## CLIPの出力の変更
+
+プロンプトを画像に反映するため、テキストの特徴量への変換を行うのがCLIP（Text Encoder）です。Stable DiffusionではCLIPの最後の層の出力を用いていますが、それを最後から二番目の層の出力を用いるよう変更できます。NovelAIによると、これによりより正確にプロンプトが反映されるようになるとのことです。
+元のまま、最後の層の出力を用いることも可能です。
+
+※Stable Diffusion 2.0では最後から二番目の層をデフォルトで使います。clip_skipオプションを指定しないでください。
+
+## 正方形以外の解像度での学習
+
+Stable Diffusionは512\*512で学習されていますが、それに加えて256\*1024や384\*640といった解像度でも学習します。これによりトリミングされる部分が減り、より正しくプロンプトと画像の関係が学習されることが期待されます。
+学習解像度はパラメータとして与えられた解像度の面積（＝メモリ使用量）を超えない範囲で、64ピクセル単位で縦横に調整、作成されます。
+
+機械学習では入力サイズをすべて統一するのが一般的ですが、特に制約があるわけではなく、実際は同一のバッチ内で統一されていれば大丈夫です。NovelAIの言うbucketingは、あらかじめ教師データを、アスペクト比に応じた学習解像度ごとに分類しておくことを指しているようです。そしてバッチを各bucket内の画像で作成することで、バッチの画像サイズを統一します。
+
+## トークン長の75から225への拡張
+
+Stable Diffusionでは最大75トークン（開始・終了を含むと77トークン）ですが、それを225トークンまで拡張します。
+ただしCLIPが受け付ける最大長は75トークンですので、225トークンの場合、単純に三分割してCLIPを呼び出してから結果を連結しています。
+
+※これが望ましい実装なのかどうかはいまひとつわかりません。とりあえず動いてはいるようです。特に2.0では何も参考になる実装がないので独自に実装してあります。
+
+※Automatic1111氏のWeb UIではカンマを意識して分割、といったこともしているようですが、私の場合はそこまでしておらず単純な分割です。
+
+# 学習の手順
+
+あらかじめこのリポジトリのREADMEを参照し、環境整備を行ってください。
+
+## データの準備
+
+[学習データの準備について](./train_README-ja.md) を参照してください。fine tuningではメタデータを用いるfine tuning方式のみ対応しています。
+
+## 学習の実行
+たとえば以下のように実行します。以下は省メモリ化のための設定です。それぞれの行を必要に応じて書き換えてください。
+
+```
+accelerate launch --num_cpu_threads_per_process 1 fine_tune.py 
+    --pretrained_model_name_or_path=<.ckptまたは.safetensordまたはDiffusers版モデルのディレクトリ> 
+    --output_dir=<学習したモデルの出力先フォルダ>  
+    --output_name=<学習したモデル出力時のファイル名> 
+    --dataset_config=<データ準備で作成した.tomlファイル> 
+    --save_model_as=safetensors 
+    --learning_rate=5e-6 --max_train_steps=10000 
+    --use_8bit_adam --xformers --gradient_checkpointing
+    --mixed_precision=fp16
+```
+
+`num_cpu_threads_per_process` には通常は1を指定するとよいようです。
+
+`pretrained_model_name_or_path` に追加学習を行う元となるモデルを指定します。Stable Diffusionのcheckpointファイル（.ckptまたは.safetensors）、Diffusersのローカルディスクにあるモデルディレクトリ、DiffusersのモデルID（"stabilityai/stable-diffusion-2"など）が指定できます。
+
+`output_dir` に学習後のモデルを保存するフォルダを指定します。`output_name` にモデルのファイル名を拡張子を除いて指定します。`save_model_as` でsafetensors形式での保存を指定しています。
+
+`dataset_config` に `.toml` ファイルを指定します。ファイル内でのバッチサイズ指定は、当初はメモリ消費を抑えるために `1` としてください。
+
+学習させるステップ数 `max_train_steps` を10000とします。学習率 `learning_rate` はここでは5e-6を指定しています。
+
+省メモリ化のため `mixed_precision="fp16"` を指定します（RTX30 シリーズ以降では `bf16` も指定できます。環境整備時にaccelerateに行った設定と合わせてください）。また `gradient_checkpointing` を指定します。
+
+オプティマイザ（モデルを学習データにあうように最適化＝学習させるクラス）にメモリ消費の少ない 8bit AdamW を使うため、 `optimizer_type="AdamW8bit"` を指定します。
+
+`xformers` オプションを指定し、xformersのCrossAttentionを用います。xformersをインストールしていない場合やエラーとなる場合（環境にもよりますが `mixed_precision="no"` の場合など）、代わりに `mem_eff_attn` オプションを指定すると省メモリ版CrossAttentionを使用します（速度は遅くなります）。
+
+ある程度メモリがある場合は、`.toml` ファイルを編集してバッチサイズをたとえば `4` くらいに増やしてください（高速化と精度向上の可能性があります）。
+
+### よく使われるオプションについて
+
+以下の場合にはオプションに関するドキュメントを参照してください。
+
+- Stable Diffusion 2.xまたはそこからの派生モデルを学習する
+- clip skipを2以上を前提としたモデルを学習する
+- 75トークンを超えたキャプションで学習する
+
+### バッチサイズについて
+
+モデル全体を学習するためLoRA等の学習に比べるとメモリ消費量は多くなります（DreamBoothと同じ）。
+
+### 学習率について
+
+1e-6から5e-6程度が一般的なようです。他のfine tuningの例なども参照してみてください。
+
+### 以前の形式のデータセット指定をした場合のコマンドライン
+
+解像度やバッチサイズをオプションで指定します。コマンドラインの例は以下の通りです。
+
+```
+accelerate launch --num_cpu_threads_per_process 1 fine_tune.py 
+    --pretrained_model_name_or_path=model.ckpt 
+    --in_json meta_lat.json 
+    --train_data_dir=train_data 
+    --output_dir=fine_tuned 
+    --shuffle_caption 
+    --train_batch_size=1 --learning_rate=5e-6 --max_train_steps=10000 
+    --use_8bit_adam --xformers --gradient_checkpointing
+    --mixed_precision=bf16
+    --save_every_n_epochs=4
+```
+
+<!-- 
+### 勾配をfp16とした学習（実験的機能）
+full_fp16オプションを指定すると勾配を通常のfloat32からfloat16（fp16）に変更して学習します（mixed precisionではなく完全なfp16学習になるようです）。これによりSD1.xの512*512サイズでは8GB未満、SD2.xの512*512サイズで12GB未満のVRAM使用量で学習できるようです。
+
+あらかじめaccelerate configでfp16を指定し、オプションでmixed_precision="fp16"としてください（bf16では動作しません）。
+
+メモリ使用量を最小化するためには、xformers、use_8bit_adam、gradient_checkpointingの各オプションを指定し、train_batch_sizeを1としてください。
+（余裕があるようならtrain_batch_sizeを段階的に増やすと若干精度が上がるはずです。）
+
+PyTorchのソースにパッチを当てて無理やり実現しています（PyTorch 1.12.1と1.13.0で確認）。精度はかなり落ちますし、途中で学習失敗する確率も高くなります。学習率やステップ数の設定もシビアなようです。それらを認識したうえで自己責任でお使いください。
+-->
+
+# fine tuning特有のその他の主なオプション
+
+すべてのオプションについては別文書を参照してください。
+
+## `train_text_encoder`
+Text Encoderも学習対象とします。メモリ使用量が若干増加します。
+
+通常のfine tuningではText Encoderは学習対象としませんが（恐らくText Encoderの出力に従うようにU-Netを学習するため）、学習データ数が少ない場合には、DreamBoothのようにText Encoder側に学習させるのも有効的なようです。
+
+## `diffusers_xformers`
+スクリプト独自のxformers置換機能ではなくDiffusersのxformers機能を利用します。Hypernetworkの学習はできなくなります。

finetune_gui.py

@@ -0,0 +1,900 @@
+import gradio as gr
+import json
+import math
+import os
+import subprocess
+import pathlib
+import argparse
+from library.common_gui import (
+    get_folder_path,
+    get_file_path,
+    get_saveasfile_path,
+    save_inference_file,
+    gradio_advanced_training,
+    run_cmd_advanced_training,
+    gradio_training,
+    run_cmd_advanced_training,
+    gradio_config,
+    gradio_source_model,
+    color_aug_changed,
+    run_cmd_training,
+    # set_legacy_8bitadam,
+    update_my_data,
+    check_if_model_exist,
+)
+from library.tensorboard_gui import (
+    gradio_tensorboard,
+    start_tensorboard,
+    stop_tensorboard,
+)
+from library.utilities import utilities_tab
+from library.sampler_gui import sample_gradio_config, run_cmd_sample
+from easygui import msgbox
+
+folder_symbol = '\U0001f4c2'  # 📂
+refresh_symbol = '\U0001f504'  # 🔄
+save_style_symbol = '\U0001f4be'  # 💾
+document_symbol = '\U0001F4C4'   # 📄
+
+PYTHON = 'python3' if os.name == 'posix' else './venv/Scripts/python.exe'
+
+
+def save_configuration(
+    save_as,
+    file_path,
+    pretrained_model_name_or_path,
+    v2,
+    v_parameterization,
+    train_dir,
+    image_folder,
+    output_dir,
+    logging_dir,
+    max_resolution,
+    min_bucket_reso,
+    max_bucket_reso,
+    batch_size,
+    flip_aug,
+    caption_metadata_filename,
+    latent_metadata_filename,
+    full_path,
+    learning_rate,
+    lr_scheduler,
+    lr_warmup,
+    dataset_repeats,
+    train_batch_size,
+    epoch,
+    save_every_n_epochs,
+    mixed_precision,
+    save_precision,
+    seed,
+    num_cpu_threads_per_process,
+    train_text_encoder,
+    create_caption,
+    create_buckets,
+    save_model_as,
+    caption_extension,
+    # use_8bit_adam,
+    xformers,
+    clip_skip,
+    save_state,
+    resume,
+    gradient_checkpointing,
+    gradient_accumulation_steps,
+    mem_eff_attn,
+    shuffle_caption,
+    output_name,
+    max_token_length,
+    max_train_epochs,
+    max_data_loader_n_workers,
+    full_fp16,
+    color_aug,
+    model_list,
+    cache_latents,
+    use_latent_files,
+    keep_tokens,
+    persistent_data_loader_workers,
+    bucket_no_upscale,
+    random_crop,
+    bucket_reso_steps,
+    caption_dropout_every_n_epochs,
+    caption_dropout_rate,
+    optimizer,
+    optimizer_args,
+    noise_offset,
+    sample_every_n_steps,
+    sample_every_n_epochs,
+    sample_sampler,
+    sample_prompts,
+    additional_parameters,
+    vae_batch_size,
+    min_snr_gamma,weighted_captions,
+):
+    # Get list of function parameters and values
+    parameters = list(locals().items())
+
+    original_file_path = file_path
+
+    save_as_bool = True if save_as.get('label') == 'True' else False
+
+    if save_as_bool:
+        print('Save as...')
+        file_path = get_saveasfile_path(file_path)
+    else:
+        print('Save...')
+        if file_path == None or file_path == '':
+            file_path = get_saveasfile_path(file_path)
+
+    # print(file_path)
+
+    if file_path == None or file_path == '':
+        return original_file_path  # In case a file_path was provided and the user decide to cancel the open action
+
+    # Return the values of the variables as a dictionary
+    variables = {
+        name: value
+        for name, value in parameters  # locals().items()
+        if name
+        not in [
+            'file_path',
+            'save_as',
+        ]
+    }
+
+    # Extract the destination directory from the file path
+    destination_directory = os.path.dirname(file_path)
+
+    # Create the destination directory if it doesn't exist
+    if not os.path.exists(destination_directory):
+        os.makedirs(destination_directory)
+
+    # Save the data to the selected file
+    with open(file_path, 'w') as file:
+        json.dump(variables, file, indent=2)
+
+    return file_path
+
+
+def open_configuration(
+    ask_for_file,
+    file_path,
+    pretrained_model_name_or_path,
+    v2,
+    v_parameterization,
+    train_dir,
+    image_folder,
+    output_dir,
+    logging_dir,
+    max_resolution,
+    min_bucket_reso,
+    max_bucket_reso,
+    batch_size,
+    flip_aug,
+    caption_metadata_filename,
+    latent_metadata_filename,
+    full_path,
+    learning_rate,
+    lr_scheduler,
+    lr_warmup,
+    dataset_repeats,
+    train_batch_size,
+    epoch,
+    save_every_n_epochs,
+    mixed_precision,
+    save_precision,
+    seed,
+    num_cpu_threads_per_process,
+    train_text_encoder,
+    create_caption,
+    create_buckets,
+    save_model_as,
+    caption_extension,
+    # use_8bit_adam,
+    xformers,
+    clip_skip,
+    save_state,
+    resume,
+    gradient_checkpointing,
+    gradient_accumulation_steps,
+    mem_eff_attn,
+    shuffle_caption,
+    output_name,
+    max_token_length,
+    max_train_epochs,
+    max_data_loader_n_workers,
+    full_fp16,
+    color_aug,
+    model_list,
+    cache_latents,
+    use_latent_files,
+    keep_tokens,
+    persistent_data_loader_workers,
+    bucket_no_upscale,
+    random_crop,
+    bucket_reso_steps,
+    caption_dropout_every_n_epochs,
+    caption_dropout_rate,
+    optimizer,
+    optimizer_args,
+    noise_offset,
+    sample_every_n_steps,
+    sample_every_n_epochs,
+    sample_sampler,
+    sample_prompts,
+    additional_parameters,
+    vae_batch_size,
+    min_snr_gamma,weighted_captions,
+):
+    # Get list of function parameters and values
+    parameters = list(locals().items())
+
+    ask_for_file = True if ask_for_file.get('label') == 'True' else False
+
+    original_file_path = file_path
+
+    if ask_for_file:
+        file_path = get_file_path(file_path)
+
+    if not file_path == '' and not file_path == None:
+        # load variables from JSON file
+        with open(file_path, 'r') as f:
+            my_data = json.load(f)
+            print('Loading config...')
+            # Update values to fix deprecated use_8bit_adam checkbox and set appropriate optimizer if it is set to True
+            my_data = update_my_data(my_data)
+    else:
+        file_path = original_file_path  # In case a file_path was provided and the user decide to cancel the open action
+        my_data = {}
+
+    values = [file_path]
+    for key, value in parameters:
+        # Set the value in the dictionary to the corresponding value in `my_data`, or the default value if not found
+        if not key in ['ask_for_file', 'file_path']:
+            values.append(my_data.get(key, value))
+    return tuple(values)
+
+
+def train_model(
+    pretrained_model_name_or_path,
+    v2,
+    v_parameterization,
+    train_dir,
+    image_folder,
+    output_dir,
+    logging_dir,
+    max_resolution,
+    min_bucket_reso,
+    max_bucket_reso,
+    batch_size,
+    flip_aug,
+    caption_metadata_filename,
+    latent_metadata_filename,
+    full_path,
+    learning_rate,
+    lr_scheduler,
+    lr_warmup,
+    dataset_repeats,
+    train_batch_size,
+    epoch,
+    save_every_n_epochs,
+    mixed_precision,
+    save_precision,
+    seed,
+    num_cpu_threads_per_process,
+    train_text_encoder,
+    generate_caption_database,
+    generate_image_buckets,
+    save_model_as,
+    caption_extension,
+    # use_8bit_adam,
+    xformers,
+    clip_skip,
+    save_state,
+    resume,
+    gradient_checkpointing,
+    gradient_accumulation_steps,
+    mem_eff_attn,
+    shuffle_caption,
+    output_name,
+    max_token_length,
+    max_train_epochs,
+    max_data_loader_n_workers,
+    full_fp16,
+    color_aug,
+    model_list,  # Keep this. Yes, it is unused here but required given the common list used
+    cache_latents,
+    use_latent_files,
+    keep_tokens,
+    persistent_data_loader_workers,
+    bucket_no_upscale,
+    random_crop,
+    bucket_reso_steps,
+    caption_dropout_every_n_epochs,
+    caption_dropout_rate,
+    optimizer,
+    optimizer_args,
+    noise_offset,
+    sample_every_n_steps,
+    sample_every_n_epochs,
+    sample_sampler,
+    sample_prompts,
+    additional_parameters,
+    vae_batch_size,
+    min_snr_gamma,weighted_captions,
+):
+    if check_if_model_exist(output_name, output_dir, save_model_as):
+        return
+    
+    if optimizer == 'Adafactor' and lr_warmup != '0':
+        msgbox("Warning: lr_scheduler is set to 'Adafactor', so 'LR warmup (% of steps)' will be considered 0.", title="Warning")
+        lr_warmup = '0'
+
+    # create caption json file
+    if generate_caption_database:
+        if not os.path.exists(train_dir):
+            os.mkdir(train_dir)
+
+        run_cmd = f'{PYTHON} finetune/merge_captions_to_metadata.py'
+        if caption_extension == '':
+            run_cmd += f' --caption_extension=".caption"'
+        else:
+            run_cmd += f' --caption_extension={caption_extension}'
+        run_cmd += f' "{image_folder}"'
+        run_cmd += f' "{train_dir}/{caption_metadata_filename}"'
+        if full_path:
+            run_cmd += f' --full_path'
+
+        print(run_cmd)
+
+        # Run the command
+        if os.name == 'posix':
+            os.system(run_cmd)
+        else:
+            subprocess.run(run_cmd)
+
+    # create images buckets
+    if generate_image_buckets:
+        run_cmd = f'{PYTHON} finetune/prepare_buckets_latents.py'
+        run_cmd += f' "{image_folder}"'
+        run_cmd += f' "{train_dir}/{caption_metadata_filename}"'
+        run_cmd += f' "{train_dir}/{latent_metadata_filename}"'
+        run_cmd += f' "{pretrained_model_name_or_path}"'
+        run_cmd += f' --batch_size={batch_size}'
+        run_cmd += f' --max_resolution={max_resolution}'
+        run_cmd += f' --min_bucket_reso={min_bucket_reso}'
+        run_cmd += f' --max_bucket_reso={max_bucket_reso}'
+        run_cmd += f' --mixed_precision={mixed_precision}'
+        # if flip_aug:
+        #     run_cmd += f' --flip_aug'
+        if full_path:
+            run_cmd += f' --full_path'
+
+        print(run_cmd)
+
+        # Run the command
+        if os.name == 'posix':
+            os.system(run_cmd)
+        else:
+            subprocess.run(run_cmd)
+
+    image_num = len(
+        [
+            f
+            for f, lower_f in (
+                (file, file.lower()) for file in os.listdir(image_folder)
+            )
+            if lower_f.endswith(('.jpg', '.jpeg', '.png', '.webp'))
+        ]
+    )
+    print(f'image_num = {image_num}')
+
+    repeats = int(image_num) * int(dataset_repeats)
+    print(f'repeats = {str(repeats)}')
+
+    # calculate max_train_steps
+    max_train_steps = int(
+        math.ceil(float(repeats) / int(train_batch_size) * int(epoch))
+    )
+
+    # Divide by two because flip augmentation create two copied of the source images
+    if flip_aug:
+        max_train_steps = int(math.ceil(float(max_train_steps) / 2))
+
+    print(f'max_train_steps = {max_train_steps}')
+
+    lr_warmup_steps = round(float(int(lr_warmup) * int(max_train_steps) / 100))
+    print(f'lr_warmup_steps = {lr_warmup_steps}')
+
+    run_cmd = f'accelerate launch --num_cpu_threads_per_process={num_cpu_threads_per_process} "./fine_tune.py"'
+    if v2:
+        run_cmd += ' --v2'
+    if v_parameterization:
+        run_cmd += ' --v_parameterization'
+    if train_text_encoder:
+        run_cmd += ' --train_text_encoder'
+    if weighted_captions:
+        run_cmd += ' --weighted_captions'
+    run_cmd += (
+        f' --pretrained_model_name_or_path="{pretrained_model_name_or_path}"'
+    )
+    if use_latent_files == 'Yes':
+        run_cmd += f' --in_json="{train_dir}/{latent_metadata_filename}"'
+    else:
+        run_cmd += f' --in_json="{train_dir}/{caption_metadata_filename}"'
+    run_cmd += f' --train_data_dir="{image_folder}"'
+    run_cmd += f' --output_dir="{output_dir}"'
+    if not logging_dir == '':
+        run_cmd += f' --logging_dir="{logging_dir}"'
+    run_cmd += f' --dataset_repeats={dataset_repeats}'
+    run_cmd += f' --learning_rate={learning_rate}'
+
+    run_cmd += ' --enable_bucket'
+    run_cmd += f' --resolution={max_resolution}'
+    run_cmd += f' --min_bucket_reso={min_bucket_reso}'
+    run_cmd += f' --max_bucket_reso={max_bucket_reso}'
+
+    if not save_model_as == 'same as source model':
+        run_cmd += f' --save_model_as={save_model_as}'
+    if int(gradient_accumulation_steps) > 1:
+        run_cmd += f' --gradient_accumulation_steps={int(gradient_accumulation_steps)}'
+    # if save_state:
+    #     run_cmd += ' --save_state'
+    # if not resume == '':
+    #     run_cmd += f' --resume={resume}'
+    if not output_name == '':
+        run_cmd += f' --output_name="{output_name}"'
+    if int(max_token_length) > 75:
+        run_cmd += f' --max_token_length={max_token_length}'
+
+    run_cmd += run_cmd_training(
+        learning_rate=learning_rate,
+        lr_scheduler=lr_scheduler,
+        lr_warmup_steps=lr_warmup_steps,
+        train_batch_size=train_batch_size,
+        max_train_steps=max_train_steps,
+        save_every_n_epochs=save_every_n_epochs,
+        mixed_precision=mixed_precision,
+        save_precision=save_precision,
+        seed=seed,
+        caption_extension=caption_extension,
+        cache_latents=cache_latents,
+        optimizer=optimizer,
+        optimizer_args=optimizer_args,
+    )
+
+    run_cmd += run_cmd_advanced_training(
+        max_train_epochs=max_train_epochs,
+        max_data_loader_n_workers=max_data_loader_n_workers,
+        max_token_length=max_token_length,
+        resume=resume,
+        save_state=save_state,
+        mem_eff_attn=mem_eff_attn,
+        clip_skip=clip_skip,
+        flip_aug=flip_aug,
+        color_aug=color_aug,
+        shuffle_caption=shuffle_caption,
+        gradient_checkpointing=gradient_checkpointing,
+        full_fp16=full_fp16,
+        xformers=xformers,
+        # use_8bit_adam=use_8bit_adam,
+        keep_tokens=keep_tokens,
+        persistent_data_loader_workers=persistent_data_loader_workers,
+        bucket_no_upscale=bucket_no_upscale,
+        random_crop=random_crop,
+        bucket_reso_steps=bucket_reso_steps,
+        caption_dropout_every_n_epochs=caption_dropout_every_n_epochs,
+        caption_dropout_rate=caption_dropout_rate,
+        noise_offset=noise_offset,
+        additional_parameters=additional_parameters,
+        vae_batch_size=vae_batch_size,
+        min_snr_gamma=min_snr_gamma,
+    )
+
+    run_cmd += run_cmd_sample(
+        sample_every_n_steps,
+        sample_every_n_epochs,
+        sample_sampler,
+        sample_prompts,
+        output_dir,
+    )
+
+    print(run_cmd)
+
+    # Run the command
+    if os.name == 'posix':
+        os.system(run_cmd)
+    else:
+        subprocess.run(run_cmd)
+
+    # check if output_dir/last is a folder... therefore it is a diffuser model
+    last_dir = pathlib.Path(f'{output_dir}/{output_name}')
+
+    if not last_dir.is_dir():
+        # Copy inference model for v2 if required
+        save_inference_file(output_dir, v2, v_parameterization, output_name)
+
+
+def remove_doublequote(file_path):
+    if file_path != None:
+        file_path = file_path.replace('"', '')
+
+    return file_path
+
+
+def finetune_tab():
+    dummy_db_true = gr.Label(value=True, visible=False)
+    dummy_db_false = gr.Label(value=False, visible=False)
+    gr.Markdown('Train a custom model using kohya finetune python code...')
+
+    (
+        button_open_config,
+        button_save_config,
+        button_save_as_config,
+        config_file_name,
+        button_load_config,
+    ) = gradio_config()
+
+    (
+        pretrained_model_name_or_path,
+        v2,
+        v_parameterization,
+        save_model_as,
+        model_list,
+    ) = gradio_source_model()
+
+    with gr.Tab('Folders'):
+        with gr.Row():
+            train_dir = gr.Textbox(
+                label='Training config folder',
+                placeholder='folder where the training configuration files will be saved',
+            )
+            train_dir_folder = gr.Button(
+                folder_symbol, elem_id='open_folder_small'
+            )
+            train_dir_folder.click(
+                get_folder_path,
+                outputs=train_dir,
+                show_progress=False,
+            )
+
+            image_folder = gr.Textbox(
+                label='Training Image folder',
+                placeholder='folder where the training images are located',
+            )
+            image_folder_input_folder = gr.Button(
+                folder_symbol, elem_id='open_folder_small'
+            )
+            image_folder_input_folder.click(
+                get_folder_path,
+                outputs=image_folder,
+                show_progress=False,
+            )
+        with gr.Row():
+            output_dir = gr.Textbox(
+                label='Model output folder',
+                placeholder='folder where the model will be saved',
+            )
+            output_dir_input_folder = gr.Button(
+                folder_symbol, elem_id='open_folder_small'
+            )
+            output_dir_input_folder.click(
+                get_folder_path,
+                outputs=output_dir,
+                show_progress=False,
+            )
+
+            logging_dir = gr.Textbox(
+                label='Logging folder',
+                placeholder='Optional: enable logging and output TensorBoard log to this folder',
+            )
+            logging_dir_input_folder = gr.Button(
+                folder_symbol, elem_id='open_folder_small'
+            )
+            logging_dir_input_folder.click(
+                get_folder_path,
+                outputs=logging_dir,
+                show_progress=False,
+            )
+        with gr.Row():
+            output_name = gr.Textbox(
+                label='Model output name',
+                placeholder='Name of the model to output',
+                value='last',
+                interactive=True,
+            )
+        train_dir.change(
+            remove_doublequote,
+            inputs=[train_dir],
+            outputs=[train_dir],
+        )
+        image_folder.change(
+            remove_doublequote,
+            inputs=[image_folder],
+            outputs=[image_folder],
+        )
+        output_dir.change(
+            remove_doublequote,
+            inputs=[output_dir],
+            outputs=[output_dir],
+        )
+    with gr.Tab('Dataset preparation'):
+        with gr.Row():
+            max_resolution = gr.Textbox(
+                label='Resolution (width,height)', value='512,512'
+            )
+            min_bucket_reso = gr.Textbox(
+                label='Min bucket resolution', value='256'
+            )
+            max_bucket_reso = gr.Textbox(
+                label='Max bucket resolution', value='1024'
+            )
+            batch_size = gr.Textbox(label='Batch size', value='1')
+        with gr.Row():
+            create_caption = gr.Checkbox(
+                label='Generate caption metadata', value=True
+            )
+            create_buckets = gr.Checkbox(
+                label='Generate image buckets metadata', value=True
+            )
+            use_latent_files = gr.Dropdown(
+                label='Use latent files',
+                choices=[
+                    'No',
+                    'Yes',
+                ],
+                value='Yes',
+            )
+        with gr.Accordion('Advanced parameters', open=False):
+            with gr.Row():
+                caption_metadata_filename = gr.Textbox(
+                    label='Caption metadata filename', value='meta_cap.json'
+                )
+                latent_metadata_filename = gr.Textbox(
+                    label='Latent metadata filename', value='meta_lat.json'
+                )
+            with gr.Row():
+                full_path = gr.Checkbox(label='Use full path', value=True)
+                weighted_captions = gr.Checkbox(
+                    label='Weighted captions', value=False
+                )
+    with gr.Tab('Training parameters'):
+        (
+            learning_rate,
+            lr_scheduler,
+            lr_warmup,
+            train_batch_size,
+            epoch,
+            save_every_n_epochs,
+            mixed_precision,
+            save_precision,
+            num_cpu_threads_per_process,
+            seed,
+            caption_extension,
+            cache_latents,
+            optimizer,
+            optimizer_args,
+        ) = gradio_training(learning_rate_value='1e-5')
+        with gr.Row():
+            dataset_repeats = gr.Textbox(label='Dataset repeats', value=40)
+            train_text_encoder = gr.Checkbox(
+                label='Train text encoder', value=True
+            )
+        with gr.Accordion('Advanced parameters', open=False):
+            with gr.Row():
+                gradient_accumulation_steps = gr.Number(
+                    label='Gradient accumulate steps', value='1'
+                )
+            (
+                # use_8bit_adam,
+                xformers,
+                full_fp16,
+                gradient_checkpointing,
+                shuffle_caption,
+                color_aug,
+                flip_aug,
+                clip_skip,
+                mem_eff_attn,
+                save_state,
+                resume,
+                max_token_length,
+                max_train_epochs,
+                max_data_loader_n_workers,
+                keep_tokens,
+                persistent_data_loader_workers,
+                bucket_no_upscale,
+                random_crop,
+                bucket_reso_steps,
+                caption_dropout_every_n_epochs,
+                caption_dropout_rate,
+                noise_offset,
+                additional_parameters,
+                vae_batch_size,
+                min_snr_gamma,
+            ) = gradio_advanced_training()
+            color_aug.change(
+                color_aug_changed,
+                inputs=[color_aug],
+                outputs=[cache_latents],  # Not applicable to fine_tune.py
+            )
+
+        (
+            sample_every_n_steps,
+            sample_every_n_epochs,
+            sample_sampler,
+            sample_prompts,
+        ) = sample_gradio_config()
+
+    button_run = gr.Button('Train model', variant='primary')
+
+    # Setup gradio tensorboard buttons
+    button_start_tensorboard, button_stop_tensorboard = gradio_tensorboard()
+
+    button_start_tensorboard.click(
+        start_tensorboard,
+        inputs=logging_dir,
+    )
+
+    button_stop_tensorboard.click(
+        stop_tensorboard,
+        show_progress=False,
+    )
+
+    settings_list = [
+        pretrained_model_name_or_path,
+        v2,
+        v_parameterization,
+        train_dir,
+        image_folder,
+        output_dir,
+        logging_dir,
+        max_resolution,
+        min_bucket_reso,
+        max_bucket_reso,
+        batch_size,
+        flip_aug,
+        caption_metadata_filename,
+        latent_metadata_filename,
+        full_path,
+        learning_rate,
+        lr_scheduler,
+        lr_warmup,
+        dataset_repeats,
+        train_batch_size,
+        epoch,
+        save_every_n_epochs,
+        mixed_precision,
+        save_precision,
+        seed,
+        num_cpu_threads_per_process,
+        train_text_encoder,
+        create_caption,
+        create_buckets,
+        save_model_as,
+        caption_extension,
+        # use_8bit_adam,
+        xformers,
+        clip_skip,
+        save_state,
+        resume,
+        gradient_checkpointing,
+        gradient_accumulation_steps,
+        mem_eff_attn,
+        shuffle_caption,
+        output_name,
+        max_token_length,
+        max_train_epochs,
+        max_data_loader_n_workers,
+        full_fp16,
+        color_aug,
+        model_list,
+        cache_latents,
+        use_latent_files,
+        keep_tokens,
+        persistent_data_loader_workers,
+        bucket_no_upscale,
+        random_crop,
+        bucket_reso_steps,
+        caption_dropout_every_n_epochs,
+        caption_dropout_rate,
+        optimizer,
+        optimizer_args,
+        noise_offset,
+        sample_every_n_steps,
+        sample_every_n_epochs,
+        sample_sampler,
+        sample_prompts,
+        additional_parameters,
+        vae_batch_size,
+        min_snr_gamma,
+        weighted_captions,
+    ]
+
+    button_run.click(train_model, inputs=settings_list)
+
+    button_open_config.click(
+        open_configuration,
+        inputs=[dummy_db_true, config_file_name] + settings_list,
+        outputs=[config_file_name] + settings_list,
+        show_progress=False,
+    )
+
+    button_load_config.click(
+        open_configuration,
+        inputs=[dummy_db_false, config_file_name] + settings_list,
+        outputs=[config_file_name] + settings_list,
+        show_progress=False,
+    )
+
+    button_save_config.click(
+        save_configuration,
+        inputs=[dummy_db_false, config_file_name] + settings_list,
+        outputs=[config_file_name],
+        show_progress=False,
+    )
+
+    button_save_as_config.click(
+        save_configuration,
+        inputs=[dummy_db_true, config_file_name] + settings_list,
+        outputs=[config_file_name],
+        show_progress=False,
+    )
+
+
+def UI(**kwargs):
+
+    css = ''
+
+    if os.path.exists('./style.css'):
+        with open(os.path.join('./style.css'), 'r', encoding='utf8') as file:
+            print('Load CSS...')
+            css += file.read() + '\n'
+
+    interface = gr.Blocks(css=css)
+
+    with interface:
+        with gr.Tab('Finetune'):
+            finetune_tab()
+        with gr.Tab('Utilities'):
+            utilities_tab(enable_dreambooth_tab=False)
+
+    # Show the interface
+    launch_kwargs = {}
+    if not kwargs.get('username', None) == '':
+        launch_kwargs['auth'] = (
+            kwargs.get('username', None),
+            kwargs.get('password', None),
+        )
+    if kwargs.get('server_port', 0) > 0:
+        launch_kwargs['server_port'] = kwargs.get('server_port', 0)
+    if kwargs.get('inbrowser', False):
+        launch_kwargs['inbrowser'] = kwargs.get('inbrowser', False)
+    print(launch_kwargs)
+    interface.launch(**launch_kwargs)
+
+
+if __name__ == '__main__':
+    # torch.cuda.set_per_process_memory_fraction(0.48)
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        '--username', type=str, default='', help='Username for authentication'
+    )
+    parser.add_argument(
+        '--password', type=str, default='', help='Password for authentication'
+    )
+    parser.add_argument(
+        '--server_port',
+        type=int,
+        default=0,
+        help='Port to run the server listener on',
+    )
+    parser.add_argument(
+        '--inbrowser', action='store_true', help='Open in browser'
+    )
+
+    args = parser.parse_args()
+
+    UI(
+        username=args.username,
+        password=args.password,
+        inbrowser=args.inbrowser,
+        server_port=args.server_port,
+    )

gui.bat

@@ -0,0 +1,16 @@
+@echo off
+
+:: Activate the virtual environment
+call .\venv\Scripts\activate.bat
+set PATH=%PATH%;%~dp0venv\Lib\site-packages\torch\lib
+
+:: Debug info about system
+python.exe .\tools\debug_info.py
+
+:: Validate the requirements and store the exit code
+python.exe .\tools\validate_requirements.py
+
+:: If the exit code is 0, run the kohya_gui.py script with the command-line arguments
+if %errorlevel% equ 0 (
+    python.exe kohya_gui.py %*
+)

gui.ps1

@@ -0,0 +1,21 @@
+# Activate the virtual environment
+& .\venv\Scripts\activate
+$env:PATH += ";$($MyInvocation.MyCommand.Path)\venv\Lib\site-packages\torch\lib"
+
+# Debug info about system
+python.exe .\tools\debug_info.py
+
+# Validate the requirements and store the exit code
+python.exe .\tools\validate_requirements.py
+
+# If the exit code is 0, read arguments from gui_parameters.txt (if it exists)
+# and run the kohya_gui.py script with the command-line arguments
+if ($LASTEXITCODE -eq 0) {
+    $argsFromFile = @()
+    if (Test-Path .\gui_parameters.txt) {
+        $argsFromFile = Get-Content .\gui_parameters.txt -Encoding UTF8 | Where-Object { $_ -notmatch "^#" } | Foreach-Object { $_ -split " " }
+    }
+    $args_combo = $argsFromFile + $args
+    Write-Host "The arguments passed to this script were: $args_combo"
+    python.exe kohya_gui.py $args_combo
+}

gui.sh

@@ -0,0 +1,15 @@
+#!/usr/bin/env bash
+
+# This gets the directory the script is run from so pathing can work relative to the script where needed.
+SCRIPT_DIR=$(cd -- "$(dirname -- "$0")" && pwd)
+
+# Step into GUI local directory
+cd "$SCRIPT_DIR"
+
+# Activate the virtual environment
+source "$SCRIPT_DIR/venv/bin/activate"
+
+# If the requirements are validated, run the kohya_gui.py script with the command-line arguments
+if python  "$SCRIPT_DIR"/tools/validate_requirements.py -r "$SCRIPT_DIR"/requirements.txt; then
+    python "$SCRIPT_DIR/kohya_gui.py" "$@"
+fi

kohya_gui.py

@@ -0,0 +1,114 @@
+import gradio as gr
+import os
+import argparse
+from dreambooth_gui import dreambooth_tab
+from finetune_gui import finetune_tab
+from textual_inversion_gui import ti_tab
+from library.utilities import utilities_tab
+from library.extract_lora_gui import gradio_extract_lora_tab
+from library.extract_lycoris_locon_gui import gradio_extract_lycoris_locon_tab
+from library.merge_lora_gui import gradio_merge_lora_tab
+from library.resize_lora_gui import gradio_resize_lora_tab
+from library.extract_lora_from_dylora_gui import gradio_extract_dylora_tab
+from library.merge_lycoris_gui import gradio_merge_lycoris_tab
+from lora_gui import lora_tab
+
+
+def UI(**kwargs):
+    css = ''
+
+    if os.path.exists('./style.css'):
+        with open(os.path.join('./style.css'), 'r', encoding='utf8') as file:
+            print('Load CSS...')
+            css += file.read() + '\n'
+
+    interface = gr.Blocks(css=css, title='Kohya_ss GUI', theme=gr.themes.Default())
+
+    with interface:
+        with gr.Tab('Dreambooth'):
+            (
+                train_data_dir_input,
+                reg_data_dir_input,
+                output_dir_input,
+                logging_dir_input,
+            ) = dreambooth_tab()
+        with gr.Tab('Dreambooth LoRA'):
+            lora_tab()
+        with gr.Tab('Dreambooth TI'):
+            ti_tab()
+        with gr.Tab('Finetune'):
+            finetune_tab()
+        with gr.Tab('Utilities'):
+            utilities_tab(
+                train_data_dir_input=train_data_dir_input,
+                reg_data_dir_input=reg_data_dir_input,
+                output_dir_input=output_dir_input,
+                logging_dir_input=logging_dir_input,
+                enable_copy_info_button=True,
+            )
+            gradio_extract_dylora_tab()
+            gradio_extract_lora_tab()
+            gradio_extract_lycoris_locon_tab()
+            gradio_merge_lora_tab()
+            gradio_merge_lycoris_tab()
+            gradio_resize_lora_tab()
+
+    # Show the interface
+    launch_kwargs = {}
+    username = kwargs.get('username')
+    password = kwargs.get('password')
+    server_port = kwargs.get('server_port', 0)
+    inbrowser = kwargs.get('inbrowser', False)
+    share = kwargs.get('share', False)
+    server_name = kwargs.get('listen')
+
+    launch_kwargs['server_name'] = server_name
+    if username and password:
+        launch_kwargs['auth'] = (username, password)
+    if server_port > 0:
+        launch_kwargs['server_port'] = server_port
+    if inbrowser:
+        launch_kwargs['inbrowser'] = inbrowser
+    if share:
+        launch_kwargs['share'] = share
+    interface.launch(**launch_kwargs)
+
+
+if __name__ == '__main__':
+    # torch.cuda.set_per_process_memory_fraction(0.48)
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        '--listen',
+        type=str,
+        default='127.0.0.1',
+        help='IP to listen on for connections to Gradio',
+    )
+    parser.add_argument(
+        '--username', type=str, default='', help='Username for authentication'
+    )
+    parser.add_argument(
+        '--password', type=str, default='', help='Password for authentication'
+    )
+    parser.add_argument(
+        '--server_port',
+        type=int,
+        default=0,
+        help='Port to run the server listener on',
+    )
+    parser.add_argument(
+        '--inbrowser', action='store_true', help='Open in browser'
+    )
+    parser.add_argument(
+        '--share', action='store_true', help='Share the gradio UI'
+    )
+
+    args = parser.parse_args()
+
+    UI(
+        username=args.username,
+        password=args.password,
+        inbrowser=args.inbrowser,
+        server_port=args.server_port,
+        share=args.share,
+        listen=args.listen,
+    )

lora_gui.py

@@ -0,0 +1,1294 @@
+# v1: initial release
+# v2: add open and save folder icons
+# v3: Add new Utilities tab for Dreambooth folder preparation
+# v3.1: Adding captionning of images to utilities
+
+import gradio as gr
+import easygui
+import json
+import math
+import os
+import subprocess
+import pathlib
+import argparse
+from library.common_gui import (
+    get_folder_path,
+    remove_doublequote,
+    get_file_path,
+    get_any_file_path,
+    get_saveasfile_path,
+    color_aug_changed,
+    save_inference_file,
+    gradio_advanced_training,
+    run_cmd_advanced_training,
+    gradio_training,
+    gradio_config,
+    gradio_source_model,
+    run_cmd_training,
+    # set_legacy_8bitadam,
+    update_my_data,
+    check_if_model_exist,
+)
+from library.dreambooth_folder_creation_gui import (
+    gradio_dreambooth_folder_creation_tab,
+)
+from library.tensorboard_gui import (
+    gradio_tensorboard,
+    start_tensorboard,
+    stop_tensorboard,
+)
+from library.dataset_balancing_gui import gradio_dataset_balancing_tab
+from library.utilities import utilities_tab
+from library.merge_lora_gui import gradio_merge_lora_tab
+from library.svd_merge_lora_gui import gradio_svd_merge_lora_tab
+from library.verify_lora_gui import gradio_verify_lora_tab
+from library.resize_lora_gui import gradio_resize_lora_tab
+from library.sampler_gui import sample_gradio_config, run_cmd_sample
+from easygui import msgbox
+
+folder_symbol = '\U0001f4c2'  # 📂
+refresh_symbol = '\U0001f504'  # 🔄
+save_style_symbol = '\U0001f4be'  # 💾
+document_symbol = '\U0001F4C4'   # 📄
+path_of_this_folder = os.getcwd()
+
+
+def save_configuration(
+    save_as,
+    file_path,
+    pretrained_model_name_or_path,
+    v2,
+    v_parameterization,
+    logging_dir,
+    train_data_dir,
+    reg_data_dir,
+    output_dir,
+    max_resolution,
+    learning_rate,
+    lr_scheduler,
+    lr_warmup,
+    train_batch_size,
+    epoch,
+    save_every_n_epochs,
+    mixed_precision,
+    save_precision,
+    seed,
+    num_cpu_threads_per_process,
+    cache_latents,
+    caption_extension,
+    enable_bucket,
+    gradient_checkpointing,
+    full_fp16,
+    no_token_padding,
+    stop_text_encoder_training,
+    # use_8bit_adam,
+    xformers,
+    save_model_as,
+    shuffle_caption,
+    save_state,
+    resume,
+    prior_loss_weight,
+    text_encoder_lr,
+    unet_lr,
+    network_dim,
+    lora_network_weights,
+    color_aug,
+    flip_aug,
+    clip_skip,
+    gradient_accumulation_steps,
+    mem_eff_attn,
+    output_name,
+    model_list,
+    max_token_length,
+    max_train_epochs,
+    max_data_loader_n_workers,
+    network_alpha,
+    training_comment,
+    keep_tokens,
+    lr_scheduler_num_cycles,
+    lr_scheduler_power,
+    persistent_data_loader_workers,
+    bucket_no_upscale,
+    random_crop,
+    bucket_reso_steps,
+    caption_dropout_every_n_epochs,
+    caption_dropout_rate,
+    optimizer,
+    optimizer_args,
+    noise_offset,
+    LoRA_type,
+    conv_dim,
+    conv_alpha,
+    sample_every_n_steps,
+    sample_every_n_epochs,
+    sample_sampler,
+    sample_prompts,
+    additional_parameters,
+    vae_batch_size,
+    min_snr_gamma,
+    down_lr_weight,mid_lr_weight,up_lr_weight,block_lr_zero_threshold,block_dims,block_alphas,conv_dims,conv_alphas,
+    weighted_captions,unit,
+):
+    # Get list of function parameters and values
+    parameters = list(locals().items())
+
+    original_file_path = file_path
+
+    save_as_bool = True if save_as.get('label') == 'True' else False
+
+    if save_as_bool:
+        print('Save as...')
+        file_path = get_saveasfile_path(file_path)
+    else:
+        print('Save...')
+        if file_path == None or file_path == '':
+            file_path = get_saveasfile_path(file_path)
+
+    # print(file_path)
+
+    if file_path == None or file_path == '':
+        return original_file_path  # In case a file_path was provided and the user decide to cancel the open action
+
+    # Return the values of the variables as a dictionary
+    variables = {
+        name: value
+        for name, value in parameters  # locals().items()
+        if name
+        not in [
+            'file_path',
+            'save_as',
+        ]
+    }
+
+    # Extract the destination directory from the file path
+    destination_directory = os.path.dirname(file_path)
+
+    # Create the destination directory if it doesn't exist
+    if not os.path.exists(destination_directory):
+        os.makedirs(destination_directory)
+
+    # Save the data to the selected file
+    with open(file_path, 'w') as file:
+        json.dump(variables, file, indent=2)
+
+    return file_path
+
+
+def open_configuration(
+    ask_for_file,
+    file_path,
+    pretrained_model_name_or_path,
+    v2,
+    v_parameterization,
+    logging_dir,
+    train_data_dir,
+    reg_data_dir,
+    output_dir,
+    max_resolution,
+    learning_rate,
+    lr_scheduler,
+    lr_warmup,
+    train_batch_size,
+    epoch,
+    save_every_n_epochs,
+    mixed_precision,
+    save_precision,
+    seed,
+    num_cpu_threads_per_process,
+    cache_latents,
+    caption_extension,
+    enable_bucket,
+    gradient_checkpointing,
+    full_fp16,
+    no_token_padding,
+    stop_text_encoder_training,
+    # use_8bit_adam,
+    xformers,
+    save_model_as,
+    shuffle_caption,
+    save_state,
+    resume,
+    prior_loss_weight,
+    text_encoder_lr,
+    unet_lr,
+    network_dim,
+    lora_network_weights,
+    color_aug,
+    flip_aug,
+    clip_skip,
+    gradient_accumulation_steps,
+    mem_eff_attn,
+    output_name,
+    model_list,
+    max_token_length,
+    max_train_epochs,
+    max_data_loader_n_workers,
+    network_alpha,
+    training_comment,
+    keep_tokens,
+    lr_scheduler_num_cycles,
+    lr_scheduler_power,
+    persistent_data_loader_workers,
+    bucket_no_upscale,
+    random_crop,
+    bucket_reso_steps,
+    caption_dropout_every_n_epochs,
+    caption_dropout_rate,
+    optimizer,
+    optimizer_args,
+    noise_offset,
+    LoRA_type,
+    conv_dim,
+    conv_alpha,
+    sample_every_n_steps,
+    sample_every_n_epochs,
+    sample_sampler,
+    sample_prompts,
+    additional_parameters,
+    vae_batch_size,
+    min_snr_gamma,
+    down_lr_weight,mid_lr_weight,up_lr_weight,block_lr_zero_threshold,block_dims,block_alphas,conv_dims,conv_alphas,
+    weighted_captions,unit,
+):
+    # Get list of function parameters and values
+    parameters = list(locals().items())
+
+    ask_for_file = True if ask_for_file.get('label') == 'True' else False
+
+    original_file_path = file_path
+
+    if ask_for_file:
+        file_path = get_file_path(file_path)
+
+    if not file_path == '' and not file_path == None:
+        # load variables from JSON file
+        with open(file_path, 'r') as f:
+            my_data = json.load(f)
+            print('Loading config...')
+
+            # Update values to fix deprecated use_8bit_adam checkbox, set appropriate optimizer if it is set to True, etc.
+            my_data = update_my_data(my_data)
+    else:
+        file_path = original_file_path  # In case a file_path was provided and the user decide to cancel the open action
+        my_data = {}
+
+    values = [file_path]
+    for key, value in parameters:
+        # Set the value in the dictionary to the corresponding value in `my_data`, or the default value if not found
+        if not key in ['ask_for_file', 'file_path']:
+            values.append(my_data.get(key, value))
+
+    # This next section is about making the LoCon parameters visible if LoRA_type = 'Standard'
+    if my_data.get('LoRA_type', 'Standard') == 'LoCon':
+        values.append(gr.Row.update(visible=True))
+    else:
+        values.append(gr.Row.update(visible=False))
+
+    return tuple(values)
+
+
+def train_model(
+    print_only,
+    pretrained_model_name_or_path,
+    v2,
+    v_parameterization,
+    logging_dir,
+    train_data_dir,
+    reg_data_dir,
+    output_dir,
+    max_resolution,
+    learning_rate,
+    lr_scheduler,
+    lr_warmup,
+    train_batch_size,
+    epoch,
+    save_every_n_epochs,
+    mixed_precision,
+    save_precision,
+    seed,
+    num_cpu_threads_per_process,
+    cache_latents,
+    caption_extension,
+    enable_bucket,
+    gradient_checkpointing,
+    full_fp16,
+    no_token_padding,
+    stop_text_encoder_training_pct,
+    # use_8bit_adam,
+    xformers,
+    save_model_as,
+    shuffle_caption,
+    save_state,
+    resume,
+    prior_loss_weight,
+    text_encoder_lr,
+    unet_lr,
+    network_dim,
+    lora_network_weights,
+    color_aug,
+    flip_aug,
+    clip_skip,
+    gradient_accumulation_steps,
+    mem_eff_attn,
+    output_name,
+    model_list,  # Keep this. Yes, it is unused here but required given the common list used
+    max_token_length,
+    max_train_epochs,
+    max_data_loader_n_workers,
+    network_alpha,
+    training_comment,
+    keep_tokens,
+    lr_scheduler_num_cycles,
+    lr_scheduler_power,
+    persistent_data_loader_workers,
+    bucket_no_upscale,
+    random_crop,
+    bucket_reso_steps,
+    caption_dropout_every_n_epochs,
+    caption_dropout_rate,
+    optimizer,
+    optimizer_args,
+    noise_offset,
+    LoRA_type,
+    conv_dim,
+    conv_alpha,
+    sample_every_n_steps,
+    sample_every_n_epochs,
+    sample_sampler,
+    sample_prompts,
+    additional_parameters,
+    vae_batch_size,
+    min_snr_gamma,
+    down_lr_weight,mid_lr_weight,up_lr_weight,block_lr_zero_threshold,block_dims,block_alphas,conv_dims,conv_alphas,
+    weighted_captions,unit,
+):
+    print_only_bool = True if print_only.get('label') == 'True' else False
+
+    if pretrained_model_name_or_path == '':
+        msgbox('Source model information is missing')
+        return
+
+    if train_data_dir == '':
+        msgbox('Image folder path is missing')
+        return
+
+    if not os.path.exists(train_data_dir):
+        msgbox('Image folder does not exist')
+        return
+
+    if reg_data_dir != '':
+        if not os.path.exists(reg_data_dir):
+            msgbox('Regularisation folder does not exist')
+            return
+
+    if output_dir == '':
+        msgbox('Output folder path is missing')
+        return
+
+    if int(bucket_reso_steps) < 1:
+        msgbox('Bucket resolution steps need to be greater than 0')
+        return
+
+    if not os.path.exists(output_dir):
+        os.makedirs(output_dir)
+
+    if stop_text_encoder_training_pct > 0:
+        msgbox(
+            'Output "stop text encoder training" is not yet supported. Ignoring'
+        )
+        stop_text_encoder_training_pct = 0
+
+    if check_if_model_exist(output_name, output_dir, save_model_as):
+        return
+    
+    if optimizer == 'Adafactor' and lr_warmup != '0':
+        msgbox("Warning: lr_scheduler is set to 'Adafactor', so 'LR warmup (% of steps)' will be considered 0.", title="Warning")
+        lr_warmup = '0'
+
+    # If string is empty set string to 0.
+    if text_encoder_lr == '':
+        text_encoder_lr = 0
+    if unet_lr == '':
+        unet_lr = 0
+
+    # if (float(text_encoder_lr) == 0) and (float(unet_lr) == 0):
+    #     msgbox(
+    #         'At least one Learning Rate value for "Text encoder" or "Unet" need to be provided'
+    #     )
+    #     return
+
+    # Get a list of all subfolders in train_data_dir
+    subfolders = [
+        f
+        for f in os.listdir(train_data_dir)
+        if os.path.isdir(os.path.join(train_data_dir, f))
+    ]
+
+    total_steps = 0
+
+    # Loop through each subfolder and extract the number of repeats
+    for folder in subfolders:
+        try:
+            # Extract the number of repeats from the folder name
+            repeats = int(folder.split('_')[0])
+            
+            # Count the number of images in the folder
+            num_images = len(
+                [
+                    f
+                    for f, lower_f in (
+                        (file, file.lower())
+                        for file in os.listdir(
+                            os.path.join(train_data_dir, folder)
+                        )
+                    )
+                    if lower_f.endswith(('.jpg', '.jpeg', '.png', '.webp'))
+                ]
+            )
+
+            print(f'Folder {folder}: {num_images} images found')
+
+            # Calculate the total number of steps for this folder
+            steps = repeats * num_images
+
+            # Print the result
+            print(f'Folder {folder}: {steps} steps')
+
+            total_steps += steps
+            
+        except ValueError:
+            # Handle the case where the folder name does not contain an underscore
+            print(f"Error: '{folder}' does not contain an underscore, skipping...")    
+
+    # calculate max_train_steps
+    max_train_steps = int(
+        math.ceil(
+            float(total_steps)
+            / int(train_batch_size)
+            * int(epoch)
+            # * int(reg_factor)
+        )
+    )
+    print(f'max_train_steps = {max_train_steps}')
+
+    # calculate stop encoder training
+    if stop_text_encoder_training_pct == None:
+        stop_text_encoder_training = 0
+    else:
+        stop_text_encoder_training = math.ceil(
+            float(max_train_steps) / 100 * int(stop_text_encoder_training_pct)
+        )
+    print(f'stop_text_encoder_training = {stop_text_encoder_training}')
+
+    lr_warmup_steps = round(float(int(lr_warmup) * int(max_train_steps) / 100))
+    print(f'lr_warmup_steps = {lr_warmup_steps}')
+
+    run_cmd = f'accelerate launch --num_cpu_threads_per_process={num_cpu_threads_per_process} "train_network.py"'
+
+    if v2:
+        run_cmd += ' --v2'
+    if v_parameterization:
+        run_cmd += ' --v_parameterization'
+    if enable_bucket:
+        run_cmd += ' --enable_bucket'
+    if no_token_padding:
+        run_cmd += ' --no_token_padding'
+    if weighted_captions:
+        run_cmd += ' --weighted_captions'
+    run_cmd += (
+        f' --pretrained_model_name_or_path="{pretrained_model_name_or_path}"'
+    )
+    run_cmd += f' --train_data_dir="{train_data_dir}"'
+    if len(reg_data_dir):
+        run_cmd += f' --reg_data_dir="{reg_data_dir}"'
+    run_cmd += f' --resolution={max_resolution}'
+    run_cmd += f' --output_dir="{output_dir}"'
+    run_cmd += f' --logging_dir="{logging_dir}"'
+    run_cmd += f' --network_alpha="{network_alpha}"'
+    if not training_comment == '':
+        run_cmd += f' --training_comment="{training_comment}"'
+    if not stop_text_encoder_training == 0:
+        run_cmd += (
+            f' --stop_text_encoder_training={stop_text_encoder_training}'
+        )
+    if not save_model_as == 'same as source model':
+        run_cmd += f' --save_model_as={save_model_as}'
+    if not float(prior_loss_weight) == 1.0:
+        run_cmd += f' --prior_loss_weight={prior_loss_weight}'
+    if LoRA_type == 'LoCon' or LoRA_type == 'LyCORIS/LoCon':
+        try:
+            import lycoris
+        except ModuleNotFoundError:
+            print(
+                "\033[1;31mError:\033[0m The required module 'lycoris_lora' is not installed. Please install by running \033[33mupgrade.ps1\033[0m before running this program."
+            )
+            return
+        run_cmd += f' --network_module=lycoris.kohya'
+        run_cmd += f' --network_args "conv_dim={conv_dim}" "conv_alpha={conv_alpha}" "algo=lora"'
+    if LoRA_type == 'LyCORIS/LoHa':
+        try:
+            import lycoris
+        except ModuleNotFoundError:
+            print(
+                "\033[1;31mError:\033[0m The required module 'lycoris_lora' is not installed. Please install by running \033[33mupgrade.ps1\033[0m before running this program."
+            )
+            return
+        run_cmd += f' --network_module=lycoris.kohya'
+        run_cmd += f' --network_args "conv_dim={conv_dim}" "conv_alpha={conv_alpha}" "algo=loha"'
+    
+        
+    if LoRA_type in ['Kohya LoCon', 'Standard']:
+        kohya_lora_var_list = ['down_lr_weight', 'mid_lr_weight', 'up_lr_weight', 'block_lr_zero_threshold', 'block_dims', 'block_alphas', 'conv_dims', 'conv_alphas']
+        
+        run_cmd += f' --network_module=networks.lora'
+        kohya_lora_vars = {key: value for key, value in vars().items() if key in kohya_lora_var_list and value}
+
+        network_args = ''
+        if LoRA_type == 'Kohya LoCon':
+            network_args += f' "conv_dim={conv_dim}" "conv_alpha={conv_alpha}"'
+
+        for key, value in kohya_lora_vars.items():
+            if value:
+                network_args += f' {key}="{value}"'
+
+        if network_args:
+            run_cmd += f' --network_args{network_args}'
+            
+    if LoRA_type in ['Kohya DyLoRA']:
+        kohya_lora_var_list = ['conv_dim', 'conv_alpha', 'down_lr_weight', 'mid_lr_weight', 'up_lr_weight', 'block_lr_zero_threshold', 'block_dims', 'block_alphas', 'conv_dims', 'conv_alphas', 'unit']
+        
+        run_cmd += f' --network_module=networks.dylora'
+        kohya_lora_vars = {key: value for key, value in vars().items() if key in kohya_lora_var_list and value}
+
+        network_args = ''
+
+        for key, value in kohya_lora_vars.items():
+            if value:
+                network_args += f' {key}="{value}"'
+
+        if network_args:
+            run_cmd += f' --network_args{network_args}'
+
+    if not (float(text_encoder_lr) == 0) or not (float(unet_lr) == 0):
+        if not (float(text_encoder_lr) == 0) and not (float(unet_lr) == 0):
+            run_cmd += f' --text_encoder_lr={text_encoder_lr}'
+            run_cmd += f' --unet_lr={unet_lr}'
+        elif not (float(text_encoder_lr) == 0):
+            run_cmd += f' --text_encoder_lr={text_encoder_lr}'
+            run_cmd += f' --network_train_text_encoder_only'
+        else:
+            run_cmd += f' --unet_lr={unet_lr}'
+            run_cmd += f' --network_train_unet_only'
+    else:
+        if float(text_encoder_lr) == 0:
+            msgbox('Please input learning rate values.')
+            return
+
+    run_cmd += f' --network_dim={network_dim}'
+
+    if not lora_network_weights == '':
+        run_cmd += f' --network_weights="{lora_network_weights}"'
+    if int(gradient_accumulation_steps) > 1:
+        run_cmd += f' --gradient_accumulation_steps={int(gradient_accumulation_steps)}'
+    if not output_name == '':
+        run_cmd += f' --output_name="{output_name}"'
+    if not lr_scheduler_num_cycles == '':
+        run_cmd += f' --lr_scheduler_num_cycles="{lr_scheduler_num_cycles}"'
+    else:
+        run_cmd += f' --lr_scheduler_num_cycles="{epoch}"'
+    if not lr_scheduler_power == '':
+        run_cmd += f' --lr_scheduler_power="{lr_scheduler_power}"'
+
+    run_cmd += run_cmd_training(
+        learning_rate=learning_rate,
+        lr_scheduler=lr_scheduler,
+        lr_warmup_steps=lr_warmup_steps,
+        train_batch_size=train_batch_size,
+        max_train_steps=max_train_steps,
+        save_every_n_epochs=save_every_n_epochs,
+        mixed_precision=mixed_precision,
+        save_precision=save_precision,
+        seed=seed,
+        caption_extension=caption_extension,
+        cache_latents=cache_latents,
+        optimizer=optimizer,
+        optimizer_args=optimizer_args,
+    )
+
+    run_cmd += run_cmd_advanced_training(
+        max_train_epochs=max_train_epochs,
+        max_data_loader_n_workers=max_data_loader_n_workers,
+        max_token_length=max_token_length,
+        resume=resume,
+        save_state=save_state,
+        mem_eff_attn=mem_eff_attn,
+        clip_skip=clip_skip,
+        flip_aug=flip_aug,
+        color_aug=color_aug,
+        shuffle_caption=shuffle_caption,
+        gradient_checkpointing=gradient_checkpointing,
+        full_fp16=full_fp16,
+        xformers=xformers,
+        # use_8bit_adam=use_8bit_adam,
+        keep_tokens=keep_tokens,
+        persistent_data_loader_workers=persistent_data_loader_workers,
+        bucket_no_upscale=bucket_no_upscale,
+        random_crop=random_crop,
+        bucket_reso_steps=bucket_reso_steps,
+        caption_dropout_every_n_epochs=caption_dropout_every_n_epochs,
+        caption_dropout_rate=caption_dropout_rate,
+        noise_offset=noise_offset,
+        additional_parameters=additional_parameters,
+        vae_batch_size=vae_batch_size,
+        min_snr_gamma=min_snr_gamma,
+    )
+
+    run_cmd += run_cmd_sample(
+        sample_every_n_steps,
+        sample_every_n_epochs,
+        sample_sampler,
+        sample_prompts,
+        output_dir,
+    )
+    
+    # if not down_lr_weight == '':
+    #     run_cmd += f' --down_lr_weight="{down_lr_weight}"'
+    # if not mid_lr_weight == '':
+    #     run_cmd += f' --mid_lr_weight="{mid_lr_weight}"'
+    # if not up_lr_weight == '':
+    #     run_cmd += f' --up_lr_weight="{up_lr_weight}"'
+    # if not block_lr_zero_threshold == '':
+    #     run_cmd += f' --block_lr_zero_threshold="{block_lr_zero_threshold}"'
+    # if not block_dims == '':
+    #     run_cmd += f' --block_dims="{block_dims}"'
+    # if not block_alphas == '':
+    #     run_cmd += f' --block_alphas="{block_alphas}"'
+    # if not conv_dims == '':
+    #     run_cmd += f' --conv_dims="{conv_dims}"'
+    # if not conv_alphas == '':
+    #     run_cmd += f' --conv_alphas="{conv_alphas}"'
+        
+
+
+
+    if print_only_bool:
+        print(
+            '\033[93m\nHere is the trainer command as a reference. It will not be executed:\033[0m\n'
+        )
+        print('\033[96m' + run_cmd + '\033[0m\n')
+    else:
+        print(run_cmd)
+        # Run the command
+        if os.name == 'posix':
+            os.system(run_cmd)
+        else:
+            subprocess.run(run_cmd)
+
+        # check if output_dir/last is a folder... therefore it is a diffuser model
+        last_dir = pathlib.Path(f'{output_dir}/{output_name}')
+
+        if not last_dir.is_dir():
+            # Copy inference model for v2 if required
+            save_inference_file(
+                output_dir, v2, v_parameterization, output_name
+            )
+
+
+def lora_tab(
+    train_data_dir_input=gr.Textbox(),
+    reg_data_dir_input=gr.Textbox(),
+    output_dir_input=gr.Textbox(),
+    logging_dir_input=gr.Textbox(),
+):
+    dummy_db_true = gr.Label(value=True, visible=False)
+    dummy_db_false = gr.Label(value=False, visible=False)
+    gr.Markdown(
+        'Train a custom model using kohya train network LoRA python code...'
+    )
+    (
+        button_open_config,
+        button_save_config,
+        button_save_as_config,
+        config_file_name,
+        button_load_config,
+    ) = gradio_config()
+
+    (
+        pretrained_model_name_or_path,
+        v2,
+        v_parameterization,
+        save_model_as,
+        model_list,
+    ) = gradio_source_model(
+        save_model_as_choices=[
+            'ckpt',
+            'safetensors',
+        ]
+    )
+
+    with gr.Tab('Folders'):
+        with gr.Row():
+            train_data_dir = gr.Textbox(
+                label='Image folder',
+                placeholder='Folder where the training folders containing the images are located',
+            )
+            train_data_dir_folder = gr.Button('📂', elem_id='open_folder_small')
+            train_data_dir_folder.click(
+                get_folder_path,
+                outputs=train_data_dir,
+                show_progress=False,
+            )
+            reg_data_dir = gr.Textbox(
+                label='Regularisation folder',
+                placeholder='(Optional) Folder where where the regularization folders containing the images are located',
+            )
+            reg_data_dir_folder = gr.Button('📂', elem_id='open_folder_small')
+            reg_data_dir_folder.click(
+                get_folder_path,
+                outputs=reg_data_dir,
+                show_progress=False,
+            )
+        with gr.Row():
+            output_dir = gr.Textbox(
+                label='Output folder',
+                placeholder='Folder to output trained model',
+            )
+            output_dir_folder = gr.Button('📂', elem_id='open_folder_small')
+            output_dir_folder.click(
+                get_folder_path,
+                outputs=output_dir,
+                show_progress=False,
+            )
+            logging_dir = gr.Textbox(
+                label='Logging folder',
+                placeholder='Optional: enable logging and output TensorBoard log to this folder',
+            )
+            logging_dir_folder = gr.Button('📂', elem_id='open_folder_small')
+            logging_dir_folder.click(
+                get_folder_path,
+                outputs=logging_dir,
+                show_progress=False,
+            )
+        with gr.Row():
+            output_name = gr.Textbox(
+                label='Model output name',
+                placeholder='(Name of the model to output)',
+                value='last',
+                interactive=True,
+            )
+            training_comment = gr.Textbox(
+                label='Training comment',
+                placeholder='(Optional) Add training comment to be included in metadata',
+                interactive=True,
+            )
+        train_data_dir.change(
+            remove_doublequote,
+            inputs=[train_data_dir],
+            outputs=[train_data_dir],
+        )
+        reg_data_dir.change(
+            remove_doublequote,
+            inputs=[reg_data_dir],
+            outputs=[reg_data_dir],
+        )
+        output_dir.change(
+            remove_doublequote,
+            inputs=[output_dir],
+            outputs=[output_dir],
+        )
+        logging_dir.change(
+            remove_doublequote,
+            inputs=[logging_dir],
+            outputs=[logging_dir],
+        )
+    with gr.Tab('Training parameters'):
+        with gr.Row():
+            LoRA_type = gr.Dropdown(
+                label='LoRA type',
+                choices=[
+                    'Kohya DyLoRA',
+                    'Kohya LoCon',
+                    # 'LoCon',
+                    'LyCORIS/LoCon',
+                    'LyCORIS/LoHa',
+                    'Standard',
+                ],
+                value='Standard',
+            )
+            lora_network_weights = gr.Textbox(
+                label='LoRA network weights',
+                placeholder='{Optional) Path to existing LoRA network weights to resume training',
+            )
+            lora_network_weights_file = gr.Button(
+                document_symbol, elem_id='open_folder_small'
+            )
+            lora_network_weights_file.click(
+                get_any_file_path,
+                inputs=[lora_network_weights],
+                outputs=lora_network_weights,
+                show_progress=False,
+            )
+        (
+            learning_rate,
+            lr_scheduler,
+            lr_warmup,
+            train_batch_size,
+            epoch,
+            save_every_n_epochs,
+            mixed_precision,
+            save_precision,
+            num_cpu_threads_per_process,
+            seed,
+            caption_extension,
+            cache_latents,
+            optimizer,
+            optimizer_args,
+        ) = gradio_training(
+            learning_rate_value='0.0001',
+            lr_scheduler_value='cosine',
+            lr_warmup_value='10',
+        )
+
+        with gr.Row():
+            text_encoder_lr = gr.Textbox(
+                label='Text Encoder learning rate',
+                value='5e-5',
+                placeholder='Optional',
+            )
+            unet_lr = gr.Textbox(
+                label='Unet learning rate',
+                value='0.0001',
+                placeholder='Optional',
+            )
+            network_dim = gr.Slider(
+                minimum=1,
+                maximum=1024,
+                label='Network Rank (Dimension)',
+                value=8,
+                step=1,
+                interactive=True,
+            )
+            network_alpha = gr.Slider(
+                minimum=0.1,
+                maximum=1024,
+                label='Network Alpha',
+                value=1,
+                step=0.1,
+                interactive=True,
+                info='alpha for LoRA weight scaling',
+            )
+        with gr.Row(visible=False) as LoCon_row:
+
+            # locon= gr.Checkbox(label='Train a LoCon instead of a general LoRA (does not support v2 base models) (may not be able to some utilities now)', value=False)
+            conv_dim = gr.Slider(
+                minimum=1,
+                maximum=512,
+                value=1,
+                step=1,
+                label='Convolution Rank (Dimension)',
+            )
+            conv_alpha = gr.Slider(
+                minimum=0.1,
+                maximum=512,
+                value=1,
+                step=0.1,
+                label='Convolution Alpha',
+            )
+        with gr.Row(visible=False) as kohya_dylora:
+            unit = gr.Slider(
+                minimum=1,
+                maximum=64,
+                label='DyLoRA Unit',
+                value=1,
+                step=1,
+                interactive=True,
+            )
+            
+        # Show of hide LoCon conv settings depending on LoRA type selection
+        def update_LoRA_settings(LoRA_type):
+            # Print a message when LoRA type is changed
+            print('LoRA type changed...')
+
+            # Determine if LoCon_row should be visible based on LoRA_type
+            LoCon_row = LoRA_type in {'LoCon', 'Kohya DyLoRA', 'Kohya LoCon', 'LyCORIS/LoHa', 'LyCORIS/LoCon'}
+
+            # Determine if LoRA_type_change should be visible based on LoRA_type
+            LoRA_type_change = LoRA_type in {'Standard', 'Kohya DyLoRA', 'Kohya LoCon'}
+
+            # Determine if kohya_dylora_visible should be visible based on LoRA_type
+            kohya_dylora_visible = LoRA_type == 'Kohya DyLoRA'
+
+            # Return the updated visibility settings for the groups
+            return (
+                gr.Group.update(visible=LoCon_row),
+                gr.Group.update(visible=LoRA_type_change),
+                gr.Group.update(visible=kohya_dylora_visible),
+            )
+
+
+        with gr.Row():
+            max_resolution = gr.Textbox(
+                label='Max resolution',
+                value='512,512',
+                placeholder='512,512',
+                info='The maximum resolution of dataset images. W,H',
+            )
+            stop_text_encoder_training = gr.Slider(
+                minimum=0,
+                maximum=100,
+                value=0,
+                step=1,
+                label='Stop text encoder training',
+                info='After what % of steps should the text encoder stop being trained. 0 = train for all steps.',
+            )
+            enable_bucket = gr.Checkbox(label='Enable buckets', value=True,
+                info='Allow non similar resolution dataset images to be trained on.',)
+            
+        with gr.Accordion('Advanced Configuration', open=False):
+            with gr.Row(visible=True) as kohya_advanced_lora:
+                with gr.Tab(label='Weights'):
+                    with gr.Row(visible=True):
+                        down_lr_weight = gr.Textbox(
+                            label='Down LR weights',
+                            placeholder='(Optional) eg: 0,0,0,0,0,0,1,1,1,1,1,1',
+                            info='Specify the learning rate weight of the down blocks of U-Net.'
+                        )
+                        mid_lr_weight = gr.Textbox(
+                            label='Mid LR weights',
+                            placeholder='(Optional) eg: 0.5',
+                            info='Specify the learning rate weight of the mid block of U-Net.'
+                        )
+                        up_lr_weight = gr.Textbox(
+                            label='Up LR weights',
+                            placeholder='(Optional) eg: 0,0,0,0,0,0,1,1,1,1,1,1',
+                            info='Specify the learning rate weight of the up blocks of U-Net. The same as down_lr_weight.'
+                        )
+                        block_lr_zero_threshold = gr.Textbox(
+                            label='Blocks LR zero threshold',
+                            placeholder='(Optional) eg: 0.1',
+                            info='If the weight is not more than this value, the LoRA module is not created. The default is 0.'
+                        )
+                with gr.Tab(label='Blocks'):
+                    with gr.Row(visible=True):
+                        block_dims = gr.Textbox(
+                            label='Block dims',
+                            placeholder='(Optional) eg: 2,2,2,2,4,4,4,4,6,6,6,6,8,6,6,6,6,4,4,4,4,2,2,2,2',
+                            info='Specify the dim (rank) of each block. Specify 25 numbers.'
+                        )
+                        block_alphas = gr.Textbox(
+                            label='Block alphas',
+                            placeholder='(Optional) eg: 2,2,2,2,4,4,4,4,6,6,6,6,8,6,6,6,6,4,4,4,4,2,2,2,2',
+                            info='Specify the alpha of each block. Specify 25 numbers as with block_dims. If omitted, the value of network_alpha is used.'
+                        )
+                with gr.Tab(label='Conv'):
+                    with gr.Row(visible=True):
+                        conv_dims = gr.Textbox(
+                            label='Conv dims',
+                            placeholder='(Optional) eg: 2,2,2,2,4,4,4,4,6,6,6,6,8,6,6,6,6,4,4,4,4,2,2,2,2',
+                            info='Expand LoRA to Conv2d 3x3 and specify the dim (rank) of each block. Specify 25 numbers.'
+                        )
+                        conv_alphas = gr.Textbox(
+                            label='Conv alphas',
+                            placeholder='(Optional) eg: 2,2,2,2,4,4,4,4,6,6,6,6,8,6,6,6,6,4,4,4,4,2,2,2,2',
+                            info='Specify the alpha of each block when expanding LoRA to Conv2d 3x3. Specify 25 numbers. If omitted, the value of conv_alpha is used.'
+                        )
+            with gr.Row():
+                no_token_padding = gr.Checkbox(
+                    label='No token padding', value=False
+                )
+                gradient_accumulation_steps = gr.Number(
+                    label='Gradient accumulate steps', value='1'
+                )
+                weighted_captions = gr.Checkbox(
+                    label='Weighted captions', value=False, info='Enable weighted captions in the standard style (token:1.3). No commas inside parens, or shuffle/dropout may break the decoder.',
+                )
+            with gr.Row():
+                prior_loss_weight = gr.Number(
+                    label='Prior loss weight', value=1.0
+                )
+                lr_scheduler_num_cycles = gr.Textbox(
+                    label='LR number of cycles',
+                    placeholder='(Optional) For Cosine with restart and polynomial only',
+                )
+
+                lr_scheduler_power = gr.Textbox(
+                    label='LR power',
+                    placeholder='(Optional) For Cosine with restart and polynomial only',
+                )
+            (
+                # use_8bit_adam,
+                xformers,
+                full_fp16,
+                gradient_checkpointing,
+                shuffle_caption,
+                color_aug,
+                flip_aug,
+                clip_skip,
+                mem_eff_attn,
+                save_state,
+                resume,
+                max_token_length,
+                max_train_epochs,
+                max_data_loader_n_workers,
+                keep_tokens,
+                persistent_data_loader_workers,
+                bucket_no_upscale,
+                random_crop,
+                bucket_reso_steps,
+                caption_dropout_every_n_epochs,
+                caption_dropout_rate,
+                noise_offset,
+                additional_parameters,
+                vae_batch_size,
+                min_snr_gamma,
+            ) = gradio_advanced_training()
+            color_aug.change(
+                color_aug_changed,
+                inputs=[color_aug],
+                outputs=[cache_latents],
+            )
+
+        (
+            sample_every_n_steps,
+            sample_every_n_epochs,
+            sample_sampler,
+            sample_prompts,
+        ) = sample_gradio_config()
+        
+        LoRA_type.change(
+            update_LoRA_settings, inputs=[LoRA_type], outputs=[LoCon_row, kohya_advanced_lora, kohya_dylora]
+        )
+
+    with gr.Tab('Tools'):
+        gr.Markdown(
+            'This section provide Dreambooth tools to help setup your dataset...'
+        )
+        gradio_dreambooth_folder_creation_tab(
+            train_data_dir_input=train_data_dir,
+            reg_data_dir_input=reg_data_dir,
+            output_dir_input=output_dir,
+            logging_dir_input=logging_dir,
+        )
+        gradio_dataset_balancing_tab()
+        gradio_merge_lora_tab()
+        gradio_svd_merge_lora_tab()
+        gradio_resize_lora_tab()
+        gradio_verify_lora_tab()
+
+    button_run = gr.Button('Train model', variant='primary')
+
+    button_print = gr.Button('Print training command')
+
+    # Setup gradio tensorboard buttons
+    button_start_tensorboard, button_stop_tensorboard = gradio_tensorboard()
+
+    button_start_tensorboard.click(
+        start_tensorboard,
+        inputs=logging_dir,
+        show_progress=False,
+    )
+
+    button_stop_tensorboard.click(
+        stop_tensorboard,
+        show_progress=False,
+    )
+
+    settings_list = [
+        pretrained_model_name_or_path,
+        v2,
+        v_parameterization,
+        logging_dir,
+        train_data_dir,
+        reg_data_dir,
+        output_dir,
+        max_resolution,
+        learning_rate,
+        lr_scheduler,
+        lr_warmup,
+        train_batch_size,
+        epoch,
+        save_every_n_epochs,
+        mixed_precision,
+        save_precision,
+        seed,
+        num_cpu_threads_per_process,
+        cache_latents,
+        caption_extension,
+        enable_bucket,
+        gradient_checkpointing,
+        full_fp16,
+        no_token_padding,
+        stop_text_encoder_training,
+        # use_8bit_adam,
+        xformers,
+        save_model_as,
+        shuffle_caption,
+        save_state,
+        resume,
+        prior_loss_weight,
+        text_encoder_lr,
+        unet_lr,
+        network_dim,
+        lora_network_weights,
+        color_aug,
+        flip_aug,
+        clip_skip,
+        gradient_accumulation_steps,
+        mem_eff_attn,
+        output_name,
+        model_list,
+        max_token_length,
+        max_train_epochs,
+        max_data_loader_n_workers,
+        network_alpha,
+        training_comment,
+        keep_tokens,
+        lr_scheduler_num_cycles,
+        lr_scheduler_power,
+        persistent_data_loader_workers,
+        bucket_no_upscale,
+        random_crop,
+        bucket_reso_steps,
+        caption_dropout_every_n_epochs,
+        caption_dropout_rate,
+        optimizer,
+        optimizer_args,
+        noise_offset,
+        LoRA_type,
+        conv_dim,
+        conv_alpha,
+        sample_every_n_steps,
+        sample_every_n_epochs,
+        sample_sampler,
+        sample_prompts,
+        additional_parameters,
+        vae_batch_size,
+        min_snr_gamma,
+        down_lr_weight,mid_lr_weight,up_lr_weight,block_lr_zero_threshold,block_dims,block_alphas,conv_dims,conv_alphas,
+        weighted_captions, unit,
+    ]
+
+    button_open_config.click(
+        open_configuration,
+        inputs=[dummy_db_true, config_file_name] + settings_list,
+        outputs=[config_file_name] + settings_list + [LoCon_row],
+        show_progress=False,
+    )
+
+    button_load_config.click(
+        open_configuration,
+        inputs=[dummy_db_false, config_file_name] + settings_list,
+        outputs=[config_file_name] + settings_list + [LoCon_row],
+        show_progress=False,
+    )
+
+    button_save_config.click(
+        save_configuration,
+        inputs=[dummy_db_false, config_file_name] + settings_list,
+        outputs=[config_file_name],
+        show_progress=False,
+    )
+
+    button_save_as_config.click(
+        save_configuration,
+        inputs=[dummy_db_true, config_file_name] + settings_list,
+        outputs=[config_file_name],
+        show_progress=False,
+    )
+
+    button_run.click(
+        train_model,
+        inputs=[dummy_db_false] + settings_list,
+        show_progress=False,
+    )
+
+    button_print.click(
+        train_model,
+        inputs=[dummy_db_true] + settings_list,
+        show_progress=False,
+    )
+
+    return (
+        train_data_dir,
+        reg_data_dir,
+        output_dir,
+        logging_dir,
+    )
+
+
+def UI(**kwargs):
+    css = ''
+
+    if os.path.exists('./style.css'):
+        with open(os.path.join('./style.css'), 'r', encoding='utf8') as file:
+            print('Load CSS...')
+            css += file.read() + '\n'
+
+    interface = gr.Blocks(css=css)
+
+    with interface:
+        with gr.Tab('LoRA'):
+            (
+                train_data_dir_input,
+                reg_data_dir_input,
+                output_dir_input,
+                logging_dir_input,
+            ) = lora_tab()
+        with gr.Tab('Utilities'):
+            utilities_tab(
+                train_data_dir_input=train_data_dir_input,
+                reg_data_dir_input=reg_data_dir_input,
+                output_dir_input=output_dir_input,
+                logging_dir_input=logging_dir_input,
+                enable_copy_info_button=True,
+            )
+
+    # Show the interface
+    launch_kwargs = {}
+    if not kwargs.get('username', None) == '':
+        launch_kwargs['auth'] = (
+            kwargs.get('username', None),
+            kwargs.get('password', None),
+        )
+    if kwargs.get('server_port', 0) > 0:
+        launch_kwargs['server_port'] = kwargs.get('server_port', 0)
+    if kwargs.get('inbrowser', False):
+        launch_kwargs['inbrowser'] = kwargs.get('inbrowser', False)
+    if kwargs.get('listen', True):
+        launch_kwargs['server_name'] = '0.0.0.0'
+    print(launch_kwargs)
+    interface.launch(**launch_kwargs)
+
+
+if __name__ == '__main__':
+    # torch.cuda.set_per_process_memory_fraction(0.48)
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        '--username', type=str, default='', help='Username for authentication'
+    )
+    parser.add_argument(
+        '--password', type=str, default='', help='Password for authentication'
+    )
+    parser.add_argument(
+        '--server_port',
+        type=int,
+        default=0,
+        help='Port to run the server listener on',
+    )
+    parser.add_argument(
+        '--inbrowser', action='store_true', help='Open in browser'
+    )
+    parser.add_argument(
+        '--listen',
+        action='store_true',
+        help='Launch gradio with server name 0.0.0.0, allowing LAN access',
+    )
+
+    args = parser.parse_args()
+
+    UI(
+        username=args.username,
+        password=args.password,
+        inbrowser=args.inbrowser,
+        server_port=args.server_port,
+    )

requirements.txt

@@ -0,0 +1,34 @@
+accelerate==0.15.0
+albumentations==1.3.0
+altair==4.2.2
+bitsandbytes==0.35.0
+dadaptation==1.5
+diffusers[torch]==0.10.2
+easygui==0.98.3
+einops==0.6.0
+ftfy==6.1.1
+gradio==3.27.0; sys_platform != 'darwin'
+# gradio==3.19.1; sys_platform != 'darwin'
+gradio==3.23.0; sys_platform == 'darwin'
+lion-pytorch==0.0.6
+opencv-python==4.7.0.68
+pytorch-lightning==1.9.0
+safetensors==0.2.6
+tensorboard==2.10.1 ; sys_platform != 'darwin'
+tensorboard==2.12.1 ; sys_platform == 'darwin'
+tk==0.1.0
+toml==0.10.2
+transformers==4.26.0
+voluptuous==0.13.1
+# for BLIP captioning
+fairscale==0.4.13
+requests==2.28.2
+timm==0.6.12
+# tensorflow<2.11
+huggingface-hub==0.13.3; sys_platform != 'darwin'
+huggingface-hub==0.13.0; sys_platform == 'darwin'
+tensorflow==2.10.1; sys_platform != 'darwin'
+# For locon support
+lycoris_lora==0.1.4
+# for kohya_ss library
+.

setup.bat

@@ -0,0 +1,48 @@
+@echo off
+
+set PYTHON_VER=3.10.9
+
+REM Check if Python version meets the recommended version
+python --version 2>nul | findstr /b /c:"Python %PYTHON_VER%" >nul
+if errorlevel 1 (
+    echo Warning: Python version %PYTHON_VER% is recommended.
+)
+
+IF NOT EXIST venv (
+    python -m venv venv
+) ELSE (
+    echo venv folder already exists, skipping creation...
+)
+call .\venv\Scripts\activate.bat
+
+echo Do you want to uninstall previous versions of torch and associated files before installing? Usefull if you are upgrading from torch 1.12.1 to torch 2.0.0 or if you are downgrading from torch 2.0.0 to torch 1.12.1.
+echo [1] - Yes
+echo [2] - No (recommanded for most)
+set /p uninstall_choice="Enter your choice (1 or 2): "
+
+if %uninstall_choice%==1 (
+    pip uninstall -y xformers
+    pip uninstall -y torch torchvision
+)
+
+echo Please choose the version of torch you want to install:
+echo [1] - v1 (torch 1.12.1) (Recommended)
+echo [2] - v2 (torch 2.0.0) (Experimental)
+set /p choice="Enter your choice (1 or 2): "
+
+if %choice%==1 (
+    pip install torch==1.12.1+cu116 torchvision==0.13.1+cu116 --extra-index-url https://download.pytorch.org/whl/cu116
+    pip install --use-pep517 --upgrade -r requirements.txt
+    pip install -U -I --no-deps https://github.com/C43H66N12O12S2/stable-diffusion-webui/releases/download/f/xformers-0.0.14.dev0-cp310-cp310-win_amd64.whl
+) else (
+    pip install torch==2.0.0+cu118 torchvision==0.15.1+cu118 --extra-index-url https://download.pytorch.org/whl/cu118
+    pip install --use-pep517 --upgrade -r requirements.txt
+    pip install --upgrade xformers==0.0.17
+    rem pip install -U -I --no-deps https://files.pythonhosted.org/packages/d6/f7/02662286419a2652c899e2b3d1913c47723fc164b4ac06a85f769c291013/xformers-0.0.17rc482-cp310-cp310-win_amd64.whl
+)
+
+copy /y .\bitsandbytes_windows\*.dll .\venv\Lib\site-packages\bitsandbytes\
+copy /y .\bitsandbytes_windows\cextension.py .\venv\Lib\site-packages\bitsandbytes\cextension.py
+copy /y .\bitsandbytes_windows\main.py .\venv\Lib\site-packages\bitsandbytes\cuda_setup\main.py
+
+accelerate config

setup.py

@@ -0,0 +1,10 @@
+from setuptools import setup, find_packages
+import subprocess
+import os
+import sys
+
+# Call the create_user_files.py script
+script_path = os.path.join("tools", "create_user_files.py")
+subprocess.run([sys.executable, script_path])
+
+setup(name="library", version="1.0.3", packages=find_packages())

setup.sh

@@ -0,0 +1,648 @@
+#!/usr/bin/env bash
+
+# This file will be the host environment setup file for all operating systems other than base Windows.
+
+# Set the required package versions here.
+# They will be appended to the requirements.txt file in the installation directory.
+TENSORFLOW_VERSION="2.12.0"
+TENSORFLOW_MACOS_VERSION="2.12.0"
+TENSORFLOW_METAL_VERSION="0.8.0"
+
+display_help() {
+  cat <<EOF
+Kohya_SS Installation Script for POSIX operating systems.
+
+Usage:
+  # Specifies custom branch, install directory, and git repo
+  setup.sh -b dev -d /workspace/kohya_ss -g https://mycustom.repo.tld/custom_fork.git
+
+  # Same as example 1, but uses long options
+  setup.sh --branch=dev --dir=/workspace/kohya_ss --git-repo=https://mycustom.repo.tld/custom_fork.git
+
+  # Maximum verbosity, fully automated installation in a runpod environment skipping the runpod env checks
+  setup.sh -vvv --skip-space-check --runpod
+
+Options:
+  -b BRANCH, --branch=BRANCH    Select which branch of kohya to check out on new installs.
+  -d DIR, --dir=DIR             The full path you want kohya_ss installed to.
+  -g REPO, --git_repo=REPO      You can optionally provide a git repo to check out for runpod installation. Useful for custom forks.
+  -h, --help                    Show this screen.
+  -i, --interactive             Interactively configure accelerate instead of using default config file.
+  -n, --no-git-update           Do not update kohya_ss repo. No git pull or clone operations.
+  -p, --public                  Expose public URL in runpod mode. Won't have an effect in other modes.
+  -r, --runpod                  Forces a runpod installation. Useful if detection fails for any reason.
+  -s, --skip-space-check        Skip the 10Gb minimum storage space check.
+  -u, --no-gui                  Skips launching the GUI.
+  -v, --verbose                 Increase verbosity levels up to 3.
+EOF
+}
+
+# Checks to see if variable is set and non-empty.
+# This is defined first, so we can use the function for some default variable values
+env_var_exists() {
+  if [[ -n "${!1}" ]]; then
+    return 0
+  else
+    return 1
+  fi
+}
+
+# Need RUNPOD to have a default value before first access
+RUNPOD=false
+if env_var_exists RUNPOD_POD_ID || env_var_exists RUNPOD_API_KEY; then
+  RUNPOD=true
+fi
+
+# This gets the directory the script is run from so pathing can work relative to the script where needed.
+SCRIPT_DIR="$(cd -- $(dirname -- "$0") && pwd)"
+
+# Variables defined before the getopts loop, so we have sane default values.
+# Default installation locations based on OS and environment
+if [[ "$OSTYPE" == "linux-gnu"* ]]; then
+  if [ "$RUNPOD" = true ]; then
+    DIR="/workspace/kohya_ss"
+  elif [ -d "$SCRIPT_DIR/.git" ]; then
+    DIR="$SCRIPT_DIR"
+  elif [ -w "/opt" ]; then
+    DIR="/opt/kohya_ss"
+  elif env_var_exists HOME; then
+    DIR="$HOME/kohya_ss"
+  else
+    # The last fallback is simply PWD
+    DIR="$(PWD)"
+  fi
+else
+  if [ -d "$SCRIPT_DIR/.git" ]; then
+    DIR="$SCRIPT_DIR"
+  elif env_var_exists HOME; then
+    DIR="$HOME/kohya_ss"
+  else
+    # The last fallback is simply PWD
+    DIR="$(PWD)"
+  fi
+fi
+
+VERBOSITY=2    #Start counting at 2 so that any increase to this will result in a minimum of file descriptor 3.  You should leave this alone.
+MAXVERBOSITY=6 #The highest verbosity we use / allow to be displayed.  Feel free to adjust.
+
+BRANCH="master"
+GIT_REPO="https://github.com/bmaltais/kohya_ss.git"
+INTERACTIVE=false
+PUBLIC=false
+SKIP_SPACE_CHECK=false
+SKIP_GIT_UPDATE=false
+SKIP_GUI=false
+
+while getopts ":vb:d:g:inprus-:" opt; do
+  # support long options: https://stackoverflow.com/a/28466267/519360
+  if [ "$opt" = "-" ]; then # long option: reformulate OPT and OPTARG
+    opt="${OPTARG%%=*}"     # extract long option name
+    OPTARG="${OPTARG#$opt}" # extract long option argument (may be empty)
+    OPTARG="${OPTARG#=}"    # if long option argument, remove assigning `=`
+  fi
+  case $opt in
+  b | branch) BRANCH="$OPTARG" ;;
+  d | dir) DIR="$OPTARG" ;;
+  g | git-repo) GIT_REPO="$OPTARG" ;;
+  i | interactive) INTERACTIVE=true ;;
+  n | no-git-update) SKIP_GIT_UPDATE=true ;;
+  p | public) PUBLIC=true ;;
+  r | runpod) RUNPOD=true ;;
+  s | skip-space-check) SKIP_SPACE_CHECK=true ;;
+  u | no-gui) SKIP_GUI=true ;;
+  v) ((VERBOSITY = VERBOSITY + 1)) ;;
+  h) display_help && exit 0 ;;
+  *) display_help && exit 0 ;;
+  esac
+done
+shift $((OPTIND - 1))
+
+# Just in case someone puts in a relative path into $DIR,
+# we're going to get the absolute path of that.
+if [[ "$DIR" != /* ]] && [[ "$DIR" != ~* ]]; then
+  DIR="$(
+    cd "$(dirname "$DIR")" || exit 1
+    pwd
+  )/$(basename "$DIR")"
+fi
+
+for v in $( #Start counting from 3 since 1 and 2 are standards (stdout/stderr).
+  seq 3 $VERBOSITY
+); do
+  (("$v" <= "$MAXVERBOSITY")) && eval exec "$v>&2" #Don't change anything higher than the maximum verbosity allowed.
+done
+
+for v in $( #From the verbosity level one higher than requested, through the maximum;
+  seq $((VERBOSITY + 1)) $MAXVERBOSITY
+); do
+  (("$v" > "2")) && eval exec "$v>/dev/null" #Redirect these to bitbucket, provided that they don't match stdout and stderr.
+done
+
+# Example of how to use the verbosity levels.
+# printf "%s\n" "This message is seen at verbosity level 1 and above." >&3
+# printf "%s\n" "This message is seen at verbosity level 2 and above." >&4
+# printf "%s\n" "This message is seen at verbosity level 3 and above." >&5
+
+# Debug variable dump at max verbosity
+echo "BRANCH: $BRANCH
+DIR: $DIR
+GIT_REPO: $GIT_REPO
+INTERACTIVE: $INTERACTIVE
+PUBLIC: $PUBLIC
+RUNPOD: $RUNPOD
+SKIP_SPACE_CHECK: $SKIP_SPACE_CHECK
+VERBOSITY: $VERBOSITY
+Script directory is ${SCRIPT_DIR}." >&5
+
+# This must be set after the getopts loop to account for $DIR changes.
+PARENT_DIR="$(dirname "${DIR}")"
+VENV_DIR="$DIR/venv"
+
+if [ -w "$PARENT_DIR" ] && [ ! -d "$DIR" ]; then
+  echo "Creating install folder ${DIR}."
+  mkdir "$DIR"
+fi
+
+if [ ! -w "$DIR" ]; then
+  echo "We cannot write to ${DIR}."
+  echo "Please ensure the install directory is accurate and you have the correct permissions."
+  exit 1
+fi
+
+# Shared functions
+# This checks for free space on the installation drive and returns that in Gb.
+size_available() {
+  local folder
+  if [ -d "$DIR" ]; then
+    folder="$DIR"
+  elif [ -d "$PARENT_DIR" ]; then
+    folder="$PARENT_DIR"
+  elif [ -d "$(echo "$DIR" | cut -d "/" -f2)" ]; then
+    folder="$(echo "$DIR" | cut -d "/" -f2)"
+  else
+    echo "We are assuming a root drive install for space-checking purposes."
+    folder='/'
+  fi
+
+  local FREESPACEINKB
+  FREESPACEINKB="$(df -Pk "$folder" | sed 1d | grep -v used | awk '{ print $4 "\t" }')"
+  echo "Detected available space in Kb: $FREESPACEINKB" >&5
+  local FREESPACEINGB
+  FREESPACEINGB=$((FREESPACEINKB / 1024 / 1024))
+  echo "$FREESPACEINGB"
+}
+
+# The expected usage is create_symlinks symlink target_file
+create_symlinks() {
+  echo "Checking symlinks now."
+  # Next line checks for valid symlink
+  if [ -L "$1" ]; then
+    # Check if the linked file exists and points to the expected file
+    if [ -e "$1" ] && [ "$(readlink "$1")" == "$2" ]; then
+      echo "$(basename "$1") symlink looks fine. Skipping."
+    else
+      if [ -f "$2" ]; then
+        echo "Broken symlink detected. Recreating $(basename "$1")."
+        rm "$1" &&
+          ln -s "$2" "$1"
+      else
+        echo "$2 does not exist. Nothing to link."
+      fi
+    fi
+  else
+    echo "Linking $(basename "$1")."
+    ln -s "$2" "$1"
+  fi
+}
+
+install_python_dependencies() {
+  # Switch to local virtual env
+  echo "Switching to virtual Python environment."
+  if ! inDocker; then
+    if command -v python3 >/dev/null; then
+      python3 -m venv "$DIR/venv"
+    elif command -v python3.10 >/dev/null; then
+      python3.10 -m venv "$DIR/venv"
+    else
+      echo "Valid python3 or python3.10 binary not found."
+      echo "Cannot proceed with the python steps."
+      return 1
+    fi
+
+    # Activate the virtual environment
+    source "$DIR/venv/bin/activate"
+  fi
+
+  # Updating pip if there is one
+  echo "Checking for pip updates before Python operations."
+  pip install --upgrade pip >&3
+
+  echo "Installing python dependencies. This could take a few minutes as it downloads files."
+  echo "If this operation ever runs too long, you can rerun this script in verbose mode to check."
+  case "$OSTYPE" in
+  "linux-gnu"*) pip install torch==1.12.1+cu116 torchvision==0.13.1+cu116 \
+    --extra-index-url https://download.pytorch.org/whl/cu116 >&3 &&
+    pip install -U -I --no-deps \
+      https://github.com/C43H66N12O12S2/stable-diffusion-webui/releases/download/linux/xformers-0.0.14.dev0-cp310-cp310-linux_x86_64.whl >&3 ;;
+  "darwin"*) pip install torch==2.0.0 torchvision==0.15.1 \
+    -f https://download.pytorch.org/whl/cpu/torch_stable.html >&3 ;;
+  "cygwin")
+    :
+    ;;
+  "msys")
+    :
+    ;;
+  esac
+
+  if [ "$RUNPOD" = true ]; then
+    echo "Installing tenssort."
+    pip install tensorrt >&3
+  fi
+
+  # DEBUG ONLY (Update this version number to whatever PyCharm recommends)
+  # pip install pydevd-pycharm~=223.8836.43
+
+  #This will copy our requirements.txt file out and make the khoya_ss lib a dynamic location then cleanup.
+  local TEMP_REQUIREMENTS_FILE="$DIR/requirements_tmp_for_setup.txt"
+  echo "Copying $DIR/requirements.txt to $TEMP_REQUIREMENTS_FILE" >&3
+  echo "Replacing the . for lib to our DIR variable in $TEMP_REQUIREMENTS_FILE." >&3
+  awk -v dir="$DIR" '/#.*kohya_ss.*library/{print; getline; sub(/^\.$/, dir)}1' "$DIR/requirements.txt" >"$TEMP_REQUIREMENTS_FILE"
+
+  # This will check if macOS is running then determine if M1+ or Intel CPU.
+  # It will append the appropriate packages to the requirements.txt file.
+  # Other OSs won't be affected and the version variables are at the top of this file.
+  if [[ "$(uname)" == "Darwin" ]]; then
+    # Check if the processor is Apple Silicon (arm64)
+    if [[ "$(uname -m)" == "arm64" ]]; then
+      echo "tensorflow-macos==$TENSORFLOW_MACOS_VERSION" >>"$TEMP_REQUIREMENTS_FILE"
+      echo "tensorflow-metal==$TENSORFLOW_METAL_VERSION" >>"$TEMP_REQUIREMENTS_FILE"
+    # Check if the processor is Intel (x86_64)
+    elif [[ "$(uname -m)" == "x86_64" ]]; then
+      echo "tensorflow==$TENSORFLOW_VERSION" >>"$TEMP_REQUIREMENTS_FILE"
+    fi
+  fi
+
+  if [ $VERBOSITY == 2 ]; then
+    python -m pip install --quiet --use-pep517 --upgrade -r "$TEMP_REQUIREMENTS_FILE" >&3
+  else
+    python -m pip install --use-pep517 --upgrade -r "$TEMP_REQUIREMENTS_FILE" >&3
+  fi
+
+  echo "Removing the temp requirements file."
+  if [ -f "$TEMP_REQUIREMENTS_FILE" ]; then
+    rm -f "$TEMP_REQUIREMENTS_FILE"
+  fi
+
+  if [ -n "$VIRTUAL_ENV" ] && ! inDocker; then
+    if command -v deactivate >/dev/null; then
+      echo "Exiting Python virtual environment."
+      deactivate
+    else
+      echo "deactivate command not found. Could still be in the Python virtual environment."
+    fi
+  fi
+}
+
+# Attempt to non-interactively install a default accelerate config file unless specified otherwise.
+# Documentation for order of precedence locations for configuration file for automated installation:
+# https://huggingface.co/docs/accelerate/basic_tutorials/launch#custom-configurations
+configure_accelerate() {
+  echo "Source accelerate config location: $DIR/config_files/accelerate/default_config.yaml" >&3
+  if [ "$INTERACTIVE" = true ]; then
+    accelerate config
+  else
+    if env_var_exists HF_HOME; then
+      if [ ! -f "$HF_HOME/accelerate/default_config.yaml" ]; then
+        mkdir -p "$HF_HOME/accelerate/" &&
+          echo "Target accelerate config location: $HF_HOME/accelerate/default_config.yaml" >&3
+        cp "$DIR/config_files/accelerate/default_config.yaml" "$HF_HOME/accelerate/default_config.yaml" &&
+          echo "Copied accelerate config file to: $HF_HOME/accelerate/default_config.yaml"
+      fi
+    elif env_var_exists XDG_CACHE_HOME; then
+      if [ ! -f "$XDG_CACHE_HOME/huggingface/accelerate" ]; then
+        mkdir -p "$XDG_CACHE_HOME/huggingface/accelerate" &&
+          echo "Target accelerate config location: $XDG_CACHE_HOME/accelerate/default_config.yaml" >&3
+        cp "$DIR/config_files/accelerate/default_config.yaml" "$XDG_CACHE_HOME/huggingface/accelerate/default_config.yaml" &&
+          echo "Copied accelerate config file to: $XDG_CACHE_HOME/huggingface/accelerate/default_config.yaml"
+      fi
+    elif env_var_exists HOME; then
+      if [ ! -f "$HOME/.cache/huggingface/accelerate" ]; then
+        mkdir -p "$HOME/.cache/huggingface/accelerate" &&
+          echo "Target accelerate config location: $HOME/accelerate/default_config.yaml" >&3
+        cp "$DIR/config_files/accelerate/default_config.yaml" "$HOME/.cache/huggingface/accelerate/default_config.yaml" &&
+          echo "Copying accelerate config file to: $HOME/.cache/huggingface/accelerate/default_config.yaml"
+      fi
+    else
+      echo "Could not place the accelerate configuration file. Please configure manually."
+      sleep 2
+      accelerate config
+    fi
+  fi
+}
+
+# Offer a warning and opportunity to cancel the installation if < 10Gb of Free Space detected
+check_storage_space() {
+  if [ "$SKIP_SPACE_CHECK" = false ]; then
+    if [ "$(size_available)" -lt 10 ]; then
+      echo "You have less than 10Gb of free space. This installation may fail."
+      MSGTIMEOUT=10 # In seconds
+      MESSAGE="Continuing in..."
+      echo "Press control-c to cancel the installation."
+      for ((i = MSGTIMEOUT; i >= 0; i--)); do
+        printf "\r${MESSAGE} %ss. " "${i}"
+        sleep 1
+      done
+    fi
+  fi
+}
+
+isContainerOrPod() {
+  local cgroup=/proc/1/cgroup
+  test -f $cgroup && (grep -qE ':cpuset:/(docker|kubepods)' $cgroup || grep -q ':/docker/' $cgroup)
+}
+
+isDockerBuildkit() {
+  local cgroup=/proc/1/cgroup
+  test -f $cgroup && grep -q ':cpuset:/docker/buildkit' $cgroup
+}
+
+isDockerContainer() {
+  [ -e /.dockerenv ]
+}
+
+inDocker() {
+  if isContainerOrPod || isDockerBuildkit || isDockerContainer; then
+    return 0
+  else
+    return 1
+  fi
+}
+
+# These are the git operations that will run to update or clone the repo
+update_kohya_ss() {
+  if [ "$SKIP_GIT_UPDATE" = false ]; then
+    if command -v git >/dev/null; then
+      # First, we make sure there are no changes that need to be made in git, so no work is lost.
+      if [ "$(git -C "$DIR" status --porcelain=v1 2>/dev/null | wc -l)" -gt 0 ] &&
+        echo "These files need to be committed or discarded: " >&4 &&
+        git -C "$DIR" status >&4; then
+        echo "There are changes that need to be committed or discarded in the repo in $DIR."
+        echo "Commit those changes or run this script with -n to skip git operations entirely."
+        exit 1
+      fi
+
+      echo "Attempting to clone $GIT_REPO."
+      if [ ! -d "$DIR/.git" ]; then
+        echo "Cloning and switching to $GIT_REPO:$BRANCH" >&4
+        git -C "$PARENT_DIR" clone -b "$BRANCH" "$GIT_REPO" "$(basename "$DIR")" >&3
+        git -C "$DIR" switch "$BRANCH" >&4
+      else
+        echo "git repo detected. Attempting to update repository instead."
+        echo "Updating: $GIT_REPO"
+        git -C "$DIR" pull "$GIT_REPO" "$BRANCH" >&3
+        if ! git -C "$DIR" switch "$BRANCH" >&4; then
+          echo "Branch $BRANCH did not exist. Creating it." >&4
+          git -C "$DIR" switch -c "$BRANCH" >&4
+        fi
+      fi
+    else
+      echo "You need to install git."
+      echo "Rerun this after installing git or run this script with -n to skip the git operations."
+    fi
+  else
+    echo "Skipping git operations."
+  fi
+}
+
+# Start OS-specific detection and work
+if [[ "$OSTYPE" == "linux-gnu"* ]]; then
+  # Check if root or sudo
+  root=false
+  if [ "$EUID" = 0 ]; then
+    root=true
+  elif command -v id >/dev/null && [ "$(id -u)" = 0 ]; then
+    root=true
+  elif [ "$UID" = 0 ]; then
+    root=true
+  fi
+
+  get_distro_name() {
+    local line
+    if [ -f /etc/os-release ]; then
+      # We search for the line starting with ID=
+      # Then we remove the ID= prefix to get the name itself
+      line="$(grep -Ei '^ID=' /etc/os-release)"
+      echo "Raw detected os-release distro line: $line" >&5
+      line=${line##*=}
+      echo "$line"
+      return 0
+    elif command -v python >/dev/null; then
+      line="$(python -mplatform)"
+      echo "$line"
+      return 0
+    elif command -v python3 >/dev/null; then
+      line="$(python3 -mplatform)"
+      echo "$line"
+      return 0
+    else
+      line="None"
+      echo "$line"
+      return 1
+    fi
+  }
+
+  # We search for the line starting with ID_LIKE=
+  # Then we remove the ID_LIKE= prefix to get the name itself
+  # This is the "type" of distro. For example, Ubuntu returns "debian".
+  get_distro_family() {
+    local line
+    if [ -f /etc/os-release ]; then
+      if grep -Eiq '^ID_LIKE=' /etc/os-release >/dev/null; then
+        line="$(grep -Ei '^ID_LIKE=' /etc/os-release)"
+        echo "Raw detected os-release distro family line: $line" >&5
+        line=${line##*=}
+        echo "$line"
+        return 0
+      else
+        line="None"
+        echo "$line"
+        return 1
+      fi
+    else
+      line="None"
+      echo "$line"
+      return 1
+    fi
+  }
+
+  check_storage_space
+  update_kohya_ss
+
+  distro=get_distro_name
+  family=get_distro_family
+  echo "Raw detected distro string: $distro" >&4
+  echo "Raw detected distro family string: $family" >&4
+
+  echo "Installing Python TK if not found on the system."
+  if "$distro" | grep -qi "Ubuntu" || "$family" | grep -qi "Ubuntu"; then
+    echo "Ubuntu detected."
+    if [ $(dpkg-query -W -f='${Status}' python3-tk 2>/dev/null | grep -c "ok installed") = 0 ]; then
+      if [ "$root" = true ]; then
+        apt update -y >&3 && apt install -y python3-tk >&3
+      else
+        echo "This script needs to be run as root or via sudo to install packages."
+        exit 1
+      fi
+    else
+      echo "Python TK found! Skipping install!"
+    fi
+  elif "$distro" | grep -Eqi "Fedora|CentOS|Redhat"; then
+    echo "Redhat or Redhat base detected."
+    if ! rpm -qa | grep -qi python3-tkinter; then
+      if [ "$root" = true ]; then
+        dnf install python3-tkinter -y >&3
+      else
+        echo "This script needs to be run as root or via sudo to install packages."
+        exit 1
+      fi
+    fi
+  elif "$distro" | grep -Eqi "arch" || "$family" | grep -qi "arch"; then
+    echo "Arch Linux or Arch base detected."
+    if ! pacman -Qi tk >/dev/null; then
+      if [ "$root" = true ]; then
+        pacman --noconfirm -S tk >&3
+      else
+        echo "This script needs to be run as root or via sudo to install packages."
+        exit 1
+      fi
+    fi
+  elif "$distro" | grep -Eqi "opensuse" || "$family" | grep -qi "opensuse"; then
+    echo "OpenSUSE detected."
+    if ! rpm -qa | grep -qi python-tk; then
+      if [ "$root" = true ]; then
+        zypper install -y python-tk >&3
+      else
+        echo "This script needs to be run as root or via sudo to install packages."
+        exit 1
+      fi
+    fi
+  elif [ "$distro" = "None" ] || [ "$family" = "None" ]; then
+    if [ "$distro" = "None" ]; then
+      echo "We could not detect your distribution of Linux. Please file a bug report on github with the contents of your /etc/os-release file."
+    fi
+
+    if [ "$family" = "None" ]; then
+      echo "We could not detect the family of your Linux distribution. Please file a bug report on github with the contents of your /etc/os-release file."
+    fi
+  fi
+
+  install_python_dependencies
+
+  # We need just a little bit more setup for non-interactive environments
+  if [ "$RUNPOD" = true ]; then
+    if inDocker; then
+      # We get the site-packages from python itself, then cut the string, so no other code changes required.
+      VENV_DIR=$(python -c "import site; print(site.getsitepackages()[0])")
+      VENV_DIR="${VENV_DIR%/lib/python3.10/site-packages}"
+    fi
+
+    # Symlink paths
+    libnvinfer_plugin_symlink="$VENV_DIR/lib/python3.10/site-packages/tensorrt/libnvinfer_plugin.so.7"
+    libnvinfer_symlink="$VENV_DIR/lib/python3.10/site-packages/tensorrt/libnvinfer.so.7"
+    libcudart_symlink="$VENV_DIR/lib/python3.10/site-packages/nvidia/cuda_runtime/lib/libcudart.so.11.0"
+
+    #Target file paths
+    libnvinfer_plugin_target="$VENV_DIR/lib/python3.10/site-packages/tensorrt/libnvinfer_plugin.so.8"
+    libnvinfer_target="$VENV_DIR/lib/python3.10/site-packages/tensorrt/libnvinfer.so.8"
+    libcudart_target="$VENV_DIR/lib/python3.10/site-packages/nvidia/cuda_runtime/lib/libcudart.so.12"
+
+    echo "Checking symlinks now."
+    create_symlinks "$libnvinfer_plugin_symlink" "$libnvinfer_plugin_target"
+    create_symlinks "$libnvinfer_symlink" "$libnvinfer_target"
+    create_symlinks "$libcudart_symlink" "$libcudart_target"
+
+    if [ -d "${VENV_DIR}/lib/python3.10/site-packages/tensorrt/" ]; then
+      export LD_LIBRARY_PATH="${LD_LIBRARY_PATH}:${VENV_DIR}/lib/python3.10/site-packages/tensorrt/"
+    else
+      echo "${VENV_DIR}/lib/python3.10/site-packages/tensorrt/ not found; not linking library."
+    fi
+
+    if [ -d "${VENV_DIR}/lib/python3.10/site-packages/tensorrt/" ]; then
+      export LD_LIBRARY_PATH="${LD_LIBRARY_PATH}:${VENV_DIR}/lib/python3.10/site-packages/nvidia/cuda_runtime/lib/"
+    else
+      echo "${VENV_DIR}/lib/python3.10/site-packages/nvidia/cuda_runtime/lib/ not found; not linking library."
+    fi
+
+    configure_accelerate
+
+    # This is a non-interactive environment, so just directly call gui.sh after all setup steps are complete.
+    if [ "$SKIP_GUI" = false ]; then
+      if command -v bash >/dev/null; then
+        if [ "$PUBLIC" = false ]; then
+          bash "$DIR"/gui.sh
+          exit 0
+        else
+          bash "$DIR"/gui.sh --share
+          exit 0
+        fi
+      else
+        # This shouldn't happen, but we're going to try to help.
+        if [ "$PUBLIC" = false ]; then
+          sh "$DIR"/gui.sh
+          exit 0
+        else
+          sh "$DIR"/gui.sh --share
+          exit 0
+        fi
+      fi
+    fi
+  fi
+
+  echo -e "Setup finished! Run \e[0;92m./gui.sh\e[0m to start."
+  echo "Please note if you'd like to expose your public server you need to run ./gui.sh --share"
+elif [[ "$OSTYPE" == "darwin"* ]]; then
+  # The initial setup script to prep the environment on macOS
+  # xformers has been omitted as that is for Nvidia GPUs only
+
+  if ! command -v brew >/dev/null; then
+    echo "Please install homebrew first. This is a requirement for the remaining setup."
+    echo "You can find that here: https://brew.sh"
+    #shellcheck disable=SC2016
+    echo 'The "brew" command should be in $PATH to be detected.'
+    exit 1
+  fi
+
+  check_storage_space
+
+  # Install base python packages
+  echo "Installing Python 3.10 if not found."
+  if ! brew ls --versions python@3.10 >/dev/null; then
+    echo "Installing Python 3.10."
+    brew install python@3.10 >&3
+  else
+    echo "Python 3.10 found!"
+  fi
+  echo "Installing Python-TK 3.10 if not found."
+  if ! brew ls --versions python-tk@3.10 >/dev/null; then
+    echo "Installing Python TK 3.10."
+    brew install python-tk@3.10 >&3
+  else
+    echo "Python Tkinter 3.10 found!"
+  fi
+
+  update_kohya_ss
+
+  if ! install_python_dependencies; then
+    echo "You may need to install Python. The command for this is brew install python@3.10."
+  fi
+
+  configure_accelerate
+  echo -e "Setup finished! Run ./gui.sh to start."
+elif [[ "$OSTYPE" == "cygwin" ]]; then
+  # Cygwin is a standalone suite of Linux utilities on Windows
+  echo "This hasn't been validated on cygwin yet."
+elif [[ "$OSTYPE" == "msys" ]]; then
+  # MinGW has the msys environment which is a standalone suite of Linux utilities on Windows
+  # "git bash" on Windows may also be detected as msys.
+  echo "This hasn't been validated in msys (mingw) on Windows yet."
+fi

style.css

@@ -0,0 +1,21 @@
+#open_folder_small{
+    height: auto;
+    min-width: auto;
+    flex-grow: 0;
+    padding-left: 0.25em;
+    padding-right: 0.25em;
+}
+
+#open_folder{
+    height: auto;
+    flex-grow: 0;
+    padding-left: 0.25em;
+    padding-right: 0.25em;
+}
+
+#number_input{
+    min-width: min-content;
+    flex-grow: 0.3;
+    padding-left: 0.75em;
+    padding-right: 0.75em;
+}

textual_inversion_gui.py

@@ -0,0 +1,1014 @@
+# v1: initial release
+# v2: add open and save folder icons
+# v3: Add new Utilities tab for Dreambooth folder preparation
+# v3.1: Adding captionning of images to utilities
+
+import gradio as gr
+import json
+import math
+import os
+import subprocess
+import pathlib
+import argparse
+from library.common_gui import (
+    get_folder_path,
+    remove_doublequote,
+    get_file_path,
+    get_any_file_path,
+    get_saveasfile_path,
+    color_aug_changed,
+    save_inference_file,
+    gradio_advanced_training,
+    run_cmd_advanced_training,
+    run_cmd_training,
+    gradio_training,
+    gradio_config,
+    gradio_source_model,
+    # set_legacy_8bitadam,
+    update_my_data,
+    check_if_model_exist,
+)
+from library.tensorboard_gui import (
+    gradio_tensorboard,
+    start_tensorboard,
+    stop_tensorboard,
+)
+from library.dreambooth_folder_creation_gui import (
+    gradio_dreambooth_folder_creation_tab,
+)
+from library.utilities import utilities_tab
+from library.sampler_gui import sample_gradio_config, run_cmd_sample
+from easygui import msgbox
+
+folder_symbol = '\U0001f4c2'  # 📂
+refresh_symbol = '\U0001f504'  # 🔄
+save_style_symbol = '\U0001f4be'  # 💾
+document_symbol = '\U0001F4C4'   # 📄
+
+
+def save_configuration(
+    save_as,
+    file_path,
+    pretrained_model_name_or_path,
+    v2,
+    v_parameterization,
+    logging_dir,
+    train_data_dir,
+    reg_data_dir,
+    output_dir,
+    max_resolution,
+    learning_rate,
+    lr_scheduler,
+    lr_warmup,
+    train_batch_size,
+    epoch,
+    save_every_n_epochs,
+    mixed_precision,
+    save_precision,
+    seed,
+    num_cpu_threads_per_process,
+    cache_latents,
+    caption_extension,
+    enable_bucket,
+    gradient_checkpointing,
+    full_fp16,
+    no_token_padding,
+    stop_text_encoder_training,
+    # use_8bit_adam,
+    xformers,
+    save_model_as,
+    shuffle_caption,
+    save_state,
+    resume,
+    prior_loss_weight,
+    color_aug,
+    flip_aug,
+    clip_skip,
+    vae,
+    output_name,
+    max_token_length,
+    max_train_epochs,
+    max_data_loader_n_workers,
+    mem_eff_attn,
+    gradient_accumulation_steps,
+    model_list,
+    token_string,
+    init_word,
+    num_vectors_per_token,
+    max_train_steps,
+    weights,
+    template,
+    keep_tokens,
+    persistent_data_loader_workers,
+    bucket_no_upscale,
+    random_crop,
+    bucket_reso_steps,
+    caption_dropout_every_n_epochs,
+    caption_dropout_rate,
+    optimizer,
+    optimizer_args,
+    noise_offset,
+    sample_every_n_steps,
+    sample_every_n_epochs,
+    sample_sampler,
+    sample_prompts,
+    additional_parameters,
+    vae_batch_size,
+    min_snr_gamma,
+):
+    # Get list of function parameters and values
+    parameters = list(locals().items())
+
+    original_file_path = file_path
+
+    save_as_bool = True if save_as.get('label') == 'True' else False
+
+    if save_as_bool:
+        print('Save as...')
+        file_path = get_saveasfile_path(file_path)
+    else:
+        print('Save...')
+        if file_path == None or file_path == '':
+            file_path = get_saveasfile_path(file_path)
+
+    # print(file_path)
+
+    if file_path == None or file_path == '':
+        return original_file_path  # In case a file_path was provided and the user decide to cancel the open action
+
+    # Return the values of the variables as a dictionary
+    variables = {
+        name: value
+        for name, value in parameters  # locals().items()
+        if name
+        not in [
+            'file_path',
+            'save_as',
+        ]
+    }
+
+    # Extract the destination directory from the file path
+    destination_directory = os.path.dirname(file_path)
+
+    # Create the destination directory if it doesn't exist
+    if not os.path.exists(destination_directory):
+        os.makedirs(destination_directory)
+
+    # Save the data to the selected file
+    with open(file_path, 'w') as file:
+        json.dump(variables, file, indent=2)
+
+    return file_path
+
+
+def open_configuration(
+    ask_for_file,
+    file_path,
+    pretrained_model_name_or_path,
+    v2,
+    v_parameterization,
+    logging_dir,
+    train_data_dir,
+    reg_data_dir,
+    output_dir,
+    max_resolution,
+    learning_rate,
+    lr_scheduler,
+    lr_warmup,
+    train_batch_size,
+    epoch,
+    save_every_n_epochs,
+    mixed_precision,
+    save_precision,
+    seed,
+    num_cpu_threads_per_process,
+    cache_latents,
+    caption_extension,
+    enable_bucket,
+    gradient_checkpointing,
+    full_fp16,
+    no_token_padding,
+    stop_text_encoder_training,
+    # use_8bit_adam,
+    xformers,
+    save_model_as,
+    shuffle_caption,
+    save_state,
+    resume,
+    prior_loss_weight,
+    color_aug,
+    flip_aug,
+    clip_skip,
+    vae,
+    output_name,
+    max_token_length,
+    max_train_epochs,
+    max_data_loader_n_workers,
+    mem_eff_attn,
+    gradient_accumulation_steps,
+    model_list,
+    token_string,
+    init_word,
+    num_vectors_per_token,
+    max_train_steps,
+    weights,
+    template,
+    keep_tokens,
+    persistent_data_loader_workers,
+    bucket_no_upscale,
+    random_crop,
+    bucket_reso_steps,
+    caption_dropout_every_n_epochs,
+    caption_dropout_rate,
+    optimizer,
+    optimizer_args,
+    noise_offset,
+    sample_every_n_steps,
+    sample_every_n_epochs,
+    sample_sampler,
+    sample_prompts,
+    additional_parameters,
+    vae_batch_size,
+    min_snr_gamma,
+):
+    # Get list of function parameters and values
+    parameters = list(locals().items())
+
+    ask_for_file = True if ask_for_file.get('label') == 'True' else False
+
+    original_file_path = file_path
+
+    if ask_for_file:
+        file_path = get_file_path(file_path)
+
+    if not file_path == '' and not file_path == None:
+        # load variables from JSON file
+        with open(file_path, 'r') as f:
+            my_data = json.load(f)
+            print('Loading config...')
+            # Update values to fix deprecated use_8bit_adam checkbox and set appropriate optimizer if it is set to True
+            my_data = update_my_data(my_data)
+    else:
+        file_path = original_file_path  # In case a file_path was provided and the user decide to cancel the open action
+        my_data = {}
+
+    values = [file_path]
+    for key, value in parameters:
+        # Set the value in the dictionary to the corresponding value in `my_data`, or the default value if not found
+        if not key in ['ask_for_file', 'file_path']:
+            values.append(my_data.get(key, value))
+    return tuple(values)
+
+
+def train_model(
+    pretrained_model_name_or_path,
+    v2,
+    v_parameterization,
+    logging_dir,
+    train_data_dir,
+    reg_data_dir,
+    output_dir,
+    max_resolution,
+    learning_rate,
+    lr_scheduler,
+    lr_warmup,
+    train_batch_size,
+    epoch,
+    save_every_n_epochs,
+    mixed_precision,
+    save_precision,
+    seed,
+    num_cpu_threads_per_process,
+    cache_latents,
+    caption_extension,
+    enable_bucket,
+    gradient_checkpointing,
+    full_fp16,
+    no_token_padding,
+    stop_text_encoder_training_pct,
+    # use_8bit_adam,
+    xformers,
+    save_model_as,
+    shuffle_caption,
+    save_state,
+    resume,
+    prior_loss_weight,
+    color_aug,
+    flip_aug,
+    clip_skip,
+    vae,
+    output_name,
+    max_token_length,
+    max_train_epochs,
+    max_data_loader_n_workers,
+    mem_eff_attn,
+    gradient_accumulation_steps,
+    model_list,  # Keep this. Yes, it is unused here but required given the common list used
+    token_string,
+    init_word,
+    num_vectors_per_token,
+    max_train_steps,
+    weights,
+    template,
+    keep_tokens,
+    persistent_data_loader_workers,
+    bucket_no_upscale,
+    random_crop,
+    bucket_reso_steps,
+    caption_dropout_every_n_epochs,
+    caption_dropout_rate,
+    optimizer,
+    optimizer_args,
+    noise_offset,
+    sample_every_n_steps,
+    sample_every_n_epochs,
+    sample_sampler,
+    sample_prompts,
+    additional_parameters,
+    vae_batch_size,
+    min_snr_gamma,
+):
+    if pretrained_model_name_or_path == '':
+        msgbox('Source model information is missing')
+        return
+
+    if train_data_dir == '':
+        msgbox('Image folder path is missing')
+        return
+
+    if not os.path.exists(train_data_dir):
+        msgbox('Image folder does not exist')
+        return
+
+    if reg_data_dir != '':
+        if not os.path.exists(reg_data_dir):
+            msgbox('Regularisation folder does not exist')
+            return
+
+    if output_dir == '':
+        msgbox('Output folder path is missing')
+        return
+
+    if token_string == '':
+        msgbox('Token string is missing')
+        return
+
+    if init_word == '':
+        msgbox('Init word is missing')
+        return
+
+    if not os.path.exists(output_dir):
+        os.makedirs(output_dir)
+
+    if check_if_model_exist(output_name, output_dir, save_model_as):
+        return
+    
+    if optimizer == 'Adafactor' and lr_warmup != '0':
+        msgbox("Warning: lr_scheduler is set to 'Adafactor', so 'LR warmup (% of steps)' will be considered 0.", title="Warning")
+        lr_warmup = '0'
+
+    # Get a list of all subfolders in train_data_dir
+    subfolders = [
+        f
+        for f in os.listdir(train_data_dir)
+        if os.path.isdir(os.path.join(train_data_dir, f))
+    ]
+
+    total_steps = 0
+
+    # Loop through each subfolder and extract the number of repeats
+    for folder in subfolders:
+        # Extract the number of repeats from the folder name
+        repeats = int(folder.split('_')[0])
+
+        # Count the number of images in the folder
+        num_images = len(
+            [
+                f
+                for f, lower_f in (
+                    (file, file.lower())
+                    for file in os.listdir(
+                        os.path.join(train_data_dir, folder)
+                    )
+                )
+                if lower_f.endswith(('.jpg', '.jpeg', '.png', '.webp'))
+            ]
+        )
+
+        # Calculate the total number of steps for this folder
+        steps = repeats * num_images
+        total_steps += steps
+
+        # Print the result
+        print(f'Folder {folder}: {steps} steps')
+
+    # Print the result
+    # print(f"{total_steps} total steps")
+
+    if reg_data_dir == '':
+        reg_factor = 1
+    else:
+        print(
+            'Regularisation images are used... Will double the number of steps required...'
+        )
+        reg_factor = 2
+
+    # calculate max_train_steps
+    if max_train_steps == '':
+        max_train_steps = int(
+            math.ceil(
+                float(total_steps)
+                / int(train_batch_size)
+                * int(epoch)
+                * int(reg_factor)
+            )
+        )
+    else:
+        max_train_steps = int(max_train_steps)
+
+    print(f'max_train_steps = {max_train_steps}')
+
+    # calculate stop encoder training
+    if stop_text_encoder_training_pct == None:
+        stop_text_encoder_training = 0
+    else:
+        stop_text_encoder_training = math.ceil(
+            float(max_train_steps) / 100 * int(stop_text_encoder_training_pct)
+        )
+    print(f'stop_text_encoder_training = {stop_text_encoder_training}')
+
+    lr_warmup_steps = round(float(int(lr_warmup) * int(max_train_steps) / 100))
+    print(f'lr_warmup_steps = {lr_warmup_steps}')
+
+    run_cmd = f'accelerate launch --num_cpu_threads_per_process={num_cpu_threads_per_process} "train_textual_inversion.py"'
+    if v2:
+        run_cmd += ' --v2'
+    if v_parameterization:
+        run_cmd += ' --v_parameterization'
+    if enable_bucket:
+        run_cmd += ' --enable_bucket'
+    if no_token_padding:
+        run_cmd += ' --no_token_padding'
+    run_cmd += (
+        f' --pretrained_model_name_or_path="{pretrained_model_name_or_path}"'
+    )
+    run_cmd += f' --train_data_dir="{train_data_dir}"'
+    if len(reg_data_dir):
+        run_cmd += f' --reg_data_dir="{reg_data_dir}"'
+    run_cmd += f' --resolution={max_resolution}'
+    run_cmd += f' --output_dir="{output_dir}"'
+    run_cmd += f' --logging_dir="{logging_dir}"'
+    if not stop_text_encoder_training == 0:
+        run_cmd += (
+            f' --stop_text_encoder_training={stop_text_encoder_training}'
+        )
+    if not save_model_as == 'same as source model':
+        run_cmd += f' --save_model_as={save_model_as}'
+    # if not resume == '':
+    #     run_cmd += f' --resume={resume}'
+    if not float(prior_loss_weight) == 1.0:
+        run_cmd += f' --prior_loss_weight={prior_loss_weight}'
+    if not vae == '':
+        run_cmd += f' --vae="{vae}"'
+    if not output_name == '':
+        run_cmd += f' --output_name="{output_name}"'
+    if int(max_token_length) > 75:
+        run_cmd += f' --max_token_length={max_token_length}'
+    if not max_train_epochs == '':
+        run_cmd += f' --max_train_epochs="{max_train_epochs}"'
+    if not max_data_loader_n_workers == '':
+        run_cmd += (
+            f' --max_data_loader_n_workers="{max_data_loader_n_workers}"'
+        )
+    if int(gradient_accumulation_steps) > 1:
+        run_cmd += f' --gradient_accumulation_steps={int(gradient_accumulation_steps)}'
+
+    run_cmd += run_cmd_training(
+        learning_rate=learning_rate,
+        lr_scheduler=lr_scheduler,
+        lr_warmup_steps=lr_warmup_steps,
+        train_batch_size=train_batch_size,
+        max_train_steps=max_train_steps,
+        save_every_n_epochs=save_every_n_epochs,
+        mixed_precision=mixed_precision,
+        save_precision=save_precision,
+        seed=seed,
+        caption_extension=caption_extension,
+        cache_latents=cache_latents,
+        optimizer=optimizer,
+        optimizer_args=optimizer_args,
+    )
+
+    run_cmd += run_cmd_advanced_training(
+        max_train_epochs=max_train_epochs,
+        max_data_loader_n_workers=max_data_loader_n_workers,
+        max_token_length=max_token_length,
+        resume=resume,
+        save_state=save_state,
+        mem_eff_attn=mem_eff_attn,
+        clip_skip=clip_skip,
+        flip_aug=flip_aug,
+        color_aug=color_aug,
+        shuffle_caption=shuffle_caption,
+        gradient_checkpointing=gradient_checkpointing,
+        full_fp16=full_fp16,
+        xformers=xformers,
+        # use_8bit_adam=use_8bit_adam,
+        keep_tokens=keep_tokens,
+        persistent_data_loader_workers=persistent_data_loader_workers,
+        bucket_no_upscale=bucket_no_upscale,
+        random_crop=random_crop,
+        bucket_reso_steps=bucket_reso_steps,
+        caption_dropout_every_n_epochs=caption_dropout_every_n_epochs,
+        caption_dropout_rate=caption_dropout_rate,
+        noise_offset=noise_offset,
+        additional_parameters=additional_parameters,
+        vae_batch_size=vae_batch_size,
+        min_snr_gamma=min_snr_gamma,
+    )
+    run_cmd += f' --token_string="{token_string}"'
+    run_cmd += f' --init_word="{init_word}"'
+    run_cmd += f' --num_vectors_per_token={num_vectors_per_token}'
+    if not weights == '':
+        run_cmd += f' --weights="{weights}"'
+    if template == 'object template':
+        run_cmd += f' --use_object_template'
+    elif template == 'style template':
+        run_cmd += f' --use_style_template'
+
+    run_cmd += run_cmd_sample(
+        sample_every_n_steps,
+        sample_every_n_epochs,
+        sample_sampler,
+        sample_prompts,
+        output_dir,
+    )
+
+    print(run_cmd)
+
+    # Run the command
+    if os.name == 'posix':
+        os.system(run_cmd)
+    else:
+        subprocess.run(run_cmd)
+
+    # check if output_dir/last is a folder... therefore it is a diffuser model
+    last_dir = pathlib.Path(f'{output_dir}/{output_name}')
+
+    if not last_dir.is_dir():
+        # Copy inference model for v2 if required
+        save_inference_file(output_dir, v2, v_parameterization, output_name)
+
+
+def ti_tab(
+    train_data_dir=gr.Textbox(),
+    reg_data_dir=gr.Textbox(),
+    output_dir=gr.Textbox(),
+    logging_dir=gr.Textbox(),
+):
+    dummy_db_true = gr.Label(value=True, visible=False)
+    dummy_db_false = gr.Label(value=False, visible=False)
+    gr.Markdown('Train a TI using kohya textual inversion python code...')
+    (
+        button_open_config,
+        button_save_config,
+        button_save_as_config,
+        config_file_name,
+        button_load_config,
+    ) = gradio_config()
+
+    (
+        pretrained_model_name_or_path,
+        v2,
+        v_parameterization,
+        save_model_as,
+        model_list,
+    ) = gradio_source_model(
+        save_model_as_choices=[
+            'ckpt',
+            'safetensors',
+        ]
+    )
+
+    with gr.Tab('Folders'):
+        with gr.Row():
+            train_data_dir = gr.Textbox(
+                label='Image folder',
+                placeholder='Folder where the training folders containing the images are located',
+            )
+            train_data_dir_input_folder = gr.Button(
+                '📂', elem_id='open_folder_small'
+            )
+            train_data_dir_input_folder.click(
+                get_folder_path,
+                outputs=train_data_dir,
+                show_progress=False,
+            )
+            reg_data_dir = gr.Textbox(
+                label='Regularisation folder',
+                placeholder='(Optional) Folder where where the regularization folders containing the images are located',
+            )
+            reg_data_dir_input_folder = gr.Button(
+                '📂', elem_id='open_folder_small'
+            )
+            reg_data_dir_input_folder.click(
+                get_folder_path,
+                outputs=reg_data_dir,
+                show_progress=False,
+            )
+        with gr.Row():
+            output_dir = gr.Textbox(
+                label='Model output folder',
+                placeholder='Folder to output trained model',
+            )
+            output_dir_input_folder = gr.Button(
+                '📂', elem_id='open_folder_small'
+            )
+            output_dir_input_folder.click(
+                get_folder_path,
+                outputs=output_dir,
+                show_progress=False,
+            )
+            logging_dir = gr.Textbox(
+                label='Logging folder',
+                placeholder='Optional: enable logging and output TensorBoard log to this folder',
+            )
+            logging_dir_input_folder = gr.Button(
+                '📂', elem_id='open_folder_small'
+            )
+            logging_dir_input_folder.click(
+                get_folder_path,
+                outputs=logging_dir,
+                show_progress=False,
+            )
+        with gr.Row():
+            output_name = gr.Textbox(
+                label='Model output name',
+                placeholder='Name of the model to output',
+                value='last',
+                interactive=True,
+            )
+        train_data_dir.change(
+            remove_doublequote,
+            inputs=[train_data_dir],
+            outputs=[train_data_dir],
+        )
+        reg_data_dir.change(
+            remove_doublequote,
+            inputs=[reg_data_dir],
+            outputs=[reg_data_dir],
+        )
+        output_dir.change(
+            remove_doublequote,
+            inputs=[output_dir],
+            outputs=[output_dir],
+        )
+        logging_dir.change(
+            remove_doublequote,
+            inputs=[logging_dir],
+            outputs=[logging_dir],
+        )
+    with gr.Tab('Training parameters'):
+        with gr.Row():
+            weights = gr.Textbox(
+                label='Resume TI training',
+                placeholder='(Optional) Path to existing TI embeding file to keep training',
+            )
+            weights_file_input = gr.Button('📂', elem_id='open_folder_small')
+            weights_file_input.click(
+                get_file_path,
+                outputs=weights,
+                show_progress=False,
+            )
+        with gr.Row():
+            token_string = gr.Textbox(
+                label='Token string',
+                placeholder='eg: cat',
+            )
+            init_word = gr.Textbox(
+                label='Init word',
+                value='*',
+            )
+            num_vectors_per_token = gr.Slider(
+                minimum=1,
+                maximum=75,
+                value=1,
+                step=1,
+                label='Vectors',
+            )
+            max_train_steps = gr.Textbox(
+                label='Max train steps',
+                placeholder='(Optional) Maximum number of steps',
+            )
+            template = gr.Dropdown(
+                label='Template',
+                choices=[
+                    'caption',
+                    'object template',
+                    'style template',
+                ],
+                value='caption',
+            )
+        (
+            learning_rate,
+            lr_scheduler,
+            lr_warmup,
+            train_batch_size,
+            epoch,
+            save_every_n_epochs,
+            mixed_precision,
+            save_precision,
+            num_cpu_threads_per_process,
+            seed,
+            caption_extension,
+            cache_latents,
+            optimizer,
+            optimizer_args,
+        ) = gradio_training(
+            learning_rate_value='1e-5',
+            lr_scheduler_value='cosine',
+            lr_warmup_value='10',
+        )
+        with gr.Row():
+            max_resolution = gr.Textbox(
+                label='Max resolution',
+                value='512,512',
+                placeholder='512,512',
+            )
+            stop_text_encoder_training = gr.Slider(
+                minimum=0,
+                maximum=100,
+                value=0,
+                step=1,
+                label='Stop text encoder training',
+            )
+            enable_bucket = gr.Checkbox(label='Enable buckets', value=True)
+        with gr.Accordion('Advanced Configuration', open=False):
+            with gr.Row():
+                no_token_padding = gr.Checkbox(
+                    label='No token padding', value=False
+                )
+                gradient_accumulation_steps = gr.Number(
+                    label='Gradient accumulate steps', value='1'
+                )
+            with gr.Row():
+                prior_loss_weight = gr.Number(
+                    label='Prior loss weight', value=1.0
+                )
+                vae = gr.Textbox(
+                    label='VAE',
+                    placeholder='(Optiona) path to checkpoint of vae to replace for training',
+                )
+                vae_button = gr.Button('📂', elem_id='open_folder_small')
+                vae_button.click(
+                    get_any_file_path,
+                    outputs=vae,
+                    show_progress=False,
+                )
+            (
+                # use_8bit_adam,
+                xformers,
+                full_fp16,
+                gradient_checkpointing,
+                shuffle_caption,
+                color_aug,
+                flip_aug,
+                clip_skip,
+                mem_eff_attn,
+                save_state,
+                resume,
+                max_token_length,
+                max_train_epochs,
+                max_data_loader_n_workers,
+                keep_tokens,
+                persistent_data_loader_workers,
+                bucket_no_upscale,
+                random_crop,
+                bucket_reso_steps,
+                caption_dropout_every_n_epochs,
+                caption_dropout_rate,
+                noise_offset,
+                additional_parameters,
+                vae_batch_size,
+                min_snr_gamma,
+            ) = gradio_advanced_training()
+            color_aug.change(
+                color_aug_changed,
+                inputs=[color_aug],
+                outputs=[cache_latents],
+            )
+
+        (
+            sample_every_n_steps,
+            sample_every_n_epochs,
+            sample_sampler,
+            sample_prompts,
+        ) = sample_gradio_config()
+
+    with gr.Tab('Tools'):
+        gr.Markdown(
+            'This section provide Dreambooth tools to help setup your dataset...'
+        )
+        gradio_dreambooth_folder_creation_tab(
+            train_data_dir_input=train_data_dir,
+            reg_data_dir_input=reg_data_dir,
+            output_dir_input=output_dir,
+            logging_dir_input=logging_dir,
+        )
+
+    button_run = gr.Button('Train model', variant='primary')
+
+    # Setup gradio tensorboard buttons
+    button_start_tensorboard, button_stop_tensorboard = gradio_tensorboard()
+
+    button_start_tensorboard.click(
+        start_tensorboard,
+        inputs=logging_dir,
+        show_progress=False,
+    )
+
+    button_stop_tensorboard.click(
+        stop_tensorboard,
+        show_progress=False,
+    )
+
+    settings_list = [
+        pretrained_model_name_or_path,
+        v2,
+        v_parameterization,
+        logging_dir,
+        train_data_dir,
+        reg_data_dir,
+        output_dir,
+        max_resolution,
+        learning_rate,
+        lr_scheduler,
+        lr_warmup,
+        train_batch_size,
+        epoch,
+        save_every_n_epochs,
+        mixed_precision,
+        save_precision,
+        seed,
+        num_cpu_threads_per_process,
+        cache_latents,
+        caption_extension,
+        enable_bucket,
+        gradient_checkpointing,
+        full_fp16,
+        no_token_padding,
+        stop_text_encoder_training,
+        # use_8bit_adam,
+        xformers,
+        save_model_as,
+        shuffle_caption,
+        save_state,
+        resume,
+        prior_loss_weight,
+        color_aug,
+        flip_aug,
+        clip_skip,
+        vae,
+        output_name,
+        max_token_length,
+        max_train_epochs,
+        max_data_loader_n_workers,
+        mem_eff_attn,
+        gradient_accumulation_steps,
+        model_list,
+        token_string,
+        init_word,
+        num_vectors_per_token,
+        max_train_steps,
+        weights,
+        template,
+        keep_tokens,
+        persistent_data_loader_workers,
+        bucket_no_upscale,
+        random_crop,
+        bucket_reso_steps,
+        caption_dropout_every_n_epochs,
+        caption_dropout_rate,
+        optimizer,
+        optimizer_args,
+        noise_offset,
+        sample_every_n_steps,
+        sample_every_n_epochs,
+        sample_sampler,
+        sample_prompts,
+        additional_parameters,
+        vae_batch_size,
+        min_snr_gamma,
+    ]
+
+    button_open_config.click(
+        open_configuration,
+        inputs=[dummy_db_true, config_file_name] + settings_list,
+        outputs=[config_file_name] + settings_list,
+        show_progress=False,
+    )
+
+    button_load_config.click(
+        open_configuration,
+        inputs=[dummy_db_false, config_file_name] + settings_list,
+        outputs=[config_file_name] + settings_list,
+        show_progress=False,
+    )
+
+    button_save_config.click(
+        save_configuration,
+        inputs=[dummy_db_false, config_file_name] + settings_list,
+        outputs=[config_file_name],
+        show_progress=False,
+    )
+
+    button_save_as_config.click(
+        save_configuration,
+        inputs=[dummy_db_true, config_file_name] + settings_list,
+        outputs=[config_file_name],
+        show_progress=False,
+    )
+
+    button_run.click(
+        train_model,
+        inputs=settings_list,
+        show_progress=False,
+    )
+
+    return (
+        train_data_dir,
+        reg_data_dir,
+        output_dir,
+        logging_dir,
+    )
+
+
+def UI(**kwargs):
+    css = ''
+
+    if os.path.exists('./style.css'):
+        with open(os.path.join('./style.css'), 'r', encoding='utf8') as file:
+            print('Load CSS...')
+            css += file.read() + '\n'
+
+    interface = gr.Blocks(css=css)
+
+    with interface:
+        with gr.Tab('Dreambooth TI'):
+            (
+                train_data_dir_input,
+                reg_data_dir_input,
+                output_dir_input,
+                logging_dir_input,
+            ) = ti_tab()
+        with gr.Tab('Utilities'):
+            utilities_tab(
+                train_data_dir_input=train_data_dir_input,
+                reg_data_dir_input=reg_data_dir_input,
+                output_dir_input=output_dir_input,
+                logging_dir_input=logging_dir_input,
+                enable_copy_info_button=True,
+            )
+
+    # Show the interface
+    launch_kwargs = {}
+    if not kwargs.get('username', None) == '':
+        launch_kwargs['auth'] = (
+            kwargs.get('username', None),
+            kwargs.get('password', None),
+        )
+    if kwargs.get('server_port', 0) > 0:
+        launch_kwargs['server_port'] = kwargs.get('server_port', 0)
+    if kwargs.get('inbrowser', False):
+        launch_kwargs['inbrowser'] = kwargs.get('inbrowser', False)
+    print(launch_kwargs)
+    interface.launch(**launch_kwargs)
+
+
+if __name__ == '__main__':
+    # torch.cuda.set_per_process_memory_fraction(0.48)
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        '--username', type=str, default='', help='Username for authentication'
+    )
+    parser.add_argument(
+        '--password', type=str, default='', help='Password for authentication'
+    )
+    parser.add_argument(
+        '--server_port',
+        type=int,
+        default=0,
+        help='Port to run the server listener on',
+    )
+    parser.add_argument(
+        '--inbrowser', action='store_true', help='Open in browser'
+    )
+
+    args = parser.parse_args()
+
+    UI(
+        username=args.username,
+        password=args.password,
+        inbrowser=args.inbrowser,
+        server_port=args.server_port,
+    )

train_README-ja.md

@@ -0,0 +1,945 @@
+__ドキュメント更新中のため記述に誤りがあるかもしれません。__
+
+# 学習について、共通編
+
+当リポジトリではモデルのfine tuning、DreamBooth、およびLoRAとTextual Inversion（[XTI:P+](https://github.com/kohya-ss/sd-scripts/pull/327)を含む）の学習をサポートします。この文書ではそれらに共通する、学習データの準備方法やオプション等について説明します。
+
+# 概要
+
+あらかじめこのリポジトリのREADMEを参照し、環境整備を行ってください。
+
+
+以下について説明します。
+
+1. 学習データの準備について（設定ファイルを用いる新形式）
+1. 学習で使われる用語のごく簡単な解説
+1. 以前の指定形式（設定ファイルを用いずコマンドラインから指定）
+1. 学習途中のサンプル画像生成
+1. 各スクリプトで共通の、よく使われるオプション
+1. fine tuning 方式のメタデータ準備：キャプションニングなど
+
+1.だけ実行すればとりあえず学習は可能です（学習については各スクリプトのドキュメントを参照）。2.以降は必要に応じて参照してください。
+
+
+# 学習データの準備について
+
+任意のフォルダ（複数でも可）に学習データの画像ファイルを用意しておきます。`.png`, `.jpg`, `.jpeg`, `.webp`, `.bmp` をサポートします。リサイズなどの前処理は基本的に必要ありません。
+
+ただし学習解像度（後述）よりも極端に小さい画像は使わないか、あらかじめ超解像AIなどで拡大しておくことをお勧めします。また極端に大きな画像（3000x3000ピクセル程度？）よりも大きな画像はエラーになる場合があるようですので事前に縮小してください。
+
+学習時には、モデルに学ばせる画像データを整理し、スクリプトに対して指定する必要があります。学習データの数、学習対象、キャプション（画像の説明）が用意できるか否かなどにより、いくつかの方法で学習データを指定できます。以下の方式があります（それぞれの名前は一般的なものではなく、当リポジトリ独自の定義です）。正則化画像については後述します。
+
+1. DreamBooth、class+identifier方式（正則化画像使用可）
+
+    特定の単語 (identifier) に学習対象を紐づけるように学習します。キャプションを用意する必要はありません。たとえば特定のキャラを学ばせる場合に使うとキャプションを用意する必要がない分、手軽ですが、髪型や服装、背景など学習データの全要素が identifier に紐づけられて学習されるため、生成時のプロンプトで服が変えられない、といった事態も起こりえます。
+
+1. DreamBooth、キャプション方式（正則化画像使用可）
+
+    画像ごとにキャプションが記録されたテキストファイルを用意して学習します。たとえば特定のキャラを学ばせると、画像の詳細をキャプションに記述することで（白い服を着たキャラA、赤い服を着たキャラA、など）キャラとそれ以外の要素が分離され、より厳密にモデルがキャラだけを学ぶことが期待できます。
+
+1. fine tuning方式（正則化画像使用不可）
+
+    あらかじめキャプションをメタデータファイルにまとめます。タグとキャプションを分けて管理したり、学習を高速化するためlatentsを事前キャッシュしたりなどの機能をサポートします（いずれも別文書で説明しています）。（fine tuning方式という名前ですが fine tuning 以外でも使えます。）
+
+学習したいものと使用できる指定方法の組み合わせは以下の通りです。
+
+| 学習対象または方法 | スクリプト | DB / class+identifier | DB / キャプション | fine tuning |
+| ----- | ----- | ----- | ----- | ----- |
+| モデルをfine tuning | `fine_tune.py`| x | x | o |
+| モデルをDreamBooth | `train_db.py`| o | o | x |
+| LoRA | `train_network.py`| o | o | o |
+| Textual Invesion | `train_textual_inversion.py`| o | o | o |
+
+## どれを選ぶか
+
+LoRA、Textual Inversionについては、手軽にキャプションファイルを用意せずに学習したい場合はDreamBooth class+identifier、用意できるならDreamBooth キャプション方式がよいでしょう。学習データの枚数が多く、かつ正則化画像を使用しない場合はfine tuning方式も検討してください。
+
+DreamBoothについても同様ですが、fine tuning方式は使えません。fine tuningの場合はfine tuning方式のみです。
+
+# 各方式の指定方法について
+
+ここではそれぞれの指定方法で典型的なパターンについてだけ説明し���す。より詳細な指定方法については [データセット設定](./config_README-ja.md) をご覧ください。
+
+# DreamBooth、class+identifier方式（正則化画像使用可）
+
+この方式では、各画像は `class identifier` というキャプションで学習されたのと同じことになります（`shs dog` など）。
+
+## step 1. identifierとclassを決める
+
+学ばせたい対象を結びつける単語identifierと、対象の属するclassを決めます。
+
+（instanceなどいろいろな呼び方がありますが、とりあえず元の論文に合わせます。）
+
+以下ごく簡単に説明します（詳しくは調べてください）。
+
+classは学習対象の一般的な種別です。たとえば特定の犬種を学ばせる場合には、classはdogになります。アニメキャラならモデルによりboyやgirl、1boyや1girlになるでしょう。
+
+identifierは学習対象を識別して学習するためのものです。任意の単語で構いませんが、元論文によると「tokinizerで1トークンになる3文字以下でレアな単語」が良いとのことです。
+
+identifierとclassを使い、たとえば「shs dog」などでモデルを学習することで、学習させたい対象をclassから識別して学習できます。
+
+画像生成時には「shs dog」とすれば学ばせた犬種の画像が生成されます。
+
+（identifierとして私が最近使っているものを参考までに挙げると、``shs sts scs cpc coc cic msm usu ici lvl cic dii muk ori hru rik koo yos wny`` などです。本当は Danbooru Tag に含まれないやつがより望ましいです。）
+
+## step 2. 正則化画像を使うか否かを決め、使う場合には正則化画像を生成する
+
+正則化画像とは、前述のclass全体が、学習対象に引っ張られることを防ぐための画像です（language drift）。正則化画像を使わないと、たとえば `shs 1girl` で特定のキャラクタを学ばせると、単なる `1girl` というプロンプトで生成してもそのキャラに似てきます。これは `1girl` が学習時のキャプションに含まれているためです。
+
+学習対象の画像と正則化画像を同時に学ばせることで、class は class のままで留まり、identifier をプロンプトにつけた時だけ学習対象が生成されるようになります。
+
+LoRAやDreamBoothで特定のキャラだけ出てくればよい場合は、正則化画像を用いなくても良いといえます。
+
+Textual Inversionでは用いなくてよいでしょう（学ばせる token string がキャプションに含まれない場合はなにも学習されないため）。
+
+正則化画像としては、学習対象のモデルで、class 名だけで生成した画像を用いるのが一般的です（たとえば `1girl`）。ただし生成画像の品質が悪い場合には、プロンプトを工夫したり、ネットから別途ダウンロードした画像を用いることもできます。
+
+（正則化画像も学習されるため、その品質はモデルに影響します。）
+
+一般的には数百枚程度、用意するのが望ましいようです（枚数が少ないと class 画像が一般化されずそれらの特徴を学んでしまいます）。
+
+生成画像を使う場合、通常、生成画像のサイズは学習解像度（より正確にはbucketの解像度、後述）にあわせてください。
+
+## step 2. 設定ファイルの記述
+
+テキストファイルを作成し、拡張子を `.toml` にします。たとえば以下のように記述します。
+
+（`#` で始まっている部分はコメントですので、このままコピペしてそのままでもよいですし、削除しても問題ありません。）
+
+```toml
+[general]
+enable_bucket = true                        # Aspect Ratio Bucketingを使うか否か
+
+[[datasets]]
+resolution = 512                            # 学習解像度
+batch_size = 4                              # バッチサイズ
+
+  [[datasets.subsets]]
+  image_dir = 'C:\hoge'                     # 学習用画像を入れたフォルダを指定
+  class_tokens = 'hoge girl'                # identifier class を指定
+  num_repeats = 10                          # 学習用画像の繰り返し回数
+
+  # 以下は正則化画像を用いる場合のみ記述する。用いない場合は削除する
+  [[datasets.subsets]]
+  is_reg = true
+  image_dir = 'C:\reg'                      # 正則化画像を入れたフォルダを指定
+  class_tokens = 'girl'                     # class を指定
+  num_repeats = 1                           # 正則化画像の繰り返し回数、基本的には1でよい
+```
+
+基本的には以下の場所のみ書き換えれば学習できます。
+
+1. 学習解像度
+
+    数値1つを指定すると正方形（`512`なら512x512）、鍵カッコカンマ区切りで2つ指定すると横×縦（`[512,768]`なら512x768）になります。SD1.x系ではもともとの学習解像度は512です。`[512,768]` 等の大きめの解像度を指定すると縦長、横長画像生成時の破綻を小さくできるかもしれません。SD2.x 768系では `768` です。
+
+1. バッチサイズ
+
+    同時に何件のデータを学習するかを指定します。GPUのVRAMサイズ、学習解像度によって変わってきます。詳しくは後述します。またfine tuning/DreamBooth/LoRA等でも変わってきますので各スクリプトの説明もご覧ください。
+
+1. フォルダ指定
+
+    学習用画像、正則化画像（使用する場合のみ）のフォルダを指定します。画像データが含まれているフォルダそのものを指定します。
+
+1. identifier と class の指定
+
+    前述のサンプルの通りです。
+
+1. 繰り返し回数
+
+    後述します。
+
+### 繰り返し回数について
+
+繰り返し回数は、正則化画像の枚数と学習用画像の枚数を調整するために用いられます。正則化画像の枚数は学習用画像よりも多いため、学習用画像を繰り返して枚数を合わせ、1対1の比率で学習できるようにします。
+
+繰り返し回数は「 __学習用画像の繰り返し回数×学習用画像の枚数≧正則化画像の繰り返し回数×正則化画像の枚数__ 」となるように指定してください。
+
+（1 epoch（データが一周すると1 epoch）のデータ数が「学習用画像の繰り返し回数×学習用画像の枚数」となります。正則化画像の枚数がそれより多いと、余った部分の正則化画像は使用されません。）
+
+## step 3. 学習
+
+それぞれのドキュメントを参考に学習を行ってください。
+
+# DreamBooth、キャプション方式（正則化画像使用可）
+
+この方式では各画像はキャプションで学習されます。
+
+## step 1. キャプションファイルを準備する
+
+学習用画像のフォルダに、画像と同じファイル名で、拡張子 `.caption`（設定で変えられます）のファイルを置いてください。それぞれのファイルは1行のみとしてください。エンコーディングは `UTF-8` です。
+
+## step 2. 正則化画像を使うか否かを決め、使う場合には正則化画像を生成する
+
+class+identifier形式と同様です。なお正則化画像にもキャプションを付けることができますが、通常は不要でしょう。
+
+## step 2. 設定ファイルの記述
+
+テキストファイルを作成し、拡張子を `.toml` にします。たとえば以下のように記述します。
+
+```toml
+[general]
+enable_bucket = true                        # Aspect Ratio Bucketingを使うか否か
+
+[[datasets]]
+resolution = 512                            # 学習解像度
+batch_size = 4                              # バッチサイズ
+
+  [[datasets.subsets]]
+  image_dir = 'C:\hoge'                     # 学習用画像を入れたフォルダを指定
+  caption_extension = '.caption'            # キャプションファイルの拡張子　.txt を使う場合には書き換える
+  num_repeats = 10                          # 学習用画像の繰り返し回数
+
+  # 以下は正則化画像を用いる場合のみ記述する。用いない場合は削除する
+  [[datasets.subsets]]
+  is_reg = true
+  image_dir = 'C:\reg'                      # 正則化画像を入れたフォルダを指定
+  class_tokens = 'girl'                     # class を指定
+  num_repeats = 1                           # 正則化画像の繰り返し回数、基本的には1でよい
+```
+
+基本的には以下を場所のみ書き換えれば学習できます。特に記述がない部分は class+identifier 方式と同じです。
+
+1. 学習解像度
+1. バッチサイズ
+1. フォルダ指定
+1. キャプションファイルの拡張子
+
+    任意の拡張子を指定できます。
+1. 繰り返し回数
+
+## step 3. 学習
+
+それぞれのドキュメントを参考に学習を行ってください。
+
+# fine tuning 方式
+
+## step 1. メタデータを準備する
+
+キャプションやタグをまとめた管理用ファイルをメタデータと呼びます。json形式で拡張子は `.json`
+ です。作成方法は長くなりますのでこの文書の末尾に書きました。
+
+## step 2. 設定ファイルの記述
+
+テキストファイルを作成し、拡張子を `.toml` にします。たとえば以下のように記述します。
+
+```toml
+[general]
+shuffle_caption = true
+keep_tokens = 1
+
+[[datasets]]
+resolution = 512                                    # 学習解像度
+batch_size = 4                                      # バッチサイズ
+
+  [[datasets.subsets]]
+  image_dir = 'C:\piyo'                             # 学習用画像を入れたフォルダを指定
+  metadata_file = 'C:\piyo\piyo_md.json'            # メタデータファイル名
+```
+
+基本的には以下を場所のみ書き換えれば学習できます。特に記述がない部分は DreamBooth, class+identifier 方式と同じです。
+
+1. 学習解像度
+1. バッチサイズ
+1. フォルダ指定
+1. メタデータファイル名
+
+    後述の方法で作成したメタデータファイルを指定します。
+
+
+## step 3. 学習
+
+それぞれのドキュメントを参考に学習を行ってください。
+
+# 学習で使われる用語のごく簡単な解説
+
+細かいことは省略していますし私も完全には理解していないため、詳しくは各自お調べください。
+
+## fine tuning（ファインチューニング）
+
+モデルを学習して微調整することを指します。使われ方によって意味が異なってきますが、狭義のfine tuningはStable Diffusionの場合、モデルを画像とキャプションで学習することです。DreamBoothは狭義のfine tuningのひとつの特殊なやり方と言えます。広義のfine tuningは、LoRAやTextual Inversion、Hypernetworksなどを含み、モデルを学習することすべてを含みます。
+
+## ステップ
+
+ざっくりいうと学習データで1回計算すると1ステップです。「学習データのキャプションを今のモデルに流してみて、出てくる画像を学習データの画像と比較し、学習データに近づくようにモデルをわずかに変更する」のが1ステップです。
+
+## バッチサイズ
+
+バッチサイズは1ステップで何件のデータをまとめて計算するかを指定する値です。まとめて計算するため速度は相対的に向上します。また一般的には精度も高くなるといわれています。
+
+`バッチサイズ×ステップ数` が学習に使われるデータの件数になります。そのため、バッチサイズを増やした分だけステップ数を減らすとよいでしょう。
+
+（ただし、たとえば「バッチサイズ1で1600ステップ」と「バッチサイズ4で400ステップ」は同じ結果にはなりません。同じ学習率の場合、一般的には後者のほうが学習不足になります。学習率を多少大きくするか（たとえば `2e-6` など）、ステップ数をたとえば500ステップにするなどして工夫してください。）
+
+バッチサイズを大きくするとその分だけGPUメモリを消費します。メモリが足りなくなるとエラーになりますし、エラーにならないギリギリでは学習速度が低下します。タスクマネージャーや `nvidia-smi` コマンドで使用メモリ量を確認しながら調整するとよいでしょう。
+
+なお、バッチは「一塊のデータ」位の意味です。
+
+## 学習率
+
+ざっくりいうと1ステップごとにどのくらい変化させるかを表します。大きな値を指定するとそれだけ速く学習が進みますが、変化しすぎてモデルが壊れたり、最適な状態にまで至れない場合があります。小さい値を指定すると学習速度は遅くなり、また最適な状態にやはり至れない場合があります。
+
+fine tuning、DreamBoooth、LoRAそれぞれで大きく異なり、また学習データや学習させたいモデル、バッチサイズやステップ数によっても変わってきます。一般的な値から初めて学習状態を見ながら増減してください。
+
+デフォルトでは学習全体を通して学習率は固定です。スケジューラの指定で学習率をどう変化させるか決められますので、それらによっても結果は変わってきます。
+
+## エポック（epoch）
+
+学習データが一通り学習されると（データが一周すると）1 epochです。繰り返し回数を指定した場合は、その繰り返し後のデータが一周すると1 epochです。
+
+1 epochのステップ数は、基本的には `データ件数÷バッチサイズ` ですが、Aspect Ratio Bucketing を使うと微妙に増えます（異なるbucketのデータは同じバッチにできないため、ステップ数が増えます）。
+
+## Aspect Ratio Bucketing
+
+Stable Diffusion のv1は512\*512で学習されていますが、それに加えて256\*1024や384\*640といった解像度でも学習します。これによりトリミングされる部分が減り、より正しくキャプションと画像の関係が学習されることが期待されます。
+
+また任意の解像度で学習するため、事前に画像データの縦横比を統一しておく必要がなくなります。
+
+設定で有効、向こうが切り替えられますが、ここまでの設定ファイルの記述例では有効になっています（`true` が設定されています）。
+
+学習解像度はパラメータとして与えられた解像度の面積（＝メモリ使用量）���超えない範囲で、64ピクセル単位（デフォルト、変更可）で縦横に調整、作成されます。
+
+機械学習では入力サイズをすべて統一するのが一般的ですが、特に制約があるわけではなく、実際は同一のバッチ内で統一されていれば大丈夫です。NovelAIの言うbucketingは、あらかじめ教師データを、アスペクト比に応じた学習解像度ごとに分類しておくことを指しているようです。そしてバッチを各bucket内の画像で作成することで、バッチの画像サイズを統一します。
+
+# 以前の指定形式（設定ファイルを用いずコマンドラインから指定）
+
+`.toml` ファイルを指定せずコマンドラインオプションで指定する方法です。DreamBooth class+identifier方式、DreamBooth キャプション方式、fine tuning方式があります。
+
+## DreamBooth、class+identifier方式
+
+フォルダ名で繰り返し回数を指定します。また `train_data_dir` オプションと `reg_data_dir` オプションを用います。
+
+### step 1. 学習用画像の準備
+
+学習用画像を格納するフォルダを作成します。 __さらにその中に__ 、以下の名前でディレクトリを作成します。
+
+```
+<繰り返し回数>_<identifier> <class>
+```
+
+間の``_``を忘れないでください。
+
+たとえば「sls frog」というプロンプトで、データを20回繰り返す場合、「20_sls frog」となります。以下のようになります。
+
+![image](https://user-images.githubusercontent.com/52813779/210770636-1c851377-5936-4c15-90b7-8ac8ad6c2074.png)
+
+### 複数class、複数対象（identifier）の学習
+
+方法は単純で、学習用画像のフォルダ内に ``繰り返し回数_<identifier> <class>`` のフォルダを複数、正則化画像フォルダにも同様に ``繰り返し回数_<class>`` のフォルダを複数、用意してください。
+
+たとえば「sls frog」と「cpc rabbit」を同時に学習する場合、以下のようになります。
+
+![image](https://user-images.githubusercontent.com/52813779/210777933-a22229db-b219-4cd8-83ca-e87320fc4192.png)
+
+classがひとつで対象が複数の場合、正則化画像フォルダはひとつで構いません。たとえば1girlにキャラAとキャラBがいる場合は次のようにします。
+
+- train_girls
+  - 10_sls 1girl
+  - 10_cpc 1girl
+- reg_girls
+  - 1_1girl
+
+### step 2. 正則化画像の準備
+
+正則化画像を使う場合の手順です。
+
+正則化画像を格納するフォルダを作成します。 __さらにその中に__  ``<繰り返し回数>_<class>`` という名前でディレクトリを作成します。
+
+たとえば「frog」というプロンプトで、データを繰り返さない（1回だけ）場合、以下のようになります。
+
+![image](https://user-images.githubusercontent.com/52813779/210770897-329758e5-3675-49f1-b345-c135f1725832.png)
+
+
+### step 3. 学習の実行
+
+各学習スクリプトを実行します。 `--train_data_dir` オプションで前述の学習用データのフォルダを（__画像を含むフォルダではなく、その親フォルダ__）、`--reg_data_dir` オプションで正則化画像のフォルダ（__画像を含むフォルダではなく、その親フォルダ__）を指定してください。
+
+## DreamBooth、キャプション方式
+
+学習用画像、正則化画像のフォルダに、画像と同じファイル名で、拡張子.caption（オプションで変えられます）のファイルを置くと、そのファイルからキャプションを読み込みプロンプトとして学習します。
+
+※それらの画像の学習に、フォルダ名（identifier class）は使用されなくなります。
+
+キャプションファイルの拡張子はデフォルトで.captionです。学習スクリプトの `--caption_extension` オプションで変更できます。`--shuffle_caption` オプションで学習時のキャプションについて、カンマ区切りの各部分をシャッフルしながら学習します。
+
+## fine tuning 方式
+
+メタデータを作るところまでは設定ファイルを使う場合と同様です。`in_json` オプションでメタデータファイルを指定します。
+
+# 学習途中でのサンプル出力
+
+学習中のモデルで試しに画像生成することで学習の進み方を確認できます。学習スクリプトに以下のオプションを指定します。
+
+- `--sample_every_n_steps` / `--sample_every_n_epochs`
+    
+    サンプル出力するステップ数またはエポック数を指定します。この数ごとにサンプル出力します。両方指定するとエポック数が優先されます。
+
+- `--sample_prompts`
+
+    サンプル出力用プロンプトのファイルを指定します。
+
+- `--sample_sampler`
+
+    サンプル出力に使うサンプラーを指定します���
+    `'ddim', 'pndm', 'heun', 'dpmsolver', 'dpmsolver++', 'dpmsingle', 'k_lms', 'k_euler', 'k_euler_a', 'k_dpm_2', 'k_dpm_2_a'`が選べます。
+
+サンプル出力を行うにはあらかじめプロンプトを記述したテキストファイルを用意しておく必要があります。1行につき1プロンプトで記述します。
+
+たとえば以下のようになります。
+
+```txt
+# prompt 1
+masterpiece, best quality, 1girl, in white shirts, upper body, looking at viewer, simple background --n low quality, worst quality, bad anatomy,bad composition, poor, low effort --w 768 --h 768 --d 1 --l 7.5 --s 28
+
+# prompt 2
+masterpiece, best quality, 1boy, in business suit, standing at street, looking back --n low quality, worst quality, bad anatomy,bad composition, poor, low effort --w 576 --h 832 --d 2 --l 5.5 --s 40
+```
+
+先頭が `#` の行はコメントになります。`--n` のように 「`--` + 英小文字」で生成画像へのオプションを指定できます。以下が使えます。
+
+- `--n` 次のオプションまでをネガティブプロンプトとします。
+- `--w` 生成画像の横幅を指定します。
+- `--h` 生成画像の高さを指定します。
+- `--d` 生成画像のseedを指定します。
+- `--l` 生成画像のCFG scaleを指定します。
+- `--s` 生成時のステップ数を指定します。
+
+
+# 各スクリプトで共通の、よく使われるオプション
+
+スクリプトの更新後、ドキュメントの更新が追い付いていない場合があります。その場合は `--help` オプションで使用できるオプションを確認してください。
+
+## 学習に使うモデル指定
+
+- `--v2` / `--v_parameterization`
+    
+    学習対象モデルとしてHugging Faceのstable-diffusion-2-base、またはそこからのfine tuningモデルを使う場合（推論時に `v2-inference.yaml` を使うように指示されているモデルの場合）は `--v2` オプションを、stable-diffusion-2や768-v-ema.ckpt、およびそれらのfine tuningモデルを使う場合（推論時に `v2-inference-v.yaml` を使うモデルの場合）は `--v2` と `--v_parameterization` の両方のオプションを指定してください。
+
+    Stable Diffusion 2.0では大きく以下の点が変わっています。
+
+    1. 使用するTokenizer
+    2. 使用するText Encoderおよび使用する出力層（2.0は最後から二番目の層を使う）
+    3. Text Encoderの出力次元数（768->1024）
+    4. U-Netの構造（CrossAttentionのhead数など）
+    5. v-parameterization（サンプリング方法が変更されているらしい）
+
+    このうちbaseでは1～4が、baseのつかない方（768-v）では1～5が採用されています。1～4を有効にするのがv2オプション、5を有効にするのがv_parameterizationオプションです。
+
+- `--pretrained_model_name_or_path` 
+    
+    追加学習を行う元となるモデルを指定します。Stable Diffusionのcheckpointファイル（.ckptまたは.safetensors）、Diffusersのローカルディスクにあるモデルディレクトリ、DiffusersのモデルID（"stabilityai/stable-diffusion-2"など）が指定できます。
+
+## 学習に関する設定
+
+- `--output_dir` 
+
+    学習後のモデルを保存するフォルダを指定します。
+    
+- `--output_name` 
+    
+    モデルのファイル名を拡張子を除いて指定します。
+    
+- `--dataset_config` 
+
+    データセットの設定を記述した `.toml` ファイルを指定します。
+
+- `--max_train_steps` / `--max_train_epochs`
+
+    学習するステップ数やエポック数を指定します。両方指定するとエポック数のほうが優先されます。
+
+- `--mixed_precision`
+
+    省メモリ化のため mixed precision （混合精度）で学習します。`--mixed_precision="fp16"` のように指定します。mixed precision なし（デフォルト）と比べて精度が低くなる可能性がありますが、学習に必要なGPUメモリ量が大きく減ります。
+    
+    （RTX30 シリーズ以降では `bf16` も指定できます。環境整備時にaccelerateに行った設定と合わせてください）。
+    
+- `--gradient_checkpointing`
+
+    学習時の重みの計算をまとめて行うのではなく少しずつ行うことで、学習に必要なGPUメモリ量を減らします。オンオフは精度には影響しませんが、オンにするとバッチサイズを大きくできるため、そちらでの影響はあります。
+    
+    また一般的にはオンにすると速度は低下しますが、バッチサイズを大きくできるので、トータルでの学習時間はむしろ速くなるかもしれません。
+
+- `--xformers` / `--mem_eff_attn`
+
+    xformersオプションを指定するとxformersのCrossAttentionを用います。xformersをインストールしていない場合やエラーとなる場合（環境にもよ��ますが `mixed_precision="no"` の場合など）、代わりに `mem_eff_attn` オプションを指定すると省メモリ版CrossAttentionを使用します（xformersよりも速度は遅くなります）。
+
+- `--save_precision`
+
+    保存時のデータ精度を指定します。save_precisionオプションにfloat、fp16、bf16のいずれかを指定すると、その形式でモデルを保存します（DreamBooth、fine tuningでDiffusers形式でモデルを保存する場合は無効です）。モデルのサイズを削減したい場合などにお使いください。
+
+- `--save_every_n_epochs` / `--save_state` / `--resume`
+    save_every_n_epochsオプションに数値を指定すると、そのエポックごとに学習途中のモデルを保存します。
+
+    save_stateオプションを同時に指定すると、optimizer等の状態も含めた学習状態を合わせて保存します（保存したモデルからも学習再開できますが、それに比べると精度の向上、学習時間の短縮が期待できます）。保存先はフォルダになります。
+    
+    学習状態は保存先フォルダに `<output_name>-??????-state`（??????はエポック数）という名前のフォルダで出力されます。長時間にわたる学習時にご利用ください。
+
+    保存された学習状態から学習を再開するにはresumeオプションを使います。学習状態のフォルダ（`output_dir` ではなくその中のstateのフォルダ）を指定してください。
+
+    なおAcceleratorの仕様により、エポック数、global stepは保存されておらず、resumeしたときにも1からになりますがご容赦ください。
+
+- `--save_model_as` （DreamBooth, fine tuning のみ）
+
+    モデルの保存形式を`ckpt, safetensors, diffusers, diffusers_safetensors` から選べます。
+    
+    `--save_model_as=safetensors` のように指定します。Stable Diffusion形式（ckptまたはsafetensors）を読み込み、Diffusers形式で保存する場合、不足する情報はHugging Faceからv1.5またはv2.1の情報を落としてきて補完します。
+    
+- `--clip_skip`
+    
+    `2` を指定すると、Text Encoder (CLIP) の後ろから二番目の層の出力を用います。1またはオプション省略時は最後の層を用います。
+
+    ※SD2.0はデフォルトで後ろから二番目の層を使うため、SD2.0の学習では指定しないでください。
+
+    学習対象のモデルがもともと二番目の層を使うように学習されている場合は、2を指定するとよいでしょう。
+
+    そうではなく最後の層を使用していた場合はモデル全体がそれを前提に学習されています。そのため改めて二番目の層を使用して学習すると、望ましい学習結果を得るにはある程度の枚数の教師データ、長めの学習が必要になるかもしれません。
+
+- `--max_token_length`
+
+    デフォルトは75です。`150` または `225` を指定することでトークン長を拡張して学習できます。長いキャプションで学習する場合に指定してください。
+    
+    ただし学習時のトークン拡張の仕様は Automatic1111 氏のWeb UIとは微妙に異なるため（分割の仕様など）、必要なければ75で学習することをお勧めします。
+
+    clip_skipと同様に、モデルの学習状態と異なる長さで学習するには、ある程度の教師データ枚数、長めの学習時間が必要になると思われます。
+
+- `--persistent_data_loader_workers`
+
+    Windows環境で指定するとエポック間の待ち時間が大幅に短縮されます。
+
+- `--max_data_loader_n_workers`
+
+    データ読み込みのプロセス数を指定します。プロセス数が多いとデータ読み込みが速くなりGPUを効率的に利用できますが、メインメモリを消費します。デフォルトは「`8` または `CPU同時実行スレッド数-1` の小さいほう」なので、メインメモリに余裕がない場合や、GPU使用率が90%程度以上なら、それらの数値を見ながら `2` または `1` 程度まで下げてください。
+
+- `--logging_dir` / `--log_prefix`
+
+    学習ログの保存に関するオプションです。logging_dirオプションにログ保存先フォルダを指定してください。TensorBoard形式のログが保存されます。
+
+    たとえば--logging_dir=logsと指定すると、作業フォルダにlogsフォルダが作成され、その中の日時フォルダにログが保存されます。
+    また--log_prefixオプションを指定すると、日時の前に指定した文字列が追加されます。「--logging_dir=logs --log_prefix=db_style1_」などとして識別用にお使いください。
+
+    TensorBoardでログを確認するには、別のコマンドプロンプトを開き、作業フォルダで以下のように入力します。
+
+    ```
+    tensorboard --logdir=logs
+    ```
+
+    （tensorboardは環境整備時にあわせてインストールされると思いますが、もし入っていないなら `pip install tensorboard` で入れてください。）
+
+    その後ブラウザを開き、http://localhost:6006/ へアクセスすると表示されます。
+
+- `--noise_offset`
+
+    こちらの記事の実装になります: https://www.crosslabs.org//blog/diffusion-with-offset-noise
+    
+    全体的に暗い、明るい画像の生成結果が良くなる可能性があるようです。LoRA学習でも有効なようです。`0.1` 程度の値を指定するとよいようです。
+
+- `--debug_dataset`
+
+    このオプションを付けることで学習を行う前に事前にどのような画像データ、キャプションで学習されるかを確認できます。Escキーを押すと終了してコマンドラインに戻ります。`S`キーで次のステップ（バッチ）、`E`キーで次のエポックに進みます。
+
+    ※Linux環境（Colabを含む）では画像は表示されません。
+
+- `--vae`
+
+    vaeオプションにStable Diffusionのcheckpoint、VAEのcheckpointファイル、DiffusesのモデルまたはVAE（ともにローカルまたはHugging FaceのモデルIDが指定できます）のいずれかを指定すると、そのVAEを使って学習します（latentsのキャッシュ時または学習中のlatents取得時）。
+
+    DreamBoothおよびfine tuningでは、保存されるモデルはこのVAEを組み込んだものになります。
+
+- `--cache_latents`
+
+    使用VRAMを減らすためVAEの出力をメインメモリにキャッシュします。`flip_aug` 以外のaugmentationは使えなくなります。また全体の学習速度が若干速くなります。
+
+- `--min_snr_gamma`
+
+    Min-SNR Weighting strategyを指定します。詳細は[こちら](https://github.com/kohya-ss/sd-scripts/pull/308)を参照してください。論文では`5`が推奨されています。
+
+## オプティマイザ関係
+
+- `--optimizer_type`
+    --オプティマイザの種類を指定します。以下が指定できます。
+    - AdamW : [torch.optim.AdamW](https://pytorch.org/docs/stable/generated/torch.optim.AdamW.html)
+    - 過去のバージョンのオプション未指定時と同じ
+    - AdamW8bit : 引数は同上
+    - 過去のバージョンの--use_8bit_adam指定時と同じ
+    - Lion : https://github.com/lucidrains/lion-pytorch
+    - 過去のバージョンの--use_lion_optimizer指定時と同じ
+    - SGDNesterov : [torch.optim.SGD](https://pytorch.org/docs/stable/generated/torch.optim.SGD.html), nesterov=True
+    - SGDNesterov8bit : 引数は同上
+    - DAdaptation : https://github.com/facebookresearch/dadaptation
+    - AdaFactor : [Transformers AdaFactor](https://huggingface.co/docs/transformers/main_classes/optimizer_schedules)
+    - 任意のオプティマイザ
+
+- `--learning_rate`
+
+    学習率を指定します。適切な学習率は学習スクリプトにより異なりますので、それぞれの説明を参照してください。
+
+- `--lr_scheduler` / `--lr_warmup_steps` / `--lr_scheduler_num_cycles` / `--lr_scheduler_power`
+  
+    学習率のスケジューラ関連の指定です。
+
+    lr_schedulerオプションで学習率のスケジューラをlinear, cosine, cosine_with_restarts, polynomial, constant, constant_with_warmup, 任意のスケジューラから選べます。デフォルトはconstantです。
+    
+    lr_warmup_stepsでスケジューラのウォームアップ（だんだん学習率を変えていく）ステップ数を指定できます。
+    
+    lr_scheduler_num_cycles は cosine with restartsスケジューラでのリスタート回数、lr_scheduler_power は polynomialスケジューラでのpolynomial power です。
+
+    詳細については各自お調べください。
+
+    任意のスケジューラを使う場合、任意のオプティマイザと同様に、`--scheduler_args`でオプション引数を指定してください。
+
+### オプティマイザの指定について
+
+オプティマイザのオプション引数は--optimizer_argsオプションで指定してください。key=valueの形式で、複数の値が指定できます。また、valueはカンマ区切りで複数の値が指定できます。たとえばAdamWオプティマイザに引数を指定する場合は、``--optimizer_args weight_decay=0.01 betas=.9,.999``のようになります。
+
+オプション引数を指定する場合は、それぞれのオプティマイザの仕様をご確認ください。
+
+一部のオプティマイザでは必須の引数があり、省略すると自動的に追加されます（SGDNesterovのmomentumなど）。コンソールの出力を確認してください。
+
+D-Adaptationオプティマイザは学習率を自動調整します。学習率のオプションに指定した値は学習率そのものではなくD-Adaptationが決定した学習率の適用率になりますので、通常は1.0を指定してください。Text EncoderにU-Netの半分の学習率を指定したい場合は、``--text_encoder_lr=0.5 --unet_lr=1.0``と指定します。
+
+AdaFactorオプティマイザはrelative_step=Trueを指定すると学習率を自動調整できます（省略時はデフォルトで追加されます）。自動調整する場合は学習率のスケジューラにはadafactor_schedulerが強制的に使用されます。またscale_parameterとwarmup_initを指定するとよいようです。
+
+自動調整する場合のオプション指定はたとえば ``--optimizer_args "relative_step=True" "scale_parameter=True" "warmup_init=True"`` のようになります。
+
+学習率を自動調整しない場合はオプション引数 ``relative_step=False`` を追加してください。その場合、学習率のスケジューラにはconstant_with_warmupが、また勾配のclip normをしないことが推奨されているようです。そのため引数は ``--optimizer_type=adafactor --optimizer_args "relative_step=False" --lr_scheduler="constant_with_warmup" --max_grad_norm=0.0`` のようになります。
+
+### 任意のオプティマイザを使う
+
+``torch.optim`` のオプティマイザを使う場合にはクラス名のみを（``--optimizer_type=RMSprop``など）、他のモジュールのオプティマイザを使う時は「モジュール名.クラス名」を指定してください（``--optimizer_type=bitsandbytes.optim.lamb.LAMB``など）。
+
+（内部でimportlibしているだけで動作は未確認です。必要ならパッケージをインストールしてください。）
+
+
+<!-- 
+## 任意サイズの画像での学習 --resolution
+正方形以外で学習できます。resolutionに「448,640」のように「幅,高さ」で指定してください。幅と高さは64で割り切れる必要があります。学習用画像、正則化画像のサイズを合わせてください。
+
+個人的には縦長の画像を生成することが多いため「448,640」などで学習することもあります。
+
+## Aspect Ratio Bucketing --enable_bucket / --min_bucket_reso / --max_bucket_reso
+enable_bucketオプションを指定すると有効になります。Stable Diffusionは512x512で学習されていますが、それに加えて256x768や384x640といった解像度でも学習します。
+
+このオプションを指定した場合は、学習用画像、正則化画像を特定の解像度に統一する必要はありません。いくつかの解像度（アスペクト比）から最適なものを選び、その解像度で学習します。
+解像度は64ピクセル単位のため、元画像とアスペクト比が完全に一致しない場合がありますが、その場合は、はみ出した部分がわずかにトリミングされます。
+
+解像度の最小サイズをmin_bucket_resoオプションで、最大サイズをmax_bucket_resoで指定できます。デフォルトはそれぞれ256、1024です。
+たとえば最小サイズに384を指定すると、256x1024や320x768などの解像度は使わなくなります。
+解像度を768x768のように大きくした場合、最大サイズに1280などを指定しても良いかもしれません。
+
+なおAspect Ratio Bucketingを有効にするときには、正則化画像についても、学習用画像と似た傾向の様々な解像度を用意した方がいいかもしれません。
+
+（ひとつのバッチ内の画像が学習用画像、正則化画像に偏らなくなるため。そこまで大きな影響はないと思いますが……。）
+
+## augmentation --color_aug / --flip_aug
+augmentationは学習時に動的にデータを変化させることで、モデルの性能を上げる手法です。color_augで色合いを微妙に変えつつ、flip_augで左右反転をしつつ、学習します。
+
+動的にデータを変化させるため、cache_latentsオプションと同時に指定できません。
+
+
+## 勾配をfp16とした学習（実験的機能） --full_fp16
+full_fp16オプションを指定すると勾配を通常のfloat32からfloat16（fp16）に変更して学習します（mixed precisionではなく完全なfp16学習になるようです）。
+これによりSD1.xの512x512サイズでは8GB未満、SD2.xの512x512サイズで12GB未満のVRAM使用量で学習できるようです。
+
+あらかじめaccelerate configでfp16を指定し、オプションで ``mixed_precision="fp16"`` としてください（bf16では動作しません）。
+
+メモリ使用量を最小化するためには、xformers、use_8bit_adam、cache_latents、gradient_checkpointingの各オプションを指定し、train_batch_sizeを1としてください。
+
+（余裕があるようならtrain_batch_sizeを段階的に増やすと若干精度が上がるはずです。）
+
+PyTorchのソースにパッチを当てて無理やり実現しています（PyTorch 1.12.1と1.13.0で確認）。精度はかなり落ちますし、途中で学習失敗する確率も高くなります。
+学習率やステップ数の設定もシビアなようです。それらを認識したうえで自己責任でお使いください。
+
+-->
+
+# メタデータファイルの作成
+
+## 教師データの用意
+
+前述のように学習させたい画像データを用意し、任意のフォルダに入れてください。
+
+たとえば以下のように画像を格納します。
+
+![教師データフォルダのスクショ](https://user-images.githubusercontent.com/52813779/208907739-8e89d5fa-6ca8-4b60-8927-f484d2a9ae04.png)
+
+## 自動キャプショニング
+
+キャプションを使わずタグだけで学習する場合はスキップしてください。
+
+また手動でキャプションを用意する場合、キャプションは教師データ画像と同じディレクトリに、同じファイル名、拡張子.caption等で用意してください。各ファイルは1行のみのテキストファイルとします。
+
+### BLIPによるキャプショニング
+
+最新版ではBLIPのダウンロード、重みのダウンロード、仮想環境の追加は不要になりました。そのままで動作します。
+
+finetuneフォルダ内のmake_captions.pyを実行します。
+
+```
+python finetune\make_captions.py --batch_size <バッチサイズ> <教師データフォルダ>
+```
+
+バッチサイズ8、教師データを親フォルダのtrain_dataに置いた場合、以下のようになります。
+
+```
+python finetune\make_captions.py --batch_size 8 ..\train_data
+```
+
+キャプションファイルが教師データ画像と同じディレクトリに、同じファイル名、拡張子.captionで作成されます。
+
+batch_sizeはGPUのVRAM容量に応じて増減してください。大きいほうが速くなります（VRAM 12GBでももう少し増やせると思います）。
+max_lengthオプションでキャプションの最大長を指定できます。デフォルトは75です。モデルをトークン長225で学習する場合には長くしても良いかもしれません。
+caption_extensionオプションでキャプションの拡張子を変更できます。デフォルトは.captionです（.txtにすると後述のDeepDanbooruと競合します）。
+
+複数の教師データフォルダがある場合には、それぞれのフォルダに対して実行してください。
+
+なお、推論にランダム性があるため、実行するたびに結果が変わります。固定する場合には--seedオプションで `--seed 42` のように乱数seedを指定してください。
+
+その他のオプションは `--help` でヘルプをご参照ください（パラメータの意味についてはドキュメントがまとまっていないようで、ソースを見るしかないようです）。
+
+デフォルトでは拡張子.captionでキャプションファイルが生成されます。
+
+![captionが生成されたフォルダ](https://user-images.githubusercontent.com/52813779/208908845-48a9d36c-f6ee-4dae-af71-9ab462d1459e.png)
+
+たとえば以下のようなキャプションが付きます。
+
+![キャプションと画像](https://user-images.githubusercontent.com/52813779/208908947-af936957-5d73-4339-b6c8-945a52857373.png)
+
+## DeepDanbooruによるタグ付け
+
+danbooruタグのタグ付け自体を行わない場合は「キャプションとタグ情報の前処理」に進んでください。
+
+タグ付けはDeepDanbooruまたはWD14Taggerで行います。WD14Taggerのほうが精度が良いようです。WD14Taggerでタグ付けする場合は、次の章へ進んでください。
+
+### 環境整備
+
+DeepDanbooru https://github.com/KichangKim/DeepDanbooru  を作業フォルダにcloneしてくるか、zipをダウンロードして展開します。私はzipで展開しました。
+またDeepDanbooruのReleasesのページ https://github.com/KichangKim/DeepDanbooru/releases  の「DeepDanbooru Pretrained Model v3-20211112-sgd-e28」のAssetsから、deepdanbooru-v3-20211112-sgd-e28.zipをダウンロードしてきてDeepDanbooruのフォルダに展開します。
+
+以下からダウンロードします。Assetsをクリックして開き、そこからダウンロードします。
+
+![DeepDanbooruダウンロードページ](https://user-images.githubusercontent.com/52813779/208909417-10e597df-7085-41ee-bd06-3e856a1339df.png)
+
+以下のようなこういうディレクトリ構造にしてください
+
+![DeepDanbooruのディレクトリ構造](https://user-images.githubusercontent.com/52813779/208909486-38935d8b-8dc6-43f1-84d3-fef99bc471aa.png)
+
+Diffusersの環境に必要なライブラリをインストールします。DeepDanbooruのフォルダに移動してインストールします（実質的にはtensorflow-ioが追加されるだけだと思います）。
+
+```
+pip install -r requirements.txt
+```
+
+続いてDeepDanbooru自体をインストールします。
+
+```
+pip install .
+```
+
+以上でタグ付けの環境整備は完了です。
+
+### タグ付けの実施
+DeepDanbooruのフォルダに移動し、deepdanbooruを実行してタグ付けを行います。
+
+```
+deepdanbooru evaluate <教師データフォルダ> --project-path deepdanbooru-v3-20211112-sgd-e28 --allow-folder --save-txt
+```
+
+教師データを親フォルダのtrain_dataに置いた場合、以下のようになります。
+
+```
+deepdanbooru evaluate ../train_data --project-path deepdanbooru-v3-20211112-sgd-e28 --allow-folder --save-txt
+```
+
+タグファイルが教師データ画像と同じディレクトリに、同じファイル名、拡張子.txtで作成されます。1件ずつ処理されるためわりと遅いです。
+
+複数の教師データフォルダがある場合には、それぞれのフォルダに対して実行してください。
+
+以下のように生成されます。
+
+![DeepDanbooruの生成ファイル](https://user-images.githubusercontent.com/52813779/208909855-d21b9c98-f2d3-4283-8238-5b0e5aad6691.png)
+
+こんな感じにタグが付きます（すごい情報量……）。
+
+![DeepDanbooruタグと画像](https://user-images.githubusercontent.com/52813779/208909908-a7920174-266e-48d5-aaef-940aba709519.png)
+
+## WD14Taggerによるタグ付け
+
+DeepDanbooruの代わりにWD14Taggerを用いる手順です。
+
+Automatic1111氏のWebUIで使用しているtaggerを利用します。こちらのgithubページ（https://github.com/toriato/stable-diffusion-webui-wd14-tagger#mrsmilingwolfs-model-aka-waifu-diffusion-14-tagger ）の情報を参考にさせていただきました。
+
+最初の環境整備で必要なモジュールはインストール済みです。また重みはHugging Faceから自動的にダウンロードしてきます。
+
+### タグ付けの実施
+
+スクリプトを実行してタグ付けを行います。
+```
+python tag_images_by_wd14_tagger.py --batch_size <バッチサイズ> <教師データフォルダ>
+```
+
+教師データを親フォルダのtrain_dataに置いた場合、以下のようになります。
+```
+python tag_images_by_wd14_tagger.py --batch_size 4 ..\train_data
+```
+
+初回起動時にはモデルファイルがwd14_tagger_modelフォルダに自動的にダウンロードされます（フォルダはオプションで変えられます）。以下のようになります。
+
+![ダウンロードされたファイル](https://user-images.githubusercontent.com/52813779/208910447-f7eb0582-90d6-49d3-a666-2b508c7d1842.png)
+
+タグファイルが教師データ画像と同じディレクトリに、同じファイル名、拡張子.txtで作成されます。
+
+![生成されたタグファイル](https://user-images.githubusercontent.com/52813779/208910534-ea514373-1185-4b7d-9ae3-61eb50bc294e.png)
+
+![タグと画像](https://user-images.githubusercontent.com/52813779/208910599-29070c15-7639-474f-b3e4-06bd5a3df29e.png)
+
+threshオプションで、判定されたタグのconfidence（確信度）がいくつ以上でタグをつけるかが指定できます。デフォルトはWD14Taggerのサンプルと同じ0.35です。値を下げるとより多くのタグが付与されますが、精度は下がります。
+
+batch_sizeはGPUのVRAM容量に応じて増減してください。大きいほうが速くなります（VRAM 12GBでももう少し増やせると思います）。caption_extensionオプションでタグファイルの拡張子を変更できます。デフォルトは.txtです。
+
+model_dirオプションでモデルの保存先フォルダを指定できます。
+
+またforce_downloadオプションを指定すると保存先フォルダがあってもモデルを再ダウンロードします。
+
+複数の教師データフォルダがある場合には、それぞれのフォルダに対して実行してください。
+
+## キャプションとタグ情報の前処理
+
+スクリプトから処理しやすいようにキャプションとタグをメタデータとしてひとつのファイルにまとめます。
+
+### キャプションの前処理
+
+キャプションをメタデータに入れるには、作業フォルダ内で以下を実行してください（キャプションを学習に使わない場合は実行不要です）（実際は1行で記述します、以下同様）。`--full_path` オプションを指定してメタデータに画像ファイルの場所をフルパスで格納します。このオプションを省略すると相対パスで記録されますが、フォルダ指定が `.toml` ファイル内で別途必要になります。
+
+```
+python merge_captions_to_metadata.py --full_path <教師データフォルダ>
+　  --in_json <読み込むメタデータファイル名> <メタデータファイル名>
+```
+
+メタデータファイル名は任意の名前です。
+教師データがtrain_data、読み込むメタデータファイルなし、メタデータファイルがmeta_cap.jsonの場合、以下のようになります。
+
+```
+python merge_captions_to_metadata.py --full_path train_data meta_cap.json
+```
+
+caption_extensionオプションでキャプションの拡張子を指定できます。
+
+複数の教師データフォルダがある場合には、full_path引数を指定しつつ、それぞれのフォルダに対して実行してください。
+
+```
+python merge_captions_to_metadata.py --full_path 
+    train_data1 meta_cap1.json
+python merge_captions_to_metadata.py --full_path --in_json meta_cap1.json 
+    train_data2 meta_cap2.json
+```
+
+in_jsonを省略すると書き込み先メタデータファイルがあるとそこから読み込み、そこに上書きします。
+
+__※in_jsonオプションと書き込み先を都度書き換えて、別のメタデータファイルへ書き出すようにすると安全です。__
+
+### タグの前処理
+
+同様にタグもメタデータにまとめます（タグを学習に使わない場合は実行不要です）。
+```
+python merge_dd_tags_to_metadata.py --full_path <教師データフォルダ> 
+    --in_json <読み込むメタデータファイル名> <書き込むメタデータファイル名>
+```
+
+先と同じディレクトリ構成で、meta_cap.jsonを読み、meta_cap_dd.jsonに書きだす場合、以下となります。
+```
+python merge_dd_tags_to_metadata.py --full_path train_data --in_json meta_cap.json meta_cap_dd.json
+```
+
+複数の教師データフォルダがある場合には、full_path引数を指定しつつ、それぞれのフォルダに対して実行してください。
+
+```
+python merge_dd_tags_to_metadata.py --full_path --in_json meta_cap2.json
+    train_data1 meta_cap_dd1.json
+python merge_dd_tags_to_metadata.py --full_path --in_json meta_cap_dd1.json 
+    train_data2 meta_cap_dd2.json
+```
+
+in_jsonを省略すると書き込み先メタデータファイルがあるとそこから読み込み、そこに上書きします。
+
+__※in_jsonオプションと書き込み先を都度書き換えて、別のメタデータファイルへ書き出すようにすると安全です。__
+
+### キャプションとタグのクリーニング
+
+ここまででメタデータファイルにキャプションとDeepDanbooruのタグがまとめられています。ただ自動キャプショニングにしたキャプションは表記ゆれなどがあり微妙（※）ですし、タグにはアンダースコアが含まれていたりratingが付いていたりしますので（DeepDanbooruの場合）、エディタの置換機能などを用いてキャプションとタグのクリーニングをしたほうがいいでしょう。
+
+※たとえばアニメ絵の少女を学習する場合、キャプションにはgirl/girls/woman/womenなどのばらつきがあります。また「anime girl」なども単に「girl」としたほうが適切かもしれません。
+
+クリーニング用のスクリプトが用意してありますので、スクリプトの内容を状況に応じて編集してお使いください。
+
+（教師データフォルダの指定は不要になりました。メタデータ内の全データをクリーニングします。）
+
+```
+python clean_captions_and_tags.py <読み込むメタデータファイル名> <書き込むメタデータファイル名>
+```
+
+--in_jsonは付きませんのでご注意ください。たとえば次のようになります。
+
+```
+python clean_captions_and_tags.py meta_cap_dd.json meta_clean.json
+```
+
+以上でキャプションとタグの前処理は完了です。
+
+## latentsの事前取得
+
+※ このステップは必須ではありません。省略しても学習時にlatentsを取得しながら学習できます。
+また学習時に `random_crop` や `color_aug` などを行う場合にはlatentsの事前取得はできません（画像を毎回変えながら学習するため）。事前取得をしない場合、ここまでのメタデータで学習できます。
+
+あらかじめ画像の潜在表現を取得しディスクに保存しておきます。それにより、学習を高速に進めることができます。あわせてbucketing（教師データをアスペクト比に応じて分類する）を行います。
+
+作業フォルダで以下のように入力してください。
+```
+python prepare_buckets_latents.py --full_path <教師データフォルダ>  
+    <読み込むメタデータファイル名> <書き込むメタデータファイル名> 
+    <fine tuningするモデル名またはcheckpoint> 
+    --batch_size <バッチサイズ> 
+    --max_resolution <解像度 幅,高さ> 
+    --mixed_precision <精度>
+```
+
+モデルがmodel.ckpt、バッチサイズ4、学習解像度は512\*512、精度no（float32）で、meta_clean.jsonからメタデータを読み込み、meta_lat.jsonに書き込む場合、以下のようになります。
+
+```
+python prepare_buckets_latents.py --full_path 
+    train_data meta_clean.json meta_lat.json model.ckpt 
+    --batch_size 4 --max_resolution 512,512 --mixed_precision no
+```
+
+教師データフォルダにnumpyのnpz形式でlatentsが保存されます。
+
+解像度の最小サイズを--min_bucket_resoオプションで、最大サイズを--max_bucket_resoで指定できます。デフォルトはそれぞれ256、1024です。たとえば最小サイズに384を指定すると、256\*1024や320\*768などの解像度は使わなくなります。
+解像度を768\*768のように大きくした場合、最大サイズに1280などを指定すると良いでしょう。
+
+--flip_augオプションを指定すると左右反転のaugmentation（データ拡張）を行います。疑似的にデータ量を二倍に増やすことができますが、データが左右対称でない場合に指定すると（例えばキャラクタの外見、髪型など）学習がうまく行かなくなります。
+
+
+（反転した画像についてもlatentsを取得し、\*\_flip.npzファイルを保存する単純な実装です。fline_tune.pyには特にオプション指定は必要ありません。\_flip付きのファイルがある場合、flip付き・なしのファイルを、ランダムに読み込みます。）
+
+バッチサイズはVRAM 12GBでももう少し増やせるかもしれません。
+解像度は64で割り切れる数字で、"幅,高さ"で指定します。解像度はfine tuning時のメモリサイズに直結します。VRAM 12GBでは512,512が限界と思われます（※）。16GBなら512,704や512,768まで上げられるかもしれません。なお256,256等にしてもVRAM 8GBでは厳しいようです（パラメータやoptimizerなどは解像度に関係せず一定のメモリが必要なため）。
+
+※batch size 1の学習で12GB VRAM、640,640で動いたとの報告もありました。
+
+以下のようにbucketingの結果が表示されます。
+
+![bucketingの結果](https://user-images.githubusercontent.com/52813779/208911419-71c00fbb-2ce6-49d5-89b5-b78d7715e441.png)
+
+複数の教師データフォルダがある場合には、full_path引数を指定しつつ、それぞれのフォルダに対して実行してください。
+```
+python prepare_buckets_latents.py --full_path  
+    train_data1 meta_clean.json meta_lat1.json model.ckpt 
+    --batch_size 4 --max_resolution 512,512 --mixed_precision no
+
+python prepare_buckets_latents.py --full_path 
+    train_data2 meta_lat1.json meta_lat2.json model.ckpt 
+    --batch_size 4 --max_resolution 512,512 --mixed_precision no
+
+```
+読み込み元と書き込み先を同じにすることも可能ですが別々の方が安全です。
+
+__※引数を都度書き換えて、別のメタデータファイルに書き込むと安全です。__
+

train_db.py

@@ -0,0 +1,437 @@
+# DreamBooth training
+# XXX dropped option: fine_tune
+
+import gc
+import time
+import argparse
+import itertools
+import math
+import os
+import toml
+from multiprocessing import Value
+
+from tqdm import tqdm
+import torch
+from accelerate.utils import set_seed
+import diffusers
+from diffusers import DDPMScheduler
+
+import library.train_util as train_util
+import library.config_util as config_util
+from library.config_util import (
+    ConfigSanitizer,
+    BlueprintGenerator,
+)
+import library.custom_train_functions as custom_train_functions
+from library.custom_train_functions import apply_snr_weight, get_weighted_text_embeddings
+
+def train(args):
+    train_util.verify_training_args(args)
+    train_util.prepare_dataset_args(args, False)
+
+    cache_latents = args.cache_latents
+
+    if args.seed is not None:
+        set_seed(args.seed)  # 乱数系列を初期化する
+
+    tokenizer = train_util.load_tokenizer(args)
+
+    blueprint_generator = BlueprintGenerator(ConfigSanitizer(True, False, True))
+    if args.dataset_config is not None:
+        print(f"Load dataset config from {args.dataset_config}")
+        user_config = config_util.load_user_config(args.dataset_config)
+        ignored = ["train_data_dir", "reg_data_dir"]
+        if any(getattr(args, attr) is not None for attr in ignored):
+            print(
+                "ignore following options because config file is found: {0} / 設定ファイルが利用されるため以下のオプションは無視されます: {0}".format(
+                    ", ".join(ignored)
+                )
+            )
+    else:
+        user_config = {
+            "datasets": [
+                {"subsets": config_util.generate_dreambooth_subsets_config_by_subdirs(args.train_data_dir, args.reg_data_dir)}
+            ]
+        }
+
+    blueprint = blueprint_generator.generate(user_config, args, tokenizer=tokenizer)
+    train_dataset_group = config_util.generate_dataset_group_by_blueprint(blueprint.dataset_group)
+
+    current_epoch = Value("i", 0)
+    current_step = Value("i", 0)
+    ds_for_collater = train_dataset_group if args.max_data_loader_n_workers == 0 else None
+    collater = train_util.collater_class(current_epoch, current_step, ds_for_collater)
+
+    if args.no_token_padding:
+        train_dataset_group.disable_token_padding()
+
+    if args.debug_dataset:
+        train_util.debug_dataset(train_dataset_group)
+        return
+
+    if cache_latents:
+        assert (
+            train_dataset_group.is_latent_cacheable()
+        ), "when caching latents, either color_aug or random_crop cannot be used / latentをキャッシュするときはcolor_augとrandom_cropは使えません"
+
+    # acceleratorを準備する
+    print("prepare accelerator")
+
+    if args.gradient_accumulation_steps > 1:
+        print(
+            f"gradient_accumulation_steps is {args.gradient_accumulation_steps}. accelerate does not support gradient_accumulation_steps when training multiple models (U-Net and Text Encoder), so something might be wrong"
+        )
+        print(
+            f"gradient_accumulation_stepsが{args.gradient_accumulation_steps}に設定されています。accelerateは複数モデル（U-NetおよびText Encoder）の学習時にgradient_accumulation_stepsをサポートしていないため結果は未知数です"
+        )
+
+    accelerator, unwrap_model = train_util.prepare_accelerator(args)
+
+    # mixed precisionに対応した型を用意しておき適宜castする
+    weight_dtype, save_dtype = train_util.prepare_dtype(args)
+
+    # モデルを読み込む
+    text_encoder, vae, unet, load_stable_diffusion_format = train_util.load_target_model(args, weight_dtype)
+
+    # verify load/save model formats
+    if load_stable_diffusion_format:
+        src_stable_diffusion_ckpt = args.pretrained_model_name_or_path
+        src_diffusers_model_path = None
+    else:
+        src_stable_diffusion_ckpt = None
+        src_diffusers_model_path = args.pretrained_model_name_or_path
+
+    if args.save_model_as is None:
+        save_stable_diffusion_format = load_stable_diffusion_format
+        use_safetensors = args.use_safetensors
+    else:
+        save_stable_diffusion_format = args.save_model_as.lower() == "ckpt" or args.save_model_as.lower() == "safetensors"
+        use_safetensors = args.use_safetensors or ("safetensors" in args.save_model_as.lower())
+
+    # モデルに xformers とか memory efficient attention を組み込む
+    train_util.replace_unet_modules(unet, args.mem_eff_attn, args.xformers)
+
+    # 学習を準備する
+    if cache_latents:
+        vae.to(accelerator.device, dtype=weight_dtype)
+        vae.requires_grad_(False)
+        vae.eval()
+        with torch.no_grad():
+            train_dataset_group.cache_latents(vae, args.vae_batch_size, args.cache_latents_to_disk, accelerator.is_main_process)
+        vae.to("cpu")
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        gc.collect()
+
+        accelerator.wait_for_everyone()
+
+    # 学習を準備する：モデルを適切な状態にする
+    train_text_encoder = args.stop_text_encoder_training is None or args.stop_text_encoder_training >= 0
+    unet.requires_grad_(True)  # 念のため追加
+    text_encoder.requires_grad_(train_text_encoder)
+    if not train_text_encoder:
+        print("Text Encoder is not trained.")
+
+    if args.gradient_checkpointing:
+        unet.enable_gradient_checkpointing()
+        text_encoder.gradient_checkpointing_enable()
+
+    if not cache_latents:
+        vae.requires_grad_(False)
+        vae.eval()
+        vae.to(accelerator.device, dtype=weight_dtype)
+
+    # 学習に必要なクラスを準備する
+    print("prepare optimizer, data loader etc.")
+    if train_text_encoder:
+        trainable_params = itertools.chain(unet.parameters(), text_encoder.parameters())
+    else:
+        trainable_params = unet.parameters()
+
+    _, _, optimizer = train_util.get_optimizer(args, trainable_params)
+
+    # dataloaderを準備する
+    # DataLoaderのプロセス数：0はメインプロセスになる
+    n_workers = min(args.max_data_loader_n_workers, os.cpu_count() - 1)  # cpu_count-1 ただし最大で指定された数まで
+    train_dataloader = torch.utils.data.DataLoader(
+        train_dataset_group,
+        batch_size=1,
+        shuffle=True,
+        collate_fn=collater,
+        num_workers=n_workers,
+        persistent_workers=args.persistent_data_loader_workers,
+    )
+
+    # 学習ステップ数を計算する
+    if args.max_train_epochs is not None:
+        args.max_train_steps = args.max_train_epochs * math.ceil(
+            len(train_dataloader) / accelerator.num_processes / args.gradient_accumulation_steps
+        )
+        print(f"override steps. steps for {args.max_train_epochs} epochs is / 指定エポックまでのステップ数: {args.max_train_steps}")
+
+    # データセット側にも学習ステップを送信
+    train_dataset_group.set_max_train_steps(args.max_train_steps)
+
+    if args.stop_text_encoder_training is None:
+        args.stop_text_encoder_training = args.max_train_steps + 1  # do not stop until end
+
+    # lr schedulerを用意する TODO gradient_accumulation_stepsの扱いが何かおかしいかもしれない。後で確認する
+    lr_scheduler = train_util.get_scheduler_fix(args, optimizer, accelerator.num_processes)
+
+    # 実験的機能：勾配も含めたfp16学習を行う　モデル全体をfp16にする
+    if args.full_fp16:
+        assert (
+            args.mixed_precision == "fp16"
+        ), "full_fp16 requires mixed precision='fp16' / full_fp16を使う場合はmixed_precision='fp16'を指定してください。"
+        print("enable full fp16 training.")
+        unet.to(weight_dtype)
+        text_encoder.to(weight_dtype)
+
+    # acceleratorがなんかよろしくやってくれるらしい
+    if train_text_encoder:
+        unet, text_encoder, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
+            unet, text_encoder, optimizer, train_dataloader, lr_scheduler
+        )
+    else:
+        unet, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(unet, optimizer, train_dataloader, lr_scheduler)
+
+    if not train_text_encoder:
+        text_encoder.to(accelerator.device, dtype=weight_dtype)  # to avoid 'cpu' vs 'cuda' error
+
+    # 実験的機能：勾配も含めたfp16学習を行う　PyTorchにパッチを当ててfp16でのgrad scaleを有効にする
+    if args.full_fp16:
+        train_util.patch_accelerator_for_fp16_training(accelerator)
+
+    # resumeする
+    train_util.resume_from_local_or_hf_if_specified(accelerator, args)
+
+    # epoch数を計算する
+    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
+    num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)
+    if (args.save_n_epoch_ratio is not None) and (args.save_n_epoch_ratio > 0):
+        args.save_every_n_epochs = math.floor(num_train_epochs / args.save_n_epoch_ratio) or 1
+
+    # 学習する
+    total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
+    print("running training / 学習開始")
+    print(f"  num train images * repeats / 学習画像の数×繰り返し回数: {train_dataset_group.num_train_images}")
+    print(f"  num reg images / 正則化画像の数: {train_dataset_group.num_reg_images}")
+    print(f"  num batches per epoch / 1epochのバッチ数: {len(train_dataloader)}")
+    print(f"  num epochs / epoch数: {num_train_epochs}")
+    print(f"  batch size per device / バッチサイズ: {args.train_batch_size}")
+    print(f"  total train batch size (with parallel & distributed & accumulation) / 総バッチサイズ（並列学習、勾配合計含む）: {total_batch_size}")
+    print(f"  gradient ccumulation steps / 勾配を合計するステップ数 = {args.gradient_accumulation_steps}")
+    print(f"  total optimization steps / 学習ステップ数: {args.max_train_steps}")
+
+    progress_bar = tqdm(range(args.max_train_steps), smoothing=0, disable=not accelerator.is_local_main_process, desc="steps")
+    global_step = 0
+
+    noise_scheduler = DDPMScheduler(
+        beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000, clip_sample=False
+    )
+
+    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}")
+        current_epoch.value = epoch + 1
+
+        # 指定したステップ数までText Encoderを学習する：epoch最初の状態
+        unet.train()
+        # train==True is required to enable gradient_checkpointing
+        if args.gradient_checkpointing or global_step < args.stop_text_encoder_training:
+            text_encoder.train()
+
+        for step, batch in enumerate(train_dataloader):
+            current_step.value = global_step
+            # 指定したステップ数でText Encoderの学習を止める
+            if global_step == args.stop_text_encoder_training:
+                print(f"stop text encoder training at step {global_step}")
+                if not args.gradient_checkpointing:
+                    text_encoder.train(False)
+                text_encoder.requires_grad_(False)
+
+            with accelerator.accumulate(unet):
+                with torch.no_grad():
+                    # latentに変換
+                    if cache_latents:
+                        latents = batch["latents"].to(accelerator.device)
+                    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)
+                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):
+                    if args.weighted_captions:
+                      encoder_hidden_states = get_weighted_text_embeddings(tokenizer,
+                        text_encoder,
+                        batch["captions"],
+                        accelerator.device,
+                        args.max_token_length // 75 if args.max_token_length else 1,
+                        clip_skip=args.clip_skip,
+                        )
+                    else:
+                      input_ids = batch["input_ids"].to(accelerator.device)
+                      encoder_hidden_states = train_util.get_hidden_states(
+                          args, input_ids, tokenizer, text_encoder, None if not args.full_fp16 else weight_dtype
+                      )
+
+                # Sample a random timestep for each image
+                timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (b_size,), device=latents.device)
+                timesteps = timesteps.long()
+
+                # Add noise to the latents according to the noise magnitude at each timestep
+                # (this is the forward diffusion process)
+                noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)
+
+                # Predict the noise residual
+                noise_pred = unet(noisy_latents, timesteps, encoder_hidden_states).sample
+
+                if args.v_parameterization:
+                    # v-parameterization training
+                    target = noise_scheduler.get_velocity(latents, noise, timesteps)
+                else:
+                    target = noise
+
+                loss = torch.nn.functional.mse_loss(noise_pred.float(), target.float(), reduction="none")
+                loss = loss.mean([1, 2, 3])
+
+                loss_weights = batch["loss_weights"]  # 各sampleごとのweight
+                loss = loss * loss_weights
+
+                if args.min_snr_gamma:
+                    loss = apply_snr_weight(loss, timesteps, noise_scheduler, args.min_snr_gamma)
+
+                loss = loss.mean()  # 平均なのでbatch_sizeで割る必要なし
+
+                accelerator.backward(loss)
+                if accelerator.sync_gradients and args.max_grad_norm != 0.0:
+                    if train_text_encoder:
+                        params_to_clip = itertools.chain(unet.parameters(), text_encoder.parameters())
+                    else:
+                        params_to_clip = unet.parameters()
+                    accelerator.clip_grad_norm_(params_to_clip, args.max_grad_norm)
+
+                optimizer.step()
+                lr_scheduler.step()
+                optimizer.zero_grad(set_to_none=True)
+
+            # Checks if the accelerator has performed an optimization step behind the scenes
+            if accelerator.sync_gradients:
+                progress_bar.update(1)
+                global_step += 1
+
+                train_util.sample_images(
+                    accelerator, args, None, global_step, accelerator.device, vae, tokenizer, text_encoder, unet
+                )
+
+            current_loss = loss.detach().item()
+            if args.logging_dir is not None:
+                logs = {"loss": current_loss, "lr": float(lr_scheduler.get_last_lr()[0])}
+                if args.optimizer_type.lower() == "DAdaptation".lower():  # tracking d*lr value
+                    logs["lr/d*lr"] = (
+                        lr_scheduler.optimizers[0].param_groups[0]["d"] * lr_scheduler.optimizers[0].param_groups[0]["lr"]
+                    )
+                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 / len(loss_list)
+            logs = {"loss": avr_loss}  # , "lr": lr_scheduler.get_last_lr()[0]}
+            progress_bar.set_postfix(**logs)
+
+            if global_step >= args.max_train_steps:
+                break
+
+        if args.logging_dir is not None:
+            logs = {"loss/epoch": loss_total / len(loss_list)}
+            accelerator.log(logs, step=epoch + 1)
+
+        accelerator.wait_for_everyone()
+
+        if args.save_every_n_epochs is not None:
+            src_path = src_stable_diffusion_ckpt if save_stable_diffusion_format else src_diffusers_model_path
+            train_util.save_sd_model_on_epoch_end(
+                args,
+                accelerator,
+                src_path,
+                save_stable_diffusion_format,
+                use_safetensors,
+                save_dtype,
+                epoch,
+                num_train_epochs,
+                global_step,
+                unwrap_model(text_encoder),
+                unwrap_model(unet),
+                vae,
+            )
+
+        train_util.sample_images(accelerator, args, epoch + 1, global_step, accelerator.device, vae, tokenizer, text_encoder, unet)
+
+    is_main_process = accelerator.is_main_process
+    if is_main_process:
+        unet = unwrap_model(unet)
+        text_encoder = unwrap_model(text_encoder)
+
+    accelerator.end_training()
+
+    if args.save_state:
+        train_util.save_state_on_train_end(args, accelerator)
+
+    del accelerator  # この後メモリを使うのでこれは消す
+
+    if is_main_process:
+        src_path = src_stable_diffusion_ckpt if save_stable_diffusion_format else src_diffusers_model_path
+        train_util.save_sd_model_on_train_end(
+            args, src_path, save_stable_diffusion_format, use_safetensors, save_dtype, epoch, global_step, text_encoder, unet, vae
+        )
+        print("model saved.")
+
+
+def setup_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser()
+
+    train_util.add_sd_models_arguments(parser)
+    train_util.add_dataset_arguments(parser, True, False, True)
+    train_util.add_training_arguments(parser, True)
+    train_util.add_sd_saving_arguments(parser)
+    train_util.add_optimizer_arguments(parser)
+    config_util.add_config_arguments(parser)
+    custom_train_functions.add_custom_train_arguments(parser)
+
+    parser.add_argument(
+        "--no_token_padding",
+        action="store_true",
+        help="disable token padding (same as Diffuser's DreamBooth) / トークンのpaddingを無効にする（Diffusers版DreamBoothと同じ動作）",
+    )
+    parser.add_argument(
+        "--stop_text_encoder_training",
+        type=int,
+        default=None,
+        help="steps to stop text encoder training, -1 for no training / Text Encoderの学習を止めるステップ数、-1で最初から学習しない",
+    )
+
+    return parser
+
+
+if __name__ == "__main__":
+    parser = setup_parser()
+
+    args = parser.parse_args()
+    args = train_util.read_config_from_file(args, parser)
+
+    train(args)

train_db_README-ja.md

@@ -0,0 +1,167 @@
+DreamBoothのガイドです。
+
+[学習についての共通ドキュメント](./train_README-ja.md) もあわせてご覧ください。
+
+# 概要
+
+DreamBoothとは、画像生成モデルに特定の主題を追加学習し、それを特定の識別子で生成する技術です。[論文はこちら](https://arxiv.org/abs/2208.12242)。
+
+具体的には、Stable Diffusionのモデルにキャラや画風などを学ばせ、それを `shs` のような特定の単語で呼び出せる（生成画像に出現させる）ことができます。
+
+スクリプトは[DiffusersのDreamBooth](https://github.com/huggingface/diffusers/tree/main/examples/dreambooth)を元にしていますが、以下のような機能追加を行っています（いくつかの機能は元のスクリプト側もその後対応しています）。
+
+スクリプトの主な機能は以下の通りです。
+
+- 8bit Adam optimizerおよびlatentのキャッシュによる省メモリ化（[Shivam Shrirao氏版](https://github.com/ShivamShrirao/diffusers/tree/main/examples/dreambooth)と同様）。
+- xformersによる省メモリ化。
+- 512x512だけではなく任意サイズでの学習。
+- augmentationによる品質の向上。
+- DreamBoothだけではなくText Encoder+U-Netのfine tuningに対応。
+- Stable Diffusion形式でのモデルの読み書き。
+- Aspect Ratio Bucketing。
+- Stable Diffusion v2.0対応。
+
+# 学習の手順
+
+あらかじめこのリポジトリのREADMEを参照し、環境整備を行ってください。
+
+## データの準備
+
+[学習データの準備について](./train_README-ja.md) を参照してください。
+
+## 学習の実行
+
+スクリプトを実行します。最大限、メモリを節約したコマンドは以下のようになります（実際には1行で入力します）。それぞれの行を必要に応じて書き換えてください。12GB程度のVRAMで動作するようです。
+
+```
+accelerate launch --num_cpu_threads_per_process 1 train_db.py 
+    --pretrained_model_name_or_path=<.ckptまたは.safetensordまたはDiffusers版モデルのディレクトリ> 
+    --dataset_config=<データ準備で作成した.tomlファイル> 
+    --output_dir=<学習したモデルの出力先フォルダ>  
+    --output_name=<学習したモデル出力時のファイル名> 
+    --save_model_as=safetensors 
+    --prior_loss_weight=1.0 
+    --max_train_steps=1600 
+    --learning_rate=1e-6 
+    --optimizer_type="AdamW8bit" 
+    --xformers 
+    --mixed_precision="fp16" 
+    --cache_latents 
+    --gradient_checkpointing
+```
+
+`num_cpu_threads_per_process` には通常は1を指定するとよいようです。
+
+`pretrained_model_name_or_path` に追加学習を行う元となるモデルを指定します。Stable Diffusionのcheckpointファイル（.ckptまたは.safetensors）、Diffusersのローカルディスクにあるモデルディレクトリ、DiffusersのモデルID（"stabilityai/stable-diffusion-2"など）が指定できます。
+
+`output_dir` に学習後のモデルを保存するフォルダを指定します。`output_name` にモデルのファイル名を拡張子を除いて指定します。`save_model_as` でsafetensors形式での保存を指定しています。
+
+`dataset_config` に `.toml` ファイルを指定します。ファイル内でのバッチサイズ指定は、当初はメモリ消費を抑えるために `1` としてください。
+
+`prior_loss_weight` は正則化画像のlossの重みです。通常は1.0を指定します。
+
+学習させるステップ数 `max_train_steps` を1600とします。学習率 `learning_rate` はここでは1e-6を指定しています。
+
+省メモリ化のため `mixed_precision="fp16"` を指定します（RTX30 シリーズ以降では `bf16` も指定できます。環境整備時にaccelerateに行った設定と合わせてください）。また `gradient_checkpointing` を指定します。
+
+オプティマイザ（モデルを学習データにあうように最適化＝学習させるクラス）にメモリ消費の少ない 8bit AdamW を使うため、 `optimizer_type="AdamW8bit"` を指定します。
+
+`xformers` オプションを指定し、xformersのCrossAttentionを用います。xformersをインストールしていない場合やエラーとなる場合（環境にもよりますが `mixed_precision="no"` の場合など）、代わりに `mem_eff_attn` オプションを指定すると省メモリ版CrossAttentionを使用します（速度は遅くなります）。
+
+省メモリ化のため `cache_latents` オプションを指定してVAEの出力をキャッシュします。
+
+ある程度メモリがある場合は、`.toml` ファイルを編集してバッチサイズをたとえば `4` くらいに増やしてください（高速化と精度向上の可能性があります）。また `cache_latents` を外すことで augmentation が可能になります。
+
+### よく使われるオプションについて
+
+以下の場合には [学習の共通ドキュメント](./train_README-ja.md) の「よく使われるオプション」を参照してください。
+
+- Stable Diffusion 2.xまたはそこからの派生モデルを学習する
+- clip skipを2以上を前提としたモデルを学習する
+- 75トークンを超えたキャプションで学習する
+
+### DreamBoothでのステップ数について
+
+当スクリプトでは省メモリ化のため、ステップ当たりの学習回数が元のスクリプトの半分になっています（対象の画像と正則化画像を同一のバッチではなく別のバッチに分割して学習するため）。
+
+元のDiffusers版やXavierXiao氏のStable Diffusion版とほぼ同じ学習を行うには、ステップ数を倍にしてください。
+
+（学習画像と正則化画像をまとめてから shuffle するため厳密にはデータの順番が変わってしまいますが、学習には大きな影響はないと思います。）
+
+### DreamBoothでのバッチサイズについて
+
+モデル全体を学習するためLoRA等の学習に比べるとメモリ消費量は多くなります（fine tuningと同じ）。
+
+### 学習率について
+
+Diffusers版では5e-6ですがStable Diffusion版は1e-6ですので、上のサンプルでは1e-6を指定しています。
+
+### 以前の形式のデータセット指定をした場合のコマンドライン
+
+解像度やバッチサイズをオプションで指定します。コマンドラインの例は以下の通りです。
+
+```
+accelerate launch --num_cpu_threads_per_process 1 train_db.py 
+    --pretrained_model_name_or_path=<.ckptまたは.safetensordまたはDiffusers版モデルのディレクトリ> 
+    --train_data_dir=<学習用データのディレクトリ> 
+    --reg_data_dir=<正則化画像のディレクトリ> 
+    --output_dir=<学習したモデルの出力先ディレクトリ> 
+    --output_name=<学習したモデル出力時のファイル名> 
+    --prior_loss_weight=1.0 
+    --resolution=512 
+    --train_batch_size=1 
+    --learning_rate=1e-6 
+    --max_train_steps=1600 
+    --use_8bit_adam 
+    --xformers 
+    --mixed_precision="bf16" 
+    --cache_latents
+    --gradient_checkpointing
+```
+
+## 学習したモデルで画像生成する
+
+学習が終わると指定したフォルダに指定した名前でsafetensorsファイルが出力されます。
+
+v1.4/1.5およびその他の派生モデルの場合、このモデルでAutomatic1111氏のWebUIなどで推論できます。models\Stable-diffusionフォルダに置いてください。
+
+v2.xモデルでWebUIで画像生成する場合、モデルの仕様が記述された.yamlファイルが別途必要になります。v2.x baseの場合はv2-inference.yamlを、768/vの場合はv2-inference-v.yamlを、同じフォルダに置き、拡張子の前の部分をモデルと同じ名前にしてください。
+
+![image](https://user-images.githubusercontent.com/52813779/210776915-061d79c3-6582-42c2-8884-8b91d2f07313.png)
+
+各yamlファイルは[Stability AIのSD2.0のリポジトリ](https://github.com/Stability-AI/stablediffusion/tree/main/configs/stable-diffusion)にあります。
+
+# DreamBooth特有のその他の主なオプション
+
+すべてのオプションについては別文書を参照してください。
+
+## Text Encoderの学習を途中から行わない --stop_text_encoder_training
+
+stop_text_encoder_trainingオプションに数値を指定すると、そのステップ数以降はText Encoderの学習を行わずU-Netだけ学習します。場合によっては精度の向上が期待できるかもしれません。
+
+（恐らくText Encoderだけ先に過学習することがあり、それを防げるのではないかと推測していますが、詳細な影響は不明です。）
+
+## Tokenizerのパディングをしない --no_token_padding
+no_token_paddingオプションを指定するとTokenizerの出力をpaddingしません（Diffusers版の旧DreamBoothと同じ動きになります）。
+
+
+<!-- 
+bucketing（後述）を利用しかつaugmentation（後述）を使う場合の例は以下のようになります。
+
+```
+accelerate launch --num_cpu_threads_per_process 8 train_db.py 
+    --pretrained_model_name_or_path=<.ckptまたは.safetensordまたはDiffusers版モデルのディレクトリ> 
+    --train_data_dir=<学習用データのディレクトリ> 
+    --reg_data_dir=<正則化画像のディレクトリ> 
+    --output_dir=<学習したモデルの出力先ディレクトリ> 
+    --resolution=768,512 
+    --train_batch_size=20 --learning_rate=5e-6 --max_train_steps=800 
+    --use_8bit_adam --xformers --mixed_precision="bf16" 
+    --save_every_n_epochs=1 --save_state --save_precision="bf16" 
+    --logging_dir=logs 
+    --enable_bucket --min_bucket_reso=384 --max_bucket_reso=1280 
+    --color_aug --flip_aug --gradient_checkpointing --seed 42
+```
+
+
+-->

train_db_README.md

@@ -0,0 +1,295 @@
+A guide to DreamBooth. The same procedure is used for training additional networks such as LoRA.
+
+# overview
+
+The main functions of the script are as follows.
+
+- Memory saving by 8bit Adam optimizer and latent cache (similar to ShivamShirao's version).
+- Saved memory by xformers.
+- Study in any size, not just 512x512.
+- Quality improvement with augmentation.
+- Supports fine tuning of Text Encoder+U-Net as well as DreamBooth.
+- Read and write models in StableDiffusion format.
+- Aspect Ratio Bucketing.
+- Supports Stable Diffusion v2.0.
+
+# learning procedure
+
+## step 1. Environment improvement
+
+See the README in this repository.
+
+
+## step 2. Determine identifier and class
+
+Decide the word identifier that connects the target you want to learn and the class to which the target belongs.
+
+(There are various names such as instance, but for the time being I will stick to the original paper.)
+
+Here's a very brief explanation (look it up for more details).
+
+class is the general type to learn. For example, if you want to learn a specific breed of dog, the class will be dog. Anime characters will be boy, girl, 1boy or 1girl depending on the model.
+
+The identifier is for identifying and learning the learning target. Any word is fine, but according to the original paper, ``a rare word with 3 letters or less that becomes one token with tokinizer'' is good.
+
+By using the identifier and class to train the model, for example, "shs dog", you can learn by identifying the object you want to learn from the class.
+
+When generating an image, if you say "shs dog", an image of the learned dog breed will be generated.
+
+(For reference, the identifier I use these days is ``shs sts scs cpc coc cic msm usu ici lvl cic dii muk ori hru rik koo yos wny``.)
+
+## step 3. Prepare images for training
+Create a folder to store training images. __In addition, create a directory with the following name:
+
+```
+<repeat count>_<identifier> <class>
+```
+
+Don't forget the ``_`` between them.
+
+The number of repetitions is specified to match the number of regularized images (described later).
+
+For example, at the prompt "sls frog", to repeat the data 20 times, it would be "20_sls frog". It will be as follows.
+
+![image](https://user-images.githubusercontent.com/52813779/210770636-1c851377-5936-4c15-90b7-8ac8ad6c2074.png)
+
+## step 4. Preparing regularized images
+This is the procedure when using a regularized image. It is also possible to learn without using the regularization image (the whole target class is affected because it is impossible to distinguish without using the regularization image).
+
+Create a folder to store the regularized images. __In addition, __ create a directory named ``<repeat count>_<class>``.
+
+For example, with the prompt "frog" and without repeating the data (just once):
+
+![image](https://user-images.githubusercontent.com/52813779/210770897-329758e5-3675-49f1-b345-c135f1725832.png)
+
+Specify the number of iterations so that " __ number of iterations of training images x number of training images ≥ number of iterations of regularization images x number of regularization images __".
+
+(The number of data in one epoch is "number of repetitions of training images x number of training images". If the number of regularization images is more than that, the remaining regularization images will not be used.)
+
+## step 5. Run training
+Run the script. The maximally memory-saving command looks like this (actually typed on one line):
+
+*The command for learning additional networks such as LoRA is ``train_network.py`` instead of ``train_db.py``. You will also need additional network_\* options, so please refer to LoRA's guide.
+
+```
+accelerate launch --num_cpu_threads_per_process 8 train_db.py
+     --pretrained_model_name_or_path=<directory of .ckpt or .safetensord or Diffusers model>
+     --train_data_dir=<training data directory>
+     --reg_data_dir=<regularized image directory>
+     --output_dir=<output destination directory for trained model>
+     --prior_loss_weight=1.0
+     --resolution=512
+     --train_batch_size=1
+     --learning_rate=1e-6
+     --max_train_steps=1600
+     --use_8bit_adam
+     --xformers
+     --mixed_precision="bf16"
+     --cache_latents
+     --gradient_checkpointing
+```
+
+It seems to be good to specify the number of CPU cores for num_cpu_threads_per_process.
+
+Specify the model to perform additional training in pretrained_model_name_or_path. You can specify a Stable Diffusion checkpoint file (.ckpt or .safetensors), a model directory on the Diffusers local disk, or a Diffusers model ID (such as "stabilityai/stable-diffusion-2"). The saved model after training will be saved in the same format as the original model by default (can be changed with the save_model_as option).
+
+prior_loss_weight is the loss weight of the regularized image. Normally, specify 1.0.
+
+resolution will be the size of the image (resolution, width and height). If bucketing (described later) is not used, use this size for training images and regularization images.
+
+train_batch_size is the training batch size. Set max_train_steps to 1600. The learning rate learning_rate is 5e-6 in the diffusers version and 1e-6 in the StableDiffusion version, so 1e-6 is specified here.
+
+Specify mixed_precision="bf16" (or "fp16") and gradient_checkpointing for memory saving.
+
+Specify the xformers option and use xformers' CrossAttention. If you don't have xformers installed, if you get an error (without mixed_precision, it was an error in my environment), specify the mem_eff_attn option instead to use the memory-saving version of CrossAttention (speed will be slower) .
+
+Cache VAE output with cache_latents option to save memory.
+
+If you have a certain amount of memory, specify it as follows, for example.
+
+```
+accelerate launch --num_cpu_threads_per_process 8 train_db.py
+     --pretrained_model_name_or_path=<directory of .ckpt or .safetensord or Diffusers model>
+     --train_data_dir=<training data directory>
+     --reg_data_dir=<regularized image directory>
+     --output_dir=<output destination directory for trained model>
+     --prior_loss_weight=1.0
+     --resolution=512
+     --train_batch_size=4
+     --learning_rate=1e-6
+     --max_train_steps=400
+     --use_8bit_adam
+     --xformers
+     --mixed_precision="bf16"
+     --cache_latents
+```
+
+Remove gradient_checkpointing to speed up (memory usage will increase). Increase the batch size to improve speed and accuracy.
+
+An example of using bucketing (see below) and using augmentation (see below) looks like this:
+
+```
+accelerate launch --num_cpu_threads_per_process 8 train_db.py
+     --pretrained_model_name_or_path=<directory of .ckpt or .safetensord or Diffusers model>
+     --train_data_dir=<training data directory>
+     --reg_data_dir=<regularized image directory>
+     --output_dir=<output destination directory for trained model>
+     --resolution=768,512
+     --train_batch_size=20 --learning_rate=5e-6 --max_train_steps=800
+     --use_8bit_adam --xformers --mixed_precision="bf16"
+     --save_every_n_epochs=1 --save_state --save_precision="bf16"
+     --logging_dir=logs
+     --enable_bucket --min_bucket_reso=384 --max_bucket_reso=1280
+     --color_aug --flip_aug --gradient_checkpointing --seed 42
+```
+
+### About the number of steps
+To save memory, the number of training steps per step is half that of train_drebooth.py (because the target image and the regularization image are divided into different batches instead of the same batch).
+Double the number of steps to get almost the same training as the original Diffusers version and XavierXiao's StableDiffusion version.
+
+(Strictly speaking, the order of the data changes due to shuffle=True, but I don't think it has a big impact on learning.)
+
+## Generate an image with the trained model
+
+Name last.ckpt in the specified folder when learning is completed will output the checkpoint (if you learned the DiffUsers version model, it will be the last folder).
+
+For v1.4/1.5 and other derived models, this model can be inferred by Automatic1111's WebUI, etc. Place it in the models\Stable-diffusion folder.
+
+When generating images with WebUI with the v2.x model, a separate .yaml file that describes the model specifications is required. Place v2-inference.yaml for v2.x base and v2-inference-v.yaml for 768/v in the same folder and make the part before the extension the same name as the model.
+
+![image](https://user-images.githubusercontent.com/52813779/210776915-061d79c3-6582-42c2-8884-8b91d2f07313.png)
+
+Each yaml file can be found at [https://github.com/Stability-AI/stablediffusion/tree/main/configs/stable-diffusion] (Stability AI SD2.0 repository).
+
+# Other study options
+
+## Supports Stable Diffusion 2.0 --v2 / --v_parameterization
+Specify the v2 option when using Hugging Face's stable-diffusion-2-base, and specify both the v2 and v_parameterization options when using stable-diffusion-2 or 768-v-ema.ckpt.
+
+In addition, learning SD 2.0 seems to be difficult with VRAM 12GB because the Text Encoder is getting bigger.
+
+The following points have changed significantly in Stable Diffusion 2.0.
+
+1. Tokenizer to use
+2. Which Text Encoder to use and which output layer to use (2.0 uses the penultimate layer)
+3. Output dimensionality of Text Encoder (768->1024)
+4. Structure of U-Net (number of heads of CrossAttention, etc.)
+5. v-parameterization (the sampling method seems to have changed)
+
+Among these, 1 to 4 are adopted for base, and 1 to 5 are adopted for the one without base (768-v). Enabling 1-4 is the v2 option, and enabling 5 is the v_parameterization option.
+
+## check training data --debug_dataset
+By adding this option, you can check what kind of image data and captions will be learned in advance before learning. Press Esc to exit and return to the command line.
+
+*Please note that it seems to hang when executed in an environment where there is no screen such as Colab.
+
+## Stop training Text Encoder --stop_text_encoder_training
+If you specify a numerical value for the stop_text_encoder_training option, after that number of steps, only the U-Net will be trained without training the Text Encoder. In some cases, the accuracy may be improved.
+
+(Probably only the Text Encoder may overfit first, and I guess that it can be prevented, but the detailed impact is unknown.)
+
+## Load and learn VAE separately --vae
+If you specify either a Stable Diffusion checkpoint, a VAE checkpoint file, a Diffuses model, or a VAE (both of which can specify a local or Hugging Face model ID) in the vae option, that VAE is used for learning (latents when caching or getting latents during learning).
+The saved model will incorporate this VAE.
+
+## save during learning --save_every_n_epochs / --save_state / --resume
+Specifying a number for the save_every_n_epochs option saves the model during training every epoch.
+
+If you specify the save_state option at the same time, the learning state including the state of the optimizer etc. will be saved together (compared to restarting learning from the checkpoint, you can expect to improve accuracy and shorten the learning time). The learning state is output in a folder named "epoch-??????-state" (?????? is the number of epochs) in the destination folder. Please use it when studying for a long time.
+
+Use the resume option to resume training from a saved training state. Please specify the learning state folder.
+
+Please note that due to the specifications of Accelerator (?), the number of epochs and global step are not saved, and it will start from 1 even when you resume.
+
+## No tokenizer padding --no_token_padding
+The no_token_padding option does not pad the output of the Tokenizer (same behavior as Diffusers version of old DreamBooth).
+
+## Training with arbitrary size images --resolution
+You can study outside the square. Please specify "width, height" like "448,640" in resolution. Width and height must be divisible by 64. Match the size of the training image and the regularization image.
+
+Personally, I often generate vertically long images, so I sometimes learn with "448, 640".
+
+## Aspect Ratio Bucketing --enable_bucket / --min_bucket_reso / --max_bucket_reso
+It is enabled by specifying the enable_bucket option. Stable Diffusion is trained at 512x512, but also at resolutions such as 256x768 and 384x640.
+
+If you specify this option, you do not need to unify the training images and regularization images to a specific resolution. Choose from several resolutions (aspect ratios) and learn at that resolution.
+Since the resolution is 64 pixels, the aspect ratio may not be exactly the same as the original image.
+
+You can specify the minimum size of the resolution with the min_bucket_reso option and the maximum size with the max_bucket_reso. The defaults are 256 and 1024 respectively.
+For example, specifying a minimum size of 384 will not use resolutions such as 256x1024 or 320x768.
+If you increase the resolution to 768x768, you may want to specify 1280 as the maximum size.
+
+When Aspect Ratio Bucketing is enabled, it may be better to prepare regularization images with various resolutions that are similar to the training images.
+
+(Because the images in one batch are not biased toward training images and regularization images.
+
+## augmentation --color_aug / --flip_aug
+Augmentation is a method of improving model performance by dynamically changing data during learning. Learn while subtly changing the hue with color_aug and flipping left and right with flip_aug.
+
+Since the data changes dynamically, it cannot be specified together with the cache_latents option.
+
+## Specify data precision when saving --save_precision
+Specifying float, fp16, or bf16 as the save_precision option will save the checkpoint in that format (only when saving in Stable Diffusion format). Please use it when you want to reduce the size of checkpoint.
+
+## save in any format --save_model_as
+Specify the save format of the model. Specify one of ckpt, safetensors, diffusers, diffusers_safetensors.
+
+When reading Stable Diffusion format (ckpt or safetensors) and saving in Diffusers format, missing information is supplemented by dropping v1.5 or v2.1 information from Hugging Face.
+
+## Save learning log --logging_dir / --log_prefix
+Specify the log save destination folder in the logging_dir option. Logs in TensorBoard format are saved.
+
+For example, if you specify --logging_dir=logs, a logs folder will be created in your working folder, and logs will be saved in the date/time folder.
+Also, if you specify the --log_prefix option, the specified string will be added before the date and time. Use "--logging_dir=logs --log_prefix=db_style1_" for identification.
+
+To check the log with TensorBoard, open another command prompt and enter the following in the working folder (I think tensorboard is installed when Diffusers is installed, but if it is not installed, pip install Please put it in tensorboard).
+
+```
+tensorboard --logdir=logs
+```
+
+Then open your browser and go to http://localhost:6006/ to see it.
+
+## scheduler related specification of learning rate --lr_scheduler / --lr_warmup_steps
+You can choose the learning rate scheduler from linear, cosine, cosine_with_restarts, polynomial, constant, constant_with_warmup with the lr_scheduler option. Default is constant. With lr_warmup_steps, you can specify the number of steps to warm up the scheduler (gradually changing the learning rate). Please do your own research for details.
+
+## Training with fp16 gradient (experimental feature) --full_fp16
+The full_fp16 option will change the gradient from normal float32 to float16 (fp16) and learn (it seems to be full fp16 learning instead of mixed precision).
+As a result, it seems that the SD1.x 512x512 size can be learned with a VRAM usage of less than 8GB, and the SD2.x 512x512 size can be learned with a VRAM usage of less than 12GB.
+
+Specify fp16 in the accelerate config beforehand and optionally set ``mixed_precision="fp16"`` (bf16 does not work).
+
+To minimize memory usage, use xformers, use_8bit_adam, cache_latents, gradient_checkpointing options and set train_batch_size to 1.
+
+(If you can afford it, increasing the train_batch_size step by step should improve the accuracy a little.)
+
+It is realized by patching the PyTorch source (confirmed with PyTorch 1.12.1 and 1.13.0). Accuracy will drop considerably, and the probability of learning failure on the way will also increase.
+The setting of the learning rate and the number of steps seems to be severe. Please be aware of them and use them at your own risk.
+
+# Other learning methods
+
+## Learning multiple classes, multiple identifiers
+The method is simple, multiple folders with ``Repetition count_<identifier> <class>`` in the training image folder, and a folder with ``Repetition count_<class>`` in the regularization image folder. Please prepare multiple
+
+For example, learning "sls frog" and "cpc rabbit" at the same time would look like this:
+
+![image](https://user-images.githubusercontent.com/52813779/210777933-a22229db-b219-4cd8-83ca-e87320fc4192.png)
+
+If you have one class and multiple targets, you can have only one regularized image folder. For example, if 1girl has character A and character B, do as follows.
+
+- train_girls
+   - 10_sls 1girl
+   - 10_cpc 1girl
+- reg_girls
+   -1_1girl
+
+If the number of data varies, it seems that good results can be obtained by adjusting the number of repetitions to unify the number of sheets for each class and identifier.
+
+## Use captions in DreamBooth
+If you put a file with the same file name as the image and the extension .caption (you can change it in the option) in the training image and regularization image folders, the caption will be read from that file and learned as a prompt.
+
+* The folder name (identifier class) will no longer be used for training those images.
+
+Adding captions to each image (you can use BLIP, etc.) may help clarify the attributes you want to learn.
+
+Caption files have a .caption extension by default. You can change it with --caption_extension. With the --shuffle_caption option, study captions during learning while shuffling each part separated by commas.

train_network.py

@@ -0,0 +1,773 @@
+from torch.nn.parallel import DistributedDataParallel as DDP
+import importlib
+import argparse
+import gc
+import math
+import os
+import random
+import time
+import json
+import toml
+from multiprocessing import Value
+
+from tqdm import tqdm
+import torch
+from accelerate.utils import set_seed
+from diffusers import DDPMScheduler
+
+import library.train_util as train_util
+from library.train_util import (
+    DreamBoothDataset,
+)
+import library.config_util as config_util
+from library.config_util import (
+    ConfigSanitizer,
+    BlueprintGenerator,
+)
+import library.huggingface_util as huggingface_util
+import library.custom_train_functions as custom_train_functions
+from library.custom_train_functions import apply_snr_weight, get_weighted_text_embeddings
+
+
+# TODO 他のスクリプトと共通化する
+def generate_step_logs(args: argparse.Namespace, current_loss, avr_loss, lr_scheduler):
+    logs = {"loss/current": current_loss, "loss/average": avr_loss}
+
+    lrs = lr_scheduler.get_last_lr()
+
+    if args.network_train_text_encoder_only or len(lrs) <= 2:  # not block lr (or single block)
+        if args.network_train_unet_only:
+            logs["lr/unet"] = float(lrs[0])
+        elif args.network_train_text_encoder_only:
+            logs["lr/textencoder"] = float(lrs[0])
+        else:
+            logs["lr/textencoder"] = float(lrs[0])
+            logs["lr/unet"] = float(lrs[-1])  # may be same to textencoder
+
+        if args.optimizer_type.lower() == "DAdaptation".lower():  # tracking d*lr value of unet.
+            logs["lr/d*lr"] = lr_scheduler.optimizers[-1].param_groups[0]["d"] * lr_scheduler.optimizers[-1].param_groups[0]["lr"]
+    else:
+        idx = 0
+        if not args.network_train_unet_only:
+            logs["lr/textencoder"] = float(lrs[0])
+            idx = 1
+
+        for i in range(idx, len(lrs)):
+            logs[f"lr/group{i}"] = float(lrs[i])
+            if args.optimizer_type.lower() == "DAdaptation".lower():
+                logs[f"lr/d*lr/group{i}"] = (
+                    lr_scheduler.optimizers[-1].param_groups[i]["d"] * lr_scheduler.optimizers[-1].param_groups[i]["lr"]
+                )
+
+    return logs
+
+
+def train(args):
+    session_id = random.randint(0, 2**32)
+    training_started_at = time.time()
+    train_util.verify_training_args(args)
+    train_util.prepare_dataset_args(args, True)
+
+    cache_latents = args.cache_latents
+    use_dreambooth_method = args.in_json is None
+    use_user_config = args.dataset_config is not None
+
+    if args.seed is None:
+        args.seed = random.randint(0, 2**32)
+    set_seed(args.seed)
+
+    tokenizer = train_util.load_tokenizer(args)
+
+    # データセットを準備する
+    blueprint_generator = BlueprintGenerator(ConfigSanitizer(True, True, True))
+    if use_user_config:
+        print(f"Load dataset config from {args.dataset_config}")
+        user_config = config_util.load_user_config(args.dataset_config)
+        ignored = ["train_data_dir", "reg_data_dir", "in_json"]
+        if any(getattr(args, attr) is not None for attr in ignored):
+            print(
+                "ignore following options because config file is found: {0} / 設定ファイルが利用されるため以下のオプションは無視されます: {0}".format(
+                    ", ".join(ignored)
+                )
+            )
+    else:
+        if use_dreambooth_method:
+            print("Use DreamBooth method.")
+            user_config = {
+                "datasets": [
+                    {"subsets": config_util.generate_dreambooth_subsets_config_by_subdirs(args.train_data_dir, args.reg_data_dir)}
+                ]
+            }
+        else:
+            print("Train with captions.")
+            user_config = {
+                "datasets": [
+                    {
+                        "subsets": [
+                            {
+                                "image_dir": args.train_data_dir,
+                                "metadata_file": args.in_json,
+                            }
+                        ]
+                    }
+                ]
+            }
+
+    blueprint = blueprint_generator.generate(user_config, args, tokenizer=tokenizer)
+    train_dataset_group = config_util.generate_dataset_group_by_blueprint(blueprint.dataset_group)
+
+    current_epoch = Value("i", 0)
+    current_step = Value("i", 0)
+    ds_for_collater = train_dataset_group if args.max_data_loader_n_workers == 0 else None
+    collater = train_util.collater_class(current_epoch, current_step, ds_for_collater)
+
+    if args.debug_dataset:
+        train_util.debug_dataset(train_dataset_group)
+        return
+    if len(train_dataset_group) == 0:
+        print(
+            "No data found. Please verify arguments (train_data_dir must be the parent of folders with images) / 画像がありません。引数指定を確認してください（train_data_dirには画像があるフォルダではなく、画像があるフォルダの親フォルダを指定する必要があります）"
+        )
+        return
+
+    if cache_latents:
+        assert (
+            train_dataset_group.is_latent_cacheable()
+        ), "when caching latents, either color_aug or random_crop cannot be used / latentをキャッシュするときはcolor_augとrandom_cropは使えません"
+
+    # acceleratorを準備する
+    print("prepare accelerator")
+    accelerator, unwrap_model = train_util.prepare_accelerator(args)
+    is_main_process = accelerator.is_main_process
+
+    # mixed precisionに対応した型を用意しておき適宜castする
+    weight_dtype, save_dtype = train_util.prepare_dtype(args)
+
+    # モデルを読み込む
+    for pi in range(accelerator.state.num_processes):
+        # TODO: modify other training scripts as well
+        if pi == accelerator.state.local_process_index:
+            print(f"loading model for process {accelerator.state.local_process_index}/{accelerator.state.num_processes}")
+
+            text_encoder, vae, unet, _ = train_util.load_target_model(
+                args, weight_dtype, accelerator.device if args.lowram else "cpu"
+            )
+
+            # work on low-ram device
+            if args.lowram:
+                text_encoder.to(accelerator.device)
+                unet.to(accelerator.device)
+                vae.to(accelerator.device)
+
+            gc.collect()
+            torch.cuda.empty_cache()
+        accelerator.wait_for_everyone()
+
+    # モデルに xformers とか memory efficient attention を組み込む
+    train_util.replace_unet_modules(unet, args.mem_eff_attn, args.xformers)
+
+    # 学習を準備する
+    if cache_latents:
+        vae.to(accelerator.device, dtype=weight_dtype)
+        vae.requires_grad_(False)
+        vae.eval()
+        with torch.no_grad():
+            train_dataset_group.cache_latents(vae, args.vae_batch_size, args.cache_latents_to_disk, accelerator.is_main_process)
+        vae.to("cpu")
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        gc.collect()
+
+        accelerator.wait_for_everyone()
+
+    # prepare network
+    import sys
+
+    sys.path.append(os.path.dirname(__file__))
+    print("import network module:", args.network_module)
+    network_module = importlib.import_module(args.network_module)
+
+    net_kwargs = {}
+    if args.network_args is not None:
+        for net_arg in args.network_args:
+            key, value = net_arg.split("=")
+            net_kwargs[key] = value
+
+    # if a new network is added in future, add if ~ then blocks for each network (;'∀')
+    network = network_module.create_network(1.0, args.network_dim, args.network_alpha, vae, text_encoder, unet, **net_kwargs)
+    if network is None:
+        return
+
+    if hasattr(network, "prepare_network"):
+        network.prepare_network(args)
+
+    train_unet = not args.network_train_text_encoder_only
+    train_text_encoder = not args.network_train_unet_only
+    network.apply_to(text_encoder, unet, train_text_encoder, train_unet)
+
+    if args.network_weights is not None:
+        info = network.load_weights(args.network_weights)
+        print(f"load network weights from {args.network_weights}: {info}")
+
+    if args.gradient_checkpointing:
+        unet.enable_gradient_checkpointing()
+        text_encoder.gradient_checkpointing_enable()
+        network.enable_gradient_checkpointing()  # may have no effect
+
+    # 学習に必要なクラスを準備する
+    print("prepare optimizer, data loader etc.")
+
+    # 後方互換性を確保するよ
+    try:
+        trainable_params = network.prepare_optimizer_params(args.text_encoder_lr, args.unet_lr, args.learning_rate)
+    except TypeError:
+        print(
+            "Deprecated: use prepare_optimizer_params(text_encoder_lr, unet_lr, learning_rate) instead of prepare_optimizer_params(text_encoder_lr, unet_lr)"
+        )
+        trainable_params = network.prepare_optimizer_params(args.text_encoder_lr, args.unet_lr)
+
+    optimizer_name, optimizer_args, optimizer = train_util.get_optimizer(args, trainable_params)
+
+    # dataloaderを準備する
+    # DataLoaderのプロセス数：0はメインプロセスになる
+    n_workers = min(args.max_data_loader_n_workers, os.cpu_count() - 1)  # cpu_count-1 ただし最大で指定された数まで
+
+    train_dataloader = torch.utils.data.DataLoader(
+        train_dataset_group,
+        batch_size=1,
+        shuffle=True,
+        collate_fn=collater,
+        num_workers=n_workers,
+        persistent_workers=args.persistent_data_loader_workers,
+    )
+
+    # 学習ステップ数を計算する
+    if args.max_train_epochs is not None:
+        args.max_train_steps = args.max_train_epochs * math.ceil(
+            len(train_dataloader) / accelerator.num_processes / args.gradient_accumulation_steps
+        )
+        if is_main_process:
+            print(f"override steps. steps for {args.max_train_epochs} epochs is / 指定エポックまでのステップ数: {args.max_train_steps}")
+
+    # データセット側にも学習ステップを送信
+    train_dataset_group.set_max_train_steps(args.max_train_steps)
+
+    # lr schedulerを用意する
+    lr_scheduler = train_util.get_scheduler_fix(args, optimizer, accelerator.num_processes)
+
+    # 実験的機能：勾配も含めたfp16学習を行う　モデル全体をfp16にする
+    if args.full_fp16:
+        assert (
+            args.mixed_precision == "fp16"
+        ), "full_fp16 requires mixed precision='fp16' / full_fp16を使う場合はmixed_precision='fp16'を指定してください。"
+        print("enable full fp16 training.")
+        network.to(weight_dtype)
+
+    # acceleratorがなんかよろしくやってくれるらしい
+    if train_unet and train_text_encoder:
+        unet, text_encoder, network, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
+            unet, text_encoder, network, optimizer, train_dataloader, lr_scheduler
+        )
+    elif train_unet:
+        unet, network, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
+            unet, network, optimizer, train_dataloader, lr_scheduler
+        )
+    elif train_text_encoder:
+        text_encoder, network, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
+            text_encoder, network, optimizer, train_dataloader, lr_scheduler
+        )
+    else:
+        network, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(network, optimizer, train_dataloader, lr_scheduler)
+
+    unet.requires_grad_(False)
+    unet.to(accelerator.device, dtype=weight_dtype)
+    text_encoder.requires_grad_(False)
+    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
+        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:
+        vae.requires_grad_(False)
+        vae.eval()
+        vae.to(accelerator.device, dtype=weight_dtype)
+
+    # 実験的機能：勾配も含めたfp16学習を行う　PyTorchにパッチを当ててfp16でのgrad scaleを有効にする
+    if args.full_fp16:
+        train_util.patch_accelerator_for_fp16_training(accelerator)
+
+    # resumeする
+    train_util.resume_from_local_or_hf_if_specified(accelerator, args)
+
+    # epoch数を計算する
+    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
+    num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)
+    if (args.save_n_epoch_ratio is not None) and (args.save_n_epoch_ratio > 0):
+        args.save_every_n_epochs = math.floor(num_train_epochs / args.save_n_epoch_ratio) or 1
+
+    # 学習する
+    # TODO: find a way to handle total batch size when there are multiple datasets
+    total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
+
+    if is_main_process:
+        print("running training / 学習開始")
+        print(f"  num train images * repeats / 学習画像の数×繰り返し回数: {train_dataset_group.num_train_images}")
+        print(f"  num reg images / 正則化画像の数: {train_dataset_group.num_reg_images}")
+        print(f"  num batches per epoch / 1epochのバッチ数: {len(train_dataloader)}")
+        print(f"  num epochs / epoch数: {num_train_epochs}")
+        print(f"  batch size per device / バッチサイズ: {', '.join([str(d.batch_size) for d in train_dataset_group.datasets])}")
+        # print(f"  total train batch size (with parallel & distributed & accumulation) / 総バッチサイズ（並列学習、勾配合計含む）: {total_batch_size}")
+        print(f"  gradient accumulation steps / 勾配を合計するステップ数 = {args.gradient_accumulation_steps}")
+        print(f"  total optimization steps / 学習ステップ数: {args.max_train_steps}")
+
+    # TODO refactor metadata creation and move to util
+    metadata = {
+        "ss_session_id": session_id,  # random integer indicating which group of epochs the model came from
+        "ss_training_started_at": training_started_at,  # unix timestamp
+        "ss_output_name": args.output_name,
+        "ss_learning_rate": args.learning_rate,
+        "ss_text_encoder_lr": args.text_encoder_lr,
+        "ss_unet_lr": args.unet_lr,
+        "ss_num_train_images": train_dataset_group.num_train_images,
+        "ss_num_reg_images": train_dataset_group.num_reg_images,
+        "ss_num_batches_per_epoch": len(train_dataloader),
+        "ss_num_epochs": num_train_epochs,
+        "ss_gradient_checkpointing": args.gradient_checkpointing,
+        "ss_gradient_accumulation_steps": args.gradient_accumulation_steps,
+        "ss_max_train_steps": args.max_train_steps,
+        "ss_lr_warmup_steps": args.lr_warmup_steps,
+        "ss_lr_scheduler": args.lr_scheduler,
+        "ss_network_module": args.network_module,
+        "ss_network_dim": args.network_dim,  # None means default because another network than LoRA may have another default dim
+        "ss_network_alpha": args.network_alpha,  # some networks may not use this value
+        "ss_mixed_precision": args.mixed_precision,
+        "ss_full_fp16": bool(args.full_fp16),
+        "ss_v2": bool(args.v2),
+        "ss_clip_skip": args.clip_skip,
+        "ss_max_token_length": args.max_token_length,
+        "ss_cache_latents": bool(args.cache_latents),
+        "ss_seed": args.seed,
+        "ss_lowram": args.lowram,
+        "ss_noise_offset": args.noise_offset,
+        "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 + (f"({optimizer_args})" if len(optimizer_args) > 0 else ""),
+        "ss_max_grad_norm": args.max_grad_norm,
+        "ss_caption_dropout_rate": args.caption_dropout_rate,
+        "ss_caption_dropout_every_n_epochs": args.caption_dropout_every_n_epochs,
+        "ss_caption_tag_dropout_rate": args.caption_tag_dropout_rate,
+        "ss_face_crop_aug_range": args.face_crop_aug_range,
+        "ss_prior_loss_weight": args.prior_loss_weight,
+        "ss_min_snr_gamma": args.min_snr_gamma,
+    }
+
+    if use_user_config:
+        # save metadata of multiple datasets
+        # NOTE: pack "ss_datasets" value as json one time
+        #   or should also pack nested collections as json?
+        datasets_metadata = []
+        tag_frequency = {}  # merge tag frequency for metadata editor
+        dataset_dirs_info = {}  # merge subset dirs for metadata editor
+
+        for dataset in train_dataset_group.datasets:
+            is_dreambooth_dataset = isinstance(dataset, DreamBoothDataset)
+            dataset_metadata = {
+                "is_dreambooth": is_dreambooth_dataset,
+                "batch_size_per_device": dataset.batch_size,
+                "num_train_images": dataset.num_train_images,  # includes repeating
+                "num_reg_images": dataset.num_reg_images,
+                "resolution": (dataset.width, dataset.height),
+                "enable_bucket": bool(dataset.enable_bucket),
+                "min_bucket_reso": dataset.min_bucket_reso,
+                "max_bucket_reso": dataset.max_bucket_reso,
+                "tag_frequency": dataset.tag_frequency,
+                "bucket_info": dataset.bucket_info,
+            }
+
+            subsets_metadata = []
+            for subset in dataset.subsets:
+                subset_metadata = {
+                    "img_count": subset.img_count,
+                    "num_repeats": subset.num_repeats,
+                    "color_aug": bool(subset.color_aug),
+                    "flip_aug": bool(subset.flip_aug),
+                    "random_crop": bool(subset.random_crop),
+                    "shuffle_caption": bool(subset.shuffle_caption),
+                    "keep_tokens": subset.keep_tokens,
+                }
+
+                image_dir_or_metadata_file = None
+                if subset.image_dir:
+                    image_dir = os.path.basename(subset.image_dir)
+                    subset_metadata["image_dir"] = image_dir
+                    image_dir_or_metadata_file = image_dir
+
+                if is_dreambooth_dataset:
+                    subset_metadata["class_tokens"] = subset.class_tokens
+                    subset_metadata["is_reg"] = subset.is_reg
+                    if subset.is_reg:
+                        image_dir_or_metadata_file = None  # not merging reg dataset
+                else:
+                    metadata_file = os.path.basename(subset.metadata_file)
+                    subset_metadata["metadata_file"] = metadata_file
+                    image_dir_or_metadata_file = metadata_file  # may overwrite
+
+                subsets_metadata.append(subset_metadata)
+
+                # merge dataset dir: not reg subset only
+                # TODO update additional-network extension to show detailed dataset config from metadata
+                if image_dir_or_metadata_file is not None:
+                    # datasets may have a certain dir multiple times
+                    v = image_dir_or_metadata_file
+                    i = 2
+                    while v in dataset_dirs_info:
+                        v = image_dir_or_metadata_file + f" ({i})"
+                        i += 1
+                    image_dir_or_metadata_file = v
+
+                    dataset_dirs_info[image_dir_or_metadata_file] = {"n_repeats": subset.num_repeats, "img_count": subset.img_count}
+
+            dataset_metadata["subsets"] = subsets_metadata
+            datasets_metadata.append(dataset_metadata)
+
+            # merge tag frequency:
+            for ds_dir_name, ds_freq_for_dir in dataset.tag_frequency.items():
+                # あるディレクトリが複数のdatasetで使用されている場合、一度だけ数える
+                # もともと繰り返し回数を指定しているので、キャプション内でのタグの出現回数と、それが学習で何度使われるかは一致しない
+                # なので、ここで複数datasetの回数を合算してもあまり意味はない
+                if ds_dir_name in tag_frequency:
+                    continue
+                tag_frequency[ds_dir_name] = ds_freq_for_dir
+
+        metadata["ss_datasets"] = json.dumps(datasets_metadata)
+        metadata["ss_tag_frequency"] = json.dumps(tag_frequency)
+        metadata["ss_dataset_dirs"] = json.dumps(dataset_dirs_info)
+    else:
+        # conserving backward compatibility when using train_dataset_dir and reg_dataset_dir
+        assert (
+            len(train_dataset_group.datasets) == 1
+        ), f"There should be a single dataset but {len(train_dataset_group.datasets)} found. This seems to be a bug. / データセットは1個だけ存在するはずですが、実際には{len(train_dataset_group.datasets)}個でした。プログラムのバグかもしれません。"
+
+        dataset = train_dataset_group.datasets[0]
+
+        dataset_dirs_info = {}
+        reg_dataset_dirs_info = {}
+        if use_dreambooth_method:
+            for subset in dataset.subsets:
+                info = reg_dataset_dirs_info if subset.is_reg else dataset_dirs_info
+                info[os.path.basename(subset.image_dir)] = {"n_repeats": subset.num_repeats, "img_count": subset.img_count}
+        else:
+            for subset in dataset.subsets:
+                dataset_dirs_info[os.path.basename(subset.metadata_file)] = {
+                    "n_repeats": subset.num_repeats,
+                    "img_count": subset.img_count,
+                }
+
+        metadata.update(
+            {
+                "ss_batch_size_per_device": args.train_batch_size,
+                "ss_total_batch_size": total_batch_size,
+                "ss_resolution": args.resolution,
+                "ss_color_aug": bool(args.color_aug),
+                "ss_flip_aug": bool(args.flip_aug),
+                "ss_random_crop": bool(args.random_crop),
+                "ss_shuffle_caption": bool(args.shuffle_caption),
+                "ss_enable_bucket": bool(dataset.enable_bucket),
+                "ss_bucket_no_upscale": bool(dataset.bucket_no_upscale),
+                "ss_min_bucket_reso": dataset.min_bucket_reso,
+                "ss_max_bucket_reso": dataset.max_bucket_reso,
+                "ss_keep_tokens": args.keep_tokens,
+                "ss_dataset_dirs": json.dumps(dataset_dirs_info),
+                "ss_reg_dataset_dirs": json.dumps(reg_dataset_dirs_info),
+                "ss_tag_frequency": json.dumps(dataset.tag_frequency),
+                "ss_bucket_info": json.dumps(dataset.bucket_info),
+            }
+        )
+
+    # add extra args
+    if args.network_args:
+        metadata["ss_network_args"] = json.dumps(net_kwargs)
+
+    # model name and hash
+    if args.pretrained_model_name_or_path is not None:
+        sd_model_name = args.pretrained_model_name_or_path
+        if os.path.exists(sd_model_name):
+            metadata["ss_sd_model_hash"] = train_util.model_hash(sd_model_name)
+            metadata["ss_new_sd_model_hash"] = train_util.calculate_sha256(sd_model_name)
+            sd_model_name = os.path.basename(sd_model_name)
+        metadata["ss_sd_model_name"] = sd_model_name
+
+    if args.vae is not None:
+        vae_name = args.vae
+        if os.path.exists(vae_name):
+            metadata["ss_vae_hash"] = train_util.model_hash(vae_name)
+            metadata["ss_new_vae_hash"] = train_util.calculate_sha256(vae_name)
+            vae_name = os.path.basename(vae_name)
+        metadata["ss_vae_name"] = vae_name
+
+    metadata = {k: str(v) for k, v in metadata.items()}
+
+    # make minimum metadata for filtering
+    minimum_keys = ["ss_network_module", "ss_network_dim", "ss_network_alpha", "ss_network_args"]
+    minimum_metadata = {}
+    for key in minimum_keys:
+        if key in metadata:
+            minimum_metadata[key] = metadata[key]
+
+    progress_bar = tqdm(range(args.max_train_steps), smoothing=0, disable=not accelerator.is_local_main_process, desc="steps")
+    global_step = 0
+
+    noise_scheduler = DDPMScheduler(
+        beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000, clip_sample=False
+    )
+    if accelerator.is_main_process:
+        accelerator.init_trackers("network_train")
+
+    loss_list = []
+    loss_total = 0.0
+    del train_dataset_group
+
+    # if hasattr(network, "on_step_start"):
+    #     on_step_start = network.on_step_start
+    # else:
+    #     on_step_start = lambda *args, **kwargs: None
+
+    for epoch in range(num_train_epochs):
+        if is_main_process:
+            print(f"epoch {epoch+1}/{num_train_epochs}")
+        current_epoch.value = epoch + 1
+
+        metadata["ss_epoch"] = str(epoch + 1)
+
+        network.on_epoch_start(text_encoder, unet)
+
+        for step, batch in enumerate(train_dataloader):
+            current_step.value = global_step
+            with accelerator.accumulate(network):
+                # on_step_start(text_encoder, unet)
+
+                with torch.no_grad():
+                    if "latents" in batch and batch["latents"] is not None:
+                        latents = batch["latents"].to(accelerator.device)
+                    else:
+                        # latentに変換
+                        latents = vae.encode(batch["images"].to(dtype=weight_dtype)).latent_dist.sample()
+                    latents = latents * 0.18215
+                b_size = latents.shape[0]
+
+                with torch.set_grad_enabled(train_text_encoder):
+                    # Get the text embedding for conditioning
+                    if args.weighted_captions:
+                        encoder_hidden_states = get_weighted_text_embeddings(
+                            tokenizer,
+                            text_encoder,
+                            batch["captions"],
+                            accelerator.device,
+                            args.max_token_length // 75 if args.max_token_length else 1,
+                            clip_skip=args.clip_skip,
+                        )
+                    else:
+                        input_ids = batch["input_ids"].to(accelerator.device)
+                        encoder_hidden_states = train_util.get_hidden_states(args, input_ids, tokenizer, text_encoder, weight_dtype)
+                # 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)
+                timesteps = timesteps.long()
+                # Add noise to the latents according to the noise magnitude at each timestep
+                # (this is the forward diffusion process)
+                noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)
+
+                # Predict the noise residual
+                with accelerator.autocast():
+                    noise_pred = unet(noisy_latents, timesteps, encoder_hidden_states).sample
+
+                if args.v_parameterization:
+                    # v-parameterization training
+                    target = noise_scheduler.get_velocity(latents, noise, timesteps)
+                else:
+                    target = noise
+
+                loss = torch.nn.functional.mse_loss(noise_pred.float(), target.float(), reduction="none")
+                loss = loss.mean([1, 2, 3])
+
+                loss_weights = batch["loss_weights"]  # 各sampleごとのweight
+                loss = loss * loss_weights
+
+                if args.min_snr_gamma:
+                    loss = apply_snr_weight(loss, timesteps, noise_scheduler, args.min_snr_gamma)
+
+                loss = loss.mean()  # 平均なのでbatch_sizeで割る必要なし
+
+                accelerator.backward(loss)
+                if accelerator.sync_gradients and args.max_grad_norm != 0.0:
+                    params_to_clip = network.get_trainable_params()
+                    accelerator.clip_grad_norm_(params_to_clip, args.max_grad_norm)
+
+                optimizer.step()
+                lr_scheduler.step()
+                optimizer.zero_grad(set_to_none=True)
+
+            # Checks if the accelerator has performed an optimization step behind the scenes
+            if accelerator.sync_gradients:
+                progress_bar.update(1)
+                global_step += 1
+
+                train_util.sample_images(
+                    accelerator, args, None, global_step, accelerator.device, vae, tokenizer, text_encoder, unet
+                )
+
+            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 / len(loss_list)
+            logs = {"loss": avr_loss}  # , "lr": lr_scheduler.get_last_lr()[0]}
+            progress_bar.set_postfix(**logs)
+
+            if args.logging_dir is not None:
+                logs = generate_step_logs(args, current_loss, avr_loss, lr_scheduler)
+                accelerator.log(logs, step=global_step)
+
+            if global_step >= args.max_train_steps:
+                break
+
+        if args.logging_dir is not None:
+            logs = {"loss/epoch": loss_total / len(loss_list)}
+            accelerator.log(logs, step=epoch + 1)
+
+        accelerator.wait_for_everyone()
+
+        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
+
+            def save_func():
+                ckpt_name = train_util.EPOCH_FILE_NAME.format(model_name, epoch + 1) + "." + args.save_model_as
+                ckpt_file = os.path.join(args.output_dir, ckpt_name)
+                metadata["ss_training_finished_at"] = str(time.time())
+                print(f"saving checkpoint: {ckpt_file}")
+                unwrap_model(network).save_weights(ckpt_file, save_dtype, minimum_metadata if args.no_metadata else metadata)
+                if args.huggingface_repo_id is not None:
+                    huggingface_util.upload(args, ckpt_file, "/" + ckpt_name)
+
+            def remove_old_func(old_epoch_no):
+                old_ckpt_name = train_util.EPOCH_FILE_NAME.format(model_name, old_epoch_no) + "." + args.save_model_as
+                old_ckpt_file = os.path.join(args.output_dir, old_ckpt_name)
+                if os.path.exists(old_ckpt_file):
+                    print(f"removing old checkpoint: {old_ckpt_file}")
+                    os.remove(old_ckpt_file)
+
+            if is_main_process:
+                saving = train_util.save_on_epoch_end(args, save_func, remove_old_func, epoch + 1, num_train_epochs)
+                if saving and args.save_state:
+                    train_util.save_state_on_epoch_end(args, accelerator, model_name, epoch + 1)
+
+        train_util.sample_images(accelerator, args, epoch + 1, global_step, accelerator.device, vae, tokenizer, text_encoder, unet)
+
+        # end of epoch
+
+    metadata["ss_epoch"] = str(num_train_epochs)
+    metadata["ss_training_finished_at"] = str(time.time())
+
+    if is_main_process:
+        network = unwrap_model(network)
+
+    accelerator.end_training()
+
+    if args.save_state:
+        train_util.save_state_on_train_end(args, accelerator)
+
+    del accelerator  # この後メモリを使うのでこれは消す
+
+    if is_main_process:
+        os.makedirs(args.output_dir, exist_ok=True)
+
+        model_name = train_util.DEFAULT_LAST_OUTPUT_NAME if args.output_name is None else args.output_name
+        ckpt_name = model_name + "." + args.save_model_as
+        ckpt_file = os.path.join(args.output_dir, ckpt_name)
+
+        print(f"save trained model to {ckpt_file}")
+        network.save_weights(ckpt_file, save_dtype, minimum_metadata if args.no_metadata else metadata)
+        if args.huggingface_repo_id is not None:
+            huggingface_util.upload(args, ckpt_file, "/" + ckpt_name, force_sync_upload=True)
+        print("model saved.")
+
+
+def setup_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser()
+
+    train_util.add_sd_models_arguments(parser)
+    train_util.add_dataset_arguments(parser, True, True, True)
+    train_util.add_training_arguments(parser, True)
+    train_util.add_optimizer_arguments(parser)
+    config_util.add_config_arguments(parser)
+    custom_train_functions.add_custom_train_arguments(parser)
+
+    parser.add_argument("--no_metadata", action="store_true", help="do not save metadata in output model / メタデータを出力先モデルに保存しない")
+    parser.add_argument(
+        "--save_model_as",
+        type=str,
+        default="safetensors",
+        choices=[None, "ckpt", "pt", "safetensors"],
+        help="format to save the model (default is .safetensors) / モデル保存時の形式（デフォルトはsafetensors）",
+    )
+
+    parser.add_argument("--unet_lr", type=float, default=None, help="learning rate for U-Net / U-Netの学習率")
+    parser.add_argument("--text_encoder_lr", type=float, default=None, help="learning rate for Text Encoder / Text Encoderの学習率")
+
+    parser.add_argument("--network_weights", type=str, default=None, help="pretrained weights for network / 学習するネットワークの初期重み")
+    parser.add_argument("--network_module", type=str, default=None, help="network module to train / 学習対象のネットワークのモジュール")
+    parser.add_argument(
+        "--network_dim", type=int, default=None, help="network dimensions (depends on each network) / モジュールの次元数（ネットワークにより定義は異なります）"
+    )
+    parser.add_argument(
+        "--network_alpha",
+        type=float,
+        default=1,
+        help="alpha for LoRA weight scaling, default 1 (same as network_dim for same behavior as old version) / LoRaの重み調整のalpha値、デフォルト1（旧バージョンと同じ動作をするにはnetwork_dimと同じ値を指定）",
+    )
+    parser.add_argument(
+        "--network_args", type=str, default=None, nargs="*", help="additional argmuments for network (key=value) / ネットワークへの追加の引数"
+    )
+    parser.add_argument("--network_train_unet_only", action="store_true", help="only training U-Net part / U-Net関連部分のみ学習する")
+    parser.add_argument(
+        "--network_train_text_encoder_only", action="store_true", help="only training Text Encoder part / Text Encoder関連部分のみ学習する"
+    )
+    parser.add_argument(
+        "--training_comment", type=str, default=None, help="arbitrary comment string stored in metadata / メタデータに記録する任意のコメント文字列"
+    )
+
+    return parser
+
+
+if __name__ == "__main__":
+    parser = setup_parser()
+
+    args = parser.parse_args()
+    args = train_util.read_config_from_file(args, parser)
+
+    train(args)

train_network_README-ja.md

@@ -0,0 +1,479 @@
+# LoRAの学習について
+
+[LoRA: Low-Rank Adaptation of Large Language Models](https://arxiv.org/abs/2106.09685)（arxiv）、[LoRA](https://github.com/microsoft/LoRA)（github）をStable Diffusionに適用したものです。
+
+[cloneofsimo氏のリポジトリ](https://github.com/cloneofsimo/lora)を大いに参考にさせていただきました。ありがとうございます。
+
+通常のLoRAは Linear およぴカーネルサイズ 1x1 の Conv2d にのみ適用されますが、カーネルサイズ 3x3 のConv2dに適用を拡大することもできます。
+
+Conv2d 3x3への拡大は [cloneofsimo氏](https://github.com/cloneofsimo/lora) が最初にリリースし、KohakuBlueleaf氏が [LoCon](https://github.com/KohakuBlueleaf/LoCon) でその有効性を明らかにしたものです。KohakuBlueleaf氏に深く感謝します。
+
+8GB VRAMでもぎりぎり動作するようです。
+
+[学習についての共通ドキュメント](./train_README-ja.md) もあわせてご覧ください。
+
+# 学習できるLoRAの種類
+
+以下の二種類をサポートします。以下は当リポジトリ内の独自の名称です。
+
+1. __LoRA-LierLa__ : (LoRA for __Li__ n __e__ a __r__  __La__ yers、リエラと読みます)
+
+    Linear およびカーネルサイズ 1x1 の Conv2d に適用されるLoRA
+
+2. __LoRA-C3Lier__ : (LoRA for __C__ olutional layers with __3__ x3 Kernel and  __Li__ n __e__ a __r__ layers、セリアと読みます)
+
+    1.に加え、カーネルサイズ 3x3 の Conv2d に適用されるLoRA
+
+LoRA-LierLaに比べ、LoRA-C3Liarは適用される層が増える分、高い精度が期待できるかもしれません。
+
+また学習時は __DyLoRA__ を使用することもできます（後述します）。
+
+## 学習したモデルに関する注意
+
+LoRA-LierLa は、AUTOMATIC1111氏のWeb UIのLoRA機能で使用することができます。
+
+LoRA-C3Liarを使いWeb UIで生成するには、こちらの[WebUI用extension](https://github.com/kohya-ss/sd-webui-additional-networks)を使ってください。
+
+いずれも学習したLoRAのモデルを、Stable Diffusionのモデルにこのリポジトリ内のスクリプトであらかじめマージすることもできます。
+
+cloneofsimo氏のリポジトリ、およびd8ahazard氏の[Dreambooth Extension for Stable-Diffusion-WebUI](https://github.com/d8ahazard/sd_dreambooth_extension)とは、現時点では互換性がありません。いくつかの機能拡張を行っているためです（後述）。
+
+# 学習の手順
+
+あらかじめこのリポジトリのREADMEを参照し、環境整備を行ってください。
+
+## データの準備
+
+[学習データの準備について](./train_README-ja.md) を参照してください。
+
+
+## 学習の実行
+
+`train_network.py`を用います。
+
+`train_network.py`では `--network_module` オプションに、学習対象のモジュール名を指定します。LoRAに対応するのは`network.lora`となりますので、それを指定してください。
+
+なお学習率は通常のDreamBoothやfine tuningよりも高めの、`1e-4`～`1e-3`程度を指定するとよいようです。
+
+以下はコマンドラインの例です。
+
+```
+accelerate launch --num_cpu_threads_per_process 1 train_network.py 
+    --pretrained_model_name_or_path=<.ckptまたは.safetensordまたはDiffusers版モデルのディレクトリ> 
+    --dataset_config=<データ準備で作成した.tomlファイル> 
+    --output_dir=<学習したモデルの出力先フォルダ>  
+    --output_name=<学習したモデル出力時のファイル名> 
+    --save_model_as=safetensors 
+    --prior_loss_weight=1.0 
+    --max_train_steps=400 
+    --learning_rate=1e-4 
+    --optimizer_type="AdamW8bit" 
+    --xformers 
+    --mixed_precision="fp16" 
+    --cache_latents 
+    --gradient_checkpointing
+    --save_every_n_epochs=1 
+    --network_module=networks.lora
+```
+
+このコマンドラインでは LoRA-LierLa が学習されます。
+
+`--output_dir` オプションで指定したフォルダに、LoRAのモデルが保存されます。他のオプション、オプティマイザ等については [学習の共通ドキュメント](./train_README-ja.md) の「よく使われるオプション」も参照してください。
+
+その他、以下のオプションが指定できます。
+
+* `--network_dim`
+  * LoRAのRANKを指定します（``--networkdim=4``など）。省略時は4になります。数が多いほど表現力は増しますが、学習に必要なメモリ、時間は増えます。また闇雲に増やしても良くないようです。
+* `--network_alpha`
+  *  アンダーフローを防ぎ安定して学習するための ``alpha`` 値を指定します。デフォルトは1です。``network_dim``と同じ値を指定すると以前のバージョンと同じ動作になります。
+* `--persistent_data_loader_workers`
+  * Windows環境で指定するとエポック間��待ち時間が大幅に短縮されます。
+* `--max_data_loader_n_workers`
+  * データ読み込みのプロセス数を指定します。プロセス数が多いとデータ読み込みが速くなりGPUを効率的に利用できますが、メインメモリを消費します。デフォルトは「`8` または `CPU同時実行スレッド数-1` の小さいほう」なので、メインメモリに余裕がない場合や、GPU使用率が90%程度以上なら、それらの数値を見ながら `2` または `1` 程度まで下げてください。
+* `--network_weights`
+  * 学習前に学習済みのLoRAの重みを読み込み、そこから追加で学習します。
+* `--network_train_unet_only`
+  * U-Netに関連するLoRAモジュールのみ有効とします。fine tuning的な学習で指定するとよいかもしれません。
+* `--network_train_text_encoder_only`
+  * Text Encoderに関連するLoRAモジュールのみ有効とします。Textual Inversion的な効果が期待できるかもしれません。
+* `--unet_lr`
+  * U-Netに関連するLoRAモジュールに、通常の学習率（--learning_rateオプションで指定）とは異なる学習率を使う時に指定します。
+* `--text_encoder_lr`
+  * Text Encoderに関連するLoRAモジュールに、通常の学習率（--learning_rateオプションで指定）とは異なる学習率を使う時に指定します。Text Encoderのほうを若干低めの学習率（5e-5など）にしたほうが良い、という話もあるようです。
+* `--network_args`
+  * 複数の引数を指定できます。後述します。
+
+`--network_train_unet_only` と `--network_train_text_encoder_only` の両方とも未指定時（デフォルト）はText EncoderとU-Netの両方のLoRAモジュールを有効にします。
+
+# その他の学習方法
+
+## LoRA-C3Lier を学習する
+
+`--network_args` に以下のように指定してください。`conv_dim` で Conv2d (3x3) の rank を、`conv_alpha` で alpha を指定してください。
+
+```
+--network_args "conv_dim=4" "conv_alpha=1"
+```
+
+以下のように alpha 省略時は1になります。
+
+```
+--network_args "conv_dim=4"
+```
+
+## DyLoRA
+
+DyLoRAはこちらの論文で提案されたものです。[DyLoRA: Parameter Efficient Tuning of Pre-trained Models using Dynamic Search-Free Low-Rank Adaptation](https://arxiv.org/abs/2210.07558)　公式実装は[こちら](https://github.com/huawei-noah/KD-NLP/tree/main/DyLoRA)です。
+
+論文によると、LoRAのrankは必ずしも高いほうが良いわけではなく、対象のモデル、データセット、タスクなどにより適切なrankを探す必要があるようです。DyLoRAを使うと、指定したdim(rank)以下のさまざまなrankで同時にLoRAを学習します。これにより最適なrankをそれぞれ学習して探す手間を省くことができます。
+
+当リポジトリの実装は公式実装をベースに独自の拡張を加えています（そのため不具合などあるかもしれません）。
+
+### 当リポジトリのDyLoRAの特徴
+
+学習後のDyLoRAのモデルファイルはLoRAと互換性があります。また、モデルファイルから指定したdim(rank)以下の複数のdimのLoRAを抽出できます。
+
+DyLoRA-LierLa、DyLoRA-C3Lierのどちらも学習できます。
+
+### DyLoRAで学習する
+
+`--network_module=networks.dylora` のように、DyLoRAに対応する`network.dylora`を指定してください。
+
+また `--network_args` に、たとえば`--network_args "unit=4"`のように`unit`を指定します。`unit`はrankを分割する単位です。たとえば`--network_dim=16 --network_args "unit=4"` のように指定します。`unit`は`network_dim`を割り切れる値（`network_dim`は`unit`の倍数）としてください。
+
+`unit`を指定しない場合は、`unit=1`として扱われます。
+
+記述例は以下です。
+
+```
+--network_module=networks.dylora --network_dim=16 --network_args "unit=4"
+
+--network_module=networks.dylora --network_dim=32 --network_alpha=16 --network_args "unit=4"
+```
+
+DyLoRA-C3Lierの場合は、`--network_args` に`"conv_dim=4"`のように`conv_dim`を指定します。通常のLoRAと異なり、`conv_dim`は`network_dim`と同じ値である必要があります。記述例は以下です。
+
+```
+--network_module=networks.dylora --network_dim=16 --network_args "conv_dim=16" "unit=4"
+
+--network_module=networks.dylora --network_dim=32 --network_alpha=16 --network_args "conv_dim=32" "conv_alpha=16" "unit=8"
+```
+
+たとえばdim=16、unit=4（後述）で学習すると、4、8、12、16の4つのrankのLoRAを学習、抽出できます。抽出した各モデルで画像を生成し、比較することで、最適なrankのLoRAを選択できます。
+
+その他のオプションは通常のLoRAと同じです。
+
+※ `unit`は当リポジトリの独自拡張で、DyLoRAでは同dim(rank)の通常LoRAに比べると学習時間が長くな��ことが予想されるため、分割単位を大きくしたものです。
+
+### DyLoRAのモデルからLoRAモデルを抽出する
+
+`networks`フォルダ内の `extract_lora_from_dylora.py`を使用します。指定した`unit`単位で、DyLoRAのモデルからLoRAのモデルを抽出します。
+
+コマンドラインはたとえば以下のようになります。
+
+```powershell
+python networks\extract_lora_from_dylora.py --model "foldername/dylora-model.safetensors" --save_to "foldername/dylora-model-split.safetensors" --unit 4
+```
+
+`--model` にはDyLoRAのモデルファイルを指定します。`--save_to` には抽出したモデルを保存するファイル名を指定します（rankの数値がファイル名に付加されます）。`--unit` にはDyLoRAの学習時の`unit`を指定します。
+
+## 階層別学習率
+
+詳細は[PR #355](https://github.com/kohya-ss/sd-scripts/pull/355) をご覧ください。
+
+フルモデルの25個のブロックの重みを指定できます。最初のブロックに該当するLoRAは存在しませんが、階層別LoRA適用等との互換性のために25個としています。またconv2d3x3に拡張しない場合も一部のブロックにはLoRAが存在しませんが、記述を統一するため常に25個の値を指定してください。
+
+`--network_args` で以下の引数を指定してください。
+
+- `down_lr_weight` : U-Netのdown blocksの学習率の重みを指定します。以下が指定可能です。
+  - ブロックごとの重み : `"down_lr_weight=0,0,0,0,0,0,1,1,1,1,1,1"` のように12個の数値を指定します。
+  - プリセットからの指定 : `"down_lr_weight=sine"` のように指定します（サインカーブで重みを指定します）。sine, cosine, linear, reverse_linear, zeros が指定可能です。また `"down_lr_weight=cosine+.25"` のように `+数値` を追加すると、指定した数値を加算します（0.25~1.25になります）。
+- `mid_lr_weight` : U-Netのmid blockの学習率の重みを指定します。`"down_lr_weight=0.5"` のように数値を一つだけ指定します。
+- `up_lr_weight` : U-Netのup blocksの学習率の重みを指定します。down_lr_weightと同様です。
+- 指定を省略した部分は1.0として扱われます。また重みを0にするとそのブロックのLoRAモジュールは作成されません。
+- `block_lr_zero_threshold` : 重みがこの値以下の場合、LoRAモジュールを作成しません。デフォルトは0です。
+
+### 階層別学習率コマンドライン指定例:
+
+```powershell
+--network_args "down_lr_weight=0.5,0.5,0.5,0.5,1.0,1.0,1.0,1.0,1.5,1.5,1.5,1.5" "mid_lr_weight=2.0" "up_lr_weight=1.5,1.5,1.5,1.5,1.0,1.0,1.0,1.0,0.5,0.5,0.5,0.5"
+
+--network_args "block_lr_zero_threshold=0.1" "down_lr_weight=sine+.5" "mid_lr_weight=1.5" "up_lr_weight=cosine+.5"
+```
+
+###  階層別学習率tomlファイル指定例:
+
+```toml
+network_args = [ "down_lr_weight=0.5,0.5,0.5,0.5,1.0,1.0,1.0,1.0,1.5,1.5,1.5,1.5", "mid_lr_weight=2.0", "up_lr_weight=1.5,1.5,1.5,1.5,1.0,1.0,1.0,1.0,0.5,0.5,0.5,0.5",]
+
+network_args = [ "block_lr_zero_threshold=0.1", "down_lr_weight=sine+.5", "mid_lr_weight=1.5", "up_lr_weight=cosine+.5", ]
+```
+
+## 階層別dim (rank)
+
+フルモデルの25個のブロックのdim (rank)を指定できます。階層別学習率と同様に一部のブロックにはLoRAが存在しない場合がありますが、常に25個の値を指定してください。
+
+`--network_args` で以下の引数を指定してください。
+
+- `block_dims` : 各ブロックのdim (rank)を指定します。`"block_dims=2,2,2,2,4,4,4,4,6,6,6,6,8,6,6,6,6,4,4,4,4,2,2,2,2"` のように25個の数値を指定します。
+- `block_alphas` : 各ブロックのalphaを指定します。block_dimsと同様に25個の数値を指定します。省略時はnetwork_alphaの値が使用されます。
+- `conv_block_dims` : LoRAをConv2d 3x3に拡張し、各ブロックのdim (rank)を指定します。
+- `conv_block_alphas` : LoRAをConv2d 3x3に拡張したときの各ブロックのalphaを指定します。省略時はconv_alphaの値が使用されます。
+
+###  階層別dim (rank)コマンドライン指定例:
+
+```powershell
+--network_args "block_dims=2,4,4,4,8,8,8,8,12,12,12,12,16,12,12,12,12,8,8,8,8,4,4,4,2"
+
+--network_args "block_dims=2,4,4,4,8,8,8,8,12,12,12,12,16,12,12,12,12,8,8,8,8,4,4,4,2" "conv_block_dims=2,2,2,2,4,4,4,4,6,6,6,6,8,6,6,6,6,4,4,4,4,2,2,2,2"
+
+--network_args "block_dims=2,4,4,4,8,8,8,8,12,12,12,12,16,12,12,12,12,8,8,8,8,4,4,4,2" "block_alphas=2,2,2,2,4,4,4,4,6,6,6,6,8,6,6,6,6,4,4,4,4,2,2,2,2"
+```
+
+###  階層別dim (rank)tomlファイル指定例:
+
+```toml
+network_args = [ "block_dims=2,4,4,4,8,8,8,8,12,12,12,12,16,12,12,12,12,8,8,8,8,4,4,4,2",]
+  
+network_args = [ "block_dims=2,4,4,4,8,8,8,8,12,12,12,12,16,12,12,12,12,8,8,8,8,4,4,4,2", "block_alphas=2,2,2,2,4,4,4,4,6,6,6,6,8,6,6,6,6,4,4,4,4,2,2,2,2",]
+```
+
+# その他のスクリプト
+
+マージ等LoRAに関連するスクリプト群です。
+
+## マージスクリプトについて
+
+merge_lora.pyでStable DiffusionのモデルにLoRAの学習結果をマージしたり、複数のLoRAモデルをマージしたりできます。
+
+### Stable DiffusionのモデルにLoRAのモデルをマージする
+
+マージ後のモデルは通常のStable Diffusionのckptと同様に扱えます。たとえば以下のようなコマンドラインになります。
+
+```
+python networks\merge_lora.py --sd_model ..\model\model.ckpt 
+    --save_to ..\lora_train1\model-char1-merged.safetensors 
+    --models ..\lora_train1\last.safetensors --ratios 0.8
+```
+
+Stable Diffusion v2.xのモデルで学習し、それにマージする場合は、--v2オプションを指定してください。
+
+--sd_modelオプションにマージの元となるStable Diffusionのモデルファイルを指定します（.ckptまたは.safetensorsのみ対応で、Diffusersは今のところ対応していません）。
+
+--save_toオプションにマージ後のモデルの保存先を指定します（.ckptまたは.safetensors、拡張子で自動判定）。
+
+--modelsに学習したLoRAのモデルファイルを指定します。複数指定も可能で、その時は順にマージします。
+
+--ratiosにそれぞれのモデルの適用率（どのくらい重みを元モデルに反映するか）を0~1.0の数値で指定します。例えば過学習に近いような場合は、適用率を下げるとマシになるかもしれません。モデルの数と同じだけ指定してください。
+
+複数指定時は以下のようになります。
+
+```
+python networks\merge_lora.py --sd_model ..\model\model.ckpt 
+    --save_to ..\lora_train1\model-char1-merged.safetensors 
+    --models ..\lora_train1\last.safetensors ..\lora_train2\last.safetensors --ratios 0.8 0.5
+```
+
+### 複数のLoRAのモデルをマージする
+
+複数のLoRAモデルをひとつずつSDモデルに適用する場合と、複数のLoRAモデルをマージしてからSDモデルにマージする場合とは、計算順序の関連で微妙に異なる結果になります。
+
+たとえば以下のようなコマンドラインになります。
+
+```
+python networks\merge_lora.py 
+    --save_to ..\lora_train1\model-char1-style1-merged.safetensors 
+    --models ..\lora_train1\last.safetensors ..\lora_train2\last.safetensors --ratios 0.6 0.4
+```
+
+--sd_modelオプションは指定不要です。
+
+--save_toオプションにマージ後のLoRAモデルの保存先を指定します（.ckptまたは.safetensors、拡張子で自動判定）。
+
+--modelsに学習したLoRAのモデルファイルを指定します。三つ以上も指定可能です。
+
+--ratiosにそれぞれのモデルの比率（どのくらい重みを元モデルに反映するか）を0~1.0の数値で指定します。二つのモデルを一対一でマージす場合は、「0.5 0.5」になります。「1.0 1.0」では合計の重みが大きくなりすぎて、恐らく結果はあまり望ましくないものになると思われます。
+
+v1で学習したLoRAとv2で学習したLoRA、rank（次元数）や``alpha``の異なるLoRAはマージできません。U-NetだけのLoRAとU-Net+Text EncoderのLoRAはマージできるはずですが、結果は未知数です。
+
+
+### その他のオプション
+
+* precision
+  * マージ計算時の精度をfloat、fp16、bf16から指定できます。省略時は精度を確保するためfloatになります。メモリ使用量を減らしたい場合はfp16/bf16を指定してください。
+* save_precision
+  * モデル保存時の精度をfloat、fp16、bf16から指定できます。省略時はprecisionと同じ精度になります。
+
+
+## 複数のrankが異なるLoRAのモデルをマージする
+
+複数のLoRAをひとつのLoRAで近似します（完全な再現はできません）。`svd_merge_lora.py`を用います。たとえば以下のようなコマンドラインになります。
+
+```
+python networks\svd_merge_lora.py 
+    --save_to ..\lora_train1\model-char1-style1-merged.safetensors 
+    --models ..\lora_train1\last.safetensors ..\lora_train2\last.safetensors 
+    --ratios 0.6 0.4 --new_rank 32 --device cuda
+```
+
+`merge_lora.py` と主なオプションは同一です。以下のオプションが追加されています。
+
+- `--new_rank`
+  - 作成するLoRAのrankを指定します。
+- `--new_conv_rank`
+  - 作成する Conv2d 3x3 LoRA の rank を指定します。省略時は `new_rank` と同じになります。
+- `--device`
+  - `--device cuda`としてcudaを指定すると計算をGPU上で行います。処理が速くなります。
+
+## 当リポジトリ内の画像生成スクリプトで生成する
+
+gen_img_diffusers.pyに、--network_module、--network_weightsの各オプションを追加してください。意味は学習時と同様です。
+
+--network_mulオプションで0~1.0の数値を指定す���と、LoRAの適用率を変えられます。
+
+## Diffusersのpipelineで生成する
+
+以下の例を参考にしてください。必要なファイルはnetworks/lora.pyのみです。Diffusersのバージョンは0.10.2以外では動作しない可能性があります。
+
+```python
+import torch
+from diffusers import StableDiffusionPipeline
+from networks.lora import LoRAModule, create_network_from_weights
+from safetensors.torch import load_file
+
+# if the ckpt is CompVis based, convert it to Diffusers beforehand with tools/convert_diffusers20_original_sd.py. See --help for more details.
+
+model_id_or_dir = r"model_id_on_hugging_face_or_dir"
+device = "cuda"
+
+# create pipe
+print(f"creating pipe from {model_id_or_dir}...")
+pipe = StableDiffusionPipeline.from_pretrained(model_id_or_dir, revision="fp16", torch_dtype=torch.float16)
+pipe = pipe.to(device)
+vae = pipe.vae
+text_encoder = pipe.text_encoder
+unet = pipe.unet
+
+# load lora networks
+print(f"loading lora networks...")
+
+lora_path1 = r"lora1.safetensors"
+sd = load_file(lora_path1)   # If the file is .ckpt, use torch.load instead.
+network1, sd = create_network_from_weights(0.5, None, vae, text_encoder,unet, sd)
+network1.apply_to(text_encoder, unet)
+network1.load_state_dict(sd)
+network1.to(device, dtype=torch.float16)
+
+# # You can merge weights instead of apply_to+load_state_dict. network.set_multiplier does not work
+# network.merge_to(text_encoder, unet, sd)
+
+lora_path2 = r"lora2.safetensors"
+sd = load_file(lora_path2) 
+network2, sd = create_network_from_weights(0.7, None, vae, text_encoder,unet, sd)
+network2.apply_to(text_encoder, unet)
+network2.load_state_dict(sd)
+network2.to(device, dtype=torch.float16)
+
+lora_path3 = r"lora3.safetensors"
+sd = load_file(lora_path3)
+network3, sd = create_network_from_weights(0.5, None, vae, text_encoder,unet, sd)
+network3.apply_to(text_encoder, unet)
+network3.load_state_dict(sd)
+network3.to(device, dtype=torch.float16)
+
+# prompts
+prompt = "masterpiece, best quality, 1girl, in white shirt, looking at viewer"
+negative_prompt = "bad quality, worst quality, bad anatomy, bad hands"
+
+# exec pipe
+print("generating image...")
+with torch.autocast("cuda"):
+    image = pipe(prompt, guidance_scale=7.5, negative_prompt=negative_prompt).images[0]
+
+# if not merged, you can use set_multiplier
+# network1.set_multiplier(0.8)
+# and generate image again...
+
+# save image
+image.save(r"by_diffusers..png")
+```
+
+## 二つのモデルの差分からLoRAモデルを作成する
+
+[こちらのディスカッション](https://github.com/cloneofsimo/lora/discussions/56)を参考に実装したものです。数式はそのまま使わせていただきました（よく理解していませんが近似には特異値分解を用いるようです）。
+
+二つのモデル（たとえばfine tuningの元モデルとfine tuning後のモデル）の差分を、LoRAで近似します。
+
+### スクリプトの実行方法
+
+以下のように指定してください。
+```
+python networks\extract_lora_from_models.py --model_org base-model.ckpt
+    --model_tuned fine-tuned-model.ckpt 
+    --save_to lora-weights.safetensors --dim 4
+```
+
+--model_orgオプションに元のStable Diffusionモデルを指定します。作成したLoRAモデルを適用する場合は、このモデルを指定して適用することになります。.ckptまたは.safetensorsが指定できます。
+
+--model_tunedオプションに差分を抽出する対象のStable Diffusionモデルを指定します。たとえばfine tuningやDreamBooth後のモデルを指定します。.ckptまたは.safetensorsが指定できます。
+
+--save_toにLoRAモデルの保存先を指定します。--dimにLoRAの次元数を指定します。
+
+生成されたLoRAモデルは、学習したLoRAモデルと同様に使用できます。
+
+Text Encoderが二つのモデルで同じ場合にはLoRAはU-NetのみのLoRAとなります。
+
+### その他のオプション
+
+- `--v2`
+  - v2.xのStable Diffusionモデルを使う場合に指定してください。
+- `--device`
+  - ``--device cuda``としてcudaを指定すると計算をGPU上で行います。処理が速くなります（CPUでもそこまで遅くないため、せいぜい倍～数倍程度のようです）。
+- `--save_precision`
+  - LoRAの保存形式を"float", "fp16", "bf16"から指定します。省略時はfloatになります。
+- `--conv_dim`
+  - 指定するとLoRAの適用範囲を Conv2d 3x3 へ拡大します。Conv2d 3x3 の rank を指定します。
+
+## 画像リサイズスクリプト
+
+（のちほどドキュメントを整理しますがとりあえずここに説明を書いておきます。）
+
+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氏のリポジトリとの違い
+
+2022/12/25時点では、当リポジトリはLoRAの適用個所をText EncoderのMLP、U-NetのFFN、Transformerのin/out projectionに拡大し、表現力が増しています。ただその代わりメモリ使用量は増え、8GBぎりぎりになりました。
+
+またモジュール入れ替え機構は全く異なります。
+
+## 将来拡張について
+
+LoRAだけでなく他の拡張にも対応可能ですので、それらも追加予定です。

train_network_README.md

@@ -0,0 +1,189 @@
+## About learning LoRA
+
+[LoRA: Low-Rank Adaptation of Large Language Models](https://arxiv.org/abs/2106.09685) (arxiv), [LoRA](https://github.com/microsoft/LoRA) (github) to Stable Applied to Diffusion.
+
+[cloneofsimo's repository](https://github.com/cloneofsimo/lora) was a great reference. Thank you very much.
+
+8GB VRAM seems to work just fine.
+
+## A Note about Trained Models
+
+Cloneofsimo's repository and d8ahazard's [Drebooth Extension for Stable-Diffusion-WebUI](https://github.com/d8ahazard/sd_drebooth_extension) are currently incompatible. Because we are doing some enhancements (see below).
+
+When generating images with WebUI, etc., merge the learned LoRA model with the learning source Stable Diffusion model in advance with the script in this repository, or click here [Extention for WebUI] (https://github .com/kohya-ss/sd-webui-additional-networks).
+
+## Learning method
+
+Use train_network.py.
+
+You can learn both the DreamBooth method (using identifiers (sks, etc.) and classes, optionally regularized images) and the fine tuning method using captions.
+
+Both methods can be learned in much the same way as existing scripts. We will discuss the differences later.
+
+### Using the DreamBooth Method
+
+Please refer to [DreamBooth guide](./train_db_README-en.md) and prepare the data.
+
+Specify train_network.py instead of train_db.py when training.
+
+Almost all options are available (except Stable Diffusion model save related), but stop_text_encoder_training is not supported.
+
+### When to use captions
+
+Please refer to [fine-tuning guide](./fine_tune_README_en.md) and perform each step.
+
+Specify train_network.py instead of fine_tune.py when training. Almost all options (except for model saving) can be used as is.
+
+In addition, it will work even if you do not perform "Pre-obtain latents". Since the latent is acquired from the VAE when learning (or caching), the learning speed will be slower, but color_aug can be used instead.
+
+### Options for Learning LoRA
+
+In train_network.py, specify the name of the module to be trained in the --network_module option. LoRA is compatible with network.lora, so please specify it.
+
+The learning rate should be set to about 1e-4, which is higher than normal DreamBooth and fine tuning.
+
+Below is an example command line (DreamBooth technique).
+
+```
+accelerate launch --num_cpu_threads_per_process 12 train_network.py
+     --pretrained_model_name_or_path=..\models\model.ckpt
+     --train_data_dir=..\data\db\char1 --output_dir=..\lora_train1
+     --reg_data_dir=..\data\db\reg1 --prior_loss_weight=1.0
+     --resolution=448,640 --train_batch_size=1 --learning_rate=1e-4
+     --max_train_steps=400 --use_8bit_adam --xformers --mixed_precision=fp16
+     --save_every_n_epochs=1 --save_model_as=safetensors --clip_skip=2 --seed=42 --color_aug
+     --network_module=networks.lora
+```
+
+The LoRA model will be saved in the directory specified by the --output_dir option.
+
+In addition, the following options can be specified.
+
+* --network_dim
+   * Specify the number of dimensions of LoRA (such as ``--networkdim=4``). Default is 4. The greater the number, the greater the expressive power, but the memory and time required for learning also increase. In addition, it seems that it is not good to increase it blindly.
+* --network_weights
+   * Load pretrained LoRA weights before training and additionally learn from them.
+* --network_train_unet_only
+   * Valid only for LoRA modules related to U-Net. It may be better to specify it in fine-tuning study.
+* --network_train_text_encoder_only
+   * Only LoRA modules related to Text Encoder are enabled. You may be able to expect a textual inversion effect.
+* --unet_lr
+   * Specify when using a learning rate different from the normal learning rate (specified with the --learning_rate option) for the LoRA module related to U-Net.
+* --text_encoder_lr
+   * Specify when using a learning rate different from the normal learning rate (specified with the --learning_rate option) for the LoRA module associated with the Text Encoder. Some people say that it is better to set the Text Encoder to a slightly lower learning rate (such as 5e-5).
+
+When neither --network_train_unet_only nor --network_train_text_encoder_only is specified (default), both Text Encoder and U-Net LoRA modules are enabled.
+
+## About the merge script
+
+merge_lora.py allows you to merge LoRA training results into a Stable Diffusion model, or merge multiple LoRA models.
+
+### Merge LoRA model into Stable Diffusion model
+
+The model after merging can be handled in the same way as normal Stable Diffusion ckpt. For example, a command line like:
+
+```
+python networks\merge_lora.py --sd_model ..\model\model.ckpt
+     --save_to ..\lora_train1\model-char1-merged.safetensors
+     --models ..\lora_train1\last.safetensors --ratios 0.8
+```
+
+Specify the --v2 option if you want to train with a Stable Diffusion v2.x model and merge with it.
+
+Specify the Stable Diffusion model file to be merged in the --sd_model option (only .ckpt or .safetensors are supported, Diffusers is not currently supported).
+
+Specify the save destination of the model after merging in the --save_to option (.ckpt or .safetensors, automatically determined by extension).
+
+Specify the LoRA model file learned in --models. It is possible to specify more than one, in which case they will be merged in order.
+
+For --ratios, specify the application rate of each model (how much weight is reflected in the original model) with a numerical value from 0 to 1.0. For example, if it is close to overfitting, it may be better if the application rate is lowered. Specify as many as the number of models.
+
+When specifying multiple, it will be as follows.
+
+```
+python networks\merge_lora.py --sd_model ..\model\model.ckpt
+     --save_to ..\lora_train1\model-char1-merged.safetensors
+     --models ..\lora_train1\last.safetensors ..\lora_train2\last.safetensors --ratios 0.8 0.5
+```
+
+### Merge multiple LoRA models
+
+Applying multiple LoRA models one by one to the SD model and merging multiple LoRA models and then merging them into the SD model yield slightly different results in relation to the calculation order.
+
+For example, a command line like:
+
+```
+python networks\merge_lora.py
+     --save_to ..\lora_train1\model-char1-style1-merged.safetensors
+     --models ..\lora_train1\last.safetensors ..\lora_train2\last.safetensors --ratios 0.6 0.4
+```
+
+The --sd_model option does not need to be specified.
+
+Specify the save destination of the merged LoRA model in the --save_to option (.ckpt or .safetensors, automatically determined by extension).
+
+Specify the LoRA model file learned in --models. Three or more can be specified.
+
+For --ratios, specify the ratio of each model (how much weight is reflected in the original model) with a numerical value from 0 to 1.0. If you merge two models one-to-one, it will be "0.5 0.5". "1.0 1.0" would give too much weight to the sum, and the result would probably be less desirable.
+
+LoRA trained with v1 and LoRA trained with v2, and LoRA with different number of dimensions cannot be merged. U-Net only LoRA and U-Net+Text Encoder LoRA should be able to merge, but the result is unknown.
+
+
+### Other Options
+
+* precision
+   * The precision for merge calculation can be specified from float, fp16, and bf16. If omitted, it will be float to ensure accuracy. Specify fp16/bf16 if you want to reduce memory usage.
+* save_precision
+   * You can specify the precision when saving the model from float, fp16, bf16. If omitted, the precision is the same as precision.
+
+## Generate with the image generation script in this repository
+
+Add options --network_module, --network_weights, --network_dim (optional) to gen_img_diffusers.py. The meaning is the same as when learning.
+
+You can change the LoRA application rate by specifying a value between 0 and 1.0 with the --network_mul option.
+
+## Create a LoRA model from the difference between two models
+
+It was implemented with reference to [this discussion](https://github.com/cloneofsimo/lora/discussions/56). I used the formula as it is (I don't understand it well, but it seems that singular value decomposition is used for approximation).
+
+LoRA approximates the difference between two models (for example, the original model after fine tuning and the model after fine tuning).
+
+### How to run scripts
+
+Please specify as follows.
+```
+python networks\extract_lora_from_models.py --model_org base-model.ckpt
+     --model_tuned fine-tuned-model.ckpt
+     --save_to lora-weights.safetensors --dim 4
+```
+
+Specify the original Stable Diffusion model for the --model_org option. When applying the created LoRA model, this model will be specified and applied. .ckpt or .safetensors can be specified.
+
+Specify the Stable Diffusion model to extract the difference in the --model_tuned option. For example, specify a model after fine tuning or DreamBooth. .ckpt or .safetensors can be specified.
+
+Specify the save destination of the LoRA model in --save_to. Specify the number of dimensions of LoRA in --dim.
+
+A generated LoRA model can be used in the same way as a trained LoRA model.
+
+If the Text Encoder is the same for both models, LoRA will be U-Net only LoRA.
+
+### Other Options
+
+--v2
+   - Please specify when using the v2.x Stable Diffusion model.
+--device
+   - If cuda is specified as ``--device cuda``, the calculation will be performed on the GPU. Processing will be faster (because even the CPU is not that slow, it seems to be at most twice or several times faster).
+--save_precision
+   - Specify the LoRA save format from "float", "fp16", "bf16". Default is float.
+
+## Additional Information
+
+### Differences from cloneofsimo's repository
+
+As of 12/25, this repository has expanded LoRA application points to Text Encoder's MLP, U-Net's FFN, and Transformer's in/out projection, increasing its expressiveness. However, the amount of memory used increased instead, and it became the last minute of 8GB.
+
+Also, the module replacement mechanism is completely different.
+
+### About Future Expansion
+
+It is possible to support not only LoRA but also other expansions, so we plan to add them as well.

train_textual_inversion.py

@@ -0,0 +1,598 @@
+import importlib
+import argparse
+import gc
+import math
+import os
+import toml
+from multiprocessing import Value
+
+from tqdm import tqdm
+import torch
+from accelerate.utils import set_seed
+import diffusers
+from diffusers import DDPMScheduler
+
+import library.train_util as train_util
+import library.huggingface_util as huggingface_util
+import library.config_util as config_util
+from library.config_util import (
+    ConfigSanitizer,
+    BlueprintGenerator,
+)
+import library.custom_train_functions as custom_train_functions
+from library.custom_train_functions import apply_snr_weight
+
+imagenet_templates_small = [
+    "a photo of a {}",
+    "a rendering of a {}",
+    "a cropped photo of the {}",
+    "the photo of a {}",
+    "a photo of a clean {}",
+    "a photo of a dirty {}",
+    "a dark photo of the {}",
+    "a photo of my {}",
+    "a photo of the cool {}",
+    "a close-up photo of a {}",
+    "a bright photo of the {}",
+    "a cropped photo of a {}",
+    "a photo of the {}",
+    "a good photo of the {}",
+    "a photo of one {}",
+    "a close-up photo of the {}",
+    "a rendition of the {}",
+    "a photo of the clean {}",
+    "a rendition of a {}",
+    "a photo of a nice {}",
+    "a good photo of a {}",
+    "a photo of the nice {}",
+    "a photo of the small {}",
+    "a photo of the weird {}",
+    "a photo of the large {}",
+    "a photo of a cool {}",
+    "a photo of a small {}",
+]
+
+imagenet_style_templates_small = [
+    "a painting in the style of {}",
+    "a rendering in the style of {}",
+    "a cropped painting in the style of {}",
+    "the painting in the style of {}",
+    "a clean painting in the style of {}",
+    "a dirty painting in the style of {}",
+    "a dark painting in the style of {}",
+    "a picture in the style of {}",
+    "a cool painting in the style of {}",
+    "a close-up painting in the style of {}",
+    "a bright painting in the style of {}",
+    "a cropped painting in the style of {}",
+    "a good painting in the style of {}",
+    "a close-up painting in the style of {}",
+    "a rendition in the style of {}",
+    "a nice painting in the style of {}",
+    "a small painting in the style of {}",
+    "a weird painting in the style of {}",
+    "a large painting in the style of {}",
+]
+
+
+def train(args):
+    if args.output_name is None:
+        args.output_name = args.token_string
+    use_template = args.use_object_template or args.use_style_template
+
+    train_util.verify_training_args(args)
+    train_util.prepare_dataset_args(args, True)
+
+    cache_latents = args.cache_latents
+
+    if args.seed is not None:
+        set_seed(args.seed)
+
+    tokenizer = train_util.load_tokenizer(args)
+
+    # acceleratorを準備する
+    print("prepare accelerator")
+    accelerator, unwrap_model = train_util.prepare_accelerator(args)
+
+    # mixed precisionに対応した型を用意しておき適宜castする
+    weight_dtype, save_dtype = train_util.prepare_dtype(args)
+
+    # モデルを読み込む
+    text_encoder, vae, unet, _ = train_util.load_target_model(args, weight_dtype)
+
+    # Convert the init_word to token_id
+    if args.init_word is not None:
+        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_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)]
+    num_added_tokens = tokenizer.add_tokens(token_strings)
+    assert (
+        num_added_tokens == args.num_vectors_per_token
+    ), f"tokenizer has same word to token string. please use another one / 指定したargs.token_stringは既に存在します。別の単語を使ってください: {args.token_string}"
+
+    token_ids = tokenizer.convert_tokens_to_ids(token_strings)
+    print(f"tokens are added: {token_ids}")
+    assert min(token_ids) == token_ids[0] and token_ids[-1] == token_ids[0] + len(token_ids) - 1, f"token ids is not ordered"
+    assert len(tokenizer) - 1 == token_ids[-1], f"token ids is not end of tokenize: {len(tokenizer)}"
+
+    # Resize the token embeddings as we are adding new special tokens to the tokenizer
+    text_encoder.resize_token_embeddings(len(tokenizer))
+
+    # 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_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
+    if args.weights is not None:
+        embeddings = load_weights(args.weights)
+        assert len(token_ids) == len(
+            embeddings
+        ), f"num_vectors_per_token is mismatch for weights / 指定した重みとnum_vectors_per_tokenの値が異なります: {len(embeddings)}"
+        # print(token_ids, embeddings.size())
+        for token_id, embedding in zip(token_ids, embeddings):
+            token_embeds[token_id] = embedding
+            # print(token_id, token_embeds[token_id].mean(), token_embeds[token_id].min())
+        print(f"weighs loaded")
+
+    print(f"create embeddings for {args.num_vectors_per_token} tokens, for {args.token_string}")
+
+    # データセットを準備する
+    blueprint_generator = BlueprintGenerator(ConfigSanitizer(True, True, False))
+    if args.dataset_config is not None:
+        print(f"Load dataset config from {args.dataset_config}")
+        user_config = config_util.load_user_config(args.dataset_config)
+        ignored = ["train_data_dir", "reg_data_dir", "in_json"]
+        if any(getattr(args, attr) is not None for attr in ignored):
+            print(
+                "ignore following options because config file is found: {0} / 設定ファイルが利用されるため以下のオプションは無視されます: {0}".format(
+                    ", ".join(ignored)
+                )
+            )
+    else:
+        use_dreambooth_method = args.in_json is None
+        if use_dreambooth_method:
+            print("Use DreamBooth method.")
+            user_config = {
+                "datasets": [
+                    {"subsets": config_util.generate_dreambooth_subsets_config_by_subdirs(args.train_data_dir, args.reg_data_dir)}
+                ]
+            }
+        else:
+            print("Train with captions.")
+            user_config = {
+                "datasets": [
+                    {
+                        "subsets": [
+                            {
+                                "image_dir": args.train_data_dir,
+                                "metadata_file": args.in_json,
+                            }
+                        ]
+                    }
+                ]
+            }
+
+    blueprint = blueprint_generator.generate(user_config, args, tokenizer=tokenizer)
+    train_dataset_group = config_util.generate_dataset_group_by_blueprint(blueprint.dataset_group)
+
+    current_epoch = Value("i", 0)
+    current_step = Value("i", 0)
+    ds_for_collater = train_dataset_group if args.max_data_loader_n_workers == 0 else None
+    collater = train_util.collater_class(current_epoch, current_step, ds_for_collater)
+
+    # make captions: tokenstring tokenstring1 tokenstring2 ...tokenstringn という文字列に書き換える超乱暴な実装
+    if use_template:
+        print("use template for training captions. is object: {args.use_object_template}")
+        templates = imagenet_templates_small if args.use_object_template else imagenet_style_templates_small
+        replace_to = " ".join(token_strings)
+        captions = []
+        for tmpl in templates:
+            captions.append(tmpl.format(replace_to))
+        train_dataset_group.add_replacement("", captions)
+
+        if args.num_vectors_per_token > 1:
+            prompt_replacement = (args.token_string, replace_to)
+        else:
+            prompt_replacement = None
+    else:
+        if args.num_vectors_per_token > 1:
+            replace_to = " ".join(token_strings)
+            train_dataset_group.add_replacement(args.token_string, replace_to)
+            prompt_replacement = (args.token_string, replace_to)
+        else:
+            prompt_replacement = None
+
+    if args.debug_dataset:
+        train_util.debug_dataset(train_dataset_group, show_input_ids=True)
+        return
+    if len(train_dataset_group) == 0:
+        print("No data found. Please verify arguments / 画像がありません。引数指定を確認してください")
+        return
+
+    if cache_latents:
+        assert (
+            train_dataset_group.is_latent_cacheable()
+        ), "when caching latents, either color_aug or random_crop cannot be used / latentをキャッシュするときはcolor_augとrandom_cropは使えません"
+
+    # モデルに xformers とか memory efficient attention を組み込む
+    train_util.replace_unet_modules(unet, args.mem_eff_attn, args.xformers)
+
+    # 学習を準備する
+    if cache_latents:
+        vae.to(accelerator.device, dtype=weight_dtype)
+        vae.requires_grad_(False)
+        vae.eval()
+        with torch.no_grad():
+            train_dataset_group.cache_latents(vae, args.vae_batch_size, args.cache_latents_to_disk, accelerator.is_main_process)
+        vae.to("cpu")
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        gc.collect()
+
+        accelerator.wait_for_everyone()
+
+    if args.gradient_checkpointing:
+        unet.enable_gradient_checkpointing()
+        text_encoder.gradient_checkpointing_enable()
+
+    # 学習に必要なクラスを準備する
+    print("prepare optimizer, data loader etc.")
+    trainable_params = text_encoder.get_input_embeddings().parameters()
+    _, _, optimizer = train_util.get_optimizer(args, trainable_params)
+
+    # dataloaderを準備する
+    # DataLoaderのプロセス数：0はメインプロセスになる
+    n_workers = min(args.max_data_loader_n_workers, os.cpu_count() - 1)  # cpu_count-1 ただし最大で指定された数まで
+    train_dataloader = torch.utils.data.DataLoader(
+        train_dataset_group,
+        batch_size=1,
+        shuffle=True,
+        collate_fn=collater,
+        num_workers=n_workers,
+        persistent_workers=args.persistent_data_loader_workers,
+    )
+
+    # 学習ステップ数を計算する
+    if args.max_train_epochs is not None:
+        args.max_train_steps = args.max_train_epochs * math.ceil(
+            len(train_dataloader) / accelerator.num_processes / args.gradient_accumulation_steps
+        )
+        print(f"override steps. steps for {args.max_train_epochs} epochs is / 指定エポックまでのステップ数: {args.max_train_steps}")
+
+    # データセット側にも学習ステップを送信
+    train_dataset_group.set_max_train_steps(args.max_train_steps)
+
+    # lr schedulerを用意する
+    lr_scheduler = train_util.get_scheduler_fix(args, optimizer, accelerator.num_processes)
+
+    # acceleratorがなんかよろしくやってくれるらしい
+    text_encoder, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
+        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))
+    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
+    text_encoder.requires_grad_(True)
+    text_encoder.text_model.encoder.requires_grad_(False)
+    text_encoder.text_model.final_layer_norm.requires_grad_(False)
+    text_encoder.text_model.embeddings.position_embedding.requires_grad_(False)
+    # text_encoder.text_model.embeddings.token_embedding.requires_grad_(True)
+
+    unet.requires_grad_(False)
+    unet.to(accelerator.device, dtype=weight_dtype)
+    if args.gradient_checkpointing:  # according to TI example in Diffusers, train is required
+        unet.train()
+    else:
+        unet.eval()
+
+    if not cache_latents:
+        vae.requires_grad_(False)
+        vae.eval()
+        vae.to(accelerator.device, dtype=weight_dtype)
+
+    # 実験的機能：勾配も含めたfp16学習を行う　PyTorchにパッチを当ててfp16でのgrad scaleを有効にする
+    if args.full_fp16:
+        train_util.patch_accelerator_for_fp16_training(accelerator)
+        text_encoder.to(weight_dtype)
+
+    # resumeする
+    train_util.resume_from_local_or_hf_if_specified(accelerator, args)
+
+    # epoch数を計算する
+    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
+    num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)
+    if (args.save_n_epoch_ratio is not None) and (args.save_n_epoch_ratio > 0):
+        args.save_every_n_epochs = math.floor(num_train_epochs / args.save_n_epoch_ratio) or 1
+
+    # 学習する
+    total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
+    print("running training / 学習開始")
+    print(f"  num train images * repeats / 学習画像の数×繰り返し回数: {train_dataset_group.num_train_images}")
+    print(f"  num reg images / 正則化画像の数: {train_dataset_group.num_reg_images}")
+    print(f"  num batches per epoch / 1epochのバッチ数: {len(train_dataloader)}")
+    print(f"  num epochs / epoch数: {num_train_epochs}")
+    print(f"  batch size per device / バッチサイズ: {args.train_batch_size}")
+    print(f"  total train batch size (with parallel & distributed & accumulation) / 総バッチサイズ（並列学習、勾配合計含む）: {total_batch_size}")
+    print(f"  gradient ccumulation steps / 勾配を合計するステップ数 = {args.gradient_accumulation_steps}")
+    print(f"  total optimization steps / 学習ステップ数: {args.max_train_steps}")
+
+    progress_bar = tqdm(range(args.max_train_steps), smoothing=0, disable=not accelerator.is_local_main_process, desc="steps")
+    global_step = 0
+
+    noise_scheduler = DDPMScheduler(
+        beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000, clip_sample=False
+    )
+
+    if accelerator.is_main_process:
+        accelerator.init_trackers("textual_inversion")
+
+    for epoch in range(num_train_epochs):
+        print(f"epoch {epoch+1}/{num_train_epochs}")
+        current_epoch.value = epoch + 1
+
+        text_encoder.train()
+
+        loss_total = 0
+
+        for step, batch in enumerate(train_dataloader):
+            current_step.value = global_step
+            with accelerator.accumulate(text_encoder):
+                with torch.no_grad():
+                    if "latents" in batch and batch["latents"] is not None:
+                        latents = batch["latents"].to(accelerator.device)
+                    else:
+                        # latentに変換
+                        latents = vae.encode(batch["images"].to(dtype=weight_dtype)).latent_dist.sample()
+                    latents = latents * 0.18215
+                b_size = latents.shape[0]
+
+                # Get the text embedding for conditioning
+                input_ids = batch["input_ids"].to(accelerator.device)
+                # 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)
+                timesteps = timesteps.long()
+
+                # Add noise to the latents according to the noise magnitude at each timestep
+                # (this is the forward diffusion process)
+                noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)
+
+                # Predict the noise residual
+                with accelerator.autocast():
+                    noise_pred = unet(noisy_latents, timesteps, encoder_hidden_states).sample
+
+                if args.v_parameterization:
+                    # v-parameterization training
+                    target = noise_scheduler.get_velocity(latents, noise, timesteps)
+                else:
+                    target = noise
+
+                loss = torch.nn.functional.mse_loss(noise_pred.float(), target.float(), reduction="none")
+                loss = loss.mean([1, 2, 3])
+
+                if args.min_snr_gamma:
+                    loss = apply_snr_weight(loss, timesteps, noise_scheduler, args.min_snr_gamma)
+
+                loss_weights = batch["loss_weights"]  # 各sampleごとのweight
+                loss = loss * loss_weights
+
+                loss = loss.mean()  # 平均なのでbatch_sizeで割る必要なし
+
+                accelerator.backward(loss)
+                if accelerator.sync_gradients and args.max_grad_norm != 0.0:
+                    params_to_clip = text_encoder.get_input_embeddings().parameters()
+                    accelerator.clip_grad_norm_(params_to_clip, args.max_grad_norm)
+
+                optimizer.step()
+                lr_scheduler.step()
+                optimizer.zero_grad(set_to_none=True)
+
+                # Let's make sure we don't update any embedding weights besides the newly added token
+                with torch.no_grad():
+                    unwrap_model(text_encoder).get_input_embeddings().weight[index_no_updates] = orig_embeds_params[
+                        index_no_updates
+                    ]
+
+            # Checks if the accelerator has performed an optimization step behind the scenes
+            if accelerator.sync_gradients:
+                progress_bar.update(1)
+                global_step += 1
+
+                train_util.sample_images(
+                    accelerator, args, None, global_step, accelerator.device, vae, tokenizer, text_encoder, unet, prompt_replacement
+                )
+
+            current_loss = loss.detach().item()
+            if args.logging_dir is not None:
+                logs = {"loss": current_loss, "lr": float(lr_scheduler.get_last_lr()[0])}
+                if args.optimizer_type.lower() == "DAdaptation".lower():  # tracking d*lr value
+                    logs["lr/d*lr"] = (
+                        lr_scheduler.optimizers[0].param_groups[0]["d"] * lr_scheduler.optimizers[0].param_groups[0]["lr"]
+                    )
+                accelerator.log(logs, step=global_step)
+
+            loss_total += current_loss
+            avr_loss = loss_total / (step + 1)
+            logs = {"loss": avr_loss}  # , "lr": lr_scheduler.get_last_lr()[0]}
+            progress_bar.set_postfix(**logs)
+
+            if global_step >= args.max_train_steps:
+                break
+
+        if args.logging_dir is not None:
+            logs = {"loss/epoch": loss_total / len(train_dataloader)}
+            accelerator.log(logs, step=epoch + 1)
+
+        accelerator.wait_for_everyone()
+
+        updated_embs = unwrap_model(text_encoder).get_input_embeddings().weight[token_ids].data.detach().clone()
+
+        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
+
+            def save_func():
+                ckpt_name = train_util.EPOCH_FILE_NAME.format(model_name, epoch + 1) + "." + args.save_model_as
+                ckpt_file = os.path.join(args.output_dir, ckpt_name)
+                print(f"saving checkpoint: {ckpt_file}")
+                save_weights(ckpt_file, updated_embs, save_dtype)
+                if args.huggingface_repo_id is not None:
+                    huggingface_util.upload(args, ckpt_file, "/" + ckpt_name)
+
+            def remove_old_func(old_epoch_no):
+                old_ckpt_name = train_util.EPOCH_FILE_NAME.format(model_name, old_epoch_no) + "." + args.save_model_as
+                old_ckpt_file = os.path.join(args.output_dir, old_ckpt_name)
+                if os.path.exists(old_ckpt_file):
+                    print(f"removing old checkpoint: {old_ckpt_file}")
+                    os.remove(old_ckpt_file)
+
+            saving = train_util.save_on_epoch_end(args, save_func, remove_old_func, epoch + 1, num_train_epochs)
+            if saving and args.save_state:
+                train_util.save_state_on_epoch_end(args, accelerator, model_name, epoch + 1)
+
+        train_util.sample_images(
+            accelerator, args, epoch + 1, global_step, accelerator.device, vae, tokenizer, text_encoder, unet, prompt_replacement
+        )
+
+        # end of epoch
+
+    is_main_process = accelerator.is_main_process
+    if is_main_process:
+        text_encoder = unwrap_model(text_encoder)
+
+    accelerator.end_training()
+
+    if args.save_state:
+        train_util.save_state_on_train_end(args, accelerator)
+
+    updated_embs = text_encoder.get_input_embeddings().weight[token_ids].data.detach().clone()
+
+    del accelerator  # この後メモリを使うのでこれは消す
+
+    if is_main_process:
+        os.makedirs(args.output_dir, exist_ok=True)
+
+        model_name = train_util.DEFAULT_LAST_OUTPUT_NAME if args.output_name is None else args.output_name
+        ckpt_name = model_name + "." + args.save_model_as
+        ckpt_file = os.path.join(args.output_dir, ckpt_name)
+
+        print(f"save trained model to {ckpt_file}")
+        save_weights(ckpt_file, updated_embs, save_dtype)
+        if args.huggingface_repo_id is not None:
+            huggingface_util.upload(args, ckpt_file, "/" + ckpt_name, force_sync_upload=True)
+        print("model saved.")
+
+
+def save_weights(file, updated_embs, save_dtype):
+    state_dict = {"emb_params": updated_embs}
+
+    if save_dtype is not None:
+        for key in list(state_dict.keys()):
+            v = state_dict[key]
+            v = v.detach().clone().to("cpu").to(save_dtype)
+            state_dict[key] = v
+
+    if os.path.splitext(file)[1] == ".safetensors":
+        from safetensors.torch import save_file
+
+        save_file(state_dict, file)
+    else:
+        torch.save(state_dict, file)  # can be loaded in Web UI
+
+
+def load_weights(file):
+    if os.path.splitext(file)[1] == ".safetensors":
+        from safetensors.torch import load_file
+
+        data = load_file(file)
+    else:
+        # compatible to Web UI's file format
+        data = torch.load(file, map_location="cpu")
+        if type(data) != dict:
+            raise ValueError(f"weight file is not dict / 重みファイルがdict形式ではありません: {file}")
+
+        if "string_to_param" in data:  # textual inversion embeddings
+            data = data["string_to_param"]
+            if hasattr(data, "_parameters"):  # support old PyTorch?
+                data = getattr(data, "_parameters")
+
+    emb = next(iter(data.values()))
+    if type(emb) != torch.Tensor:
+        raise ValueError(f"weight file does not contains Tensor / 重みファイルのデータがTensorではありません: {file}")
+
+    if len(emb.size()) == 1:
+        emb = emb.unsqueeze(0)
+
+    return emb
+
+
+def setup_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser()
+
+    train_util.add_sd_models_arguments(parser)
+    train_util.add_dataset_arguments(parser, True, True, False)
+    train_util.add_training_arguments(parser, True)
+    train_util.add_optimizer_arguments(parser)
+    config_util.add_config_arguments(parser)
+    custom_train_functions.add_custom_train_arguments(parser, False)
+
+    parser.add_argument(
+        "--save_model_as",
+        type=str,
+        default="pt",
+        choices=[None, "ckpt", "pt", "safetensors"],
+        help="format to save the model (default is .pt) / モデル保存時の形式（デフォルトはpt）",
+    )
+
+    parser.add_argument("--weights", type=str, default=None, help="embedding weights to initialize / 学習するネットワークの初期重み")
+    parser.add_argument(
+        "--num_vectors_per_token", type=int, default=1, help="number of vectors per token / トークンに割り当てるembeddingsの要素数"
+    )
+    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="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",
+        help="ignore caption and use default templates for stype / キャプションは使わずデフォルトのスタイル用テンプレートで学習する",
+    )
+
+    return parser
+
+
+if __name__ == "__main__":
+    parser = setup_parser()
+
+    args = parser.parse_args()
+    args = train_util.read_config_from_file(args, parser)
+
+    train(args)

train_textual_inversion_XTI.py

@@ -0,0 +1,650 @@
+import importlib
+import argparse
+import gc
+import math
+import os
+import toml
+from multiprocessing import Value
+
+from tqdm import tqdm
+import torch
+from accelerate.utils import set_seed
+import diffusers
+from diffusers import DDPMScheduler
+
+import library.train_util as train_util
+import library.huggingface_util as huggingface_util
+import library.config_util as config_util
+from library.config_util import (
+    ConfigSanitizer,
+    BlueprintGenerator,
+)
+import library.custom_train_functions as custom_train_functions
+from library.custom_train_functions import apply_snr_weight
+from XTI_hijack import unet_forward_XTI, downblock_forward_XTI, upblock_forward_XTI
+
+imagenet_templates_small = [
+    "a photo of a {}",
+    "a rendering of a {}",
+    "a cropped photo of the {}",
+    "the photo of a {}",
+    "a photo of a clean {}",
+    "a photo of a dirty {}",
+    "a dark photo of the {}",
+    "a photo of my {}",
+    "a photo of the cool {}",
+    "a close-up photo of a {}",
+    "a bright photo of the {}",
+    "a cropped photo of a {}",
+    "a photo of the {}",
+    "a good photo of the {}",
+    "a photo of one {}",
+    "a close-up photo of the {}",
+    "a rendition of the {}",
+    "a photo of the clean {}",
+    "a rendition of a {}",
+    "a photo of a nice {}",
+    "a good photo of a {}",
+    "a photo of the nice {}",
+    "a photo of the small {}",
+    "a photo of the weird {}",
+    "a photo of the large {}",
+    "a photo of a cool {}",
+    "a photo of a small {}",
+]
+
+imagenet_style_templates_small = [
+    "a painting in the style of {}",
+    "a rendering in the style of {}",
+    "a cropped painting in the style of {}",
+    "the painting in the style of {}",
+    "a clean painting in the style of {}",
+    "a dirty painting in the style of {}",
+    "a dark painting in the style of {}",
+    "a picture in the style of {}",
+    "a cool painting in the style of {}",
+    "a close-up painting in the style of {}",
+    "a bright painting in the style of {}",
+    "a cropped painting in the style of {}",
+    "a good painting in the style of {}",
+    "a close-up painting in the style of {}",
+    "a rendition in the style of {}",
+    "a nice painting in the style of {}",
+    "a small painting in the style of {}",
+    "a weird painting in the style of {}",
+    "a large painting in the style of {}",
+]
+
+
+def train(args):
+    if args.output_name is None:
+        args.output_name = args.token_string
+    use_template = args.use_object_template or args.use_style_template
+
+    train_util.verify_training_args(args)
+    train_util.prepare_dataset_args(args, True)
+
+    if args.sample_every_n_steps is not None or args.sample_every_n_epochs is not None:
+        print(
+            "sample_every_n_steps and sample_every_n_epochs are not supported in this script currently / sample_every_n_stepsとsample_every_n_epochsは現在このスクリプトではサポートされていません"
+        )
+
+    cache_latents = args.cache_latents
+
+    if args.seed is not None:
+        set_seed(args.seed)
+
+    tokenizer = train_util.load_tokenizer(args)
+
+    # acceleratorを準備する
+    print("prepare accelerator")
+    accelerator, unwrap_model = train_util.prepare_accelerator(args)
+
+    # mixed precisionに対応した型を用意しておき適宜castする
+    weight_dtype, save_dtype = train_util.prepare_dtype(args)
+
+    # モデルを読み込む
+    text_encoder, vae, unet, _ = train_util.load_target_model(args, weight_dtype)
+
+    # Convert the init_word to token_id
+    if args.init_word is not None:
+        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_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)]
+    num_added_tokens = tokenizer.add_tokens(token_strings)
+    assert (
+        num_added_tokens == args.num_vectors_per_token
+    ), f"tokenizer has same word to token string. please use another one / 指定したargs.token_stringは既に存在します。別の単語を使ってください: {args.token_string}"
+
+    token_ids = tokenizer.convert_tokens_to_ids(token_strings)
+    print(f"tokens are added: {token_ids}")
+    assert min(token_ids) == token_ids[0] and token_ids[-1] == token_ids[0] + len(token_ids) - 1, f"token ids is not ordered"
+    assert len(tokenizer) - 1 == token_ids[-1], f"token ids is not end of tokenize: {len(tokenizer)}"
+
+    token_strings_XTI = []
+    XTI_layers = [
+        "IN01",
+        "IN02",
+        "IN04",
+        "IN05",
+        "IN07",
+        "IN08",
+        "MID",
+        "OUT03",
+        "OUT04",
+        "OUT05",
+        "OUT06",
+        "OUT07",
+        "OUT08",
+        "OUT09",
+        "OUT10",
+        "OUT11",
+    ]
+    for layer_name in XTI_layers:
+        token_strings_XTI += [f"{t}_{layer_name}" for t in token_strings]
+
+    tokenizer.add_tokens(token_strings_XTI)
+    token_ids_XTI = tokenizer.convert_tokens_to_ids(token_strings_XTI)
+    print(f"tokens are added (XTI): {token_ids_XTI}")
+    # Resize the token embeddings as we are adding new special tokens to the tokenizer
+    text_encoder.resize_token_embeddings(len(tokenizer))
+
+    # 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_ids is not None:
+        for i, token_id in enumerate(token_ids_XTI):
+            token_embeds[token_id] = token_embeds[init_token_ids[(i // 16) % len(init_token_ids)]]
+            # print(token_id, token_embeds[token_id].mean(), token_embeds[token_id].min())
+
+    # load weights
+    if args.weights is not None:
+        embeddings = load_weights(args.weights)
+        assert len(token_ids) == len(
+            embeddings
+        ), f"num_vectors_per_token is mismatch for weights / 指定した重みとnum_vectors_per_tokenの値が異なります: {len(embeddings)}"
+        # print(token_ids, embeddings.size())
+        for token_id, embedding in zip(token_ids_XTI, embeddings):
+            token_embeds[token_id] = embedding
+            # print(token_id, token_embeds[token_id].mean(), token_embeds[token_id].min())
+        print(f"weighs loaded")
+
+    print(f"create embeddings for {args.num_vectors_per_token} tokens, for {args.token_string}")
+
+    # データセットを準備する
+    blueprint_generator = BlueprintGenerator(ConfigSanitizer(True, True, False))
+    if args.dataset_config is not None:
+        print(f"Load dataset config from {args.dataset_config}")
+        user_config = config_util.load_user_config(args.dataset_config)
+        ignored = ["train_data_dir", "reg_data_dir", "in_json"]
+        if any(getattr(args, attr) is not None for attr in ignored):
+            print(
+                "ignore following options because config file is found: {0} / 設定ファイルが利用されるため以下のオプションは無視されます: {0}".format(
+                    ", ".join(ignored)
+                )
+            )
+    else:
+        use_dreambooth_method = args.in_json is None
+        if use_dreambooth_method:
+            print("Use DreamBooth method.")
+            user_config = {
+                "datasets": [
+                    {"subsets": config_util.generate_dreambooth_subsets_config_by_subdirs(args.train_data_dir, args.reg_data_dir)}
+                ]
+            }
+        else:
+            print("Train with captions.")
+            user_config = {
+                "datasets": [
+                    {
+                        "subsets": [
+                            {
+                                "image_dir": args.train_data_dir,
+                                "metadata_file": args.in_json,
+                            }
+                        ]
+                    }
+                ]
+            }
+
+    blueprint = blueprint_generator.generate(user_config, args, tokenizer=tokenizer)
+    train_dataset_group = config_util.generate_dataset_group_by_blueprint(blueprint.dataset_group)
+    train_dataset_group.enable_XTI(XTI_layers, token_strings=token_strings)
+    current_epoch = Value("i", 0)
+    current_step = Value("i", 0)
+    ds_for_collater = train_dataset_group if args.max_data_loader_n_workers == 0 else None
+    collater = train_util.collater_class(current_epoch, current_step, ds_for_collater)
+
+    # make captions: tokenstring tokenstring1 tokenstring2 ...tokenstringn という文字列に書き換える超乱暴な実装
+    if use_template:
+        print("use template for training captions. is object: {args.use_object_template}")
+        templates = imagenet_templates_small if args.use_object_template else imagenet_style_templates_small
+        replace_to = " ".join(token_strings)
+        captions = []
+        for tmpl in templates:
+            captions.append(tmpl.format(replace_to))
+        train_dataset_group.add_replacement("", captions)
+
+        if args.num_vectors_per_token > 1:
+            prompt_replacement = (args.token_string, replace_to)
+        else:
+            prompt_replacement = None
+    else:
+        if args.num_vectors_per_token > 1:
+            replace_to = " ".join(token_strings)
+            train_dataset_group.add_replacement(args.token_string, replace_to)
+            prompt_replacement = (args.token_string, replace_to)
+        else:
+            prompt_replacement = None
+
+    if args.debug_dataset:
+        train_util.debug_dataset(train_dataset_group, show_input_ids=True)
+        return
+    if len(train_dataset_group) == 0:
+        print("No data found. Please verify arguments / 画像がありません。引数指定を確認してください")
+        return
+
+    if cache_latents:
+        assert (
+            train_dataset_group.is_latent_cacheable()
+        ), "when caching latents, either color_aug or random_crop cannot be used / latentをキャッシュするときはcolor_augとrandom_cropは使えません"
+
+    # モデルに xformers とか memory efficient attention を組み込む
+    train_util.replace_unet_modules(unet, args.mem_eff_attn, args.xformers)
+    diffusers.models.UNet2DConditionModel.forward = unet_forward_XTI
+    diffusers.models.unet_2d_blocks.CrossAttnDownBlock2D.forward = downblock_forward_XTI
+    diffusers.models.unet_2d_blocks.CrossAttnUpBlock2D.forward = upblock_forward_XTI
+
+    # 学習を準備する
+    if cache_latents:
+        vae.to(accelerator.device, dtype=weight_dtype)
+        vae.requires_grad_(False)
+        vae.eval()
+        with torch.no_grad():
+            train_dataset_group.cache_latents(vae, args.vae_batch_size, args.cache_latents_to_disk, accelerator.is_main_process)
+        vae.to("cpu")
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        gc.collect()
+
+        accelerator.wait_for_everyone()
+
+    if args.gradient_checkpointing:
+        unet.enable_gradient_checkpointing()
+        text_encoder.gradient_checkpointing_enable()
+
+    # 学習に必要なクラスを準備する
+    print("prepare optimizer, data loader etc.")
+    trainable_params = text_encoder.get_input_embeddings().parameters()
+    _, _, optimizer = train_util.get_optimizer(args, trainable_params)
+
+    # dataloaderを準備する
+    # DataLoaderのプロセス数：0はメインプロセスになる
+    n_workers = min(args.max_data_loader_n_workers, os.cpu_count() - 1)  # cpu_count-1 ただし最大で指定された数まで
+    train_dataloader = torch.utils.data.DataLoader(
+        train_dataset_group,
+        batch_size=1,
+        shuffle=True,
+        collate_fn=collater,
+        num_workers=n_workers,
+        persistent_workers=args.persistent_data_loader_workers,
+    )
+
+    # 学習ステップ数を計算する
+    if args.max_train_epochs is not None:
+        args.max_train_steps = args.max_train_epochs * math.ceil(
+            len(train_dataloader) / accelerator.num_processes / args.gradient_accumulation_steps
+        )
+        print(f"override steps. steps for {args.max_train_epochs} epochs is / 指定エポックまでのステップ数: {args.max_train_steps}")
+
+    # データセット側にも学習ステップを送信
+    train_dataset_group.set_max_train_steps(args.max_train_steps)
+
+    # lr schedulerを用意する
+    lr_scheduler = train_util.get_scheduler_fix(args, optimizer, accelerator.num_processes)
+
+    # acceleratorがなんかよろしくやってくれるらしい
+    text_encoder, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
+        text_encoder, optimizer, train_dataloader, lr_scheduler
+    )
+
+    index_no_updates = torch.arange(len(tokenizer)) < token_ids_XTI[0]
+    # 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
+    text_encoder.requires_grad_(True)
+    text_encoder.text_model.encoder.requires_grad_(False)
+    text_encoder.text_model.final_layer_norm.requires_grad_(False)
+    text_encoder.text_model.embeddings.position_embedding.requires_grad_(False)
+    # text_encoder.text_model.embeddings.token_embedding.requires_grad_(True)
+
+    unet.requires_grad_(False)
+    unet.to(accelerator.device, dtype=weight_dtype)
+    if args.gradient_checkpointing:  # according to TI example in Diffusers, train is required
+        unet.train()
+    else:
+        unet.eval()
+
+    if not cache_latents:
+        vae.requires_grad_(False)
+        vae.eval()
+        vae.to(accelerator.device, dtype=weight_dtype)
+
+    # 実験的機能：勾配も含めたfp16学習を行う　PyTorchにパッチを当ててfp16でのgrad scaleを有効にする
+    if args.full_fp16:
+        train_util.patch_accelerator_for_fp16_training(accelerator)
+        text_encoder.to(weight_dtype)
+
+    # resumeする
+    train_util.resume_from_local_or_hf_if_specified(accelerator, args)
+
+    # epoch数を計算する
+    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
+    num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)
+    if (args.save_n_epoch_ratio is not None) and (args.save_n_epoch_ratio > 0):
+        args.save_every_n_epochs = math.floor(num_train_epochs / args.save_n_epoch_ratio) or 1
+
+    # 学習する
+    total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
+    print("running training / 学習開始")
+    print(f"  num train images * repeats / 学習画像の数×繰り返し回数: {train_dataset_group.num_train_images}")
+    print(f"  num reg images / 正則化画像の数: {train_dataset_group.num_reg_images}")
+    print(f"  num batches per epoch / 1epochのバッチ数: {len(train_dataloader)}")
+    print(f"  num epochs / epoch数: {num_train_epochs}")
+    print(f"  batch size per device / バッチサイズ: {args.train_batch_size}")
+    print(f"  total train batch size (with parallel & distributed & accumulation) / 総バッチサイズ（並列学習、勾配合計含む）: {total_batch_size}")
+    print(f"  gradient ccumulation steps / 勾配を合計するステップ数 = {args.gradient_accumulation_steps}")
+    print(f"  total optimization steps / 学習ステップ数: {args.max_train_steps}")
+
+    progress_bar = tqdm(range(args.max_train_steps), smoothing=0, disable=not accelerator.is_local_main_process, desc="steps")
+    global_step = 0
+
+    noise_scheduler = DDPMScheduler(
+        beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000, clip_sample=False
+    )
+
+    if accelerator.is_main_process:
+        accelerator.init_trackers("textual_inversion")
+
+    for epoch in range(num_train_epochs):
+        print(f"epoch {epoch+1}/{num_train_epochs}")
+        current_epoch.value = epoch + 1
+
+        text_encoder.train()
+
+        loss_total = 0
+
+        for step, batch in enumerate(train_dataloader):
+            current_step.value = global_step
+            with accelerator.accumulate(text_encoder):
+                with torch.no_grad():
+                    if "latents" in batch and batch["latents"] is not None:
+                        latents = batch["latents"].to(accelerator.device)
+                    else:
+                        # latentに変換
+                        latents = vae.encode(batch["images"].to(dtype=weight_dtype)).latent_dist.sample()
+                    latents = latents * 0.18215
+                b_size = latents.shape[0]
+
+                # Get the text embedding for conditioning
+                input_ids = batch["input_ids"].to(accelerator.device)
+                # weight_dtype) use float instead of fp16/bf16 because text encoder is float
+                encoder_hidden_states = torch.stack(
+                    [
+                        train_util.get_hidden_states(args, s, tokenizer, text_encoder, weight_dtype)
+                        for s in torch.split(input_ids, 1, dim=1)
+                    ]
+                )
+
+                # 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)
+                timesteps = timesteps.long()
+
+                # Add noise to the latents according to the noise magnitude at each timestep
+                # (this is the forward diffusion process)
+                noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)
+
+                # Predict the noise residual
+                with accelerator.autocast():
+                    noise_pred = unet(noisy_latents, timesteps, encoder_hidden_states=encoder_hidden_states).sample
+
+                if args.v_parameterization:
+                    # v-parameterization training
+                    target = noise_scheduler.get_velocity(latents, noise, timesteps)
+                else:
+                    target = noise
+
+                loss = torch.nn.functional.mse_loss(noise_pred.float(), target.float(), reduction="none")
+                loss = loss.mean([1, 2, 3])
+
+                if args.min_snr_gamma:
+                    loss = apply_snr_weight(loss, timesteps, noise_scheduler, args.min_snr_gamma)
+
+                loss_weights = batch["loss_weights"]  # 各sampleごとのweight
+                loss = loss * loss_weights
+
+                loss = loss.mean()  # 平均なのでbatch_sizeで割る必要なし
+
+                accelerator.backward(loss)
+                if accelerator.sync_gradients and args.max_grad_norm != 0.0:
+                    params_to_clip = text_encoder.get_input_embeddings().parameters()
+                    accelerator.clip_grad_norm_(params_to_clip, args.max_grad_norm)
+
+                optimizer.step()
+                lr_scheduler.step()
+                optimizer.zero_grad(set_to_none=True)
+
+                # Let's make sure we don't update any embedding weights besides the newly added token
+                with torch.no_grad():
+                    unwrap_model(text_encoder).get_input_embeddings().weight[index_no_updates] = orig_embeds_params[
+                        index_no_updates
+                    ]
+
+            # Checks if the accelerator has performed an optimization step behind the scenes
+            if accelerator.sync_gradients:
+                progress_bar.update(1)
+                global_step += 1
+                # TODO: fix sample_images
+                # train_util.sample_images(
+                #     accelerator, args, None, global_step, accelerator.device, vae, tokenizer, text_encoder, unet, prompt_replacement
+                # )
+
+            current_loss = loss.detach().item()
+            if args.logging_dir is not None:
+                logs = {"loss": current_loss, "lr": float(lr_scheduler.get_last_lr()[0])}
+                if args.optimizer_type.lower() == "DAdaptation".lower():  # tracking d*lr value
+                    logs["lr/d*lr"] = (
+                        lr_scheduler.optimizers[0].param_groups[0]["d"] * lr_scheduler.optimizers[0].param_groups[0]["lr"]
+                    )
+                accelerator.log(logs, step=global_step)
+
+            loss_total += current_loss
+            avr_loss = loss_total / (step + 1)
+            logs = {"loss": avr_loss}  # , "lr": lr_scheduler.get_last_lr()[0]}
+            progress_bar.set_postfix(**logs)
+
+            if global_step >= args.max_train_steps:
+                break
+
+        if args.logging_dir is not None:
+            logs = {"loss/epoch": loss_total / len(train_dataloader)}
+            accelerator.log(logs, step=epoch + 1)
+
+        accelerator.wait_for_everyone()
+
+        updated_embs = unwrap_model(text_encoder).get_input_embeddings().weight[token_ids_XTI].data.detach().clone()
+
+        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
+
+            def save_func():
+                ckpt_name = train_util.EPOCH_FILE_NAME.format(model_name, epoch + 1) + "." + args.save_model_as
+                ckpt_file = os.path.join(args.output_dir, ckpt_name)
+                print(f"saving checkpoint: {ckpt_file}")
+                save_weights(ckpt_file, updated_embs, save_dtype)
+                if args.huggingface_repo_id is not None:
+                    huggingface_util.upload(args, ckpt_file, "/" + ckpt_name)
+
+            def remove_old_func(old_epoch_no):
+                old_ckpt_name = train_util.EPOCH_FILE_NAME.format(model_name, old_epoch_no) + "." + args.save_model_as
+                old_ckpt_file = os.path.join(args.output_dir, old_ckpt_name)
+                if os.path.exists(old_ckpt_file):
+                    print(f"removing old checkpoint: {old_ckpt_file}")
+                    os.remove(old_ckpt_file)
+
+            saving = train_util.save_on_epoch_end(args, save_func, remove_old_func, epoch + 1, num_train_epochs)
+            if saving and args.save_state:
+                train_util.save_state_on_epoch_end(args, accelerator, model_name, epoch + 1)
+
+        # TODO: fix sample_images
+        # train_util.sample_images(
+        #     accelerator, args, epoch + 1, global_step, accelerator.device, vae, tokenizer, text_encoder, unet, prompt_replacement
+        # )
+
+        # end of epoch
+
+    is_main_process = accelerator.is_main_process
+    if is_main_process:
+        text_encoder = unwrap_model(text_encoder)
+
+    accelerator.end_training()
+
+    if args.save_state:
+        train_util.save_state_on_train_end(args, accelerator)
+
+    updated_embs = text_encoder.get_input_embeddings().weight[token_ids_XTI].data.detach().clone()
+
+    del accelerator  # この後メモリを使うのでこれは消す
+
+    if is_main_process:
+        os.makedirs(args.output_dir, exist_ok=True)
+
+        model_name = train_util.DEFAULT_LAST_OUTPUT_NAME if args.output_name is None else args.output_name
+        ckpt_name = model_name + "." + args.save_model_as
+        ckpt_file = os.path.join(args.output_dir, ckpt_name)
+
+        print(f"save trained model to {ckpt_file}")
+        save_weights(ckpt_file, updated_embs, save_dtype)
+        if args.huggingface_repo_id is not None:
+            huggingface_util.upload(args, ckpt_file, "/" + ckpt_name, force_sync_upload=True)
+        print("model saved.")
+
+
+def save_weights(file, updated_embs, save_dtype):
+    updated_embs = updated_embs.reshape(16, -1, updated_embs.shape[-1])
+    updated_embs = updated_embs.chunk(16)
+    XTI_layers = [
+        "IN01",
+        "IN02",
+        "IN04",
+        "IN05",
+        "IN07",
+        "IN08",
+        "MID",
+        "OUT03",
+        "OUT04",
+        "OUT05",
+        "OUT06",
+        "OUT07",
+        "OUT08",
+        "OUT09",
+        "OUT10",
+        "OUT11",
+    ]
+    state_dict = {}
+    for i, layer_name in enumerate(XTI_layers):
+        state_dict[layer_name] = updated_embs[i].squeeze(0).detach().clone().to("cpu").to(save_dtype)
+
+    # if save_dtype is not None:
+    #     for key in list(state_dict.keys()):
+    #         v = state_dict[key]
+    #         v = v.detach().clone().to("cpu").to(save_dtype)
+    #         state_dict[key] = v
+
+    if os.path.splitext(file)[1] == ".safetensors":
+        from safetensors.torch import save_file
+
+        save_file(state_dict, file)
+    else:
+        torch.save(state_dict, file)  # can be loaded in Web UI
+
+
+def load_weights(file):
+    if os.path.splitext(file)[1] == ".safetensors":
+        from safetensors.torch import load_file
+
+        data = load_file(file)
+    else:
+        raise ValueError(f"NOT XTI: {file}")
+
+    if len(data.values()) != 16:
+        raise ValueError(f"NOT XTI: {file}")
+
+    emb = torch.concat([x for x in data.values()])
+
+    return emb
+
+
+def setup_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser()
+
+    train_util.add_sd_models_arguments(parser)
+    train_util.add_dataset_arguments(parser, True, True, False)
+    train_util.add_training_arguments(parser, True)
+    train_util.add_optimizer_arguments(parser)
+    config_util.add_config_arguments(parser)
+    custom_train_functions.add_custom_train_arguments(parser, False)
+
+    parser.add_argument(
+        "--save_model_as",
+        type=str,
+        default="pt",
+        choices=[None, "ckpt", "pt", "safetensors"],
+        help="format to save the model (default is .pt) / モデル保存時の形式（デフォルトはpt）",
+    )
+
+    parser.add_argument("--weights", type=str, default=None, help="embedding weights to initialize / 学習するネットワークの初期重み")
+    parser.add_argument(
+        "--num_vectors_per_token", type=int, default=1, help="number of vectors per token / トークンに割り当てるembeddingsの要素数"
+    )
+    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="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",
+        help="ignore caption and use default templates for stype / キャプションは使わずデフォルトのスタイル用テンプレートで学習する",
+    )
+
+    return parser
+
+
+if __name__ == "__main__":
+    parser = setup_parser()
+
+    args = parser.parse_args()
+    args = train_util.read_config_from_file(args, parser)
+
+    train(args)

train_ti_README-ja.md

@@ -0,0 +1,105 @@
+[Textual Inversion](https://textual-inversion.github.io/) の学習についての説明です。
+
+[学習についての共通ドキュメント](./train_README-ja.md) もあわせてご覧ください。
+
+実装に当たっては https://github.com/huggingface/diffusers/tree/main/examples/textual_inversion を大いに参考にしました。
+
+学習したモデルはWeb UIでもそのまま使えます。
+
+# 学習の手順
+
+あらかじめこのリポジトリのREADMEを参照し、環境整備を行ってください。
+
+## データの準備
+
+[学習データの準備について](./train_README-ja.md) を参照してください。
+
+## 学習の実行
+
+``train_textual_inversion.py`` を用います。以下はコマンドラインの例です（DreamBooth手法）。
+
+```
+accelerate launch --num_cpu_threads_per_process 1 train_textual_inversion.py 
+    --dataset_config=<データ準備で作成した.tomlファイル> 
+    --output_dir=<学習したモデルの出力先フォルダ>  
+    --output_name=<学習したモデル出力時のファイル名> 
+    --save_model_as=safetensors 
+    --prior_loss_weight=1.0 
+    --max_train_steps=1600 
+    --learning_rate=1e-6 
+    --optimizer_type="AdamW8bit" 
+    --xformers 
+    --mixed_precision="fp16" 
+    --cache_latents 
+    --gradient_checkpointing
+    --token_string=mychar4 --init_word=cute --num_vectors_per_token=4
+```
+
+``--token_string`` に学習時のトークン文字列を指定します。__学習時のプロンプトは、この文字列を含むようにしてください（token_stringがmychar4なら、``mychar4 1girl`` など）__。プロンプトのこの文字列の部分が、Textual Inversionの新しいtokenに置換されて学習されます。DreamBooth, class+identifier形式のデータセットとして、`token_string` をトークン文字列にするのが最も簡単で確実です。
+
+プロンプトにトークン文字列が含まれているかどうかは、``--debug_dataset`` で置換後のtoken idが表示されますので、以下のように ``49408`` 以降のtokenが存在するかどうかで確認できます。
+
+```
+input ids: tensor([[49406, 49408, 49409, 49410, 49411, 49412, 49413, 49414, 49415, 49407,
+         49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,
+         49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,
+         49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,
+         49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,
+         49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,
+         49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,
+         49407, 49407, 49407, 49407, 49407, 49407, 49407]])
+```
+
+tokenizerがすでに持っている単語（一般的な単語）は使用できません。
+
+``--init_word`` にembeddingsを初期化するときのコピー元トークンの文字列を指定します。学ばせたい概念が近いものを選ぶとよいようです。二つ以上のトークンになる文字列は指定できません。
+
+``--num_vectors_per_token`` にいくつのトークンをこの学習で使うかを指定します。多いほうが表現力が増しますが、その分多くのトークンを消費します。たとえばnum_vectors_per_token=8の場合、指定したトークン文字列は（一般的なプロンプトの77トークン制限のうち）8トークンを消費します。
+
+以上がTextual Inversionのための主なオプションです。以降は他の学習スクリプトと同様です。
+
+`num_cpu_threads_per_process` には通常は1を指定するとよいようです。
+
+`pretrained_model_name_or_path` に追加学習を行う元となるモデルを指定します。Stable Diffusionのcheckpointファイル（.ckptまたは.safetensors）、Diffusersのローカルディスクにあるモデルディレクトリ、DiffusersのモデルID（"stabilityai/stable-diffusion-2"など）が指定できます。
+
+`output_dir` に学習後のモデルを保存するフォルダを指定します。`output_name` にモデルのファイル名を拡張子を除いて指定します。`save_model_as` でsafetensors形式での保存を指定しています。
+
+`dataset_config` に `.toml` ファイルを指定します。ファイル内でのバッチサイズ指定は、当初はメモリ消費を抑えるために `1` としてください。
+
+学習させるステップ数 `max_train_steps` を10000とします。学習率 `learning_rate` はここでは5e-6を指定しています。
+
+省メモリ化のため `mixed_precision="fp16"` を指定します（RTX30 シリーズ以降では `bf16` も指定できます。環境整備時にaccelerateに行った設定と合わせてください）。また `gradient_checkpointing` を指定します。
+
+オプティマイザ（モデルを学習データにあうように最適化＝学習させるクラス）にメモリ消費の少ない 8bit AdamW を使うため、 `optimizer_type="AdamW8bit"` を指定します。
+
+`xformers` オプションを指定し、xformersのCrossAttentionを用います。xformersをインストールしていない場合やエラーとなる場合（環境にもよりますが `mixed_precision="no"` の場合など）、代わりに `mem_eff_attn` オプションを指定すると省メモリ版CrossAttentionを使用します（速度は遅くなります）。
+
+ある程度メモリがある場合は、`.toml` ファイルを編集してバッチサイズをたとえば `8` くらいに増やしてください（高速化と精度向上の可能性があります）。
+
+### よく使われるオプションについて
+
+以下の場合にはオプションに関するドキュメントを参照してください。
+
+- Stable Diffusion 2.xまたはそこからの派生モデルを学習する
+- clip skipを2以上を前提としたモデルを学習する
+- 75トークンを超えたキャプションで学習する
+
+### Textual Inversionでのバッチサイズについて
+
+モデル全体を学習するDreamBoothやfine tuningに比べてメモリ使用量が少ないため、バッチサイズは大きめにできます。
+
+# Textual Inversionのその他の主なオプション
+
+すべてのオプションについては別文書を参照してください。
+
+* `--weights`
+  * 学習前に学習済みのembeddingsを読み込み、そこから追加で学習します。
+* `--use_object_template`
+  * キャプションではなく既定の物体用テンプレート文字列（``a photo of a {}``など）で学習します。公式実装と同じになります。キャプションは無視されます。
+* `--use_style_template`
+  * キャプションではなく既定のスタイル用テンプレート文字列で学習します（``a painting in the style of {}``など）。公式実装と同じになります。キャプションは無視されます。
+
+## 当リポジトリ内の画像生成スクリプトで生成する
+
+gen_img_diffusers.pyに、``--textual_inversion_embeddings`` オプションで学習したembeddingsファイルを指定してください（複数可）。プロンプトでembeddingsファイルのファイル名（拡張子を除く）を使うと、そのembeddingsが適用されます。
+

train_ti_README.md

@@ -0,0 +1,62 @@
+## About learning Textual Inversion
+
+[Textual Inversion](https://textual-inversion.github.io/). I heavily referenced https://github.com/huggingface/diffusers/tree/main/examples/textual_inversion for the implementation.
+
+The trained model can be used as is on the Web UI.
+
+In addition, it is probably compatible with SD2.x, but it has not been tested at this time.
+
+## Learning method
+
+Use ``train_textual_inversion.py``.
+
+Data preparation is exactly the same as ``train_network.py``, so please refer to [their document](./train_network_README-en.md).
+
+## options
+
+Below is an example command line (DreamBooth technique).
+
+```
+accelerate launch --num_cpu_threads_per_process 1 train_textual_inversion.py
+     --pretrained_model_name_or_path=..\models\model.ckpt
+     --train_data_dir=..\data\db\char1 --output_dir=..\ti_train1
+     --resolution=448,640 --train_batch_size=1 --learning_rate=1e-4
+     --max_train_steps=400 --use_8bit_adam --xformers --mixed_precision=fp16
+     --save_every_n_epochs=1 --save_model_as=safetensors --clip_skip=2 --seed=42 --color_aug
+     --token_string=mychar4 --init_word=cute --num_vectors_per_token=4
+```
+
+``--token_string`` specifies the token string for learning. __The learning prompt should contain this string (eg ``mychar4 1girl`` if token_string is mychar4)__. This string part of the prompt is replaced with a new token for Textual Inversion and learned.
+
+``--debug_dataset`` will display the token id after substitution, so you can check if the token string after ``49408`` exists as shown below. I can confirm.
+
+```
+input ids: tensor([[49406, 49408, 49409, 49410, 49411, 49412, 49413, 49414, 49415, 49407,
+          49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,
+          49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,
+          49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,
+          49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,
+          49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,
+          49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,
+          49407, 49407, 49407, 49407, 49407, 49407, 49407]])
+```
+
+Words that the tokenizer already has (common words) cannot be used.
+
+In ``--init_word``, specify the string of the copy source token when initializing embeddings. It seems to be a good idea to choose something that has a similar concept to what you want to learn. You cannot specify a character string that becomes two or more tokens.
+
+``--num_vectors_per_token`` specifies how many tokens to use for this training. The higher the number, the more expressive it is, but it consumes more tokens. For example, if num_vectors_per_token=8, then the specified token string will consume 8 tokens (out of the 77 token limit for a typical prompt).
+
+
+In addition, the following options can be specified.
+
+* --weights
+   * Load learned embeddings before learning and learn additionally from there.
+* --use_object_template
+   * Learn with default object template strings (such as ``a photo of a {}``) instead of captions. It will be the same as the official implementation. Captions are ignored.
+* --use_style_template
+   * Learn with default style template strings instead of captions (such as ``a painting in the style of {}``). It will be the same as the official implementation. Captions are ignored.
+
+## Generate with the image generation script in this repository
+
+In gen_img_diffusers.py, specify the learned embeddings file with the ``--textual_inversion_embeddings`` option. Using the filename (without the extension) of the embeddings file at the prompt will apply the embeddings.

upgrade.bat

@@ -0,0 +1,16 @@
+@echo off
+:: Check if there are any changes that need to be committed
+git status --short
+if %errorlevel%==1 (
+    echo There are changes that need to be committed. Please stash or undo your changes before running this script.
+    exit
+)
+
+:: Pull the latest changes from the remote repository
+git pull
+
+:: Activate the virtual environment
+call .\venv\Scripts\activate.bat
+
+:: Upgrade the required packages
+pip install --use-pep517 --upgrade -r requirements.txt

upgrade.ps1

@@ -0,0 +1,14 @@
+# Check if there are any changes that need to be committed
+if (git status --short) {
+    Write-Error "There are changes that need to be committed. Please stash or undo your changes before running this script."
+    return
+}
+
+# Pull the latest changes from the remote repository
+git pull
+
+# Activate the virtual environment
+.\venv\Scripts\activate
+
+# Upgrade the required packages
+pip install --use-pep517 --upgrade -r requirements.txt

upgrade.sh

@@ -0,0 +1,16 @@
+#!/bin/bash
+
+# Check if there are any changes that need to be committed
+if git status --short | grep -q "^[^ ?][^?]*"; then
+    echo "There are changes that need to be committed. Please stash or undo your changes before running this script."
+    exit 1
+fi
+
+# Pull the latest changes from the remote repository
+git pull
+
+# Activate the virtual environment
+source venv/bin/activate
+
+# Upgrade the required packages
+pip install --use-pep517 --upgrade -r requirements.txt

utilities.cmd

@@ -0,0 +1 @@
+.\venv\Scripts\python.exe library\utilities.py