Update README.md

Add LoRA weights checking script

.gitignore

@@ -1,3 +1,6 @@
 logs
 __pycache__
 wd14_tagger_model
+venv
+*.egg-info
+build

README-ja.md

@@ -14,10 +14,118 @@ GUIやPowerShellスクリプトなど、より使いやすくする機能が[bma
 
 ## 使用法について
 
-note.comに記事がありますのでそちらをご覧ください（将来的にはこちらへ移すかもしれません）。
+当リポジトリ内およびnote.comに記事がありますのでそちらをご覧ください（将来的にはすべてこちらへ移すかもしれません）。
+
+* note.com [環境整備とDreamBooth学習スクリプトについて](https://note.com/kohya_ss/n/nba4eceaa4594)
+* [fine-tuningのガイド](./fine_tune_README_ja.md):
+BLIPによるキャプショニングと、DeepDanbooruまたはWD14 taggerによるタグ付けを含みます
+* 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 --system-site-packages 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
+```
+
+コマンドプロンプトでは以下になります。
+
+
+```bat
+git clone https://github.com/kohya-ss/sd-scripts.git
+cd sd-scripts
+
+python -m venv --system-site-packages 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
+```
+
+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が使われます。）
+
+## アップグレード
+
+新しいリリースがあった場合、以下のコマンドで更新できます。
+
+```powershell
+cd sd-scripts
+git pull
+.\venv\Scripts\activate
+pip install --upgrade -r <requirement file name>
+```
+
+コマンドが成功すれば新しいバージョンが使用できます。
+
+## 謝意
+
+LoRAの実装は[cloneofsimo氏のリポジトリ](https://github.com/cloneofsimo/lora)を基にしたものです。感謝申し上げます。
+
+## ライセンス
+
+スクリプトのライセンスは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
+
 
-* [環境整備とDreamBooth学習スクリプトについて](https://note.com/kohya_ss/n/nee3ed1649fb6)
-* [fine-tuningスクリプト](https://note.com/kohya_ss/n/nbf7ce8d80f29):
-Including BLIP captioning and tagging by DeepDanbooru or WD14 tagger
-* [画像生成スクリプト](https://note.com/kohya_ss/n/n2693183a798e)
-* [モデル変換スクリプト](https://note.com/kohya_ss/n/n374f316fe4ad)

README.md

@@ -1,5 +1,18 @@
 This repository contains training, generation and utility scripts for Stable Diffusion.
 
+## Updates
+
+- 22 Jan. 2023, 2023/1/22
+  - Fix script to check LoRA weights ``check_lora_weights.py``. Some layer weights were shown as ``0.0`` even if the layer is trained, because of the overflow of ``torch.mean``. Sorry for the confusion.
+  - Noe the script shows the mean of the absolute values of the weights, and the minimum of the absolute values of the weights.
+  - LoRAの重みをチェックするスクリプト ``check_lora_weights.py`` を修正しました。一部のレイヤーで学習されているにもかかわらず重みが ``0.0`` と表示されていました。混乱を招き申し訳ありません。
+  - スクリプトを「重みの絶対の平均」と「重みの絶対値の最小値」を表示するよう修正しました。
+
+Please read [Releases](https://github.com/kohya-ss/sd-scripts/releases) for recent updates.
+最近の更新情報は [Release](https://github.com/kohya-ss/sd-scripts/releases) をご覧ください。
+
+##
+
 [日本語版README](./README-ja.md)
 
 For easier use (GUI and PowerShell scripts etc...), please visit [the repository maintained by bmaltais](https://github.com/bmaltais/kohya_ss). Thanks to @bmaltais!
@@ -8,21 +21,103 @@ This repository contains the scripts for:
 
 * DreamBooth training, including U-Net and Text Encoder
 * fine-tuning (native training), including U-Net and Text Encoder
+* LoRA training
 * image generation
 * model conversion (supports 1.x and 2.x, Stable Diffision ckpt/safetensors and Diffusers)
 
-## About requirements_*.txt
+## About requirements.txt
 
-These files do not contain requirements for PyTorch and Diffusers. Because the versions of them depend on your environment. Please install PyTorch at first, then Diffusers. 
+These files do not contain requirements for PyTorch. Because the versions of them depend on your environment. Please install PyTorch at first (see installation guide below.) 
 
-The scripts is tested with PyTorch 1.12.1 and 1.13.0, Diffusers 0.10.2.
+The scripts are tested with PyTorch 1.12.1 and 1.13.0, Diffusers 0.10.2.
 
 ## Links to how-to-use documents
 
 All documents are in Japanese currently, and CUI based.
 
-* [Environment setup and DreamBooth training guide](https://note.com/kohya_ss/n/nee3ed1649fb6)
-* [Fine-tuning step-by-step guide](https://note.com/kohya_ss/n/nbf7ce8d80f29):
+* [DreamBooth training guide](./train_db_README-ja.md)
+* [Step by Step fine-tuning guide](./fine_tune_README_ja.md):
 Including BLIP captioning and tagging by DeepDanbooru or WD14 tagger
-* [Image generation](https://note.com/kohya_ss/n/n2693183a798e)
-* [Model conversion](https://note.com/kohya_ss/n/n374f316fe4ad)
+* [training LoRA](./train_network_README-ja.md)
+* note.com [Image generation](https://note.com/kohya_ss/n/n2693183a798e)
+* note.com [Model conversion](https://note.com/kohya_ss/n/n374f316fe4ad)
+
+## Windows Required Dependencies
+
+Python 3.10.6 and 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
+
+Give unrestricted script access to powershell so venv can work:
+
+- Open an administrator powershell window
+- Type `Set-ExecutionPolicy Unrestricted` and answer A
+- Close admin powershell window
+
+## Windows Installation
+
+Open a regular Powershell terminal and type the following inside:
+
+```powershell
+git clone https://github.com/kohya-ss/sd-scripts.git
+cd sd-scripts
+
+python -m venv --system-site-packages 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
+
+```
+
+Answers to accelerate config:
+
+```txt
+- This machine
+- No distributed training
+- NO
+- NO
+- NO
+- all
+- fp16
+```
+
+note: Some user reports ``ValueError: fp16 mixed precision requires a GPU`` is occurred in training. In this case, answer `0` for the 6th question: 
+``What GPU(s) (by id) should be used for training on this machine as a comma-separated list? [all]:`` 
+
+(Single GPU with id `0` will be used.)
+
+## Upgrade
+
+When a new release comes out you can upgrade your repo with the following command:
+
+```powershell
+cd sd-scripts
+git pull
+.\venv\Scripts\activate
+pip install --upgrade -r requirements.txt
+```
+
+Once the commands have completed successfully you should be ready to use the new version.
+
+## Credits
+
+The implementation for LoRA is based on [cloneofsimo's repo](https://github.com/cloneofsimo/lora). Thank you for great work!!!
+
+## License
+
+The majority of scripts is licensed under ASL 2.0 (including codes from Diffusers, cloneofsimo's), however portions of the project are available under separate license terms:
+
+[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

bitsandbytes_windows/cextension.py

@@ -0,0 +1,54 @@
+import ctypes as ct
+from pathlib import Path
+from warnings import warn
+
+from .cuda_setup.main import evaluate_cuda_setup
+
+
+class CUDALibrary_Singleton(object):
+    _instance = None
+
+    def __init__(self):
+        raise RuntimeError("Call get_instance() instead")
+
+    def initialize(self):
+        binary_name = evaluate_cuda_setup()
+        package_dir = Path(__file__).parent
+        binary_path = package_dir / binary_name
+
+        if not binary_path.exists():
+            print(f"CUDA SETUP: TODO: compile library for specific version: {binary_name}")
+            legacy_binary_name = "libbitsandbytes.so"
+            print(f"CUDA SETUP: Defaulting to {legacy_binary_name}...")
+            binary_path = package_dir / legacy_binary_name
+            if not binary_path.exists():
+                print('CUDA SETUP: CUDA detection failed. Either CUDA driver not installed, CUDA not installed, or you have multiple conflicting CUDA libraries!')
+                print('CUDA SETUP: If you compiled from source, try again with `make CUDA_VERSION=DETECTED_CUDA_VERSION` for example, `make CUDA_VERSION=113`.')
+                raise Exception('CUDA SETUP: Setup Failed!')
+            # self.lib = ct.cdll.LoadLibrary(binary_path)
+            self.lib = ct.cdll.LoadLibrary(str(binary_path))            # $$$
+        else:
+            print(f"CUDA SETUP: Loading binary {binary_path}...")
+            # self.lib = ct.cdll.LoadLibrary(binary_path)
+            self.lib = ct.cdll.LoadLibrary(str(binary_path))            # $$$
+
+    @classmethod
+    def get_instance(cls):
+        if cls._instance is None:
+            cls._instance = cls.__new__(cls)
+            cls._instance.initialize()
+        return cls._instance
+
+
+lib = CUDALibrary_Singleton.get_instance().lib
+try:
+    lib.cadam32bit_g32
+    lib.get_context.restype = ct.c_void_p
+    lib.get_cusparse.restype = ct.c_void_p
+    COMPILED_WITH_CUDA = True
+except AttributeError:
+    warn(
+        "The installed version of bitsandbytes was compiled without GPU support. "
+        "8-bit optimizers and GPU quantization are unavailable."
+    )
+    COMPILED_WITH_CUDA = False

bitsandbytes_windows/main.py

@@ -0,0 +1,166 @@
+"""
+extract factors the build is dependent on:
+[X] compute capability
+    [ ] TODO: Q - What if we have multiple GPUs of different makes?
+- CUDA version
+- Software:
+    - CPU-only: only CPU quantization functions (no optimizer, no matrix multiple)
+    - CuBLAS-LT: full-build 8-bit optimizer
+    - no CuBLAS-LT: no 8-bit matrix multiplication (`nomatmul`)
+
+evaluation:
+    - if paths faulty, return meaningful error
+    - else:
+        - determine CUDA version
+        - determine capabilities
+        - based on that set the default path
+"""
+
+import ctypes
+
+from .paths import determine_cuda_runtime_lib_path
+
+
+def check_cuda_result(cuda, result_val):
+    # 3. Check for CUDA errors
+    if result_val != 0:
+        error_str = ctypes.c_char_p()
+        cuda.cuGetErrorString(result_val, ctypes.byref(error_str))
+        print(f"CUDA exception! Error code: {error_str.value.decode()}")
+
+def get_cuda_version(cuda, cudart_path):
+    # https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART____VERSION.html#group__CUDART____VERSION
+    try:
+        cudart = ctypes.CDLL(cudart_path)
+    except OSError:
+        # TODO: shouldn't we error or at least warn here?
+        print(f'ERROR: libcudart.so could not be read from path: {cudart_path}!')
+        return None
+
+    version = ctypes.c_int()
+    check_cuda_result(cuda, cudart.cudaRuntimeGetVersion(ctypes.byref(version)))
+    version = int(version.value)
+    major = version//1000
+    minor = (version-(major*1000))//10
+
+    if major < 11:
+       print('CUDA SETUP: CUDA version lower than 11 are currently not supported for LLM.int8(). You will be only to use 8-bit optimizers and quantization routines!!')
+
+    return f'{major}{minor}'
+
+
+def get_cuda_lib_handle():
+    # 1. find libcuda.so library (GPU driver) (/usr/lib)
+    try:
+        cuda = ctypes.CDLL("libcuda.so")
+    except OSError:
+        # TODO: shouldn't we error or at least warn here?
+        print('CUDA SETUP: WARNING! libcuda.so not found! Do you have a CUDA driver installed? If you are on a cluster, make sure you are on a CUDA machine!')
+        return None
+    check_cuda_result(cuda, cuda.cuInit(0))
+
+    return cuda
+
+
+def get_compute_capabilities(cuda):
+    """
+    1. find libcuda.so library (GPU driver) (/usr/lib)
+       init_device -> init variables -> call function by reference
+    2. call extern C function to determine CC
+       (https://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__DEVICE__DEPRECATED.html)
+    3. Check for CUDA errors
+       https://stackoverflow.com/questions/14038589/what-is-the-canonical-way-to-check-for-errors-using-the-cuda-runtime-api
+    # bits taken from https://gist.github.com/f0k/63a664160d016a491b2cbea15913d549
+    """
+
+
+    nGpus = ctypes.c_int()
+    cc_major = ctypes.c_int()
+    cc_minor = ctypes.c_int()
+
+    device = ctypes.c_int()
+
+    check_cuda_result(cuda, cuda.cuDeviceGetCount(ctypes.byref(nGpus)))
+    ccs = []
+    for i in range(nGpus.value):
+        check_cuda_result(cuda, cuda.cuDeviceGet(ctypes.byref(device), i))
+        ref_major = ctypes.byref(cc_major)
+        ref_minor = ctypes.byref(cc_minor)
+        # 2. call extern C function to determine CC
+        check_cuda_result(
+            cuda, cuda.cuDeviceComputeCapability(ref_major, ref_minor, device)
+        )
+        ccs.append(f"{cc_major.value}.{cc_minor.value}")
+
+    return ccs
+
+
+# def get_compute_capability()-> Union[List[str, ...], None]: # FIXME: error
+def get_compute_capability(cuda):
+    """
+    Extracts the highest compute capbility from all available GPUs, as compute
+    capabilities are downwards compatible. If no GPUs are detected, it returns
+    None.
+    """
+    ccs = get_compute_capabilities(cuda)
+    if ccs is not None:
+        # TODO: handle different compute capabilities; for now, take the max
+        return ccs[-1]
+    return None
+
+
+def evaluate_cuda_setup():
+    print('')
+    print('='*35 + 'BUG REPORT' + '='*35)
+    print('Welcome to bitsandbytes. For bug reports, please submit your error trace to: https://github.com/TimDettmers/bitsandbytes/issues')
+    print('For effortless bug reporting copy-paste your error into this form: https://docs.google.com/forms/d/e/1FAIpQLScPB8emS3Thkp66nvqwmjTEgxp8Y9ufuWTzFyr9kJ5AoI47dQ/viewform?usp=sf_link')
+    print('='*80)
+    return "libbitsandbytes_cuda116.dll"            # $$$
+    
+    binary_name = "libbitsandbytes_cpu.so"
+    #if not torch.cuda.is_available():
+        #print('No GPU detected. Loading CPU library...')
+        #return binary_name
+
+    cudart_path = determine_cuda_runtime_lib_path()
+    if cudart_path is None:
+        print(
+            "WARNING: No libcudart.so found! Install CUDA or the cudatoolkit package (anaconda)!"
+        )
+        return binary_name
+
+    print(f"CUDA SETUP: CUDA runtime path found: {cudart_path}")
+    cuda = get_cuda_lib_handle()
+    cc = get_compute_capability(cuda)
+    print(f"CUDA SETUP: Highest compute capability among GPUs detected: {cc}")
+    cuda_version_string = get_cuda_version(cuda, cudart_path)
+
+
+    if cc == '':
+        print(
+            "WARNING: No GPU detected! Check your CUDA paths. Processing to load CPU-only library..."
+        )
+        return binary_name
+
+    # 7.5 is the minimum CC vor cublaslt
+    has_cublaslt = cc in ["7.5", "8.0", "8.6"]
+
+    # TODO:
+    # (1) CUDA missing cases (no CUDA installed by CUDA driver (nvidia-smi accessible)
+    # (2) Multiple CUDA versions installed
+
+    # we use ls -l instead of nvcc to determine the cuda version
+    # since most installations will have the libcudart.so installed, but not the compiler
+    print(f'CUDA SETUP: Detected CUDA version {cuda_version_string}')
+
+    def get_binary_name():
+        "if not has_cublaslt (CC < 7.5), then we have to choose  _nocublaslt.so"
+        bin_base_name = "libbitsandbytes_cuda"
+        if has_cublaslt:
+            return f"{bin_base_name}{cuda_version_string}.so"
+        else:
+            return f"{bin_base_name}{cuda_version_string}_nocublaslt.so"
+
+    binary_name = get_binary_name()
+
+    return binary_name

fine_tune_README_ja.md

@@ -0,0 +1,465 @@
+NovelAIの提案した学習手法、自動キャプションニング、タグ付け、Windows＋VRAM 12GB（v1.4/1.5の場合）環境等に対応したfine tuningです。
+
+## 概要
+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](./README-ja.md)を参照してください。
+
+## 教師データの用意
+
+学習させたい画像データを用意し、任意のフォルダに入れてください。リサイズ等の事前の準備は必要ありません。
+ただし学習解像度よりもサイズが小さい画像については、超解像などで品質を保ったまま拡大しておくことをお勧めします。
+
+複数の教師データフォルダにも対応しています。前処理をそれぞれのフォルダに対して実行する形となります。
+
+たとえば以下のように画像を格納します。
+
+![教師データフォルダのスクショ](https://user-images.githubusercontent.com/52813779/208907739-8e89d5fa-6ca8-4b60-8927-f484d2a9ae04.png)
+
+## 自動キャプショニング
+キャプションを使わずタグだけで学習する場合はスキップしてください。
+
+また手動でキャプションを用意する場合、キャプションは教師データ画像と同じディレクトリに、同じファイル名、拡張子.caption等で用意してください。各ファイルは1行のみのテキストファイルとします。
+
+### BLIPによるキャプショニング
+
+最新版ではBLIPのダウンロード、重みのダウンロード、仮想環境の追加は不要になりました。そのままで動作します。
+
+finetuneフォルダ内のmake_captions.pyを実行します。
+
+```
+python finetune\make_captions.py --batch_size <バッチサイズ> <教師データフォルダ>
+```
+
+バッチサイズ8、教師データを親フォルダのtrain_dataに置いた場合、以下のようになります。
+
+```
+python finetune\make_captions.py --batch_size 8 ..\train_data
+```
+
+キャプションファイルが教師データ画像と同じディレクトリに、同じファイル名、拡張子.captionで作成されます。
+
+batch_sizeはGPUのVRAM容量に応じて増減してください。大きいほうが速くなります（VRAM 12GBでももう少し増やせると思います）。
+max_lengthオプションでキャプションの最大長を指定できます。デフォルトは75です。モデルをトークン長225で学習する場合には長くしても良いかもしれません。
+caption_extensionオプションでキャプションの拡張子を変更できます。デフォルトは.captionです（.txtにすると後述のDeepDanbooruと競合します）。
+
+複数の教師データフォルダがある場合には、それぞれのフォルダに対して実行してください。
+
+なお、推論にランダム性があるため、実行するたびに結果が変わります。固定する場合には--seedオプションで「--seed 42」のように乱数seedを指定してください。
+
+その他のオプションは--helpでヘルプをご参照ください（パラメータの意味についてはドキュメントがまとまっていないようで、ソースを見るしかないようです）。
+
+デフォルトでは拡張子.captionでキャプションファイルが生成されます。
+
+![captionが生成されたフォルダ](https://user-images.githubusercontent.com/52813779/208908845-48a9d36c-f6ee-4dae-af71-9ab462d1459e.png)
+
+たとえば以下のようなキャプションが付きます。
+
+![キャプションと画像](https://user-images.githubusercontent.com/52813779/208908947-af936957-5d73-4339-b6c8-945a52857373.png)
+
+## DeepDanbooruによるタグ付け
+danbooruタグのタグ付け自体を行わない場合は「キャプションとタグ情報の前処理」に進んでください。
+
+タグ付けはDeepDanbooruまたはWD14Taggerで行います。WD14Taggerのほうが精度が良いようです。WD14Taggerでタグ付けする場合は、次の章へ進んでください。
+
+### 環境整備
+DeepDanbooru https://github.com/KichangKim/DeepDanbooru  を作業フォルダにcloneしてくるか、zipをダウンロードして展開します。私はzipで展開しました。
+またDeepDanbooruのReleasesのページ https://github.com/KichangKim/DeepDanbooru/releases  の「DeepDanbooru Pretrained Model v3-20211112-sgd-e28」のAssetsから、deepdanbooru-v3-20211112-sgd-e28.zipをダウンロードしてきてDeepDanbooruのフォルダに展開します。
+
+以下からダウンロードします。Assetsをクリックして開き、そこからダウンロードします。
+
+![DeepDanbooruダウンロードページ](https://user-images.githubusercontent.com/52813779/208909417-10e597df-7085-41ee-bd06-3e856a1339df.png)
+
+以下のようなこういうディレクトリ構造にしてください
+
+![DeepDanbooruのディレクトリ構造](https://user-images.githubusercontent.com/52813779/208909486-38935d8b-8dc6-43f1-84d3-fef99bc471aa.png)
+
+Diffusersの環境に必要なライブラリをインストールします。DeepDanbooruのフォルダに移動してインストールします（実質的にはtensorflow-ioが追加されるだけだと思います）。
+
+```
+pip install -r requirements.txt
+```
+
+続いてDeepDanbooru自体をインストールします。
+
+```
+pip install .
+```
+
+以上でタグ付けの環境整備は完了です。
+
+### タグ付けの実施
+DeepDanbooruのフォルダに移動し、deepdanbooruを実行してタグ付けを行います。
+
+```
+deepdanbooru evaluate <教師データフォルダ> --project-path deepdanbooru-v3-20211112-sgd-e28 --allow-folder --save-txt
+```
+
+教師データを親フォルダのtrain_dataに置いた場合、以下のようになります。
+
+```
+deepdanbooru evaluate ../train_data --project-path deepdanbooru-v3-20211112-sgd-e28 --allow-folder --save-txt
+```
+
+タグファイルが教師データ画像と同じディレクトリに、同じファイル名、拡張子.txtで作成されます。1件ずつ処理されるためわりと遅いです。
+
+複数の教師データフォルダがある場合には、それぞれのフォルダに対して実行してください。
+
+以下のように生成されます。
+
+![DeepDanbooruの生成ファイル](https://user-images.githubusercontent.com/52813779/208909855-d21b9c98-f2d3-4283-8238-5b0e5aad6691.png)
+
+こんな感じにタグが付きます（すごい情報量……）。
+
+![DeepDanbooruタグと画像](https://user-images.githubusercontent.com/52813779/208909908-a7920174-266e-48d5-aaef-940aba709519.png)
+
+## WD14Taggerによるタグ付け
+DeepDanbooruの代わりにWD14Taggerを用いる手順です。
+
+Automatic1111氏のWebUIで使用しているtaggerを利用します。こちらのgithubページ（https://github.com/toriato/stable-diffusion-webui-wd14-tagger#mrsmilingwolfs-model-aka-waifu-diffusion-14-tagger ）の情報を参考にさせていただきました。
+
+最初の環境整備で必要なモジュールはインストール済みです。また重みはHugging Faceから自動的にダウンロードしてきます。
+
+### タグ付けの実施
+スクリプトを実行してタグ付けを行います。
+```
+python tag_images_by_wd14_tagger.py --batch_size <バッチサイズ> <教師データフォルダ>
+```
+
+教師データを親フォルダのtrain_dataに置いた場合、以下のようになります。
+```
+python tag_images_by_wd14_tagger.py --batch_size 4 ..\train_data
+```
+
+初回起動時にはモデルファイルがwd14_tagger_modelフォルダに自動的にダウンロードされます（フォルダはオプションで変えられます）。以下のようになります。
+
+![ダウンロードされたファイル](https://user-images.githubusercontent.com/52813779/208910447-f7eb0582-90d6-49d3-a666-2b508c7d1842.png)
+
+タグファイルが教師データ画像と同じディレクトリに、同じファイル名、拡張子.txtで作成されます。
+
+![生成されたタグファイル](https://user-images.githubusercontent.com/52813779/208910534-ea514373-1185-4b7d-9ae3-61eb50bc294e.png)
+
+![タグと画像](https://user-images.githubusercontent.com/52813779/208910599-29070c15-7639-474f-b3e4-06bd5a3df29e.png)
+
+threshオプションで、判定されたタグのconfidence（確信度）がいくつ以上でタグをつけるかが指定できます。デフォルトはWD14Taggerのサンプルと同じ0.35です。値を下げるとより多くのタグが付与されますが、精度は下がります。
+batch_sizeはGPUのVRAM容量に応じて増減してください。大きいほうが速くなります（VRAM 12GBでももう少し増やせると思います）。caption_extensionオプションでタグファイルの拡張子を変更できます。デフォルトは.txtです。
+model_dirオプションでモデルの保存先フォルダを指定できます。
+またforce_downloadオプションを指定すると保存先フォルダがあってもモデルを再ダウンロードします。
+
+複数の教師データフォルダがある場合には、それぞれのフォルダに対して実行してください。
+
+## キャプションとタグ情報の前処理
+
+スクリプトから処理しやすいようにキャプションとタグをメタデータとしてひとつのファイルにまとめます。
+
+### キャプションの前処理
+
+キャプションをメタデータに入れるには、作業フォルダ内で以下を実行してください（キャプションを学習に使わない場合は実行不要です）（実際は1行で記述します、以下同様）。
+
+```
+python merge_captions_to_metadata.py <教師データフォルダ>
+　  --in_json <読み込むメタデータファイル名> 
+    <メタデータファイル名>
+```
+
+メタデータファイル名は任意の名前です。
+教師データがtrain_data、読み込むメタデータファイルなし、メタデータファイルがmeta_cap.jsonの場合、以下のようになります。
+
+```
+python merge_captions_to_metadata.py train_data meta_cap.json
+```
+
+caption_extensionオプションでキャプションの拡張子を指定できます。
+
+複数の教師データフォルダがある場合には、full_path引数を指定してください（メタデータにフルパスで情報を持つようになります）。そして、それぞれのフォルダに対して実行してください。
+
+```
+python merge_captions_to_metadata.py --full_path 
+    train_data1 meta_cap1.json
+python merge_captions_to_metadata.py --full_path --in_json meta_cap1.json 
+    train_data2 meta_cap2.json
+```
+
+in_jsonを省略すると書き込み先メタデータファイルがあるとそこから読み込み、そこに上書きします。
+
+__※in_jsonオプションと書き込み先を都度書き換えて、別のメタデータファイルへ書き出すようにすると安全です。__
+
+### タグの前処理
+
+同様にタグもメタデータにまとめます（タグを学習に使わない場合は実行不要です）。
+```
+python merge_dd_tags_to_metadata.py <教師データフォルダ> 
+    --in_json <読み込むメタデータファイル名>
+    <書き込むメタデータファイル名>
+```
+
+先と同じディレクトリ構成で、meta_cap.jsonを読み、meta_cap_dd.jsonに書きだす場合、以下となります。
+```
+python merge_dd_tags_to_metadata.py train_data --in_json meta_cap.json meta_cap_dd.json
+```
+
+複数の教師データフォルダがある場合には、full_path引数を指定してください。そして、それぞれのフォルダに対して実行してください。
+
+```
+python merge_dd_tags_to_metadata.py --full_path --in_json meta_cap2.json
+    train_data1 meta_cap_dd1.json
+python merge_dd_tags_to_metadata.py --full_path --in_json meta_cap_dd1.json 
+    train_data2 meta_cap_dd2.json
+```
+
+in_jsonを省略すると書き込み先メタデータファイルがあるとそこから読み込み、そこに上書きします。
+
+__※in_jsonオプションと書き込み先を都度書き換えて、別のメタデータファイルへ書き出すようにすると安全です。__
+
+### キャプションとタグのクリーニング
+ここまででメタデータファイルにキャプションとDeepDanbooruのタグがまとめられています。ただ自動キャプショニングにしたキャプションは表記ゆれなどがあり微妙（※）ですし、タグにはアンダースコアが含まれていたりratingが付いていたりしますので（DeepDanbooruの場合）、エディタの置換機能などを用いてキャプションとタグのクリーニングをしたほうがいいでしょう。
+
+※たとえばアニメ絵の少女を学習する場合、キャプションにはgirl/girls/woman/womenなどのばらつきがあります。また「anime girl」なども単に「girl」としたほうが適切かもしれません。
+
+クリーニング用のスクリプトが用意してありますので、スクリプトの内容を状況に応じて編集してお使いください。
+
+（教師データフォルダの指定は不要になりました。メタデータ内の全データをクリーニングします。）
+
+```
+python clean_captions_and_tags.py <読み込むメタデータファイル名> <書き込むメタデータファイル名>
+```
+
+--in_jsonは付きませんのでご注意ください。たとえば次のようになります。
+
+```
+python clean_captions_and_tags.py meta_cap_dd.json meta_clean.json
+```
+
+以上でキャプションとタグの前処理は完了です。
+
+## latentsの事前取得
+
+学習を高速に進めるためあらかじめ画像の潜在表現を取得しディスクに保存しておきます。あわせてbucketing（教師データをアスペクト比に応じて分類する）を行います。
+
+作業フォルダで以下のように入力してください。
+```
+python prepare_buckets_latents.py <教師データフォルダ>  
+    <読み込むメタデータファイル名> <書き込むメタデータファイル名> 
+    <fine tuningするモデル名またはcheckpoint> 
+    --batch_size <バッチサイズ> 
+    --max_resolution <解像度 幅,高さ> 
+    --mixed_precision <精度>
+```
+
+モデルがmodel.ckpt、バッチサイズ4、学習解像度は512\*512、精度no（float32）で、meta_clean.jsonからメタデータを読み込み、meta_lat.jsonに書き込む場合、以下のようになります。
+
+```
+python prepare_buckets_latents.py 
+    train_data meta_clean.json meta_lat.json model.ckpt 
+    --batch_size 4 --max_resolution 512,512 --mixed_precision no
+```
+
+教師データフォルダにnumpyのnpz形式でlatentsが保存されます。
+
+Stable Diffusion 2.0のモデルを読み込む場合は--v2オプションを指定してください（--v_parameterizationは不要です）。
+
+解像度の最小サイズを--min_bucket_resoオプションで、最大サイズを--max_bucket_resoで指定できます。デフォルトはそれぞれ256、1024です。たとえば最小サイズに384を指定すると、256\*1024や320\*768などの解像度は使わなくなります。
+解像度を768\*768のように大きくした場合、最大サイズに1280などを指定すると良いでしょう。
+
+--flip_augオプションを指定すると左右反転のaugmentation（データ拡張）を行います。疑似的にデータ量を二倍に増やすことができますが、データが左右対称でない場合に指定すると（例えばキャラクタの外見、髪型など）学習がうまく行かなくなります。
+（反転した画像についてもlatentsを取得し、\*\_flip.npzファイルを保存する単純な実装です。fline_tune.pyには特にオプション指定は必要ありません。\_flip付きのファイルがある場合、flip付き・なしのファイルを、ランダムに読み込みます。）
+
+バッチサイズはVRAM 12GBでももう少し増やせるかもしれません。
+解像度は64で割り切れる数字で、"幅,高さ"で指定します。解像度はfine tuning時のメモリサイズに直結します。VRAM 12GBでは512,512が限界と思われます（※）。16GBなら512,704や512,768まで上げられるかもしれません。なお256,256等にしてもVRAM 8GBでは厳しいようです（パラメータやoptimizerなどは解像度に関係せず一定のメモリが必要なため）。
+
+※batch size 1の学習で12GB VRAM、640,640で動いたとの報告もありました。
+
+以下のようにbucketingの結果が表示されます。
+
+![bucketingの結果](https://user-images.githubusercontent.com/52813779/208911419-71c00fbb-2ce6-49d5-89b5-b78d7715e441.png)
+
+複数の教師データフォルダがある場合には、full_path引数を指定してください。そして、それぞれのフォルダに対して実行してください。
+```
+python prepare_buckets_latents.py --full_path  
+    train_data1 meta_clean.json meta_lat1.json model.ckpt 
+    --batch_size 4 --max_resolution 512,512 --mixed_precision no
+
+python prepare_buckets_latents.py --full_path 
+    train_data2 meta_lat1.json meta_lat2.json model.ckpt 
+    --batch_size 4 --max_resolution 512,512 --mixed_precision no
+
+```
+読み込み元と書き込み先を同じにすることも可能ですが別々の方が安全です。
+
+__※引数を都度書き換えて、別のメタデータファイルに書き込むと安全です。__
+
+
+## 学習の実行
+たとえば以下のように実行します。以下は省メモリ化のための設定です。
+```
+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
+```
+
+accelerateのnum_cpu_threads_per_processにはCPUのコア数を指定するとよいようです。
+
+pretrained_model_name_or_pathに学習対象のモデルを指定します（Stable DiffusionのcheckpointかDiffusersのモデル）。Stable Diffusionのcheckpointは.ckptと.safetensorsに対応しています（拡張子で自動判定）。
+
+in_jsonにlatentをキャッシュしたときのメタデータファイルを指定します。
+
+train_data_dirに教師データのフォルダを、output_dirに学習後のモデルの出力先フォルダを指定します。
+
+shuffle_captionを指定すると、キャプション、タグをカンマ区切りされた単位でシャッフルして学習します（Waifu Diffusion v1.3で行っている手法です）。
+（先頭のトークンのいくつかをシャッフルせずに固定できます。その他のオプションのkeep_tokensをご覧ください。）
+
+train_batch_sizeにバッチサイズを指定します。VRAM 12GBでは1か2程度を指定してください。解像度によっても指定可能な数は変わってきます。
+学習に使用される実際のデータ量は「バッチサイズ×ステップ数」です。バッチサイズを増やした時には、それに応じてステップ数を下げることが可能です。
+
+learning_rateに学習率を指定します。たとえばWaifu Diffusion v1.3は5e-6のようです。
+max_train_stepsにステップ数を指定します。
+
+use_8bit_adamを指定すると8-bit Adam Optimizerを使用します。省メモリ化、高速化されますが精度は下がる可能性があります。
+
+xformersを指定するとCrossAttentionを置換して省メモリ化、高速化します。
+※11/9時点ではfloat32の学習ではxformersがエラーになるため、bf16/fp16を使うか、代わりにmem_eff_attnを指定して省メモリ版CrossAttentionを使ってください（速度はxformersに劣ります）。
+
+gradient_checkpointingで勾配の途中保存を有効にします。速度は遅くなりますが使用メモリ量が減ります。
+
+mixed_precisionで混合精度を使うか否かを指定します。"fp16"または"bf16"を指定すると省メモリになりますが精度は劣ります。
+"fp16"と"bf16"は使用メモリ量はほぼ同じで、bf16の方が学習結果は良くなるとの話もあります（試した範囲ではあまり違いは感じられませんでした）。
+"no"を指定すると使用しません（float32になります）。
+
+※bf16で学習したcheckpointをAUTOMATIC1111氏のWeb UIで読み込むとエラーになるようです。これはデータ型のbfloat16がWeb UIのモデルsafety checkerでエラーとなるためのようです。save_precisionオプションを指定してfp16またはfloat32形式で保存してください。またはsafetensors形式で保管しても良さそうです。
+
+save_every_n_epochsを指定するとそのエポックだけ経過するたびに学習中のモデルを保存します。
+
+### Stable Diffusion 2.0対応
+Hugging Faceのstable-diffusion-2-baseを使う場合は--v2オプションを、stable-diffusion-2または768-v-ema.ckptを使う場合は--v2と--v_parameterizationの両方のオプションを指定してください。
+
+### メモリに余裕がある場合に精度や速度を上げる
+まずgradient_checkpointingを外すと速度が上がります。ただし設定できるバッチサイズが減りますので、精度と速度のバランスを見ながら設定してください。
+
+バッチサイズを増やすと速度、精度が上がります。メモリが足りる範囲で、1データ当たりの速度を確認しながら増やしてください（メモリがぎりぎりになるとかえって速度が落ちることがあります）。
+
+### 使用するCLIP出力の変更
+clip_skipオプションに2を指定すると、後ろから二番目の層の出力を用います。1またはオプション省略時は最後の層を用います。
+学習したモデルはAutomatic1111氏のWeb UIで推論できるはずです。
+
+※SD2.0はデフォルトで後ろから二番目の層を使うため、SD2.0の学習では指定しないでください。
+
+学習対象のモデルがもともと二番目の層を使うように学習されている場合は、2を指定するとよいでしょう。
+
+そうではなく最後の層を使用していた場合はモデル全体がそれを前提に学習されています。そのため改めて二番目の層を使用して学習すると、望ましい学習結果を得るにはある程度の枚数の教師データ、長めの学習が必要になるかもしれません。
+
+### トークン長の拡張
+max_token_lengthに150または225を指定することでトークン長を拡張して学習できます。
+学習したモデルはAutomatic1111氏のWeb UIで推論できるはずです。
+
+clip_skipと同様に、モデルの学習状態と異なる長さで学習するには、ある程度の教師データ枚数、長めの学習時間が必要になると思われます。
+
+### 学習ログの保存
+logging_dirオプションにログ保存先フォルダを指定してください。TensorBoard形式のログが保存されます。
+
+たとえば--logging_dir=logsと指定すると、作業フォルダにlogsフォルダが作成され、その中の日時フォルダにログが保存されます。
+また--log_prefixオプションを指定すると、日時の前に指定した文字列が追加されます。「--logging_dir=logs --log_prefix=fine_tune_style1」などとして識別用にお使いください。
+
+TensorBoardでログを確認するには、別のコマンドプロンプトを開き、作業フォルダで以下のように入力します（tensorboardはDiffusersのインストール時にあわせてインストールされると思いますが、もし入っていないならpip install tensorboardで入れてください）。
+```
+tensorboard --logdir=logs
+```
+
+### Hypernetworkの学習
+別の記事で解説予定です。
+
+### 勾配をfp16とした学習（実験的機能）
+full_fp16オプションを指定すると勾配を通常のfloat32からfloat16（fp16）に変更して学習します（mixed precisionではなく完全なfp16学習になるようです）。これによりSD1.xの512*512サイズでは8GB未満、SD2.xの512*512サイズで12GB未満のVRAM使用量で学習できるようです。
+
+あらかじめ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で確認）。精度はかなり落ちますし、途中で学習失敗する確率も高くなります。学習率やステップ数の設定もシビアなようです。それらを認識したうえで自己責任でお使いください。
+
+### その他のオプション
+
+#### keep_tokens
+数値を指定するとキャプションの先頭から、指定した数だけのトークン（カンマ区切りの文字列）をシャッフルせず固定します。
+
+キャプションとタグが両方ある場合、学習時のプロンプトは「キャプション,タグ1,タグ2……」のように連結されますので、「--keep_tokens=1」とすれば、学習時にキャプションが必ず先頭に来るようになります。
+
+#### dataset_repeats
+データセットの枚数が極端に少ない場合、epochがすぐに終わってしまうため（epochの区切りで少し時間が掛かります）、数値を指定してデータを何倍かしてepochを長めにしてください。
+
+#### train_text_encoder
+Text Encoderも学習対象とします。メモリ使用量が若干増加します。
+
+通常のfine tuningではText Encoderは学習対象としませんが（恐らくText Encoderの出力に従うようにU-Netを学習するため）、学習データ数が少ない場合には、DreamBoothのようにText Encoder側に学習させるのも有効的なようです。
+
+#### save_precision
+checkpoint保存時のデータ形式をfloat、fp16、bf16から指定できます（未指定時は学習中のデータ形式と同じ）。ディスク容量が節約できますがモデルによる生成結果は変わってきます。またfloatやfp16を指定すると、1111氏のWeb UIでも読めるようになるはずです。
+
+※VAEについては元のcheckpointのデータ形式のままになりますので、fp16でもモデルサイズが2GB強まで小さくならない場合があります。
+
+#### save_model_as
+モデルの保存形式を指定します。ckpt、safetensors、diffusers、diffusers_safetensorsのいずれかを指定してください。
+
+Stable Diffusion形式（ckptまたはsafetensors）を読み込み、Diffusers形式で保存する場合、不足する情報はHugging Faceからv1.5またはv2.1の情報を落としてきて補完します。
+
+#### use_safetensors
+このオプションを指定するとsafetensors形式でcheckpointを保存します。保存形式はデフォルト（読み込んだ形式と同じ）になります。
+
+#### save_stateとresume
+save_stateオプションで、途中保存時および最終保存時に、checkpointに加えてoptimizer等の学習状態をフォルダに保存します。これにより中断してから学習再開したときの精度低下が避けられます（optimizerは状態を持ちながら最適化をしていくため、その状態がリセットされると再び初期状態から最適化を行わなくてはなりません）。なお、Accelerateの仕様でステップ数は保存されません。
+
+スクリプト起動時、resumeオプションで状態の保存されたフォルダを指定すると再開できます。
+
+学習状態は一回の保存あたり5GB程度になりますのでディスク容量にご注意ください。
+
+#### gradient_accumulation_steps
+指定したステップ数だけまとめて勾配を更新します。バッチサイズを増やすのと同様の効果がありますが、メモリを若干消費します。
+
+※Accelerateの仕様で学習モデルが複数の場合には対応していないとのことですので、Text Encoderを学習対象にして、このオプションに2以上の値を指定するとエラーになるかもしれません。
+
+#### lr_scheduler / lr_warmup_steps
+lr_schedulerオプションで学習率のスケジューラをlinear, cosine, cosine_with_restarts, polynomial, constant, constant_with_warmupから選べます。デフォルトはconstantです。
+
+lr_warmup_stepsでスケジューラのウォームアップ（だんだん学習率を変えていく）ステップ数を指定できます。詳細については各自お調べください。
+
+#### diffusers_xformers
+スクリプト独自のxformers置換機能ではなくDiffusersのxformers機能を利用します。Hypernetworkの学習はできなくなります。

finetune/blip/blip.py

@@ -0,0 +1,240 @@
+'''
+ * Copyright (c) 2022, salesforce.com, inc.
+ * All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ * For full license text, see LICENSE.txt file in the repo root or https://opensource.org/licenses/BSD-3-Clause
+ * By Junnan Li
+'''
+import warnings
+warnings.filterwarnings("ignore")
+
+# from models.vit import VisionTransformer, interpolate_pos_embed
+# from models.med import BertConfig, BertModel, BertLMHeadModel
+from blip.vit import VisionTransformer, interpolate_pos_embed
+from blip.med import BertConfig, BertModel, BertLMHeadModel
+from transformers import BertTokenizer
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+import os
+from urllib.parse import urlparse
+from timm.models.hub import download_cached_file
+
+class BLIP_Base(nn.Module):
+    def __init__(self,                 
+                 med_config = 'configs/med_config.json',  
+                 image_size = 224,
+                 vit = 'base',
+                 vit_grad_ckpt = False,
+                 vit_ckpt_layer = 0,                 
+                 ):
+        """
+        Args:
+            med_config (str): path for the mixture of encoder-decoder model's configuration file
+            image_size (int): input image size
+            vit (str): model size of vision transformer
+        """               
+        super().__init__()
+        
+        self.visual_encoder, vision_width = create_vit(vit,image_size, vit_grad_ckpt, vit_ckpt_layer)
+        self.tokenizer = init_tokenizer()   
+        med_config = BertConfig.from_json_file(med_config)
+        med_config.encoder_width = vision_width
+        self.text_encoder = BertModel(config=med_config, add_pooling_layer=False)  
+
+        
+    def forward(self, image, caption, mode):
+        
+        assert mode in ['image', 'text', 'multimodal'], "mode parameter must be image, text, or multimodal"
+        text = self.tokenizer(caption, return_tensors="pt").to(image.device) 
+        
+        if mode=='image':    
+            # return image features
+            image_embeds = self.visual_encoder(image)             
+            return image_embeds
+        
+        elif mode=='text':
+            # return text features
+            text_output = self.text_encoder(text.input_ids, attention_mask = text.attention_mask,                      
+                                            return_dict = True, mode = 'text')  
+            return text_output.last_hidden_state
+        
+        elif mode=='multimodal':
+            # return multimodel features
+            image_embeds = self.visual_encoder(image)    
+            image_atts = torch.ones(image_embeds.size()[:-1],dtype=torch.long).to(image.device)      
+            
+            text.input_ids[:,0] = self.tokenizer.enc_token_id
+            output = self.text_encoder(text.input_ids,
+                                       attention_mask = text.attention_mask,
+                                       encoder_hidden_states = image_embeds,
+                                       encoder_attention_mask = image_atts,      
+                                       return_dict = True,
+                                      )              
+            return output.last_hidden_state
+        
+        
+        
+class BLIP_Decoder(nn.Module):
+    def __init__(self,                 
+                 med_config = 'configs/med_config.json',  
+                 image_size = 384,
+                 vit = 'base',
+                 vit_grad_ckpt = False,
+                 vit_ckpt_layer = 0,
+                 prompt = 'a picture of ',
+                 ):
+        """
+        Args:
+            med_config (str): path for the mixture of encoder-decoder model's configuration file
+            image_size (int): input image size
+            vit (str): model size of vision transformer
+        """            
+        super().__init__()
+        
+        self.visual_encoder, vision_width = create_vit(vit,image_size, vit_grad_ckpt, vit_ckpt_layer)
+        self.tokenizer = init_tokenizer()   
+        med_config = BertConfig.from_json_file(med_config)
+        med_config.encoder_width = vision_width
+        self.text_decoder = BertLMHeadModel(config=med_config)    
+        
+        self.prompt = prompt
+        self.prompt_length = len(self.tokenizer(self.prompt).input_ids)-1
+
+        
+    def forward(self, image, caption):
+        
+        image_embeds = self.visual_encoder(image) 
+        image_atts = torch.ones(image_embeds.size()[:-1],dtype=torch.long).to(image.device)
+        
+        text = self.tokenizer(caption, padding='longest', truncation=True, max_length=40, return_tensors="pt").to(image.device) 
+        
+        text.input_ids[:,0] = self.tokenizer.bos_token_id
+        
+        decoder_targets = text.input_ids.masked_fill(text.input_ids == self.tokenizer.pad_token_id, -100)         
+        decoder_targets[:,:self.prompt_length] = -100
+     
+        decoder_output = self.text_decoder(text.input_ids, 
+                                           attention_mask = text.attention_mask, 
+                                           encoder_hidden_states = image_embeds,
+                                           encoder_attention_mask = image_atts,                  
+                                           labels = decoder_targets,
+                                           return_dict = True,   
+                                          )   
+        loss_lm = decoder_output.loss
+        
+        return loss_lm
+        
+    def generate(self, image, sample=False, num_beams=3, max_length=30, min_length=10, top_p=0.9, repetition_penalty=1.0):
+        image_embeds = self.visual_encoder(image)
+
+        if not sample:
+            image_embeds = image_embeds.repeat_interleave(num_beams,dim=0)
+            
+        image_atts = torch.ones(image_embeds.size()[:-1],dtype=torch.long).to(image.device)
+        model_kwargs = {"encoder_hidden_states": image_embeds, "encoder_attention_mask":image_atts}
+        
+        prompt = [self.prompt] * image.size(0)
+        input_ids = self.tokenizer(prompt, return_tensors="pt").input_ids.to(image.device) 
+        input_ids[:,0] = self.tokenizer.bos_token_id
+        input_ids = input_ids[:, :-1] 
+
+        if sample:
+            #nucleus sampling
+            outputs = self.text_decoder.generate(input_ids=input_ids,
+                                                  max_length=max_length,
+                                                  min_length=min_length,
+                                                  do_sample=True,
+                                                  top_p=top_p,
+                                                  num_return_sequences=1,
+                                                  eos_token_id=self.tokenizer.sep_token_id,
+                                                  pad_token_id=self.tokenizer.pad_token_id, 
+                                                  repetition_penalty=1.1,                                            
+                                                  **model_kwargs)
+        else:
+            #beam search
+            outputs = self.text_decoder.generate(input_ids=input_ids,
+                                                  max_length=max_length,
+                                                  min_length=min_length,
+                                                  num_beams=num_beams,
+                                                  eos_token_id=self.tokenizer.sep_token_id,
+                                                  pad_token_id=self.tokenizer.pad_token_id,     
+                                                  repetition_penalty=repetition_penalty,
+                                                  **model_kwargs)            
+            
+        captions = []    
+        for output in outputs:
+            caption = self.tokenizer.decode(output, skip_special_tokens=True)    
+            captions.append(caption[len(self.prompt):])
+        return captions
+    
+
+def blip_decoder(pretrained='',**kwargs):
+    model = BLIP_Decoder(**kwargs)
+    if pretrained:
+        model,msg = load_checkpoint(model,pretrained)
+        assert(len(msg.missing_keys)==0)
+    return model    
+    
+def blip_feature_extractor(pretrained='',**kwargs):
+    model = BLIP_Base(**kwargs)
+    if pretrained:
+        model,msg = load_checkpoint(model,pretrained)
+        assert(len(msg.missing_keys)==0)
+    return model        
+
+def init_tokenizer():
+    tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
+    tokenizer.add_special_tokens({'bos_token':'[DEC]'})
+    tokenizer.add_special_tokens({'additional_special_tokens':['[ENC]']})       
+    tokenizer.enc_token_id = tokenizer.additional_special_tokens_ids[0]  
+    return tokenizer
+
+
+def create_vit(vit, image_size, use_grad_checkpointing=False, ckpt_layer=0, drop_path_rate=0):
+        
+    assert vit in ['base', 'large'], "vit parameter must be base or large"
+    if vit=='base':
+        vision_width = 768
+        visual_encoder = VisionTransformer(img_size=image_size, patch_size=16, embed_dim=vision_width, depth=12, 
+                                           num_heads=12, use_grad_checkpointing=use_grad_checkpointing, ckpt_layer=ckpt_layer,
+                                           drop_path_rate=0 or drop_path_rate
+                                          )   
+    elif vit=='large':
+        vision_width = 1024
+        visual_encoder = VisionTransformer(img_size=image_size, patch_size=16, embed_dim=vision_width, depth=24, 
+                                           num_heads=16, use_grad_checkpointing=use_grad_checkpointing, ckpt_layer=ckpt_layer,
+                                           drop_path_rate=0.1 or drop_path_rate
+                                          )   
+    return visual_encoder, vision_width
+
+def is_url(url_or_filename):
+    parsed = urlparse(url_or_filename)
+    return parsed.scheme in ("http", "https")
+
+def load_checkpoint(model,url_or_filename):
+    if is_url(url_or_filename):
+        cached_file = download_cached_file(url_or_filename, check_hash=False, progress=True)
+        checkpoint = torch.load(cached_file, map_location='cpu') 
+    elif os.path.isfile(url_or_filename):        
+        checkpoint = torch.load(url_or_filename, map_location='cpu') 
+    else:
+        raise RuntimeError('checkpoint url or path is invalid')
+        
+    state_dict = checkpoint['model']
+    
+    state_dict['visual_encoder.pos_embed'] = interpolate_pos_embed(state_dict['visual_encoder.pos_embed'],model.visual_encoder) 
+    if 'visual_encoder_m.pos_embed' in model.state_dict().keys():
+        state_dict['visual_encoder_m.pos_embed'] = interpolate_pos_embed(state_dict['visual_encoder_m.pos_embed'],
+                                                                         model.visual_encoder_m)    
+    for key in model.state_dict().keys():
+        if key in state_dict.keys():
+            if state_dict[key].shape!=model.state_dict()[key].shape:
+                del state_dict[key]
+    
+    msg = model.load_state_dict(state_dict,strict=False)
+    print('load checkpoint from %s'%url_or_filename)  
+    return model,msg
+    

finetune/blip/med.py

@@ -0,0 +1,955 @@
+'''
+ * Copyright (c) 2022, salesforce.com, inc.
+ * All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ * For full license text, see LICENSE.txt file in the repo root or https://opensource.org/licenses/BSD-3-Clause
+ * By Junnan Li
+ * Based on huggingface code base
+ * https://github.com/huggingface/transformers/blob/v4.15.0/src/transformers/models/bert
+'''
+
+import math
+import os
+import warnings
+from dataclasses import dataclass
+from typing import Optional, Tuple
+
+import torch
+from torch import Tensor, device, dtype, nn
+import torch.utils.checkpoint
+from torch import nn
+from torch.nn import CrossEntropyLoss
+import torch.nn.functional as F
+
+from transformers.activations import ACT2FN
+from transformers.file_utils import (
+    ModelOutput,
+)
+from transformers.modeling_outputs import (
+    BaseModelOutputWithPastAndCrossAttentions,
+    BaseModelOutputWithPoolingAndCrossAttentions,
+    CausalLMOutputWithCrossAttentions,
+    MaskedLMOutput,
+    MultipleChoiceModelOutput,
+    NextSentencePredictorOutput,
+    QuestionAnsweringModelOutput,
+    SequenceClassifierOutput,
+    TokenClassifierOutput,
+)
+from transformers.modeling_utils import (
+    PreTrainedModel,
+    apply_chunking_to_forward,
+    find_pruneable_heads_and_indices,
+    prune_linear_layer,
+)
+from transformers.utils import logging
+from transformers.models.bert.configuration_bert import BertConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+class BertEmbeddings(nn.Module):
+    """Construct the embeddings from word and position embeddings."""
+
+    def __init__(self, config):
+        super().__init__()
+        self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
+        self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)
+
+        # self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be able to load
+        # any TensorFlow checkpoint file
+        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+        # position_ids (1, len position emb) is contiguous in memory and exported when serialized
+        self.register_buffer("position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)))
+        self.position_embedding_type = getattr(config, "position_embedding_type", "absolute")
+        
+        self.config = config
+
+    def forward(
+        self, input_ids=None, position_ids=None, inputs_embeds=None, past_key_values_length=0
+    ):
+        if input_ids is not None:
+            input_shape = input_ids.size()
+        else:
+            input_shape = inputs_embeds.size()[:-1]
+
+        seq_length = input_shape[1]
+
+        if position_ids is None:
+            position_ids = self.position_ids[:, past_key_values_length : seq_length + past_key_values_length]
+
+        if inputs_embeds is None:
+            inputs_embeds = self.word_embeddings(input_ids)
+
+        embeddings = inputs_embeds
+
+        if self.position_embedding_type == "absolute":
+            position_embeddings = self.position_embeddings(position_ids)
+            embeddings += position_embeddings
+        embeddings = self.LayerNorm(embeddings)
+        embeddings = self.dropout(embeddings)
+        return embeddings
+
+
+class BertSelfAttention(nn.Module):
+    def __init__(self, config, is_cross_attention):
+        super().__init__()
+        self.config = config
+        if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"):
+            raise ValueError(
+                "The hidden size (%d) is not a multiple of the number of attention "
+                "heads (%d)" % (config.hidden_size, config.num_attention_heads)
+            )
+        
+        self.num_attention_heads = config.num_attention_heads
+        self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
+        self.all_head_size = self.num_attention_heads * self.attention_head_size
+
+        self.query = nn.Linear(config.hidden_size, self.all_head_size)
+        if is_cross_attention:
+            self.key = nn.Linear(config.encoder_width, self.all_head_size)
+            self.value = nn.Linear(config.encoder_width, self.all_head_size)
+        else:
+            self.key = nn.Linear(config.hidden_size, self.all_head_size)
+            self.value = nn.Linear(config.hidden_size, self.all_head_size)
+
+        self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
+        self.position_embedding_type = getattr(config, "position_embedding_type", "absolute")
+        if self.position_embedding_type == "relative_key" or self.position_embedding_type == "relative_key_query":
+            self.max_position_embeddings = config.max_position_embeddings
+            self.distance_embedding = nn.Embedding(2 * config.max_position_embeddings - 1, self.attention_head_size)
+        self.save_attention = False   
+            
+    def save_attn_gradients(self, attn_gradients):
+        self.attn_gradients = attn_gradients
+        
+    def get_attn_gradients(self):
+        return self.attn_gradients
+    
+    def save_attention_map(self, attention_map):
+        self.attention_map = attention_map
+        
+    def get_attention_map(self):
+        return self.attention_map
+    
+    def transpose_for_scores(self, x):
+        new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
+        x = x.view(*new_x_shape)
+        return x.permute(0, 2, 1, 3)
+
+    def forward(
+        self,
+        hidden_states,
+        attention_mask=None,
+        head_mask=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        past_key_value=None,
+        output_attentions=False,
+    ):
+        mixed_query_layer = self.query(hidden_states)
+
+        # If this is instantiated as a cross-attention module, the keys
+        # and values come from an encoder; the attention mask needs to be
+        # such that the encoder's padding tokens are not attended to.
+        is_cross_attention = encoder_hidden_states is not None
+
+        if is_cross_attention:
+            key_layer = self.transpose_for_scores(self.key(encoder_hidden_states))
+            value_layer = self.transpose_for_scores(self.value(encoder_hidden_states))
+            attention_mask = encoder_attention_mask
+        elif past_key_value is not None:
+            key_layer = self.transpose_for_scores(self.key(hidden_states))
+            value_layer = self.transpose_for_scores(self.value(hidden_states))
+            key_layer = torch.cat([past_key_value[0], key_layer], dim=2)
+            value_layer = torch.cat([past_key_value[1], value_layer], dim=2)
+        else:
+            key_layer = self.transpose_for_scores(self.key(hidden_states))
+            value_layer = self.transpose_for_scores(self.value(hidden_states))
+
+        query_layer = self.transpose_for_scores(mixed_query_layer)
+
+        past_key_value = (key_layer, value_layer)
+
+        # Take the dot product between "query" and "key" to get the raw attention scores.
+        attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
+
+        if self.position_embedding_type == "relative_key" or self.position_embedding_type == "relative_key_query":
+            seq_length = hidden_states.size()[1]
+            position_ids_l = torch.arange(seq_length, dtype=torch.long, device=hidden_states.device).view(-1, 1)
+            position_ids_r = torch.arange(seq_length, dtype=torch.long, device=hidden_states.device).view(1, -1)
+            distance = position_ids_l - position_ids_r
+            positional_embedding = self.distance_embedding(distance + self.max_position_embeddings - 1)
+            positional_embedding = positional_embedding.to(dtype=query_layer.dtype)  # fp16 compatibility
+
+            if self.position_embedding_type == "relative_key":
+                relative_position_scores = torch.einsum("bhld,lrd->bhlr", query_layer, positional_embedding)
+                attention_scores = attention_scores + relative_position_scores
+            elif self.position_embedding_type == "relative_key_query":
+                relative_position_scores_query = torch.einsum("bhld,lrd->bhlr", query_layer, positional_embedding)
+                relative_position_scores_key = torch.einsum("bhrd,lrd->bhlr", key_layer, positional_embedding)
+                attention_scores = attention_scores + relative_position_scores_query + relative_position_scores_key
+
+        attention_scores = attention_scores / math.sqrt(self.attention_head_size)
+        if attention_mask is not None:
+            # Apply the attention mask is (precomputed for all layers in BertModel forward() function)
+            attention_scores = attention_scores + attention_mask
+
+        # Normalize the attention scores to probabilities.
+        attention_probs = nn.Softmax(dim=-1)(attention_scores)
+        
+        if is_cross_attention and self.save_attention:
+            self.save_attention_map(attention_probs)
+            attention_probs.register_hook(self.save_attn_gradients)         
+
+        # This is actually dropping out entire tokens to attend to, which might
+        # seem a bit unusual, but is taken from the original Transformer paper.
+        attention_probs_dropped = self.dropout(attention_probs)
+
+        # Mask heads if we want to
+        if head_mask is not None:
+            attention_probs_dropped = attention_probs_dropped * head_mask
+
+        context_layer = torch.matmul(attention_probs_dropped, value_layer)
+
+        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
+        new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
+        context_layer = context_layer.view(*new_context_layer_shape)
+
+        outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)
+
+        outputs = outputs + (past_key_value,)
+        return outputs
+
+
+class BertSelfOutput(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+    def forward(self, hidden_states, input_tensor):
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.dropout(hidden_states)
+        hidden_states = self.LayerNorm(hidden_states + input_tensor)
+        return hidden_states
+
+
+class BertAttention(nn.Module):
+    def __init__(self, config, is_cross_attention=False):
+        super().__init__()
+        self.self = BertSelfAttention(config, is_cross_attention)
+        self.output = BertSelfOutput(config)
+        self.pruned_heads = set()
+
+    def prune_heads(self, heads):
+        if len(heads) == 0:
+            return
+        heads, index = find_pruneable_heads_and_indices(
+            heads, self.self.num_attention_heads, self.self.attention_head_size, self.pruned_heads
+        )
+
+        # Prune linear layers
+        self.self.query = prune_linear_layer(self.self.query, index)
+        self.self.key = prune_linear_layer(self.self.key, index)
+        self.self.value = prune_linear_layer(self.self.value, index)
+        self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)
+
+        # Update hyper params and store pruned heads
+        self.self.num_attention_heads = self.self.num_attention_heads - len(heads)
+        self.self.all_head_size = self.self.attention_head_size * self.self.num_attention_heads
+        self.pruned_heads = self.pruned_heads.union(heads)
+
+    def forward(
+        self,
+        hidden_states,
+        attention_mask=None,
+        head_mask=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        past_key_value=None,
+        output_attentions=False,
+    ):
+        self_outputs = self.self(
+            hidden_states,
+            attention_mask,
+            head_mask,
+            encoder_hidden_states,
+            encoder_attention_mask,
+            past_key_value,
+            output_attentions,
+        )
+        attention_output = self.output(self_outputs[0], hidden_states)
+        outputs = (attention_output,) + self_outputs[1:]  # add attentions if we output them
+        return outputs
+
+
+class BertIntermediate(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
+        if isinstance(config.hidden_act, str):
+            self.intermediate_act_fn = ACT2FN[config.hidden_act]
+        else:
+            self.intermediate_act_fn = config.hidden_act
+
+    def forward(self, hidden_states):
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.intermediate_act_fn(hidden_states)
+        return hidden_states
+
+
+class BertOutput(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
+        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+    def forward(self, hidden_states, input_tensor):
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.dropout(hidden_states)
+        hidden_states = self.LayerNorm(hidden_states + input_tensor)
+        return hidden_states
+
+
+class BertLayer(nn.Module):
+    def __init__(self, config, layer_num):
+        super().__init__()
+        self.config = config
+        self.chunk_size_feed_forward = config.chunk_size_feed_forward
+        self.seq_len_dim = 1
+        self.attention = BertAttention(config)      
+        self.layer_num = layer_num          
+        if self.config.add_cross_attention:
+            self.crossattention = BertAttention(config, is_cross_attention=self.config.add_cross_attention)
+        self.intermediate = BertIntermediate(config)
+        self.output = BertOutput(config)
+
+    def forward(
+        self,
+        hidden_states,
+        attention_mask=None,
+        head_mask=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        past_key_value=None,
+        output_attentions=False,
+        mode=None,
+    ):
+        # decoder uni-directional self-attention cached key/values tuple is at positions 1,2
+        self_attn_past_key_value = past_key_value[:2] if past_key_value is not None else None
+        self_attention_outputs = self.attention(
+            hidden_states,
+            attention_mask,
+            head_mask,
+            output_attentions=output_attentions,
+            past_key_value=self_attn_past_key_value,
+        )
+        attention_output = self_attention_outputs[0]
+
+        outputs = self_attention_outputs[1:-1]
+        present_key_value = self_attention_outputs[-1]
+
+        if mode=='multimodal':
+            assert encoder_hidden_states is not None, "encoder_hidden_states must be given for cross-attention layers"
+
+            cross_attention_outputs = self.crossattention(
+                attention_output,
+                attention_mask,
+                head_mask,
+                encoder_hidden_states,
+                encoder_attention_mask,
+                output_attentions=output_attentions,
+            )
+            attention_output = cross_attention_outputs[0]
+            outputs = outputs + cross_attention_outputs[1:-1]  # add cross attentions if we output attention weights                               
+        layer_output = apply_chunking_to_forward(
+            self.feed_forward_chunk, self.chunk_size_feed_forward, self.seq_len_dim, attention_output
+        )
+        outputs = (layer_output,) + outputs
+
+        outputs = outputs + (present_key_value,)
+
+        return outputs
+
+    def feed_forward_chunk(self, attention_output):
+        intermediate_output = self.intermediate(attention_output)
+        layer_output = self.output(intermediate_output, attention_output)
+        return layer_output
+
+
+class BertEncoder(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.layer = nn.ModuleList([BertLayer(config,i) for i in range(config.num_hidden_layers)])
+        self.gradient_checkpointing = False
+
+    def forward(
+        self,
+        hidden_states,
+        attention_mask=None,
+        head_mask=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        past_key_values=None,
+        use_cache=None,
+        output_attentions=False,
+        output_hidden_states=False,
+        return_dict=True,
+        mode='multimodal',
+    ):
+        all_hidden_states = () if output_hidden_states else None
+        all_self_attentions = () if output_attentions else None
+        all_cross_attentions = () if output_attentions and self.config.add_cross_attention else None
+
+        next_decoder_cache = () if use_cache else None
+               
+        for i in range(self.config.num_hidden_layers):
+            layer_module = self.layer[i]
+            if output_hidden_states:
+                all_hidden_states = all_hidden_states + (hidden_states,)
+
+            layer_head_mask = head_mask[i] if head_mask is not None else None
+            past_key_value = past_key_values[i] if past_key_values is not None else None
+
+            if self.gradient_checkpointing and self.training:
+
+                if use_cache:
+                    logger.warn(
+                        "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
+                    )
+                    use_cache = False
+
+                def create_custom_forward(module):
+                    def custom_forward(*inputs):
+                        return module(*inputs, past_key_value, output_attentions)
+
+                    return custom_forward
+
+                layer_outputs = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(layer_module),
+                    hidden_states,
+                    attention_mask,
+                    layer_head_mask,
+                    encoder_hidden_states,
+                    encoder_attention_mask,
+                    mode=mode,
+                )
+            else:
+                layer_outputs = layer_module(
+                    hidden_states,
+                    attention_mask,
+                    layer_head_mask,
+                    encoder_hidden_states,
+                    encoder_attention_mask,
+                    past_key_value,
+                    output_attentions,
+                    mode=mode,
+                )
+
+            hidden_states = layer_outputs[0]
+            if use_cache:
+                next_decoder_cache += (layer_outputs[-1],)
+            if output_attentions:
+                all_self_attentions = all_self_attentions + (layer_outputs[1],)
+
+        if output_hidden_states:
+            all_hidden_states = all_hidden_states + (hidden_states,)
+
+        if not return_dict:
+            return tuple(
+                v
+                for v in [
+                    hidden_states,
+                    next_decoder_cache,
+                    all_hidden_states,
+                    all_self_attentions,
+                    all_cross_attentions,
+                ]
+                if v is not None
+            )
+        return BaseModelOutputWithPastAndCrossAttentions(
+            last_hidden_state=hidden_states,
+            past_key_values=next_decoder_cache,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attentions,
+            cross_attentions=all_cross_attentions,
+        )
+
+
+class BertPooler(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        self.activation = nn.Tanh()
+
+    def forward(self, hidden_states):
+        # We "pool" the model by simply taking the hidden state corresponding
+        # to the first token.
+        first_token_tensor = hidden_states[:, 0]
+        pooled_output = self.dense(first_token_tensor)
+        pooled_output = self.activation(pooled_output)
+        return pooled_output
+
+
+class BertPredictionHeadTransform(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        if isinstance(config.hidden_act, str):
+            self.transform_act_fn = ACT2FN[config.hidden_act]
+        else:
+            self.transform_act_fn = config.hidden_act
+        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+
+    def forward(self, hidden_states):
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.transform_act_fn(hidden_states)
+        hidden_states = self.LayerNorm(hidden_states)
+        return hidden_states
+
+
+class BertLMPredictionHead(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.transform = BertPredictionHeadTransform(config)
+
+        # The output weights are the same as the input embeddings, but there is
+        # an output-only bias for each token.
+        self.decoder = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        self.bias = nn.Parameter(torch.zeros(config.vocab_size))
+
+        # Need a link between the two variables so that the bias is correctly resized with `resize_token_embeddings`
+        self.decoder.bias = self.bias
+
+    def forward(self, hidden_states):
+        hidden_states = self.transform(hidden_states)
+        hidden_states = self.decoder(hidden_states)
+        return hidden_states
+
+
+class BertOnlyMLMHead(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.predictions = BertLMPredictionHead(config)
+
+    def forward(self, sequence_output):
+        prediction_scores = self.predictions(sequence_output)
+        return prediction_scores
+
+
+class BertPreTrainedModel(PreTrainedModel):
+    """
+    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
+    models.
+    """
+
+    config_class = BertConfig
+    base_model_prefix = "bert"
+    _keys_to_ignore_on_load_missing = [r"position_ids"]
+
+    def _init_weights(self, module):
+        """ Initialize the weights """
+        if isinstance(module, (nn.Linear, nn.Embedding)):
+            # Slightly different from the TF version which uses truncated_normal for initialization
+            # cf https://github.com/pytorch/pytorch/pull/5617
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+        elif isinstance(module, nn.LayerNorm):
+            module.bias.data.zero_()
+            module.weight.data.fill_(1.0)
+        if isinstance(module, nn.Linear) and module.bias is not None:
+            module.bias.data.zero_()
+
+
+class BertModel(BertPreTrainedModel):
+    """
+    The model can behave as an encoder (with only self-attention) as well as a decoder, in which case a layer of
+    cross-attention is added between the self-attention layers, following the architecture described in `Attention is
+    all you need <https://arxiv.org/abs/1706.03762>`__ by Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit,
+    Llion Jones, Aidan N. Gomez, Lukasz Kaiser and Illia Polosukhin.
+    argument and :obj:`add_cross_attention` set to :obj:`True`; an :obj:`encoder_hidden_states` is then expected as an
+    input to the forward pass.
+    """
+
+    def __init__(self, config, add_pooling_layer=True):
+        super().__init__(config)
+        self.config = config
+
+        self.embeddings = BertEmbeddings(config)
+        
+        self.encoder = BertEncoder(config)
+
+        self.pooler = BertPooler(config) if add_pooling_layer else None
+
+        self.init_weights()
+ 
+
+    def get_input_embeddings(self):
+        return self.embeddings.word_embeddings
+
+    def set_input_embeddings(self, value):
+        self.embeddings.word_embeddings = value
+
+    def _prune_heads(self, heads_to_prune):
+        """
+        Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
+        class PreTrainedModel
+        """
+        for layer, heads in heads_to_prune.items():
+            self.encoder.layer[layer].attention.prune_heads(heads)
+
+    
+    def get_extended_attention_mask(self, attention_mask: Tensor, input_shape: Tuple[int], device: device, is_decoder: bool) -> Tensor:
+        """
+        Makes broadcastable attention and causal masks so that future and masked tokens are ignored.
+
+        Arguments:
+            attention_mask (:obj:`torch.Tensor`):
+                Mask with ones indicating tokens to attend to, zeros for tokens to ignore.
+            input_shape (:obj:`Tuple[int]`):
+                The shape of the input to the model.
+            device: (:obj:`torch.device`):
+                The device of the input to the model.
+
+        Returns:
+            :obj:`torch.Tensor` The extended attention mask, with a the same dtype as :obj:`attention_mask.dtype`.
+        """
+        # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
+        # ourselves in which case we just need to make it broadcastable to all heads.
+        if attention_mask.dim() == 3:
+            extended_attention_mask = attention_mask[:, None, :, :]
+        elif attention_mask.dim() == 2:
+            # Provided a padding mask of dimensions [batch_size, seq_length]
+            # - if the model is a decoder, apply a causal mask in addition to the padding mask
+            # - if the model is an encoder, make the mask broadcastable to [batch_size, num_heads, seq_length, seq_length]
+            if is_decoder:
+                batch_size, seq_length = input_shape
+
+                seq_ids = torch.arange(seq_length, device=device)
+                causal_mask = seq_ids[None, None, :].repeat(batch_size, seq_length, 1) <= seq_ids[None, :, None]
+                # in case past_key_values are used we need to add a prefix ones mask to the causal mask
+                # causal and attention masks must have same type with pytorch version < 1.3
+                causal_mask = causal_mask.to(attention_mask.dtype)
+   
+                if causal_mask.shape[1] < attention_mask.shape[1]:
+                    prefix_seq_len = attention_mask.shape[1] - causal_mask.shape[1]
+                    causal_mask = torch.cat(
+                        [
+                            torch.ones((batch_size, seq_length, prefix_seq_len), device=device, dtype=causal_mask.dtype),
+                            causal_mask,
+                        ],
+                        axis=-1,
+                    )                     
+
+                extended_attention_mask = causal_mask[:, None, :, :] * attention_mask[:, None, None, :]
+            else:
+                extended_attention_mask = attention_mask[:, None, None, :]
+        else:
+            raise ValueError(
+                "Wrong shape for input_ids (shape {}) or attention_mask (shape {})".format(
+                    input_shape, attention_mask.shape
+                )
+            )
+
+        # Since attention_mask is 1.0 for positions we want to attend and 0.0 for
+        # masked positions, this operation will create a tensor which is 0.0 for
+        # positions we want to attend and -10000.0 for masked positions.
+        # Since we are adding it to the raw scores before the softmax, this is
+        # effectively the same as removing these entirely.
+        extended_attention_mask = extended_attention_mask.to(dtype=self.dtype)  # fp16 compatibility
+        extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0
+        return extended_attention_mask
+    
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        position_ids=None,
+        head_mask=None,
+        inputs_embeds=None,
+        encoder_embeds=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        past_key_values=None,
+        use_cache=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+        is_decoder=False,
+        mode='multimodal',
+    ):
+        r"""
+        encoder_hidden_states  (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`, `optional`):
+            Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if
+            the model is configured as a decoder.
+        encoder_attention_mask (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`):
+            Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in
+            the cross-attention if the model is configured as a decoder. Mask values selected in ``[0, 1]``:
+            - 1 for tokens that are **not masked**,
+            - 0 for tokens that are **masked**.
+        past_key_values (:obj:`tuple(tuple(torch.FloatTensor))` of length :obj:`config.n_layers` with each tuple having 4 tensors of shape :obj:`(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`):
+            Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding.
+            If :obj:`past_key_values` are used, the user can optionally input only the last :obj:`decoder_input_ids`
+            (those that don't have their past key value states given to this model) of shape :obj:`(batch_size, 1)`
+            instead of all :obj:`decoder_input_ids` of shape :obj:`(batch_size, sequence_length)`.
+        use_cache (:obj:`bool`, `optional`):
+            If set to :obj:`True`, :obj:`past_key_values` key value states are returned and can be used to speed up
+            decoding (see :obj:`past_key_values`).
+        """
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if is_decoder:
+            use_cache = use_cache if use_cache is not None else self.config.use_cache
+        else:
+            use_cache = False
+
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
+        elif input_ids is not None:
+            input_shape = input_ids.size()
+            batch_size, seq_length = input_shape
+            device = input_ids.device
+        elif inputs_embeds is not None:
+            input_shape = inputs_embeds.size()[:-1]
+            batch_size, seq_length = input_shape
+            device = inputs_embeds.device
+        elif encoder_embeds is not None:    
+            input_shape = encoder_embeds.size()[:-1]
+            batch_size, seq_length = input_shape 
+            device = encoder_embeds.device
+        else:
+            raise ValueError("You have to specify either input_ids or inputs_embeds or encoder_embeds")
+
+        # past_key_values_length
+        past_key_values_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0
+
+        if attention_mask is None:
+            attention_mask = torch.ones(((batch_size, seq_length + past_key_values_length)), device=device)
+            
+        # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
+        # ourselves in which case we just need to make it broadcastable to all heads.
+        extended_attention_mask: torch.Tensor = self.get_extended_attention_mask(attention_mask, input_shape, 
+                                                                                 device, is_decoder)
+
+        # If a 2D or 3D attention mask is provided for the cross-attention
+        # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
+        if encoder_hidden_states is not None:
+            if type(encoder_hidden_states) == list:
+                encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states[0].size()
+            else:
+                encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size()
+            encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length)
+            
+            if type(encoder_attention_mask) == list:
+                encoder_extended_attention_mask = [self.invert_attention_mask(mask) for mask in encoder_attention_mask]
+            elif encoder_attention_mask is None:
+                encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device)
+                encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask)
+            else:    
+                encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask)
+        else:
+            encoder_extended_attention_mask = None
+
+        # Prepare head mask if needed
+        # 1.0 in head_mask indicate we keep the head
+        # attention_probs has shape bsz x n_heads x N x N
+        # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
+        # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
+        head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)
+        
+        if encoder_embeds is None:
+            embedding_output = self.embeddings(
+                input_ids=input_ids,
+                position_ids=position_ids,
+                inputs_embeds=inputs_embeds,
+                past_key_values_length=past_key_values_length,
+            )
+        else:
+            embedding_output = encoder_embeds
+            
+        encoder_outputs = self.encoder(
+            embedding_output,
+            attention_mask=extended_attention_mask,
+            head_mask=head_mask,
+            encoder_hidden_states=encoder_hidden_states,
+            encoder_attention_mask=encoder_extended_attention_mask,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            mode=mode,
+        )
+        sequence_output = encoder_outputs[0]
+        pooled_output = self.pooler(sequence_output) if self.pooler is not None else None
+
+        if not return_dict:
+            return (sequence_output, pooled_output) + encoder_outputs[1:]
+
+        return BaseModelOutputWithPoolingAndCrossAttentions(
+            last_hidden_state=sequence_output,
+            pooler_output=pooled_output,
+            past_key_values=encoder_outputs.past_key_values,
+            hidden_states=encoder_outputs.hidden_states,
+            attentions=encoder_outputs.attentions,
+            cross_attentions=encoder_outputs.cross_attentions,
+        )
+
+
+
+class BertLMHeadModel(BertPreTrainedModel):
+
+    _keys_to_ignore_on_load_unexpected = [r"pooler"]
+    _keys_to_ignore_on_load_missing = [r"position_ids", r"predictions.decoder.bias"]
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.bert = BertModel(config, add_pooling_layer=False)
+        self.cls = BertOnlyMLMHead(config)
+
+        self.init_weights()
+
+    def get_output_embeddings(self):
+        return self.cls.predictions.decoder
+
+    def set_output_embeddings(self, new_embeddings):
+        self.cls.predictions.decoder = new_embeddings
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        position_ids=None,
+        head_mask=None,
+        inputs_embeds=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        labels=None,
+        past_key_values=None,
+        use_cache=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+        return_logits=False,            
+        is_decoder=True,
+        reduction='mean',
+        mode='multimodal', 
+    ):
+        r"""
+        encoder_hidden_states  (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`, `optional`):
+            Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if
+            the model is configured as a decoder.
+        encoder_attention_mask (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`):
+            Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in
+            the cross-attention if the model is configured as a decoder. Mask values selected in ``[0, 1]``:
+            - 1 for tokens that are **not masked**,
+            - 0 for tokens that are **masked**.
+        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`):
+            Labels for computing the left-to-right language modeling loss (next word prediction). Indices should be in
+            ``[-100, 0, ..., config.vocab_size]`` (see ``input_ids`` docstring) Tokens with indices set to ``-100`` are
+            ignored (masked), the loss is only computed for the tokens with labels n ``[0, ..., config.vocab_size]``
+        past_key_values (:obj:`tuple(tuple(torch.FloatTensor))` of length :obj:`config.n_layers` with each tuple having 4 tensors of shape :obj:`(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`):
+            Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding.
+            If :obj:`past_key_values` are used, the user can optionally input only the last :obj:`decoder_input_ids`
+            (those that don't have their past key value states given to this model) of shape :obj:`(batch_size, 1)`
+            instead of all :obj:`decoder_input_ids` of shape :obj:`(batch_size, sequence_length)`.
+        use_cache (:obj:`bool`, `optional`):
+            If set to :obj:`True`, :obj:`past_key_values` key value states are returned and can be used to speed up
+            decoding (see :obj:`past_key_values`).
+        Returns:
+        Example::
+            >>> from transformers import BertTokenizer, BertLMHeadModel, BertConfig
+            >>> import torch
+            >>> tokenizer = BertTokenizer.from_pretrained('bert-base-cased')
+            >>> config = BertConfig.from_pretrained("bert-base-cased")
+            >>> model = BertLMHeadModel.from_pretrained('bert-base-cased', config=config)
+            >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt")
+            >>> outputs = model(**inputs)
+            >>> prediction_logits = outputs.logits
+        """
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        if labels is not None:
+            use_cache = False
+
+        outputs = self.bert(
+            input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            encoder_hidden_states=encoder_hidden_states,
+            encoder_attention_mask=encoder_attention_mask,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            is_decoder=is_decoder,
+            mode=mode,
+        )
+        
+        sequence_output = outputs[0]
+        prediction_scores = self.cls(sequence_output)
+        
+        if return_logits:
+            return prediction_scores[:, :-1, :].contiguous()  
+
+        lm_loss = None
+        if labels is not None:
+            # we are doing next-token prediction; shift prediction scores and input ids by one
+            shifted_prediction_scores = prediction_scores[:, :-1, :].contiguous()
+            labels = labels[:, 1:].contiguous()
+            loss_fct = CrossEntropyLoss(reduction=reduction, label_smoothing=0.1) 
+            lm_loss = loss_fct(shifted_prediction_scores.view(-1, self.config.vocab_size), labels.view(-1))
+            if reduction=='none':
+                lm_loss = lm_loss.view(prediction_scores.size(0),-1).sum(1)               
+
+        if not return_dict:
+            output = (prediction_scores,) + outputs[2:]
+            return ((lm_loss,) + output) if lm_loss is not None else output
+
+        return CausalLMOutputWithCrossAttentions(
+            loss=lm_loss,
+            logits=prediction_scores,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+            cross_attentions=outputs.cross_attentions,
+        )
+
+    def prepare_inputs_for_generation(self, input_ids, past=None, attention_mask=None, **model_kwargs):
+        input_shape = input_ids.shape
+        # if model is used as a decoder in encoder-decoder model, the decoder attention mask is created on the fly
+        if attention_mask is None:
+            attention_mask = input_ids.new_ones(input_shape)
+
+        # cut decoder_input_ids if past is used
+        if past is not None:
+            input_ids = input_ids[:, -1:]
+
+        return {
+            "input_ids": input_ids, 
+            "attention_mask": attention_mask, 
+            "past_key_values": past,
+            "encoder_hidden_states": model_kwargs.get("encoder_hidden_states", None),
+            "encoder_attention_mask": model_kwargs.get("encoder_attention_mask", None),
+            "is_decoder": True,
+        }
+
+    def _reorder_cache(self, past, beam_idx):
+        reordered_past = ()
+        for layer_past in past:
+            reordered_past += (tuple(past_state.index_select(0, beam_idx) for past_state in layer_past),)
+        return reordered_past

finetune/blip/med_config.json

@@ -0,0 +1,22 @@
+{
+    "architectures": [
+      "BertModel"
+    ],
+    "attention_probs_dropout_prob": 0.1,
+    "hidden_act": "gelu",
+    "hidden_dropout_prob": 0.1,
+    "hidden_size": 768,
+    "initializer_range": 0.02,
+    "intermediate_size": 3072,
+    "layer_norm_eps": 1e-12,
+    "max_position_embeddings": 512,
+    "model_type": "bert",
+    "num_attention_heads": 12,
+    "num_hidden_layers": 12,
+    "pad_token_id": 0,
+    "type_vocab_size": 2,
+    "vocab_size": 30524,
+    "encoder_width": 768,
+    "add_cross_attention": true   
+  }
+  

finetune/blip/vit.py

@@ -0,0 +1,305 @@
+'''
+ * Copyright (c) 2022, salesforce.com, inc.
+ * All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ * For full license text, see LICENSE.txt file in the repo root or https://opensource.org/licenses/BSD-3-Clause
+ * By Junnan Li
+ * Based on timm code base
+ * https://github.com/rwightman/pytorch-image-models/tree/master/timm
+'''
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from functools import partial
+
+from timm.models.vision_transformer import _cfg, PatchEmbed
+from timm.models.registry import register_model
+from timm.models.layers import trunc_normal_, DropPath
+from timm.models.helpers import named_apply, adapt_input_conv
+
+from fairscale.nn.checkpoint.checkpoint_activations import checkpoint_wrapper
+
+class Mlp(nn.Module):
+    """ MLP as used in Vision Transformer, MLP-Mixer and related networks
+    """
+    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+class Attention(nn.Module):
+    def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        # NOTE scale factor was wrong in my original version, can set manually to be compat with prev weights
+        self.scale = qk_scale or head_dim ** -0.5
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+        self.attn_gradients = None
+        self.attention_map = None
+        
+    def save_attn_gradients(self, attn_gradients):
+        self.attn_gradients = attn_gradients
+        
+    def get_attn_gradients(self):
+        return self.attn_gradients
+    
+    def save_attention_map(self, attention_map):
+        self.attention_map = attention_map
+        
+    def get_attention_map(self):
+        return self.attention_map
+    
+    def forward(self, x, register_hook=False):
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]   # make torchscript happy (cannot use tensor as tuple)
+
+        attn = (q @ k.transpose(-2, -1)) * self.scale
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+                
+        if register_hook:
+            self.save_attention_map(attn)
+            attn.register_hook(self.save_attn_gradients)        
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class Block(nn.Module):
+
+    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
+                 drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, use_grad_checkpointing=False):
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.attn = Attention(
+            dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
+        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
+
+        if use_grad_checkpointing:
+            self.attn = checkpoint_wrapper(self.attn)
+            self.mlp = checkpoint_wrapper(self.mlp)
+
+    def forward(self, x, register_hook=False):
+        x = x + self.drop_path(self.attn(self.norm1(x), register_hook=register_hook))
+        x = x + self.drop_path(self.mlp(self.norm2(x)))
+        return x
+
+    
+class VisionTransformer(nn.Module):
+    """ Vision Transformer
+    A PyTorch impl of : `An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale`  -
+        https://arxiv.org/abs/2010.11929
+    """
+    def __init__(self, img_size=224, patch_size=16, in_chans=3, num_classes=1000, embed_dim=768, depth=12,
+                 num_heads=12, mlp_ratio=4., qkv_bias=True, qk_scale=None, representation_size=None,
+                 drop_rate=0., attn_drop_rate=0., drop_path_rate=0., norm_layer=None, 
+                 use_grad_checkpointing=False, ckpt_layer=0):
+        """
+        Args:
+            img_size (int, tuple): input image size
+            patch_size (int, tuple): patch size
+            in_chans (int): number of input channels
+            num_classes (int): number of classes for classification head
+            embed_dim (int): embedding dimension
+            depth (int): depth of transformer
+            num_heads (int): number of attention heads
+            mlp_ratio (int): ratio of mlp hidden dim to embedding dim
+            qkv_bias (bool): enable bias for qkv if True
+            qk_scale (float): override default qk scale of head_dim ** -0.5 if set
+            representation_size (Optional[int]): enable and set representation layer (pre-logits) to this value if set
+            drop_rate (float): dropout rate
+            attn_drop_rate (float): attention dropout rate
+            drop_path_rate (float): stochastic depth rate
+            norm_layer: (nn.Module): normalization layer
+        """
+        super().__init__()
+        self.num_features = self.embed_dim = embed_dim  # num_features for consistency with other models
+        norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6)
+
+        self.patch_embed = PatchEmbed(
+            img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)
+
+        num_patches = self.patch_embed.num_patches
+
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))
+        self.pos_drop = nn.Dropout(p=drop_rate)
+
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule
+        self.blocks = nn.ModuleList([
+            Block(
+                dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,
+                drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer,
+                use_grad_checkpointing=(use_grad_checkpointing and i>=depth-ckpt_layer)
+            )
+            for i in range(depth)])
+        self.norm = norm_layer(embed_dim)
+
+        trunc_normal_(self.pos_embed, std=.02)
+        trunc_normal_(self.cls_token, std=.02)
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        return {'pos_embed', 'cls_token'}
+
+    def forward(self, x, register_blk=-1):
+        B = x.shape[0]
+        x = self.patch_embed(x)
+
+        cls_tokens = self.cls_token.expand(B, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
+        x = torch.cat((cls_tokens, x), dim=1)
+  
+        x = x + self.pos_embed[:,:x.size(1),:]
+        x = self.pos_drop(x)
+
+        for i,blk in enumerate(self.blocks):
+            x = blk(x, register_blk==i)
+        x = self.norm(x)
+        
+        return x
+
+    @torch.jit.ignore()
+    def load_pretrained(self, checkpoint_path, prefix=''):
+        _load_weights(self, checkpoint_path, prefix)
+        
+
+@torch.no_grad()
+def _load_weights(model: VisionTransformer, checkpoint_path: str, prefix: str = ''):
+    """ Load weights from .npz checkpoints for official Google Brain Flax implementation
+    """
+    import numpy as np
+
+    def _n2p(w, t=True):
+        if w.ndim == 4 and w.shape[0] == w.shape[1] == w.shape[2] == 1:
+            w = w.flatten()
+        if t:
+            if w.ndim == 4:
+                w = w.transpose([3, 2, 0, 1])
+            elif w.ndim == 3:
+                w = w.transpose([2, 0, 1])
+            elif w.ndim == 2:
+                w = w.transpose([1, 0])
+        return torch.from_numpy(w)
+
+    w = np.load(checkpoint_path)
+    if not prefix and 'opt/target/embedding/kernel' in w:
+        prefix = 'opt/target/'
+
+    if hasattr(model.patch_embed, 'backbone'):
+        # hybrid
+        backbone = model.patch_embed.backbone
+        stem_only = not hasattr(backbone, 'stem')
+        stem = backbone if stem_only else backbone.stem
+        stem.conv.weight.copy_(adapt_input_conv(stem.conv.weight.shape[1], _n2p(w[f'{prefix}conv_root/kernel'])))
+        stem.norm.weight.copy_(_n2p(w[f'{prefix}gn_root/scale']))
+        stem.norm.bias.copy_(_n2p(w[f'{prefix}gn_root/bias']))
+        if not stem_only:
+            for i, stage in enumerate(backbone.stages):
+                for j, block in enumerate(stage.blocks):
+                    bp = f'{prefix}block{i + 1}/unit{j + 1}/'
+                    for r in range(3):
+                        getattr(block, f'conv{r + 1}').weight.copy_(_n2p(w[f'{bp}conv{r + 1}/kernel']))
+                        getattr(block, f'norm{r + 1}').weight.copy_(_n2p(w[f'{bp}gn{r + 1}/scale']))
+                        getattr(block, f'norm{r + 1}').bias.copy_(_n2p(w[f'{bp}gn{r + 1}/bias']))
+                    if block.downsample is not None:
+                        block.downsample.conv.weight.copy_(_n2p(w[f'{bp}conv_proj/kernel']))
+                        block.downsample.norm.weight.copy_(_n2p(w[f'{bp}gn_proj/scale']))
+                        block.downsample.norm.bias.copy_(_n2p(w[f'{bp}gn_proj/bias']))
+        embed_conv_w = _n2p(w[f'{prefix}embedding/kernel'])
+    else:
+        embed_conv_w = adapt_input_conv(
+            model.patch_embed.proj.weight.shape[1], _n2p(w[f'{prefix}embedding/kernel']))
+    model.patch_embed.proj.weight.copy_(embed_conv_w)
+    model.patch_embed.proj.bias.copy_(_n2p(w[f'{prefix}embedding/bias']))
+    model.cls_token.copy_(_n2p(w[f'{prefix}cls'], t=False))
+    pos_embed_w = _n2p(w[f'{prefix}Transformer/posembed_input/pos_embedding'], t=False)
+    if pos_embed_w.shape != model.pos_embed.shape:
+        pos_embed_w = resize_pos_embed(  # resize pos embedding when different size from pretrained weights
+            pos_embed_w, model.pos_embed, getattr(model, 'num_tokens', 1), model.patch_embed.grid_size)
+    model.pos_embed.copy_(pos_embed_w)
+    model.norm.weight.copy_(_n2p(w[f'{prefix}Transformer/encoder_norm/scale']))
+    model.norm.bias.copy_(_n2p(w[f'{prefix}Transformer/encoder_norm/bias']))
+#     if isinstance(model.head, nn.Linear) and model.head.bias.shape[0] == w[f'{prefix}head/bias'].shape[-1]:
+#         model.head.weight.copy_(_n2p(w[f'{prefix}head/kernel']))
+#         model.head.bias.copy_(_n2p(w[f'{prefix}head/bias']))
+#     if isinstance(getattr(model.pre_logits, 'fc', None), nn.Linear) and f'{prefix}pre_logits/bias' in w:
+#         model.pre_logits.fc.weight.copy_(_n2p(w[f'{prefix}pre_logits/kernel']))
+#         model.pre_logits.fc.bias.copy_(_n2p(w[f'{prefix}pre_logits/bias']))
+    for i, block in enumerate(model.blocks.children()):
+        block_prefix = f'{prefix}Transformer/encoderblock_{i}/'
+        mha_prefix = block_prefix + 'MultiHeadDotProductAttention_1/'
+        block.norm1.weight.copy_(_n2p(w[f'{block_prefix}LayerNorm_0/scale']))
+        block.norm1.bias.copy_(_n2p(w[f'{block_prefix}LayerNorm_0/bias']))
+        block.attn.qkv.weight.copy_(torch.cat([
+            _n2p(w[f'{mha_prefix}{n}/kernel'], t=False).flatten(1).T for n in ('query', 'key', 'value')]))
+        block.attn.qkv.bias.copy_(torch.cat([
+            _n2p(w[f'{mha_prefix}{n}/bias'], t=False).reshape(-1) for n in ('query', 'key', 'value')]))
+        block.attn.proj.weight.copy_(_n2p(w[f'{mha_prefix}out/kernel']).flatten(1))
+        block.attn.proj.bias.copy_(_n2p(w[f'{mha_prefix}out/bias']))
+        for r in range(2):
+            getattr(block.mlp, f'fc{r + 1}').weight.copy_(_n2p(w[f'{block_prefix}MlpBlock_3/Dense_{r}/kernel']))
+            getattr(block.mlp, f'fc{r + 1}').bias.copy_(_n2p(w[f'{block_prefix}MlpBlock_3/Dense_{r}/bias']))
+        block.norm2.weight.copy_(_n2p(w[f'{block_prefix}LayerNorm_2/scale']))
+        block.norm2.bias.copy_(_n2p(w[f'{block_prefix}LayerNorm_2/bias']))
+
+            
+def interpolate_pos_embed(pos_embed_checkpoint, visual_encoder):        
+    # interpolate position embedding
+    embedding_size = pos_embed_checkpoint.shape[-1]
+    num_patches = visual_encoder.patch_embed.num_patches
+    num_extra_tokens = visual_encoder.pos_embed.shape[-2] - num_patches
+    # height (== width) for the checkpoint position embedding
+    orig_size = int((pos_embed_checkpoint.shape[-2] - num_extra_tokens) ** 0.5)
+    # height (== width) for the new position embedding
+    new_size = int(num_patches ** 0.5)
+
+    if orig_size!=new_size:
+        # class_token and dist_token are kept unchanged
+        extra_tokens = pos_embed_checkpoint[:, :num_extra_tokens]
+        # only the position tokens are interpolated
+        pos_tokens = pos_embed_checkpoint[:, num_extra_tokens:]
+        pos_tokens = pos_tokens.reshape(-1, orig_size, orig_size, embedding_size).permute(0, 3, 1, 2)
+        pos_tokens = torch.nn.functional.interpolate(
+            pos_tokens, size=(new_size, new_size), mode='bicubic', align_corners=False)
+        pos_tokens = pos_tokens.permute(0, 2, 3, 1).flatten(1, 2)
+        new_pos_embed = torch.cat((extra_tokens, pos_tokens), dim=1)
+        print('reshape position embedding from %d to %d'%(orig_size ** 2,new_size ** 2))
+        
+        return new_pos_embed    
+    else:
+        return pos_embed_checkpoint

finetune/make_captions.py

@@ -1,10 +1,8 @@
-# このスクリプトのライセンスは、Apache License 2.0とします
-# (c) 2022 Kohya S. @kohya_ss
-
 import argparse
 import glob
 import os
 import json
+import random
 
 from PIL import Image
 from tqdm import tqdm
@@ -12,20 +10,34 @@ import numpy as np
 import torch
 from torchvision import transforms
 from torchvision.transforms.functional import InterpolationMode
-from models.blip import blip_decoder
+from blip.blip import blip_decoder
 # from Salesforce_BLIP.models.blip import blip_decoder
 
 DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
 
 def main(args):
+  # fix the seed for reproducibility
+  seed = args.seed # + utils.get_rank()
+  torch.manual_seed(seed)
+  np.random.seed(seed)
+  random.seed(seed)
+    
+  if not os.path.exists("blip"):
+    args.train_data_dir = os.path.abspath(args.train_data_dir)        # convert to absolute path
+
+    cwd = os.getcwd()
+    print('Current Working Directory is: ', cwd)
+    os.chdir('finetune')
+
+  print(f"load images from {args.train_data_dir}")
   image_paths = glob.glob(os.path.join(args.train_data_dir, "*.jpg")) + \
       glob.glob(os.path.join(args.train_data_dir, "*.png")) + glob.glob(os.path.join(args.train_data_dir, "*.webp"))
   print(f"found {len(image_paths)} images.")
 
   print(f"loading BLIP caption: {args.caption_weights}")
   image_size = 384
-  model = blip_decoder(pretrained=args.caption_weights, image_size=image_size, vit='large')
+  model = blip_decoder(pretrained=args.caption_weights, image_size=image_size, vit='large', med_config="./blip/med_config.json")
   model.eval()
   model = model.to(DEVICE)
   print("BLIP loaded")
@@ -75,7 +87,7 @@ def main(args):
 if __name__ == '__main__':
   parser = argparse.ArgumentParser()
   parser.add_argument("train_data_dir", type=str, help="directory for train images / 学習画像データのディレクトリ")
-  parser.add_argument("caption_weights", type=str,
+  parser.add_argument("--caption_weights", type=str, default="https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model_large_caption.pth",
                       help="BLIP caption weights (model_large_caption.pth) / BLIP captionの重みファイル(model_large_caption.pth)")
   parser.add_argument("--caption_extention", type=str, default=None,
                       help="extension of caption file (for backward compatibility) / 出力されるキャプションファイルの拡張子（スペルミスしていたのを残してあります）")
@@ -87,6 +99,7 @@ if __name__ == '__main__':
   parser.add_argument("--top_p", type=float, default=0.9, help="top_p in Nucleus sampling / Nucleus sampling時のtop_p")
   parser.add_argument("--max_length", type=int, default=75, help="max length of caption / captionの最大長")
   parser.add_argument("--min_length", type=int, default=5, help="min length of caption / captionの最小長")
+  parser.add_argument('--seed', default=42, type=int, help='seed for reproducibility / 再現性を確保するための乱数seed')
   parser.add_argument("--debug", action="store_true", help="debug mode")
 
   args = parser.parse_args()

finetune/prepare_buckets_latents.py

@@ -14,7 +14,7 @@ import cv2
 import torch
 from torchvision import transforms
 
-import model_util
+import library.model_util as model_util
 
 DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
@@ -130,14 +130,16 @@ def main(args):
         latents = get_latents(vae, [img for _, _, img in bucket], weight_dtype)
 
         for (image_key, reso, _), latent in zip(bucket, latents):
-          np.savez(os.path.join(args.train_data_dir, os.path.splitext(os.path.basename(image_key))[0]), latent)
+          npz_file_name = os.path.splitext(os.path.basename(image_key))[0] if args.full_path else image_key
+          np.savez(os.path.join(args.train_data_dir, npz_file_name), latent)
 
         # flip
         if args.flip_aug:
           latents = get_latents(vae, [img[:, ::-1].copy() for _, _, img in bucket], weight_dtype)   # copyがないとTensor変換できない
 
           for (image_key, reso, _), latent in zip(bucket, latents):
-            np.savez(os.path.join(args.train_data_dir, os.path.splitext(os.path.basename(image_key))[0] + '_flip'), latent)
+            npz_file_name = os.path.splitext(os.path.basename(image_key))[0] if args.full_path else image_key
+            np.savez(os.path.join(args.train_data_dir, npz_file_name + '_flip'), latent)
 
         bucket.clear()
 

gen_img_diffusers.py

@@ -1,38 +1,3 @@
-# txt2img with Diffusers: supports SD checkpoints, EulerScheduler, clip-skip, 225 tokens, Hypernetwork etc...
-
-# v2: CLIP guided Stable Diffusion, Image guided Stable Diffusion, highres. fix
-# v3: Add dpmsolver/dpmsolver++, add VAE loading, add upscale, add 'bf16', fix the issue hypernetwork_mul is not working
-# v4: SD2.0 support (new U-Net/text encoder/tokenizer), simplify by DiffUsers 0.9.0, no_preview in interactive mode
-# v5: fix clip_sample=True for scheduler, add VGG guidance
-# v6: refactor to use model util, load VAE without vae folder, support safe tensors
-# v7: add use_original_file_name and iter_same_seed option, change vgg16 guide input image size, 
-# Diffusers 0.10.0 (support new schedulers (dpm_2, dpm_2_a, heun, dpmsingle), supports all scheduler in v-prediction)
-# v8: accept wildcard for ckpt name (when only one file is matched), fix a bug app crushes because PIL image doesn't have filename attr sometimes,
-# v9: sort file names, fix an issue in img2img when prompt from metadata with images_per_prompt>1
-# v10: fix app crashes when different image size in prompts
-
-# Copyright 2022 kohya_ss @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.
-
-# license of included scripts:
-
-# FlashAttention: based on https://github.com/lucidrains/memory-efficient-attention-pytorch/blob/main/memory_efficient_attention_pytorch/flash_attention.py
-# MIT https://github.com/lucidrains/memory-efficient-attention-pytorch/blob/main/LICENSE
-
-# Diffusers (model conversion, CLIP guided stable diffusion, schedulers etc.):
-# ASL 2.0 https://github.com/huggingface/diffusers/blob/main/LICENSE
-
 """
 VGG(
   (features): Sequential(
@@ -81,11 +46,13 @@ VGG(
 )
 """
 
+import json
 from typing import List, Optional, Union
 import glob
 import importlib
 import inspect
 import time
+import zipfile
 from diffusers.utils import deprecate
 from diffusers.configuration_utils import FrozenDict
 import argparse
@@ -113,7 +80,7 @@ import PIL
 from PIL import Image
 from PIL.PngImagePlugin import PngInfo
 
-import model_util
+import library.model_util as model_util
 
 # Tokenizer: checkpointから読み込むのではなくあらかじめ提供されているものを使う
 TOKENIZER_PATH = "openai/clip-vit-large-patch14"
@@ -333,7 +300,7 @@ def replace_unet_modules(unet: diffusers.models.unet_2d_condition.UNet2DConditio
 
 
 def replace_unet_cross_attn_to_memory_efficient():
-  print("Replace CrossAttention.forward to use Hypernetwork and FlashAttention")
+  print("Replace CrossAttention.forward to use NAI style Hypernetwork and FlashAttention")
   flash_func = FlashAttentionFunction
 
   def forward_flash_attn(self, x, context=None, mask=None):
@@ -373,7 +340,7 @@ def replace_unet_cross_attn_to_memory_efficient():
 
 
 def replace_unet_cross_attn_to_xformers():
-  print("Replace CrossAttention.forward to use Hypernetwork and xformers")
+  print("Replace CrossAttention.forward to use NAI style Hypernetwork and xformers")
   try:
     import xformers.ops
   except ImportError:
@@ -517,7 +484,7 @@ class PipelineLike():
       self.vgg16_feat_model = torchvision.models._utils.IntermediateLayerGetter(vgg16_model.features, return_layers=return_layers)
       self.vgg16_normalize = transforms.Normalize(mean=VGG16_IMAGE_MEAN, std=VGG16_IMAGE_STD)
 
-# region xformersとか使う部分：独自に書き換えるので関係なし
+  # region xformersとか使う部分：独自に書き換えるので関係なし
   def enable_xformers_memory_efficient_attention(self):
     r"""
     Enable memory efficient attention as implemented in xformers.
@@ -590,6 +557,7 @@ class PipelineLike():
       width: int = 512,
       num_inference_steps: int = 50,
       guidance_scale: float = 7.5,
+      negative_scale: float = None,
       strength: float = 0.8,
       # num_images_per_prompt: Optional[int] = 1,
       eta: float = 0.0,
@@ -708,6 +676,11 @@ class PipelineLike():
     # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
     # corresponds to doing no classifier free guidance.
     do_classifier_free_guidance = guidance_scale > 1.0
+
+    if not do_classifier_free_guidance and negative_scale is not None:
+      print(f"negative_scale is ignored if guidance scalle <= 1.0")
+      negative_scale = None
+
     # get unconditional embeddings for classifier free guidance
     if negative_prompt is None:
       negative_prompt = [""] * batch_size
@@ -729,8 +702,21 @@ class PipelineLike():
         **kwargs,
     )
 
+    if negative_scale is not None:
+      _, real_uncond_embeddings, _ = get_weighted_text_embeddings(
+          pipe=self,
+          prompt=prompt,                                   # こちらのトークン長に合わせてuncondを作るので75トークン超で必須
+          uncond_prompt=[""]*batch_size,
+          max_embeddings_multiples=max_embeddings_multiples,
+          clip_skip=self.clip_skip,
+          **kwargs,
+      )
+
     if do_classifier_free_guidance:
-      text_embeddings = torch.cat([uncond_embeddings, text_embeddings])
+      if negative_scale is None:
+        text_embeddings = torch.cat([uncond_embeddings, text_embeddings])
+      else:
+        text_embeddings = torch.cat([uncond_embeddings, text_embeddings, real_uncond_embeddings])
 
     # CLIP guidanceで使用するembeddingsを取得する
     if self.clip_guidance_scale > 0:
@@ -861,22 +847,28 @@ class PipelineLike():
     if accepts_eta:
       extra_step_kwargs["eta"] = eta
 
+    num_latent_input = (3 if negative_scale is not None else 2) if do_classifier_free_guidance else 1
     for i, t in enumerate(tqdm(timesteps)):
       # expand the latents if we are doing classifier free guidance
-      latent_model_input = latents.repeat((2, 1, 1, 1)) if do_classifier_free_guidance else latents
+      latent_model_input = latents.repeat((num_latent_input, 1, 1, 1))
       latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
-
       # predict the noise residual
       noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample
 
       # perform guidance
       if do_classifier_free_guidance:
-        noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
-        noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
+        if negative_scale is None:
+          noise_pred_uncond, noise_pred_text = noise_pred.chunk(num_latent_input)        # uncond by negative prompt
+          noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
+        else:
+          noise_pred_negative, noise_pred_text, noise_pred_uncond = noise_pred.chunk(num_latent_input)       # uncond is real uncond
+          noise_pred = noise_pred_uncond + guidance_scale * \
+              (noise_pred_text - noise_pred_uncond) - negative_scale * (noise_pred_negative - noise_pred_uncond)
 
       # perform clip guidance
       if self.clip_guidance_scale > 0 or self.clip_image_guidance_scale > 0 or self.vgg16_guidance_scale > 0:
-        text_embeddings_for_guidance = (text_embeddings.chunk(2)[1] if do_classifier_free_guidance else text_embeddings)
+        text_embeddings_for_guidance = (text_embeddings.chunk(num_latent_input)[
+                                        1] if do_classifier_free_guidance else text_embeddings)
 
         if self.clip_guidance_scale > 0:
           noise_pred, latents = self.cond_fn(latents, t, i, text_embeddings_for_guidance, noise_pred,
@@ -1867,25 +1859,6 @@ def main(args):
   if not args.diffusers_xformers:
     replace_unet_modules(unet, not args.xformers, args.xformers)
 
-  # hypernetworkを組み込む
-  if args.hypernetwork_module is not None:
-    assert not args.diffusers_xformers, "cannot use hypernetwork with diffusers_xformers / diffusers_xformers指定時はHypernetworkは利用できません"
-
-    print("import hypernetwork module:", args.hypernetwork_module)
-    hyp_module = importlib.import_module(args.hypernetwork_module)
-
-    hypernetwork = hyp_module.Hypernetwork(args.hypernetwork_mul)
-
-    print("load hypernetwork weights from:", args.hypernetwork_weights)
-    hyp_sd = torch.load(args.hypernetwork_weights, map_location='cpu')
-    success = hypernetwork.load_from_state_dict(hyp_sd)
-    assert success, "hypernetwork weights loading failed."
-
-    if args.opt_channels_last:
-      hypernetwork.to(memory_format=torch.channels_last)
-  else:
-    hypernetwork = None
-
   # tokenizerを読み込む
   print("loading tokenizer")
   if use_stable_diffusion_format:
@@ -2000,10 +1973,51 @@ def main(args):
   if vgg16_model is not None:
     vgg16_model.to(dtype).to(device)
 
-  if hypernetwork is not None:
-    hypernetwork.to(dtype).to(device)
-    print("apply hypernetwork")
-    hypernetwork.apply_to_diffusers(vae, text_encoder, unet)
+  # networkを組み込む
+  if args.network_module:
+    networks = []
+    for i, network_module in enumerate(args.network_module):
+      print("import network module:", network_module)
+      imported_module = importlib.import_module(network_module)
+
+      network_mul = 1.0 if args.network_mul is None or len(args.network_mul) <= i else args.network_mul[i]
+      network_dim = None if args.network_dim is None or len(args.network_dim) <= i else args.network_dim[i]
+
+      net_kwargs = {}
+      if args.network_args and i < len(args.network_args):
+        network_args = args.network_args[i]
+        # TODO escape special chars
+        network_args = network_args.split(";")
+        for net_arg in network_args:
+          key, value = net_arg.split("=")
+          net_kwargs[key] = value
+
+      network = imported_module.create_network(network_mul, network_dim, vae, text_encoder, unet, **net_kwargs)
+      if network is None:
+        return
+
+      if args.network_weights and i < len(args.network_weights):
+        network_weight = args.network_weights[i]
+        print("load network weights from:", network_weight)
+
+        if os.path.splitext(network_weight)[1] == '.safetensors':
+          from safetensors.torch import safe_open
+          with safe_open(network_weight, framework="pt") as f:
+            metadata = f.metadata()
+          if metadata is not None:
+            print(f"metadata for: {network_weight}: {metadata}")
+
+        network.load_weights(network_weight)
+
+      network.apply_to(text_encoder, unet)
+
+      if args.opt_channels_last:
+        network.to(memory_format=torch.channels_last)
+      network.to(dtype).to(device)
+
+      networks.append(network)
+  else:
+    networks = []
 
   if args.opt_channels_last:
     print(f"set optimizing: channels last")
@@ -2012,8 +2026,9 @@ def main(args):
     unet.to(memory_format=torch.channels_last)
     if clip_model is not None:
       clip_model.to(memory_format=torch.channels_last)
-    if hypernetwork is not None:
-      hypernetwork.to(memory_format=torch.channels_last)
+    if networks:
+      for network in networks:
+        network.to(memory_format=torch.channels_last)
     if vgg16_model is not None:
       vgg16_model.to(memory_format=torch.channels_last)
 
@@ -2055,7 +2070,7 @@ def main(args):
         print(f"convert image to RGB from {image.mode}: {p}")
         image = image.convert("RGB")
       images.append(image)
-    
+
     return images
 
   def resize_images(imgs, size):
@@ -2156,12 +2171,12 @@ def main(args):
         # 1st stageのバッチを作成して呼び出す
         print("process 1st stage1")
         batch_1st = []
-        for params1, (width, height, steps, scale, strength) in batch:
+        for params1, (width, height, steps, scale, negative_scale, strength) in batch:
           width_1st = int(width * args.highres_fix_scale + .5)
           height_1st = int(height * args.highres_fix_scale + .5)
           width_1st = width_1st - width_1st % 32
           height_1st = height_1st - height_1st % 32
-          batch_1st.append((params1, (width_1st, height_1st, args.highres_fix_steps, scale, strength)))
+          batch_1st.append((params1, (width_1st, height_1st, args.highres_fix_steps, scale, negative_scale, strength)))
         images_1st = process_batch(batch_1st, True, True)
 
         # 2nd stageのバッチを作成して以下処理する
@@ -2173,7 +2188,8 @@ def main(args):
           batch_2nd.append(((step, prompt, negative_prompt, seed+1, image, None, clip_prompt, guide_image), params2))
         batch = batch_2nd
 
-      (step_first, _, _, _, init_image, mask_image, _, guide_image), (width, height, steps, scale, strength) = batch[0]
+      (step_first, _, _, _, init_image, mask_image, _, guide_image), (width,
+                                                                      height, steps, scale, negative_scale, strength) = batch[0]
       noise_shape = (LATENT_CHANNELS, height // DOWNSAMPLING_FACTOR, width // DOWNSAMPLING_FACTOR)
 
       prompts = []
@@ -2249,7 +2265,7 @@ def main(args):
         guide_images = guide_images[0]
 
       # generate
-      images = pipe(prompts, negative_prompts, init_images, mask_images, height, width, steps, scale, strength, latents=start_code,
+      images = pipe(prompts, negative_prompts, init_images, mask_images, height, width, steps, scale, negative_scale, strength, latents=start_code,
                     output_type='pil', max_embeddings_multiples=max_embeddings_multiples, img2img_noise=i2i_noises, clip_prompts=clip_prompts, clip_guide_images=guide_images)[0]
       if highres_1st and not args.highres_fix_save_1st:
         return images
@@ -2266,6 +2282,8 @@ def main(args):
         metadata.add_text("scale", str(scale))
         if negative_prompt is not None:
           metadata.add_text("negative-prompt", negative_prompt)
+        if negative_scale is not None:
+          metadata.add_text("negative-scale", str(negative_scale))
         if clip_prompt is not None:
           metadata.add_text("clip-prompt", clip_prompt)
 
@@ -2318,6 +2336,7 @@ def main(args):
       width = args.W
       height = args.H
       scale = args.scale
+      negative_scale = args.negative_scale
       steps = args.steps
       seeds = None
       strength = 0.8 if args.strength is None else args.strength
@@ -2360,6 +2379,15 @@ def main(args):
             print(f"scale: {scale}")
             continue
 
+          m = re.match(r'nl ([\d\.]+|none|None)', parg, re.IGNORECASE)
+          if m:               # negative scale
+            if m.group(1).lower() == 'none':
+              negative_scale = None
+            else:
+              negative_scale = float(m.group(1))
+            print(f"negative scale: {negative_scale}")
+            continue
+
           m = re.match(r't ([\d\.]+)', parg, re.IGNORECASE)
           if m:               # strength
             strength = float(m.group(1))
@@ -2422,8 +2450,9 @@ def main(args):
             print("Use previous image as guide image.")
             guide_image = prev_image
 
+        # TODO named tupleか何かにする
         b1 = ((global_step, prompt, negative_prompt, seed, init_image, mask_image, clip_prompt, guide_image),
-              (width, height, steps, scale, strength))
+              (width, height, steps, scale, negative_scale, strength))
         if len(batch_data) > 0 and batch_data[-1][1] != b1[1]:  # バッチ分割必要？
           process_batch(batch_data, highres_fix)
           batch_data.clear()
@@ -2483,17 +2512,24 @@ if __name__ == '__main__':
   #                     help="Replace CLIP (Text Encoder) to l/14@336 / CLIP(Text Encoder)をl/14@336に入れ替える")
   parser.add_argument("--seed", type=int, default=None,
                       help="seed, or seed of seeds in multiple generation / 1枚生成時のseed、または複数枚生成時の乱数seedを決めるためのseed")
-  parser.add_argument("--iter_same_seed", action='store_true', help='use same seed for all prompts in iteration if no seed specified / 乱数seedの指定がないとき繰り返し内はすべて同じseedを使う（プロンプト間の差異の比較用）')
+  parser.add_argument("--iter_same_seed", action='store_true',
+                      help='use same seed for all prompts in iteration if no seed specified / 乱数seedの指定がないとき繰り返し内はすべて同じseedを使う（プロンプト間の差異の比較用）')
   parser.add_argument("--fp16", action='store_true', help='use fp16 / fp16を指定し省メモリ化する')
   parser.add_argument("--bf16", action='store_true', help='use bfloat16 / bfloat16を指定し省メモリ化する')
   parser.add_argument("--xformers", action='store_true', help='use xformers / xformersを使用し高速化する')
   parser.add_argument("--diffusers_xformers", action='store_true',
-                      help='use xformers by diffusers (Hypernetworks doen\'t work) / Diffusersでxformersを使用する（Hypernetwork利用不可）')
+                      help='use xformers by diffusers (Hypernetworks doesn\'t work) / Diffusersでxformersを使用する（Hypernetwork利用不可）')
   parser.add_argument("--opt_channels_last", action='store_true',
-                      help='set channels last option to model / モデルにchannles lastを指定し最適化する')
-  parser.add_argument("--hypernetwork_module", type=str, default=None, help='Hypernetwork module to use / Hypernetworkを使う時そのモジュール名')
-  parser.add_argument("--hypernetwork_weights", type=str, default=None, help='Hypernetwork weights to load / Hypernetworkの重み')
-  parser.add_argument("--hypernetwork_mul", type=float, default=1.0, help='Hypernetwork multiplier / Hypernetworkの効果の倍率')
+                      help='set channels last option to model / モデルにchannels lastを指定し最適化する')
+  parser.add_argument("--network_module", type=str, default=None, nargs='*',
+                      help='Hypernetwork module to use / Hypernetworkを使う時そのモジュール名')
+  parser.add_argument("--network_weights", type=str, default=None, nargs='*',
+                      help='Hypernetwork weights to load / Hypernetworkの重み')
+  parser.add_argument("--network_mul", type=float, default=None, nargs='*', help='Hypernetwork multiplier / Hypernetworkの効果の倍率')
+  parser.add_argument("--network_dim", type=int, default=None, nargs='*',
+                      help='network dimensions (depends on each network) / モジュールの次元数（ネットワークにより定義は異なります）')
+  parser.add_argument("--network_args", type=str, default=None, nargs='*',
+                      help='additional argmuments for network (key=value) / ネットワークへの追加の引数')
   parser.add_argument("--clip_skip", type=int, default=None, help='layer number from bottom to use in CLIP / CLIPの後ろからn層目の出力を使う')
   parser.add_argument("--max_embeddings_multiples", type=int, default=None,
                       help='max embeding multiples, max token length is 75 * multiples / トークン長をデフォルトの何倍とするか 75*この値 がトークン長となる')
@@ -2512,6 +2548,8 @@ if __name__ == '__main__':
                       help="1st stage steps for highres fix / highres fixの最初のステージのステップ数")
   parser.add_argument("--highres_fix_save_1st", action='store_true',
                       help="save 1st stage images for highres fix / highres fixの最初のステージの画像を保存する")
+  parser.add_argument("--negative_scale", type=float, default=None,
+                      help="set another guidance scale for negative prompt / ネガティブプロンプトのscaleを指定する")
 
   args = parser.parse_args()
   main(args)

library/model_util.py

@@ -624,8 +624,16 @@ def convert_ldm_clip_checkpoint_v2(checkpoint, max_length):
       new_sd[key_pfx + "k_proj" + key_suffix] = values[1]
       new_sd[key_pfx + "v_proj" + key_suffix] = values[2]
 
-  # position_idsの追加
-  new_sd["text_model.embeddings.position_ids"] = torch.Tensor([list(range(max_length))]).to(torch.int64)
+  # rename or add position_ids
+  ANOTHER_POSITION_IDS_KEY = "text_model.encoder.text_model.embeddings.position_ids"
+  if ANOTHER_POSITION_IDS_KEY in new_sd:
+    # waifu diffusion v1.4
+    position_ids = new_sd[ANOTHER_POSITION_IDS_KEY]
+    del new_sd[ANOTHER_POSITION_IDS_KEY]
+  else:
+    position_ids = torch.Tensor([list(range(max_length))]).to(torch.int64)
+
+  new_sd["text_model.embeddings.position_ids"] = position_ids
   return new_sd
 
 # endregion
@@ -878,7 +886,7 @@ def load_models_from_stable_diffusion_checkpoint(v2, ckpt_path, dtype=None):
 
   vae = AutoencoderKL(**vae_config)
   info = vae.load_state_dict(converted_vae_checkpoint)
-  print("loadint vae:", info)
+  print("loading vae:", info)
 
   # convert text_model
   if v2:
@@ -1097,12 +1105,12 @@ def load_vae(vae_id, dtype):
 
   if vae_id.endswith(".bin"):
     # SD 1.5 VAE on Huggingface
-    vae_sd = torch.load(vae_id, map_location="cpu")
-    converted_vae_checkpoint = vae_sd
+    converted_vae_checkpoint = torch.load(vae_id, map_location="cpu")
   else:
     # StableDiffusion
-    vae_model = torch.load(vae_id, map_location="cpu")
-    vae_sd = vae_model['state_dict']
+    vae_model = (load_file(vae_id, "cpu") if is_safetensors(vae_id)
+                 else torch.load(vae_id, map_location="cpu"))
+    vae_sd = vae_model['state_dict'] if 'state_dict' in vae_model else vae_model
 
     # vae only or full model
     full_model = False
@@ -1124,15 +1132,6 @@ def load_vae(vae_id, dtype):
   vae.load_state_dict(converted_vae_checkpoint)
   return vae
 
-
-def get_epoch_ckpt_name(use_safetensors, epoch):
-  return f"epoch-{epoch:06d}" + (".safetensors" if use_safetensors else ".ckpt")
-
-
-def get_last_ckpt_name(use_safetensors):
-  return f"last" + (".safetensors" if use_safetensors else ".ckpt")
-
-
 # endregion
 
 

networks/check_lora_weights.py

@@ -0,0 +1,32 @@
+import argparse
+import os
+import torch
+from safetensors.torch import load_file
+
+
+def main(file):
+  print(f"loading: {file}")
+  if os.path.splitext(file)[1] == '.safetensors':
+    sd = load_file(file)
+  else:
+    sd = torch.load(file, map_location='cpu')
+
+  values = []
+
+  keys = list(sd.keys())
+  for key in keys:
+    if 'lora_up' in key or 'lora_down' in key:
+      values.append((key, sd[key]))
+  print(f"number of LoRA modules: {len(values)}")
+
+  for key, value in values:
+    value = value.to(torch.float32)
+    print(f"{key},{torch.mean(torch.abs(value))},{torch.min(torch.abs(value))}")
+
+
+if __name__ == '__main__':
+  parser = argparse.ArgumentParser()
+  parser.add_argument("file", type=str, help="model file to check / 重みを確認するモデルファイル")
+  args = parser.parse_args()
+
+  main(args.file)

networks/extract_lora_from_models.py

@@ -0,0 +1,158 @@
+# extract approximating LoRA by svd from two SD models
+# The code is based on https://github.com/cloneofsimo/lora/blob/develop/lora_diffusion/cli_svd.py
+# Thanks to cloneofsimo!
+
+import argparse
+import os
+import torch
+from safetensors.torch import load_file, save_file
+from tqdm import tqdm
+import library.model_util as model_util
+import lora
+
+
+CLAMP_QUANTILE = 0.99
+MIN_DIFF = 1e-6
+
+
+def save_to_file(file_name, model, state_dict, dtype):
+  if dtype is not None:
+    for key in list(state_dict.keys()):
+      if type(state_dict[key]) == torch.Tensor:
+        state_dict[key] = state_dict[key].to(dtype)
+
+  if os.path.splitext(file_name)[1] == '.safetensors':
+    save_file(model, file_name)
+  else:
+    torch.save(model, file_name)
+
+
+def svd(args):
+  def str_to_dtype(p):
+    if p == 'float':
+      return torch.float
+    if p == 'fp16':
+      return torch.float16
+    if p == 'bf16':
+      return torch.bfloat16
+    return None
+
+  save_dtype = str_to_dtype(args.save_precision)
+
+  print(f"loading SD model : {args.model_org}")
+  text_encoder_o, _, unet_o = model_util.load_models_from_stable_diffusion_checkpoint(args.v2, args.model_org)
+  print(f"loading SD model : {args.model_tuned}")
+  text_encoder_t, _, unet_t = model_util.load_models_from_stable_diffusion_checkpoint(args.v2, args.model_tuned)
+
+  # create LoRA network to extract weights
+  lora_network_o = lora.create_network(1.0, args.dim, None, text_encoder_o, unet_o)
+  lora_network_t = lora.create_network(1.0, args.dim, None, text_encoder_t, unet_t)
+  assert len(lora_network_o.text_encoder_loras) == len(
+      lora_network_t.text_encoder_loras), f"model version is different (SD1.x vs SD2.x) / それぞれのモデルのバージョンが違います（SD1.xベースとSD2.xベース） "
+
+  # get diffs
+  diffs = {}
+  text_encoder_different = False
+  for i, (lora_o, lora_t) in enumerate(zip(lora_network_o.text_encoder_loras, lora_network_t.text_encoder_loras)):
+    lora_name = lora_o.lora_name
+    module_o = lora_o.org_module
+    module_t = lora_t.org_module
+    diff = module_t.weight - module_o.weight
+
+    # Text Encoder might be same
+    if torch.max(torch.abs(diff)) > MIN_DIFF:
+      text_encoder_different = True
+
+    diff = diff.float()
+    diffs[lora_name] = diff
+
+  if not text_encoder_different:
+    print("Text encoder is same. Extract U-Net only.")
+    lora_network_o.text_encoder_loras = []
+    diffs = {}
+
+  for i, (lora_o, lora_t) in enumerate(zip(lora_network_o.unet_loras, lora_network_t.unet_loras)):
+    lora_name = lora_o.lora_name
+    module_o = lora_o.org_module
+    module_t = lora_t.org_module
+    diff = module_t.weight - module_o.weight
+    diff = diff.float()
+    
+    if args.device:
+      diff = diff.to(args.device)
+      
+    diffs[lora_name] = diff
+
+  # make LoRA with svd
+  print("calculating by svd")
+  rank = args.dim
+  lora_weights = {}
+  with torch.no_grad():
+    for lora_name, mat in tqdm(list(diffs.items())):
+      conv2d = (len(mat.size()) == 4)
+      if conv2d:
+        mat = mat.squeeze()
+
+      U, S, Vh = torch.linalg.svd(mat)
+
+      U = U[:, :rank]
+      S = S[:rank]
+      U = U @ torch.diag(S)
+
+      Vh = Vh[:rank, :]
+
+      dist = torch.cat([U.flatten(), Vh.flatten()])
+      hi_val = torch.quantile(dist, CLAMP_QUANTILE)
+      low_val = -hi_val
+
+      U = U.clamp(low_val, hi_val)
+      Vh = Vh.clamp(low_val, hi_val)
+
+      lora_weights[lora_name] = (U, Vh)
+
+  # make state dict for LoRA
+  lora_network_o.apply_to(text_encoder_o, unet_o, text_encoder_different, True)   # to make state dict
+  lora_sd = lora_network_o.state_dict()
+  print(f"LoRA has {len(lora_sd)} weights.")
+
+  for key in list(lora_sd.keys()):
+    lora_name = key.split('.')[0]
+    i = 0 if "lora_up" in key else 1
+
+    weights = lora_weights[lora_name][i]
+    # print(key, i, weights.size(), lora_sd[key].size())
+    if len(lora_sd[key].size()) == 4:
+      weights = weights.unsqueeze(2).unsqueeze(3)
+
+    assert weights.size() == lora_sd[key].size()
+    lora_sd[key] = weights
+
+  # load state dict to LoRA and save it
+  info = lora_network_o.load_state_dict(lora_sd)
+  print(f"Loading extracted LoRA weights: {info}")
+
+  dir_name = os.path.dirname(args.save_to)
+  if dir_name and not os.path.exists(dir_name):
+    os.makedirs(dir_name, exist_ok=True)
+
+  lora_network_o.save_weights(args.save_to, save_dtype, {})
+  print(f"LoRA weights are saved to: {args.save_to}")
+
+
+if __name__ == '__main__':
+  parser = argparse.ArgumentParser()
+  parser.add_argument("--v2", action='store_true',
+                      help='load Stable Diffusion v2.x model / Stable Diffusion 2.xのモデルを読み込む')
+  parser.add_argument("--save_precision", type=str, default=None,
+                      choices=[None, "float", "fp16", "bf16"], help="precision in saving, same to merging if omitted / 保存時に精度を変更して保存する、省略時はfloat")
+  parser.add_argument("--model_org", type=str, default=None,
+                      help="Stable Diffusion original model: ckpt or safetensors file / 元モデル、ckptまたはsafetensors")
+  parser.add_argument("--model_tuned", type=str, default=None,
+                      help="Stable Diffusion tuned model, LoRA is difference of `original to tuned`: ckpt or safetensors file / 派生モデル（生成されるLoRAは元→派生の差分になります）、ckptまたはsafetensors")
+  parser.add_argument("--save_to", type=str, default=None,
+                      help="destination file name: ckpt or safetensors file / 保存先のファイル名、ckptまたはsafetensors")
+  parser.add_argument("--dim", type=int, default=4, help="dimension of LoRA (default 4) / LoRAの次元数（デフォルト4）")
+  parser.add_argument("--device", type=str, default=None, help="device to use, 'cuda' for GPU / 計算を行うデバイス、'cuda'でGPUを使う")
+
+  args = parser.parse_args()
+  svd(args)

networks/lora.py

@@ -0,0 +1,193 @@
+# LoRA network module
+# reference:
+# https://github.com/microsoft/LoRA/blob/main/loralib/layers.py
+# https://github.com/cloneofsimo/lora/blob/master/lora_diffusion/lora.py
+
+import math
+import os
+import torch
+
+
+class LoRAModule(torch.nn.Module):
+  """
+  replaces forward method of the original Linear, instead of replacing the original Linear module.
+  """
+
+  def __init__(self, lora_name, org_module: torch.nn.Module, multiplier=1.0, lora_dim=4):
+    super().__init__()
+    self.lora_name = lora_name
+
+    if org_module.__class__.__name__ == 'Conv2d':
+      in_dim = org_module.in_channels
+      out_dim = org_module.out_channels
+      self.lora_down = torch.nn.Conv2d(in_dim, lora_dim, (1, 1), bias=False)
+      self.lora_up = torch.nn.Conv2d(lora_dim, out_dim, (1, 1), bias=False)
+    else:
+      in_dim = org_module.in_features
+      out_dim = org_module.out_features
+      self.lora_down = torch.nn.Linear(in_dim, lora_dim, bias=False)
+      self.lora_up = torch.nn.Linear(lora_dim, out_dim, bias=False)
+
+    # same as microsoft's
+    torch.nn.init.kaiming_uniform_(self.lora_down.weight, a=math.sqrt(5))
+    torch.nn.init.zeros_(self.lora_up.weight)
+
+    self.multiplier = multiplier
+    self.org_module = org_module                  # remove in applying
+
+  def apply_to(self):
+    self.org_forward = self.org_module.forward
+    self.org_module.forward = self.forward
+    del self.org_module
+
+  def forward(self, x):
+    return self.org_forward(x) + self.lora_up(self.lora_down(x)) * self.multiplier
+
+
+def create_network(multiplier, network_dim, vae, text_encoder, unet, **kwargs):
+  if network_dim is None:
+    network_dim = 4                     # default
+  network = LoRANetwork(text_encoder, unet, multiplier=multiplier, lora_dim=network_dim)
+  return network
+
+
+class LoRANetwork(torch.nn.Module):
+  UNET_TARGET_REPLACE_MODULE = ["Transformer2DModel", "Attention"]
+  TEXT_ENCODER_TARGET_REPLACE_MODULE = ["CLIPAttention", "CLIPMLP"]
+  LORA_PREFIX_UNET = 'lora_unet'
+  LORA_PREFIX_TEXT_ENCODER = 'lora_te'
+
+  def __init__(self, text_encoder, unet, multiplier=1.0, lora_dim=4) -> None:
+    super().__init__()
+    self.multiplier = multiplier
+    self.lora_dim = lora_dim
+
+    # create module instances
+    def create_modules(prefix, root_module: torch.nn.Module, target_replace_modules) -> list[LoRAModule]:
+      loras = []
+      for name, module in root_module.named_modules():
+        if module.__class__.__name__ in target_replace_modules:
+          for child_name, child_module in module.named_modules():
+            if child_module.__class__.__name__ == "Linear" or (child_module.__class__.__name__ == "Conv2d" and child_module.kernel_size == (1, 1)):
+              lora_name = prefix + '.' + name + '.' + child_name
+              lora_name = lora_name.replace('.', '_')
+              lora = LoRAModule(lora_name, child_module, self.multiplier, self.lora_dim)
+              loras.append(lora)
+      return loras
+
+    self.text_encoder_loras = create_modules(LoRANetwork.LORA_PREFIX_TEXT_ENCODER,
+                                             text_encoder, LoRANetwork.TEXT_ENCODER_TARGET_REPLACE_MODULE)
+    print(f"create LoRA for Text Encoder: {len(self.text_encoder_loras)} modules.")
+
+    self.unet_loras = create_modules(LoRANetwork.LORA_PREFIX_UNET, unet, LoRANetwork.UNET_TARGET_REPLACE_MODULE)
+    print(f"create LoRA for U-Net: {len(self.unet_loras)} modules.")
+
+    self.weights_sd = None
+
+    # assertion
+    names = set()
+    for lora in self.text_encoder_loras + self.unet_loras:
+      assert lora.lora_name not in names, f"duplicated lora name: {lora.lora_name}"
+      names.add(lora.lora_name)
+
+  def load_weights(self, file):
+    if os.path.splitext(file)[1] == '.safetensors':
+      from safetensors.torch import load_file, safe_open
+      self.weights_sd = load_file(file)
+    else:
+      self.weights_sd = torch.load(file, map_location='cpu')
+
+  def apply_to(self, text_encoder, unet, apply_text_encoder=None, apply_unet=None):
+    if self.weights_sd:
+      weights_has_text_encoder = weights_has_unet = False
+      for key in self.weights_sd.keys():
+        if key.startswith(LoRANetwork.LORA_PREFIX_TEXT_ENCODER):
+          weights_has_text_encoder = True
+        elif key.startswith(LoRANetwork.LORA_PREFIX_UNET):
+          weights_has_unet = True
+
+      if apply_text_encoder is None:
+        apply_text_encoder = weights_has_text_encoder
+      else:
+        assert apply_text_encoder == weights_has_text_encoder, f"text encoder weights: {weights_has_text_encoder} but text encoder flag: {apply_text_encoder} / 重みとText Encoderのフラグが矛盾しています"
+
+      if apply_unet is None:
+        apply_unet = weights_has_unet
+      else:
+        assert apply_unet == weights_has_unet, f"u-net weights: {weights_has_unet} but u-net flag: {apply_unet} / 重みとU-Netのフラグが矛盾しています"
+    else:
+      assert apply_text_encoder is not None and apply_unet is not None, f"internal error: flag not set"
+
+    if apply_text_encoder:
+      print("enable LoRA for text encoder")
+    else:
+      self.text_encoder_loras = []
+
+    if apply_unet:
+      print("enable LoRA for U-Net")
+    else:
+      self.unet_loras = []
+
+    for lora in self.text_encoder_loras + self.unet_loras:
+      lora.apply_to()
+      self.add_module(lora.lora_name, lora)
+
+    if self.weights_sd:
+      # if some weights are not in state dict, it is ok because initial LoRA does nothing (lora_up is initialized by zeros)
+      info = self.load_state_dict(self.weights_sd, False)
+      print(f"weights are loaded: {info}")
+
+  def enable_gradient_checkpointing(self):
+    # not supported
+    pass
+
+  def prepare_optimizer_params(self, text_encoder_lr, unet_lr):
+    def enumerate_params(loras):
+      params = []
+      for lora in loras:
+        params.extend(lora.parameters())
+      return params
+
+    self.requires_grad_(True)
+    params = []
+
+    if self.text_encoder_loras:
+      param_data = {'params': enumerate_params(self.text_encoder_loras)}
+      if text_encoder_lr is not None:
+        param_data['lr'] = text_encoder_lr
+      params.append(param_data)
+
+    if self.unet_loras:
+      param_data = {'params': enumerate_params(self.unet_loras)}
+      if unet_lr is not None:
+        param_data['lr'] = unet_lr
+      params.append(param_data)
+
+    return params
+
+  def prepare_grad_etc(self, text_encoder, unet):
+    self.requires_grad_(True)
+
+  def on_epoch_start(self, text_encoder, unet):
+    self.train()
+
+  def get_trainable_params(self):
+    return self.parameters()
+
+  def save_weights(self, file, dtype, metadata):
+    if metadata is not None and len(metadata) == 0:
+      metadata = None
+
+    state_dict = self.state_dict()
+
+    if dtype is not None:
+      for key in list(state_dict.keys()):
+        v = state_dict[key]
+        v = v.detach().clone().to("cpu").to(dtype)
+        state_dict[key] = v
+
+    if os.path.splitext(file)[1] == '.safetensors':
+      from safetensors.torch import save_file
+      save_file(state_dict, file, metadata)
+    else:
+      torch.save(state_dict, file)

networks/merge_lora.py

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

requirements.txt

@@ -0,0 +1,23 @@
+accelerate==0.15.0
+transformers==4.25.1
+ftfy
+albumentations
+opencv-python
+einops
+diffusers[torch]==0.10.2
+pytorch_lightning
+bitsandbytes==0.35.0
+tensorboard
+safetensors==0.2.6
+gradio
+altair
+easygui
+# for BLIP captioning
+requests
+timm==0.4.12
+fairscale==0.4.4
+# for WD14 captioning
+tensorflow<2.11
+huggingface-hub
+# for kohya_ss library
+.

requirements_blip.txt

@@ -1,3 +0,0 @@
-timm==0.4.12
-transformers==4.16.2
-fairscale==0.4.4

requirements_db_finetune.txt

@@ -1,8 +0,0 @@
-accelerate==0.14.0
-transformers>=4.21.0
-ftfy
-albumentations
-opencv-python
-einops
-pytorch_lightning
-safetensors

requirements_wd14_tagger.txt

@@ -1,2 +0,0 @@
-tensorflow<2.11
-huggingface-hub

setup.py

@@ -0,0 +1,3 @@
+from setuptools import setup, find_packages
+ 
+setup(name = "library", packages = find_packages())

tools/convert_diffusers20_original_sd.py

@@ -9,7 +9,7 @@ import os
 import torch
 from diffusers import StableDiffusionPipeline
 
-import model_util
+import library.model_util as model_util
 
 
 def convert(args):
@@ -48,7 +48,7 @@ def convert(args):
       v2_model = unet.config.cross_attention_dim == 1024
       print("checking model version: model is " + ('v2' if v2_model else 'v1'))
     else:
-      v2_model = args.v1
+      v2_model = not args.v1
 
   # 変換して保存する
   msg = ("checkpoint" + ("" if save_dtype is None else f" in {save_dtype}")) if is_save_ckpt else "Diffusers"

tools/detect_face_rotate.py

@@ -5,7 +5,7 @@
 
 # v2: extract max face if multiple faces are found
 # v3: add crop_ratio option
-# v4: add multple faces extraction and min/max size
+# v4: add multiple faces extraction and min/max size
 
 import argparse
 import math

train_db.py

@@ -0,0 +1,337 @@
+# DreamBooth training
+# XXX dropped option: fine_tune
+
+import gc
+import time
+import argparse
+import itertools
+import math
+import os
+
+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
+from library.train_util import DreamBoothDataset
+
+
+def collate_fn(examples):
+  return examples[0]
+
+
+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)
+
+  train_dataset = DreamBoothDataset(args.train_batch_size, args.train_data_dir, args.reg_data_dir,
+                                    tokenizer, args.max_token_length, args.caption_extension, args.shuffle_caption, args.keep_tokens,
+                                    args.resolution, args.enable_bucket, args.min_bucket_reso, args.max_bucket_reso, args.prior_loss_weight,
+                                    args.flip_aug, args.color_aug, args.face_crop_aug_range, args.random_crop, args.debug_dataset)
+  if args.no_token_padding:
+    train_dataset.disable_token_padding()
+  train_dataset.make_buckets()
+
+  if args.debug_dataset:
+    train_util.debug_dataset(train_dataset)
+    return
+
+  # 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.cache_latents(vae)
+    vae.to("cpu")
+    if torch.cuda.is_available():
+      torch.cuda.empty_cache()
+    gc.collect()
+
+  # 学習を準備する：モデルを適切な状態にする
+  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.")
+
+  # 8-bit Adamを使う
+  if args.use_8bit_adam:
+    try:
+      import bitsandbytes as bnb
+    except ImportError:
+      raise ImportError("No bitsand bytes / bitsandbytesがインストールされていないようです")
+    print("use 8-bit Adam optimizer")
+    optimizer_class = bnb.optim.AdamW8bit
+  else:
+    optimizer_class = torch.optim.AdamW
+
+  if train_text_encoder:
+    trainable_params = (itertools.chain(unet.parameters(), text_encoder.parameters()))
+  else:
+    trainable_params = unet.parameters()
+
+  # betaやweight decayはdiffusers DreamBoothもDreamBooth SDもデフォルト値のようなのでオプションはとりあえず省略
+  optimizer = optimizer_class(trainable_params, lr=args.learning_rate)
+
+  # 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, batch_size=1, shuffle=False, collate_fn=collate_fn, num_workers=n_workers)
+
+  # 学習ステップ数を計算する
+  if args.max_train_epochs is not None:
+    args.max_train_steps = args.max_train_epochs * len(train_dataloader)
+    print(f"override steps. steps for {args.max_train_epochs} epochs is / 指定エポックまでのステップ数: {args.max_train_steps}")
+
+  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を用意する
+  lr_scheduler = diffusers.optimization.get_scheduler(
+      args.lr_scheduler, optimizer, num_warmup_steps=args.lr_warmup_steps, num_training_steps=args.max_train_steps)
+
+  # 実験的機能：勾配も含めた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する
+  if args.resume is not None:
+    print(f"resume training from state: {args.resume}")
+    accelerator.load_state(args.resume)
+
+  # 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)
+
+  # 学習する
+  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.num_train_images}")
+  print(f"  num reg images / 正則化画像の数: {train_dataset.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")
+
+  for epoch in range(num_train_epochs):
+    print(f"epoch {epoch+1}/{num_train_epochs}")
+
+    # 指定したステップ数まで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()
+
+    loss_total = 0
+    for step, batch in enumerate(train_dataloader):
+      # 指定したステップ数で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
+
+        # Sample noise that we'll add to the latents
+        noise = torch.randn_like(latents, device=latents.device)
+        b_size = latents.shape[0]
+
+        # Get the text embedding for conditioning
+        with torch.set_grad_enabled(global_step < args.stop_text_encoder_training):
+          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
+
+        loss = loss.mean()                # 平均なのでbatch_sizeで割る必要なし
+
+        accelerator.backward(loss)
+        if accelerator.sync_gradients:
+          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, 1.0)  # 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
+
+      current_loss = loss.detach().item()
+      if args.logging_dir is not None:
+        logs = {"loss": current_loss, "lr": lr_scheduler.get_last_lr()[0]}
+        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 = {"epoch_loss": 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)
+
+  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.")
+
+
+if __name__ == '__main__':
+  parser = argparse.ArgumentParser()
+
+  train_util.add_sd_models_arguments(parser)
+  train_util.add_dataset_arguments(parser, True, False)
+  train_util.add_training_arguments(parser, True)
+  train_util.add_sd_saving_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で最初から学習しない")
+
+  args = parser.parse_args()
+  train(args)

train_db_README-ja.md

@@ -0,0 +1,296 @@
+DreamBoothのガイドです。LoRA等の追加ネットワークの学習にも同じ手順を使います。
+
+# 概要
+
+スクリプトの主な機能は以下の通りです。
+
+- 8bit Adam optimizerおよびlatentのキャッシュによる省メモリ化（ShivamShrirao氏版と同様）。
+- xformersによる省メモリ化。
+- 512x512だけではなく任意サイズでの学習。
+- augmentationによる品質の向上。
+- DreamBoothだけではなくText Encoder+U-Netのfine tuningに対応。
+- StableDiffusion形式でのモデルの読み書き。
+- Aspect Ratio Bucketing。
+- Stable Diffusion v2.0対応。
+
+# 学習の手順
+
+## step 1. 環境整備
+
+このリポジトリのREADMEを参照してください。
+
+
+## step 2. identifierとclassを決める
+
+学ばせたい対象を結びつける単語identifierと、対象の属するclassを決めます。
+
+（instanceなどいろいろな呼び方がありますが、とりあえず元の論文に合わせます。）
+
+以下ごく簡単に説明します（詳しくは調べてください）。
+
+classは学習対象の一般的な種別です。たとえば特定の犬種を学ばせる場合には、classはdogになります。アニメキャラならモデルによりboyやgirl、1boyや1girlになるでしょう。
+
+identifierは学習対象を識別して学習するためのものです。任意の単語で構いませんが、元論文によると「tokinizerで1トークンになる3文字以下でレアな単語」が良いとのことです。
+
+identifierとclassを使い、たとえば「shs dog」などでモデルを学習することで、学習させたい対象をclassから識別して学習できます。
+
+画像生成時には「shs dog」とすれば学ばせた犬種の画像が生成されます。
+
+（identifierとして私が最近使っているものを参考までに挙げると、``shs sts scs cpc coc cic msm usu ici lvl cic dii muk ori hru rik koo yos wny`` などです。）
+
+## step 3. 学習用画像の準備
+学習用画像を格納するフォルダを作成します。 __さらにその中に__ 、以下の名前でディレクトリを作成します。
+
+```
+<繰り返し回数>_<identifier> <class>
+```
+
+間の``_``を忘れないでください。
+
+繰り返し回数は、正則化画像と枚数を合わせるために指定します（後述します）。
+
+たとえば「sls frog」というプロンプトで、データを20回繰り返す場合、「20_sls frog」となります。以下のようになります。
+
+![image](https://user-images.githubusercontent.com/52813779/210770636-1c851377-5936-4c15-90b7-8ac8ad6c2074.png)
+
+## step 4. 正則化画像の準備
+正則化画像を使う場合の手順です。使わずに学習することもできます（正則化画像を使わないと区別ができなくなるので対象class全体が影響を受けます）。
+
+正則化画像を格納するフォルダを作成します。 __さらにその中に__  ``<繰り返し回数>_<class>`` という名前でディレクトリを作成します。
+
+たとえば「frog」というプロンプトで、データを繰り返さない（1回だけ）場合、以下のようになります。
+
+![image](https://user-images.githubusercontent.com/52813779/210770897-329758e5-3675-49f1-b345-c135f1725832.png)
+
+繰り返し回数は「 __学習用画像の繰り返し回数×学習用画像の枚数≧正則化画像の繰り返し回数×正則化画像の枚数__ 」となるように指定してください。
+
+（1 epochのデータ数が「学習用画像の繰り返し回数×学習用画像の枚数」となります。正則化画像の枚数がそれより多いと、余った部分の正則化画像は使用されません。）
+
+## step 5. 学習の実行
+スクリプトを実行します。最大限、メモリを節約したコマンドは以下のようになります（実際には1行で入力します）。
+
+※LoRA等の追加ネットワークを学習する場合のコマンドは ``train_db.py`` ではなく ``train_network.py`` となります。また追加でnetwork_\*オプションが必要となりますので、LoRAのガイドを参照してください。
+
+```
+accelerate launch --num_cpu_threads_per_process 8 train_db.py 
+    --pretrained_model_name_or_path=<.ckptまたは.safetensordまたはDiffusers版モデルのディレクトリ> 
+    --train_data_dir=<学習用データのディレクトリ> 
+    --reg_data_dir=<正則化画像のディレクトリ> 
+    --output_dir=<学習したモデルの出力先ディレクトリ> 
+    --prior_loss_weight=1.0 
+    --resolution=512 
+    --train_batch_size=1 
+    --learning_rate=1e-6 
+    --max_train_steps=1600 
+    --use_8bit_adam 
+    --xformers 
+    --mixed_precision="bf16" 
+    --cache_latents
+    --gradient_checkpointing
+```
+
+num_cpu_threads_per_processにはCPUコア数を指定するとよいようです。
+
+pretrained_model_name_or_pathに追加学習を行う元となるモデルを指定します。Stable Diffusionのcheckpointファイル（.ckptまたは.safetensors）、Diffusersのローカルディスクにあるモデルディレクトリ、DiffusersのモデルID（"stabilityai/stable-diffusion-2"など）が指定できます。学習後のモデルの保存形式はデフォルトでは元のモデルと同じになります（save_model_asオプションで変更できます）。
+
+prior_loss_weightは正則化画像のlossの重みです。通常は1.0を指定します。
+
+resolutionは画像のサイズ（解像度、幅と高さ）になります。bucketing（後述）を用いない場合、学習用画像、正則化画像はこのサイズとしてください。
+
+train_batch_sizeは学習時のバッチサイズです。max_train_stepsを1600とします。学習率learning_rateは、diffusers版では5e-6ですがStableDiffusion版は1e-6ですのでここでは1e-6を指定しています。
+
+省メモリ化のためmixed_precision="bf16"（または"fp16"）、およびgradient_checkpointing を指定します。
+
+xformersオプションを指定し、xformersのCrossAttentionを用います。xformersをインストールしていない場合、エラーとなる場合（mixed_precisionなしの場合、私の環境ではエラーとなりました）、代わりにmem_eff_attnオプションを指定すると省メモリ版CrossAttentionを使用します（速度は遅くなります）。
+
+省メモリ化のためcache_latentsオプションを指定してVAEの出力をキャッシュします。
+
+ある程度メモリがある場合はたとえば以下のように指定します。
+
+```
+accelerate launch --num_cpu_threads_per_process 8 train_db.py 
+    --pretrained_model_name_or_path=<.ckptまたは.safetensordまたはDiffusers版モデルのディレクトリ> 
+    --train_data_dir=<学習用データのディレクトリ> 
+    --reg_data_dir=<正則化画像のディレクトリ> 
+    --output_dir=<学習したモデルの出力先ディレクトリ> 
+    --prior_loss_weight=1.0 
+    --resolution=512 
+    --train_batch_size=4 
+    --learning_rate=1e-6 
+    --max_train_steps=400 
+    --use_8bit_adam 
+    --xformers 
+    --mixed_precision="bf16" 
+    --cache_latents
+```
+
+gradient_checkpointingを外し高速化します（メモリ使用量は増えます）。バッチサイズを増やし、高速化と精度向上を図ります。
+
+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_dreambooth.pyの半分になっています（対象の画像と正則化画像を同一のバッチではなく別のバッチに分割して学習するため）。
+元のDiffusers版やXavierXiao氏のStableDiffusion版とほぼ同じ学習を行うには、ステップ数を倍にしてください。
+
+（shuffle=Trueのため厳密にはデータの順番が変わってしまいますが、学習には大きな影響はないと思います。）
+
+## 学習したモデルで画像生成する
+
+学習が終わると指定したフォルダにlast.ckptという名前でcheckpointが出力されます（DiffUsers版モデルを学習した場合はlastフォルダになります）。
+
+v1.4/1.5およびその他の派生モデルの場合、このモデルでAutomatic1111氏のWebUIなどで推論できます。models\Stable-diffusionフォルダに置いてください。
+
+v2.xモデルでWebUIで画像生成する場合、モデルの仕様が記述された.yamlファイルが別途必要になります。v2.x baseの場合はv2-inference.yamlを、768/vの場合はv2-inference-v.yamlを、同じフォルダに置き、拡張子の前の部分をモデルと同じ名前にしてください。
+
+![image](https://user-images.githubusercontent.com/52813779/210776915-061d79c3-6582-42c2-8884-8b91d2f07313.png)
+
+各yamlファイルは[https://github.com/Stability-AI/stablediffusion/tree/main/configs/stable-diffusion](Stability AIのSD2.0のリポジトリ)にあります。
+
+# その他の学習オプション
+
+## Stable Diffusion 2.0対応 --v2 / --v_parameterization
+Hugging Faceのstable-diffusion-2-baseを使う場合はv2オプションを、stable-diffusion-2または768-v-ema.ckptを使う場合はv2とv_parameterizationの両方のオプションを指定してください。
+
+なおSD 2.0の学習はText Encoderが大きくなっているためVRAM 12GBでは厳しいようです。
+
+Stable Diffusion 2.0では大きく以下の点が変わっています。
+
+1. 使用するTokenizer
+2. 使用するText Encoderおよび使用する出力層（2.0は最後から二番目の層を使う）
+3. Text Encoderの出力次元数（768->1024）
+4. U-Netの構造（CrossAttentionのhead数など）
+5. v-parameterization（サンプリング方法が変更されているらしい）
+
+このうちbaseでは1～4が、baseのつかない方（768-v）では1～5が採用されています。1～4を有効にするのがv2オプション、5を有効にするのがv_parameterizationオプションです。
+
+## 学習データの確認 --debug_dataset
+このオプションを付けることで学習を行う前に事前にどのような画像データ、キャプションで学習されるかを確認できます。Escキーを押すと終了してコマンドラインに戻ります。
+
+※Colabなど画面が存在しない環境で実行するとハングするようですのでご注意ください。
+
+## Text Encoderの学習を途中から行わない --stop_text_encoder_training
+stop_text_encoder_trainingオプションに数値を指定すると、そのステップ数以降はText Encoderの学習を行わずU-Netだけ学習します。場合によっては精度の向上が期待できるかもしれません。
+
+（恐らくText Encoderだけ先に過学習することがあり、それを防げるのではないかと推測していますが、詳細な影響は不明です。）
+
+## VAEを別途読み込んで学習する --vae
+vaeオプションにStable Diffusionのcheckpoint、VAEのcheckpointファイル、DiffusesのモデルまたはVAE（ともにローカルまたはHugging FaceのモデルIDが指定できます）のいずれかを指定すると、そのVAEを使って学習します（latentsのキャッシュ時または学習中のlatents取得時）。
+保存されるモデルはこのVAEを組み込んだものになります。
+
+## 学習途中での保存 --save_every_n_epochs / --save_state / --resume
+save_every_n_epochsオプションに数値を指定すると、そのエポックごとに学習途中のモデルを保存します。
+
+save_stateオプションを同時に指定すると、optimizer等の状態も含めた学習状態を合わせて保存します（checkpointから学習再開するのに比べて、精度の向上、学習時間の短縮が期待できます）。学習状態は保存先フォルダに"epoch-??????-state"（??????はエポック数）という名前のフォルダで出力されます。長時間にわたる学習時にご利用ください。
+
+保存された学習状態から学習を再開するにはresumeオプションを使います。学習状態のフォルダを指定してください。
+
+なおAcceleratorの仕様により(?)、エポック数、global stepは保存されておらず、resumeしたときにも1からになりますがご容赦ください。
+
+## Tokenizerのパディングをしない --no_token_padding
+no_token_paddingオプションを指定するとTokenizerの出力をpaddingしません（Diffusers版の旧DreamBoothと同じ動きになります）。
+
+## 任意サイズの画像での学習 --resolution
+正方形以外で学習できます。resolutionに「448,640」のように「幅,高さ」で指定してください。幅と高さは64で割り切れる必要があります。学習用画像、正則化画像のサイズを合わせてください。
+
+個人的には縦長の画像を生成することが多いため「448,640」などで学習することもあります。
+
+## Aspect Ratio Bucketing --enable_bucket / --min_bucket_reso / --max_bucket_reso
+enable_bucketオプションを指定すると有効になります。Stable Diffusionは512x512で学習されていますが、それに加えて256x768や384x640といった解像度でも学習します。
+
+このオプションを指定した場合は、学習用画像、正則化画像を特定の解像度に統一する必要はありません。いくつかの解像度（アスペクト比）から最適なものを選び、その解像度で学習します。
+解像度は64ピクセル単位のため、元画像とアスペクト比が完全に一致しない場合がありますが、その場合は、はみ出した部分がわずかにトリミングされます。
+
+解像度の最小サイズをmin_bucket_resoオプションで、最大サイズをmax_bucket_resoで指定できます。デフォルトはそれぞれ256、1024です。
+たとえば最小サイズに384を指定すると、256x1024や320x768などの解像度は使わなくなります。
+解像度を768x768のように大きくした場合、最大サイズに1280などを指定しても良いかもしれません。
+
+なおAspect Ratio Bucketingを有効にするときには、正則化画像についても、学習用画像と似た傾向の様々な解像度を用意した方がいいかもしれません。
+
+（ひとつのバッチ内の画像が学習用画像、正則化画像に偏らなくなるため。そこまで大きな影響はないと思いますが……。）
+
+## augmentation --color_aug / --flip_aug
+augmentationは学習時に動的にデータを変化させることで、モデルの性能を上げる手法です。color_augで色合いを微妙に変えつつ、flip_augで左右反転をしつつ、学習します。
+
+動的にデータを変化させるため、cache_latentsオプションと同時に指定できません。
+
+## 保存時のデータ精度の指定 --save_precision
+save_precisionオプションにfloat、fp16、bf16のいずれかを指定すると、その形式でcheckpointを保存します（Stable Diffusion形式で保存する場合のみ）。checkpointのサイズを削減したい場合などにお使いください。
+
+## 任意の形式で保存する --save_model_as
+モデルの保存形式を指定します。ckpt、safetensors、diffusers、diffusers_safetensorsのいずれかを指定してください。
+
+Stable Diffusion形式（ckptまたはsafetensors）を読み込み、Diffusers形式で保存する場合、不足する情報はHugging Faceからv1.5またはv2.1の情報を落としてきて補完します。
+
+## 学習ログの保存 --logging_dir / --log_prefix
+logging_dirオプションにログ保存先フォルダを指定してください。TensorBoard形式のログが保存されます。
+
+たとえば--logging_dir=logsと指定すると、作業フォルダにlogsフォルダが作成され、その中の日時フォルダにログが保存されます。
+また--log_prefixオプションを指定すると、日時の前に指定した文字列が追加されます。「--logging_dir=logs --log_prefix=db_style1_」などとして識別用にお使いください。
+
+TensorBoardでログを確認するには、別のコマンドプロンプトを開き、作業フォルダで以下のように入力します（tensorboardはDiffusersのインストール時にあわせてインストールされると思いますが、もし入っていないならpip install tensorboardで入れてください）。
+
+```
+tensorboard --logdir=logs
+```
+
+その後ブラウザを開き、http://localhost:6006/ へアクセスすると表示されます。
+
+## 学習率のスケジューラ関連の指定 --lr_scheduler / --lr_warmup_steps
+lr_schedulerオプションで学習率のスケジューラをlinear, cosine, cosine_with_restarts, polynomial, constant, constant_with_warmupから選べます。デフォルトはconstantです。lr_warmup_stepsでスケジューラのウォームアップ（だんだん学習率を変えていく）ステップ数を指定できます。詳細については各自お調べください。
+
+## 勾配をfp16とした学習（実験的機能） --full_fp16
+full_fp16オプションを指定すると勾配を通常のfloat32からfloat16（fp16）に変更して学習します（mixed precisionではなく完全なfp16学習になるようです）。
+これによりSD1.xの512x512サイズでは8GB未満、SD2.xの512x512サイズで12GB未満のVRAM使用量で学習できるようです。
+
+あらかじめaccelerate configでfp16を指定し、オプションで ``mixed_precision="fp16"`` としてください（bf16では動作しません）。
+
+メモリ使用量を最小化するためには、xformers、use_8bit_adam、cache_latents、gradient_checkpointingの各オプションを指定し、train_batch_sizeを1としてください。
+
+（余裕があるようならtrain_batch_sizeを段階的に増やすと若干精度が上がるはずです。）
+
+PyTorchのソースにパッチを当てて無理やり実現しています（PyTorch 1.12.1と1.13.0で確認）。精度はかなり落ちますし、途中で学習失敗する確率も高くなります。
+学習率やステップ数の設定もシビアなようです。それらを認識したうえで自己責任でお使いください。
+
+# その他の学習方法
+
+## 複数class、複数対象（identifier）の学習
+方法は単純で、学習用画像のフォルダ内に ``繰り返し回数_<identifier> <class>`` のフォルダを複数、正則化画像フォルダにも同様に ``繰り返し回数_<class>`` のフォルダを複数、用意してください。
+
+たとえば「sls frog」と「cpc rabbit」を同時に学習する場合、以下のようになります。
+
+![image](https://user-images.githubusercontent.com/52813779/210777933-a22229db-b219-4cd8-83ca-e87320fc4192.png)
+
+classがひとつで対象が複数の場合、正則化画像フォルダはひとつで構いません。たとえば1girlにキャラAとキャラBがいる場合は次のようにします。
+
+- train_girls
+  - 10_sls 1girl
+  - 10_cpc 1girl
+- reg_girls
+  - 1_1girl
+
+データ数にばらつきがある場合、繰り返し回数を調整してclass、identifierごとの枚数を統一すると良い結果が得られることがあるようです。
+
+## DreamBoothでキャプションを使う
+学習用画像、正則化画像のフォルダに、画像と同じファイル名で、拡張子.caption（オプションで変えられます）のファイルを置くと、そのファイルからキャプションを読み込みプロンプトとして学習します。
+
+※それらの画像の学習に、フォルダ名（identifier class）は使用されなくなります。
+
+各画像にキャプションを付けることで（BLIP等を使っても良いでしょう）、学習したい属性をより明確にできるかもしれません。
+
+キャプションファイルの拡張子はデフォルトで.captionです。--caption_extensionで変更できます。--shuffle_captionオプションで学習時のキャプションについて、カンマ区切りの各部分をシャッフルしながら学習します。
+

train_network.py

@@ -0,0 +1,427 @@
+import importlib
+import argparse
+import gc
+import math
+import os
+
+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
+from library.train_util import DreamBoothDataset, FineTuningDataset
+
+
+def collate_fn(examples):
+  return examples[0]
+
+
+def train(args):
+  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
+
+  if args.seed is not None:
+    set_seed(args.seed)
+
+  tokenizer = train_util.load_tokenizer(args)
+
+  # データセットを準備する
+  if use_dreambooth_method:
+    print("Use DreamBooth method.")
+    train_dataset = DreamBoothDataset(args.train_batch_size, args.train_data_dir, args.reg_data_dir,
+                                      tokenizer, args.max_token_length, args.caption_extension, args.shuffle_caption, args.keep_tokens,
+                                      args.resolution, args.enable_bucket, args.min_bucket_reso, args.max_bucket_reso, args.prior_loss_weight,
+                                      args.flip_aug, args.color_aug, args.face_crop_aug_range, args.random_crop, args.debug_dataset)
+  else:
+    print("Train with captions.")
+    train_dataset = FineTuningDataset(args.in_json, args.train_batch_size, args.train_data_dir,
+                                      tokenizer, args.max_token_length, args.shuffle_caption, args.keep_tokens,
+                                      args.resolution, args.enable_bucket, args.min_bucket_reso, args.max_bucket_reso,
+                                      args.flip_aug, args.color_aug, args.face_crop_aug_range, args.random_crop,
+                                      args.dataset_repeats, args.debug_dataset)
+  train_dataset.make_buckets()
+
+  if args.debug_dataset:
+    train_util.debug_dataset(train_dataset)
+    return
+  if len(train_dataset) == 0:
+    print("No data found. Please verify arguments / 画像がありません。引数指定を確認してください")
+    return
+
+  # 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)
+
+  # モデルに 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.cache_latents(vae)
+    vae.to("cpu")
+    if torch.cuda.is_available():
+      torch.cuda.empty_cache()
+    gc.collect()
+
+  # prepare network
+  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
+
+  network = network_module.create_network(1.0, args.network_dim, vae, text_encoder, unet, **net_kwargs)
+  if network is None:
+    return
+
+  if args.network_weights is not None:
+    print("load network weights from:", args.network_weights)
+    network.load_weights(args.network_weights)
+
+  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.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.")
+
+  # 8-bit Adamを使う
+  if args.use_8bit_adam:
+    try:
+      import bitsandbytes as bnb
+    except ImportError:
+      raise ImportError("No bitsand bytes / bitsandbytesがインストールされていないようです")
+    print("use 8-bit Adam optimizer")
+    optimizer_class = bnb.optim.AdamW8bit
+  else:
+    optimizer_class = torch.optim.AdamW
+
+  trainable_params = network.prepare_optimizer_params(args.text_encoder_lr, args.unet_lr)
+
+  # betaやweight decayはdiffusers DreamBoothもDreamBooth SDもデフォルト値のようなのでオプションはとりあえず省略
+  optimizer = optimizer_class(trainable_params, lr=args.learning_rate)
+
+  # 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, batch_size=1, shuffle=False, collate_fn=collate_fn, num_workers=n_workers)
+
+  # 学習ステップ数を計算する
+  if args.max_train_epochs is not None:
+    args.max_train_steps = args.max_train_epochs * len(train_dataloader)
+    print(f"override steps. steps for {args.max_train_epochs} epochs is / 指定エポックまでのステップ数: {args.max_train_steps}")
+
+  # lr schedulerを用意する
+  lr_scheduler = diffusers.optimization.get_scheduler(
+      args.lr_scheduler, optimizer, num_warmup_steps=args.lr_warmup_steps, num_training_steps=args.max_train_steps * args.gradient_accumulation_steps)
+
+  # 実験的機能：勾配も含めた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, dtype=weight_dtype)
+  if args.gradient_checkpointing:                       # according to TI example in Diffusers, train is required
+    unet.train()
+    text_encoder.train()
+
+    # set top parameter requires_grad = True for gradient checkpointing works
+    text_encoder.text_model.embeddings.requires_grad_(True)
+  else:
+    unet.eval()
+    text_encoder.eval()
+
+  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する
+  if args.resume is not None:
+    print(f"resume training from state: {args.resume}")
+    accelerator.load_state(args.resume)
+
+  # 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)
+
+  # 学習する
+  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.num_train_images}")
+  print(f"  num reg images / 正則化画像の数: {train_dataset.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 accumulation steps / 勾配を合計するステップ数 = {args.gradient_accumulation_steps}")
+  print(f"  total optimization steps / 学習ステップ数: {args.max_train_steps}")
+
+  metadata = {
+      "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.num_train_images,          # includes repeating TODO more detailed data
+      "ss_num_reg_images": train_dataset.num_reg_images,
+      "ss_num_batches_per_epoch": len(train_dataloader),
+      "ss_num_epochs": num_train_epochs,
+      "ss_batch_size_per_device": args.train_batch_size,
+      "ss_total_batch_size": total_batch_size,
+      "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_mixed_precision": args.mixed_precision,
+      "ss_full_fp16": bool(args.full_fp16),
+      "ss_v2": bool(args.v2),
+      "ss_resolution": args.resolution,
+      "ss_clip_skip": args.clip_skip,
+      "ss_max_token_length": args.max_token_length,
+      "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_cache_latents": bool(args.cache_latents),
+      "ss_enable_bucket": bool(train_dataset.enable_bucket),        # TODO move to BaseDataset from DB/FT
+      "ss_min_bucket_reso": args.min_bucket_reso,                   # TODO get from dataset
+      "ss_max_bucket_reso": args.max_bucket_reso,
+      "ss_seed": args.seed
+  }
+
+  # uncomment if another network is added
+  # for key, value in net_kwargs.items():
+  #   metadata["ss_arg_" + key] = value
+
+  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)
+      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)
+      vae_name = os.path.basename(vae_name)
+    metadata["ss_vae_name"] = vae_name
+
+  metadata = {k: str(v) for k, v in metadata.items()}
+
+  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")
+
+  for epoch in range(num_train_epochs):
+    print(f"epoch {epoch+1}/{num_train_epochs}")
+    metadata["ss_epoch"] = str(epoch+1)
+
+    network.on_epoch_start(text_encoder, unet)
+
+    loss_total = 0
+    for step, batch in enumerate(train_dataloader):
+      with accelerator.accumulate(network):
+        with torch.no_grad():
+          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
+          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)
+
+        # 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
+
+        loss = loss.mean()                # 平均なのでbatch_sizeで割る必要なし
+
+        accelerator.backward(loss)
+        if accelerator.sync_gradients:
+          params_to_clip = network.get_trainable_params()
+          accelerator.clip_grad_norm_(params_to_clip, 1.0)  # 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
+
+      current_loss = loss.detach().item()
+      if args.logging_dir is not None:
+        logs = {"loss": current_loss, "lr": lr_scheduler.get_last_lr()[0]}
+        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 = {"epoch_loss": loss_total / len(train_dataloader)}
+      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)
+        print(f"saving checkpoint: {ckpt_file}")
+        unwrap_model(network).save_weights(ckpt_file, save_dtype, None if args.no_metadata else metadata)
+
+      def remove_old_func(old_epoch_no):
+        old_ckpt_name = train_util.EPOCH_FILE_NAME.format(model_name, old_epoch_no) + '.' + args.save_model_as
+        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)
+
+    # end of epoch
+
+  metadata["ss_epoch"] = str(num_train_epochs)
+
+  is_main_process = accelerator.is_main_process
+  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, None if args.no_metadata else metadata)
+    print("model saved.")
+
+
+if __name__ == '__main__':
+  parser = argparse.ArgumentParser()
+
+  train_util.add_sd_models_arguments(parser)
+  train_util.add_dataset_arguments(parser, True, True)
+  train_util.add_training_arguments(parser, True)
+
+  parser.add_argument("--no_metadata", action='store_true', help="do not save metadata in output model / メタデータを出力先モデルに保存しない")
+  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("--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_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関連部分のみ学習する")
+
+  args = parser.parse_args()
+  train(args)

train_network_README-ja.md

@@ -0,0 +1,189 @@
+## 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)を大いに参考にさせていただきました。ありがとうございます。
+
+8GB VRAMでもぎりぎり動作するようです。
+
+## 学習したモデルに関する注意
+
+cloneofsimo氏のリポジトリ、およびd8ahazard氏の[Dreambooth Extension for Stable-Diffusion-WebUI](https://github.com/d8ahazard/sd_dreambooth_extension)とは、現時点では互換性がありません。いくつかの機能拡張を行っているためです（後述）。
+
+WebUI等で画像生成する場合には、学習したLoRAのモデルを学習元のStable Diffusionのモデルにこのリポジトリ内のスクリプトであらかじめマージしておくか、こちらの[WebUI用extention](https://github.com/kohya-ss/sd-webui-additional-networks)を使ってください。
+
+## 学習方法
+
+train_network.pyを用います。
+
+DreamBoothの手法（identifier（sksなど）とclass、オプションで正則化画像を用いる）と、キャプションを用いるfine tuningの手法の両方で学習できます。
+
+どちらの方法も既存のスクリプトとほぼ同じ方法で学習できます。異なる点については後述します。
+
+### DreamBoothの手法を用いる場合
+
+[DreamBoothのガイド](./train_db_README-ja.md) を参照してデータを用意してください。
+
+学習するとき、train_db.pyの代わりにtrain_network.pyを指定してください。
+
+ほぼすべてのオプション（Stable Diffusionのモデル保存関係を除く）が使えますが、stop_text_encoder_trainingはサポートしていません。
+
+### キャプションを用いる場合
+
+[fine-tuningのガイド](./fine_tune_README_ja.md) を参照し、各手順を実行してください。
+
+学習するとき、fine_tune.pyの代わりにtrain_network.pyを指定してください。ほぼすべてのオプション（モデル保存関係を除く）がそのまま使えます。
+
+なお「latentsの事前取得」は行わなくても動作します。VAEから学習時（またはキャッシュ時）にlatentを取得するため学習速度は遅くなりますが、代わりにcolor_augが使えるようになります。
+
+### LoRAの学習のためのオプション
+
+train_network.pyでは--network_moduleオプションに、学習対象のモジュール名を指定します。LoRAに対応するのはnetwork.loraとなりますので、それを指定してください。
+
+なお学習率は通常のDreamBoothやfine tuningよりも高めの、1e-4程度を指定するとよいようです。
+
+以下はコマンドラインの例です（DreamBooth手法）。
+
+```
+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
+```
+
+--output_dirオプションで指定したディレクトリに、LoRAのモデルが保存されます。
+
+その他、以下のオプションが指定できます。
+
+* --network_dim
+  * LoRAの次元数を指定します（``--networkdim=4``など）。省略時は4になります。数が多いほど表現力は増しますが、学習に必要なメモリ、時間は増えます。また闇雲に増やしても良くないようです。
+* --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_train_unet_onlyと--network_train_text_encoder_onlyの両方とも未指定時（デフォルト）はText EncoderとU-Netの両方の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、次元数の異なるLoRAはマージできません。U-NetだけのLoRAとU-Net+Text EncoderのLoRAはマージできるはずですが、結果は未知数です。
+
+
+### その他のオプション
+
+* precision
+  * マージ計算時の精度をfloat、fp16、bf16から指定できます。省略時は精度を確保するためfloatになります。メモリ使用量を減らしたい場合はfp16/bf16を指定してください。
+* save_precision
+  * モデル保存時の精度をfloat、fp16、bf16から指定できます。省略時はprecisionと同じ精度になります。
+
+## 当リポジトリ内の画像生成スクリプトで生成する
+
+gen_img_diffusers.pyに、--network_module、--network_weights、--network_dim（省略可）の各オプションを追加してください。意味は学習時と同様です。
+
+--network_mulオプションで0~1.0の数値を指定すると、LoRAの適用率を変えられます。
+
+## 二つのモデルの差分から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になります。
+
+## 追加情報
+
+### cloneofsimo氏のリポジトリとの違い
+
+12/25時点では、当リポジトリはLoRAの適用個所をText EncoderのMLP、U-NetのFFN、Transformerのin/out projectionに拡大し、表現力が増しています。ただその代わりメモリ使用量は増え、8GBぎりぎりになりました。
+
+またモジュール入れ替え機構は全く異なります。
+
+### 将来拡張について
+
+LoRAだけでなく他の拡張にも対応可能ですので、それらも追加予定です。