Repository: KohakuBlueleaf/LyCORIS (language: python)

Issue:
Fix: use float for constrain
This PR fixes constrain being rounded to an int (my fault), which would be a problem if you want a very small constraint as discussed in the paper.

This PR also changes constrain to be ε * _d_, where _d == out_dim_ of the module (pg. 6, Figure 5). In cases where constrain is not specified in args, constrain defaults to 0 (int) so it should still be OFT.


Discussion / hints:

Fix patch:
diff --git a/.gitignore b/.gitignore
index 931a8e5..3312b69 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,11 +1,17 @@
+# example
+data/
+
 # test
 test_model/
 huggingface/
+logs/
+output/
+exps/
+*.ps1
 *.safetensors
 *.pt
 *.ckpt
-logs/
-output/
+test*.py
 
 # Editors
 .vscode/
diff --git a/Change.md b/Change.md
index 5e18fd3..03d7cde 100644
--- a/Change.md
+++ b/Change.md
@@ -1,5 +1,25 @@
 # Change Log
 
+## 2023/12/02 update to 2.0.0
+* Start supporting [HCP-Diffusion](https://github.com/IrisRainbowNeko/HCP-Diffusion) (The reason to name this version "2.0.0")
+  * Now LyCORIS support LoHa/LoKr/Diag-OFT algorithm in HCP-Diffusion
+  * Add Pivotal tuning utilities
+  * Add hcp convert utilities
+  * Have no plan at this time to support full/lora and train_norms since HCP can do them natively
+* Add Diag-OFT modules
+* Add standalone usage support
+  * Can wrap any pytorch module which contains Linear/Conv2d/LayerNorm/GroupNorm modules
+  * Will support more module in the future
+* Add SDXL support in Merge script
+* Add SDXL support in Extract-locon
+* More efficient (speed/vram) implementation for full module
+* Better implementation of custom state_dict
+* Fix errors of dropouts
+* Fix errors of apply_max_norms
+* Fix errors of resume
+
+---
+
 ## 2023/09/27 update to 1.9.0
 * Add norm modules (for training LayerNorm and GroupNorm, which should be good for style)
 * Add full modules (So you can "native finetune" with lycoris now, should be convinient to try different weight)
diff --git a/README.md b/README.md
index dc561d7..af12c19 100644
--- a/README.md
+++ b/README.md
@@ -59,58 +59,93 @@ However, newer model types may not always be supported. If you encounter this is
 
 ### Training
 
-For the time being LyCORIS is mainly trained with [kohya-ss/sd-scripts](https://github.com/kohya-ss/sd-scripts) (see a list of compatible graphical interfaces and colabs at the end of the section). Supports for other trainers are coming soon.
+There are three different ways to train LyCORIS models.
 
-A detilaed description of the network arguments is provided [here](docs/Network-Args.md).
+- With [kohya-ss/sd-scripts](https://github.com/kohya-ss/sd-scripts) (see a list of compatible graphical interfaces and colabs at the end of the section)
+- With [HCP-Diffusion](https://github.com/IrisRainbowNeko/HCP-Diffusion)
+- With your own script by using LyCORIS as standalone wrappers for any pytorch modules.
+
+In any case, please install this package in the corresponding virtual environment. You can either install it
+- through pip
+    ```bash
+    pip install lycoris_lora
+    ```
+
+- or from source
+    ```bash
+    git clone https://github.com/KohakuBlueleaf/LyCORIS
+    cd LyCORIS
+    pip install .
+    ```
+
+A detilaed description of the network arguments is provided in [docs/Network-Args.md](docs/Network-Args.md).
 
 #### kohya script
 
-1. Activate sd-scripts' venv with
+You can use this package's kohya module to run kohya's training script to train lycoris module for SD models
+
+- with command line arguments
     ```bash
-    source PATH_TO_SDSCRIPTS_VENV/Scripts/activate
+    accelerate launch train_network.py \
+      --network_module lycoris.kohya \
+      --network_dim "DIM_FOR_LINEAR" --network_alpha "ALPHA_FOR_LINEAR"\
+      --network_args "conv_dim=DIM_FOR_CONV" "conv_alpha=ALPHA_FOR_CONV" \
+      "dropout=DROPOUT_RATE" "algo=locon" \
     ```
-    or
-    ```powershell
-    PATH_TO_SDSCRIPTS_VENV\Scripts\Activate.ps1 # or .bat     for cmd
+
+- with `toml` files
+    ```bash
+    accelerate launch train_network.py \
+      --config_file example_configs/training_configs/kohya/loha_config.toml
     ```
+    For your convenience, some example `toml` files for kohya LyCORIS training are provided in [example/training_configs/kohya](example_configs/training_configs/kohya).
 
-2. Install this package
-    * through pip
-        ```bash
-        pip install lycoris_lora
-        ```
 
-    * from source
-        ```bash
-        git clone https://github.com/KohakuBlueleaf/LyCORIS
-        cd LyCORIS
-        pip install .
-        ````
+#### HCP-Diffusion
 
-3. Use this package's kohya module to run kohya's training script to train lycoris module for SD models
+You can use this package's hcp module to run HCP-Diffusion's training script to train lycoris module for SD models
 
-    - with command line arguments
-        ```bash
-        python3 sd-scripts/train_network.py \
-          --network_module lycoris.kohya \
-          --network_dim "DIM_FOR_LINEAR" --network_alpha "ALPHA_FOR_LINEAR"\
-          --network_args "conv_dim=DIM_FOR_CONV" "conv_alpha=ALPHA_FOR_CONV" \
-          "dropout=DROPOUT_RATE" "algo=locon" \
-        ```
+```bash
+accelerate launch -m hcpdiff.train_ac_single \
+  --cfg example_configs/training_configs/hcp/hcp_diag_oft.yaml
+```
+For your convenience, some example `yaml` files for HCP LyCORIS training are provided in [example/training_configs/hcp](example_configs/training_configs/hcp).
+
+For the moment being the outputs of HCP-Diffusion are not directly compatible with a1111/sdwebui.
+You can perform conversion with [tools/batch_hcp_convert.py](tools/batch_hcp_convert.py).
 
-    - with `toml` files
-        ```bash
-        python train_network.py --config_file XXX.toml
-        ```
-        For your convenience, some example `toml` files for LyCORIS training are provided in [example/training_configs](example_configs/training_configs).
+In the case of pivotal tuning, [tools/batch_bundle_convert.py](tools/batch_bundle_convert.py) can be further used to convert to and from bundle formats.
 
 
-* Tips:
-  * Use network_dim=0 or conv_dim=0 to disable linear/conv layer
-  * LoHa/LoKr/(IA)^3 doesn't support dropout yet.
+#### As standalone wrappers
+See `standalone_example.py` for full example.
 
+Import `create_lycoris` and `LycorisNetwork` from `lycoris` library, put your preset to `LycorisNetwork` and then use `create_lycoris` to create LyCORIS module for your pytorch module.
 
-#### Graphical Interfaces and Colabs
+For example:
+```py
+from lycoris import create_lycoris, LycorisNetwork
+
+LycorisNetwork.apply_preset(
+    {"target_name": [".*attn.*"]}
+)
+lycoris_net = create_lycoris(
+    your_model, 
+    1.0, 
+    linear_dim=16, 
+    linear_alpha=2.0, 
+    algo="lokr"
+)
+lycoris_net.apply_to()
+
+# after apply_to(), your_model() will run with LyCORIS net
+lycoris_param = lycoris_net.parameters()
+forward_with_lyco = your_model(x)
+```
+
+You can check my [HakuPhi](https://github.com/KohakuBlueleaf/HakuPhi) project to see how I utilize LyCORIS to finetune the Phi-1.5 models.
+
+#### Graphical Interfaces and Colabs (via kohya trainer)
 
 You can also train LyCORIS with the following graphical interfaces
 * [bmaltais/kohya_ss](https://github.com/bmaltais/kohya_ss)
@@ -121,7 +156,7 @@ and colabs (please help us complete the list!)
 * [hollowstrawberry/kohya-colab](https://github.com/hollowstrawberry/kohya-colab)
 * [Linaqruf/kohya-trainer](https://github.com/Linaqruf/kohya-trainer)
 
-However, they are not guaranteed to be up-to-date. In particular, newer types may not be supported. Consider requesting the developpers for support or simply use the original kohya script in this case.
+However, they are not guaranteed to be up-to-date. In particular, newer types may not be supported. Consider requesting the developers for support or simply use the original kohya script in this case.
 
 
 ## Utilities
@@ -134,9 +169,10 @@ python3 extract_locon.py <settings> <base_model> <db_model> <output>
 Use --help to get more info
 ```
 $ python3 extract_locon.py --help
-usage: extract_locon.py [-h] [--is_v2] [--device DEVICE] [--mode MODE] [--safetensors] [--linear_dim LINEAR_DIM] [--conv_dim CONV_DIM]
-                        [--linear_threshold LINEAR_THRESHOLD] [--conv_threshold CONV_THRESHOLD] [--linear_ratio LINEAR_RATIO] [--conv_ratio CONV_RATIO]
-                        [--linear_percentile LINEAR_PERCENTILE] [--conv_percentile CONV_PERCENTILE]
+usage: extract_locon.py [-h] [--is_v2] [--is_sdxl] [--device DEVICE] [--mode MODE] [--safetensors] [--linear_dim LINEAR_DIM]
+                        [--conv_dim CONV_DIM] [--linear_threshold LINEAR_THRESHOLD] [--conv_threshold CONV_THRESHOLD]
+                        [--linear_ratio LINEAR_RATIO] [--conv_ratio CONV_RATIO] [--linear_quantile LINEAR_QUANTILE]
+                        [--conv_quantile CONV_QUANTILE] [--use_sparse_bias] [--sparsity SPARSITY] [--disable_cp]
                         base_model db_model output_name
 ```
 
@@ -148,7 +184,7 @@ python3 merge.py <settings> <base_model> <lycoris_model> <output>
 Use --help to get more info
 ```
 $ python3 merge.py --help
-usage: merge.py [-h] [--is_v2] [--device DEVICE] [--dtype DTYPE] [--weight WEIGHT] base_model lycoris_model output_name
+usage: merge.py [-h] [--is_v2] [--is_sdxl] [--device DEVICE] [--dtype DTYPE] [--weight WEIGHT] base_model lycoris_model output_name
 ```
 
 
@@ -156,15 +192,24 @@ usage: merge.py [-h] [--is_v2] [--device DEVICE] [--dtype DTYPE] [--weight WEIGH
 ## Change Log
 For full log, please see [Change.md](Change.md)
 
-### 2023/09/27 update to 1.9.0
-* Add norm modules (for training LayerNorm and GroupNorm, which should be good for style)
-* Add full modules (So you can "native fine-tune" with LyCORIS now, should be convenient to try different weight)
-* Add preset config system
-* Add custom config system
-* Merge script support norm and full modules
-* Fix errors with optional requirements
-* Fix errors with not necessary import
-* Fix wrong factorization behaviors
+
+## 2023/12/02 update to 2.0.0
+* Start supporting [HCP-Diffusion](https://github.com/IrisRainbowNeko/HCP-Diffusion) (The reason to name this version "2.0.0")
+  * Now LyCORIS support LoHa/LoKr/Diag-OFT algorithm in HCP-Diffusion
+  * Add Pivotal tuning utilities
+  * Add hcp convert utilities
+  * Have no plan at this time to support full/lora and train_norms since HCP can do them natively
+* Add Diag-OFT modules
+* Add standalone usage support
+  * Can wrap any pytorch module which contains Linear/Conv2d/LayerNorm/GroupNorm modules
+  * Will support more module in the future
+* Add SDXL support in Merge script
+* Add SDXL support in Extract-locon
+* More efficient (speed/vram) implementation for full module
+* Better implementation of custom state_dict
+* Fix errors of dropouts
+* Fix errors of apply_max_norms
+* Fix errors of resume
 
 
 ## Todo list
diff --git a/docs/Algo-List.md b/docs/Algo-List.md
index 18a6708..b4342ed 100644
--- a/docs/Algo-List.md
+++ b/docs/Algo-List.md
@@ -2,6 +2,9 @@
 
 See [Algo-Details.md](Algo-Details.md) and [Demo.md](Demo.md) for more examples and explanation
 
+All the methods below except for GLoKr are supported in a1111/sd-webui.
+However, newer methods may only be available in the latest release / the dev branch.
+
 ### Conventional LoRA
 * Trigged by `algo=lora`
 * Includes Conv layer implementation from LoCon.
@@ -34,12 +37,10 @@ See [Algo-Details.md](Algo-Details.md) and [Demo.md](Demo.md) for more examples
 
 ### Native Fine-Tuning
 * Trigged by `algo=full`
-* **Only supported from a1111/sd-webui 1.6.0**
 * Also known as dreambooth. This just uses full matrix without further decomposition.
 * With our implementation you can use it as LoRA, change base model, and merge on the fly.
-* Pre-decomposing the difference makes training slower, so you may want to train the full model with another trainer (like [EveryDream](https://github.com/victorchall/EveryDream2trainer), [HCP-Diffusion](https://github.com/7eu7d7/HCP-Diffusion)) before extracting it to the LyCORIS full format.
 * Can potentially give the best result if everything is tuned correctly, but is probably also the most sensitive to hyperparameter adjustment.
-* Produces the biggest file, but does _not_ train slower than other methods. especially if you consider direct implementation in other trainers.
+* Produces the biggest file, but does _not_ train slower than other methods.
 
 ### (IA)^3
 * Trigged by `algo=ia3`
@@ -48,6 +49,7 @@ See [Algo-Details.md](Algo-Details.md) and [Demo.md](Demo.md) for more examples
 * This method is good at learning style, but hard to transfer.
 * This method produces very tiny file (less than 1 MB).
 * No network arguments for this method (even preset is not applicable here).
+* **Can be regarded as a special case of Diag-OFT listed below.**
 * Ref: [Few-Shot Parameter-Efficient Fine-Tuning is Better and Cheaper than In-Context Learning](https://arxiv.org/abs/2205.05638)
 
 ### DyLoRA
@@ -64,8 +66,20 @@ See [Algo-Details.md](Algo-Details.md) and [Demo.md](Demo.md) for more examples
 
 ### GLoRA
 
-- TODO
+* Triggered by `algo=glora`
+* Ref: [One-for-All: Generalized LoRA for Parameter-Efficient Fine-tuning](https://arxiv.org/abs/2306.07967)
+* TODO
 
 ### GLoKr
 
 - TODO
+
+
+### Diag-OFT
+
+* Triggered by `algo=diag-oft`
+* It preserves the hyperspherical energy by training orthogonal transformations that apply to outputs of each layer.
+* It converges faster than LoRA according to the original paper, but experiments are still needed.
+* `dim` corresponds to block size: we fix block size instead of block number here to make it more comparable to LoRA
+* Set `constrain` to get COFT and set `rescaled` to get rescaled OFT
+* Ref: [Controlling Text-to-Image Diffusion by Orthogonal Finetuning](https://arxiv.org/abs/2306.07280)
diff --git a/docs/Network-Args.md b/docs/Network-Args.md
index 9179939..1e63a0c 100644
--- a/docs/Network-Args.md
+++ b/docs/Network-Args.md
@@ -1,6 +1,6 @@
 # Network Arguments
 
-Arguments to put in `network_args` for kohya sd  scripts
+Arguments to put in `network_args` for kohya sd scripts
 
 ### Algo
 - Set with `algo=ALGO_NAME`
@@ -8,26 +8,37 @@ Arguments to put in `network_args` for kohya sd  scripts
 
 ### Preset
 - Set with `preset=PRESET/CONFIG_FILE`
-- Pre-impelmented: `full` (default), `attn-mlp`, `attn-only` etc.
+- Pre-implemented: `full` (default), `attn-mlp`, `attn-only` etc.
 - Valid for all but (IA)^3
 - Use `preset=xxx.toml` to choose config file (for LyCORIS module settings)
 - More info in [Preset](Preset.md)
 
 ### Dimension
-- Set with `dim=INT`
+- Dimension of the linear layers is set with the _script argument_ `network_dim`
+- Dimension of the convolutional layers is set with `conv_dim=INT`
 - Valid for all but (IA)^3 and native fine-tuning
 - For LoKr, setting dimension to sufficiently large value (>10240/2) prevents the second block from being further decomposed
 
 ### Alpha
-- Set with `alpha=NUMBER`
+- Alpha of the linear layers is set with the _script argument_ `network_alpha`
+- Alpha of the convolutional layers is set with `conv_alpha=FLOAT`
 - Valid for all but (IA)^3 and native fine-tuning, ignored by full dimension LoKr as well
 - Merge ratio is alpha/dimension, check Appendix B.1 of our [paper](https://arxiv.org/abs/2309.14859) for relation between alpha and learning rate / initialization
 
+### Dropouts
+- Set with `dropout=FLOAT`, `rank_dropout=FLOAT`, `module_dropout=FLOAT`
+- Set the dropout rate, the types of dropout that are valid could vary from method to method
+
 ### Factor
 - Set with `factor=INT`
 - Valid for LoKr
 - Use `-1` to get the smallest decomposition
 
+### Decompose both
+- Enabled with `decompose_both=True`
+- Valid for LoKr
+- Perform LoRA decomposition of both matrices resulting from LoKr decomposition (by default only the larger matrix is decomposed)
+
 ### Block Size
 - Set with `block_size=INT`
 - Valid for DyLoRA
@@ -40,8 +51,18 @@ Arguments to put in `network_args` for kohya sd  scripts
 
 ### Scalar
 - Enabled with `use_scalar=True`
+- Valid for LoRA, LoHa, and LoKr.
 - Train an additional scalar in front of the weight difference
 - Use a different weight initialization strategy
 
 ### Normalization Layers
 - Enabled with `train_norm=True`
+- Valid for all but (IA)^3
+
+### Rescaled OFT
+- Enabled with `rescaled=True`
+- Valid for Diag-OFT
+
+### Constrained OFT
+- Enabled with `constrain=FLOAT`
+- Valid for Diag-OFT
diff --git a/example_configs/training_configs/hcp/base_dataset.yaml b/example_configs/training_configs/hcp/base_dataset.yaml
new file mode 100644
index 0000000..ba429a3
--- /dev/null
+++ b/example_configs/training_configs/hcp/base_dataset.yaml
@@ -0,0 +1,32 @@
+data:
+  dataset1:
+    _target_: hcpdiff.data.TextImagePairDataset
+    _partial_: True # Not directly instantiate the object here. There are other parameters to be added in the runtime.
+    batch_size: 4
+    cache_latents: True
+    att_mask_encode: False
+    loss_weight: 1.0
+
+    source:
+      data_source1:
+        _target_: hcpdiff.data.source.Text2ImageAttMapSource
+        img_root: 'imgs/'
+        prompt_template: 'prompt_tuning_template/object.txt'
+        caption_file: null # path to image captions (file_words)
+        att_mask: null
+        bg_color: [ 255, 255, 255 ] # RGB; for ARGB -> RGB
+
+        word_names: {}
+
+        text_transforms:
+          _target_: torchvision.transforms.Compose
+          transforms:
+            - _target_: hcpdiff.utils.caption_tools.TagShuffle
+            - _target_: hcpdiff.utils.caption_tools.TagDropout
+              p: 0.1
+            - _target_: hcpdiff.utils.caption_tools.TemplateFill
+              word_names: ${....word_names}
+    bucket:
+      _target_: hcpdiff.data.bucket.RatioBucket.from_files # aspect ratio bucket
+      target_area: ${hcp.eval:"512*512"}
+      num_bucket: 5
\ No newline at end of file
diff --git a/example_configs/training_configs/hcp/hcp_caption.txt b/example_configs/training_configs/hcp/hcp_caption.txt
new file mode 100644
index 0000000..f4026e6
--- /dev/null
+++ b/example_configs/training_configs/hcp/hcp_caption.txt
@@ -0,0 +1,1 @@
+{caption}
\ No newline at end of file
diff --git a/example_configs/training_configs/hcp/hcp_diag_oft.yaml b/example_configs/training_configs/hcp/hcp_diag_oft.yaml
new file mode 100644
index 0000000..b443aa2
--- /dev/null
+++ b/example_configs/training_configs/hcp/hcp_diag_oft.yaml
@@ -0,0 +1,69 @@
+_base_:
+  - example_configs/training_configs/hcp/base_dataset.yaml
+  - example_configs/training_configs/hcp/train_base.yaml
+  - example_configs/training_configs/hcp/tuning_base.yaml
+
+
+plugin_unet:
+  diag_oft:
+    _target_: lycoris.hcp.DiagOFTBlock.wrap_model
+    _partial_: True
+    lr: 3e-4
+    dim: 4
+    layers:
+      - 're:.*\.attn.?$'
+      - 're:.*\.ff$'
+
+
+plugin_TE:
+  diag_oft:
+    _target_: lycoris.hcp.DiagOFTBlock.wrap_model
+    _partial_: True
+    lr: 3e-4
+    dim: 4
+    layers:
+      - 're:.*self_attn$'
+      - 're:.*mlp$'
+
+
+tokenizer_pt:
+  train: null
+
+train:
+  train_steps: 100
+  gradient_accumulation_steps: 1
+  save_step: 100
+
+  scheduler:
+    name: 'constant'
+    num_training_steps: 100
+
+  loss:
+    criterion: # min SNR loss
+      _target_: hcpdiff.loss.MinSNRLoss
+      gamma: 5.0
+
+
+model:
+  pretrained_model_name_or_path: 'KBlueLeaf/kohaku-v4-rev1.2'
+  clip_skip: 1
+  tokenizer_repeats: 1
+  ema_unet: 0
+  ema_text_encoder: 0
+
+
+data:
+  dataset1:
+    batch_size: 4
+    cache_latents: False
+
+    source:
+      data_source1:
+        img_root: 'F:\dataset\yog_lora\10_yog_pre'
+        prompt_template: 'example_configs/training_configs/hcp/hcp_caption.txt'
+        caption_file: 'F:\dataset\yog_lora\10_yog_pre'
+
+    bucket:
+      _target_: hcpdiff.data.bucket.RatioBucket.from_files # aspect ratio bucket
+      target_area: ${hcp.eval:"512*512"}
+      num_bucket: 12
\ No newline at end of file
diff --git a/example_configs/training_configs/hcp/hcp_loha.yaml b/example_configs/training_configs/hcp/hcp_loha.yaml
new file mode 100644
index 0000000..23cefe2
--- /dev/null
+++ b/example_configs/training_configs/hcp/hcp_loha.yaml
@@ -0,0 +1,53 @@
+_base_:
+  - example_configs/training_configs/hcp/base_dataset.yaml
+  - example_configs/training_configs/hcp/train_base.yaml
+  - example_configs/training_configs/hcp/tuning_base.yaml
+
+
+plugin_unet:
+  loha:
+    _target_: lycoris.hcp.LohaBlock
+    _partial_: True
+    lr: 5e-4
+    dim: 32
+    alpha: 16
+    layers:
+      - 're:.*\.to_q$'
+      - 're:.*\.to_k$'
+      - 're:.*\.to_v$'
+
+
+tokenizer_pt:
+  train: null
+
+train:
+  gradient_accumulation_steps: 1
+  save_step: 100
+
+  scheduler:
+    name: 'constant'
+    num_training_steps: 100
+
+
+model:
+  pretrained_model_name_or_path: 'KBlueLeaf/kohaku-v4-rev1.2'
+  tokenizer_repeats: 1
+  ema_unet: 0
+  ema_text_encoder: 0
+
+
+data:
+  dataset1:
+    batch_size: 4
+    cache_latents: False
+
+    source:
+      data_source1:
+        img_root: 'C:\Users\apoll\Desktop\AI\training_data\fuzi-choko\5_author_27517'
+        prompt_template: 'example_configs/training_configs/hcp/hcp_caption.txt'
+        caption_file: 'C:\Users\apoll\Desktop\AI\training_data\fuzi-choko\5_author_27517'
+
+    bucket:
+      _target_: hcpdiff.data.bucket.RatioBucket.from_files # aspect ratio bucket
+      target_area: ${hcp.eval:"512*512"}
+      num_bucket: 12
\ No newline at end of file
diff --git a/example_configs/training_configs/hcp/hcp_lokr.yaml b/example_configs/training_configs/hcp/hcp_lokr.yaml
new file mode 100644
index 0000000..e63a2c3
--- /dev/null
+++ b/example_configs/training_configs/hcp/hcp_lokr.yaml
@@ -0,0 +1,73 @@
+_base_:
+  - example_configs/training_configs/hcp/base_dataset.yaml
+  - example_configs/training_configs/hcp/train_base.yaml
+  - example_configs/training_configs/hcp/tuning_base.yaml
+
+
+plugin_unet:
+  lokr:
+    _target_: lycoris.hcp.LokrBlock.wrap_model
+    _partial_: True
+    lr: 1e-4
+    dim: 10000
+    alpha: 0
+    factor: 8
+    layers:
+      - 're:.*\.attn.?$'
+      - 're:.*\.ff$'
+
+
+plugin_TE:
+  lokr:
+    _target_: lycoris.hcp.LokrBlock.wrap_model
+    _partial_: True
+    lr: 1e-4
+    dim: 10000
+    alpha: 0
+    factor: 8
+    layers:
+      - 're:.*self_attn$'
+      - 're:.*mlp$'
+
+
+tokenizer_pt:
+  train: null
+
+train:
+  train_steps: 100
+  gradient_accumulation_steps: 1
+  save_step: 100
+
+  scheduler:
+    name: 'constant'
+    num_training_steps: 100
+
+  loss:
+    criterion: # min SNR loss
+      _target_: hcpdiff.loss.MinSNRLoss
+      gamma: 5.0
+
+
+model:
+  pretrained_model_name_or_path: 'KBlueLeaf/kohaku-v4-rev1.2'
+  clip_skip: 1
+  tokenizer_repeats: 1
+  ema_unet: 0
+  ema_text_encoder: 0
+
+
+data:
+  dataset1:
+    batch_size: 4
+    cache_latents: False
+
+    source:
+      data_source1:
+        img_root: 'F:\dataset\yog_lora\10_yog_pre'
+        prompt_template: 'example_configs/training_configs/hcp/hcp_caption.txt'
+        caption_file: 'F:\dataset\yog_lora\10_yog_pre'
+
+    bucket:
+      _target_: hcpdiff.data.bucket.RatioBucket.from_files # aspect ratio bucket
+      target_area: ${hcp.eval:"512*512"}
+      num_bucket: 12
\ No newline at end of file
diff --git a/example_configs/training_configs/hcp/text2img.yaml b/example_configs/training_configs/hcp/text2img.yaml
new file mode 100644
index 0000000..bec62e8
--- /dev/null
+++ b/example_configs/training_configs/hcp/text2img.yaml
@@ -0,0 +1,51 @@
+# base_state*base_model_alpha + (lora_state[i]*lora_scale[i]*lora_alpha[i]) + (part_state[k]*part_alpha[k])
+
+pretrained_model: ''
+prompt: ''
+neg_prompt: 'lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry'
+emb_dir: 'embs/'
+N_repeats: 1
+clip_skip: 0
+clip_final_norm: True
+bs: 4
+num: 1
+seed: null
+dtype: 'fp16'
+
+condition: null
+
+ex_input: {}
+
+# Syntactic sugar for interface
+save:
+  out_dir: 'output/'
+  save_cfg: True
+  image_type: png
+  quality: 95
+#  image_type: webp
+#  quality: 75
+
+offload: null
+
+#vae_optimize: null
+vae_optimize:
+  tiling: False
+  slicing: False
+
+interface:
+  - _target_: hcpdiff.vis.DiskInterface
+    show_steps: 0
+    save_root: ${save.out_dir}
+    save_cfg: ${save.save_cfg}
+    image_type: ${save.image_type}
+    quality: ${save.quality}
+
+infer_args:
+  width: 512
+  height: 512
+  guidance_scale: 7.5
+  num_inference_steps: 50
+
+new_components: {}
+
+merge: null
\ No newline at end of file
diff --git a/example_configs/training_configs/hcp/train_base.yaml b/example_configs/training_configs/hcp/train_base.yaml
new file mode 100644
index 0000000..e7ece8a
--- /dev/null
+++ b/example_configs/training_configs/hcp/train_base.yaml
@@ -0,0 +1,81 @@
+exp_dir: exps/${hcp.time:}
+mixed_precision: 'fp16'
+allow_tf32: False
+seed: 114514
+ckpt_type: 'safetensors' # [torch, safetensors]
+
+vis_info:
+  prompt: null
+  negative_prompt: ''
+
+train:
+  train_steps: 1000
+  train_epochs: null # Choose one of [train_steps, train_epochs]
+  gradient_accumulation_steps: 1
+  workers: 4
+  max_grad_norm: 1.0
+  set_grads_to_none: False
+  save_step: 100
+  cfg_scale: '1.0' # for DreamArtist
+
+  resume: null
+#  resume:
+#    ckpt_path:
+#      unet: []
+#      TE: []
+#      words: {}
+#    start_step: 0
+
+  loss:
+    criterion:
+      _target_: torch.nn.MSELoss
+      _partial_: True
+      reduction: 'none' # support for attention mask
+    type: 'eps' # 'eps' or 'sample'
+
+  optimizer:
+    _target_: torch.optim.AdamW
+    _partial_: True
+    weight_decay: 1e-3
+
+  optimizer_pt:
+    _target_: torch.optim.AdamW
+    _partial_: True
+    weight_decay: 5e-4
+
+  scale_lr: True # auto scale lr with total batch size
+  scheduler:
+    name: 'one_cycle'
+    num_warmup_steps: 200
+    num_training_steps: 1000
+    scheduler_kwargs: {} # args for scheduler
+
+  scale_lr_pt: True
+  scheduler_pt: ${.scheduler}
+
+logger:
+  -
+    _target_: hcpdiff.loggers.CLILogger
+    _partial_: True
+    out_path: 'train.log'
+    log_step: 20
+
+model:
+  revision: null
+  pretrained_model_name_or_path: null
+  tokenizer_repeats: 2
+  enable_xformers: True
+  gradient_checkpointing: True
+  force_cast_precision: False
+  ema_unet: 0 # 0 to disable
+  ema_text_encoder: 0 # 0 to disable
+  clip_skip: 0
+  clip_final_norm: True
+
+  tokenizer: null
+  noise_scheduler: null
+  unet: null
+  text_encoder: null
+  vae: null
+
+previewer: null
\ No newline at end of file
diff --git a/example_configs/training_configs/hcp/tuning_base.yaml b/example_configs/training_configs/hcp/tuning_base.yaml
new file mode 100644
index 0000000..2eaf072
--- /dev/null
+++ b/example_configs/training_configs/hcp/tuning_base.yaml
@@ -0,0 +1,15 @@
+tokenizer_pt:
+  emb_dir: './'
+  replace: False
+  train: []
+  #train:
+  #  - {name: pt1, lr: 1e-3}
+
+unet: null
+lora_unet: null
+
+text_encoder: null
+lora_text_encoder: null
+
+plugin_unet: null
+plugin_TE: null
\ No newline at end of file
diff --git a/example_configs/training_configs/loha_config.toml b/example_configs/training_configs/kohya/loha_config.toml
similarity index 100%
rename from example_configs/training_configs/loha_config.toml
rename to example_configs/training_configs/kohya/loha_config.toml
diff --git a/example_configs/training_configs/lokr_config.toml b/example_configs/training_configs/kohya/lokr_config.toml
similarity index 100%
rename from example_configs/training_configs/lokr_config.toml
rename to example_configs/training_configs/kohya/lokr_config.toml
diff --git a/example_configs/training_configs/lora_config.toml b/example_configs/training_configs/kohya/lora_config.toml
similarity index 100%
rename from example_configs/training_configs/lora_config.toml
rename to example_configs/training_configs/kohya/lora_config.toml
diff --git a/example_configs/training_configs/lycoris_full_config.toml b/example_configs/training_configs/kohya/lycoris_full_config.toml
similarity index 100%
rename from example_configs/training_configs/lycoris_full_config.toml
rename to example_configs/training_configs/kohya/lycoris_full_config.toml
diff --git a/experiments/better_conv.py b/experiments/better_conv.py
index 9cc87ae..72e6f59 100644
--- a/experiments/better_conv.py
+++ b/experiments/better_conv.py
@@ -7,47 +7,47 @@ class HadaWeightCP(torch.autograd.Function):
     @staticmethod
     def forward(ctx, orig_weight, t1, w1a, w1b, t2, w2a, w2b, scale=torch.tensor(1)):
         ctx.save_for_backward(t1, w1a, w1b, t2, w2a, w2b, scale)
-        
-        temp = torch.einsum('i j k l, j r -> i r k l', t1, w1b)
-        rebuild1 = torch.einsum('i j k l, i r -> r j k l', temp, w1a)
-        
-        temp = torch.einsum('i j k l, j r -> i r k l', t2, w2b)
-        rebuild2 = torch.einsum('i j k l, i r -> r j k l', temp, w2a)
-        
-        return orig_weight + rebuild1*rebuild2*scale
+
+        temp = torch.einsum("i j k l, j r -> i r k l", t1, w1b)
+        rebuild1 = torch.einsum("i j k l, i r -> r j k l", temp, w1a)
+
+        temp = torch.einsum("i j k l, j r -> i r k l", t2, w2b)
+        rebuild2 = torch.einsum("i j k l, i r -> r j k l", temp, w2a)
+
+        return orig_weight + rebuild1 * rebuild2 * scale
 
     @staticmethod
     def backward(ctx, grad_out):
         (t1, w1a, w1b, t2, w2a, w2b, scale) = ctx.saved_tensors
-        
-        grad_out = grad_out*scale
-        
-        temp = torch.einsum('i j k l, j r -> i r k l', t2, w2b)
-        rebuild = torch.einsum('i j k l, i r -> r j k l', temp, w2a)
-        
-        grad_w = rebuild*grad_out
+
+        grad_out = grad_out * scale
+
+        temp = torch.einsum("i j k l, j r -> i r k l", t2, w2b)
+        rebuild = torch.einsum("i j k l, i r -> r j k l", temp, w2a)
+
+        grad_w = rebuild * grad_out
         del rebuild
-        
-        grad_w1a = torch.einsum('r j k l, i j k l -> r i', temp, grad_w)
-        grad_temp = torch.einsum('i j k l, i r -> r j k l', grad_w, w1a.T)
+
+        grad_w1a = torch.einsum("r j k l, i j k l -> r i", temp, grad_w)
+        grad_temp = torch.einsum("i j k l, i r -> r j k l", grad_w, w1a.T)
         del grad_w, temp
-        
-        grad_w1b = torch.einsum('i r k l, i j k l -> r j', t1, grad_temp)
-        grad_t1 = torch.einsum('i j k l, j r -> i r k l', grad_temp, w1b.T)
+
+        grad_w1b = torch.einsum("i r k l, i j k l -> r j", t1, grad_temp)
+        grad_t1 = torch.einsum("i j k l, j r -> i r k l", grad_temp, w1b.T)
         del grad_temp
-        
-        temp = torch.einsum('i j k l, j r -> i r k l', t1, w1b)
-        rebuild = torch.einsum('i j k l, i r -> r j k l', temp, w1a)
-        
-        grad_w = rebuild*grad_out
+
+        temp = torch.einsum("i j k l, j r -> i r k l", t1, w1b)
+        rebuild = torch.einsum("i j k l, i r -> r j k l", temp, w1a)
+
+        grad_w = rebuild * grad_out
         del rebuild
-        
-        grad_w2a = torch.einsum('r j k l, i j k l -> r i', temp, grad_w)
-        grad_temp = torch.einsum('i j k l, i r -> r j k l', grad_w, w2a.T)
+
+        grad_w2a = torch.einsum("r j k l, i j k l -> r i", temp, grad_w)
+        grad_temp = torch.einsum("i j k l, i r -> r j k l", grad_w, w2a.T)
         del grad_w, temp
-        
-        grad_w2b = torch.einsum('i r k l, i j k l -> r j', t2, grad_temp)
-        grad_t2 = torch.einsum('i j k l, j r -> i r k l', grad_temp, w2b.T)
+
+        grad_w2b = torch.einsum("i r k l, i j k l -> r j", t2, grad_temp)
+        grad_t2 = torch.einsum("i j k l, j r -> i r k l", grad_temp, w2b.T)
         del grad_temp
         return grad_out, grad_t1, grad_w1a, grad_w1b, grad_t2, grad_w2a, grad_w2b, None
 
@@ -57,13 +57,13 @@ def make_weight_cp(orig_weight, t1, w1a, w1b, t2, w2a, w2b, scale=torch.tensor(0
 
 
 def make_cp(orig_weight, t1, w1a, w1b, t2, w2a, w2b, scale=torch.tensor(0.25)):
-    temp = torch.einsum('i j k l, j r -> i r k l', t1, w1b)
-    rebuild1 = torch.einsum('i j k l, i r -> r j k l', temp, w1a)
-    
-    temp = torch.einsum('i j k l, j r -> i r k l', t2, w2b)
-    rebuild2 = torch.einsum('i j k l, i r -> r j k l', temp, w2a)
-    
-    return orig_weight + rebuild1*rebuild2 * scale
+    temp = torch.einsum("i j k l, j r -> i r k l", t1, w1b)
+    rebuild1 = torch.einsum("i j k l, i r -> r j k l", temp, w1a)
+
+    temp = torch.einsum("i j k l, j r -> i r k l", t2, w2b)
+    rebuild2 = torch.einsum("i j k l, i r -> r j k l", temp, w2a)
+
+    return orig_weight + rebuild1 * rebuild2 * scale
 
 
 KERNEL_SIZE = 3
@@ -100,6 +100,6 @@ def make_cp(orig_weight, t1, w1a, w1b, t2, w2a, w2b, scale=torch.tensor(0.25)):
 grad2 = t1.grad.clone()
 
 
-print('MSE Loss: ', F.mse_loss(grad1, grad2))
-print('L1 Loss : ', F.l1_loss(grad1, grad2))
-print('Distance: ', torch.dist(grad1, grad2))
\ No newline at end of file
+print("MSE Loss: ", F.mse_loss(grad1, grad2))
+print("L1 Loss : ", F.l1_loss(grad1, grad2))
+print("Distance: ", torch.dist(grad1, grad2))
diff --git a/experiments/better_conv_extract.py b/experiments/better_conv_extract.py
index 002108a..df9f166 100644
--- a/experiments/better_conv_extract.py
+++ b/experiments/better_conv_extract.py
@@ -1,4 +1,5 @@
 import os, sys
+
 sys.path.insert(0, os.getcwd())
 import torch
 import torch.nn as nn
@@ -9,21 +10,19 @@
 
 
 def make_sparse(t, sparsity=0.95):
-    sparse_t = t.masked_fill(
-        torch.abs(t) < torch.quantile(torch.abs(t), sparsity), 
-        0
-    )
+    sparse_t = t.masked_fill(torch.abs(t) < torch.quantile(torch.abs(t), sparsity), 0)
     return sparse_t
 
+
 def make_cp(t, x1, x2):
     x1 = x1.reshape(x1.size(0), -1)
     x2 = x2.reshape(x2.size(0), -1)
-    
+
     # [rank, rank, k, k] * [rank, out]
-    temp = torch.einsum('n m k l, i n -> i m k l', t, x2)
-    
+    temp = torch.einsum("n m k l, i n -> i m k l", t, x2)
+
     # [out, rank, k, k] * [rank, in]
-    result = torch.einsum('i m k l, m j -> i j k l', temp, x1)
+    result = torch.einsum("i m k l, m j -> i j k l", temp, x1)
     print(result.shape)
     return result
 
@@ -40,28 +39,27 @@ def make_cp(t, x1, x2):
 
 conv_orig = nn.Conv2d(IN_CH, OUT_CH, 3, 1, 1, bias=False).cuda()
 conv_orig.weight = nn.Parameter(
-    (torch.randn(OUT_CH, LORA_RANK) 
-     @ torch.randn(LORA_RANK, IN_CH*9) 
-     * 0.05
-     ).reshape(conv_orig.weight.shape).cuda()
+    (torch.randn(OUT_CH, LORA_RANK) @ torch.randn(LORA_RANK, IN_CH * 9) * 0.05)
+    .reshape(conv_orig.weight.shape)
+    .cuda()
 )
 conv_rebd = nn.Conv2d(IN_CH, OUT_CH, 3, 1, 1, bias=False).cuda()
 
-extract_a, extract_b, diff = extract_conv(conv_orig.weight, 'fixed', LORA_RANK, 'cuda')
+extract_a, extract_b, diff = extract_conv(conv_orig.weight, "fixed", LORA_RANK, "cuda")
 out_weight = extract_b.reshape(OUT_CH, -1) @ extract_a.reshape(LORA_RANK, -1)
 out_weight = out_weight.reshape(conv_orig.weight.shape)
 
-extract_a, extract_c, diff = extract_conv(extract_a.transpose(0, 1), 'fixed', LORA_RANK, 'cuda')
+extract_a, extract_c, diff = extract_conv(
+    extract_a.transpose(0, 1), "fixed", LORA_RANK, "cuda"
+)
 extract_a = extract_a.transpose(0, 1)
 extract_c = extract_c.transpose(0, 1)
 print(extract_a.shape, extract_b.shape, extract_c.shape)
 
 
-conv_rebd.weight = nn.Parameter(make_cp(
-    extract_a, extract_c, extract_b
-))
+conv_rebd.weight = nn.Parameter(make_cp(extract_a, extract_c, extract_b))
 print()
-print('without sparse bias')
-print('MSE Loss: ', F.mse_loss(conv_orig.weight, conv_rebd.weight))
-print('L1 Loss : ', F.l1_loss(conv_orig.weight, conv_rebd.weight))
-print('Distance: ', torch.dist(conv_orig.weight, conv_rebd.weight))
\ No newline at end of file
+print("without sparse bias")
+print("MSE Loss: ", F.mse_loss(conv_orig.weight, conv_rebd.weight))
+print("L1 Loss : ", F.l1_loss(conv_orig.weight, conv_rebd.weight))
+print("Distance: ", torch.dist(conv_orig.weight, conv_rebd.weight))
diff --git a/experiments/compare_norm.py b/experiments/compare_norm.py
index 168326f..8b0822a 100644
--- a/experiments/compare_norm.py
+++ b/experiments/compare_norm.py
@@ -9,13 +9,13 @@
 with safe_open("./sd1.safetensors", framework="pt", device=0) as f:
     for k in f.keys():
         tensor = f.get_tensor(k)
-        if len(tensor.shape)>1 and max(tensor.shape)>1000:
+        if len(tensor.shape) > 1 and max(tensor.shape) > 1000:
             print(torch.norm(tensor), tensor.shape)
 
-print('='*100)
-print('-'*100)
+print("=" * 100)
+print("-" * 100)
 with safe_open("./xl.safetensors", framework="pt", device=0) as f:
     for k in f.keys():
         tensor = f.get_tensor(k)
-        if len(tensor.shape)>1 and max(tensor.shape)>1000:
-            print(torch.norm(tensor), tensor.shape)
\ No newline at end of file
+        if len(tensor.shape) > 1 and max(tensor.shape) > 1000:
+            print(torch.norm(tensor), tensor.shape)
diff --git a/experiments/concept.py b/experiments/concept.py
new file mode 100644
index 0000000..5d9e50c
--- /dev/null
+++ b/experiments/concept.py
@@ -0,0 +1,285 @@
+import os, sys
+
+sys.path.insert(0, os.getcwd())
+import math
+from warnings import warn
+import os
+from typing import List
+import torch
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from ..lycoris.kohya_utils import *
+
+
+class ConceptModule(nn.Module):
+    def __init__(
+        self,
+        lora_name,
+        org_module: nn.Module,
+        multiplier=1.0,
+    ):
+        """if alpha == 0 or None, alpha is rank (no scaling)."""
+        super().__init__()
+        self.lora_name = lora_name
+
+        if org_module.__class__.__name__ == "Conv2d":
+            # For general LoCon
+            stride = org_module.stride
+            padding = org_module.padding
+            self.op = F.conv2d
+            self.extra_args = {"stride": stride, "padding": padding}
+        else:
+            self.op = F.linear
+            self.extra_args = {}
+
+        self.concept_mask = nn.Parameter(torch.ones_like(org_module.weight))
+        # self.concept_neuron = None
+
+        self.multiplier = multiplier
+        self.org_module = [org_module]
+
+    def apply_to(self):
+        self.org_module[0].forward = self.forward
+
+    def forward(self, x):
+        bias = None if self.org_module[0].bias is None else self.org_module[0].bias.data
+        return self.op(
+            x,
+            self.org_module[0].weight.data
+            * (1 + self.concept_mask * self.concept_neuron),
+            bias,
+            **self.extra_args,
+        )
+        if torch.max(self.concept_mask) - torch.min(self.concept_mask) < 0.1:
+            return self.op(
+                x,
+                self.org_module[0].weight.data * self.concept_mask,
+                bias,
+                **self.extra_args,
+            )
+        else:
+            if self.concept_neuron is None:
+                print("Found neurons")
+                mask = self.concept_mask
+                threshold = torch.quantile(mask, 0.05)
+                neuron = mask < threshold
+                self.concept_mask.requires_grad_(False)
+                self.concept_mask[mask >= threshold] = 0
+                self.concept_mask[neuron] = mask[neuron] - 1
+                self.concept_mask.requires_grad_(True)
+                self.concept_neuron = neuron
+            self.concept_neuron = self.concept_mask * self.concept_neuron
+            return self.op(
+                x,
+                self.org_module[0].weight.data
+                * (1 + self.concept_mask * self.concept_neuron),
+                bias,
+                **self.extra_args,
+            )
+
+
+class ConceptNetwork(torch.nn.Module):
+    """
+    LoRA + LoCon
+    """
+
+    # Ignore proj_in or proj_out, their channels is only a few.
+    UNET_TARGET_REPLACE_MODULE = [
+        "Transformer2DModel",
+        "Attention",
+    ]
+    TEXT_ENCODER_TARGET_REPLACE_MODULE = ["CLIPAttention"]
+    LORA_PREFIX_UNET = "lora_unet"
+    LORA_PREFIX_TEXT_ENCODER = "lora_te"
+
+    def __init__(self, text_encoder, unet, multiplier=1.0, *args, **kwargs) -> None:
+        super().__init__()
+        self.multiplier = multiplier
+
+        total_params = 0
+
+        # create module instances
+        def create_modules(
+            prefix, root_module: torch.nn.Module, target_replace_modules
+        ) -> List[ConceptModule]:
+            nonlocal total_params
+            print("Create Concept Module")
+            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_name.rsplit(".", 1)[-1] not in [
+                            "to_k",
+                            "to_v",
+                            "k_proj",
+                            "v_proj",
+                        ]:
+                            continue
+                        lora_name = prefix + "." + name + "." + child_name
+                        lora_name = lora_name.replace(".", "_")
+                        total_params += child_module.weight.size().numel()
+                        lora = ConceptModule(
+                            lora_name,
+                            child_module,
+                            self.multiplier,
+                        )
+                        loras.append(lora)
+            return loras
+
+        self.text_encoder_loras = create_modules(
+            ConceptNetwork.LORA_PREFIX_TEXT_ENCODER,
+            text_encoder,
+            ConceptNetwork.TEXT_ENCODER_TARGET_REPLACE_MODULE,
+        )
+        print(f"create LoRA for Text Encoder: {len(self.text_encoder_loras)} modules.")
+
+        self.unet_loras = create_modules(
+            ConceptNetwork.LORA_PREFIX_UNET,
+            unet,
+            ConceptNetwork.UNET_TARGET_REPLACE_MODULE,
+        )
+        print(f"create LoRA for U-Net: {len(self.unet_loras)} modules.")
+
+        self.weights_sd = None
+
+        print(f"Total concept param: {total_params}")
+
+        # 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 set_multiplier(self, multiplier):
+        self.multiplier = multiplier
+        for lora in self.text_encoder_loras + self.unet_loras:
+            lora.multiplier = self.multiplier
+
+    def load_weights(self, file):
+        if os.path.splitext(file)[1] == ".safetensors":
+            from safetensors.torch import load_file, safe_open
+
+            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(ConceptNetwork.LORA_PREFIX_TEXT_ENCODER):
+                    weights_has_text_encoder = True
+                elif key.startswith(ConceptNetwork.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
+        def make_ckpt(module):
+            if isinstance(module, torch.nn.Module):
+                module.grad_ckpt = True
+
+        self.apply(make_ckpt)
+        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)
+        all_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
+            all_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
+            all_params.append(param_data)
+
+        return all_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
+
+            # Precalculate model hashes to save time on indexing
+            if metadata is None:
+                metadata = {}
+            model_hash, legacy_hash = precalculate_safetensors_hashes(
+                state_dict, metadata
+            )
+            metadata["sshs_model_hash"] = model_hash
+            metadata["sshs_legacy_hash"] = legacy_hash
+
+            save_file(state_dict, file, metadata)
+        else:
+            torch.save(state_dict, file)
diff --git a/experiments/concept_neuron.py b/experiments/concept_neuron.py
index 3b7d010..289fa31 100644
--- a/experiments/concept_neuron.py
+++ b/experiments/concept_neuron.py
@@ -1,4 +1,5 @@
 import os, sys
+
 sys.path.insert(0, os.getcwd())
 import math
 from warnings import warn
@@ -18,31 +19,27 @@
 
 class ConceptModule(nn.Module):
     def __init__(
-        self, 
-        lora_name, org_module: nn.Module, 
+        self,
+        lora_name,
+        org_module: nn.Module,
         multiplier=1.0,
     ):
-        """ if alpha == 0 or None, alpha is rank (no scaling). """
+        """if alpha == 0 or None, alpha is rank (no scaling)."""
         super().__init__()
         self.lora_name = lora_name
 
-        if org_module.__class__.__name__ == 'Conv2d':
+        if org_module.__class__.__name__ == "Conv2d":
             # For general LoCon
             stride = org_module.stride
             padding = org_module.padding
             self.op = F.conv2d
-            self.extra_args = {
-                'stride': stride,
-                'padding': padding
-            }
+            self.extra_args = {"stride": stride, "padding": padding}
         else:
             self.op = F.linear
             self.extra_args = {}
-        
-        self.concept_mask = nn.Parameter(torch.ones_like(
-            org_module.weight
-        ))
-        
+
+        self.concept_mask = nn.Parameter(torch.ones_like(org_module.weight))
+
         self.multiplier = multiplier
         self.org_module = [org_module]
 
@@ -52,84 +49,102 @@ def apply_to(self):
     def forward(self, x):
         bias = None if self.org_module[0].bias is None else self.org_module[0].bias.data
         return self.op(
-            x, self.org_module[0].weight.data * (1 + (self.concept_mask-1) * self.multiplier),
-            bias, **self.extra_args
+            x,
+            self.org_module[0].weight.data
+            * (1 + (self.concept_mask - 1) * self.multiplier),
+            bias,
+            **self.extra_args,
         )
 
 
 class ConceptNetwork(torch.nn.Module):
-    '''
+    """
     LoRA + LoCon
-    '''
+    """
+
     # Ignore proj_in or proj_out, their channels is only a few.
     UNET_TARGET_REPLACE_MODULE = [
-        "Transformer2DModel", 
+        "Transformer2DModel",
         "Attention",
     ]
     TEXT_ENCODER_TARGET_REPLACE_MODULE = ["CLIPAttention"]
-    LORA_PREFIX_UNET = 'lora_unet'
-    LORA_PREFIX_TEXT_ENCODER = 'lora_te'
+    LORA_PREFIX_UNET = "lora_unet"
+    LORA_PREFIX_TEXT_ENCODER = "lora_te"
 
-    def __init__(
-        self, 
-        text_encoder, unet, 
-        multiplier=1.0,
-        **kwargs
-    ) -> None:
+    def __init__(self, text_encoder, unet, multiplier=1.0, **kwargs) -> None:
         super().__init__()
         self.multiplier = multiplier
-        
+
         total_params = 0
+
         # create module instances
-        def create_modules(prefix, root_module: torch.nn.Module, target_replace_modules) -> List[ConceptModule]:
+        def create_modules(
+            prefix, root_module: torch.nn.Module, target_replace_modules
+        ) -> List[ConceptModule]:
             nonlocal total_params
-            print('Create Concept Module')
+            print("Create Concept Module")
             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_name.rsplit('.', 1)[-1] not in ['to_k', 'to_v', 'k_proj', 'v_proj']:
+                        if child_name.rsplit(".", 1)[-1] not in [
+                            "to_k",
+                            "to_v",
+                            "k_proj",
+                            "v_proj",
+                        ]:
                             continue
-                        lora_name = prefix + '.' + name + '.' + child_name
-                        lora_name = lora_name.replace('.', '_')
+                        lora_name = prefix + "." + name + "." + child_name
+                        lora_name = lora_name.replace(".", "_")
                         total_params += child_module.weight.size().numel()
                         lora = ConceptModule(
-                            lora_name, child_module, self.multiplier, 
+                            lora_name,
+                            child_module,
+                            self.multiplier,
                         )
                         loras.append(lora)
             return loras
 
         self.text_encoder_loras = create_modules(
             ConceptNetwork.LORA_PREFIX_TEXT_ENCODER,
-            text_encoder, 
-            ConceptNetwork.TEXT_ENCODER_TARGET_REPLACE_MODULE
+            text_encoder,
+            ConceptNetwork.TEXT_ENCODER_TARGET_REPLACE_MODULE,
+        )
+        print(
+            f"create Concept for Text Encoder: {len(self.text_encoder_loras)} modules."
         )
-        print(f"create Concept for Text Encoder: {len(self.text_encoder_loras)} modules.")
 
-        self.unet_loras = create_modules(ConceptNetwork.LORA_PREFIX_UNET, unet, ConceptNetwork.UNET_TARGET_REPLACE_MODULE)
+        self.unet_loras = create_modules(
+            ConceptNetwork.LORA_PREFIX_UNET,
+            unet,
+            ConceptNetwork.UNET_TARGET_REPLACE_MODULE,
+        )
         print(f"create Concept for U-Net: {len(self.unet_loras)} modules.")
 
         self.weights_sd = None
-        
-        print(f'Total concept param: {total_params}')
-        
+
+        print(f"Total concept param: {total_params}")
+
         # 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}"
+            assert (
+                lora.lora_name not in names
+            ), f"duplicated lora name: {lora.lora_name}"
             names.add(lora.lora_name)
 
     def set_multiplier(self, multiplier):
         self.multiplier = multiplier
         for lora in self.text_encoder_loras + self.unet_loras:
             lora.multiplier = self.multiplier
-            
+
     def load_weights(self, file):
-        if os.path.splitext(file)[1] == '.safetensors':
+        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')
+            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:
@@ -143,14 +158,20 @@ def apply_to(self, text_encoder, unet, apply_text_encoder=None, apply_unet=None)
             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のフラグが矛盾しています"
+                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のフラグが矛盾しています"
+                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"
+            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")
@@ -176,6 +197,7 @@ def enable_gradient_checkpointing(self):
         def make_ckpt(module):
             if isinstance(module, torch.nn.Module):
                 module.grad_ckpt = True
+
         self.apply(make_ckpt)
         pass
 
@@ -190,15 +212,15 @@ def enumerate_params(loras):
         all_params = []
 
         if self.text_encoder_loras:
-            param_data = {'params': enumerate_params(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
+                param_data["lr"] = text_encoder_lr
             all_params.append(param_data)
 
         if self.unet_loras:
-            param_data = {'params': enumerate_params(self.unet_loras)}
+            param_data = {"params": enumerate_params(self.unet_loras)}
             if unet_lr is not None:
-                param_data['lr'] = unet_lr
+                param_data["lr"] = unet_lr
             all_params.append(param_data)
 
         return all_params
@@ -224,16 +246,18 @@ def save_weights(self, file, dtype, metadata):
                 v = v.detach().clone().to("cpu").to(dtype)
                 state_dict[key] = v
 
-        if os.path.splitext(file)[1] == '.safetensors':
+        if os.path.splitext(file)[1] == ".safetensors":
             from safetensors.torch import save_file
 
             # Precalculate model hashes to save time on indexing
             if metadata is None:
                 metadata = {}
-            model_hash, legacy_hash = precalculate_safetensors_hashes(state_dict, metadata)
+            model_hash, legacy_hash = precalculate_safetensors_hashes(
+                state_dict, metadata
+            )
             metadata["sshs_model_hash"] = model_hash
             metadata["sshs_legacy_hash"] = legacy_hash
 
             save_file(state_dict, file, metadata)
         else:
-            torch.save(state_dict, file)
\ No newline at end of file
+            torch.save(state_dict, file)
diff --git a/experiments/einsum_einops.py b/experiments/einsum_einops.py
new file mode 100644
index 0000000..d30b732
--- /dev/null
+++ b/experiments/einsum_einops.py
@@ -0,0 +1,11 @@
+import torch
+from einops import einsum, rearrange
+
+
+diag_oft = torch.zeros(3, 4, 4)
+weight = torch.randn(12, 16)
+eye = torch.eye(4).unsqueeze(0)
+
+weight = rearrange(weight, "(k n) ... -> k n ...", k=3, n=4)
+result = torch.einsum("k n m, k n ... -> k m ...", diag_oft + eye, weight)
+print(result)
diff --git a/experiments/flops_calc.py b/experiments/flops_calc.py
new file mode 100644
index 0000000..fb937ed
--- /dev/null
+++ b/experiments/flops_calc.py
@@ -0,0 +1,79 @@
+def mac_matmul(n, p, m):
+    return n * p * m
+
+
+def linear_layer_mac(seq_len, in_dim, out_dim):
+    return seq_len * in_dim * out_dim
+
+
+def conv_layer_mac(in_h, in_w, in_dim, out_dim, kernel_size, stride, padding):
+    conv_ops = (in_h + padding * 2) * (in_w + padding * 2) // stride
+    filter_mac = kernel_size**2 * in_dim * out_dim
+    return 2 * conv_ops * filter_mac
+
+
+def mac_for_lin_bax(seq_len, in_dim, out_dim, rank):
+    return linear_layer_mac(seq_len, in_dim, rank) + linear_layer_mac(
+        seq_len, rank, out_dim
+    )
+
+
+def mac_for_conv_bax(in_h, in_w, in_dim, out_dim, rank):
+    return conv_layer_mac(in_h, in_w, in_dim, rank, 3, 1, 1) + conv_layer_mac(
+        in_h, in_w, rank, out_dim, 1, 1, 0
+    )
+
+
+def mac_for_lin_wba(in_dim, out_dim, rank):
+    return mac_matmul(in_dim, rank, out_dim)
+
+
+def mac_for_conv_wba(in_dim, out_dim, rank):
+    return mac_matmul(in_dim, rank, out_dim * 3**2)
+
+
+# shape list for 64x64 input of SD Unet
+shape_list = [(320, 96, 96), (640, 48, 48), (1280, 24, 24), (1280, 12, 12)]
+
+mac_lin_bax = 0
+mac_conv_bax = 0
+mac_lin_wba = 0
+mac_conv_wba = 0
+rank = 8
+
+
+for shape in shape_list:
+    print(shape)
+    original_lin = linear_layer_mac(shape[1] * shape[2], shape[0], shape[0])
+    original_conv = conv_layer_mac(shape[1], shape[2], shape[0], shape[0], 3, 1, 1)
+
+    mac_lin_bax = mac_for_lin_bax(shape[1] * shape[2], shape[0], shape[0], rank)
+    mac_lin_backward_bax = (
+        mac_matmul(rank, shape[1] * shape[2], shape[0])  # dB = dY*(AX)
+        + mac_matmul(shape[1] * shape[2], shape[0], rank)  # d(AX) = B*dY
+        + mac_matmul(rank, shape[1] * shape[2], shape[0])  # dA = d(AX)*X
+        + mac_matmul(shape[1] * shape[2], rank, shape[0])  # dX = A*d(AX)
+    )
+    mac_lin_wba = mac_for_lin_wba(shape[0], shape[0], rank)
+    mac_lin_backward_wba = mac_matmul(rank, shape[0], shape[0]) * 2
+
+    mac_conv_bax = mac_for_conv_bax(shape[1], shape[2], shape[0], shape[0], rank)
+    mac_conv_backward_bax = ()  # 2 times of forward I guess, since for linear it is 2 times.
+    mac_conv_wba = mac_for_conv_wba(shape[0], shape[0], rank)
+    mac_conv_backward_wba = mac_matmul(rank, shape[0] * 3**2, shape[0]) + mac_matmul(
+        rank, shape[0], shape[0] * 3**2
+    )
+
+    print(f"lin")
+    print(f"lin w size    : {shape[0]**2}")
+    print(f"original_lin  : {original_lin}")
+    print(f"mac_lin_bax   : {mac_lin_bax}")
+    print(f"mac_lin_bwax  : {mac_lin_backward_bax}")
+    print(f"mac_lin_wba   : {mac_lin_wba}")
+    print(f"mac_lin_bwba  : {mac_lin_backward_wba}")
+    print(f"conv w size   : {shape[0]**2*3**2}")
+    print(f"original_conv : {original_conv}")
+    print(f"mac_conv_bax  : {mac_conv_bax}")
+    print(f"mac_conv_wba  : {mac_conv_wba}")
+    print(f"mac_conv_bwba : {mac_conv_backward_wba}")
+    print()
diff --git a/experiments/loha_svd.py b/experiments/loha_svd.py
index 47803c9..8375a91 100644
--- a/experiments/loha_svd.py
+++ b/experiments/loha_svd.py
@@ -3,7 +3,7 @@
 
 # constants
 loha_dim = 16
-lora_dim = loha_dim*2 # same parameter count
+lora_dim = loha_dim * 2  # same parameter count
 weight_rank = 256
 
 # make norm=1 low rank weight for test
@@ -15,9 +15,14 @@
 e, s, v = torch.svd(test_weight)
 a = e[:, :lora_dim]
 b = torch.diag(s[:lora_dim]) @ v[:, :lora_dim].T
-diff = test_weight - a@b
+diff = test_weight - a @ b
 
-print(torch.mean(diff**2), torch.mean(diff), torch.max(torch.abs(diff)), torch.min(torch.abs(diff)))
+print(
+    torch.mean(diff**2),
+    torch.mean(diff),
+    torch.max(torch.abs(diff)),
+    torch.min(torch.abs(diff)),
+)
 
 
 # hadamard product
@@ -29,6 +34,11 @@
 
 e, s, v = torch.svd(b)
 b = e[:, :loha_dim] @ torch.diag(s[:loha_dim]) @ v[:, :loha_dim].T
-diff = test_weight - a*b
-
-print(torch.mean(diff**2), torch.mean(diff), torch.max(torch.abs(diff)), torch.min(torch.abs(diff)))
\ No newline at end of file
+diff = test_weight - a * b
+
+print(
+    torch.mean(diff**2),
+    torch.mean(diff),
+    torch.max(torch.abs(diff)),
+    torch.min(torch.abs(diff)),
+)
diff --git a/experiments/lokr_linear.py b/experiments/lokr_linear.py
index a9c2947..b919386 100644
--- a/experiments/lokr_linear.py
+++ b/experiments/lokr_linear.py
@@ -9,16 +9,16 @@
 def lokr_linear(h, a, b, c):
     vq = a.size(1)
     uq = c.size(1)
-    
-    h_in_group = rearrange(h, 'b ... (uq vq) -> b ... uq vq', uq=uq, vq=vq)
+
+    h_in_group = rearrange(h, "b ... (uq vq) -> b ... uq vq", uq=uq, vq=vq)
     ha = F.linear(h_in_group, a)
     hb = F.linear(ha, b)
-    
+
     h_cross_group = hb.transpose(-1, -2)
     hc = F.linear(h_cross_group, c)
-    
-    h = rearrange(hc, 'b ... vp up -> b ... (up vp)')
-    
+
+    h = rearrange(hc, "b ... vp up -> b ... (up vp)")
+
     return h
 
 
@@ -35,4 +35,4 @@ def lokr_rebuild(h, a, b, c):
 result1 = lokr_linear(h, a, b, c)
 result2 = lokr_rebuild(h, a, b, c)
 
-print(F.mse_loss(result1, result2))
\ No newline at end of file
+print(F.mse_loss(result1, result2))
diff --git a/experiments/module_count.py b/experiments/module_count.py
index f199632..ea93b02 100644
--- a/experiments/module_count.py
+++ b/experiments/module_count.py
@@ -6,10 +6,12 @@
 import lycoris.kohya as network_module
 
 
-clip, vae, unet = load_models_from_stable_diffusion_checkpoint(False, './test_model/ACertainty.ckpt')
-net = network_module.create_network(1, 1, 1, vae, clip, unet, preset='full')
-net = network_module.create_network(1, 1, 1, vae, clip, unet, preset='attn-mlp')
-net = network_module.create_network(1, 1, 1, vae, clip, unet, preset='attn-only')
+clip, vae, unet = load_models_from_stable_diffusion_checkpoint(
+    False, "./test_model/ACertainty.ckpt"
+)
+net = network_module.create_network(1, 1, 1, vae, clip, unet, preset="full")
+net = network_module.create_network(1, 1, 1, vae, clip, unet, preset="attn-mlp")
+net = network_module.create_network(1, 1, 1, vae, clip, unet, preset="attn-only")
 
 total = 0
 for name, module in unet.named_modules():
@@ -26,4 +28,4 @@
     if isinstance(module, nn.Conv2d):
         total += 1
 
-print(total)
\ No newline at end of file
+print(total)
diff --git a/experiments/weakre_speed.py b/experiments/weakre_speed.py
index 3d6eef1..e480226 100644
--- a/experiments/weakre_speed.py
+++ b/experiments/weakre_speed.py
@@ -5,10 +5,11 @@
 class data:
     data = 985761298734287
 
+
 o = data()
 d1 = [o]
 d2 = ref(o)
 
 
-print(timeit('d1[0]', globals=globals(), number=10000000))
-print(timeit('d2()', globals=globals(), number=10000000))
\ No newline at end of file
+print(timeit("d1[0]", globals=globals(), number=10000000))
+print(timeit("d2()", globals=globals(), number=10000000))
diff --git a/hyperdream_gen_weight.py b/hyperdream_gen_weight.py
index 0c09d86..1640b95 100644
--- a/hyperdream_gen_weight.py
+++ b/hyperdream_gen_weight.py
@@ -1,4 +1,5 @@
 import os, sys
+
 sys.path.insert(0, os.getcwd())
 import argparse
 
@@ -6,44 +7,46 @@
 def get_args():
     parser = argparse.ArgumentParser()
     parser.add_argument(
-        "base_model", help="The model you want to build hypernetwork on it",
-        default='', type=str
-    )
-    parser.add_argument(
-        "hyperdream_model", help="the hypernetwork you want to use to generate lora weight",
-        default='', type=str
+        "base_model",
+        help="The model you want to build hypernetwork on it",
+        default="",
+        type=str,
     )
     parser.add_argument(
-        "image_path", help="the image for generate weight",
-        default='', type=str
+        "hyperdream_model",
+        help="the hypernetwork you want to use to generate lora weight",
+        default="",
+        type=str,
     )
     parser.add_argument(
-        "output_name", help="the output model",
-        default='./out.pt', type=str
+        "image_path", help="the image for generate weight", default="", type=str
     )
     parser.add_argument(
-        "--is_v2", help="Your base model is sd v2 or not",
-        default=False, action="store_true"
+        "output_name", help="the output model", default="./out.pt", type=str
     )
     parser.add_argument(
-        "--device", help="Which device you want to use to merge the weight",
-        default='cpu', type=str
+        "--is_v2",
+        help="Your base model is sd v2 or not",
+        default=False,
+        action="store_true",
     )
     parser.add_argument(
-        "--dtype", help='dtype to save',
-        default='float', type=str
+        "--device",
+        help="Which device you want to use to merge the weight",
+        default="cpu",
+        type=str,
     )
+    parser.add_argument("--dtype", help="dtype to save", default="float", type=str)
     parser.add_argument(
-        "--weight", help='weight for the lyco model to merge',
-        default='1.0', type=float
+        "--weight", help="weight for the lyco model to merge", default="1.0", type=float
     )
     return parser.parse_args()
+
+
 ARGS = get_args()
 
 
-from lycoris.kohya.model_utils import (
-    load_file
-)
+from lycoris.kohya.model_utils import load_file
 from library.model_util import load_models_from_stable_diffusion_checkpoint
 from lycoris.kohya import create_hypernetwork
 
@@ -58,70 +61,81 @@ def get_args():
 
 @torch.no_grad()
 def main():
-    te, vae, unet = load_models_from_stable_diffusion_checkpoint(ARGS.is_v2, ARGS.base_model)
-    if ARGS.hyperdream_model.rsplit('.', 1)[-1] == 'safetensors':
+    te, vae, unet = load_models_from_stable_diffusion_checkpoint(
+        ARGS.is_v2, ARGS.base_model
+    )
+    if ARGS.hyperdream_model.rsplit(".", 1)[-1] == "safetensors":
         lyco = load_file(ARGS.hyperdream_model)
     else:
         lyco = torch.load(ARGS.hyperdream_model)
-    
-    dtype_str = ARGS.dtype.replace('fp', 'float').replace('loat', '').strip()
+
+    dtype_str = ARGS.dtype.replace("fp", "float").replace("loat", "").strip()
     dtype = {
-        'f': torch.float,
-        'f16': torch.float16,
-        'f32': torch.float32,
-        'f64': torch.float64,
-        'bf': torch.bfloat16,
-        'bf16': torch.bfloat16,
+        "f": torch.float,
+        "f16": torch.float16,
+        "f32": torch.float32,
+        "f64": torch.float64,
+        "bf": torch.bfloat16,
+        "bf16": torch.bfloat16,
     }.get(dtype_str, None)
     if dtype is None:
         raise ValueError(f'Cannot Find the dtype "{ARGS.dtype}"')
-    
-    hyperdream = create_hypernetwork(
-        1.0, 4, 1, vae, te, unet, down_dim=128, up_dim=64, delta_iters = 4, decoder_blocks=4,
-    ).to(dtype).to(ARGS.device)
+
+    hyperdream = (
+        create_hypernetwork(
+            1.0,
+            4,
+            1,
+            vae,
+            te,
+            unet,
+            down_dim=128,
+            up_dim=64,
+            delta_iters=4,
+            decoder_blocks=4,
+        )
+        .to(dtype)
+        .to(ARGS.device)
+    )
     (missing_keys, unexpected_keys) = hyperdream.load_state_dict(lyco, strict=False)
-    
+
     if unexpected_keys:
-        print(f'Unexpected keys: {unexpected_keys}')
-    
+        print(f"Unexpected keys: {unexpected_keys}")
+
     for k in missing_keys:
-        if (k=='checkpoint'
-            or 'block_pos_emb' in k
-            or 'encoder_model' in k):
+        if k == "checkpoint" or "block_pos_emb" in k or "encoder_model" in k:
             continue
-        assert False, f'Cannot find {k}, only keys in encoder_model can be missing'
-    
-    ref_img = Image.open(
-        ARGS.image_path
-    ).convert('RGB')
+        assert False, f"Cannot find {k}, only keys in encoder_model can be missing"
+
+    ref_img = Image.open(ARGS.image_path).convert("RGB")
     ref_img = resize(ref_img, hyperdream.weight_generater.ref_size)
     ref_img = to_tensor(ref_img).unsqueeze(0).to(dtype).to(ARGS.device) * 2 - 1
-    
+
     with torch.autocast(ARGS.device, dtype=dtype):
         hyperdream.update_reference(ref_img, 7)
-    
+
     state_dict = {}
     for lora in hyperdream.loras:
         down, up = lora.make_lightweight(*lora.data)
-        if down.dim()==3:
+        if down.dim() == 3:
             down = down.mean(dim=0)
-        if up.dim()==3:
+        if up.dim() == 3:
             up = up.mean(dim=0)
-        
-        test_updown = up@down
+
+        test_updown = up @ down
         print(torch.norm(test_updown))
         if isinstance(lora.lora_down, nn.Conv2d):
             down = down.unsqueeze(-1).unsqueeze(-1)
             up = up.unsqueeze(-1).unsqueeze(-1)
-        state_dict[f'{lora.lora_name}.lora_down.weight'] = down
-        state_dict[f'{lora.lora_name}.lora_up.weight'] = up
-        state_dict[f'{lora.lora_name}.alpha'] = torch.tensor(1)
-    
-    if ARGS.output_name.endswith('.safetensors'):
+        state_dict[f"{lora.lora_name}.lora_down.weight"] = down
+        state_dict[f"{lora.lora_name}.lora_up.weight"] = up
+        state_dict[f"{lora.lora_name}.alpha"] = torch.tensor(1)
+
+    if ARGS.output_name.endswith(".safetensors"):
         save_file(state_dict, ARGS.output_name)
     else:
         torch.save(state_dict, ARGS.output_name)
 
 
-if __name__ == '__main__':
-    main()
\ No newline at end of file
+if __name__ == "__main__":
+    main()
diff --git a/lycoris/__init__.py b/lycoris/__init__.py
index 79b2065..599047c 100644
--- a/lycoris/__init__.py
+++ b/lycoris/__init__.py
@@ -1,5 +1,21 @@
 from . import (
     kohya,
     modules,
-    utils
-)
\ No newline at end of file
+    utils,
+)
+
+from .modules.locon import LoConModule
+from .modules.loha import LohaModule
+from .modules.lokr import LokrModule
+from .modules.dylora import DyLoraModule
+from .modules.glora import GLoRAModule
+from .modules.norms import NormModule
+from .modules.full import FullModule
+from .modules.diag_oft import DiagOFTModule
+from .modules import make_module
+
+from .wrapper import (
+    LycorisNetwork,
+    create_lycoris,
+    create_lycoris_from_weights,
+)
diff --git a/lycoris/config.py b/lycoris/config.py
index 763db55..fc22a65 100644
--- a/lycoris/config.py
+++ b/lycoris/config.py
@@ -1,98 +1,86 @@
 PRESET = {
-    'full':{
-        'enable_conv': True,
-        'unet_target_module':[
-            "Transformer2DModel", 
-            "ResnetBlock2D", 
-            "Downsample2D", 
-            "Upsample2D"
+    "full": {
+        "enable_conv": True,
+        "unet_target_module": [
+            "Transformer2DModel",
+            "ResnetBlock2D",
+            "Downsample2D",
+            "Upsample2D",
         ],
-        'unet_target_name':[
+        "unet_target_name": [
             "conv_in",
             "conv_out",
             "time_embedding.linear_1",
             "time_embedding.linear_2",
         ],
-        'text_encoder_target_module': [
-            "CLIPAttention", "CLIPMLP"
-        ],
-        'text_encoder_target_name': [],
+        "text_encoder_target_module": ["CLIPAttention", "CLIPMLP"],
+        "text_encoder_target_name": [],
     },
-    'full-lin':{
-        'enable_conv': False,
-        'unet_target_module':[
-            "Transformer2DModel", 
+    "full-lin": {
+        "enable_conv": False,
+        "unet_target_module": [
+            "Transformer2DModel",
             "ResnetBlock2D",
         ],
-        'unet_target_name':[
+        "unet_target_name": [
             "time_embedding.linear_1",
             "time_embedding.linear_2",
         ],
-        'text_encoder_target_module': [
-            "CLIPAttention", "CLIPMLP"
-        ],
-        'text_encoder_target_name': [],
+        "text_encoder_target_module": ["CLIPAttention", "CLIPMLP"],
+        "text_encoder_target_name": [],
     },
-    'attn-mlp':{
-        'enable_conv': False,
-        'unet_target_module':[
-            "Transformer2DModel", 
-        ],
-        'unet_target_name':[],
-        'text_encoder_target_module': [
-            "CLIPAttention", "CLIPMLP"
+    "attn-mlp": {
+        "enable_conv": False,
+        "unet_target_module": [
+            "Transformer2DModel",
         ],
-        'text_encoder_target_name': [],
+        "unet_target_name": [],
+        "text_encoder_target_module": ["CLIPAttention", "CLIPMLP"],
+        "text_encoder_target_name": [],
     },
-    'attn-only':{
-        'enable_conv': False,
-        'unet_target_module':[
-            "CrossAttention", 
+    "attn-only": {
+        "enable_conv": False,
+        "unet_target_module": [
+            "CrossAttention",
         ],
-        'unet_target_name':[],
-        'text_encoder_target_module': [
-            "CLIPAttention"
-        ],
-        'text_encoder_target_name': [],
+        "unet_target_name": [],
+        "text_encoder_target_module": ["CLIPAttention"],
+        "text_encoder_target_name": [],
     },
-    'unet-only':{
-        'enable_conv': True,
-        'unet_target_module':[
-            "Transformer2DModel", 
-            "ResnetBlock2D", 
-            "Downsample2D", 
-            "Upsample2D"
+    "unet-only": {
+        "enable_conv": True,
+        "unet_target_module": [
+            "Transformer2DModel",
+            "ResnetBlock2D",
+            "Downsample2D",
+            "Upsample2D",
         ],
-        'unet_target_name':[
+        "unet_target_name": [
             "conv_in",
             "conv_out",
             "time_embedding.linear_1",
             "time_embedding.linear_2",
         ],
-        'text_encoder_target_module': [],
-        'text_encoder_target_name': [],
+        "text_encoder_target_module": [],
+        "text_encoder_target_name": [],
     },
-    'unet-transformer-only':{
-        'enable_conv': False,
-        'unet_target_module':[
-            "Transformer2DModel", 
+    "unet-transformer-only": {
+        "enable_conv": False,
+        "unet_target_module": [
+            "Transformer2DModel",
         ],
-        'unet_target_name':[],
-        'text_encoder_target_module': [],
-        'text_encoder_target_name': [],
+        "unet_target_name": [],
+        "text_encoder_target_module": [],
+        "text_encoder_target_name": [],
     },
-    'unet-convblock-only':{
-        'enable_conv': True,
-        'unet_target_module':[
-            "ResnetBlock2D", 
-            "Downsample2D", 
-            "Upsample2D"
-        ],
-        'unet_target_name':[
+    "unet-convblock-only": {
+        "enable_conv": True,
+        "unet_target_module": ["ResnetBlock2D", "Downsample2D", "Upsample2D"],
+        "unet_target_name": [
             "conv_in",
             "conv_out",
         ],
-        'text_encoder_target_module': [],
-        'text_encoder_target_name': [],
+        "text_encoder_target_module": [],
+        "text_encoder_target_name": [],
     },
 }
diff --git a/lycoris/hcp/__init__.py b/lycoris/hcp/__init__.py
new file mode 100644
index 0000000..9bd7dbf
--- /dev/null
+++ b/lycoris/hcp/__init__.py
@@ -0,0 +1,4 @@
+from .base import LycorisPluginContainer, LycorisPluginBlock
+from .loha import LohaBlock
+from .lokr import LokrBlock
+from .diag_oft import DiagOFTBlock
diff --git a/lycoris/hcp/base.py b/lycoris/hcp/base.py
new file mode 100644
index 0000000..937d536
--- /dev/null
+++ b/lycoris/hcp/base.py
@@ -0,0 +1,153 @@
+from typing import Tuple, Callable, Any, Optional
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from diffusers.models.lora import LoRACompatibleLinear, LoRACompatibleConv
+
+try:
+    import hcpdiff
+    from hcpdiff.models.plugin import PatchPluginContainer, PatchPluginBlock
+    from hcpdiff.utils.utils import isinstance_list
+except ImportError:
+    print(
+        "HCP-Diffusion is not installed, hcp feature will be disabled. "
+        "Use `pip install hcpdiff` to install hcpdiff."
+    )
+
+from ..utils.logger import logger
+
+
+class LycorisPluginContainer(PatchPluginContainer):
+    def __init__(self, host_name: str, host: nn.Module, parent_block: nn.Module):
+        super(LycorisPluginContainer, self).__init__(host_name, host, parent_block)
+
+        self.op = None
+        self.module_type = ""
+        self.extra_args = {}
+        if isinstance(host, nn.Conv2d):
+            self.op = F.conv2d
+            self.module_type = "conv"
+            self.extra_args = {
+                "stride": host.stride,
+                "padding": host.padding,
+                "dilation": host.dilation,
+                "groups": host.groups,
+            }
+        elif isinstance(host, nn.Linear):
+            self.op = F.linear
+            self.module_type = "linear"
+            self.extra_args = {}
+        elif isinstance(host, nn.LayerNorm):
+            self.op = F.layer_norm
+            self.module_type = "norm"
+            self.extra_args = {
+                "normalized_shape": host.normalized_shape,
+                "eps": host.eps,
+            }
+        elif isinstance(host, nn.GroupNorm):
+            self.op = F.group_norm
+            self.module_type = "norm"
+            self.extra_args = {"num_groups": host.num_groups, "eps": host.eps}
+        else:
+            logger.warning(
+                f"Unsupported host type: '{type(host)}' for LyCORIS plugin. "
+                "Will ignore this block."
+            )
+
+    def forward(self, *args, **kwargs):
+        if self.op is None:
+            return self._host(*args, **kwargs)
+
+        if (
+            isinstance(self._host, (LoRACompatibleLinear, LoRACompatibleConv))
+            and len(args) == 2
+        ):
+            scale = args[1]
+            args = [args[0]]
+        else:
+            scale = 1
+
+        org_weight = self._host.weight
+        org_bias = getattr(self._host, "bias", None)
+        new_weight = org_weight
+        new_bias = org_bias
+        total_diff_output = 0
+
+        for name in self.plugin_names:
+            lyco_plugin_block = getattr(self, name)
+            (
+                diff_weight,
+                diff_bias,
+                diff_output,
+                update_weight,
+                update_bias,
+            ) = lyco_plugin_block(
+                org_weight, org_bias, new_weight, new_bias, *args, **kwargs
+            )
+            # logger.debug(f'diff weight norm: {torch.norm(diff_weight).detach().cpu().item()}')
+            if diff_weight is not None:
+                if update_weight:
+                    new_weight = new_weight + diff_weight * scale
+                else:
+                    new_weight = diff_weight
+            if diff_bias is not None:
+                if update_bias:
+                    new_bias = (
+                        diff_bias * scale
+                        if new_bias is None
+                        else new_bias + diff_bias * scale
+                    )
+                else:
+                    new_bias = diff_bias
+            if diff_output is not None:
+                total_diff_output = total_diff_output + diff_output * scale
+
+        weight_dict = {
+            "weight": new_weight,
+            "bias": new_bias,
+        }
+        output = self.op(*args, **weight_dict, **self.extra_args)
+        return output + total_diff_output
+
+
+class LycorisPluginBlock(PatchPluginBlock):
+    container_cls = LycorisPluginContainer
+
+    def __init__(
+        self,
+        *args,
+        dim=4,
+        alpha=1.0,
+        dropout=0.0,
+        rank_dropout=0.0,
+        module_dropout=0.0,
+        use_tucker=False,
+        **kwargs,
+    ):
+        super(LycorisPluginBlock, self).__init__(*args, **kwargs)
+        self.module_type = self.container().module_type
+        self.dim = dim
+        self.alpha = alpha
+        self.dropout = dropout
+        self.rank_dropout = rank_dropout
+        self.module_dropout = module_dropout
+        self.tucker = False
+
+        if self.module_type == "conv":
+            self.out_dim = self.host().out_channels
+            self.in_dim = self.host().in_channels
+            self.k_size = self.host().kernel_size
+            self.tucker = use_tucker
+            self.shape = [self.out_dim, self.in_dim, *self.k_size]
+        elif self.module_type == "linear":
+            self.out_dim = self.host().out_features
+            self.in_dim = self.host().in_features
+            self.shape = [self.out_dim, self.in_dim]
+        elif self.module_type == "norm":
+            self.dim = self.host().weight.shape[0]
+        else:
+            raise NotImplementedError
+
+    def forward(self, org_weight, org_bias, new_weight, new_bias, *args, **kwargs):
+        raise NotImplementedError
diff --git a/lycoris/hcp/diag_oft.py b/lycoris/hcp/diag_oft.py
new file mode 100644
index 0000000..fa7ba6e
--- /dev/null
+++ b/lycoris/hcp/diag_oft.py
@@ -0,0 +1,54 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from einops import rearrange
+
+from .base import LycorisPluginBlock
+from ..modules.lokr import factorization
+from ..modules.diag_oft import DiagOFTModule
+
+
+class DiagOFTBlock(DiagOFTModule, LycorisPluginBlock):
+    wrapable_classes = (nn.Conv2d, nn.Linear)
+
+    def __init__(self, *args, constrain=0, rescaled=False, **kwargs):
+        LycorisPluginBlock.__init__(self, *args, **kwargs)
+
+        out_dim = self.shape[0]
+        self.block_size, self.block_num = factorization(out_dim, self.dim)
+        self.scale = 1.0
+        # block_num > block_size
+        self.rescaled = rescaled
+        self.constrain = constrain * out_dim
+        self.oft_blocks = nn.Parameter(
+            torch.zeros(self.block_num, self.block_size, self.block_size)
+        )
+        if rescaled:
+            self.rescale = nn.Parameter(torch.ones(self.block_num, self.block_size, 1))
+
+    def forward(self, orig_weight, org_bias, new_weight, new_bias, *args, **kwargs):
+        device = self.oft_blocks.device
+        if self.rank_dropout and self.training:
+            drop = (torch.rand(self.oft_blocks, device=device) < self.rank_dropout).to(
+                self.oft_blocks.dtype
+            )
+            if self.rank_dropout_scale:
+                drop /= drop.mean()
+        else:
+            drop = 1
+
+        r = self.get_r()
+        org_weight = orig_weight.to(device, dtype=r.dtype)
+        org_weight = rearrange(
+            org_weight, "(k n) ... -> k n ...", k=self.block_num, n=self.block_size
+        )
+        # Init R=0, so add I on it to ensure the output of step0 is original model output
+        weight = torch.einsum(
+            "k n m, k n ... -> k m ...",
+            r * self.scale + (1 - self.scale) * self.I,
+            org_weight,
+        )
+        weight = rearrange(weight, "k m ... -> (k m) ...")
+        # disable update weight for replace the weight directly
+        return weight, None, None, False, False
diff --git a/lycoris/hcp/loha.py b/lycoris/hcp/loha.py
new file mode 100644
index 0000000..63c1b23
--- /dev/null
+++ b/lycoris/hcp/loha.py
@@ -0,0 +1,85 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .base import LycorisPluginBlock
+from ..modules.loha import make_weight, make_weight_cp
+
+
+class LohaBlock(LycorisPluginBlock):
+    wrapable_classes = (nn.Conv2d, nn.Linear)
+
+    def __init__(self, *args, **kwargs):
+        super(LohaBlock, self).__init__(*args, **kwargs)
+
+        if self.tucker:
+            self.hada_t1 = nn.Parameter(
+                torch.empty(self.dim, self.dim, self.shape[2], self.shape[3])
+            )
+            self.hada_w1_a = nn.Parameter(
+                torch.empty(self.dim, self.shape[0])
+            )  # out_dim, 1-mode
+            self.hada_w1_b = nn.Parameter(
+                torch.empty(self.dim, self.shape[1])
+            )  # in_dim , 2-mode
+
+            self.hada_t2 = nn.Parameter(
+                torch.empty(self.dim, self.dim, self.shape[2], self.shape[3])
+            )
+            self.hada_w2_a = nn.Parameter(
+                torch.empty(self.dim, self.shape[0])
+            )  # out_dim, 1-mode
+            self.hada_w2_b = nn.Parameter(
+                torch.empty(self.dim, self.shape[1])
+            )  # in_dim , 2-mode
+        else:
+            self.hada_w1_a = nn.Parameter(torch.empty(self.shape[0], self.dim))
+            self.hada_w1_b = nn.Parameter(torch.empty(self.dim, self.shape[1]))
+
+            self.hada_w2_a = nn.Parameter(torch.empty(self.shape[0], self.dim))
+            self.hada_w2_b = nn.Parameter(torch.empty(self.dim, self.shape[1]))
+
+        if type(self.alpha) == torch.Tensor:
+            self.alpha = (
+                self.alpha.detach().float().numpy()
+            )  # without casting, bf16 causes error
+        alpha = self.dim if self.alpha is None or self.alpha == 0 else self.alpha
+        self.scale = alpha / self.dim
+        delattr(self, "alpha")
+        self.register_buffer("alpha", torch.tensor(alpha))  # 定数として扱える
+
+        if self.tucker:
+            torch.nn.init.normal_(self.hada_t1, std=0.1)
+            torch.nn.init.normal_(self.hada_t2, std=0.1)
+        torch.nn.init.normal_(self.hada_w1_b, std=1)
+        torch.nn.init.normal_(self.hada_w1_a, std=0.1)
+        torch.nn.init.normal_(self.hada_w2_b, std=1)
+        torch.nn.init.constant_(self.hada_w2_a, 0)
+
+    def forward(self, orig_weight, org_bias, new_weight, new_bias, *args, **kwargs):
+        if self.tucker:
+            weight = make_weight_cp(
+                self.hada_t1,
+                self.hada_w1_a,
+                self.hada_w1_b,
+                self.hada_t2,
+                self.hada_w2_a,
+                self.hada_w2_b,
+                scale=torch.tensor(self.scale),
+            )
+        else:
+            weight = make_weight(
+                self.hada_w1_a,
+                self.hada_w1_b,
+                self.hada_w2_a,
+                self.hada_w2_b,
+                scale=torch.tensor(self.scale),
+            )
+        if orig_weight is not None:
+            weight = weight.reshape(orig_weight.shape)
+        if self.training and self.rank_dropout:
+            drop = (torch.rand(weight.size(0)) < self.rank_dropout).to(weight.dtype)
+            drop = drop.view(-1, *[1] * len(weight.shape[1:])).to(weight.device)
+            drop /= drop.mean()
+            weight *= drop
+        return weight, None, None, True, True
diff --git a/lycoris/hcp/lokr.py b/lycoris/hcp/lokr.py
new file mode 100644
index 0000000..e093cde
--- /dev/null
+++ b/lycoris/hcp/lokr.py
@@ -0,0 +1,127 @@
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .base import LycorisPluginBlock
+from ..modules.lokr import make_kron, make_weight_cp, factorization
+
+
+class LokrBlock(LycorisPluginBlock):
+    wrapable_classes = (nn.Conv2d, nn.Linear)
+
+    def __init__(self, *args, factor=-1, decompose_both=False, **kwargs):
+        super(LokrBlock, self).__init__(*args, **kwargs)
+
+        in_m, in_n = factorization(self.shape[1], factor)
+        out_l, out_k = factorization(self.shape[0], factor)
+
+        self.shape[0] = (out_l, out_k)
+        self.shape[1] = (in_m, in_n)
+
+        if decompose_both and self.dim < max(self.shape[0][0], self.shape[1][0]) / 2:
+            self.lokr_w1_a = nn.Parameter(torch.empty(self.shape[0][0], self.dim))
+            self.lokr_w1_b = nn.Parameter(torch.empty(self.dim, self.shape[1][0]))
+        else:
+            self.use_w1 = True
+            self.lokr_w1 = nn.Parameter(
+                torch.empty(self.shape[0][0], self.shape[1][0])
+            )  # a*c, 1-mode
+
+        if self.module_type == "conv":
+            if self.dim >= max(self.shape[0][1], self.shape[1][1]) / 2:
+                self.use_w2 = True
+                self.lokr_w2 = nn.Parameter(
+                    torch.empty(self.shape[0][1], self.shape[1][1], *self.k_size)
+                )
+            elif self.tucker:
+                self.lokr_t2 = nn.Parameter(
+                    torch.empty(self.dim, self.dim, self.shape[2], self.shape[3])
+                )
+                self.lokr_w2_a = nn.Parameter(
+                    torch.empty(self.dim, self.shape[0][1])
+                )  # b, 1-mode
+                self.lokr_w2_b = nn.Parameter(
+                    torch.empty(self.dim, self.shape[1][1])
+                )  # d, 2-mode
+            else:  # Conv2d not cp
+                # bigger part. weight and LoRA. [b, dim] x [dim, d*k1*k2]
+                self.lokr_w2_a = nn.Parameter(torch.empty(self.shape[0][1], self.dim))
+                self.lokr_w2_b = nn.Parameter(
+                    torch.empty(
+                        self.dim, self.shape[1][1] * self.shape[2] * self.shape[3]
+                    )
+                )
+                # w1 ⊗ (w2_a x w2_b) = (a, b)⊗((c, dim)x(dim, d*k1*k2)) = (a, b)⊗(c, d*k1*k2) = (ac, bd*k1*k2)
+        elif self.module_type == "linear":
+            if self.dim < max(self.shape[0][1], self.shape[1][1]) / 2:
+                # bigger part. weight and LoRA. [b, dim] x [dim, d]
+                self.lokr_w2_a = nn.Parameter(torch.empty(self.shape[0][1], self.dim))
+                self.lokr_w2_b = nn.Parameter(torch.empty(self.dim, self.shape[1][1]))
+                # w1 ⊗ (w2_a x w2_b) = (a, b)⊗((c, dim)x(dim, d)) = (a, b)⊗(c, d) = (ac, bd)
+            else:
+                self.use_w2 = True
+                self.lokr_w2 = nn.Parameter(
+                    torch.empty(self.shape[0][1], self.shape[1][1])
+                )
+        else:
+            raise NotImplementedError
+
+        if self.use_w2:
+            torch.nn.init.constant_(self.lokr_w2, 0)
+        else:
+            if self.tucker:
+                torch.nn.init.kaiming_uniform_(self.lokr_t2, a=math.sqrt(5))
+            torch.nn.init.kaiming_uniform_(self.lokr_w2_a, a=math.sqrt(5))
+            torch.nn.init.constant_(self.lokr_w2_b, 0)
+
+        if self.use_w1:
+            torch.nn.init.kaiming_uniform_(self.lokr_w1, a=math.sqrt(5))
+        else:
+            torch.nn.init.kaiming_uniform_(self.lokr_w1_a, a=math.sqrt(5))
+            torch.nn.init.kaiming_uniform_(self.lokr_w1_b, a=math.sqrt(5))
+
+        if type(self.alpha) == torch.Tensor:
+            self.alpha = (
+                self.alpha.detach().float().numpy()
+            )  # without casting, bf16 causes error
+        alpha = self.dim if self.alpha is None or self.alpha == 0 else self.alpha
+        self.scale = alpha / self.dim
+        print(self.scale)
+        delattr(self, "alpha")
+        self.register_buffer("alpha", torch.tensor(alpha))  # 定数として扱える
+
+        weight = make_kron(
+            self.lokr_w1 if self.use_w1 else self.lokr_w1_a @ self.lokr_w1_b,
+            (
+                self.lokr_w2
+                if self.use_w2
+                else make_weight_cp(self.lokr_t2, self.lokr_w2_a, self.lokr_w2_b)
+                if self.cp
+                else self.lokr_w2_a @ self.lokr_w2_b
+            ),
+            torch.tensor(self.scale),
+        )
+        assert torch.sum(torch.isnan(weight)) == 0, "lokr init weight is nan"
+
+    def forward(self, orig_weight, org_bias, new_weight, new_bias, *args, **kwargs):
+        weight = make_kron(
+            self.lokr_w1 if self.use_w1 else self.lokr_w1_a @ self.lokr_w1_b,
+            (
+                self.lokr_w2
+                if self.use_w2
+                else make_weight_cp(self.lokr_t2, self.lokr_w2_a, self.lokr_w2_b)
+                if self.tucker
+                else self.lokr_w2_a @ self.lokr_w2_b
+            ),
+            torch.tensor(self.scale),
+        )
+        if orig_weight is not None:
+            weight = weight.reshape(orig_weight.shape)
+        if self.training and self.rank_dropout:
+            drop = (torch.rand(weight.size(0)) < self.rank_dropout).to(weight.dtype)
+            drop = drop.view(-1, *[1] * len(weight.shape[1:])).to(weight.device)
+            drop /= drop.mean()
+            weight *= drop
+        return weight, None, None, True, True
diff --git a/lycoris/kohya/__init__.py b/lycoris/kohya/__init__.py
index 4fd129b..7a64107 100644
--- a/lycoris/kohya/__init__.py
+++ b/lycoris/kohya/__init__.py
@@ -8,6 +8,7 @@
 from warnings import warn
 import os
 import sys
+
 sys.setrecursionlimit(10000)
 from typing import List
 import torch
@@ -22,103 +23,134 @@
 from ..modules.glora import GLoRAModule
 from ..modules.norms import NormModule
 from ..modules.full import FullModule
+from ..modules.diag_oft import DiagOFTModule
 from ..modules import make_module
 
 from ..config import PRESET
 from ..utils.preset import read_preset
-from ..utils import get_module
+from ..utils import get_module, str_bool
 
 
 network_module_dict = {
-    'lora': LoConModule,
-    'locon': LoConModule,
-    'loha': LohaModule,
-    'ia3':  IA3Module,
-    'lokr': LokrModule,
-    'dylora': DyLoraModule,
-    'glora': GLoRAModule,
-    'full': FullModule,
+    "lora": LoConModule,
+    "locon": LoConModule,
+    "loha": LohaModule,
+    "ia3": IA3Module,
+    "lokr": LokrModule,
+    "dylora": DyLoraModule,
+    "glora": GLoRAModule,
+    "full": FullModule,
+    "diag-oft": DiagOFTModule,
 }
 
 
-def create_network(multiplier, network_dim, network_alpha, vae, text_encoder, unet, **kwargs):
+def create_network(
+    multiplier, network_dim, network_alpha, vae, text_encoder, unet, **kwargs
+):
     if network_dim is None:
-        network_dim = 4                     # default
-    conv_dim = int(kwargs.get('conv_dim', network_dim) or network_dim)
-    conv_alpha = float(kwargs.get('conv_alpha', network_alpha) or network_alpha)
-    dropout = float(kwargs.get('dropout', 0.) or 0.)
-    rank_dropout = float(kwargs.get("rank_dropout", 0.) or 0.)
-    module_dropout = float(kwargs.get("module_dropout", 0.) or 0.)
-    algo = (kwargs.get('algo', 'lora') or 'lora').lower()
-    use_tucker = (not kwargs.get('disable_conv_cp', True)
-              or kwargs.get('use_conv_cp', False)
-              or kwargs.get('use_cp', False)
-              or kwargs.get('use_tucker', False))
-    if 'disable_conv_cp' in kwargs or 'use_cp' in kwargs or 'use_conv_cp' in kwargs:
-        warn("disable_conv_cp and use_cp are deprecated. Please use use_tucker instead.", stacklevel=2)
-    use_scalar = kwargs.get('use_scalar', False)
-    block_size = int(kwargs.get('block_size', 4) or 4)
-    train_norm = kwargs.get('train_norm', False)
-    
-    if algo == 'glora' and conv_dim>0:
+        network_dim = 4  # default
+    conv_dim = int(kwargs.get("conv_dim", network_dim) or network_dim)
+    conv_alpha = float(kwargs.get("conv_alpha", network_alpha) or network_alpha)
+    dropout = float(kwargs.get("dropout", 0.0) or 0.0)
+    rank_dropout = float(kwargs.get("rank_dropout", 0.0) or 0.0)
+    module_dropout = float(kwargs.get("module_dropout", 0.0) or 0.0)
+    algo = (kwargs.get("algo", "lora") or "lora").lower()
+    use_tucker = str_bool(
+        not kwargs.get("disable_conv_cp", True)
+        or kwargs.get("use_conv_cp", False)
+        or kwargs.get("use_cp", False)
+        or kwargs.get("use_tucker", False)
+    )
+    if "disable_conv_cp" in kwargs or "use_cp" in kwargs or "use_conv_cp" in kwargs:
+        warn(
+            "disable_conv_cp and use_cp are deprecated. Please use use_tucker instead.",
+            stacklevel=2,
+        )
+    use_scalar = str_bool(kwargs.get("use_scalar", False))
+    block_size = int(kwargs.get("block_size", 4) or 4)
+    train_norm = str_bool(kwargs.get("train_norm", False))
+    constrain = float(kwargs.get("constrain", 0) or 0)
+    rescaled = str_bool(kwargs.get("rescaled", False))
+
+    if algo == "glora" and conv_dim > 0:
         conv_dim = 0
-        print('Disable conv layer for GLoRA')
-    
-    preset = kwargs.get('preset', 'full')
+        print("Disable conv layer for GLoRA")
+
+    preset = kwargs.get("preset", "full")
     if preset not in PRESET:
         preset = read_preset(preset)
     else:
         preset = PRESET[preset]
     assert preset is not None
     LycorisNetwork.apply_preset(preset)
-    
-    print(f'Using rank adaptation algo: {algo}')
-    
-    if ((algo == 'loha')
-        and not kwargs.get('no_dim_warn', False) 
-        and (network_dim>64 or conv_dim>64)):
-        print('='*20 + 'WARNING' + '='*20)
-        warning_type ={
-            'loha': "Hadamard Product representation"
-        }
+
+    print(f"Using rank adaptation algo: {algo}")
+
+    if (
+        (algo == "loha")
+        and not kwargs.get("no_dim_warn", False)
+        and (network_dim > 64 or conv_dim > 64)
+    ):
+        print("=" * 20 + "WARNING" + "=" * 20)
+        warning_type = {"loha": "Hadamard Product representation"}
         warning_msg = f"""You are not supposed to use dim>64 (64*64 = 4096, it already has enough rank)\n
             in {warning_type[algo]}!\n
             Please consider use lower dim or disable this warning with --network_args no_dim_warn=True\n
             If you just want to use high dim {algo}, please consider use lower lr.
         """
         warn(warning_msg, stacklevel=2)
-        print('='*20 + 'WARNING' + '='*20)
-    
-    if algo == 'ia3':
+        print("=" * 20 + "WARNING" + "=" * 20)
+
+    if algo == "ia3":
         network = IA3Network(
-            text_encoder, unet,
-            multiplier = multiplier, 
+            text_encoder,
+            unet,
+            multiplier=multiplier,
         )
     else:
         network = LycorisNetwork(
-            text_encoder, unet, 
-            multiplier=multiplier, 
-            lora_dim=network_dim, conv_lora_dim=conv_dim, 
-            alpha=network_alpha, conv_alpha=conv_alpha,
-            dropout=dropout, rank_dropout=rank_dropout, module_dropout=module_dropout,
-            use_tucker=use_tucker, use_scalar=use_scalar,
-            network_module=algo, train_norm=train_norm,
-            decompose_both=kwargs.get('decompose_both', False),
-            factor=kwargs.get('factor', -1),
-            block_size = block_size
+            text_encoder,
+            unet,
+            multiplier=multiplier,
+            lora_dim=network_dim,
+            conv_lora_dim=conv_dim,
+            alpha=network_alpha,
+            conv_alpha=conv_alpha,
+            dropout=dropout,
+            rank_dropout=rank_dropout,
+            module_dropout=module_dropout,
+            use_tucker=use_tucker,
+            use_scalar=use_scalar,
+            network_module=algo,
+            train_norm=train_norm,
+            decompose_both=kwargs.get("decompose_both", False),
+            factor=kwargs.get("factor", -1),
+            block_size=block_size,
+            constrain=constrain,
+            rescaled=rescaled,
         )
-    
-    if algo=='dylora':
-        #dylora didn't support scale weight norm yet
-        delattr(type(network), 'apply_max_norm_regularization')
-    
+
+    if algo == "dylora":
+        # dylora didn't support scale weight norm yet
+        delattr(type(network), "apply_max_norm_regularization")
+
     return network
 
 
-def create_network_from_weights(multiplier, file, vae, text_encoder, unet, weights_sd=None, for_inference=False, **kwargs):
+def create_network_from_weights(
+    multiplier,
+    file,
+    vae,
+    text_encoder,
+    unet,
+    weights_sd=None,
+    for_inference=False,
+    **kwargs,
+):
     if weights_sd is None:
         if os.path.splitext(file)[1] == ".safetensors":
             from safetensors.torch import load_file, safe_open
+
             weights_sd = load_file(file)
         else:
             weights_sd = torch.load(file, map_location="cpu")
@@ -135,21 +167,23 @@ def create_network_from_weights(multiplier, file, vae, text_encoder, unet, weigh
             unet_loras[lora_name] = None
         elif lora_name.startswith(LycorisNetwork.LORA_PREFIX_TEXT_ENCODER):
             te_loras[lora_name] = None
-    
+
     for name, modules in unet.named_modules():
-        lora_name = f'{LycorisNetwork.LORA_PREFIX_UNET}_{name}'.replace('.','_')
+        lora_name = f"{LycorisNetwork.LORA_PREFIX_UNET}_{name}".replace(".", "_")
         if lora_name in unet_loras:
             unet_loras[lora_name] = modules
-    
+
     for name, modules in text_encoder.named_modules():
-        lora_name = f'{LycorisNetwork.LORA_PREFIX_TEXT_ENCODER}_{name}'.replace('.','_')
+        lora_name = f"{LycorisNetwork.LORA_PREFIX_TEXT_ENCODER}_{name}".replace(
+            ".", "_"
+        )
         if lora_name in te_loras:
             te_loras[lora_name] = modules
-    
+
     network = LycorisNetwork(text_encoder, unet)
     network.unet_loras = []
     network.text_encoder_loras = []
-    
+
     for lora_name, orig_modules in unet_loras.items():
         if orig_modules is None:
             continue
@@ -157,7 +191,7 @@ def create_network_from_weights(multiplier, file, vae, text_encoder, unet, weigh
         module = make_module(lyco_type, params, lora_name, orig_modules)
         if module is not None:
             network.unet_loras.append(module)
-    
+
     for lora_name, orig_modules in te_loras.items():
         if orig_modules is None:
             continue
@@ -165,64 +199,86 @@ def create_network_from_weights(multiplier, file, vae, text_encoder, unet, weigh
         module = make_module(lyco_type, params, lora_name, orig_modules)
         if module is not None:
             network.text_encoder_loras.append(module)
-    
+
     for lora in network.unet_loras + network.text_encoder_loras:
         lora.multiplier = multiplier
-    
+
     return network, weights_sd
 
 
-def create_hypernetwork(multiplier, network_dim, network_alpha, vae, text_encoder, unet, vocab_size, **kwargs):
+def create_hypernetwork(
+    multiplier,
+    network_dim,
+    network_alpha,
+    vae,
+    text_encoder,
+    unet,
+    vocab_size,
+    **kwargs,
+):
     if network_dim is None:
         network_dim = 4
-    dropout = float(kwargs.get('dropout', 0.) or 0.)
-    rank_dropout = float(kwargs.get("rank_dropout", 0.) or 0.)
-    module_dropout = float(kwargs.get("module_dropout", 0.) or 0.)
-    algo = (kwargs.get('algo', 'lora') or 'lora').lower()
-    use_tucker = (not kwargs.get('disable_conv_cp', True)
-              or kwargs.get('use_conv_cp', False)
-              or kwargs.get('use_cp', False)
-              or kwargs.get('use_tucker', False))
-    if 'disable_conv_cp' in kwargs or 'use_cp' in kwargs or 'use_conv_cp' in kwargs:
-        warn("disable_conv_cp and use_cp are deprecated. Please use use_tucker instead.", stacklevel=2)
-    block_size = int(kwargs.get('block_size', 4) or 4)
-    down_dim = int(kwargs.get('down_dim', 128) or 128)
-    up_dim = int(kwargs.get('up_dim', 64) or 64)
-    delta_iters = int(kwargs.get('delta_iters', 5) or 5)
-    decoder_blocks = int(kwargs.get('decoder_blocks', 4) or 4)
+    dropout = float(kwargs.get("dropout", 0.0) or 0.0)
+    rank_dropout = float(kwargs.get("rank_dropout", 0.0) or 0.0)
+    module_dropout = float(kwargs.get("module_dropout", 0.0) or 0.0)
+    algo = (kwargs.get("algo", "lora") or "lora").lower()
+    use_tucker = str_bool(
+        not kwargs.get("disable_conv_cp", True)
+        or kwargs.get("use_conv_cp", False)
+        or kwargs.get("use_cp", False)
+        or kwargs.get("use_tucker", False)
+    )
+    if "disable_conv_cp" in kwargs or "use_cp" in kwargs or "use_conv_cp" in kwargs:
+        warn(
+            "disable_conv_cp and use_cp are deprecated. Please use use_tucker instead.",
+            stacklevel=2,
+        )
+    block_size = int(kwargs.get("block_size", 4) or 4)
+    down_dim = int(kwargs.get("down_dim", 128) or 128)
+    up_dim = int(kwargs.get("up_dim", 64) or 64)
+    delta_iters = int(kwargs.get("delta_iters", 5) or 5)
+    decoder_blocks = int(kwargs.get("decoder_blocks", 4) or 4)
     network_module = {
-        'lora': LoConModule,
-        'locon': LoConModule,
+        "lora": LoConModule,
+        "locon": LoConModule,
     }[algo]
-    
-    print(f'Using rank adaptation algo: {algo}')
-    
+
+    print(f"Using rank adaptation algo: {algo}")
+
     return HyperDreamNetwork(
-        text_encoder, unet, 
-        multiplier=multiplier, 
-        lora_dim=network_dim, alpha=network_alpha,
+        text_encoder,
+        unet,
+        multiplier=multiplier,
+        lora_dim=network_dim,
+        alpha=network_alpha,
         use_tucker=use_tucker,
-        dropout=dropout, rank_dropout=rank_dropout, module_dropout=module_dropout,
+        dropout=dropout,
+        rank_dropout=rank_dropout,
+        module_dropout=module_dropout,
         network_module=network_module,
-        down_dim=down_dim, up_dim=up_dim, delta_iters=delta_iters, 
-        decoder_blocks=decoder_blocks, vocab_size=vocab_size,
-        decompose_both=kwargs.get('decompose_both', False),
-        factor=kwargs.get('factor', -1),
-        block_size = block_size
+        down_dim=down_dim,
+        up_dim=up_dim,
+        delta_iters=delta_iters,
+        decoder_blocks=decoder_blocks,
+        vocab_size=vocab_size,
+        decompose_both=kwargs.get("decompose_both", False),
+        factor=kwargs.get("factor", -1),
+        block_size=block_size,
     )
 
 
 class LycorisNetwork(torch.nn.Module):
-    '''
+    """
     LoRA + LoCon
-    '''
+    """
+
     # Ignore proj_in or proj_out, their channels is only a few.
     ENABLE_CONV = True
     UNET_TARGET_REPLACE_MODULE = [
-        "Transformer2DModel", 
-        "ResnetBlock2D", 
-        "Downsample2D", 
-        "Upsample2D"
+        "Transformer2DModel",
+        "ResnetBlock2D",
+        "Downsample2D",
+        "Upsample2D",
     ]
     UNET_TARGET_REPLACE_NAME = [
         "conv_in",
@@ -231,99 +287,111 @@ class LycorisNetwork(torch.nn.Module):
         "time_embedding.linear_2",
     ]
     TEXT_ENCODER_TARGET_REPLACE_MODULE = ["CLIPAttention", "CLIPMLP"]
-    LORA_PREFIX_UNET = 'lora_unet'
-    LORA_PREFIX_TEXT_ENCODER = 'lora_te'
+    LORA_PREFIX_UNET = "lora_unet"
+    LORA_PREFIX_TEXT_ENCODER = "lora_te"
     MODULE_ALGO_MAP = {}
     NAME_ALGO_MAP = {}
 
     @classmethod
     def apply_preset(cls, preset):
-        if 'enable_conv' in preset:
-            cls.ENABLE_CONV = preset['enable_conv']
-        if 'unet_target_module' in preset:
-            cls.UNET_TARGET_REPLACE_MODULE = preset['unet_target_module']
-        if 'unet_target_name' in preset:
-            cls.UNET_TARGET_REPLACE_NAME = preset['unet_target_name']
-        if 'text_encoder_target_module' in preset:
-            cls.TEXT_ENCODER_TARGET_REPLACE_MODULE = preset['text_encoder_target_module']
-        if 'text_encoder_target_name' in preset:
-            cls.TEXT_ENCODER_TARGET_REPLACE_NAME = preset['text_encoder_target_name']
-        if 'module_algo_map' in preset:
-            cls.MODULE_ALGO_MAP = preset['module_algo_map']
-        if 'name_algo_map' in preset:
-            cls.NAME_ALGO_MAP = preset['name_algo_map']
+        if "enable_conv" in preset:
+            cls.ENABLE_CONV = preset["enable_conv"]
+        if "unet_target_module" in preset:
+            cls.UNET_TARGET_REPLACE_MODULE = preset["unet_target_module"]
+        if "unet_target_name" in preset:
+            cls.UNET_TARGET_REPLACE_NAME = preset["unet_target_name"]
+        if "text_encoder_target_module" in preset:
+            cls.TEXT_ENCODER_TARGET_REPLACE_MODULE = preset[
+                "text_encoder_target_module"
+            ]
+        if "text_encoder_target_name" in preset:
+            cls.TEXT_ENCODER_TARGET_REPLACE_NAME = preset["text_encoder_target_name"]
+        if "module_algo_map" in preset:
+            cls.MODULE_ALGO_MAP = preset["module_algo_map"]
+        if "name_algo_map" in preset:
+            cls.NAME_ALGO_MAP = preset["name_algo_map"]
         return cls
 
     def __init__(
-        self, 
-        text_encoder, unet, 
-        multiplier=1.0, 
-        lora_dim=4, conv_lora_dim=4, 
-        alpha=1, conv_alpha=1,
-        use_tucker = False,
-        dropout = 0, rank_dropout = 0, module_dropout = 0,
-        network_module:str = 'locon', 
-        norm_modules = NormModule, train_norm = False,
+        self,
+        text_encoder,
+        unet,
+        multiplier=1.0,
+        lora_dim=4,
+        conv_lora_dim=4,
+        alpha=1,
+        conv_alpha=1,
+        use_tucker=False,
+        dropout=0,
+        rank_dropout=0,
+        module_dropout=0,
+        network_module: str = "locon",
+        norm_modules=NormModule,
+        train_norm=False,
         **kwargs,
     ) -> None:
         super().__init__()
         root_kwargs = kwargs
         self.multiplier = multiplier
         self.lora_dim = lora_dim
-        
+
         if not self.ENABLE_CONV:
             conv_lora_dim = 0
-        
+
         self.conv_lora_dim = int(conv_lora_dim)
-        if self.conv_lora_dim and self.conv_lora_dim != self.lora_dim: 
-            print('Apply different lora dim for conv layer')
-            print(f'Conv Dim: {conv_lora_dim}, Linear Dim: {lora_dim}')
+        if self.conv_lora_dim and self.conv_lora_dim != self.lora_dim:
+            print("Apply different lora dim for conv layer")
+            print(f"Conv Dim: {conv_lora_dim}, Linear Dim: {lora_dim}")
         elif self.conv_lora_dim == 0:
-            print('Disable conv layer')
-        
+            print("Disable conv layer")
+
         self.alpha = alpha
         self.conv_alpha = float(conv_alpha)
-        if self.conv_lora_dim and self.alpha != self.conv_alpha: 
-            print('Apply different alpha value for conv layer')
-            print(f'Conv alpha: {conv_alpha}, Linear alpha: {alpha}')
-        
+        if self.conv_lora_dim and self.alpha != self.conv_alpha:
+            print("Apply different alpha value for conv layer")
+            print(f"Conv alpha: {conv_alpha}, Linear alpha: {alpha}")
+
         if 1 >= dropout >= 0:
-            print(f'Use Dropout value: {dropout}')
+            print(f"Use Dropout value: {dropout}")
         self.dropout = dropout
         self.rank_dropout = rank_dropout
         self.module_dropout = module_dropout
-        
+
         self.use_tucker = use_tucker
-        
+
         def create_single_module(
-            lora_name: str, 
-            module: torch.nn.Module, 
+            lora_name: str,
+            module: torch.nn.Module,
             algo_name,
-            dim = None,
-            alpha = None,
-            use_tucker = self.use_tucker,
+            dim=None,
+            alpha=None,
+            use_tucker=self.use_tucker,
             **kwargs,
         ):
             for k, v in root_kwargs.items():
-                if k in kwargs: continue
+                if k in kwargs:
+                    continue
                 kwargs[k] = v
-            
-            if train_norm and 'Norm' in module.__class__.__name__:
+
+            if train_norm and "Norm" in module.__class__.__name__:
                 return norm_modules(
-                    lora_name, module, self.multiplier, 
-                    self.rank_dropout, self.module_dropout, 
-                    **kwargs
+                    lora_name,
+                    module,
+                    self.multiplier,
+                    self.rank_dropout,
+                    self.module_dropout,
+                    **kwargs,
                 )
             lora = None
-            if module.__class__.__name__ == 'Linear' and lora_dim>0:
+            if module.__class__.__name__ == "Linear" and lora_dim > 0:
                 dim = dim or lora_dim
                 alpha = alpha or self.alpha
-            elif module.__class__.__name__ == 'Conv2d':
+            elif module.__class__.__name__ == "Conv2d":
                 k_size, *_ = module.kernel_size
-                if k_size==1 and lora_dim>0:
+                if k_size == 1 and lora_dim > 0:
                     dim = dim or lora_dim
                     alpha = alpha or self.alpha
-                elif conv_lora_dim>0 or dim:
+                elif conv_lora_dim > 0 or dim:
                     dim = dim or conv_lora_dim
                     alpha = alpha or self.conv_alpha
                 else:
@@ -331,56 +399,61 @@ def create_single_module(
             else:
                 return None
             lora = network_module_dict[algo_name](
-                lora_name, module, self.multiplier, 
-                dim, alpha, 
-                self.dropout, self.rank_dropout, self.module_dropout,
+                lora_name,
+                module,
+                self.multiplier,
+                dim,
+                alpha,
+                self.dropout,
+                self.rank_dropout,
+                self.module_dropout,
                 use_tucker,
-                **kwargs
+                **kwargs,
             )
             return lora
-        
+
         def create_modules_(
             prefix: str,
             root_module: torch.nn.Module,
             algo,
-            configs = {},
+            configs={},
         ):
             loras = {}
             lora_names = []
             for name, module in root_module.named_modules():
-                if module is root_module: continue
+                if module is root_module:
+                    continue
                 module_name = module.__class__.__name__
                 if module_name in self.MODULE_ALGO_MAP:
                     next_config = self.MODULE_ALGO_MAP[module_name]
-                    next_algo = next_config.get('algo', algo)
+                    next_algo = next_config.get("algo", algo)
                     new_loras, new_lora_names = create_modules_(
-                        f'{prefix}_{name}', module, next_algo, next_config
+                        f"{prefix}_{name}", module, next_algo, next_config
                     )
                     for lora_name, lora in zip(new_lora_names, new_loras):
                         if lora_name not in loras:
                             loras[lora_name] = lora
                             lora_names.append(lora_name)
                     continue
-                lora_name = prefix + '.' + name
-                lora_name = lora_name.replace('.', '_')
-                if lora_name in loras: continue
-                
-                lora = create_single_module(
-                    lora_name, module, algo, **configs
-                )
+                lora_name = prefix + "." + name
+                lora_name = lora_name.replace(".", "_")
+                if lora_name in loras:
+                    continue
+
+                lora = create_single_module(lora_name, module, algo, **configs)
                 if lora is not None:
                     loras[lora_name] = lora
                     lora_names.append(lora_name)
             return [loras[lora_name] for lora_name in lora_names], lora_names
-        
+
         # create module instances
         def create_modules(
-            prefix, 
-            root_module: torch.nn.Module, 
+            prefix,
+            root_module: torch.nn.Module,
             target_replace_modules,
-            target_replace_names = []
+            target_replace_names=[],
         ) -> List:
-            print('Create LyCORIS Module')
+            print("Create LyCORIS Module")
             loras = []
             next_config = {}
             for name, module in root_module.named_modules():
@@ -388,34 +461,38 @@ def create_modules(
                 if module_name in target_replace_modules:
                     if module_name in self.MODULE_ALGO_MAP:
                         next_config = self.MODULE_ALGO_MAP[module_name]
-                        algo = next_config.get('algo', network_module)
+                        algo = next_config.get("algo", network_module)
                     else:
                         algo = network_module
-                    loras.extend(create_modules_(f'{prefix}_{name}', module, algo, next_config)[0])
+                    loras.extend(
+                        create_modules_(f"{prefix}_{name}", module, algo, next_config)[
+                            0
+                        ]
+                    )
                     next_config = {}
                 elif name in target_replace_names:
                     if name in self.NAME_ALGO_MAP:
                         next_config = self.NAME_ALGO_MAP[name]
-                        algo = next_config.get('algo', network_module)
+                        algo = next_config.get("algo", network_module)
                     elif module_name in self.MODULE_ALGO_MAP:
                         next_config = self.MODULE_ALGO_MAP[module_name]
-                        algo = next_config.get('algo', network_module)
+                        algo = next_config.get("algo", network_module)
                     else:
                         algo = network_module
-                    lora_name = prefix + '.' + name
-                    lora_name = lora_name.replace('.', '_')
+                    lora_name = prefix + "." + name
+                    lora_name = lora_name.replace(".", "_")
                     lora = create_single_module(lora_name, module, algo, **next_config)
                     next_config = {}
                     if lora is not None:
                         loras.append(lora)
             return loras
-        
+
         if network_module == GLoRAModule:
-            print('GLoRA enabled, only train transformer')
+            print("GLoRA enabled, only train transformer")
             # only train transformer (for GLoRA)
             LycorisNetwork.UNET_TARGET_REPLACE_MODULE = [
-                "Transformer2DModel", 
-                "Attention", 
+                "Transformer2DModel",
+                "Attention",
             ]
             LycorisNetwork.UNET_TARGET_REPLACE_NAME = []
 
@@ -428,25 +505,32 @@ def create_modules(
 
         self.text_encoder_loras = []
         for i, te in enumerate(text_encoders):
-            self.text_encoder_loras.extend(create_modules(
-                LycorisNetwork.LORA_PREFIX_TEXT_ENCODER + (f'{i+1}' if use_index else ''),
-                te, 
-                LycorisNetwork.TEXT_ENCODER_TARGET_REPLACE_MODULE,
-                LycorisNetwork.TEXT_ENCODER_TARGET_REPLACE_NAME
-            ))
-        print(f"create LyCORIS for Text Encoder: {len(self.text_encoder_loras)} modules.")
+            self.text_encoder_loras.extend(
+                create_modules(
+                    LycorisNetwork.LORA_PREFIX_TEXT_ENCODER
+                    + (f"{i+1}" if use_index else ""),
+                    te,
+                    LycorisNetwork.TEXT_ENCODER_TARGET_REPLACE_MODULE,
+                    LycorisNetwork.TEXT_ENCODER_TARGET_REPLACE_NAME,
+                )
+            )
+        print(
+            f"create LyCORIS for Text Encoder: {len(self.text_encoder_loras)} modules."
+        )
 
         self.unet_loras = create_modules(
-            LycorisNetwork.LORA_PREFIX_UNET, 
-            unet, 
+            LycorisNetwork.LORA_PREFIX_UNET,
+            unet,
             LycorisNetwork.UNET_TARGET_REPLACE_MODULE,
-            LycorisNetwork.UNET_TARGET_REPLACE_NAME
+            LycorisNetwork.UNET_TARGET_REPLACE_NAME,
         )
         print(f"create LyCORIS for U-Net: {len(self.unet_loras)} modules.")
-        
+
         algo_table = {}
         for lora in self.text_encoder_loras + self.unet_loras:
-            algo_table[lora.__class__.__name__] = algo_table.get(lora.__class__.__name__, 0) + 1
+            algo_table[lora.__class__.__name__] = (
+                algo_table.get(lora.__class__.__name__, 0) + 1
+            )
         print(f"module type table: {algo_table}")
 
         self.weights_sd = None
@@ -454,41 +538,35 @@ def create_modules(
         # 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}"
+            assert (
+                lora.lora_name not in names
+            ), f"duplicated lora name: {lora.lora_name}"
             names.add(lora.lora_name)
 
     def set_multiplier(self, multiplier):
         self.multiplier = multiplier
         for lora in self.text_encoder_loras + self.unet_loras:
             lora.multiplier = self.multiplier
-    
+
     def load_weights(self, file):
-        if os.path.splitext(file)[1] == '.safetensors':
+        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')
+            self.weights_sd = torch.load(file, map_location="cpu")
+        missing, unexpected = self.load_state_dict(self.weights_sd, strict=False)
+        state = {}
+        if missing:
+            state["missing keys"] = missing
+        if unexpected:
+            state["unexpected keys"] = unexpected
+        return state
 
     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(LycorisNetwork.LORA_PREFIX_TEXT_ENCODER):
-                    weights_has_text_encoder = True
-                elif key.startswith(LycorisNetwork.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"
+        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 LyCORIS for text encoder")
@@ -508,32 +586,33 @@ def apply_to(self, text_encoder, unet, apply_text_encoder=None, apply_unet=None)
             # 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 apply_max_norm_regularization(self, max_norm_value, device):
         key_scaled = 0
         norms = []
         for model in self.unet_loras:
-            if hasattr(model, 'apply_max_norm'):
+            if hasattr(model, "apply_max_norm"):
                 scaled, norm = model.apply_max_norm(max_norm_value, device)
                 norms.append(norm)
                 key_scaled += scaled
-        
+
         for model in self.text_encoder_loras:
-            if hasattr(model, 'apply_max_norm'):
+            if hasattr(model, "apply_max_norm"):
                 scaled, norm = model.apply_max_norm(max_norm_value, device)
                 norms.append(norm)
                 key_scaled += scaled
-        
-        if key_scaled==0:
+
+        if key_scaled == 0:
             return key_scaled, 0, 0
-        
-        return key_scaled, sum(norms)/len(norms), max(norms)
+
+        return key_scaled, sum(norms) / len(norms), max(norms)
 
     def enable_gradient_checkpointing(self):
         # not supported
         def make_ckpt(module):
             if isinstance(module, torch.nn.Module):
                 module.grad_ckpt = True
+
         self.apply(make_ckpt)
         pass
 
@@ -548,15 +627,15 @@ def enumerate_params(loras):
         all_params = []
 
         if self.text_encoder_loras:
-            param_data = {'params': enumerate_params(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
+                param_data["lr"] = text_encoder_lr
             all_params.append(param_data)
 
         if self.unet_loras:
-            param_data = {'params': enumerate_params(self.unet_loras)}
+            param_data = {"params": enumerate_params(self.unet_loras)}
             if unet_lr is not None:
-                param_data['lr'] = unet_lr
+                param_data["lr"] = unet_lr
             all_params.append(param_data)
 
         return all_params
@@ -582,113 +661,119 @@ def save_weights(self, file, dtype, metadata):
                 v = v.detach().clone().to("cpu").to(dtype)
                 state_dict[key] = v
 
-        if os.path.splitext(file)[1] == '.safetensors':
+        if os.path.splitext(file)[1] == ".safetensors":
             from safetensors.torch import save_file
 
             # Precalculate model hashes to save time on indexing
             if metadata is None:
                 metadata = {}
-            model_hash, legacy_hash = precalculate_safetensors_hashes(state_dict, metadata)
+            model_hash, legacy_hash = precalculate_safetensors_hashes(
+                state_dict, metadata
+            )
             metadata["sshs_model_hash"] = model_hash
             metadata["sshs_legacy_hash"] = legacy_hash
 
             save_file(state_dict, file, metadata)
         else:
             torch.save(state_dict, file)
-    
-    def apply_max_norm_regularization(self, max_norm_value, device):
-        norms = []
-        scaled = 0
-        
-        for lora in self.unet_loras:
-            if hasattr(lora, 'apply_max_norm'):
-                scaled, norm = lora.apply_max_norm(max_norm_value, device)
-                norms.append(norm)
-                scaled += int(scaled)
-        
-        for lora in self.text_encoder_loras:
-            if hasattr(lora, 'apply_max_norm'):
-                scaled, norm = lora.apply_max_norm(max_norm_value, device)
-                norms.append(norm)
-                scaled += int(scaled)
-        
-        return scaled, sum(norms)/len(norms), max(norms)
 
 
 class HyperDreamNetwork(torch.nn.Module):
-    '''
+    """
     HyperDreamBooth hypernetwork part
     only train Attention right now
-    '''
+    """
+
     UNET_TARGET_REPLACE_MODULE = [
-        "Attention", 
+        "Attention",
     ]
     UNET_TARGET_REPLACE_NAME = []
     TEXT_ENCODER_TARGET_REPLACE_MODULE = ["CLIPAttention"]
-    LORA_PREFIX_UNET = 'lora_unet'
-    LORA_PREFIX_TEXT_ENCODER = 'lora_te'
+    LORA_PREFIX_UNET = "lora_unet"
+    LORA_PREFIX_TEXT_ENCODER = "lora_te"
 
     def __init__(
-        self, 
-        text_encoder, unet, 
-        multiplier=1.0, 
-        lora_dim=4, alpha=1,
-        use_tucker = False,
-        dropout = 0, rank_dropout = 0, module_dropout = 0,
-        network_module = LoConModule,
-        down_dim = 100, up_dim = 50, delta_iters = 5, decoder_blocks = 4, vocab_size = 49408,
+        self,
+        text_encoder,
+        unet,
+        multiplier=1.0,
+        lora_dim=4,
+        alpha=1,
+        use_tucker=False,
+        dropout=0,
+        rank_dropout=0,
+        module_dropout=0,
+        network_module=LoConModule,
+        down_dim=100,
+        up_dim=50,
+        delta_iters=5,
+        decoder_blocks=4,
+        vocab_size=49408,
         **kwargs,
     ) -> None:
         from ..modules.hypernet import ImgWeightGenerator, TextWeightGenerator
+
         super().__init__()
         self.gradient_ckpt = False
         self.multiplier = multiplier
         self.lora_dim = lora_dim
         self.alpha = alpha
-        
+
         if 1 >= dropout >= 0:
-            print(f'Use Dropout value: {dropout}')
+            print(f"Use Dropout value: {dropout}")
         if network_module != LoConModule:
-            print('HyperDreamBooth only support LoRA at this time')
+            print("HyperDreamBooth only support LoRA at this time")
             raise NotImplementedError
-        if lora_dim*(down_dim+up_dim) > 4096:
-            print('weight elements > 4096 (dim * (down_dim + up_dim)) is not recommended!')
-        
+        if lora_dim * (down_dim + up_dim) > 4096:
+            print(
+                "weight elements > 4096 (dim * (down_dim + up_dim)) is not recommended!"
+            )
+
         self.dropout = dropout
         self.rank_dropout = rank_dropout
         self.module_dropout = module_dropout
-        
+
         # create module instances
         def create_modules(
-            prefix, 
-            root_module: torch.nn.Module, 
+            prefix,
+            root_module: torch.nn.Module,
             target_replace_modules,
-            target_replace_names = []
+            target_replace_names=[],
         ) -> List[network_module]:
-            print('Create LyCORIS Module')
+            print("Create LyCORIS Module")
             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():
-                        lora_name = prefix + '.' + name + '.' + child_name
-                        lora_name = lora_name.replace('.', '_')
-                        if child_module.__class__.__name__ == 'Linear' and lora_dim>0:
+                        lora_name = prefix + "." + name + "." + child_name
+                        lora_name = lora_name.replace(".", "_")
+                        if child_module.__class__.__name__ == "Linear" and lora_dim > 0:
                             lora = network_module(
-                                lora_name, child_module, self.multiplier, 
-                                self.lora_dim, self.alpha, 
-                                self.dropout, self.rank_dropout, self.module_dropout, 
+                                lora_name,
+                                child_module,
+                                self.multiplier,
+                                self.lora_dim,
+                                self.alpha,
+                                self.dropout,
+                                self.rank_dropout,
+                                self.module_dropout,
                                 use_tucker,
-                                **kwargs
+                                **kwargs,
                             )
-                        elif child_module.__class__.__name__ == 'Conv2d':
+                        elif child_module.__class__.__name__ == "Conv2d":
                             k_size, *_ = child_module.kernel_size
-                            if k_size==1 and lora_dim>0:
+                            if k_size == 1 and lora_dim > 0:
                                 lora = network_module(
-                                    lora_name, child_module, self.multiplier, 
-                                    self.lora_dim, self.alpha, 
-                                    self.dropout, self.rank_dropout, self.module_dropout, 
+                                    lora_name,
+                                    child_module,
+                                    self.multiplier,
+                                    self.lora_dim,
+                                    self.alpha,
+                                    self.dropout,
+                                    self.rank_dropout,
+                                    self.module_dropout,
                                     use_tucker,
-                                    **kwargs
+                                    **kwargs,
                                 )
                             else:
                                 continue
@@ -696,25 +781,35 @@ def create_modules(
                             continue
                         loras.append(lora)
                 elif name in target_replace_names:
-                    lora_name = prefix + '.' + name
-                    lora_name = lora_name.replace('.', '_')
-                    if module.__class__.__name__ == 'Linear' and lora_dim>0:
+                    lora_name = prefix + "." + name
+                    lora_name = lora_name.replace(".", "_")
+                    if module.__class__.__name__ == "Linear" and lora_dim > 0:
                         lora = network_module(
-                            lora_name, module, self.multiplier, 
-                            self.lora_dim, self.alpha, 
-                            self.dropout, self.rank_dropout, self.module_dropout, 
+                            lora_name,
+                            module,
+                            self.multiplier,
+                            self.lora_dim,
+                            self.alpha,
+                            self.dropout,
+                            self.rank_dropout,
+                            self.module_dropout,
                             use_tucker,
-                            **kwargs
+                            **kwargs,
                         )
-                    elif module.__class__.__name__ == 'Conv2d':
+                    elif module.__class__.__name__ == "Conv2d":
                         k_size, *_ = module.kernel_size
-                        if k_size==1 and lora_dim>0:
+                        if k_size == 1 and lora_dim > 0:
                             lora = network_module(
-                                lora_name, module, self.multiplier, 
-                                self.lora_dim, self.alpha, 
-                                self.dropout, self.rank_dropout, self.module_dropout, 
+                                lora_name,
+                                module,
+                                self.multiplier,
+                                self.lora_dim,
+                                self.alpha,
+                                self.dropout,
+                                self.rank_dropout,
+                                self.module_dropout,
                                 use_tucker,
-                                **kwargs
+                                **kwargs,
                             )
                         else:
                             continue
@@ -732,30 +827,39 @@ def create_modules(
 
         self.text_encoder_loras = []
         for i, te in enumerate(text_encoders):
-            self.text_encoder_loras.extend(create_modules(
-                LycorisNetwork.LORA_PREFIX_TEXT_ENCODER + (f'{i+1}' if use_index else ''),
-                te, 
-                LycorisNetwork.TEXT_ENCODER_TARGET_REPLACE_MODULE
-            ))
-        print(f"create LyCORIS for Text Encoder: {len(self.text_encoder_loras)} modules.")
-
-        self.unet_loras = create_modules(LycorisNetwork.LORA_PREFIX_UNET, unet, LycorisNetwork.UNET_TARGET_REPLACE_MODULE)
+            self.text_encoder_loras.extend(
+                create_modules(
+                    LycorisNetwork.LORA_PREFIX_TEXT_ENCODER
+                    + (f"{i+1}" if use_index else ""),
+                    te,
+                    LycorisNetwork.TEXT_ENCODER_TARGET_REPLACE_MODULE,
+                )
+            )
+        print(
+            f"create LyCORIS for Text Encoder: {len(self.text_encoder_loras)} modules."
+        )
+
+        self.unet_loras = create_modules(
+            LycorisNetwork.LORA_PREFIX_UNET,
+            unet,
+            LycorisNetwork.UNET_TARGET_REPLACE_MODULE,
+        )
         print(f"create LyCORIS for U-Net: {len(self.unet_loras)} modules.")
-        
+
         self.loras: list[LoConModule] = self.text_encoder_loras + self.unet_loras
         self.img_weight_generater = ImgWeightGenerator(
-            weight_dim=(down_dim+up_dim)*lora_dim,
+            weight_dim=(down_dim + up_dim) * lora_dim,
             weight_num=len(self.unet_loras),
             sample_iters=delta_iters,
             decoder_blocks=decoder_blocks,
         )
         self.text_weight_generater = TextWeightGenerator(
-            weight_dim=(down_dim+up_dim)*lora_dim,
+            weight_dim=(down_dim + up_dim) * lora_dim,
             weight_num=len(self.text_encoder_loras),
             sample_iters=delta_iters,
             decoder_blocks=decoder_blocks,
         )
-        self.split = (down_dim* lora_dim, up_dim  * lora_dim)
+        self.split = (down_dim * lora_dim, up_dim * lora_dim)
         self.lora_dim = lora_dim
 
         self.weights_sd = None
@@ -763,27 +867,34 @@ def create_modules(
         # 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}"
+            assert (
+                lora.lora_name not in names
+            ), f"duplicated lora name: {lora.lora_name}"
             names.add(lora.lora_name)
-        
+
         self.checkpoint = torch.nn.Parameter(torch.tensor(0.0))
-        
+
         with torch.no_grad():
             self.update_reference(
-                torch.randn(1, 3, *self.img_weight_generater.ref_size),
-                ["test"]
+                torch.randn(1, 3, *self.img_weight_generater.ref_size), ["test"]
             )
-        
+
         # for lora in self.loras:
         #     assert torch.all(lora.data[0]==0)
 
     def gen_weight(self, ref_img, caption, iter=None, ensure_grad=0):
         unet_weights = self.img_weight_generater(ref_img, iter, ensure_grad=ensure_grad)
         unet_weights = unet_weights + self.checkpoint
-        unet_weights =  [i.split(self.split, dim=-1) for i in unet_weights.split(1, dim=1)]
-        text_weights = self.text_weight_generater(caption, iter, ensure_grad=ensure_grad)
+        unet_weights = [
+            i.split(self.split, dim=-1) for i in unet_weights.split(1, dim=1)
+        ]
+        text_weights = self.text_weight_generater(
+            caption, iter, ensure_grad=ensure_grad
+        )
         text_weights = text_weights + self.checkpoint
-        text_weights =  [i.split(self.split, dim=-1) for i in text_weights.split(1, dim=1)]
+        text_weights = [
+            i.split(self.split, dim=-1) for i in text_weights.split(1, dim=1)
+        ]
         return unet_weights, text_weights
 
     def update_reference(self, ref_img, caption, iter=None):
@@ -794,17 +905,25 @@ def update_reference(self, ref_img, caption, iter=None):
                 self.gen_weight, ref_img, caption, iter, ensure_grad
             )
         else:
-            unet_weights_list, text_weights_list = self.gen_weight(ref_img, caption, iter)
-            
+            unet_weights_list, text_weights_list = self.gen_weight(
+                ref_img, caption, iter
+            )
+
         for idx, (lora, weight) in enumerate(zip(self.unet_loras, unet_weights_list)):
-            assert lora.multiplier > 0, f"multiplier must be positive: {lora.multiplier}"
+            assert (
+                lora.multiplier > 0
+            ), f"multiplier must be positive: {lora.multiplier}"
             # weight: [batch, 1, weight_dim]
             # if weight.dim()==3:
             #     weight = weight.squeeze(1)
             lora.update_weights(*weight, idx)
-            
-        for idx, (lora, weight) in enumerate(zip(self.text_encoder_loras, text_weights_list)):
-            assert lora.multiplier > 0, f"multiplier must be positive: {lora.multiplier}"
+
+        for idx, (lora, weight) in enumerate(
+            zip(self.text_encoder_loras, text_weights_list)
+        ):
+            assert (
+                lora.multiplier > 0
+            ), f"multiplier must be positive: {lora.multiplier}"
             # weight: [batch, 1, weight_dim]
             # if weight.dim()==3:
             #     weight = weight.squeeze(1)
@@ -814,13 +933,14 @@ def set_multiplier(self, multiplier):
         self.multiplier = multiplier
         for lora in self.text_encoder_loras + self.unet_loras:
             lora.multiplier = self.multiplier
-    
+
     def load_weights(self, file):
-        if os.path.splitext(file)[1] == '.safetensors':
+        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')
+            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:
@@ -834,14 +954,20 @@ def apply_to(self, text_encoder, unet, apply_text_encoder=None, apply_unet=None)
             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のフラグが矛盾しています"
+                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のフラグが矛盾しています"
+                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"
+            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 LyCORIS for text encoder")
@@ -862,29 +988,63 @@ def enable_gradient_checkpointing(self):
     def prepare_optimizer_params(self, text_encoder_lr, unet_lr, learning_rate):
         self.requires_grad_(True)
         all_params = []
-        
+
         if self.text_encoder_loras:
-            all_params.append({
-                'params': (
-                    [p for p in self.text_weight_generater.decoder_model.parameters()]
-                    + [p for p in self.text_weight_generater.pos_emb_proj.parameters()]
-                    + [p for p in self.text_weight_generater.feature_proj.parameters()]
-                    + ([p for p in self.text_weight_generater.encoder_model.parameters()] 
-                    if self.text_weight_generater.train_encoder else [])
-                ), 
-                'lr': text_encoder_lr
-            })
+            all_params.append(
+                {
+                    "params": (
+                        [
+                            p
+                            for p in self.text_weight_generater.decoder_model.parameters()
+                        ]
+                        + [
+                            p
+                            for p in self.text_weight_generater.pos_emb_proj.parameters()
+                        ]
+                        + [
+                            p
+                            for p in self.text_weight_generater.feature_proj.parameters()
+                        ]
+                        + (
+                            [
+                                p
+                                for p in self.text_weight_generater.encoder_model.parameters()
+                            ]
+                            if self.text_weight_generater.train_encoder
+                            else []
+                        )
+                    ),
+                    "lr": text_encoder_lr,
+                }
+            )
         if self.unet_loras:
-            all_params.append({
-                'params': (
-                    [p for p in self.img_weight_generater.decoder_model.parameters()]
-                    + [p for p in self.img_weight_generater.pos_emb_proj.parameters()]
-                    + [p for p in self.img_weight_generater.feature_proj.parameters()]
-                    + ([p for p in self.img_weight_generater.encoder_model.parameters()] 
-                    if self.img_weight_generater.train_encoder else [])
-                ), 
-                'lr': unet_lr
-            })
+            all_params.append(
+                {
+                    "params": (
+                        [
+                            p
+                            for p in self.img_weight_generater.decoder_model.parameters()
+                        ]
+                        + [
+                            p
+                            for p in self.img_weight_generater.pos_emb_proj.parameters()
+                        ]
+                        + [
+                            p
+                            for p in self.img_weight_generater.feature_proj.parameters()
+                        ]
+                        + (
+                            [
+                                p
+                                for p in self.img_weight_generater.encoder_model.parameters()
+                            ]
+                            if self.img_weight_generater.train_encoder
+                            else []
+                        )
+                    ),
+                    "lr": unet_lr,
+                }
+            )
         return all_params
 
     def prepare_grad_etc(self, text_encoder, unet):
@@ -903,14 +1063,14 @@ def save_weights(self, file, dtype, metadata):
         state_dict = self.img_weight_generater.state_dict()
         if not self.img_weight_generater.train_encoder:
             for k in self.img_weight_generater.encoder_model.state_dict().keys():
-                state_dict.pop(f'encoder_model.{k}')
-        state_dict = {f'img_weight_generater.{i}': v for i, v in state_dict.items()}
+                state_dict.pop(f"encoder_model.{k}")
+        state_dict = {f"img_weight_generater.{i}": v for i, v in state_dict.items()}
 
         state_dict = self.text_weight_generater.state_dict()
         if not self.text_weight_generater.train_encoder:
             for k in self.text_weight_generater.encoder_model.state_dict().keys():
-                state_dict.pop(f'encoder_model.{k}')
-        state_dict = {f'text_weight_generater.{i}': v for i, v in state_dict.items()}
+                state_dict.pop(f"encoder_model.{k}")
+        state_dict = {f"text_weight_generater.{i}": v for i, v in state_dict.items()}
 
         if dtype is not None:
             for key in list(state_dict.keys()):
@@ -918,13 +1078,15 @@ def save_weights(self, file, dtype, metadata):
                 v = v.detach().clone().to("cpu").to(dtype)
                 state_dict[key] = v
 
-        if os.path.splitext(file)[1] == '.safetensors':
+        if os.path.splitext(file)[1] == ".safetensors":
             from safetensors.torch import save_file
 
             # Precalculate model hashes to save time on indexing
             if metadata is None:
                 metadata = {}
-            model_hash, legacy_hash = precalculate_safetensors_hashes(state_dict, metadata)
+            model_hash, legacy_hash = precalculate_safetensors_hashes(
+                state_dict, metadata
+            )
             metadata["sshs_model_hash"] = model_hash
             metadata["sshs_legacy_hash"] = legacy_hash
 
@@ -934,55 +1096,61 @@ def save_weights(self, file, dtype, metadata):
 
 
 class IA3Network(torch.nn.Module):
-    '''
+    """
     IA3 network
-    '''
+    """
+
     # Ignore proj_in or proj_out, their channels is only a few.
     UNET_TARGET_REPLACE_MODULE = []
     UNET_TARGET_REPLACE_NAME = ["to_k", "to_v", "ff.net.2"]
     TEXT_ENCODER_TARGET_REPLACE_MODULE = []
-    TEXT_ENCODER_TARGET_REPLACE_NAME= ["k_proj", "v_proj", "mlp.fc2"]
+    TEXT_ENCODER_TARGET_REPLACE_NAME = ["k_proj", "v_proj", "mlp.fc2"]
     TRAIN_INPUT = ["mlp.fc2", "ff.net.2"]
-    LORA_PREFIX_UNET = 'lora_unet'
-    LORA_PREFIX_TEXT_ENCODER = 'lora_te'
+    LORA_PREFIX_UNET = "lora_unet"
+    LORA_PREFIX_TEXT_ENCODER = "lora_te"
 
     def __init__(
-        self, 
-        text_encoder, unet, 
+        self,
+        text_encoder,
+        unet,
         multiplier=1.0,
         **kwargs,
     ) -> None:
         super().__init__()
         self.multiplier = multiplier
-        
+
         # create module instances
         def create_modules(
-            prefix, 
-            root_module: torch.nn.Module, 
+            prefix,
+            root_module: torch.nn.Module,
             target_replace_modules,
-            target_replace_names = [],
-            target_train_input = []
+            target_replace_names=[],
+            target_train_input=[],
         ) -> List[IA3Module]:
-            print('Create LyCORIS Module')
+            print("Create LyCORIS Module")
             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():
-                        lora_name = prefix + '.' + name + '.' + child_name
-                        lora_name = lora_name.replace('.', '_')
-                        if child_module.__class__.__name__ in {'Linear', 'Conv2d'}:
+                        lora_name = prefix + "." + name + "." + child_name
+                        lora_name = lora_name.replace(".", "_")
+                        if child_module.__class__.__name__ in {"Linear", "Conv2d"}:
                             lora = IA3Module(
-                                lora_name, child_module, self.multiplier,
+                                lora_name,
+                                child_module,
+                                self.multiplier,
                                 name in target_train_input,
                                 **kwargs,
                             )
                             loras.append(lora)
                 elif any(i in name for i in target_replace_names):
-                    lora_name = prefix + '.' + name
-                    lora_name = lora_name.replace('.', '_')
-                    if module.__class__.__name__ in {'Linear', 'Conv2d'}:
+                    lora_name = prefix + "." + name
+                    lora_name = lora_name.replace(".", "_")
+                    if module.__class__.__name__ in {"Linear", "Conv2d"}:
                         lora = IA3Module(
-                            lora_name, module, self.multiplier,
+                            lora_name,
+                            module,
+                            self.multiplier,
                             name in target_train_input,
                             **kwargs,
                         )
@@ -998,21 +1166,26 @@ def create_modules(
 
         self.text_encoder_loras = []
         for i, te in enumerate(text_encoders):
-            self.text_encoder_loras.extend(create_modules(
-                IA3Network.LORA_PREFIX_TEXT_ENCODER + (f'{i+1}' if use_index else ''),
-                te, 
-                IA3Network.TEXT_ENCODER_TARGET_REPLACE_MODULE,
-                IA3Network.TEXT_ENCODER_TARGET_REPLACE_NAME,
-                IA3Network.TRAIN_INPUT
-            ))
-        print(f"create LyCORIS for Text Encoder: {len(self.text_encoder_loras)} modules.")
+            self.text_encoder_loras.extend(
+                create_modules(
+                    IA3Network.LORA_PREFIX_TEXT_ENCODER
+                    + (f"{i+1}" if use_index else ""),
+                    te,
+                    IA3Network.TEXT_ENCODER_TARGET_REPLACE_MODULE,
+                    IA3Network.TEXT_ENCODER_TARGET_REPLACE_NAME,
+                    IA3Network.TRAIN_INPUT,
+                )
+            )
+        print(
+            f"create LyCORIS for Text Encoder: {len(self.text_encoder_loras)} modules."
+        )
 
         self.unet_loras = create_modules(
-            IA3Network.LORA_PREFIX_UNET, 
-            unet, 
+            IA3Network.LORA_PREFIX_UNET,
+            unet,
             IA3Network.UNET_TARGET_REPLACE_MODULE,
             IA3Network.UNET_TARGET_REPLACE_NAME,
-            IA3Network.TRAIN_INPUT
+            IA3Network.TRAIN_INPUT,
         )
         print(f"create LyCORIS for U-Net: {len(self.unet_loras)} modules.")
 
@@ -1021,20 +1194,23 @@ def create_modules(
         # 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}"
+            assert (
+                lora.lora_name not in names
+            ), f"duplicated lora name: {lora.lora_name}"
             names.add(lora.lora_name)
 
     def set_multiplier(self, multiplier):
         self.multiplier = multiplier
         for lora in self.text_encoder_loras + self.unet_loras:
             lora.multiplier = self.multiplier
-            
+
     def load_weights(self, file):
-        if os.path.splitext(file)[1] == '.safetensors':
+        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')
+            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:
@@ -1048,14 +1224,20 @@ def apply_to(self, text_encoder, unet, apply_text_encoder=None, apply_unet=None)
             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のフラグが矛盾しています"
+                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のフラグが矛盾しています"
+                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"
+            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 LyCORIS for text encoder")
@@ -1081,6 +1263,7 @@ def enable_gradient_checkpointing(self):
         def make_ckpt(module):
             if isinstance(module, torch.nn.Module):
                 module.grad_ckpt = True
+
         self.apply(make_ckpt)
         pass
 
@@ -1095,15 +1278,15 @@ def enumerate_params(loras):
         all_params = []
 
         if self.text_encoder_loras:
-            param_data = {'params': enumerate_params(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
+                param_data["lr"] = text_encoder_lr
             all_params.append(param_data)
 
         if self.unet_loras:
-            param_data = {'params': enumerate_params(self.unet_loras)}
+            param_data = {"params": enumerate_params(self.unet_loras)}
             if unet_lr is not None:
-                param_data['lr'] = unet_lr
+                param_data["lr"] = unet_lr
             all_params.append(param_data)
 
         return all_params
@@ -1129,13 +1312,15 @@ def save_weights(self, file, dtype, metadata):
                 v = v.detach().clone().to("cpu").to(dtype)
                 state_dict[key] = v
 
-        if os.path.splitext(file)[1] == '.safetensors':
+        if os.path.splitext(file)[1] == ".safetensors":
             from safetensors.torch import save_file
 
             # Precalculate model hashes to save time on indexing
             if metadata is None:
                 metadata = {}
-            model_hash, legacy_hash = precalculate_safetensors_hashes(state_dict, metadata)
+            model_hash, legacy_hash = precalculate_safetensors_hashes(
+                state_dict, metadata
+            )
             metadata["sshs_model_hash"] = model_hash
             metadata["sshs_legacy_hash"] = legacy_hash
 
diff --git a/lycoris/kohya/model_utils.py b/lycoris/kohya/model_utils.py
index 76a6961..0a37d4f 100644
--- a/lycoris/kohya/model_utils.py
+++ b/lycoris/kohya/model_utils.py
@@ -7,7 +7,11 @@
 import torch
 import diffusers
 from transformers import CLIPTextModel, CLIPTokenizer, CLIPTextConfig, logging
-from diffusers import AutoencoderKL, DDIMScheduler, StableDiffusionPipeline  # , UNet2DConditionModel
+from diffusers import (
+    AutoencoderKL,
+    DDIMScheduler,
+    StableDiffusionPipeline,
+)  # , UNet2DConditionModel
 from safetensors.torch import load_file, save_file
 from .original_unet import UNet2DConditionModel
 
@@ -74,7 +78,9 @@ def renew_resnet_paths(old_list, n_shave_prefix_segments=0):
         new_item = new_item.replace("emb_layers.1", "time_emb_proj")
         new_item = new_item.replace("skip_connection", "conv_shortcut")
 
-        new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments)
+        new_item = shave_segments(
+            new_item, n_shave_prefix_segments=n_shave_prefix_segments
+        )
 
         mapping.append({"old": old_item, "new": new_item})
 
@@ -90,7 +96,9 @@ def renew_vae_resnet_paths(old_list, n_shave_prefix_segments=0):
         new_item = old_item
 
         new_item = new_item.replace("nin_shortcut", "conv_shortcut")
-        new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments)
+        new_item = shave_segments(
+            new_item, n_shave_prefix_segments=n_shave_prefix_segments
+        )
 
         mapping.append({"old": old_item, "new": new_item})
 
@@ -154,7 +162,9 @@ def renew_vae_attention_paths(old_list, n_shave_prefix_segments=0):
             new_item = new_item.replace("proj_out.weight", "to_out.0.weight")
             new_item = new_item.replace("proj_out.bias", "to_out.0.bias")
 
-        new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments)
+        new_item = shave_segments(
+            new_item, n_shave_prefix_segments=n_shave_prefix_segments
+        )
 
         mapping.append({"old": old_item, "new": new_item})
 
@@ -162,7 +172,12 @@ def renew_vae_attention_paths(old_list, n_shave_prefix_segments=0):
 
 
 def assign_to_checkpoint(
-    paths, checkpoint, old_checkpoint, attention_paths_to_split=None, additional_replacements=None, config=None
+    paths,
+    checkpoint,
+    old_checkpoint,
+    attention_paths_to_split=None,
+    additional_replacements=None,
+    config=None,
 ):
     """
     This does the final conversion step: take locally converted weights and apply a global renaming
@@ -171,7 +186,9 @@ def assign_to_checkpoint(
 
     Assigns the weights to the new checkpoint.
     """
-    assert isinstance(paths, list), "Paths should be a list of dicts containing 'old' and 'new' keys."
+    assert isinstance(
+        paths, list
+    ), "Paths should be a list of dicts containing 'old' and 'new' keys."
 
     # Splits the attention layers into three variables.
     if attention_paths_to_split is not None:
@@ -183,7 +200,9 @@ def assign_to_checkpoint(
 
             num_heads = old_tensor.shape[0] // config["num_head_channels"] // 3
 
-            old_tensor = old_tensor.reshape((num_heads, 3 * channels // num_heads) + old_tensor.shape[1:])
+            old_tensor = old_tensor.reshape(
+                (num_heads, 3 * channels // num_heads) + old_tensor.shape[1:]
+            )
             query, key, value = old_tensor.split(channels // num_heads, dim=1)
 
             checkpoint[path_map["query"]] = query.reshape(target_shape)
@@ -194,7 +213,10 @@ def assign_to_checkpoint(
         new_path = path["new"]
 
         # These have already been assigned
-        if attention_paths_to_split is not None and new_path in attention_paths_to_split:
+        if (
+            attention_paths_to_split is not None
+            and new_path in attention_paths_to_split
+        ):
             continue
 
         # Global renaming happens here
@@ -212,7 +234,11 @@ def assign_to_checkpoint(
             if "proj_attn.weight" in new_path:
                 reshaping = True
         else:
-            if ".attentions." in new_path and ".0.to_" in new_path and old_checkpoint[path["old"]].ndim > 2:
+            if (
+                ".attentions." in new_path
+                and ".0.to_" in new_path
+                and old_checkpoint[path["old"]].ndim > 2
+            ):
                 reshaping = True
 
         if reshaping:
@@ -257,10 +283,18 @@ def convert_ldm_unet_checkpoint(v2, checkpoint, config):
 
     new_checkpoint = {}
 
-    new_checkpoint["time_embedding.linear_1.weight"] = unet_state_dict["time_embed.0.weight"]
-    new_checkpoint["time_embedding.linear_1.bias"] = unet_state_dict["time_embed.0.bias"]
-    new_checkpoint["time_embedding.linear_2.weight"] = unet_state_dict["time_embed.2.weight"]
-    new_checkpoint["time_embedding.linear_2.bias"] = unet_state_dict["time_embed.2.bias"]
+    new_checkpoint["time_embedding.linear_1.weight"] = unet_state_dict[
+        "time_embed.0.weight"
+    ]
+    new_checkpoint["time_embedding.linear_1.bias"] = unet_state_dict[
+        "time_embed.0.bias"
+    ]
+    new_checkpoint["time_embedding.linear_2.weight"] = unet_state_dict[
+        "time_embed.2.weight"
+    ]
+    new_checkpoint["time_embedding.linear_2.bias"] = unet_state_dict[
+        "time_embed.2.bias"
+    ]
 
     new_checkpoint["conv_in.weight"] = unet_state_dict["input_blocks.0.0.weight"]
     new_checkpoint["conv_in.bias"] = unet_state_dict["input_blocks.0.0.bias"]
@@ -271,44 +305,91 @@ def convert_ldm_unet_checkpoint(v2, checkpoint, config):
     new_checkpoint["conv_out.bias"] = unet_state_dict["out.2.bias"]
 
     # Retrieves the keys for the input blocks only
-    num_input_blocks = len({".".join(layer.split(".")[:2]) for layer in unet_state_dict if "input_blocks" in layer})
+    num_input_blocks = len(
+        {
+            ".".join(layer.split(".")[:2])
+            for layer in unet_state_dict
+            if "input_blocks" in layer
+        }
+    )
     input_blocks = {
-        layer_id: [key for key in unet_state_dict if f"input_blocks.{layer_id}." in key] for layer_id in range(num_input_blocks)
+        layer_id: [key for key in unet_state_dict if f"input_blocks.{layer_id}." in key]
+        for layer_id in range(num_input_blocks)
     }
 
     # Retrieves the keys for the middle blocks only
-    num_middle_blocks = len({".".join(layer.split(".")[:2]) for layer in unet_state_dict if "middle_block" in layer})
+    num_middle_blocks = len(
+        {
+            ".".join(layer.split(".")[:2])
+            for layer in unet_state_dict
+            if "middle_block" in layer
+        }
+    )
     middle_blocks = {
-        layer_id: [key for key in unet_state_dict if f"middle_block.{layer_id}." in key] for layer_id in range(num_middle_blocks)
+        layer_id: [key for key in unet_state_dict if f"middle_block.{layer_id}." in key]
+        for layer_id in range(num_middle_blocks)
     }
 
     # Retrieves the keys for the output blocks only
-    num_output_blocks = len({".".join(layer.split(".")[:2]) for layer in unet_state_dict if "output_blocks" in layer})
+    num_output_blocks = len(
+        {
+            ".".join(layer.split(".")[:2])
+            for layer in unet_state_dict
+            if "output_blocks" in layer
+        }
+    )
     output_blocks = {
-        layer_id: [key for key in unet_state_dict if f"output_blocks.{layer_id}." in key] for layer_id in range(num_output_blocks)
+        layer_id: [
+            key for key in unet_state_dict if f"output_blocks.{layer_id}." in key
+        ]
+        for layer_id in range(num_output_blocks)
     }
 
     for i in range(1, num_input_blocks):
         block_id = (i - 1) // (config["layers_per_block"] + 1)
         layer_in_block_id = (i - 1) % (config["layers_per_block"] + 1)
 
-        resnets = [key for key in input_blocks[i] if f"input_blocks.{i}.0" in key and f"input_blocks.{i}.0.op" not in key]
+        resnets = [
+            key
+            for key in input_blocks[i]
+            if f"input_blocks.{i}.0" in key and f"input_blocks.{i}.0.op" not in key
+        ]
         attentions = [key for key in input_blocks[i] if f"input_blocks.{i}.1" in key]
 
         if f"input_blocks.{i}.0.op.weight" in unet_state_dict:
-            new_checkpoint[f"down_blocks.{block_id}.downsamplers.0.conv.weight"] = unet_state_dict.pop(
-                f"input_blocks.{i}.0.op.weight"
-            )
-            new_checkpoint[f"down_blocks.{block_id}.downsamplers.0.conv.bias"] = unet_state_dict.pop(f"input_blocks.{i}.0.op.bias")
+            new_checkpoint[
+                f"down_blocks.{block_id}.downsamplers.0.conv.weight"
+            ] = unet_state_dict.pop(f"input_blocks.{i}.0.op.weight")
+            new_checkpoint[
+                f"down_blocks.{block_id}.downsamplers.0.conv.bias"
+            ] = unet_state_dict.pop(f"input_blocks.{i}.0.op.bias")
 
         paths = renew_resnet_paths(resnets)
-        meta_path = {"old": f"input_blocks.{i}.0", "new": f"down_blocks.{block_id}.resnets.{layer_in_block_id}"}
-        assign_to_checkpoint(paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config)
+        meta_path = {
+            "old": f"input_blocks.{i}.0",
+            "new": f"down_blocks.{block_id}.resnets.{layer_in_block_id}",
+        }
+        assign_to_checkpoint(
+            paths,
+            new_checkpoint,
+            unet_state_dict,
+            additional_replacements=[meta_path],
+            config=config,
+        )
 
         if len(attentions):
             paths = renew_attention_paths(attentions)
-            meta_path = {"old": f"input_blocks.{i}.1", "new": f"down_blocks.{block_id}.attentions.{layer_in_block_id}"}
-            assign_to_checkpoint(paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config)
+            meta_path = {
+                "old": f"input_blocks.{i}.1",
+                "new": f"down_blocks.{block_id}.attentions.{layer_in_block_id}",
+            }
+            assign_to_checkpoint(
+                paths,
+                new_checkpoint,
+                unet_state_dict,
+                additional_replacements=[meta_path],
+                config=config,
+            )
 
     resnet_0 = middle_blocks[0]
     attentions = middle_blocks[1]
@@ -322,7 +403,13 @@ def convert_ldm_unet_checkpoint(v2, checkpoint, config):
 
     attentions_paths = renew_attention_paths(attentions)
     meta_path = {"old": "middle_block.1", "new": "mid_block.attentions.0"}
-    assign_to_checkpoint(attentions_paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config)
+    assign_to_checkpoint(
+        attentions_paths,
+        new_checkpoint,
+        unet_state_dict,
+        additional_replacements=[meta_path],
+        config=config,
+    )
 
     for i in range(num_output_blocks):
         block_id = i // (config["layers_per_block"] + 1)
@@ -339,13 +426,24 @@ def convert_ldm_unet_checkpoint(v2, checkpoint, config):
 
         if len(output_block_list) > 1:
             resnets = [key for key in output_blocks[i] if f"output_blocks.{i}.0" in key]
-            attentions = [key for key in output_blocks[i] if f"output_blocks.{i}.1" in key]
+            attentions = [
+                key for key in output_blocks[i] if f"output_blocks.{i}.1" in key
+            ]
 
             resnet_0_paths = renew_resnet_paths(resnets)
             paths = renew_resnet_paths(resnets)
 
-            meta_path = {"old": f"output_blocks.{i}.0", "new": f"up_blocks.{block_id}.resnets.{layer_in_block_id}"}
-            assign_to_checkpoint(paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config)
+            meta_path = {
+                "old": f"output_blocks.{i}.0",
+                "new": f"up_blocks.{block_id}.resnets.{layer_in_block_id}",
+            }
+            assign_to_checkpoint(
+                paths,
+                new_checkpoint,
+                unet_state_dict,
+                additional_replacements=[meta_path],
+                config=config,
+            )
 
             # オリジナル：
             # if ["conv.weight", "conv.bias"] in output_block_list.values():
@@ -356,13 +454,15 @@ def convert_ldm_unet_checkpoint(v2, checkpoint, config):
                 l.sort()
 
             if ["conv.bias", "conv.weight"] in output_block_list.values():
-                index = list(output_block_list.values()).index(["conv.bias", "conv.weight"])
-                new_checkpoint[f"up_blocks.{block_id}.upsamplers.0.conv.bias"] = unet_state_dict[
-                    f"output_blocks.{i}.{index}.conv.bias"
-                ]
-                new_checkpoint[f"up_blocks.{block_id}.upsamplers.0.conv.weight"] = unet_state_dict[
-                    f"output_blocks.{i}.{index}.conv.weight"
-                ]
+                index = list(output_block_list.values()).index(
+                    ["conv.bias", "conv.weight"]
+                )
+                new_checkpoint[
+                    f"up_blocks.{block_id}.upsamplers.0.conv.bias"
+                ] = unet_state_dict[f"output_blocks.{i}.{index}.conv.bias"]
+                new_checkpoint[
+                    f"up_blocks.{block_id}.upsamplers.0.conv.weight"
+                ] = unet_state_dict[f"output_blocks.{i}.{index}.conv.weight"]
 
                 # Clear attentions as they have been attributed above.
                 if len(attentions) == 2:
@@ -374,12 +474,28 @@ def convert_ldm_unet_checkpoint(v2, checkpoint, config):
                     "old": f"output_blocks.{i}.1",
                     "new": f"up_blocks.{block_id}.attentions.{layer_in_block_id}",
                 }
-                assign_to_checkpoint(paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config)
+                assign_to_checkpoint(
+                    paths,
+                    new_checkpoint,
+                    unet_state_dict,
+                    additional_replacements=[meta_path],
+                    config=config,
+                )
         else:
-            resnet_0_paths = renew_resnet_paths(output_block_layers, n_shave_prefix_segments=1)
+            resnet_0_paths = renew_resnet_paths(
+                output_block_layers, n_shave_prefix_segments=1
+            )
             for path in resnet_0_paths:
                 old_path = ".".join(["output_blocks", str(i), path["old"]])
-                new_path = ".".join(["up_blocks", str(block_id), "resnets", str(layer_in_block_id), path["new"]])
+                new_path = ".".join(
+                    [
+                        "up_blocks",
+                        str(block_id),
+                        "resnets",
+                        str(layer_in_block_id),
+                        path["new"],
+                    ]
+                )
 
                 new_checkpoint[new_path] = unet_state_dict[old_path]
 
@@ -407,17 +523,29 @@ def convert_ldm_vae_checkpoint(checkpoint, config):
 
     new_checkpoint["encoder.conv_in.weight"] = vae_state_dict["encoder.conv_in.weight"]
     new_checkpoint["encoder.conv_in.bias"] = vae_state_dict["encoder.conv_in.bias"]
-    new_checkpoint["encoder.conv_out.weight"] = vae_state_dict["encoder.conv_out.weight"]
+    new_checkpoint["encoder.conv_out.weight"] = vae_state_dict[
+        "encoder.conv_out.weight"
+    ]
     new_checkpoint["encoder.conv_out.bias"] = vae_state_dict["encoder.conv_out.bias"]
-    new_checkpoint["encoder.conv_norm_out.weight"] = vae_state_dict["encoder.norm_out.weight"]
-    new_checkpoint["encoder.conv_norm_out.bias"] = vae_state_dict["encoder.norm_out.bias"]
+    new_checkpoint["encoder.conv_norm_out.weight"] = vae_state_dict[
+        "encoder.norm_out.weight"
+    ]
+    new_checkpoint["encoder.conv_norm_out.bias"] = vae_state_dict[
+        "encoder.norm_out.bias"
+    ]
 
     new_checkpoint["decoder.conv_in.weight"] = vae_state_dict["decoder.conv_in.weight"]
     new_checkpoint["decoder.conv_in.bias"] = vae_state_dict["decoder.conv_in.bias"]
-    new_checkpoint["decoder.conv_out.weight"] = vae_state_dict["decoder.conv_out.weight"]
+    new_checkpoint["decoder.conv_out.weight"] = vae_state_dict[
+        "decoder.conv_out.weight"
+    ]
     new_checkpoint["decoder.conv_out.bias"] = vae_state_dict["decoder.conv_out.bias"]
-    new_checkpoint["decoder.conv_norm_out.weight"] = vae_state_dict["decoder.norm_out.weight"]
-    new_checkpoint["decoder.conv_norm_out.bias"] = vae_state_dict["decoder.norm_out.bias"]
+    new_checkpoint["decoder.conv_norm_out.weight"] = vae_state_dict[
+        "decoder.norm_out.weight"
+    ]
+    new_checkpoint["decoder.conv_norm_out.bias"] = vae_state_dict[
+        "decoder.norm_out.bias"
+    ]
 
     new_checkpoint["quant_conv.weight"] = vae_state_dict["quant_conv.weight"]
     new_checkpoint["quant_conv.bias"] = vae_state_dict["quant_conv.bias"]
@@ -425,27 +553,55 @@ def convert_ldm_vae_checkpoint(checkpoint, config):
     new_checkpoint["post_quant_conv.bias"] = vae_state_dict["post_quant_conv.bias"]
 
     # Retrieves the keys for the encoder down blocks only
-    num_down_blocks = len({".".join(layer.split(".")[:3]) for layer in vae_state_dict if "encoder.down" in layer})
-    down_blocks = {layer_id: [key for key in vae_state_dict if f"down.{layer_id}" in key] for layer_id in range(num_down_blocks)}
+    num_down_blocks = len(
+        {
+            ".".join(layer.split(".")[:3])
+            for layer in vae_state_dict
+            if "encoder.down" in layer
+        }
+    )
+    down_blocks = {
+        layer_id: [key for key in vae_state_dict if f"down.{layer_id}" in key]
+        for layer_id in range(num_down_blocks)
+    }
 
     # Retrieves the keys for the decoder up blocks only
-    num_up_blocks = len({".".join(layer.split(".")[:3]) for layer in vae_state_dict if "decoder.up" in layer})
-    up_blocks = {layer_id: [key for key in vae_state_dict if f"up.{layer_id}" in key] for layer_id in range(num_up_blocks)}
+    num_up_blocks = len(
+        {
+            ".".join(layer.split(".")[:3])
+            for layer in vae_state_dict
+            if "decoder.up" in layer
+        }
+    )
+    up_blocks = {
+        layer_id: [key for key in vae_state_dict if f"up.{layer_id}" in key]
+        for layer_id in range(num_up_blocks)
+    }
 
     for i in range(num_down_blocks):
-        resnets = [key for key in down_blocks[i] if f"down.{i}" in key and f"down.{i}.downsample" not in key]
+        resnets = [
+            key
+            for key in down_blocks[i]
+            if f"down.{i}" in key and f"down.{i}.downsample" not in key
+        ]
 
         if f"encoder.down.{i}.downsample.conv.weight" in vae_state_dict:
-            new_checkpoint[f"encoder.down_blocks.{i}.downsamplers.0.conv.weight"] = vae_state_dict.pop(
-                f"encoder.down.{i}.downsample.conv.weight"
-            )
-            new_checkpoint[f"encoder.down_blocks.{i}.downsamplers.0.conv.bias"] = vae_state_dict.pop(
-                f"encoder.down.{i}.downsample.conv.bias"
-            )
+            new_checkpoint[
+                f"encoder.down_blocks.{i}.downsamplers.0.conv.weight"
+            ] = vae_state_dict.pop(f"encoder.down.{i}.downsample.conv.weight")
+            new_checkpoint[
+                f"encoder.down_blocks.{i}.downsamplers.0.conv.bias"
+            ] = vae_state_dict.pop(f"encoder.down.{i}.downsample.conv.bias")
 
         paths = renew_vae_resnet_paths(resnets)
         meta_path = {"old": f"down.{i}.block", "new": f"down_blocks.{i}.resnets"}
-        assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)
+        assign_to_checkpoint(
+            paths,
+            new_checkpoint,
+            vae_state_dict,
+            additional_replacements=[meta_path],
+            config=config,
+        )
 
     mid_resnets = [key for key in vae_state_dict if "encoder.mid.block" in key]
     num_mid_res_blocks = 2
@@ -454,29 +610,51 @@ def convert_ldm_vae_checkpoint(checkpoint, config):
 
         paths = renew_vae_resnet_paths(resnets)
         meta_path = {"old": f"mid.block_{i}", "new": f"mid_block.resnets.{i - 1}"}
-        assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)
+        assign_to_checkpoint(
+            paths,
+            new_checkpoint,
+            vae_state_dict,
+            additional_replacements=[meta_path],
+            config=config,
+        )
 
     mid_attentions = [key for key in vae_state_dict if "encoder.mid.attn" in key]
     paths = renew_vae_attention_paths(mid_attentions)
     meta_path = {"old": "mid.attn_1", "new": "mid_block.attentions.0"}
-    assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)
+    assign_to_checkpoint(
+        paths,
+        new_checkpoint,
+        vae_state_dict,
+        additional_replacements=[meta_path],
+        config=config,
+    )
     conv_attn_to_linear(new_checkpoint)
 
     for i in range(num_up_blocks):
         block_id = num_up_blocks - 1 - i
-        resnets = [key for key in up_blocks[block_id] if f"up.{block_id}" in key and f"up.{block_id}.upsample" not in key]
+        resnets = [
+            key
+            for key in up_blocks[block_id]
+            if f"up.{block_id}" in key and f"up.{block_id}.upsample" not in key
+        ]
 
         if f"decoder.up.{block_id}.upsample.conv.weight" in vae_state_dict:
-            new_checkpoint[f"decoder.up_blocks.{i}.upsamplers.0.conv.weight"] = vae_state_dict[
-                f"decoder.up.{block_id}.upsample.conv.weight"
-            ]
-            new_checkpoint[f"decoder.up_blocks.{i}.upsamplers.0.conv.bias"] = vae_state_dict[
-                f"decoder.up.{block_id}.upsample.conv.bias"
-            ]
+            new_checkpoint[
+                f"decoder.up_blocks.{i}.upsamplers.0.conv.weight"
+            ] = vae_state_dict[f"decoder.up.{block_id}.upsample.conv.weight"]
+            new_checkpoint[
+                f"decoder.up_blocks.{i}.upsamplers.0.conv.bias"
+            ] = vae_state_dict[f"decoder.up.{block_id}.upsample.conv.bias"]
 
         paths = renew_vae_resnet_paths(resnets)
         meta_path = {"old": f"up.{block_id}.block", "new": f"up_blocks.{i}.resnets"}
-        assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)
+        assign_to_checkpoint(
+            paths,
+            new_checkpoint,
+            vae_state_dict,
+            additional_replacements=[meta_path],
+            config=config,
+        )
 
     mid_resnets = [key for key in vae_state_dict if "decoder.mid.block" in key]
     num_mid_res_blocks = 2
@@ -485,12 +663,24 @@ def convert_ldm_vae_checkpoint(checkpoint, config):
 
         paths = renew_vae_resnet_paths(resnets)
         meta_path = {"old": f"mid.block_{i}", "new": f"mid_block.resnets.{i - 1}"}
-        assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)
+        assign_to_checkpoint(
+            paths,
+            new_checkpoint,
+            vae_state_dict,
+            additional_replacements=[meta_path],
+            config=config,
+        )
 
     mid_attentions = [key for key in vae_state_dict if "decoder.mid.attn" in key]
     paths = renew_vae_attention_paths(mid_attentions)
     meta_path = {"old": "mid.attn_1", "new": "mid_block.attentions.0"}
-    assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)
+    assign_to_checkpoint(
+        paths,
+        new_checkpoint,
+        vae_state_dict,
+        additional_replacements=[meta_path],
+        config=config,
+    )
     conv_attn_to_linear(new_checkpoint)
     return new_checkpoint
 
@@ -501,19 +691,29 @@ def create_unet_diffusers_config(v2, use_linear_projection_in_v2=False):
     """
     # unet_params = original_config.model.params.unet_config.params
 
-    block_out_channels = [UNET_PARAMS_MODEL_CHANNELS * mult for mult in UNET_PARAMS_CHANNEL_MULT]
+    block_out_channels = [
+        UNET_PARAMS_MODEL_CHANNELS * mult for mult in UNET_PARAMS_CHANNEL_MULT
+    ]
 
     down_block_types = []
     resolution = 1
     for i in range(len(block_out_channels)):
-        block_type = "CrossAttnDownBlock2D" if resolution in UNET_PARAMS_ATTENTION_RESOLUTIONS else "DownBlock2D"
+        block_type = (
+            "CrossAttnDownBlock2D"
+            if resolution in UNET_PARAMS_ATTENTION_RESOLUTIONS
+            else "DownBlock2D"
+        )
         down_block_types.append(block_type)
         if i != len(block_out_channels) - 1:
             resolution *= 2
 
     up_block_types = []
     for i in range(len(block_out_channels)):
-        block_type = "CrossAttnUpBlock2D" if resolution in UNET_PARAMS_ATTENTION_RESOLUTIONS else "UpBlock2D"
+        block_type = (
+            "CrossAttnUpBlock2D"
+            if resolution in UNET_PARAMS_ATTENTION_RESOLUTIONS
+            else "UpBlock2D"
+        )
         up_block_types.append(block_type)
         resolution //= 2
 
@@ -525,8 +725,12 @@ def create_unet_diffusers_config(v2, use_linear_projection_in_v2=False):
         up_block_types=tuple(up_block_types),
         block_out_channels=tuple(block_out_channels),
         layers_per_block=UNET_PARAMS_NUM_RES_BLOCKS,
-        cross_attention_dim=UNET_PARAMS_CONTEXT_DIM if not v2 else V2_UNET_PARAMS_CONTEXT_DIM,
-        attention_head_dim=UNET_PARAMS_NUM_HEADS if not v2 else V2_UNET_PARAMS_ATTENTION_HEAD_DIM,
+        cross_attention_dim=UNET_PARAMS_CONTEXT_DIM
+        if not v2
+        else V2_UNET_PARAMS_CONTEXT_DIM,
+        attention_head_dim=UNET_PARAMS_NUM_HEADS
+        if not v2
+        else V2_UNET_PARAMS_ATTENTION_HEAD_DIM,
         # use_linear_projection=UNET_PARAMS_USE_LINEAR_PROJECTION if not v2 else V2_UNET_PARAMS_USE_LINEAR_PROJECTION,
     )
     if v2 and use_linear_projection_in_v2:
@@ -563,7 +767,9 @@ def convert_ldm_clip_checkpoint_v1(checkpoint):
     text_model_dict = {}
     for key in keys:
         if key.startswith("cond_stage_model.transformer"):
-            text_model_dict[key[len("cond_stage_model.transformer.") :]] = checkpoint[key]
+            text_model_dict[key[len("cond_stage_model.transformer.") :]] = checkpoint[
+                key
+            ]
     return text_model_dict
 
 
@@ -592,13 +798,17 @@ def convert_key(key):
             else:
                 raise ValueError(f"unexpected key in SD: {key}")
         elif ".positional_embedding" in key:
-            key = key.replace(".positional_embedding", ".embeddings.position_embedding.weight")
+            key = key.replace(
+                ".positional_embedding", ".embeddings.position_embedding.weight"
+            )
         elif ".text_projection" in key:
             key = None  # 使われない???
         elif ".logit_scale" in key:
             key = None  # 使われない???
         elif ".token_embedding" in key:
-            key = key.replace(".token_embedding.weight", ".embeddings.token_embedding.weight")
+            key = key.replace(
+                ".token_embedding.weight", ".embeddings.token_embedding.weight"
+            )
         elif ".ln_final" in key:
             key = key.replace(".ln_final", ".final_layer_norm")
         return key
@@ -623,7 +833,10 @@ def convert_key(key):
             values = torch.chunk(checkpoint[key], 3)
 
             key_suffix = ".weight" if "weight" in key else ".bias"
-            key_pfx = key.replace("cond_stage_model.model.transformer.resblocks.", "text_model.encoder.layers.")
+            key_pfx = key.replace(
+                "cond_stage_model.model.transformer.resblocks.",
+                "text_model.encoder.layers.",
+            )
             key_pfx = key_pfx.replace("_weight", "")
             key_pfx = key_pfx.replace("_bias", "")
             key_pfx = key_pfx.replace(".attn.in_proj", ".self_attn.")
@@ -699,7 +912,9 @@ def convert_unet_state_dict_to_sd(v2, unet_state_dict):
                 # no attention layers in down_blocks.3
                 hf_down_atn_prefix = f"down_blocks.{i}.attentions.{j}."
                 sd_down_atn_prefix = f"input_blocks.{3*i + j + 1}.1."
-                unet_conversion_map_layer.append((sd_down_atn_prefix, hf_down_atn_prefix))
+                unet_conversion_map_layer.append(
+                    (sd_down_atn_prefix, hf_down_atn_prefix)
+                )
 
         for j in range(3):
             # loop over resnets/attentions for upblocks
@@ -717,7 +932,9 @@ def convert_unet_state_dict_to_sd(v2, unet_state_dict):
             # no downsample in down_blocks.3
             hf_downsample_prefix = f"down_blocks.{i}.downsamplers.0.conv."
             sd_downsample_prefix = f"input_blocks.{3*(i+1)}.0.op."
-            unet_conversion_map_layer.append((sd_downsample_prefix, hf_downsample_prefix))
+            unet_conversion_map_layer.append(
+                (sd_downsample_prefix, hf_downsample_prefix)
+            )
 
             # no upsample in up_blocks.3
             hf_upsample_prefix = f"up_blocks.{i}.upsamplers.0."
@@ -755,6 +972,7 @@ def convert_unet_state_dict_to_sd(v2, unet_state_dict):
 
     return new_state_dict
 
+
 def controlnet_conversion_map():
     unet_conversion_map = [
         ("time_embed.0.weight", "time_embedding.linear_1.weight"),
@@ -786,12 +1004,16 @@ def controlnet_conversion_map():
             if i < 3:
                 hf_down_atn_prefix = f"down_blocks.{i}.attentions.{j}."
                 sd_down_atn_prefix = f"input_blocks.{3*i + j + 1}.1."
-                unet_conversion_map_layer.append((sd_down_atn_prefix, hf_down_atn_prefix))
+                unet_conversion_map_layer.append(
+                    (sd_down_atn_prefix, hf_down_atn_prefix)
+                )
 
         if i < 3:
             hf_downsample_prefix = f"down_blocks.{i}.downsamplers.0.conv."
             sd_downsample_prefix = f"input_blocks.{3*(i+1)}.0.op."
-            unet_conversion_map_layer.append((sd_downsample_prefix, hf_downsample_prefix))
+            unet_conversion_map_layer.append(
+                (sd_downsample_prefix, hf_downsample_prefix)
+            )
 
     hf_mid_atn_prefix = "mid_block.attentions.0."
     sd_mid_atn_prefix = "middle_block.1."
@@ -819,7 +1041,11 @@ def controlnet_conversion_map():
 
 
 def convert_controlnet_state_dict_to_sd(controlnet_state_dict):
-    unet_conversion_map, unet_conversion_map_resnet, unet_conversion_map_layer = controlnet_conversion_map()
+    (
+        unet_conversion_map,
+        unet_conversion_map_resnet,
+        unet_conversion_map_layer,
+    ) = controlnet_conversion_map()
 
     mapping = {k: k for k in controlnet_state_dict.keys()}
     for sd_name, diffusers_name in unet_conversion_map:
@@ -836,8 +1062,13 @@ def convert_controlnet_state_dict_to_sd(controlnet_state_dict):
     new_state_dict = {v: controlnet_state_dict[k] for k, v in mapping.items()}
     return new_state_dict
 
+
 def convert_controlnet_state_dict_to_diffusers(controlnet_state_dict):
-    unet_conversion_map, unet_conversion_map_resnet, unet_conversion_map_layer = controlnet_conversion_map()
+    (
+        unet_conversion_map,
+        unet_conversion_map_resnet,
+        unet_conversion_map_layer,
+    ) = controlnet_conversion_map()
 
     mapping = {k: k for k in controlnet_state_dict.keys()}
     for sd_name, diffusers_name in unet_conversion_map:
@@ -854,6 +1085,7 @@ def convert_controlnet_state_dict_to_diffusers(controlnet_state_dict):
     new_state_dict = {v: controlnet_state_dict[k] for k, v in mapping.items()}
     return new_state_dict
 
+
 # ================#
 # VAE Conversion #
 # ================#
@@ -953,9 +1185,18 @@ def is_safetensors(path):
 def load_checkpoint_with_text_encoder_conversion(ckpt_path, device="cpu"):
     # text encoderの格納形式が違うモデルに対応する ('text_model'がない)
     TEXT_ENCODER_KEY_REPLACEMENTS = [
-        ("cond_stage_model.transformer.embeddings.", "cond_stage_model.transformer.text_model.embeddings."),
-        ("cond_stage_model.transformer.encoder.", "cond_stage_model.transformer.text_model.encoder."),
-        ("cond_stage_model.transformer.final_layer_norm.", "cond_stage_model.transformer.text_model.final_layer_norm."),
+        (
+            "cond_stage_model.transformer.embeddings.",
+            "cond_stage_model.transformer.text_model.embeddings.",
+        ),
+        (
+            "cond_stage_model.transformer.encoder.",
+            "cond_stage_model.transformer.text_model.encoder.",
+        ),
+        (
+            "cond_stage_model.transformer.final_layer_norm.",
+            "cond_stage_model.transformer.text_model.final_layer_norm.",
+        ),
     ]
 
     if is_safetensors(ckpt_path):
@@ -984,7 +1225,9 @@ def load_checkpoint_with_text_encoder_conversion(ckpt_path, device="cpu"):
 
 
 # TODO dtype指定の動作が怪しいので確認する text_encoderを指定形式で作れるか未確認
-def load_models_from_stable_diffusion_checkpoint(v2, ckpt_path, device="cpu", dtype=None, unet_use_linear_projection_in_v2=True):
+def load_models_from_stable_diffusion_checkpoint(
+    v2, ckpt_path, device="cpu", dtype=None, unet_use_linear_projection_in_v2=True
+):
     _, state_dict = load_checkpoint_with_text_encoder_conversion(ckpt_path, device)
 
     # Convert the UNet2DConditionModel model.
@@ -1005,7 +1248,9 @@ def load_models_from_stable_diffusion_checkpoint(v2, ckpt_path, device="cpu", dt
 
     # convert text_model
     if v2:
-        converted_text_encoder_checkpoint = convert_ldm_clip_checkpoint_v2(state_dict, 77)
+        converted_text_encoder_checkpoint = convert_ldm_clip_checkpoint_v2(
+            state_dict, 77
+        )
         cfg = CLIPTextConfig(
             vocab_size=49408,
             hidden_size=1024,
@@ -1028,7 +1273,9 @@ def load_models_from_stable_diffusion_checkpoint(v2, ckpt_path, device="cpu", dt
             transformers_version="4.25.0.dev0",
         )
         text_model = CLIPTextModel._from_config(cfg)
-        info = text_model.load_state_dict(converted_text_encoder_checkpoint)
+        info = text_model.load_state_dict(
+            converted_text_encoder_checkpoint, strict=False
+        )
     else:
         converted_text_encoder_checkpoint = convert_ldm_clip_checkpoint_v1(state_dict)
 
@@ -1057,11 +1304,14 @@ def load_models_from_stable_diffusion_checkpoint(v2, ckpt_path, device="cpu", dt
             torch_dtype="float32",
         )
         text_model = CLIPTextModel._from_config(cfg)
-        info = text_model.load_state_dict(converted_text_encoder_checkpoint)
+        info = text_model.load_state_dict(
+            converted_text_encoder_checkpoint, strict=False
+        )
     print("loading text encoder:", info)
 
     return text_model, vae, unet
 
+
 def get_model_version_str_for_sd1_sd2(v2, v_parameterization):
     # only for reference
     version_str = "sd"
@@ -1073,6 +1323,7 @@ def get_model_version_str_for_sd1_sd2(v2, v_parameterization):
         version_str += "_v"
     return version_str
 
+
 def convert_text_encoder_state_dict_to_sd_v2(checkpoint, make_dummy_weights=False):
     def convert_key(key):
         # position_idsの除去
@@ -1097,9 +1348,13 @@ def convert_key(key):
             else:
                 raise ValueError(f"unexpected key in DiffUsers model: {key}")
         elif ".position_embedding" in key:
-            key = key.replace("embeddings.position_embedding.weight", "positional_embedding")
+            key = key.replace(
+                "embeddings.position_embedding.weight", "positional_embedding"
+            )
         elif ".token_embedding" in key:
-            key = key.replace("embeddings.token_embedding.weight", "token_embedding.weight")
+            key = key.replace(
+                "embeddings.token_embedding.weight", "token_embedding.weight"
+            )
         elif "final_layer_norm" in key:
             key = key.replace("final_layer_norm", "ln_final")
         return key
@@ -1125,7 +1380,9 @@ def convert_key(key):
             value_v = checkpoint[key_v]
             value = torch.cat([value_q, value_k, value_v])
 
-            new_key = key.replace("text_model.encoder.layers.", "transformer.resblocks.")
+            new_key = key.replace(
+                "text_model.encoder.layers.", "transformer.resblocks."
+            )
             new_key = new_key.replace(".self_attn.q_proj.", ".attn.in_proj_")
             new_sd[new_key] = value
 
@@ -1135,16 +1392,30 @@ def convert_key(key):
         keys = list(new_sd.keys())
         for key in keys:
             if key.startswith("transformer.resblocks.22."):
-                new_sd[key.replace(".22.", ".23.")] = new_sd[key].clone()  # copyしないとsafetensorsの保存で落ちる
+                new_sd[key.replace(".22.", ".23.")] = new_sd[
+                    key
+                ].clone()  # copyしないとsafetensorsの保存で落ちる
 
         # Diffusersに含まれない重みを作っておく
-        new_sd["text_projection"] = torch.ones((1024, 1024), dtype=new_sd[keys[0]].dtype, device=new_sd[keys[0]].device)
+        new_sd["text_projection"] = torch.ones(
+            (1024, 1024), dtype=new_sd[keys[0]].dtype, device=new_sd[keys[0]].device
+        )
         new_sd["logit_scale"] = torch.tensor(1)
 
     return new_sd
 
 
-def save_stable_diffusion_checkpoint(v2, output_file, text_encoder, unet, ckpt_path, epochs, steps, save_dtype=None, vae=None):
+def save_stable_diffusion_checkpoint(
+    v2,
+    output_file,
+    text_encoder,
+    unet,
+    ckpt_path,
+    epochs,
+    steps,
+    save_dtype=None,
+    vae=None,
+):
     if ckpt_path is not None:
         # epoch/stepを参照する。またVAEがメモリ上にないときなど、もう一度VAEを含めて読み込む
         checkpoint, state_dict = load_checkpoint_with_text_encoder_conversion(ckpt_path)
@@ -1157,7 +1428,9 @@ def save_stable_diffusion_checkpoint(v2, output_file, text_encoder, unet, ckpt_p
             del state_dict["state_dict"]
     else:
         # 新しく作る
-        assert vae is not None, "VAE is required to save a checkpoint without a given checkpoint"
+        assert (
+            vae is not None
+        ), "VAE is required to save a checkpoint without a given checkpoint"
         checkpoint = {}
         state_dict = {}
         strict = False
@@ -1176,8 +1449,12 @@ def update_sd(prefix, sd):
 
     # Convert the text encoder model
     if v2:
-        make_dummy = ckpt_path is None  # 参照元のcheckpointがない場合は最後の層を前の層から複製して作るなどダミーの重みを入れる
-        text_enc_dict = convert_text_encoder_state_dict_to_sd_v2(text_encoder.state_dict(), make_dummy)
+        make_dummy = (
+            ckpt_path is None
+        )  # 参照元のcheckpointがない場合は最後の層を前の層から複製して作るなどダミーの重みを入れる
+        text_enc_dict = convert_text_encoder_state_dict_to_sd_v2(
+            text_encoder.state_dict(), make_dummy
+        )
         update_sd("cond_stage_model.model.", text_enc_dict)
     else:
         text_enc_dict = text_encoder.state_dict()
@@ -1213,7 +1490,15 @@ def update_sd(prefix, sd):
     return key_count
 
 
-def save_diffusers_checkpoint(v2, output_dir, text_encoder, unet, pretrained_model_name_or_path, vae=None, use_safetensors=False):
+def save_diffusers_checkpoint(
+    v2,
+    output_dir,
+    text_encoder,
+    unet,
+    pretrained_model_name_or_path,
+    vae=None,
+    use_safetensors=False,
+):
     if pretrained_model_name_or_path is None:
         # load default settings for v1/v2
         if v2:
@@ -1221,10 +1506,16 @@ def save_diffusers_checkpoint(v2, output_dir, text_encoder, unet, pretrained_mod
         else:
             pretrained_model_name_or_path = DIFFUSERS_REF_MODEL_ID_V1
 
-    scheduler = DDIMScheduler.from_pretrained(pretrained_model_name_or_path, subfolder="scheduler")
-    tokenizer = CLIPTokenizer.from_pretrained(pretrained_model_name_or_path, subfolder="tokenizer")
+    scheduler = DDIMScheduler.from_pretrained(
+        pretrained_model_name_or_path, subfolder="scheduler"
+    )
+    tokenizer = CLIPTokenizer.from_pretrained(
+        pretrained_model_name_or_path, subfolder="tokenizer"
+    )
     if vae is None:
-        vae = AutoencoderKL.from_pretrained(pretrained_model_name_or_path, subfolder="vae")
+        vae = AutoencoderKL.from_pretrained(
+            pretrained_model_name_or_path, subfolder="vae"
+        )
 
     pipeline = StableDiffusionPipeline(
         unet=unet,
@@ -1247,11 +1538,15 @@ def load_vae(vae_id, dtype):
     if os.path.isdir(vae_id) or not os.path.isfile(vae_id):
         # Diffusers local/remote
         try:
-            vae = AutoencoderKL.from_pretrained(vae_id, subfolder=None, torch_dtype=dtype)
+            vae = AutoencoderKL.from_pretrained(
+                vae_id, subfolder=None, torch_dtype=dtype
+            )
         except EnvironmentError as e:
             print(f"exception occurs in loading vae: {e}")
             print("retry with subfolder='vae'")
-            vae = AutoencoderKL.from_pretrained(vae_id, subfolder="vae", torch_dtype=dtype)
+            vae = AutoencoderKL.from_pretrained(
+                vae_id, subfolder="vae", torch_dtype=dtype
+            )
         return vae
 
     # local
@@ -1262,7 +1557,11 @@ def load_vae(vae_id, dtype):
         converted_vae_checkpoint = torch.load(vae_id, map_location="cpu")
     else:
         # StableDiffusion
-        vae_model = load_file(vae_id, "cpu") if is_safetensors(vae_id) else torch.load(vae_id, map_location="cpu")
+        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
diff --git a/lycoris/kohya/original_unet.py b/lycoris/kohya/original_unet.py
index 4def788..2d1ff36 100644
--- a/lycoris/kohya/original_unet.py
+++ b/lycoris/kohya/original_unet.py
@@ -128,8 +128,18 @@
 NORM_EPS: float = 1e-5
 TRANSFORMER_NORM_NUM_GROUPS = 32
 
-DOWN_BLOCK_TYPES = ["CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "DownBlock2D"]
-UP_BLOCK_TYPES = ["UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D"]
+DOWN_BLOCK_TYPES = [
+    "CrossAttnDownBlock2D",
+    "CrossAttnDownBlock2D",
+    "CrossAttnDownBlock2D",
+    "DownBlock2D",
+]
+UP_BLOCK_TYPES = [
+    "UpBlock2D",
+    "CrossAttnUpBlock2D",
+    "CrossAttnUpBlock2D",
+    "CrossAttnUpBlock2D",
+]
 
 
 # region memory effcient attention
@@ -171,7 +181,9 @@ def forward(ctx, q, k, v, mask, causal, q_bucket_size, k_bucket_size):
 
         o = torch.zeros_like(q)
         all_row_sums = torch.zeros((*q.shape[:-1], 1), dtype=dtype, device=device)
-        all_row_maxes = torch.full((*q.shape[:-1], 1), max_neg_value, dtype=dtype, device=device)
+        all_row_maxes = torch.full(
+            (*q.shape[:-1], 1), max_neg_value, dtype=dtype, device=device
+        )
 
         scale = q.shape[-1] ** -0.5
 
@@ -200,15 +212,17 @@ def forward(ctx, q, k, v, mask, causal, q_bucket_size, k_bucket_size):
             for k_ind, (kc, vc) in enumerate(col_splits):
                 k_start_index = k_ind * k_bucket_size
 
-                attn_weights = torch.einsum("... i d, ... j d -> ... i j", qc, kc) * scale
+                attn_weights = (
+                    torch.einsum("... i d, ... j d -> ... i j", qc, kc) * scale
+                )
 
                 if exists(row_mask):
                     attn_weights.masked_fill_(~row_mask, max_neg_value)
 
                 if causal and q_start_index < (k_start_index + k_bucket_size - 1):
-                    causal_mask = torch.ones((qc.shape[-2], kc.shape[-2]), dtype=torch.bool, device=device).triu(
-                        q_start_index - k_start_index + 1
-                    )
+                    causal_mask = torch.ones(
+                        (qc.shape[-2], kc.shape[-2]), dtype=torch.bool, device=device
+                    ).triu(q_start_index - k_start_index + 1)
                     attn_weights.masked_fill_(causal_mask, max_neg_value)
 
                 block_row_maxes = attn_weights.amax(dim=-1, keepdims=True)
@@ -218,18 +232,27 @@ def forward(ctx, q, k, v, mask, causal, q_bucket_size, k_bucket_size):
                 if exists(row_mask):
                     exp_weights.masked_fill_(~row_mask, 0.0)
 
-                block_row_sums = exp_weights.sum(dim=-1, keepdims=True).clamp(min=EPSILON)
+                block_row_sums = exp_weights.sum(dim=-1, keepdims=True).clamp(
+                    min=EPSILON
+                )
 
                 new_row_maxes = torch.maximum(block_row_maxes, row_maxes)
 
-                exp_values = torch.einsum("... i j, ... j d -> ... i d", exp_weights, vc)
+                exp_values = torch.einsum(
+                    "... i j, ... j d -> ... i d", exp_weights, vc
+                )
 
                 exp_row_max_diff = torch.exp(row_maxes - new_row_maxes)
                 exp_block_row_max_diff = torch.exp(block_row_maxes - new_row_maxes)
 
-                new_row_sums = exp_row_max_diff * row_sums + exp_block_row_max_diff * block_row_sums
+                new_row_sums = (
+                    exp_row_max_diff * row_sums
+                    + exp_block_row_max_diff * block_row_sums
+                )
 
-                oc.mul_((row_sums / new_row_sums) * exp_row_max_diff).add_((exp_block_row_max_diff / new_row_sums) * exp_values)
+                oc.mul_((row_sums / new_row_sums) * exp_row_max_diff).add_(
+                    (exp_block_row_max_diff / new_row_sums) * exp_values
+                )
 
                 row_maxes.copy_(new_row_maxes)
                 row_sums.copy_(new_row_sums)
@@ -279,12 +302,14 @@ def backward(ctx, do):
             for k_ind, (kc, vc, dkc, dvc) in enumerate(col_splits):
                 k_start_index = k_ind * k_bucket_size
 
-                attn_weights = torch.einsum("... i d, ... j d -> ... i j", qc, kc) * scale
+                attn_weights = (
+                    torch.einsum("... i d, ... j d -> ... i j", qc, kc) * scale
+                )
 
                 if causal and q_start_index < (k_start_index + k_bucket_size - 1):
-                    causal_mask = torch.ones((qc.shape[-2], kc.shape[-2]), dtype=torch.bool, device=device).triu(
-                        q_start_index - k_start_index + 1
-                    )
+                    causal_mask = torch.ones(
+                        (qc.shape[-2], kc.shape[-2]), dtype=torch.bool, device=device
+                    ).triu(q_start_index - k_start_index + 1)
                     attn_weights.masked_fill_(causal_mask, max_neg_value)
 
                 exp_attn_weights = torch.exp(attn_weights - mc)
@@ -340,7 +365,9 @@ def get_timestep_embedding(
     assert len(timesteps.shape) == 1, "Timesteps should be a 1d-array"
 
     half_dim = embedding_dim // 2
-    exponent = -math.log(max_period) * torch.arange(start=0, end=half_dim, dtype=torch.float32, device=timesteps.device)
+    exponent = -math.log(max_period) * torch.arange(
+        start=0, end=half_dim, dtype=torch.float32, device=timesteps.device
+    )
     exponent = exponent / (half_dim - downscale_freq_shift)
 
     emb = torch.exp(exponent)
@@ -368,7 +395,13 @@ def __init__(self, sample):
 
 
 class TimestepEmbedding(nn.Module):
-    def __init__(self, in_channels: int, time_embed_dim: int, act_fn: str = "silu", out_dim: int = None):
+    def __init__(
+        self,
+        in_channels: int,
+        time_embed_dim: int,
+        act_fn: str = "silu",
+        out_dim: int = None,
+    ):
         super().__init__()
 
         self.linear_1 = nn.Linear(in_channels, time_embed_dim)
@@ -395,7 +428,9 @@ def forward(self, sample):
 
 
 class Timesteps(nn.Module):
-    def __init__(self, num_channels: int, flip_sin_to_cos: bool, downscale_freq_shift: float):
+    def __init__(
+        self, num_channels: int, flip_sin_to_cos: bool, downscale_freq_shift: float
+    ):
         super().__init__()
         self.num_channels = num_channels
         self.flip_sin_to_cos = flip_sin_to_cos
@@ -421,14 +456,22 @@ def __init__(
         self.in_channels = in_channels
         self.out_channels = out_channels
 
-        self.norm1 = torch.nn.GroupNorm(num_groups=NORM_GROUPS, num_channels=in_channels, eps=NORM_EPS, affine=True)
+        self.norm1 = torch.nn.GroupNorm(
+            num_groups=NORM_GROUPS, num_channels=in_channels, eps=NORM_EPS, affine=True
+        )
 
-        self.conv1 = torch.nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
+        self.conv1 = torch.nn.Conv2d(
+            in_channels, out_channels, kernel_size=3, stride=1, padding=1
+        )
 
         self.time_emb_proj = torch.nn.Linear(TIME_EMBED_DIM, out_channels)
 
-        self.norm2 = torch.nn.GroupNorm(num_groups=NORM_GROUPS, num_channels=out_channels, eps=NORM_EPS, affine=True)
-        self.conv2 = torch.nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
+        self.norm2 = torch.nn.GroupNorm(
+            num_groups=NORM_GROUPS, num_channels=out_channels, eps=NORM_EPS, affine=True
+        )
+        self.conv2 = torch.nn.Conv2d(
+            out_channels, out_channels, kernel_size=3, stride=1, padding=1
+        )
 
         # if non_linearity == "swish":
         self.nonlinearity = lambda x: F.silu(x)
@@ -437,7 +480,9 @@ def __init__(
 
         self.conv_shortcut = None
         if self.use_in_shortcut:
-            self.conv_shortcut = torch.nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
+            self.conv_shortcut = torch.nn.Conv2d(
+                in_channels, out_channels, kernel_size=1, stride=1, padding=0
+            )
 
     def forward(self, input_tensor, temb):
         hidden_states = input_tensor
@@ -510,7 +555,9 @@ def custom_forward(*inputs):
 
                     return custom_forward
 
-                hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(resnet), hidden_states, temb)
+                hidden_states = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(resnet), hidden_states, temb
+                )
             else:
                 hidden_states = resnet(hidden_states, temb)
 
@@ -552,7 +599,9 @@ def __init__(
     ):
         super().__init__()
         inner_dim = dim_head * heads
-        cross_attention_dim = cross_attention_dim if cross_attention_dim is not None else query_dim
+        cross_attention_dim = (
+            cross_attention_dim if cross_attention_dim is not None else query_dim
+        )
         self.upcast_attention = upcast_attention
 
         self.scale = dim_head**-0.5
@@ -581,14 +630,18 @@ def reshape_heads_to_batch_dim(self, tensor):
         batch_size, seq_len, dim = tensor.shape
         head_size = self.heads
         tensor = tensor.reshape(batch_size, seq_len, head_size, dim // head_size)
-        tensor = tensor.permute(0, 2, 1, 3).reshape(batch_size * head_size, seq_len, dim // head_size)
+        tensor = tensor.permute(0, 2, 1, 3).reshape(
+            batch_size * head_size, seq_len, dim // head_size
+        )
         return tensor
 
     def reshape_batch_dim_to_heads(self, tensor):
         batch_size, seq_len, dim = tensor.shape
         head_size = self.heads
         tensor = tensor.reshape(batch_size // head_size, head_size, seq_len, dim)
-        tensor = tensor.permute(0, 2, 1, 3).reshape(batch_size // head_size, seq_len, dim * head_size)
+        tensor = tensor.permute(0, 2, 1, 3).reshape(
+            batch_size // head_size, seq_len, dim * head_size
+        )
         return tensor
 
     def forward(self, hidden_states, context=None, mask=None):
@@ -621,7 +674,13 @@ def _attention(self, query, key, value):
             key = key.float()
 
         attention_scores = torch.baddbmm(
-            torch.empty(query.shape[0], query.shape[1], key.shape[1], dtype=query.dtype, device=query.device),
+            torch.empty(
+                query.shape[0],
+                query.shape[1],
+                key.shape[1],
+                dtype=query.dtype,
+                device=query.device,
+            ),
             query,
             key.transpose(-1, -2),
             beta=0,
@@ -650,13 +709,17 @@ def forward_memory_efficient_xformers(self, x, context=None, mask=None):
         k_in = self.to_k(context)
         v_in = self.to_v(context)
 
-        q, k, v = map(lambda t: rearrange(t, "b n (h d) -> b n h d", h=h), (q_in, k_in, v_in))
+        q, k, v = map(
+            lambda t: rearrange(t, "b n (h d) -> b n h d", h=h), (q_in, k_in, v_in)
+        )
         del q_in, k_in, v_in
 
         q = q.contiguous()
         k = k.contiguous()
         v = v.contiguous()
-        out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=None)  # 最適なのを選んでくれる
+        out = xformers.ops.memory_efficient_attention(
+            q, k, v, attn_bias=None
+        )  # 最適なのを選んでくれる
 
         out = rearrange(out, "b n h d -> b n (h d)", h=h)
 
@@ -694,10 +757,14 @@ def forward_sdpa(self, x, context=None, mask=None):
         k_in = self.to_k(context)
         v_in = self.to_v(context)
 
-        q, k, v = map(lambda t: rearrange(t, "b n (h d) -> b h n d", h=h), (q_in, k_in, v_in))
+        q, k, v = map(
+            lambda t: rearrange(t, "b n (h d) -> b h n d", h=h), (q_in, k_in, v_in)
+        )
         del q_in, k_in, v_in
 
-        out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False)
+        out = F.scaled_dot_product_attention(
+            q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False
+        )
 
         out = rearrange(out, "b h n d -> b n (h d)", h=h)
 
@@ -754,7 +821,12 @@ def forward(self, hidden_states):
 
 class BasicTransformerBlock(nn.Module):
     def __init__(
-        self, dim: int, num_attention_heads: int, attention_head_dim: int, cross_attention_dim: int, upcast_attention: bool = False
+        self,
+        dim: int,
+        num_attention_heads: int,
+        attention_head_dim: int,
+        cross_attention_dim: int,
+        upcast_attention: bool = False,
     ):
         super().__init__()
 
@@ -824,12 +896,19 @@ def __init__(
         inner_dim = num_attention_heads * attention_head_dim
         self.use_linear_projection = use_linear_projection
 
-        self.norm = torch.nn.GroupNorm(num_groups=TRANSFORMER_NORM_NUM_GROUPS, num_channels=in_channels, eps=1e-6, affine=True)
+        self.norm = torch.nn.GroupNorm(
+            num_groups=TRANSFORMER_NORM_NUM_GROUPS,
+            num_channels=in_channels,
+            eps=1e-6,
+            affine=True,
+        )
 
         if use_linear_projection:
             self.proj_in = nn.Linear(in_channels, inner_dim)
         else:
-            self.proj_in = nn.Conv2d(in_channels, inner_dim, kernel_size=1, stride=1, padding=0)
+            self.proj_in = nn.Conv2d(
+                in_channels, inner_dim, kernel_size=1, stride=1, padding=0
+            )
 
         self.transformer_blocks = nn.ModuleList(
             [
@@ -846,7 +925,9 @@ def __init__(
         if use_linear_projection:
             self.proj_out = nn.Linear(in_channels, inner_dim)
         else:
-            self.proj_out = nn.Conv2d(inner_dim, in_channels, kernel_size=1, stride=1, padding=0)
+            self.proj_out = nn.Conv2d(
+                inner_dim, in_channels, kernel_size=1, stride=1, padding=0
+            )
 
     def set_use_memory_efficient_attention(self, xformers, mem_eff):
         for transformer in self.transformer_blocks:
@@ -856,7 +937,13 @@ def set_use_sdpa(self, sdpa):
         for transformer in self.transformer_blocks:
             transformer.set_use_sdpa(sdpa)
 
-    def forward(self, hidden_states, encoder_hidden_states=None, timestep=None, return_dict: bool = True):
+    def forward(
+        self,
+        hidden_states,
+        encoder_hidden_states=None,
+        timestep=None,
+        return_dict: bool = True,
+    ):
         # 1. Input
         batch, _, height, weight = hidden_states.shape
         residual = hidden_states
@@ -865,23 +952,37 @@ def forward(self, hidden_states, encoder_hidden_states=None, timestep=None, retu
         if not self.use_linear_projection:
             hidden_states = self.proj_in(hidden_states)
             inner_dim = hidden_states.shape[1]
-            hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(batch, height * weight, inner_dim)
+            hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(
+                batch, height * weight, inner_dim
+            )
         else:
             inner_dim = hidden_states.shape[1]
-            hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(batch, height * weight, inner_dim)
+            hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(
+                batch, height * weight, inner_dim
+            )
             hidden_states = self.proj_in(hidden_states)
 
         # 2. Blocks
         for block in self.transformer_blocks:
-            hidden_states = block(hidden_states, context=encoder_hidden_states, timestep=timestep)
+            hidden_states = block(
+                hidden_states, context=encoder_hidden_states, timestep=timestep
+            )
 
         # 3. Output
         if not self.use_linear_projection:
-            hidden_states = hidden_states.reshape(batch, height, weight, inner_dim).permute(0, 3, 1, 2).contiguous()
+            hidden_states = (
+                hidden_states.reshape(batch, height, weight, inner_dim)
+                .permute(0, 3, 1, 2)
+                .contiguous()
+            )
             hidden_states = self.proj_out(hidden_states)
         else:
             hidden_states = self.proj_out(hidden_states)
-            hidden_states = hidden_states.reshape(batch, height, weight, inner_dim).permute(0, 3, 1, 2).contiguous()
+            hidden_states = (
+                hidden_states.reshape(batch, height, weight, inner_dim)
+                .permute(0, 3, 1, 2)
+                .contiguous()
+            )
 
         output = hidden_states + residual
 
@@ -912,7 +1013,9 @@ def __init__(
         for i in range(LAYERS_PER_BLOCK):
             in_channels = in_channels if i == 0 else out_channels
 
-            resnets.append(ResnetBlock2D(in_channels=in_channels, out_channels=out_channels))
+            resnets.append(
+                ResnetBlock2D(in_channels=in_channels, out_channels=out_channels)
+            )
             attentions.append(
                 Transformer2DModel(
                     attn_num_head_channels,
@@ -927,7 +1030,9 @@ def __init__(
         self.resnets = nn.ModuleList(resnets)
 
         if add_downsample:
-            self.downsamplers = nn.ModuleList([Downsample2D(out_channels, out_channels)])
+            self.downsamplers = nn.ModuleList(
+                [Downsample2D(out_channels, out_channels)]
+            )
         else:
             self.downsamplers = None
 
@@ -956,13 +1061,19 @@ def custom_forward(*inputs):
 
                     return custom_forward
 
-                hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(resnet), hidden_states, temb)
                 hidden_states = torch.utils.checkpoint.checkpoint(
-                    create_custom_forward(attn, return_dict=False), hidden_states, encoder_hidden_states
+                    create_custom_forward(resnet), hidden_states, temb
+                )
+                hidden_states = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(attn, return_dict=False),
+                    hidden_states,
+                    encoder_hidden_states,
                 )[0]
             else:
                 hidden_states = resnet(hidden_states, temb)
-                hidden_states = attn(hidden_states, encoder_hidden_states=encoder_hidden_states).sample
+                hidden_states = attn(
+                    hidden_states, encoder_hidden_states=encoder_hidden_states
+                ).sample
 
             output_states += (hidden_states,)
 
@@ -1039,10 +1150,14 @@ def custom_forward(*inputs):
 
                 if attn is not None:
                     hidden_states = torch.utils.checkpoint.checkpoint(
-                        create_custom_forward(attn, return_dict=False), hidden_states, encoder_hidden_states
+                        create_custom_forward(attn, return_dict=False),
+                        hidden_states,
+                        encoder_hidden_states,
                     )[0]
 
-                hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(resnet), hidden_states, temb)
+                hidden_states = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(resnet), hidden_states, temb
+                )
             else:
                 if attn is not None:
                     hidden_states = attn(hidden_states, encoder_hidden_states).sample
@@ -1074,9 +1189,13 @@ def forward(self, hidden_states, output_size):
 
         # if `output_size` is passed we force the interpolation output size and do not make use of `scale_factor=2`
         if output_size is None:
-            hidden_states = F.interpolate(hidden_states, scale_factor=2.0, mode="nearest")
+            hidden_states = F.interpolate(
+                hidden_states, scale_factor=2.0, mode="nearest"
+            )
         else:
-            hidden_states = F.interpolate(hidden_states, size=output_size, mode="nearest")
+            hidden_states = F.interpolate(
+                hidden_states, size=output_size, mode="nearest"
+            )
 
         # If the input is bfloat16, we cast back to bfloat16
         if dtype == torch.bfloat16:
@@ -1101,7 +1220,9 @@ def __init__(
         resnets = []
 
         for i in range(LAYERS_PER_BLOCK_UP):
-            res_skip_channels = in_channels if (i == LAYERS_PER_BLOCK_UP - 1) else out_channels
+            res_skip_channels = (
+                in_channels if (i == LAYERS_PER_BLOCK_UP - 1) else out_channels
+            )
             resnet_in_channels = prev_output_channel if i == 0 else out_channels
 
             resnets.append(
@@ -1126,7 +1247,9 @@ def set_use_memory_efficient_attention(self, xformers, mem_eff):
     def set_use_sdpa(self, sdpa):
         pass
 
-    def forward(self, hidden_states, res_hidden_states_tuple, temb=None, upsample_size=None):
+    def forward(
+        self, hidden_states, res_hidden_states_tuple, temb=None, upsample_size=None
+    ):
         for resnet in self.resnets:
             # pop res hidden states
             res_hidden_states = res_hidden_states_tuple[-1]
@@ -1141,7 +1264,9 @@ def custom_forward(*inputs):
 
                     return custom_forward
 
-                hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(resnet), hidden_states, temb)
+                hidden_states = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(resnet), hidden_states, temb
+                )
             else:
                 hidden_states = resnet(hidden_states, temb)
 
@@ -1172,7 +1297,9 @@ def __init__(
         self.attn_num_head_channels = attn_num_head_channels
 
         for i in range(LAYERS_PER_BLOCK_UP):
-            res_skip_channels = in_channels if (i == LAYERS_PER_BLOCK_UP - 1) else out_channels
+            res_skip_channels = (
+                in_channels if (i == LAYERS_PER_BLOCK_UP - 1) else out_channels
+            )
             resnet_in_channels = prev_output_channel if i == 0 else out_channels
 
             resnets.append(
@@ -1235,13 +1362,19 @@ def custom_forward(*inputs):
 
                     return custom_forward
 
-                hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(resnet), hidden_states, temb)
                 hidden_states = torch.utils.checkpoint.checkpoint(
-                    create_custom_forward(attn, return_dict=False), hidden_states, encoder_hidden_states
+                    create_custom_forward(resnet), hidden_states, temb
+                )
+                hidden_states = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(attn, return_dict=False),
+                    hidden_states,
+                    encoder_hidden_states,
                 )[0]
             else:
                 hidden_states = resnet(hidden_states, temb)
-                hidden_states = attn(hidden_states, encoder_hidden_states=encoder_hidden_states).sample
+                hidden_states = attn(
+                    hidden_states, encoder_hidden_states=encoder_hidden_states
+                ).sample
 
         if self.upsamplers is not None:
             for upsampler in self.upsamplers:
@@ -1337,10 +1470,14 @@ def __init__(
         # state_dictの書式が変わるのでmoduleの持ち方は変えられない
 
         # input
-        self.conv_in = nn.Conv2d(IN_CHANNELS, BLOCK_OUT_CHANNELS[0], kernel_size=3, padding=(1, 1))
+        self.conv_in = nn.Conv2d(
+            IN_CHANNELS, BLOCK_OUT_CHANNELS[0], kernel_size=3, padding=(1, 1)
+        )
 
         # time
-        self.time_proj = Timesteps(BLOCK_OUT_CHANNELS[0], TIME_EMBED_FLIP_SIN_TO_COS, TIME_EMBED_FREQ_SHIFT)
+        self.time_proj = Timesteps(
+            BLOCK_OUT_CHANNELS[0], TIME_EMBED_FLIP_SIN_TO_COS, TIME_EMBED_FREQ_SHIFT
+        )
 
         self.time_embedding = TimestepEmbedding(TIMESTEP_INPUT_DIM, TIME_EMBED_DIM)
 
@@ -1390,7 +1527,9 @@ def __init__(
 
             prev_output_channel = output_channel
             output_channel = reversed_block_out_channels[i]
-            input_channel = reversed_block_out_channels[min(i + 1, len(BLOCK_OUT_CHANNELS) - 1)]
+            input_channel = reversed_block_out_channels[
+                min(i + 1, len(BLOCK_OUT_CHANNELS) - 1)
+            ]
 
             # add upsample block for all BUT final layer
             if not is_final_block:
@@ -1414,9 +1553,13 @@ def __init__(
             prev_output_channel = output_channel
 
         # out
-        self.conv_norm_out = nn.GroupNorm(num_channels=BLOCK_OUT_CHANNELS[0], num_groups=NORM_GROUPS, eps=NORM_EPS)
+        self.conv_norm_out = nn.GroupNorm(
+            num_channels=BLOCK_OUT_CHANNELS[0], num_groups=NORM_GROUPS, eps=NORM_EPS
+        )
         self.conv_act = nn.SiLU()
-        self.conv_out = nn.Conv2d(BLOCK_OUT_CHANNELS[0], OUT_CHANNELS, kernel_size=3, padding=1)
+        self.conv_out = nn.Conv2d(
+            BLOCK_OUT_CHANNELS[0], OUT_CHANNELS, kernel_size=3, padding=1
+        )
 
     # region diffusers compatibility
     def prepare_config(self):
@@ -1436,7 +1579,10 @@ def set_attention_slice(self, slice_size):
         raise NotImplementedError("Attention slicing is not supported for this model.")
 
     def is_gradient_checkpointing(self) -> bool:
-        return any(hasattr(m, "gradient_checkpointing") and m.gradient_checkpointing for m in self.modules())
+        return any(
+            hasattr(m, "gradient_checkpointing") and m.gradient_checkpointing
+            for m in self.modules()
+        )
 
     def enable_gradient_checkpointing(self):
         self.set_gradient_checkpointing(value=True)
@@ -1544,7 +1690,9 @@ def forward(
             down_block_res_samples = tuple(down_block_res_samples)
 
         # 4. mid
-        sample = self.mid_block(sample, emb, encoder_hidden_states=encoder_hidden_states)
+        sample = self.mid_block(
+            sample, emb, encoder_hidden_states=encoder_hidden_states
+        )
 
         # ControlNetの出力を追加する
         if mid_block_additional_residual is not None:
@@ -1555,7 +1703,9 @@ def forward(
             is_final_block = i == len(self.up_blocks) - 1
 
             res_samples = down_block_res_samples[-len(upsample_block.resnets) :]
-            down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)]  # skip connection
+            down_block_res_samples = down_block_res_samples[
+                : -len(upsample_block.resnets)
+            ]  # skip connection
 
             # if we have not reached the final block and need to forward the upsample size, we do it here
             # 前述のように最後のブロック以外ではupsample_sizeを伝える
@@ -1572,7 +1722,10 @@ def forward(
                 )
             else:
                 sample = upsample_block(
-                    hidden_states=sample, temb=emb, res_hidden_states_tuple=res_samples, upsample_size=upsample_size
+                    hidden_states=sample,
+                    temb=emb,
+                    res_hidden_states_tuple=res_samples,
+                    upsample_size=upsample_size,
                 )
 
         # 6. post-process
diff --git a/lycoris/kohya/sdxl_model_util.py b/lycoris/kohya/sdxl_model_util.py
index 6abb228..f98ace4 100644
--- a/lycoris/kohya/sdxl_model_util.py
+++ b/lycoris/kohya/sdxl_model_util.py
@@ -1,6 +1,11 @@
 import torch
 from safetensors.torch import load_file, save_file
-from transformers import CLIPTextModel, CLIPTextConfig, CLIPTextModelWithProjection, CLIPTokenizer
+from transformers import (
+    CLIPTextModel,
+    CLIPTextConfig,
+    CLIPTextModelWithProjection,
+    CLIPTokenizer,
+)
 from diffusers import AutoencoderKL, EulerDiscreteScheduler, UNet2DConditionModel
 from . import model_utils as model_util
 from . import sdxl_original_unet
@@ -86,13 +91,17 @@ def convert_key(key):
             else:
                 raise ValueError(f"unexpected key in SD: {key}")
         elif ".positional_embedding" in key:
-            key = key.replace(".positional_embedding", ".embeddings.position_embedding.weight")
+            key = key.replace(
+                ".positional_embedding", ".embeddings.position_embedding.weight"
+            )
         elif ".text_projection" in key:
             key = key.replace("text_model.text_projection", "text_projection.weight")
         elif ".logit_scale" in key:
             key = None  # 後で処理する
         elif ".token_embedding" in key:
-            key = key.replace(".token_embedding.weight", ".embeddings.token_embedding.weight")
+            key = key.replace(
+                ".token_embedding.weight", ".embeddings.token_embedding.weight"
+            )
         elif ".ln_final" in key:
             key = key.replace(".ln_final", ".final_layer_norm")
         # ckpt from comfy has this key: text_model.encoder.text_model.embeddings.position_ids
@@ -115,7 +124,9 @@ def convert_key(key):
             values = torch.chunk(checkpoint[key], 3)
 
             key_suffix = ".weight" if "weight" in key else ".bias"
-            key_pfx = key.replace(SDXL_KEY_PREFIX + "transformer.resblocks.", "text_model.encoder.layers.")
+            key_pfx = key.replace(
+                SDXL_KEY_PREFIX + "transformer.resblocks.", "text_model.encoder.layers."
+            )
             key_pfx = key_pfx.replace("_weight", "")
             key_pfx = key_pfx.replace("_bias", "")
             key_pfx = key_pfx.replace(".attn.in_proj", ".self_attn.")
@@ -224,15 +235,19 @@ def load_models_from_sdxl_checkpoint(model_version, ckpt_path, map_location):
     te2_sd = {}
     for k in list(state_dict.keys()):
         if k.startswith("conditioner.embedders.0.transformer."):
-            te1_sd[k.replace("conditioner.embedders.0.transformer.", "")] = state_dict.pop(k)
+            te1_sd[
+                k.replace("conditioner.embedders.0.transformer.", "")
+            ] = state_dict.pop(k)
         elif k.startswith("conditioner.embedders.1.model."):
             te2_sd[k] = state_dict.pop(k)
 
-    info1 = text_model1.load_state_dict(te1_sd)
+    info1 = text_model1.load_state_dict(te1_sd, strict=False)
     print("text encoder 1:", info1)
 
-    converted_sd, logit_scale = convert_sdxl_text_encoder_2_checkpoint(te2_sd, max_length=77)
-    info2 = text_model2.load_state_dict(converted_sd)
+    converted_sd, logit_scale = convert_sdxl_text_encoder_2_checkpoint(
+        te2_sd, max_length=77
+    )
+    info2 = text_model2.load_state_dict(converted_sd, strict=False)
     print("text encoder 2:", info2)
 
     # prepare vae
@@ -241,7 +256,9 @@ def load_models_from_sdxl_checkpoint(model_version, ckpt_path, map_location):
     vae = AutoencoderKL(**vae_config)  # .to(device)
 
     print("loading VAE from checkpoint")
-    converted_vae_checkpoint = model_util.convert_ldm_vae_checkpoint(state_dict, vae_config)
+    converted_vae_checkpoint = model_util.convert_ldm_vae_checkpoint(
+        state_dict, vae_config
+    )
     info = vae.load_state_dict(converted_vae_checkpoint)
     print("VAE:", info)
 
@@ -264,7 +281,9 @@ def make_unet_conversion_map():
                 # no attention layers in down_blocks.3
                 hf_down_atn_prefix = f"down_blocks.{i}.attentions.{j}."
                 sd_down_atn_prefix = f"input_blocks.{3*i + j + 1}.1."
-                unet_conversion_map_layer.append((sd_down_atn_prefix, hf_down_atn_prefix))
+                unet_conversion_map_layer.append(
+                    (sd_down_atn_prefix, hf_down_atn_prefix)
+                )
 
         for j in range(3):
             # loop over resnets/attentions for upblocks
@@ -282,7 +301,9 @@ def make_unet_conversion_map():
             # no downsample in down_blocks.3
             hf_downsample_prefix = f"down_blocks.{i}.downsamplers.0.conv."
             sd_downsample_prefix = f"input_blocks.{3*(i+1)}.0.op."
-            unet_conversion_map_layer.append((sd_downsample_prefix, hf_downsample_prefix))
+            unet_conversion_map_layer.append(
+                (sd_downsample_prefix, hf_downsample_prefix)
+            )
 
             # no upsample in up_blocks.3
             hf_upsample_prefix = f"up_blocks.{i}.upsamplers.0."
@@ -389,9 +410,13 @@ def convert_key(key):
             else:
                 raise ValueError(f"unexpected key in DiffUsers model: {key}")
         elif ".position_embedding" in key:
-            key = key.replace("embeddings.position_embedding.weight", "positional_embedding")
+            key = key.replace(
+                "embeddings.position_embedding.weight", "positional_embedding"
+            )
         elif ".token_embedding" in key:
-            key = key.replace("embeddings.token_embedding.weight", "token_embedding.weight")
+            key = key.replace(
+                "embeddings.token_embedding.weight", "token_embedding.weight"
+            )
         elif "text_projection" in key:  # no dot in key
             key = key.replace("text_projection.weight", "text_projection")
         elif "final_layer_norm" in key:
@@ -419,7 +444,9 @@ def convert_key(key):
             value_v = checkpoint[key_v]
             value = torch.cat([value_q, value_k, value_v])
 
-            new_key = key.replace("text_model.encoder.layers.", "transformer.resblocks.")
+            new_key = key.replace(
+                "text_model.encoder.layers.", "transformer.resblocks."
+            )
             new_key = new_key.replace(".self_attn.q_proj.", ".attn.in_proj_")
             new_sd[new_key] = value
 
@@ -456,7 +483,9 @@ def update_sd(prefix, sd):
     # Convert the text encoders
     update_sd("conditioner.embedders.0.transformer.", text_encoder1.state_dict())
 
-    text_enc2_dict = convert_text_encoder_2_state_dict_to_sdxl(text_encoder2.state_dict(), logit_scale)
+    text_enc2_dict = convert_text_encoder_2_state_dict_to_sdxl(
+        text_encoder2.state_dict(), logit_scale
+    )
     update_sd("conditioner.embedders.1.model.", text_enc2_dict)
 
     # Convert the VAE
@@ -484,7 +513,14 @@ def update_sd(prefix, sd):
 
 
 def save_diffusers_checkpoint(
-    output_dir, text_encoder1, text_encoder2, unet, pretrained_model_name_or_path, vae=None, use_safetensors=False, save_dtype=None
+    output_dir,
+    text_encoder1,
+    text_encoder2,
+    unet,
+    pretrained_model_name_or_path,
+    vae=None,
+    use_safetensors=False,
+    save_dtype=None,
 ):
     from diffusers import StableDiffusionXLPipeline
 
@@ -501,11 +537,19 @@ def save_diffusers_checkpoint(
     if pretrained_model_name_or_path is None:
         pretrained_model_name_or_path = DIFFUSERS_REF_MODEL_ID_SDXL
 
-    scheduler = EulerDiscreteScheduler.from_pretrained(pretrained_model_name_or_path, subfolder="scheduler")
-    tokenizer1 = CLIPTokenizer.from_pretrained(pretrained_model_name_or_path, subfolder="tokenizer")
-    tokenizer2 = CLIPTokenizer.from_pretrained(pretrained_model_name_or_path, subfolder="tokenizer_2")
+    scheduler = EulerDiscreteScheduler.from_pretrained(
+        pretrained_model_name_or_path, subfolder="scheduler"
+    )
+    tokenizer1 = CLIPTokenizer.from_pretrained(
+        pretrained_model_name_or_path, subfolder="tokenizer"
+    )
+    tokenizer2 = CLIPTokenizer.from_pretrained(
+        pretrained_model_name_or_path, subfolder="tokenizer_2"
+    )
     if vae is None:
-        vae = AutoencoderKL.from_pretrained(pretrained_model_name_or_path, subfolder="vae")
+        vae = AutoencoderKL.from_pretrained(
+            pretrained_model_name_or_path, subfolder="vae"
+        )
 
     # prevent local path from being saved
     def remove_name_or_path(model):
diff --git a/lycoris/kohya/sdxl_original_unet.py b/lycoris/kohya/sdxl_original_unet.py
index 6ea4bc3..9beeb73 100644
--- a/lycoris/kohya/sdxl_original_unet.py
+++ b/lycoris/kohya/sdxl_original_unet.py
@@ -79,7 +79,9 @@ def forward(ctx, q, k, v, mask, causal, q_bucket_size, k_bucket_size):
 
         o = torch.zeros_like(q)
         all_row_sums = torch.zeros((*q.shape[:-1], 1), dtype=dtype, device=device)
-        all_row_maxes = torch.full((*q.shape[:-1], 1), max_neg_value, dtype=dtype, device=device)
+        all_row_maxes = torch.full(
+            (*q.shape[:-1], 1), max_neg_value, dtype=dtype, device=device
+        )
 
         scale = q.shape[-1] ** -0.5
 
@@ -108,15 +110,17 @@ def forward(ctx, q, k, v, mask, causal, q_bucket_size, k_bucket_size):
             for k_ind, (kc, vc) in enumerate(col_splits):
                 k_start_index = k_ind * k_bucket_size
 
-                attn_weights = torch.einsum("... i d, ... j d -> ... i j", qc, kc) * scale
+                attn_weights = (
+                    torch.einsum("... i d, ... j d -> ... i j", qc, kc) * scale
+                )
 
                 if exists(row_mask):
                     attn_weights.masked_fill_(~row_mask, max_neg_value)
 
                 if causal and q_start_index < (k_start_index + k_bucket_size - 1):
-                    causal_mask = torch.ones((qc.shape[-2], kc.shape[-2]), dtype=torch.bool, device=device).triu(
-                        q_start_index - k_start_index + 1
-                    )
+                    causal_mask = torch.ones(
+                        (qc.shape[-2], kc.shape[-2]), dtype=torch.bool, device=device
+                    ).triu(q_start_index - k_start_index + 1)
                     attn_weights.masked_fill_(causal_mask, max_neg_value)
 
                 block_row_maxes = attn_weights.amax(dim=-1, keepdims=True)
@@ -126,18 +130,27 @@ def forward(ctx, q, k, v, mask, causal, q_bucket_size, k_bucket_size):
                 if exists(row_mask):
                     exp_weights.masked_fill_(~row_mask, 0.0)
 
-                block_row_sums = exp_weights.sum(dim=-1, keepdims=True).clamp(min=EPSILON)
+                block_row_sums = exp_weights.sum(dim=-1, keepdims=True).clamp(
+                    min=EPSILON
+                )
 
                 new_row_maxes = torch.maximum(block_row_maxes, row_maxes)
 
-                exp_values = torch.einsum("... i j, ... j d -> ... i d", exp_weights, vc)
+                exp_values = torch.einsum(
+                    "... i j, ... j d -> ... i d", exp_weights, vc
+                )
 
                 exp_row_max_diff = torch.exp(row_maxes - new_row_maxes)
                 exp_block_row_max_diff = torch.exp(block_row_maxes - new_row_maxes)
 
-                new_row_sums = exp_row_max_diff * row_sums + exp_block_row_max_diff * block_row_sums
+                new_row_sums = (
+                    exp_row_max_diff * row_sums
+                    + exp_block_row_max_diff * block_row_sums
+                )
 
-                oc.mul_((row_sums / new_row_sums) * exp_row_max_diff).add_((exp_block_row_max_diff / new_row_sums) * exp_values)
+                oc.mul_((row_sums / new_row_sums) * exp_row_max_diff).add_(
+                    (exp_block_row_max_diff / new_row_sums) * exp_values
+                )
 
                 row_maxes.copy_(new_row_maxes)
                 row_sums.copy_(new_row_sums)
@@ -187,12 +200,14 @@ def backward(ctx, do):
             for k_ind, (kc, vc, dkc, dvc) in enumerate(col_splits):
                 k_start_index = k_ind * k_bucket_size
 
-                attn_weights = torch.einsum("... i d, ... j d -> ... i j", qc, kc) * scale
+                attn_weights = (
+                    torch.einsum("... i d, ... j d -> ... i j", qc, kc) * scale
+                )
 
                 if causal and q_start_index < (k_start_index + k_bucket_size - 1):
-                    causal_mask = torch.ones((qc.shape[-2], kc.shape[-2]), dtype=torch.bool, device=device).triu(
-                        q_start_index - k_start_index + 1
-                    )
+                    causal_mask = torch.ones(
+                        (qc.shape[-2], kc.shape[-2]), dtype=torch.bool, device=device
+                    ).triu(q_start_index - k_start_index + 1)
                     attn_weights.masked_fill_(causal_mask, max_neg_value)
 
                 exp_attn_weights = torch.exp(attn_weights - mc)
@@ -247,7 +262,9 @@ def get_timestep_embedding(
     assert len(timesteps.shape) == 1, "Timesteps should be a 1d-array"
 
     half_dim = embedding_dim // 2
-    exponent = -math.log(max_period) * torch.arange(start=0, end=half_dim, dtype=torch.float32, device=timesteps.device)
+    exponent = -math.log(max_period) * torch.arange(
+        start=0, end=half_dim, dtype=torch.float32, device=timesteps.device
+    )
     exponent = exponent / (half_dim - downscale_freq_shift)
 
     emb = torch.exp(exponent)
@@ -288,7 +305,9 @@ def __init__(
             nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1),
         )
 
-        self.emb_layers = nn.Sequential(nn.SiLU(), nn.Linear(TIME_EMBED_DIM, out_channels))
+        self.emb_layers = nn.Sequential(
+            nn.SiLU(), nn.Linear(TIME_EMBED_DIM, out_channels)
+        )
 
         self.out_layers = nn.Sequential(
             GroupNorm32(32, out_channels),
@@ -298,7 +317,9 @@ def __init__(
         )
 
         if in_channels != out_channels:
-            self.skip_connection = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
+            self.skip_connection = nn.Conv2d(
+                in_channels, out_channels, kernel_size=1, stride=1, padding=0
+            )
         else:
             self.skip_connection = nn.Identity()
 
@@ -322,7 +343,9 @@ def custom_forward(*inputs):
 
                 return custom_forward
 
-            x = torch.utils.checkpoint.checkpoint(create_custom_forward(self.forward_body), x, emb)
+            x = torch.utils.checkpoint.checkpoint(
+                create_custom_forward(self.forward_body), x, emb
+            )
         else:
             x = self.forward_body(x, emb)
 
@@ -356,7 +379,9 @@ def custom_forward(*inputs):
 
                 return custom_forward
 
-            hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(self.forward_body), hidden_states)
+            hidden_states = torch.utils.checkpoint.checkpoint(
+                create_custom_forward(self.forward_body), hidden_states
+            )
         else:
             hidden_states = self.forward_body(hidden_states)
 
@@ -374,7 +399,9 @@ def __init__(
     ):
         super().__init__()
         inner_dim = dim_head * heads
-        cross_attention_dim = cross_attention_dim if cross_attention_dim is not None else query_dim
+        cross_attention_dim = (
+            cross_attention_dim if cross_attention_dim is not None else query_dim
+        )
         self.upcast_attention = upcast_attention
 
         self.scale = dim_head**-0.5
@@ -403,14 +430,18 @@ def reshape_heads_to_batch_dim(self, tensor):
         batch_size, seq_len, dim = tensor.shape
         head_size = self.heads
         tensor = tensor.reshape(batch_size, seq_len, head_size, dim // head_size)
-        tensor = tensor.permute(0, 2, 1, 3).reshape(batch_size * head_size, seq_len, dim // head_size)
+        tensor = tensor.permute(0, 2, 1, 3).reshape(
+            batch_size * head_size, seq_len, dim // head_size
+        )
         return tensor
 
     def reshape_batch_dim_to_heads(self, tensor):
         batch_size, seq_len, dim = tensor.shape
         head_size = self.heads
         tensor = tensor.reshape(batch_size // head_size, head_size, seq_len, dim)
-        tensor = tensor.permute(0, 2, 1, 3).reshape(batch_size // head_size, seq_len, dim * head_size)
+        tensor = tensor.permute(0, 2, 1, 3).reshape(
+            batch_size // head_size, seq_len, dim * head_size
+        )
         return tensor
 
     def forward(self, hidden_states, context=None, mask=None):
@@ -443,7 +474,13 @@ def _attention(self, query, key, value):
             key = key.float()
 
         attention_scores = torch.baddbmm(
-            torch.empty(query.shape[0], query.shape[1], key.shape[1], dtype=query.dtype, device=query.device),
+            torch.empty(
+                query.shape[0],
+                query.shape[1],
+                key.shape[1],
+                dtype=query.dtype,
+                device=query.device,
+            ),
             query,
             key.transpose(-1, -2),
             beta=0,
@@ -472,13 +509,17 @@ def forward_memory_efficient_xformers(self, x, context=None, mask=None):
         k_in = self.to_k(context)
         v_in = self.to_v(context)
 
-        q, k, v = map(lambda t: rearrange(t, "b n (h d) -> b n h d", h=h), (q_in, k_in, v_in))
+        q, k, v = map(
+            lambda t: rearrange(t, "b n (h d) -> b n h d", h=h), (q_in, k_in, v_in)
+        )
         del q_in, k_in, v_in
 
         q = q.contiguous()
         k = k.contiguous()
         v = v.contiguous()
-        out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=None)  # 最適なのを選んでくれる
+        out = xformers.ops.memory_efficient_attention(
+            q, k, v, attn_bias=None
+        )  # 最適なのを選んでくれる
         del q, k, v
 
         out = rearrange(out, "b n h d -> b n (h d)", h=h)
@@ -517,10 +558,14 @@ def forward_sdpa(self, x, context=None, mask=None):
         k_in = self.to_k(context)
         v_in = self.to_v(context)
 
-        q, k, v = map(lambda t: rearrange(t, "b n (h d) -> b h n d", h=h), (q_in, k_in, v_in))
+        q, k, v = map(
+            lambda t: rearrange(t, "b n (h d) -> b h n d", h=h), (q_in, k_in, v_in)
+        )
         del q_in, k_in, v_in
 
-        out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False)
+        out = F.scaled_dot_product_attention(
+            q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False
+        )
 
         out = rearrange(out, "b h n d -> b n (h d)", h=h)
 
@@ -577,7 +622,12 @@ def forward(self, hidden_states):
 
 class BasicTransformerBlock(nn.Module):
     def __init__(
-        self, dim: int, num_attention_heads: int, attention_head_dim: int, cross_attention_dim: int, upcast_attention: bool = False
+        self,
+        dim: int,
+        num_attention_heads: int,
+        attention_head_dim: int,
+        cross_attention_dim: int,
+        upcast_attention: bool = False,
     ):
         super().__init__()
 
@@ -641,7 +691,12 @@ def custom_forward(*inputs):
 
                 return custom_forward
 
-            output = torch.utils.checkpoint.checkpoint(create_custom_forward(self.forward_body), hidden_states, context, timestep)
+            output = torch.utils.checkpoint.checkpoint(
+                create_custom_forward(self.forward_body),
+                hidden_states,
+                context,
+                timestep,
+            )
         else:
             output = self.forward_body(hidden_states, context, timestep)
 
@@ -666,13 +721,17 @@ def __init__(
         inner_dim = num_attention_heads * attention_head_dim
         self.use_linear_projection = use_linear_projection
 
-        self.norm = torch.nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
+        self.norm = torch.nn.GroupNorm(
+            num_groups=32, num_channels=in_channels, eps=1e-6, affine=True
+        )
         # self.norm = GroupNorm32(32, in_channels, eps=1e-6, affine=True)
 
         if use_linear_projection:
             self.proj_in = nn.Linear(in_channels, inner_dim)
         else:
-            self.proj_in = nn.Conv2d(in_channels, inner_dim, kernel_size=1, stride=1, padding=0)
+            self.proj_in = nn.Conv2d(
+                in_channels, inner_dim, kernel_size=1, stride=1, padding=0
+            )
 
         blocks = []
         for _ in range(num_transformer_layers):
@@ -691,7 +750,9 @@ def __init__(
         if use_linear_projection:
             self.proj_out = nn.Linear(in_channels, inner_dim)
         else:
-            self.proj_out = nn.Conv2d(inner_dim, in_channels, kernel_size=1, stride=1, padding=0)
+            self.proj_out = nn.Conv2d(
+                inner_dim, in_channels, kernel_size=1, stride=1, padding=0
+            )
 
         self.gradient_checkpointing = False
 
@@ -712,23 +773,37 @@ def forward(self, hidden_states, encoder_hidden_states=None, timestep=None):
         if not self.use_linear_projection:
             hidden_states = self.proj_in(hidden_states)
             inner_dim = hidden_states.shape[1]
-            hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(batch, height * weight, inner_dim)
+            hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(
+                batch, height * weight, inner_dim
+            )
         else:
             inner_dim = hidden_states.shape[1]
-            hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(batch, height * weight, inner_dim)
+            hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(
+                batch, height * weight, inner_dim
+            )
             hidden_states = self.proj_in(hidden_states)
 
         # 2. Blocks
         for block in self.transformer_blocks:
-            hidden_states = block(hidden_states, context=encoder_hidden_states, timestep=timestep)
+            hidden_states = block(
+                hidden_states, context=encoder_hidden_states, timestep=timestep
+            )
 
         # 3. Output
         if not self.use_linear_projection:
-            hidden_states = hidden_states.reshape(batch, height, weight, inner_dim).permute(0, 3, 1, 2).contiguous()
+            hidden_states = (
+                hidden_states.reshape(batch, height, weight, inner_dim)
+                .permute(0, 3, 1, 2)
+                .contiguous()
+            )
             hidden_states = self.proj_out(hidden_states)
         else:
             hidden_states = self.proj_out(hidden_states)
-            hidden_states = hidden_states.reshape(batch, height, weight, inner_dim).permute(0, 3, 1, 2).contiguous()
+            hidden_states = (
+                hidden_states.reshape(batch, height, weight, inner_dim)
+                .permute(0, 3, 1, 2)
+                .contiguous()
+            )
 
         output = hidden_states + residual
 
@@ -760,9 +835,13 @@ def forward_body(self, hidden_states, output_size=None):
 
         # if `output_size` is passed we force the interpolation output size and do not make use of `scale_factor=2`
         if output_size is None:
-            hidden_states = F.interpolate(hidden_states, scale_factor=2.0, mode="nearest")
+            hidden_states = F.interpolate(
+                hidden_states, scale_factor=2.0, mode="nearest"
+            )
         else:
-            hidden_states = F.interpolate(hidden_states, size=output_size, mode="nearest")
+            hidden_states = F.interpolate(
+                hidden_states, size=output_size, mode="nearest"
+            )
 
         # If the input is bfloat16, we cast back to bfloat16
         if dtype == torch.bfloat16:
@@ -782,7 +861,9 @@ def custom_forward(*inputs):
 
                 return custom_forward
 
-            hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(self.forward_body), hidden_states, output_size)
+            hidden_states = torch.utils.checkpoint.checkpoint(
+                create_custom_forward(self.forward_body), hidden_states, output_size
+            )
         else:
             hidden_states = self.forward_body(hidden_states, output_size)
 
@@ -827,7 +908,12 @@ def __init__(
         self.input_blocks = nn.ModuleList(
             [
                 nn.Sequential(
-                    nn.Conv2d(self.in_channels, self.model_channels, kernel_size=3, padding=(1, 1)),
+                    nn.Conv2d(
+                        self.in_channels,
+                        self.model_channels,
+                        kernel_size=3,
+                        padding=(1, 1),
+                    ),
                 )
             ]
         )
@@ -925,7 +1011,8 @@ def __init__(
         for i in range(3):
             layers = [
                 ResnetBlock2D(
-                    in_channels=4 * self.model_channels + (4 if i <= 1 else 2) * self.model_channels,
+                    in_channels=4 * self.model_channels
+                    + (4 if i <= 1 else 2) * self.model_channels,
                     out_channels=4 * self.model_channels,
                 ),
                 Transformer2DModel(
@@ -951,7 +1038,8 @@ def __init__(
         for i in range(3):
             layers = [
                 ResnetBlock2D(
-                    in_channels=2 * self.model_channels + (4 if i == 0 else (2 if i == 1 else 1)) * self.model_channels,
+                    in_channels=2 * self.model_channels
+                    + (4 if i == 0 else (2 if i == 1 else 1)) * self.model_channels,
                     out_channels=2 * self.model_channels,
                 ),
                 Transformer2DModel(
@@ -977,7 +1065,8 @@ def __init__(
         for i in range(3):
             layers = [
                 ResnetBlock2D(
-                    in_channels=1 * self.model_channels + (2 if i == 0 else 1) * self.model_channels,
+                    in_channels=1 * self.model_channels
+                    + (2 if i == 0 else 1) * self.model_channels,
                     out_channels=1 * self.model_channels,
                 ),
             ]
@@ -986,7 +1075,11 @@ def __init__(
 
         # output
         self.out = nn.ModuleList(
-            [GroupNorm32(32, self.model_channels), nn.SiLU(), nn.Conv2d(self.model_channels, self.out_channels, 3, padding=1)]
+            [
+                GroupNorm32(32, self.model_channels),
+                nn.SiLU(),
+                nn.Conv2d(self.model_channels, self.out_channels, 3, padding=1),
+            ]
         )
 
     # region diffusers compatibility
@@ -1007,7 +1100,10 @@ def set_attention_slice(self, slice_size):
         raise NotImplementedError("Attention slicing is not supported for this model.")
 
     def is_gradient_checkpointing(self) -> bool:
-        return any(hasattr(m, "gradient_checkpointing") and m.gradient_checkpointing for m in self.modules())
+        return any(
+            hasattr(m, "gradient_checkpointing") and m.gradient_checkpointing
+            for m in self.modules()
+        )
 
     def enable_gradient_checkpointing(self):
         self.gradient_checkpointing = True
@@ -1047,11 +1143,15 @@ def forward(self, x, timesteps=None, context=None, y=None, **kwargs):
         timesteps = timesteps.expand(x.shape[0])
 
         hs = []
-        t_emb = get_timestep_embedding(timesteps, self.model_channels)  # , repeat_only=False)
+        t_emb = get_timestep_embedding(
+            timesteps, self.model_channels
+        )  # , repeat_only=False)
         t_emb = t_emb.to(x.dtype)
         emb = self.time_embed(t_emb)
 
-        assert x.shape[0] == y.shape[0], f"batch size mismatch: {x.shape[0]} != {y.shape[0]}"
+        assert (
+            x.shape[0] == y.shape[0]
+        ), f"batch size mismatch: {x.shape[0]} != {y.shape[0]}"
         assert x.dtype == y.dtype, f"dtype mismatch: {x.dtype} != {y.dtype}"
         # assert x.dtype == self.dtype
         emb = emb + self.label_emb(y)
@@ -1109,7 +1209,9 @@ def call_module(module, h, emb, context):
 
     import transformers
 
-    optimizer = transformers.optimization.Adafactor(unet.parameters(), relative_step=True)  # working at 22.2GB with torch2
+    optimizer = transformers.optimization.Adafactor(
+        unet.parameters(), relative_step=True
+    )  # working at 22.2GB with torch2
 
     scaler = torch.cuda.amp.GradScaler(enabled=True)
 
diff --git a/lycoris/kohya/utils.py b/lycoris/kohya/utils.py
index 3b3f69d..892b93f 100644
--- a/lycoris/kohya/utils.py
+++ b/lycoris/kohya/utils.py
@@ -45,4 +45,4 @@ def precalculate_safetensors_hashes(tensors, metadata):
 
     model_hash = addnet_hash_safetensors(b)
     legacy_hash = addnet_hash_legacy(b)
-    return model_hash, legacy_hash
\ No newline at end of file
+    return model_hash, legacy_hash
diff --git a/lycoris/modules/__init__.py b/lycoris/modules/__init__.py
index 391a60d..65e4cdd 100644
--- a/lycoris/modules/__init__.py
+++ b/lycoris/modules/__init__.py
@@ -11,23 +11,24 @@
 @torch.no_grad()
 def make_module(lyco_type, params, lora_name, orig_module):
     module = None
-    if lyco_type == 'locon':
+    if lyco_type == "locon":
         up, down, mid, alpha = params
         module = LoConModule(
-            lora_name, orig_module, 1, 
-            down.size(0), torch.tensor(alpha).item(),
-            use_cp = mid is not None
+            lora_name,
+            orig_module,
+            1,
+            down.size(0),
+            torch.tensor(alpha).item(),
+            use_cp=mid is not None,
         )
         module.lora_up.weight.data.copy_(up)
         module.lora_down.weight.data.copy_(down)
         if mid is not None:
             module.lora_mid.weight.data.copy_(mid)
-    elif lyco_type == 'hada':
+    elif lyco_type == "hada":
         w1a, w1b, w2a, w2b, t1, t2, alpha = params
         module = LohaModule(
-            lora_name, orig_module, 1,
-            w1b.size(0), alpha.item(),
-            use_cp = t1 is not None
+            lora_name, orig_module, 1, w1b.size(0), alpha.item(), use_cp=t1 is not None
         )
         module.hada_w1_a.copy_(w1a)
         module.hada_w1_b.copy_(w1b)
@@ -36,7 +37,7 @@ def make_module(lyco_type, params, lora_name, orig_module):
         if t1 is not None:
             module.hada_t1.copy_(t1)
             module.hada_t2.copy_(t2)
-    elif lyco_type == 'kron':
+    elif lyco_type == "kron":
         w1, w1a, w1b, w2, w2a, w2b, _, t2, alpha = params
         if w1a is not None:
             lora_dim = w1a.size(1)
@@ -44,34 +45,39 @@ def make_module(lyco_type, params, lora_name, orig_module):
             lora_dim = w2a.size(1)
         else:
             lora_dim = 10000000000000
-        
+
         if w1 is None:
             out_dim = w1a.size(0)
             in_dim = w1b.size(1)
         else:
             out_dim, in_dim = w1.shape
-        
+
         shape_s = [out_dim, in_dim]
-        
+
         if w2 is None:
             out_dim *= w2a.size(0)
             in_dim *= w2b.size(1)
         else:
             out_dim *= w2.size(0)
             in_dim *= w2.size(1)
-        
-        if (shape_s[0] == factorization(out_dim, -1)[0] 
-            and shape_s[1] == factorization(in_dim, -1)[0]):
+
+        if (
+            shape_s[0] == factorization(out_dim, -1)[0]
+            and shape_s[1] == factorization(in_dim, -1)[0]
+        ):
             factor = -1
         else:
             factor = max(w1.shape) if w1 is not None else max(w1a.size(0), w1b.size(1))
-        
+
         module = LokrModule(
-            lora_name, orig_module, 1,
-            lora_dim, torch.tensor(alpha).item(), 
-            use_cp = t2 is not None,
-            decompose_both = w1 is None and w2 is None,
-            factor = factor
+            lora_name,
+            orig_module,
+            1,
+            lora_dim,
+            torch.tensor(alpha).item(),
+            use_cp=t2 is not None,
+            decompose_both=w1 is None and w2 is None,
+            factor=factor,
         )
         if w1 is not None:
             module.lokr_w1.copy_(w1)
@@ -85,24 +91,28 @@ def make_module(lyco_type, params, lora_name, orig_module):
             module.lokr_w2_b.copy_(w2b)
         if t2 is not None:
             module.lokr_t2.copy_(t2)
-    elif lyco_type == 'norm':
+    elif lyco_type == "norm":
         w_norm, b_norm = params
         module = NormModule(
-            lora_name, orig_module, 1,
+            lora_name,
+            orig_module,
+            1,
         )
         module.w_norm.copy_(w_norm)
         module.b_norm.copy_(b_norm)
-    elif lyco_type == 'full':
+    elif lyco_type == "full":
         diff, diff_b = params
         module = FullModule(
-            lora_name, orig_module, 1,
+            lora_name,
+            orig_module,
+            1,
         )
         module.diff.copy_(diff)
         if diff_b is not None:
             module.diff_b.copy_(diff_b)
-    elif lyco_type == 'ia3':
+    elif lyco_type == "ia3":
         pass
     else:
         return None
-    
-    return module
\ No newline at end of file
+
+    return module
diff --git a/lycoris/modules/attention.py b/lycoris/modules/attention.py
index 2ccef4a..c8f340d 100644
--- a/lycoris/modules/attention.py
+++ b/lycoris/modules/attention.py
@@ -7,17 +7,14 @@
 from einops import rearrange, repeat
 
 from lycoris.utils import default
-from lycoris.utils.xformers_utils import (
-    XFORMERS_AVAIL,
-    memory_efficient_attention
-)
+from lycoris.utils.xformers_utils import XFORMERS_AVAIL, memory_efficient_attention
 
 
 class GEGLU(nn.Module):
     def __init__(self, dim_in, dim_out):
         super().__init__()
         self.proj = nn.Linear(dim_in, dim_out * 2)
-    
+
     def forward(self, x):
         x, gate = self.proj(x).chunk(2, dim=-1)
         return x * F.gelu(gate)
@@ -28,18 +25,16 @@ def __init__(self, dim, dim_out=None, mult=4, glu=False):
         super().__init__()
         inner_dim = int(dim * mult)
         dim_out = default(dim_out, dim)
-        project_in = nn.Sequential(
-            nn.Linear(dim, inner_dim),
-            nn.GELU()
-        ) if not glu else GEGLU(dim, inner_dim)
-    
-        self.net = nn.Sequential(
-            project_in,
-            nn.Linear(inner_dim, dim_out)
+        project_in = (
+            nn.Sequential(nn.Linear(dim, inner_dim), nn.GELU())
+            if not glu
+            else GEGLU(dim, inner_dim)
         )
+
+        self.net = nn.Sequential(project_in, nn.Linear(inner_dim, dim_out))
         # nn.init.constant_(self.net[-1].weight, 0)
         # nn.init.constant_(self.net[-1].bias, 0)
-    
+
     def forward(self, x):
         return self.net(x)
 
@@ -49,43 +44,44 @@ def vanilla_attention(q, k, v, mask, scale=None):
         scale = math.sqrt(q.size(-1))
     scores = torch.bmm(q, k.transpose(-1, -2)) / scale
     if mask is not None:
-        mask = rearrange(mask, 'b ... -> b (...)')
+        mask = rearrange(mask, "b ... -> b (...)")
         max_neg_value = -torch.finfo(scores.dtype).max
-        mask = repeat(mask, 'b j -> (b h) j', h=q.size(-3))
+        mask = repeat(mask, "b j -> (b h) j", h=q.size(-3))
         scores = scores.masked_fill(~mask, max_neg_value)
     p_attn = F.softmax(scores, dim=-1)
     return torch.bmm(p_attn, v)
 
 
 MEMORY_LAYOUTS = {
-    'torch': (
-        'b n (h d) -> b h n d',
-        'b h n d -> b n (h d)',
+    "torch": (
+        "b n (h d) -> b h n d",
+        "b h n d -> b n (h d)",
         lambda x: (1, x, 1, 1),
     ),
-    'xformers': (
-        'b n (h d) -> b n h d',
-        'b n h d -> b n (h d)',
+    "xformers": (
+        "b n (h d) -> b n h d",
+        "b n h d -> b n (h d)",
         lambda x: (1, 1, x, 1),
     ),
-    'vanilla': (
-        'b n (h d) -> b h n d',
-        'b h n d -> b n (h d)',
+    "vanilla": (
+        "b n (h d) -> b h n d",
+        "b h n d -> b n (h d)",
         lambda x: (1, x, 1, 1),
-    )
+    ),
 }
 ATTN_FUNCTION = {
-    'vanilla': vanilla_attention,
-    'torch': getattr(F, 'scaled_dot_product_attention', None),
-    'xformers': memory_efficient_attention
+    "vanilla": vanilla_attention,
+    "torch": getattr(F, "scaled_dot_product_attention", None),
+    "xformers": memory_efficient_attention,
 }
 
 
 class Attention(nn.Module):
-    '''
+    """
     Attention Class without norm and residual
     (You need to wrap them by your self)
-    '''
+    """
+
     def __init__(
         self,
         in_ch,
@@ -95,18 +91,18 @@ def __init__(
         self_cross=False,
         double_attn=False,
         single_kv_head=False,
-        attn_backend='torch',
+        attn_backend="torch",
         cosine_attn=False,
-        qk_head_ch=-1
+        qk_head_ch=-1,
     ):
         super().__init__()
-        if heads==-1:
-            assert in_ch%head_ch==0
-            heads = in_ch//head_ch
-        if head_ch==-1:
-            assert in_ch%heads==0
-            head_ch = in_ch//heads
-        if qk_head_ch==-1:
+        if heads == -1:
+            assert in_ch % head_ch == 0
+            heads = in_ch // head_ch
+        if head_ch == -1:
+            assert in_ch % heads == 0
+            head_ch = in_ch // heads
+        if qk_head_ch == -1:
             qk_head_ch = head_ch
         q_ch = heads * qk_head_ch
         k_ch = (1 if single_kv_head else heads) * qk_head_ch
@@ -115,12 +111,12 @@ def __init__(
         assert inner_ch == in_ch
         use_context = context_ch is not None
         context_ch = default(context_ch, in_ch)
-        
-        if attn_backend == 'xformers':
+
+        if attn_backend == "xformers":
             assert XFORMERS_AVAIL
-        if attn_backend == 'torch':
-            assert torch.version.__version__ >= '2.0.0'
-        
+        if attn_backend == "torch":
+            assert torch.version.__version__ >= "2.0.0"
+
         self.heads = heads
         self.self_cross = self_cross
         self.double_attn = double_attn
@@ -128,12 +124,14 @@ def __init__(
         self.attn = ATTN_FUNCTION[attn_backend]
         self.memory_layout = MEMORY_LAYOUTS[attn_backend]
         self.cosine_attn = cosine_attn
-        
+
         if cosine_attn:
-            self.scale = nn.Parameter(torch.ones(MEMORY_LAYOUTS[attn_backend][2](heads)))
+            self.scale = nn.Parameter(
+                torch.ones(MEMORY_LAYOUTS[attn_backend][2](heads))
+            )
         else:
             self.scale = None
-        
+
         self.q = nn.Linear(in_ch, q_ch, bias=False)
         if self_cross and use_context:
             self.k = nn.Linear(in_ch, k_ch, bias=False)
@@ -144,13 +142,13 @@ def __init__(
             assert double_attn == False
             self.k = nn.Linear(context_ch, k_ch, bias=False)
             self.v = nn.Linear(context_ch, v_ch, bias=False)
-        
+
         if double_attn:
-            self.out = nn.Linear(inner_ch*2, in_ch)
+            self.out = nn.Linear(inner_ch * 2, in_ch)
         else:
             self.out = nn.Linear(inner_ch, in_ch)
-    
-    def forward(self, x:torch.Tensor, context=None, mask=None):
+
+    def forward(self, x: torch.Tensor, context=None, mask=None):
         # Input Projection
         heads = self.heads
         q = self.q(x)
@@ -168,38 +166,34 @@ def forward(self, x:torch.Tensor, context=None, mask=None):
             ctx = default(context, x)
             k = self.k(ctx)
             v = self.v(ctx)
-        
+
         # Rearrange for Attention
         q = rearrange(q, self.memory_layout[0], h=heads)
         if self.single_kv_head:
             k = k.unsqueeze(1)
             v = v.unsqueeze(1)
-            
+
             b, _, seq, _ = k.shape
             k = k.expand(b, heads, seq, k.size(3))
             v = v.expand(b, heads, seq, v.size(3))
         else:
             k = rearrange(k, self.memory_layout[0], h=heads)
             v = rearrange(v, self.memory_layout[0], h=heads)
-        
+
         if self.cosine_attn:
             q = (F.normalize(q, dim=-1) * math.sqrt(q.size(-1))).to(v.dtype)
             k = (F.normalize(k, dim=-1) * self.scale).to(v.dtype)
             if ck is not None and self.double_attn:
                 ck = (F.normalize(ck, dim=-1) * self.scale).to(v.dtype)
-        
+
         # Attention
-        out = self.attn(
-            q.contiguous(), k.contiguous(), v.contiguous(), mask
-        )
+        out = self.attn(q.contiguous(), k.contiguous(), v.contiguous(), mask)
         out = rearrange(out, self.memory_layout[1], h=heads)
         if self.double_attn:
-            out2 = self.attn(
-                q.contiguous(), ck.contiguous(), cv.contiguous(), mask
-            )
+            out2 = self.attn(q.contiguous(), ck.contiguous(), cv.contiguous(), mask)
             out2 = rearrange(out2, self.memory_layout[1], h=heads)
             out = torch.cat([out, out2], dim=-1)
-        
+
         # Output Projection
         return self.out(out)
 
@@ -214,7 +208,7 @@ def __init__(
         gated_ff=True,
         self_cross=False,
         single_kv_head=False,
-        attn_backend='torch',
+        attn_backend="torch",
         cosine_attn=False,
         qk_head_ch=-1,
         disable_self_attn=False,
@@ -223,24 +217,43 @@ def __init__(
         super().__init__()
         self.single_attn = single_attn
         self.disable_self_attn = disable_self_attn or single_attn
-        
+
         self.norm1 = nn.LayerNorm(dim)
         self.attn1 = Attention(
-            dim, context_dim if self.disable_self_attn else None, n_heads, d_head,
-            self_cross, single_kv_head, attn_backend, cosine_attn, qk_head_ch
+            dim,
+            context_dim if self.disable_self_attn else None,
+            n_heads,
+            d_head,
+            self_cross,
+            single_kv_head,
+            attn_backend,
+            cosine_attn,
+            qk_head_ch,
         )
         if not single_attn:
             self.norm2 = nn.LayerNorm(dim)
             self.attn2 = Attention(
-                dim, context_dim, n_heads, d_head,
-                self_cross, single_kv_head, attn_backend, cosine_attn, qk_head_ch
+                dim,
+                context_dim,
+                n_heads,
+                d_head,
+                self_cross,
+                single_kv_head,
+                attn_backend,
+                cosine_attn,
+                qk_head_ch,
             )
         self.norm3 = nn.LayerNorm(dim)
         self.ff = FeedForward(dim, glu=gated_ff)
-    
+
     def forward(self, x, context=None):
-        x = self.attn1(self.norm1(x), context=context if self.disable_self_attn else None) + x
+        x = (
+            self.attn1(
+                self.norm1(x), context=context if self.disable_self_attn else None
+            )
+            + x
+        )
         if not self.single_attn and (context is not None or self.attn2.self_cross):
             x = self.attn2(self.norm2(x), context=context) + x
         x = self.ff(self.norm3(x)) + x
-        return x
\ No newline at end of file
+        return x
diff --git a/lycoris/modules/base.py b/lycoris/modules/base.py
new file mode 100644
index 0000000..465bcab
--- /dev/null
+++ b/lycoris/modules/base.py
@@ -0,0 +1,37 @@
+from collections import OrderedDict
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class ModuleCustomSD(nn.Module):
+    def custom_state_dict(self):
+        return None
+
+    def state_dict(self, *args, destination=None, prefix="", keep_vars=False):
+        # TODO: Remove `args` and the parsing logic when BC allows.
+        if len(args) > 0:
+            if destination is None:
+                destination = args[0]
+            if len(args) > 1 and prefix == "":
+                prefix = args[1]
+            if len(args) > 2 and keep_vars is False:
+                keep_vars = args[2]
+            # DeprecationWarning is ignored by default
+
+        if destination is None:
+            destination = OrderedDict()
+            destination._metadata = OrderedDict()
+
+        local_metadata = dict(version=self._version)
+        if hasattr(destination, "_metadata"):
+            destination._metadata[prefix[:-1]] = local_metadata
+
+        if (custom_sd := self.custom_state_dict()) is not None:
+            for k, v in custom_sd.items():
+                destination[f"{prefix}{k}"] = v
+            return destination
+        else:
+            return super().state_dict(
+                *args, destination=destination, prefix=prefix, keep_vars=keep_vars
+            )
diff --git a/lycoris/modules/diag_oft.py b/lycoris/modules/diag_oft.py
new file mode 100644
index 0000000..87ae1ee
--- /dev/null
+++ b/lycoris/modules/diag_oft.py
@@ -0,0 +1,146 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from einops import rearrange
+
+from .base import ModuleCustomSD
+from .lokr import factorization
+
+
+class DiagOFTModule(ModuleCustomSD):
+    """
+    modifed from kohya-ss/sd-scripts/networks/lora:LoRAModule
+    """
+
+    def __init__(
+        self,
+        lora_name,
+        org_module: nn.Module,
+        multiplier=1.0,
+        lora_dim=4,
+        alpha=1,
+        dropout=0.0,
+        rank_dropout=0.0,
+        module_dropout=0.0,
+        use_tucker=False,
+        use_scalar=False,
+        rank_dropout_scale=False,
+        constrain=0,
+        rescaled=False,
+        **kwargs,
+    ):
+        """if alpha == 0 or None, alpha is rank (no scaling)."""
+        super().__init__()
+        self.lora_name = lora_name
+
+        if isinstance(org_module, nn.Linear):
+            self.op = F.linear
+            self.dim = org_module.out_features
+            self.kw_dict = {}
+        elif isinstance(org_module, nn.Conv2d):
+            self.op = F.conv2d
+            self.dim = org_module.out_channels
+            self.kw_dict = {
+                "stride": org_module.stride,
+                "padding": org_module.padding,
+                "dilation": org_module.dilation,
+                "groups": org_module.groups,
+            }
+        else:
+            raise NotImplementedError
+
+        out_dim = org_module.weight.shape[0]
+        self.block_size, self.block_num = factorization(out_dim, lora_dim)
+        # block_num > block_size
+        self.rescaled = rescaled
+        self.constrain = constrain * out_dim
+        self.oft_blocks = nn.Parameter(
+            torch.zeros(self.block_num, self.block_size, self.block_size)
+        )
+        if rescaled:
+            self.rescale = nn.Parameter(torch.ones(self.block_num, self.block_size, 1))
+
+        self.rank_dropout = rank_dropout
+        self.rank_dropout_scale = rank_dropout_scale
+        self.module_dropout = module_dropout
+
+        self.multiplier = multiplier
+        self.org_module = [org_module]
+
+    @property
+    def I(self):
+        return torch.eye(self.block_size, device=next(self.parameters()).device)
+
+    def apply_to(self, **kwargs):
+        self.org_forward = self.org_module[0].forward
+        self.org_module[0].forward = self.forward
+
+    def custom_state_dict(self):
+        return {"oft_diag": self.get_r()}
+
+    def get_r(self):
+        I = self.I
+        # for Q = -Q^T
+        q = self.oft_blocks - self.oft_blocks.transpose(1, 2)
+        normed_q = q
+        if self.constrain > 0:
+            q_norm = torch.norm(q) + 1e-8
+            if q_norm > self.constrain:
+                normed_q = q * self.constrain / q_norm
+        # use float() to prevent unsupported type
+        r = (I + normed_q) @ (I - normed_q).float().inverse()
+
+        if self.rescaled:
+            # Noted: not implemented in Kohya
+            r = self.rescale * r
+        return r
+
+    def make_weight(self, scale=1, device=None):
+        if self.rank_dropout and self.training:
+            drop = torch.rand(self.dim, device=device) < self.rank_dropout
+            drop = drop.to(self.oft_blocks.dtype)
+            if self.rank_dropout_scale:
+                drop /= drop.mean()
+        else:
+            drop = 1
+
+        r = self.get_r()
+        org_weight = self.org_module[0].weight.to(device, dtype=r.dtype)
+        org_weight = rearrange(
+            org_weight,
+            "(k n) ... -> k n ...",
+            k=self.block_num,
+            n=self.block_size,
+        )
+        # Init R=0, so add I on it to ensure the output of step0 is original model output
+        weight = torch.einsum(
+            "k n m, k n ... -> k m ...",
+            r * scale + (1 - scale) * self.I,
+            org_weight,
+        )
+        weight = rearrange(weight, "k m ... -> (k m) ...")
+        return weight
+
+    @torch.no_grad()
+    def apply_max_norm(self, max_norm, device=None):
+        orig_norm = self.oft_diag.to(device).norm()
+        norm = torch.clamp(orig_norm, max_norm / 2)
+        desired = torch.clamp(norm, max=max_norm)
+        ratio = desired / norm
+
+        scaled = ratio.item() != 1.0
+        if scaled:
+            self.oft_blocks *= ratio
+
+        return scaled, orig_norm * ratio
+
+    def forward(self, x):
+        if self.module_dropout and self.training:
+            if torch.rand(1) < self.module_dropout:
+                return self.org_forward(x)
+        scale = self.multiplier
+
+        w = self.make_weight(scale, x.device)
+        kw_dict = self.kw_dict | {"weight": w, "bias": self.org_module[0].bias}
+        return self.op(x, **kw_dict)
diff --git a/lycoris/modules/dylora.py b/lycoris/modules/dylora.py
index e5b6d7d..d0922dc 100644
--- a/lycoris/modules/dylora.py
+++ b/lycoris/modules/dylora.py
@@ -1,49 +1,52 @@
 import math
 import random
 
-from collections import OrderedDict, abc as container_abcs
-
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
+from .base import ModuleCustomSD
+
 
-class DyLoraModule(nn.Module):
+class DyLoraModule(ModuleCustomSD):
     """
     Hadamard product Implementaion for Dynamic Low Rank adaptation
     """
 
     def __init__(
-        self, 
-        lora_name, 
-        org_module: nn.Module, 
-        multiplier=1.0, 
-        lora_dim=4, alpha=1, 
-        dropout=0., rank_dropout=0., module_dropout=0.,
-        use_tucker=False, 
+        self,
+        lora_name,
+        org_module: nn.Module,
+        multiplier=1.0,
+        lora_dim=4,
+        alpha=1,
+        dropout=0.0,
+        rank_dropout=0.0,
+        module_dropout=0.0,
+        use_tucker=False,
         block_size=1,
         **kwargs,
     ):
-        """ if alpha == 0 or None, alpha is rank (no scaling). """
+        """if alpha == 0 or None, alpha is rank (no scaling)."""
         super().__init__()
         self.lora_name = lora_name
         self.lora_dim = lora_dim
-        assert lora_dim % block_size == 0, 'lora_dim must be a multiple of block_size'
-        self.block_count = lora_dim//block_size
+        assert lora_dim % block_size == 0, "lora_dim must be a multiple of block_size"
+        self.block_count = lora_dim // block_size
         self.block_size = block_size
-        
+
         self.shape = org_module.weight.shape
-        if org_module.__class__.__name__ == 'Conv2d':
+        if org_module.__class__.__name__ == "Conv2d":
             in_dim = org_module.in_channels
             k_size = org_module.kernel_size
             out_dim = org_module.out_channels
-            shape = (out_dim, in_dim*k_size[0]*k_size[1])
+            shape = (out_dim, in_dim * k_size[0] * k_size[1])
             self.op = F.conv2d
             self.extra_args = {
                 "stride": org_module.stride,
                 "padding": org_module.padding,
                 "dilation": org_module.dilation,
-                "groups": org_module.groups
+                "groups": org_module.groups,
             }
         else:
             in_dim = org_module.in_features
@@ -51,62 +54,43 @@ def __init__(
             shape = (out_dim, in_dim)
             self.op = F.linear
             self.extra_args = {}
-        
+
         self.lora_dim = lora_dim
-        self.up_list = nn.ParameterList([
-            torch.empty(shape[0], 1)
-            for i in range(lora_dim)
-        ])
-        
-        self.down_list = nn.ParameterList([
-            torch.empty(1, shape[1])
-            for i in range(lora_dim)
-        ])
-        
+        self.up_list = nn.ParameterList(
+            [torch.empty(shape[0], 1) for i in range(lora_dim)]
+        )
+
+        self.down_list = nn.ParameterList(
+            [torch.empty(1, shape[1]) for i in range(lora_dim)]
+        )
+
         self.index = 0
-        
+
         if type(alpha) == torch.Tensor:
             alpha = alpha.detach().float().numpy()  # without casting, bf16 causes error
         alpha = lora_dim if alpha is None or alpha == 0 else alpha
         self.scale = alpha / self.lora_dim
-        self.register_buffer('alpha', torch.tensor(alpha)) # 定数として扱える
+        self.register_buffer("alpha", torch.tensor(alpha))  # 定数として扱える
 
         # Need more experiences on init method
-        
+
         for v in self.down_list:
             torch.nn.init.kaiming_uniform_(v, a=math.sqrt(5))
         for v in self.up_list:
             torch.nn.init.zeros_(v)
 
         self.multiplier = multiplier
-        self.org_module = [org_module] # remove in applying
+        self.org_module = [org_module]  # remove in applying
         self.grad_ckpt = False
         self.state_dict()
-    
-    def state_dict(self, *args, destination=None, prefix='', keep_vars=False):
-        # TODO: Remove `args` and the parsing logic when BC allows.
-        if len(args) > 0:
-            if destination is None:
-                destination = args[0]
-            if len(args) > 1 and prefix == '':
-                prefix = args[1]
-            if len(args) > 2 and keep_vars is False:
-                keep_vars = args[2]
-            # DeprecationWarning is ignored by default
-
-        if destination is None:
-            destination = OrderedDict()
-            destination._metadata = OrderedDict()
-
-        local_metadata = dict(version=self._version)
-        if hasattr(destination, "_metadata"):
-            destination._metadata[prefix[:-1]] = local_metadata
-
-        destination[f'{prefix}alpha'] = self.alpha
-        destination[f'{prefix}lora_up.weight'] = nn.Parameter(
+
+    def custom_state_dict(self):
+        destination = {}
+        destination["alpha"] = self.alpha
+        destination["lora_up.weight"] = nn.Parameter(
             torch.concat(list(self.up_list), dim=1)
         )
-        destination[f'{prefix}lora_down.weight'] = nn.Parameter(
+        destination["lora_down.weight"] = nn.Parameter(
             torch.concat(list(self.down_list))
         )
         return destination
@@ -116,22 +100,23 @@ def apply_to(self):
 
     @torch.enable_grad()
     def forward(self, x):
-        b = random.randint(0, self.block_count-1)
-        
+        b = random.randint(0, self.block_count - 1)
+
         down = torch.concat(
-            list(i.data for i in self.down_list[:b*self.block_size]) 
-            + list(self.down_list[b*self.block_size:(b+1)*self.block_size])
+            list(i.data for i in self.down_list[: b * self.block_size])
+            + list(self.down_list[b * self.block_size : (b + 1) * self.block_size])
         )
         up = torch.concat(
-            list(i.data for i in self.up_list[:b*self.block_size]) 
-            + list(self.up_list[b*self.block_size:(b+1)*self.block_size]),
-            dim=1
+            list(i.data for i in self.up_list[: b * self.block_size])
+            + list(self.up_list[b * self.block_size : (b + 1) * self.block_size]),
+            dim=1,
         )
-        
+
         bias = None if self.org_module[0].bias is None else self.org_module[0].bias.data
         return self.op(
-            x, 
-            self.org_module[0].weight + (up@down).view(self.shape) * self.alpha/(b+1),
+            x,
+            self.org_module[0].weight
+            + (up @ down).view(self.shape) * self.alpha / (b + 1),
             bias,
-            **self.extra_args
-        )
\ No newline at end of file
+            **self.extra_args,
+        )
diff --git a/lycoris/modules/full.py b/lycoris/modules/full.py
index a381e50..1e48f40 100644
--- a/lycoris/modules/full.py
+++ b/lycoris/modules/full.py
@@ -1,27 +1,26 @@
-import math
-from weakref import ref
-from collections import OrderedDict, abc as container_abcs
-
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
+from .base import ModuleCustomSD
 
-class FullModule(nn.Module):
-    """
-    modifed from kohya-ss/sd-scripts/networks/lora:LoRAModule
-    """
 
+class FullModule(ModuleCustomSD):
     def __init__(
-        self, 
-        lora_name, org_module: nn.Module, 
-        multiplier=1.0, 
-        lora_dim=4, alpha=1, 
-        dropout=0., rank_dropout=0., module_dropout=0.,
-        use_tucker=False, use_scalar=False,
+        self,
+        lora_name,
+        org_module: nn.Module,
+        multiplier=1.0,
+        lora_dim=4,
+        alpha=1,
+        dropout=0.0,
+        rank_dropout=0.0,
+        module_dropout=0.0,
+        use_tucker=False,
+        use_scalar=False,
+        rank_dropout_scale=False,
         **kwargs,
     ):
-        """ if alpha == 0 or None, alpha is rank (no scaling). """
         super().__init__()
         self.lora_name = lora_name
 
@@ -40,60 +39,54 @@ def __init__(
             }
         else:
             raise NotImplementedError
-        self.diff = nn.Parameter(torch.zeros_like(org_module.weight))
+        self.weight = nn.Parameter(torch.zeros_like(org_module.weight))
         if org_module.bias is not None:
-            self.diff_b = nn.Parameter(torch.zeros_like(org_module.bias))
+            self.bias = nn.Parameter(torch.zeros_like(org_module.bias))
         else:
-            self.diff_b = None
-        
+            self.bias = None
+
         self.rank_dropout = rank_dropout
         self.module_dropout = module_dropout
-        
+
         self.multiplier = multiplier
         self.org_module = [org_module]
 
     def apply_to(self, **kwargs):
         self.org_forward = self.org_module[0].forward
-        self.org_module[0].forward = self.forward
+        self.weight.data.copy_(self.org_module[0].weight.data)
+        if self.org_module[0].bias is not None:
+            self.bias.data.copy_(self.org_module[0].bias.data)
 
-    def make_weight(self, scale = 1, device=None):
+    def custom_state_dict(self):
+        sd = {"diff": self.weight.data.cpu() - self.org_module[0].weight.data.cpu()}
+        if self.bias is not None:
+            sd["diff_b"] = self.bias.data.cpu() - self.org_module[0].bias.data.cpu()
+        return sd
+
+    def make_weight(self, scale=1, device=None):
         drop = (
-            torch.rand(self.dim, device=device) < self.rank_dropout 
-            if self.rank_dropout and self.training 
+            torch.rand(self.dim, device=device) > self.rank_dropout
+            if self.rank_dropout and self.training
             else 1
         )
-        org_weight = self.org_module[0].weight.to(device, dtype=self.diff.dtype)
-        weight = self.diff.to(device) * drop * scale
-        weight = weight + org_weight
-        
-        if self.diff_b is not None:
-            org_bias = self.org_module[0].bias.to(device, dtype=self.diff_b.dtype)
-            bias = self.diff_b.to(device) * drop * scale
-            bias = bias + org_bias
+        if drop != 1 or scale != 1:
+            org_weight = self.org_module[0].weight.to(device, dtype=self.weight.dtype)
+            diff = self.weight.to(device) - org_weight
+            weight = diff * drop * scale + org_weight
+            if self.bias is not None:
+                org_bias = self.org_module[0].bias.to(device, dtype=self.bias.dtype)
+                diff_b = self.bias.to(device) - org_bias
+                bias = diff_b * drop * scale + org_bias
         else:
-            bias = None
+            weight = self.weight
+            bias = self.bias
         return weight, bias
 
-    @torch.no_grad()
-    def apply_max_norm(self, max_norm, device=None):
-        orig_norm = self.diff.to(device).norm()
-        norm = torch.clamp(orig_norm, max_norm/2)
-        desired = torch.clamp(norm, max=max_norm)
-        ratio = desired/norm
-        
-        scaled = ratio.item() != 1.0
-        if scaled:
-            self.diff *= ratio
-            self.diff_b *= ratio
-        
-        return scaled, orig_norm*ratio
-
     def forward(self, x):
         if self.module_dropout and self.training:
             if torch.rand(1) < self.module_dropout:
                 return self.org_forward(x)
         scale = self.multiplier
-        
-        w, b = self.make_weight(scale, x.device)
-        kw_dict = self.kw_dict | {"weight": w, "bias": b}
+        weight, bias = self.make_weight(scale)
+        kw_dict = self.kw_dict | {"weight": weight, "bias": bias}
         return self.op(x, **kw_dict)
diff --git a/lycoris/modules/glokr.py b/lycoris/modules/glokr.py
index c95da7a..bbcf448 100644
--- a/lycoris/modules/glokr.py
+++ b/lycoris/modules/glokr.py
@@ -1,20 +1,20 @@
-import math
-
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
+from .base import ModuleCustomSD
+
 
-def factorization(dimension: int, factor:int=-1) -> tuple[int, int]:
-    '''
+def factorization(dimension: int, factor: int = -1) -> tuple[int, int]:
+    """
     return a tuple of two value of input dimension decomposed by the number closest to factor
     second value is higher or equal than first value.
-    
+
     In LoRA with Kroneckor Product, first value is a value for weight scale.
     secon value is a value for weight.
-    
+
     Becuase of non-commutative property, A⊗B ≠ B⊗A. Meaning of two matrices is slightly different.
-    
+
     examples)
     factor
         -1               2                4               8               16               ...
@@ -24,8 +24,8 @@ def factorization(dimension: int, factor:int=-1) -> tuple[int, int]:
     360 -> 45, 8    360 -> 180, 2    360 -> 90, 4    360 -> 45, 8    360 -> 45, 8
     512 -> 32, 16   512 -> 256, 2    512 -> 128, 4   512 -> 64, 8    512 -> 32, 16
     1024 -> 32, 32  1024 -> 512, 2   1024 -> 256, 4  1024 -> 128, 8  1024 -> 64, 16
-    '''
-    
+    """
+
     if factor > 0 and (dimension % factor) == 0:
         m = factor
         n = dimension // factor
@@ -34,12 +34,12 @@ def factorization(dimension: int, factor:int=-1) -> tuple[int, int]:
         factor = dimension
     m, n = 1, dimension
     length = m + n
-    while m<n:
+    while m < n:
         new_m = m + 1
-        while dimension%new_m != 0:
+        while dimension % new_m != 0:
             new_m += 1
         new_n = dimension // new_m
-        if new_m + new_n > length or new_m>factor:
+        if new_m + new_n > length or new_m > factor:
             break
         else:
             m, n = new_m, new_n
@@ -49,7 +49,7 @@ def factorization(dimension: int, factor:int=-1) -> tuple[int, int]:
 
 
 def make_weight_cp(t, wa, wb):
-    rebuild2 = torch.einsum('i j k l, i p, j r -> p r k l', t, wa, wb) # [c, d, k1, k2]
+    rebuild2 = torch.einsum("i j k l, i p, j r -> p r k l", t, wa, wb)  # [c, d, k1, k2]
     return rebuild2
 
 
@@ -58,11 +58,11 @@ def make_kron(w1, w2, scale):
         w1 = w1.unsqueeze(2).unsqueeze(2)
     w2 = w2.contiguous()
     rebuild = torch.kron(w1, w2)
-    
-    return rebuild*scale
+
+    return rebuild * scale
 
 
-class LokrModule(nn.Module):
+class LokrModule(ModuleCustomSD):
     """
     modifed from kohya-ss/sd-scripts/networks/lora:LoRAModule
         and from KohakuBlueleaf/LyCORIS/lycoris:loha:LoHaModule
@@ -70,17 +70,22 @@ class LokrModule(nn.Module):
     """
 
     def __init__(
-        self, 
-        lora_name, org_module: nn.Module, 
-        multiplier=1.0, 
-        lora_dim=4, alpha=1, 
-        dropout=0., rank_dropout=0., module_dropout=0.,
+        self,
+        lora_name,
+        org_module: nn.Module,
+        multiplier=1.0,
+        lora_dim=4,
+        alpha=1,
+        dropout=0.0,
+        rank_dropout=0.0,
+        module_dropout=0.0,
         use_tucker=False,
-        decompose_both = False,
-        factor:int=-1, # factorization factor
+        decompose_both=False,
+        factor: int = -1,  # factorization factor
+        rank_dropout_scale=False,
         **kwargs,
     ):
-        """ if alpha == 0 or None, alpha is rank (no scaling). """
+        """if alpha == 0 or None, alpha is rank (no scaling)."""
         super().__init__()
         factor = int(factor)
         self.lora_name = lora_name
@@ -90,61 +95,78 @@ def __init__(
         self.use_w2 = False
 
         self.shape = org_module.weight.shape
-        if org_module.__class__.__name__ == 'Conv2d':
+        if org_module.__class__.__name__ == "Conv2d":
             in_dim = org_module.in_channels
             k_size = org_module.kernel_size
             out_dim = org_module.out_channels
-            
+
             in_m, in_n = factorization(in_dim, factor)
             out_l, out_k = factorization(out_dim, factor)
-            shape = ((out_l, out_k), (in_m, in_n), *k_size) # ((a, b), (c, d), *k_size)
-            
-            self.tucker = use_tucker and k_size!=(1, 1)
-            if decompose_both and lora_dim < max(shape[0][0], shape[1][0])/2:
+            shape = ((out_l, out_k), (in_m, in_n), *k_size)  # ((a, b), (c, d), *k_size)
+
+            self.tucker = use_tucker and k_size != (1, 1)
+            if decompose_both and lora_dim < max(shape[0][0], shape[1][0]) / 2:
                 self.lokr_w1_a = nn.Parameter(torch.empty(shape[0][0], lora_dim))
                 self.lokr_w1_b = nn.Parameter(torch.empty(lora_dim, shape[1][0]))
             else:
                 self.use_w1 = True
-                self.lokr_w1 = nn.Parameter(torch.empty(shape[0][0], shape[1][0]))  # a*c, 1-mode
-            
-            if lora_dim >= max(shape[0][1], shape[1][1])/2:
+                self.lokr_w1 = nn.Parameter(
+                    torch.empty(shape[0][0], shape[1][0])
+                )  # a*c, 1-mode
+
+            if lora_dim >= max(shape[0][1], shape[1][1]) / 2:
                 self.use_w2 = True
-                self.lokr_w2 = nn.Parameter(torch.empty(shape[0][1], shape[1][1], *k_size))
+                self.lokr_w2 = nn.Parameter(
+                    torch.empty(shape[0][1], shape[1][1], *k_size)
+                )
             elif self.tucker:
-                self.lokr_t2 = nn.Parameter(torch.empty(lora_dim, lora_dim, shape[2], shape[3]))
-                self.lokr_w2_a = nn.Parameter(torch.empty(lora_dim, shape[0][1])) # b, 1-mode
-                self.lokr_w2_b = nn.Parameter(torch.empty(lora_dim, shape[1][1])) # d, 2-mode
-            else: # Conv2d not cp
+                self.lokr_t2 = nn.Parameter(
+                    torch.empty(lora_dim, lora_dim, shape[2], shape[3])
+                )
+                self.lokr_w2_a = nn.Parameter(
+                    torch.empty(lora_dim, shape[0][1])
+                )  # b, 1-mode
+                self.lokr_w2_b = nn.Parameter(
+                    torch.empty(lora_dim, shape[1][1])
+                )  # d, 2-mode
+            else:  # Conv2d not cp
                 # bigger part. weight and LoRA. [b, dim] x [dim, d*k1*k2]
                 self.lokr_w2_a = nn.Parameter(torch.empty(shape[0][1], lora_dim))
-                self.lokr_w2_b = nn.Parameter(torch.empty(lora_dim, shape[1][1]*shape[2]*shape[3]))
+                self.lokr_w2_b = nn.Parameter(
+                    torch.empty(lora_dim, shape[1][1] * shape[2] * shape[3])
+                )
                 # w1 ⊗ (w2_a x w2_b) = (a, b)⊗((c, dim)x(dim, d*k1*k2)) = (a, b)⊗(c, d*k1*k2) = (ac, bd*k1*k2)
-            
+
             self.op = F.conv2d
             self.extra_args = {
                 "stride": org_module.stride,
                 "padding": org_module.padding,
                 "dilation": org_module.dilation,
-                "groups": org_module.groups
+                "groups": org_module.groups,
             }
 
-        else: # Linear
+        else:  # Linear
             in_dim = org_module.in_features
             out_dim = org_module.out_features
-            
+
             in_m, in_n = factorization(in_dim, factor)
             out_l, out_k = factorization(out_dim, factor)
-            shape = ((out_l, out_k), (in_m, in_n)) # ((a, b), (c, d)), out_dim = a*c, in_dim = b*d
-            
+            shape = (
+                (out_l, out_k),
+                (in_m, in_n),
+            )  # ((a, b), (c, d)), out_dim = a*c, in_dim = b*d
+
             # smaller part. weight scale
-            if decompose_both and lora_dim < max(shape[0][0], shape[1][0])/2:
+            if decompose_both and lora_dim < max(shape[0][0], shape[1][0]) / 2:
                 self.lokr_w1_a = nn.Parameter(torch.empty(shape[0][0], lora_dim))
                 self.lokr_w1_b = nn.Parameter(torch.empty(lora_dim, shape[1][0]))
             else:
                 self.use_w1 = True
-                self.lokr_w1 = nn.Parameter(torch.empty(shape[0][0], shape[1][0]))  # a*c, 1-mode
+                self.lokr_w1 = nn.Parameter(
+                    torch.empty(shape[0][0], shape[1][0])
+                )  # a*c, 1-mode
 
-            if lora_dim < max(shape[0][1], shape[1][1])/2:
+            if lora_dim < max(shape[0][1], shape[1][1]) / 2:
                 # bigger part. weight and LoRA. [b, dim] x [dim, d]
                 self.lokr_w2_a = nn.Parameter(torch.empty(shape[0][1], lora_dim))
                 self.lokr_w2_b = nn.Parameter(torch.empty(lora_dim, shape[1][1]))
@@ -155,21 +177,22 @@ def __init__(
 
             self.op = F.linear
             self.extra_args = {}
-        
+
         self.dropout = dropout
         if dropout:
             print("[WARN]LoHa/LoKr haven't implemented normal dropout yet.")
         self.rank_dropout = rank_dropout
+        self.rank_dropout_scale = rank_dropout_scale
         self.module_dropout = module_dropout
-        
+
         if isinstance(alpha, torch.Tensor):
             alpha = alpha.detach().float().numpy()  # without casting, bf16 causes error
         alpha = lora_dim if alpha is None or alpha == 0 else alpha
         if self.use_w2 and self.use_w1:
-            #use scale = 1
+            # use scale = 1
             alpha = lora_dim
         self.scale = alpha / self.lora_dim
-        self.register_buffer('alpha', torch.tensor(alpha)) # 定数として扱える
+        self.register_buffer("alpha", torch.tensor(alpha))  # 定数として扱える
 
         if self.use_w2:
             torch.nn.init.constant_(self.lokr_w2, 0)
@@ -178,7 +201,7 @@ def __init__(
                 torch.nn.init.normal_(self.lokr_t2, std=0.1)
             torch.nn.init.normal_(self.lokr_w2_a, std=1)
             torch.nn.init.constant_(self.lokr_w2_b, 0)
-        
+
         if self.use_w1:
             torch.nn.init.normal_(self.lokr_w1, std=1)
         else:
@@ -188,11 +211,17 @@ def __init__(
         self.multiplier = multiplier
         self.org_module = [org_module]
         weight = make_kron(
-            self.lokr_w1 if self.use_w1 else self.lokr_w1_a@self.lokr_w1_b,
-            (self.lokr_w2 if self.use_w2 
-             else make_weight_cp(self.lokr_t2, self.lokr_w2_a, self.lokr_w2_b) if self.tucker 
-             else self.lokr_w2_a@self.lokr_w2_b),
-            torch.tensor(self.multiplier * self.scale)
+            self.lokr_w1 if self.use_w1 else self.lokr_w1_a @ self.lokr_w1_b,
+            (
+                self.lokr_w2
+                if self.use_w2
+                else (
+                    make_weight_cp(self.lokr_t2, self.lokr_w2_a, self.lokr_w2_b)
+                    if self.tucker
+                    else self.lokr_w2_a @ self.lokr_w2_b
+                )
+            ),
+            torch.tensor(self.multiplier * self.scale),
         )
         assert torch.sum(torch.isnan(weight)) == 0, "weight is nan"
 
@@ -200,47 +229,54 @@ def __init__(
     def apply_to(self):
         self.org_forward = self.org_module[0].forward
         self.org_module[0].forward = self.forward
-    
-    def get_weight(self, orig_weight = None):
+
+    def get_weight(self, orig_weight=None):
         weight = make_kron(
-            self.lokr_w1 if self.use_w1 else self.lokr_w1_a@self.lokr_w1_b,
-            (self.lokr_w2 if self.use_w2 
-             else make_weight_cp(self.lokr_t2, self.lokr_w2_a, self.lokr_w2_b) if self.tucker 
-             else self.lokr_w2_a@self.lokr_w2_b),
-            torch.tensor(self.scale)
+            self.lokr_w1 if self.use_w1 else self.lokr_w1_a @ self.lokr_w1_b,
+            (
+                self.lokr_w2
+                if self.use_w2
+                else make_weight_cp(self.lokr_t2, self.lokr_w2_a, self.lokr_w2_b)
+                if self.tucker
+                else self.lokr_w2_a @ self.lokr_w2_b
+            ),
+            torch.tensor(self.scale),
         )
         if orig_weight is not None:
             weight = weight.reshape(orig_weight.shape)
         if self.training and self.rank_dropout:
-            drop = torch.rand(weight.size(0)) < self.rank_dropout
-            weight *= drop.view(-1, [1]*len(weight.shape[1:])).to(weight.device)
+            drop = (torch.rand(weight.size(0)) < self.rank_dropout).to(weight.dtype)
+            drop = drop.view(-1, *[1] * len(weight.shape[1:])).to(weight.device)
+            if self.rank_dropout_scale:
+                drop /= drop.mean()
+            weight *= drop
         return weight
 
     @torch.no_grad()
     def apply_max_norm(self, max_norm, device=None):
         orig_norm = self.get_weight().norm()
-        norm = torch.clamp(orig_norm, max_norm/2)
+        norm = torch.clamp(orig_norm, max_norm / 2)
         desired = torch.clamp(norm, max=max_norm)
-        ratio = desired.cpu()/norm.cpu()
-        
+        ratio = desired.cpu() / norm.cpu()
+
         scaled = ratio.item() != 1.0
         if scaled:
-            modules = (4 - self.use_w1 - self.use_w2 + (not self.use_w2 and self.tucker))
+            modules = 4 - self.use_w1 - self.use_w2 + (not self.use_w2 and self.tucker)
             if self.use_w1:
-                self.lokr_w1 *= ratio**(1/modules)
+                self.lokr_w1 *= ratio ** (1 / modules)
             else:
-                self.lokr_w1_a *= ratio**(1/modules)
-                self.lokr_w1_b *= ratio**(1/modules)
-            
+                self.lokr_w1_a *= ratio ** (1 / modules)
+                self.lokr_w1_b *= ratio ** (1 / modules)
+
             if self.use_w2:
-                self.lokr_w2 *= ratio**(1/modules)
+                self.lokr_w2 *= ratio ** (1 / modules)
             else:
                 if self.tucker:
-                    self.lokr_t2 *= ratio**(1/modules)
-                self.lokr_w2_a  *= ratio**(1/modules)
-                self.lokr_w2_b  *= ratio**(1/modules)
-        
-        return scaled, orig_norm*ratio
+                    self.lokr_t2 *= ratio ** (1 / modules)
+                self.lokr_w2_a *= ratio ** (1 / modules)
+                self.lokr_w2_b *= ratio ** (1 / modules)
+
+        return scaled, orig_norm * ratio
 
     def forward(self, x):
         if self.module_dropout and self.training:
@@ -248,16 +284,13 @@ def forward(self, x):
                 return self.op(
                     x,
                     self.org_module[0].weight.data,
-                    None if self.org_module[0].bias is None else self.org_module[0].bias.data
+                    None
+                    if self.org_module[0].bias is None
+                    else self.org_module[0].bias.data,
                 )
         weight = (
-            self.org_module[0].weight.data 
+            self.org_module[0].weight.data
             + self.get_weight(self.org_module[0].weight.data) * self.multiplier
         )
         bias = None if self.org_module[0].bias is None else self.org_module[0].bias.data
-        return self.op(
-            x, 
-            weight.view(self.shape),
-            bias,
-            **self.extra_args
-        )
+        return self.op(x, weight.view(self.shape), bias, **self.extra_args)
diff --git a/lycoris/modules/glora.py b/lycoris/modules/glora.py
index 98d7ade..f85165f 100644
--- a/lycoris/modules/glora.py
+++ b/lycoris/modules/glora.py
@@ -1,33 +1,39 @@
 import math
-from weakref import ref
 
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
+from .base import ModuleCustomSD
 
-class GLoRAModule(nn.Module):
+
+class GLoRAModule(ModuleCustomSD):
     """
     modifed from kohya-ss/sd-scripts/networks/lora:LoRAModule
     """
 
     def __init__(
-        self, 
-        lora_name, org_module: nn.Module, 
-        multiplier=1.0, 
-        lora_dim=4, alpha=1, 
-        dropout=0., rank_dropout=0., module_dropout=0.,
+        self,
+        lora_name,
+        org_module: nn.Module,
+        multiplier=1.0,
+        lora_dim=4,
+        alpha=1,
+        dropout=0.0,
+        rank_dropout=0.0,
+        module_dropout=0.0,
+        rank_dropout_scale=False,
         *args,
         **kwargs,
     ):
-        """ if alpha == 0 or None, alpha is rank (no scaling). """
+        """if alpha == 0 or None, alpha is rank (no scaling)."""
         super().__init__()
         self.lora_name = lora_name
         self.lora_dim = lora_dim
         self.cp = False
 
         if isinstance(org_module, nn.Conv2d):
-            assert org_module.kernel_size == (1,1)
+            assert org_module.kernel_size == (1, 1)
             in_dim = org_module.in_channels
             out_dim = org_module.out_channels
             self.a1 = nn.Conv2d(in_dim, lora_dim, (1, 1), bias=False)
@@ -44,19 +50,20 @@ def __init__(
         else:
             raise NotImplementedError
         self.shape = org_module.weight.shape
-        
+
         if dropout:
             self.dropout = nn.Dropout(dropout)
         else:
             self.dropout = nn.Identity()
         self.rank_dropout = rank_dropout
+        self.rank_dropout_scale = rank_dropout_scale
         self.module_dropout = module_dropout
-        
+
         if type(alpha) == torch.Tensor:
             alpha = alpha.detach().float().numpy()  # without casting, bf16 causes error
         alpha = lora_dim if alpha is None or alpha == 0 else alpha
         self.scale = alpha / self.lora_dim
-        self.register_buffer('alpha', torch.tensor(alpha)) # 定数として扱える
+        self.register_buffer("alpha", torch.tensor(alpha))  # 定数として扱える
 
         # same as microsoft's
         torch.nn.init.kaiming_uniform_(self.a1.weight, a=math.sqrt(5))
@@ -77,27 +84,38 @@ def make_weight(self, device=None):
         wb1 = self.b1.weight.view(self.b1.weight.size(0), -1)
         wb2 = self.b2.weight.view(self.b2.weight.size(0), -1)
         orig = self.org_module[0].weight.view(self.org_module[0].weight.size(0), -1)
-        return ((wb2 @ wb1) + ((orig @ wa2) @ wa1))
+        return (wb2 @ wb1) + ((orig @ wa2) @ wa1)
 
     def forward(self, x):
         if self.module_dropout and self.training:
             if torch.rand(1) < self.module_dropout:
                 return self.org_forward(x)
         scale = self.scale * self.multiplier
-        
+
         ax_mid = self.a1(x) * scale
         bx_mid = self.b1(x) * scale
-        
+
         if self.rank_dropout and self.training:
-            drop_a = torch.rand(self.lora_dim, device=ax_mid.device) < self.rank_dropout
-            drop_b = torch.rand(self.lora_dim, device=bx_mid.device) < self.rank_dropout
-            if (dims:=len(x.shape)) == 4:
+            drop_a = (
+                torch.rand(self.lora_dim, device=ax_mid.device) < self.rank_dropout
+            ).to(ax_mid.dtype)
+            if self.rank_dropout_scale:
+                drop_a /= drop_a.mean()
+            drop_b = (
+                torch.rand(self.lora_dim, device=bx_mid.device) < self.rank_dropout
+            ).to(bx_mid.dtype)
+            if self.rank_dropout_scale:
+                drop_b /= drop_b.mean()
+            if (dims := len(x.shape)) == 4:
                 drop_a = drop_a.view(1, -1, 1, 1)
                 drop_b = drop_b.view(1, -1, 1, 1)
             else:
-                drop_a = drop_a.view(*[1]*(dims-1), -1)
-                drop_b = drop_b.view(*[1]*(dims-1), -1)
+                drop_a = drop_a.view(*[1] * (dims - 1), -1)
+                drop_b = drop_b.view(*[1] * (dims - 1), -1)
             ax_mid = ax_mid * drop_a
             bx_mid = bx_mid * drop_b
-        
-        return self.org_forward(x + self.dropout(self.a2(ax_mid))*self.scale) + self.dropout(self.b2(bx_mid))*self.scale
+
+        return (
+            self.org_forward(x + self.dropout(self.a2(ax_mid)) * self.scale)
+            + self.dropout(self.b2(bx_mid)) * self.scale
+        )
diff --git a/lycoris/modules/hyperlycoris.py b/lycoris/modules/hyperlycoris.py
deleted file mode 100644
index 7961d00..0000000
--- a/lycoris/modules/hyperlycoris.py
+++ /dev/null
@@ -1,3 +0,0 @@
-'''
-Based on HyperDreamBooth, specific for LyCORIS modules
-'''
\ No newline at end of file
diff --git a/lycoris/modules/hypernet/__init__.py b/lycoris/modules/hypernet/__init__.py
index c6e24c2..243c210 100644
--- a/lycoris/modules/hypernet/__init__.py
+++ b/lycoris/modules/hypernet/__init__.py
@@ -1,5 +1,1 @@
-from .generater import (
-    WeightDecoder,
-    ImgWeightGenerator,
-    TextWeightGenerator
-)
\ No newline at end of file
+from .generater import WeightDecoder, ImgWeightGenerator, TextWeightGenerator
diff --git a/lycoris/modules/hypernet/generater.py b/lycoris/modules/hypernet/generater.py
index c8722cc..b0c24bd 100644
--- a/lycoris/modules/hypernet/generater.py
+++ b/lycoris/modules/hypernet/generater.py
@@ -13,14 +13,19 @@
 
 
 def _get_sinusoid_encoding_table(n_position, d_hid):
-    ''' Sinusoid position encoding table '''
+    """Sinusoid position encoding table"""
     # TODO: make it with torch instead of numpy
 
     def get_position_angle_vec(position):
         # this part calculate the position In brackets
-        return [position / np.power(10000, 2 * (hid_j // 2) / d_hid) for hid_j in range(d_hid)]
+        return [
+            position / np.power(10000, 2 * (hid_j // 2) / d_hid)
+            for hid_j in range(d_hid)
+        ]
 
-    sinusoid_table = np.array([get_position_angle_vec(pos_i) for pos_i in range(n_position)])
+    sinusoid_table = np.array(
+        [get_position_angle_vec(pos_i) for pos_i in range(n_position)]
+    )
     # [:, 0::2] are all even subscripts, is dim_2i
     sinusoid_table[:, 0::2] = np.sin(sinusoid_table[:, 0::2])  # dim 2i
     sinusoid_table[:, 1::2] = np.cos(sinusoid_table[:, 1::2])  # dim 2i+1
@@ -29,47 +34,55 @@ def get_position_angle_vec(position):
 
 
 class WeightDecoder(nn.Module):
-    def __init__(self, weight_dim: int = 150, weight_num: int = 54, decoder_blocks: int = 4):
+    def __init__(
+        self, weight_dim: int = 150, weight_num: int = 54, decoder_blocks: int = 4
+    ):
         super(WeightDecoder, self).__init__()
         self.weight_num = weight_num
         self.weight_dim = weight_dim
-        
+
         self.register_buffer(
-            'block_pos_emb', 
-            _get_sinusoid_encoding_table(weight_num*2, weight_dim)
+            "block_pos_emb", _get_sinusoid_encoding_table(weight_num * 2, weight_dim)
         )
-        
+
         # calc heads for mem-eff or flash_attn
         heads = 1
-        while weight_dim % heads==0 and weight_dim // heads > 64:
+        while weight_dim % heads == 0 and weight_dim // heads > 64:
             heads *= 2
-        
+
         self.pos_emb_proj = nn.Linear(weight_dim, weight_dim, bias=False)
         self.decoder_model = nn.ModuleList(
-            TransformerBlock(weight_dim, heads, weight_dim//heads, context_dim=weight_dim, gated_ff=False)
+            TransformerBlock(
+                weight_dim,
+                heads,
+                weight_dim // heads,
+                context_dim=weight_dim,
+                gated_ff=False,
+            )
             for _ in range(decoder_blocks)
         )
         # self.delta_proj = nn.Linear(weight_dim, weight_dim, bias=False)
         self.delta_proj = nn.Sequential(
-            nn.LayerNorm(weight_dim),
-            nn.Linear(weight_dim, weight_dim, bias=False)
+            nn.LayerNorm(weight_dim), nn.Linear(weight_dim, weight_dim, bias=False)
         )
         self.init_weights()
-    
+
     def init_weights(self):
         def basic_init(module):
             if isinstance(module, nn.Linear):
                 nn.init.xavier_uniform_(module.weight)
                 if module.bias is not None:
                     nn.init.constant_(module.bias, 0)
-        
+
         self.apply(basic_init)
         torch.nn.init.constant_(self.delta_proj[1].weight, 0)
         # advice from Nataniel Ruiz, looks like 1e-3 is small enough
         # torch.nn.init.normal_(self.delta_proj[1].weight, std=1e-3)
-    
+
     def forward(self, weight, features):
-        pos_emb = self.pos_emb_proj(self.block_pos_emb[:, :weight.size(1)].clone().detach())
+        pos_emb = self.pos_emb_proj(
+            self.block_pos_emb[:, : weight.size(1)].clone().detach()
+        )
         h = weight + pos_emb
         for decoder in self.decoder_model:
             h = decoder(h, context=features)
@@ -79,12 +92,12 @@ def forward(self, weight, features):
 
 class ImgWeightGenerator(nn.Module):
     def __init__(
-        self, 
+        self,
         encoder_model_name: str = "vit_base_patch16_224",
         train_encoder: bool = False,
-        reference_size: Tuple[int] = (224, 224), 
-        weight_dim: int = 150,                      # 100+50 in paper
-        weight_num: int = 54,                       # 21*2 in SD UNet + 12 in CLIP-L TE
+        reference_size: Tuple[int] = (224, 224),
+        weight_dim: int = 150,  # 100+50 in paper
+        weight_num: int = 54,  # 21*2 in SD UNet + 12 in CLIP-L TE
         decoder_blocks: int = 4,
         sample_iters: int = 1,
     ):
@@ -94,16 +107,17 @@ def __init__(
         self.weight_dim = weight_dim
         self.sample_iters = sample_iters
         self.train_encoder = train_encoder
-        
+
         self.register_buffer(
-            'block_pos_emb', 
-            _get_sinusoid_encoding_table(weight_num*2, weight_dim)
+            "block_pos_emb", _get_sinusoid_encoding_table(weight_num * 2, weight_dim)
+        )
+
+        self.encoder_model: nn.Module = create_model(
+            encoder_model_name, pretrained=True
         )
-        
-        self.encoder_model: nn.Module = create_model(encoder_model_name, pretrained=True)
         for p in self.encoder_model.parameters():
             p.requires_grad_(train_encoder)
-        
+
         test_input = torch.randn(1, 3, *reference_size)
         test_output = self.encoder_model.forward_features(test_input)
         if isinstance(test_output, list):
@@ -111,11 +125,11 @@ def __init__(
         if len(test_output.shape) == 4:
             # B, C, H, W -> B, L, C
             test_output = test_output.view(1, test_output.size(1), -1).transpose(1, 2)
-        
+
         self.feature_proj = nn.Linear(test_output.shape[-1], weight_dim, bias=False)
         self.pos_emb_proj = nn.Linear(weight_dim, weight_dim, bias=False)
         self.decoder_model = WeightDecoder(weight_dim, weight_num, decoder_blocks)
-    
+
     def forward(self, ref_img, iters=None, weight=None, ensure_grad=0):
         ref_img = resize(ref_img, self.ref_size, antialias=True)
         if not self.train_encoder:
@@ -127,14 +141,16 @@ def forward(self, ref_img, iters=None, weight=None, ensure_grad=0):
             features = features[-1]
         if len(features.shape) == 4:
             # B, C, H, W -> B, L, C
-            features = features.view(features.size(0), features.size(1), -1).transpose(1, 2)
-        features = self.feature_proj(features+ensure_grad)
-        
+            features = features.view(features.size(0), features.size(1), -1).transpose(
+                1, 2
+            )
+        features = self.feature_proj(features + ensure_grad)
+
         if weight is None:
             weight = torch.zeros(
                 ref_img.size(0), self.weight_num, self.weight_dim, device=ref_img.device
             )
-        
+
         for iter in range(iters or self.sample_iters):
             weight = self.decoder_model(weight, features)
         return weight
@@ -142,31 +158,31 @@ def forward(self, ref_img, iters=None, weight=None, ensure_grad=0):
 
 class TextWeightGenerator(nn.Module):
     def __init__(
-        self, 
+        self,
         train_encoder: bool = False,
-        reference_size: Tuple[int] = (224, 224), 
-        weight_dim: int = 150,                      # 100+50 in paper
-        weight_num: int = 54,                       # 21*2 in SD UNet + 12 in CLIP-L TE
+        reference_size: Tuple[int] = (224, 224),
+        weight_dim: int = 150,  # 100+50 in paper
+        weight_num: int = 54,  # 21*2 in SD UNet + 12 in CLIP-L TE
         decoder_blocks: int = 4,
         sample_iters: int = 1,
     ):
         from .text_encoder import FrozenOpenCLIPEmbedder
+
         super(TextWeightGenerator, self).__init__()
         self.ref_size = reference_size
         self.weight_num = weight_num
         self.weight_dim = weight_dim
         self.sample_iters = sample_iters
         self.train_encoder = train_encoder
-        
+
         self.register_buffer(
-            'block_pos_emb', 
-            _get_sinusoid_encoding_table(weight_num*2, weight_dim)
+            "block_pos_emb", _get_sinusoid_encoding_table(weight_num * 2, weight_dim)
         )
-        
+
         self.encoder_model: nn.Module = FrozenOpenCLIPEmbedder()
         for p in self.encoder_model.parameters():
             p.requires_grad_(train_encoder)
-        
+
         test_input = ["test"]
         test_output = self.encoder_model(test_input)
         if isinstance(test_output, list):
@@ -174,11 +190,11 @@ def __init__(
         if len(test_output.shape) == 4:
             # B, C, H, W -> B, L, C
             test_output = test_output.view(1, test_output.size(1), -1).transpose(1, 2)
-        
+
         self.feature_proj = nn.Linear(test_output.shape[-1], weight_dim, bias=False)
         self.pos_emb_proj = nn.Linear(weight_dim, weight_dim, bias=False)
         self.decoder_model = WeightDecoder(weight_dim, weight_num, decoder_blocks)
-    
+
     def forward(self, caption, iters=None, weight=None, ensure_grad=0):
         if not self.train_encoder:
             with torch.no_grad():
@@ -189,14 +205,19 @@ def forward(self, caption, iters=None, weight=None, ensure_grad=0):
             features = features[-1]
         if len(features.shape) == 4:
             # B, C, H, W -> B, L, C
-            features = features.view(features.size(0), features.size(1), -1).transpose(1, 2)
-        features = self.feature_proj(features+ensure_grad)
-        
+            features = features.view(features.size(0), features.size(1), -1).transpose(
+                1, 2
+            )
+        features = self.feature_proj(features + ensure_grad)
+
         if weight is None:
             weight = torch.zeros(
-                features.size(0), self.weight_num, self.weight_dim, device=features.device
+                features.size(0),
+                self.weight_num,
+                self.weight_dim,
+                device=features.device,
             )
-        
+
         for iter in range(iters or self.sample_iters):
             weight = self.decoder_model(weight, features)
-        return weight
\ No newline at end of file
+        return weight
diff --git a/lycoris/modules/hypernet/text_encoder.py b/lycoris/modules/hypernet/text_encoder.py
index 036816b..fa4e5ec 100644
--- a/lycoris/modules/hypernet/text_encoder.py
+++ b/lycoris/modules/hypernet/text_encoder.py
@@ -1,6 +1,6 @@
-'''
+"""
 https://github.com/Stability-AI/stablediffusion/blob/main/ldm/modules/encoders/modules.py
-'''
+"""
 
 
 import torch
@@ -26,7 +26,7 @@ def encode(self, x):
 
 
 class ClassEmbedder(nn.Module):
-    def __init__(self, embed_dim, n_classes=1000, key='class', ucg_rate=0.1):
+    def __init__(self, embed_dim, n_classes=1000, key="class", ucg_rate=0.1):
         super().__init__()
         self.key = key
         self.embedding = nn.Embedding(n_classes, embed_dim)
@@ -38,15 +38,17 @@ def forward(self, batch, key=None, disable_dropout=False):
             key = self.key
         # this is for use in crossattn
         c = batch[key][:, None]
-        if self.ucg_rate > 0. and not disable_dropout:
-            mask = 1. - torch.bernoulli(torch.ones_like(c) * self.ucg_rate)
-            c = mask * c + (1-mask) * torch.ones_like(c)*(self.n_classes-1)
+        if self.ucg_rate > 0.0 and not disable_dropout:
+            mask = 1.0 - torch.bernoulli(torch.ones_like(c) * self.ucg_rate)
+            c = mask * c + (1 - mask) * torch.ones_like(c) * (self.n_classes - 1)
             c = c.long()
         c = self.embedding(c)
         return c
 
     def get_unconditional_conditioning(self, bs, device="cuda"):
-        uc_class = self.n_classes - 1  # 1000 classes --> 0 ... 999, one extra class for ucg (class 1000)
+        uc_class = (
+            self.n_classes - 1
+        )  # 1000 classes --> 0 ... 999, one extra class for ucg (class 1000)
         uc = torch.ones((bs,), device=device) * uc_class
         uc = {self.key: uc}
         return uc
@@ -60,24 +62,34 @@ def disabled_train(self, mode=True):
 
 class FrozenT5Embedder(AbstractEncoder):
     """Uses the T5 transformer encoder for text"""
-    def __init__(self, version="google/t5-v1_1-large", device="cuda", max_length=77, freeze=True):  # others are google/t5-v1_1-xl and google/t5-v1_1-xxl
+
+    def __init__(
+        self, version="google/t5-v1_1-large", device="cuda", max_length=77, freeze=True
+    ):  # others are google/t5-v1_1-xl and google/t5-v1_1-xxl
         super().__init__()
         self.tokenizer = T5Tokenizer.from_pretrained(version)
         self.transformer = T5EncoderModel.from_pretrained(version)
         self.device = device
-        self.max_length = max_length   # TODO: typical value?
+        self.max_length = max_length  # TODO: typical value?
         if freeze:
             self.freeze()
 
     def freeze(self):
         self.transformer = self.transformer.eval()
-        #self.train = disabled_train
+        # self.train = disabled_train
         for param in self.parameters():
             param.requires_grad = False
 
     def forward(self, text):
-        batch_encoding = self.tokenizer(text, truncation=True, max_length=self.max_length, return_length=True,
-                                        return_overflowing_tokens=False, padding="max_length", return_tensors="pt")
+        batch_encoding = self.tokenizer(
+            text,
+            truncation=True,
+            max_length=self.max_length,
+            return_length=True,
+            return_overflowing_tokens=False,
+            padding="max_length",
+            return_tensors="pt",
+        )
         tokens = batch_encoding["input_ids"].to(self.device)
         outputs = self.transformer(input_ids=tokens)
 
@@ -90,14 +102,18 @@ def encode(self, text):
 
 class FrozenCLIPEmbedder(AbstractEncoder):
     """Uses the CLIP transformer encoder for text (from huggingface)"""
-    LAYERS = [
-        "last",
-        "pooled",
-        "last_with_pooled",
-        "hidden"
-    ]
-    def __init__(self, version="openai/clip-vit-large-patch14", device="cuda", max_length=77,
-                 freeze=True, layer="last", layer_idx=None):  # clip-vit-base-patch32
+
+    LAYERS = ["last", "pooled", "last_with_pooled", "hidden"]
+
+    def __init__(
+        self,
+        version="openai/clip-vit-large-patch14",
+        device="cuda",
+        max_length=77,
+        freeze=True,
+        layer="last",
+        layer_idx=None,
+    ):  # clip-vit-base-patch32
         super().__init__()
         assert layer in self.LAYERS
         self.tokenizer = CLIPTokenizer.from_pretrained(version)
@@ -114,27 +130,31 @@ def __init__(self, version="openai/clip-vit-large-patch14", device="cuda", max_l
 
     def freeze(self):
         self.transformer = self.transformer.eval()
-        #self.train = disabled_train
+        # self.train = disabled_train
         for param in self.parameters():
             param.requires_grad = False
 
     def forward(self, text):
         batch_encoding = self.tokenizer(
-            text, truncation=True, max_length=self.max_length, return_length=True,
-            return_overflowing_tokens=False, padding="max_length", return_tensors="pt"
+            text,
+            truncation=True,
+            max_length=self.max_length,
+            return_length=True,
+            return_overflowing_tokens=False,
+            padding="max_length",
+            return_tensors="pt",
         )
         tokens = batch_encoding["input_ids"].to(self.device)
-        outputs = self.transformer(input_ids=tokens, output_hidden_states=self.layer=="hidden")
+        outputs = self.transformer(
+            input_ids=tokens, output_hidden_states=self.layer == "hidden"
+        )
         # print(outputs)
         if self.layer == "last":
             z = outputs.last_hidden_state
         elif self.layer == "pooled":
             z = outputs.pooler_output[:, None, :]
         elif self.layer == "last_with_pooled":
-            z = (
-                outputs.last_hidden_state,
-                outputs.pooler_output
-            )
+            z = (outputs.last_hidden_state, outputs.pooler_output)
         else:
             z = outputs.hidden_states[self.layer_idx]
         return z
@@ -147,27 +167,27 @@ class FrozenOpenCLIPEmbedder(AbstractEncoder):
     """
     Uses the OpenCLIP transformer encoder for text
     """
+
     LAYERS = [
-        #"pooled",
+        # "pooled",
         "last",
-        "penultimate"
+        "penultimate",
     ]
+
     def __init__(
-        self, 
-        arch="ViT-L-14", 
-        version="datacomp_xl_s13b_b90k", 
+        self,
+        arch="ViT-L-14",
+        version="datacomp_xl_s13b_b90k",
         max_length=77,
-        freeze=True, 
-        layer="last"
+        freeze=True,
+        layer="last",
     ):
         super().__init__()
         assert layer in self.LAYERS
-        model, _, _ = open_clip.create_model_and_transforms(
-            arch, pretrained=version
-        )
+        model, _, _ = open_clip.create_model_and_transforms(arch, pretrained=version)
         del model.visual
         self.model = model
-        
+
         self.max_length = max_length
         if freeze:
             self.freeze()
@@ -186,7 +206,9 @@ def freeze(self):
 
     def forward(self, text):
         tokens = open_clip.tokenize(text)
-        z = self.encode_with_transformer(tokens.to(self.model.token_embedding.weight.device))
+        z = self.encode_with_transformer(
+            tokens.to(self.model.token_embedding.weight.device)
+        )
         return z
 
     def encode_with_transformer(self, text):
@@ -198,11 +220,14 @@ def encode_with_transformer(self, text):
         x = self.model.ln_final(x)
         return x
 
-    def text_transformer_forward(self, x: torch.Tensor, attn_mask = None):
+    def text_transformer_forward(self, x: torch.Tensor, attn_mask=None):
         for i, r in enumerate(self.model.transformer.resblocks):
             if i == len(self.model.transformer.resblocks) - self.layer_idx:
                 break
-            if self.model.transformer.grad_checkpointing and not torch.jit.is_scripting():
+            if (
+                self.model.transformer.grad_checkpointing
+                and not torch.jit.is_scripting()
+            ):
                 x = checkpoint(r, x, attn_mask)
             else:
                 x = r(x, attn_mask=attn_mask)
@@ -214,15 +239,17 @@ def encode(self, text):
 
 class FrozenCLIPT5Encoder(AbstractEncoder):
     def __init__(
-        self, 
-        clip_version = "openai/clip-vit-large-patch14", 
-        t5_version = "google/t5-v1_1-xl", 
-        device = "cuda",
-        clip_max_length =77, 
-        t5_max_length = 77
+        self,
+        clip_version="openai/clip-vit-large-patch14",
+        t5_version="google/t5-v1_1-xl",
+        device="cuda",
+        clip_max_length=77,
+        t5_max_length=77,
     ):
         super().__init__()
-        self.clip_encoder = FrozenCLIPEmbedder(clip_version, device, max_length=clip_max_length)
+        self.clip_encoder = FrozenCLIPEmbedder(
+            clip_version, device, max_length=clip_max_length
+        )
         self.t5_encoder = FrozenT5Embedder(t5_version, device, max_length=t5_max_length)
 
     def encode(self, text):
@@ -234,7 +261,7 @@ def forward(self, text):
         return [clip_z, t5_z]
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     clip = FrozenCLIPEmbedder(layer="last_with_pooled").cuda()
-    ctx, embed = clip('test')
-    print(ctx.shape, embed.shape)
\ No newline at end of file
+    ctx, embed = clip("test")
+    print(ctx.shape, embed.shape)
diff --git a/lycoris/modules/ia3.py b/lycoris/modules/ia3.py
index e679945..42e9ed5 100644
--- a/lycoris/modules/ia3.py
+++ b/lycoris/modules/ia3.py
@@ -1,30 +1,30 @@
-import math
-
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
+from .base import ModuleCustomSD
+
 
-class IA3Module(nn.Module):
+class IA3Module(ModuleCustomSD):
     """
     Hadamard product Implementaion for Low Rank Adaptation
     """
 
     def __init__(
-        self, 
-        lora_name, 
-        org_module: nn.Module, 
-        multiplier=1.0, 
-        train_on_input=False, 
+        self,
+        lora_name,
+        org_module: nn.Module,
+        multiplier=1.0,
+        train_on_input=False,
         **kwargs
     ):
-        """ if alpha == 0 or None, alpha is rank (no scaling). """
+        """if alpha == 0 or None, alpha is rank (no scaling)."""
         super().__init__()
         self.lora_name = lora_name
-        self.cp=False
-        
+        self.cp = False
+
         self.shape = org_module.weight.shape
-        if org_module.__class__.__name__ == 'Conv2d':
+        if org_module.__class__.__name__ == "Conv2d":
             in_dim = org_module.in_channels
             out_dim = org_module.out_channels
             if train_on_input:
@@ -39,7 +39,7 @@ def __init__(
                 train_dim = in_dim
             else:
                 train_dim = out_dim
-            
+
             self.weight = nn.Parameter(torch.empty(train_dim))
 
         # Need more experiences on init method
@@ -47,10 +47,10 @@ def __init__(
 
         self.multiplier = multiplier
         self.org_forward = None
-        self.org_module = [org_module] # remove in applying
+        self.org_module = [org_module]  # remove in applying
         self.grad_ckpt = False
         self.train_input = train_on_input
-        self.register_buffer('on_input', torch.tensor(int(train_on_input)))
+        self.register_buffer("on_input", torch.tensor(int(train_on_input)))
 
     def apply_to(self):
         self.org_forward = self.org_module[0].forward
@@ -65,4 +65,4 @@ def forward(self, x):
         if not self.train_input:
             out = out * (1 + self.weight * self.multiplier)
             out = out.to(dtype)
-        return out
\ No newline at end of file
+        return out
diff --git a/lycoris/modules/locon.py b/lycoris/modules/locon.py
index 23f93cc..0e85334 100644
--- a/lycoris/modules/locon.py
+++ b/lycoris/modules/locon.py
@@ -1,27 +1,33 @@
 import math
-from weakref import ref
-from collections import OrderedDict, abc as container_abcs
 
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
+from .base import ModuleCustomSD
 
-class LoConModule(nn.Module):
+
+class LoConModule(ModuleCustomSD):
     """
     modifed from kohya-ss/sd-scripts/networks/lora:LoRAModule
     """
 
     def __init__(
-        self, 
-        lora_name, org_module: nn.Module, 
-        multiplier=1.0, 
-        lora_dim=4, alpha=1, 
-        dropout=0., rank_dropout=0., module_dropout=0.,
-        use_tucker=False, use_scalar=False,
+        self,
+        lora_name,
+        org_module: nn.Module,
+        multiplier=1.0,
+        lora_dim=4,
+        alpha=1,
+        dropout=0.0,
+        rank_dropout=0.0,
+        module_dropout=0.0,
+        use_tucker=False,
+        use_scalar=False,
+        rank_dropout_scale=False,
         **kwargs,
     ):
-        """ if alpha == 0 or None, alpha is rank (no scaling). """
+        """if alpha == 0 or None, alpha is rank (no scaling)."""
         super().__init__()
         self.lora_name = lora_name
         self.lora_dim = lora_dim
@@ -39,10 +45,14 @@ def __init__(
             self.up_op = F.conv2d
             if use_tucker and k_size != (1, 1):
                 self.lora_down = nn.Conv2d(in_dim, lora_dim, (1, 1), bias=False)
-                self.lora_mid = nn.Conv2d(lora_dim, lora_dim, k_size, stride, padding, bias=False)
+                self.lora_mid = nn.Conv2d(
+                    lora_dim, lora_dim, k_size, stride, padding, bias=False
+                )
                 self.tucker = True
             else:
-                self.lora_down = nn.Conv2d(in_dim, lora_dim, k_size, stride, padding, bias=False)
+                self.lora_down = nn.Conv2d(
+                    in_dim, lora_dim, k_size, stride, padding, bias=False
+                )
             self.lora_up = nn.Conv2d(lora_dim, out_dim, (1, 1), bias=False)
         elif isinstance(org_module, nn.Linear):
             self.isconv = False
@@ -55,19 +65,20 @@ def __init__(
         else:
             raise NotImplementedError
         self.shape = org_module.weight.shape
-        
+
         if dropout:
             self.dropout = nn.Dropout(dropout)
         else:
             self.dropout = nn.Identity()
         self.rank_dropout = rank_dropout
+        self.rank_dropout_scale = rank_dropout_scale
         self.module_dropout = module_dropout
-        
+
         if type(alpha) == torch.Tensor:
             alpha = alpha.detach().float().numpy()  # without casting, bf16 causes error
         alpha = lora_dim if alpha is None or alpha == 0 else alpha
         self.scale = alpha / self.lora_dim
-        self.register_buffer('alpha', torch.tensor(alpha)) # 定数として扱える
+        self.register_buffer("alpha", torch.tensor(alpha))  # 定数として扱える
 
         if use_scalar:
             self.scalar = nn.Parameter(torch.tensor(0.0))
@@ -85,7 +96,7 @@ def __init__(
         self.multiplier = multiplier
         self.org_module = [org_module]
         self.register_load_state_dict_post_hook(self.load_weight_hook)
-    
+
     def load_weight_hook(self, *args, **kwargs):
         self.scalar = nn.Parameter(torch.ones_like(self.scalar))
 
@@ -99,46 +110,31 @@ def apply_to(self, is_hypernet=False, **kwargs):
     def make_weight(self, device=None):
         wa = self.lora_up.weight.to(device)
         wb = self.lora_down.weight.to(device)
-        return (wa.view(wa.size(0), -1) @ wb.view(wb.size(0), -1)).view(self.shape) * self.scalar
-    
-    def state_dict(self, *args, destination=None, prefix='', keep_vars=False):
-        # TODO: Remove `args` and the parsing logic when BC allows.
-        if len(args) > 0:
-            if destination is None:
-                destination = args[0]
-            if len(args) > 1 and prefix == '':
-                prefix = args[1]
-            if len(args) > 2 and keep_vars is False:
-                keep_vars = args[2]
-            # DeprecationWarning is ignored by default
-
-        if destination is None:
-            destination = OrderedDict()
-            destination._metadata = OrderedDict()
-
-        local_metadata = dict(version=self._version)
-        if hasattr(destination, "_metadata"):
-            destination._metadata[prefix[:-1]] = local_metadata
-
-        destination[f'{prefix}alpha'] = self.alpha
-        destination[f'{prefix}lora_up.weight'] = self.lora_up.weight * self.scalar
-        destination[f'{prefix}lora_down.weight'] = self.lora_down.weight
+        return (wa.view(wa.size(0), -1) @ wb.view(wb.size(0), -1)).view(
+            self.shape
+        ) * self.scalar
+
+    def custom_state_dict(self):
+        destination = {}
+        destination["alpha"] = self.alpha
+        destination["lora_up.weight"] = self.lora_up.weight * self.scalar
+        destination["lora_down.weight"] = self.lora_down.weight
         if self.tucker:
-            destination[f'{prefix}lora_mid.weight'] = self.lora_mid.weight
+            destination["lora_mid.weight"] = self.lora_mid.weight
         return destination
 
     @torch.no_grad()
     def apply_max_norm(self, max_norm, device=None):
-        orig_norm = self.make_weight(device).norm()*self.scale
-        norm = torch.clamp(orig_norm, max_norm/2)
+        orig_norm = self.make_weight(device).norm() * self.scale
+        norm = torch.clamp(orig_norm, max_norm / 2)
         desired = torch.clamp(norm, max=max_norm)
-        ratio = desired.cpu()/norm.cpu()
-        
+        ratio = desired.cpu() / norm.cpu()
+
         scaled = ratio.item() != 1.0
         if scaled:
             self.scalar *= ratio
-        
-        return scaled, orig_norm*ratio
+
+        return scaled, orig_norm * ratio
 
     def update_weights(self, down, up, idx):
         self.down, self.up = self.make_lightweight(down.squeeze(1), up.squeeze(1), idx)
@@ -157,8 +153,14 @@ def make_lightweight(self, down, up, seed=None, down_aux=None, up_aux=None):
             rng_state = torch.get_rng_state()
             torch.manual_seed(seed)
             if down_aux is None or up_aux is None:
-                down_aux = torch.empty(down.size(down.dim()-1), self.lora_down.weight.size(1), device=down.device)
-                up_aux = torch.empty(self.lora_up.weight.size(0), up.size(up.dim()-2), device=up.device)
+                down_aux = torch.empty(
+                    down.size(down.dim() - 1),
+                    self.lora_down.weight.size(1),
+                    device=down.device,
+                )
+                up_aux = torch.empty(
+                    self.lora_up.weight.size(0), up.size(up.dim() - 2), device=up.device
+                )
                 nn.init.orthogonal_(down_aux)
                 nn.init.orthogonal_(up_aux)
                 # print(up_aux.shape)
@@ -167,7 +169,7 @@ def make_lightweight(self, down, up, seed=None, down_aux=None, up_aux=None):
             down = down.squeeze(0)
         if up.dim() == 3 and up.size(0) == 1:
             up = up.squeeze(0)
-        down = down + 1 # avoid zero grad or slow training, give it a constant
+        down = down + 1  # avoid zero grad or slow training, give it a constant
         return (down @ down_aux), (up_aux @ up)
 
     def apply_lightweight(self, down, up, seed=None, down_aux=None, up_aux=None):
@@ -181,48 +183,56 @@ def hypernet_forward(self, x):
             if torch.rand(1) < self.module_dropout:
                 return self.org_forward(x)
         scale = self.scale * self.multiplier
-        
+
         down_weight = self.down
         up_weight = self.up
-        
+
         x_batch = None
         if down_weight.dim() == 3:
             if x.size(0) != down_weight.size(0):
-                assert self.isconv == False, "Convolutional hypernet with batch size mismatch is not supported"
+                assert (
+                    self.isconv == False
+                ), "Convolutional hypernet with batch size mismatch is not supported"
                 x_batch = x.size(0)
                 x = x.view(down_weight.size(0), -1, *x.shape[1:])
-            
+
             if self.isconv:
-                mid = torch.einsum('ijk, ik... -> ij...', down_weight, x)
+                mid = torch.einsum("ijk, ik... -> ij...", down_weight, x)
             else:
-                mid = torch.einsum('ijk, i...k -> i...j', down_weight, x)
+                mid = torch.einsum("ijk, i...k -> i...j", down_weight, x)
         else:
             if self.isconv:
                 weight = down_weight.unsqueeze(-1).unsqueeze(-1)
             else:
                 weight = down_weight
             mid = self.down_op(x, weight)
-        
+
         if self.rank_dropout and self.training:
-            drop = torch.rand(self.lora_dim, device=mid.device) < self.rank_dropout
-            if (dims:=len(x.shape)) == 4:
+            drop = (
+                torch.rand(self.lora_dim, device=mid.device) < self.rank_dropout
+            ).to(mid.dtype)
+            if self.rank_dropout_scale:
+                drop /= drop.mean()
+            if (dims := len(x.shape)) == 4:
                 drop = drop.view(1, -1, 1, 1)
             else:
-                drop = drop.view(*[1]*(dims-1), -1)
+                drop = drop.view(*[1] * (dims - 1), -1)
             mid = mid * drop
-        
+
         if up_weight.dim() == 3:
             mid_batch = None
             if mid.size(0) != up_weight.size(0):
-                assert self.isconv == False, "Convolutional hypernet with batch size mismatch is not supported"
+                assert (
+                    self.isconv == False
+                ), "Convolutional hypernet with batch size mismatch is not supported"
                 mid_batch = mid.size(0)
                 mid = mid.view(up_weight.size(0), -1, *mid.shape[1:])
-            
+
             if self.isconv:
-                up = torch.einsum('ijk, ik... -> ij...', up_weight, mid)
+                up = torch.einsum("ijk, ik... -> ij...", up_weight, mid)
             else:
-                up = torch.einsum('ijk, i...k -> i...j', up_weight, mid)
-            
+                up = torch.einsum("ijk, i...k -> i...j", up_weight, mid)
+
             if mid_batch is not None:
                 up = up.view(mid_batch, *up.shape[2:])
         else:
@@ -231,10 +241,10 @@ def hypernet_forward(self, x):
             else:
                 weight = up_weight
             up = self.up_op(mid, weight)
-        
+
         if x_batch is not None:
             up = up.view(x_batch, *up.shape[2:])
-        
+
         org_out = self.org_forward(x)
         # print(x.shape, org_out.shape, up.shape)
         return org_out + self.dropout(up) * scale
@@ -248,13 +258,19 @@ def forward(self, x):
             mid = self.lora_mid(self.lora_down(x))
         else:
             mid = self.lora_down(x)
-        
+
         if self.rank_dropout and self.training:
-            drop = torch.rand(self.lora_dim, device=mid.device) < self.rank_dropout
-            if (dims:=len(x.shape)) == 4:
+            drop = (
+                torch.rand(self.lora_dim, device=mid.device) > self.rank_dropout
+            ).to(mid.dtype)
+            if self.rank_dropout_scale:
+                drop /= drop.mean()
+            if (dims := len(x.shape)) == 4:
                 drop = drop.view(1, -1, 1, 1)
             else:
-                drop = drop.view(*[1]*(dims-1), -1)
+                drop = drop.view(*[1] * (dims - 1), -1)
             mid = mid * drop
-        
-        return self.org_forward(x) + self.dropout(self.lora_up(mid) * self.scalar * scale)
+
+        return self.org_forward(x) + self.dropout(
+            self.lora_up(mid) * self.scalar * scale
+        )
diff --git a/lycoris/modules/loha.py b/lycoris/modules/loha.py
index 67a7824..904a4cc 100644
--- a/lycoris/modules/loha.py
+++ b/lycoris/modules/loha.py
@@ -1,30 +1,29 @@
-import math
-from collections import OrderedDict, abc as container_abcs
-
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
+from .base import ModuleCustomSD
+
 
 class HadaWeight(torch.autograd.Function):
     @staticmethod
     def forward(ctx, w1a, w1b, w2a, w2b, scale=torch.tensor(1)):
         ctx.save_for_backward(w1a, w1b, w2a, w2b, scale)
-        diff_weight = ((w1a@w1b)*(w2a@w2b)) * scale
+        diff_weight = ((w1a @ w1b) * (w2a @ w2b)) * scale
         return diff_weight
 
     @staticmethod
     def backward(ctx, grad_out):
         (w1a, w1b, w2a, w2b, scale) = ctx.saved_tensors
         grad_out = grad_out * scale
-        temp = grad_out*(w2a@w2b)
+        temp = grad_out * (w2a @ w2b)
         grad_w1a = temp @ w1b.T
         grad_w1b = w1a.T @ temp
 
-        temp = grad_out * (w1a@w1b)
+        temp = grad_out * (w1a @ w1b)
         grad_w2a = temp @ w2b.T
         grad_w2b = w2a.T @ temp
-        
+
         del temp
         return grad_w1a, grad_w1b, grad_w2a, grad_w2b, None
 
@@ -33,43 +32,43 @@ class HadaWeightCP(torch.autograd.Function):
     @staticmethod
     def forward(ctx, t1, w1a, w1b, t2, w2a, w2b, scale=torch.tensor(1)):
         ctx.save_for_backward(t1, w1a, w1b, t2, w2a, w2b, scale)
-        
-        rebuild1 = torch.einsum('i j k l, j r, i p -> p r k l', t1, w1b, w1a)
-        rebuild2 = torch.einsum('i j k l, j r, i p -> p r k l', t2, w2b, w2a)
-        
-        return rebuild1*rebuild2*scale
+
+        rebuild1 = torch.einsum("i j k l, j r, i p -> p r k l", t1, w1b, w1a)
+        rebuild2 = torch.einsum("i j k l, j r, i p -> p r k l", t2, w2b, w2a)
+
+        return rebuild1 * rebuild2 * scale
 
     @staticmethod
     def backward(ctx, grad_out):
         (t1, w1a, w1b, t2, w2a, w2b, scale) = ctx.saved_tensors
         grad_out = grad_out * scale
-        
-        temp = torch.einsum('i j k l, j r -> i r k l', t2, w2b)
-        rebuild = torch.einsum('i j k l, i r -> r j k l', temp, w2a)
-        
-        grad_w = rebuild*grad_out
+
+        temp = torch.einsum("i j k l, j r -> i r k l", t2, w2b)
+        rebuild = torch.einsum("i j k l, i r -> r j k l", temp, w2a)
+
+        grad_w = rebuild * grad_out
         del rebuild
-        
-        grad_w1a = torch.einsum('r j k l, i j k l -> r i', temp, grad_w)
-        grad_temp = torch.einsum('i j k l, i r -> r j k l', grad_w, w1a.T)
+
+        grad_w1a = torch.einsum("r j k l, i j k l -> r i", temp, grad_w)
+        grad_temp = torch.einsum("i j k l, i r -> r j k l", grad_w, w1a.T)
         del grad_w, temp
-        
-        grad_w1b = torch.einsum('i r k l, i j k l -> r j', t1, grad_temp)
-        grad_t1 = torch.einsum('i j k l, j r -> i r k l', grad_temp, w1b.T)
+
+        grad_w1b = torch.einsum("i r k l, i j k l -> r j", t1, grad_temp)
+        grad_t1 = torch.einsum("i j k l, j r -> i r k l", grad_temp, w1b.T)
         del grad_temp
-        
-        temp = torch.einsum('i j k l, j r -> i r k l', t1, w1b)
-        rebuild = torch.einsum('i j k l, i r -> r j k l', temp, w1a)
-        
-        grad_w = rebuild*grad_out
+
+        temp = torch.einsum("i j k l, j r -> i r k l", t1, w1b)
+        rebuild = torch.einsum("i j k l, i r -> r j k l", temp, w1a)
+
+        grad_w = rebuild * grad_out
         del rebuild
-        
-        grad_w2a = torch.einsum('r j k l, i j k l -> r i', temp, grad_w)
-        grad_temp = torch.einsum('i j k l, i r -> r j k l', grad_w, w2a.T)
+
+        grad_w2a = torch.einsum("r j k l, i j k l -> r i", temp, grad_w)
+        grad_temp = torch.einsum("i j k l, i r -> r j k l", grad_w, w2a.T)
         del grad_w, temp
-        
-        grad_w2b = torch.einsum('i r k l, i j k l -> r j', t2, grad_temp)
-        grad_t2 = torch.einsum('i j k l, j r -> i r k l', grad_temp, w2b.T)
+
+        grad_w2b = torch.einsum("i r k l, i j k l -> r j", t2, grad_temp)
+        grad_t2 = torch.einsum("i j k l, j r -> i r k l", grad_temp, w2b.T)
         del grad_temp
         return grad_t1, grad_w1a, grad_w1b, grad_t2, grad_w2a, grad_w2b, None
 
@@ -82,43 +81,48 @@ def make_weight_cp(t1, w1a, w1b, t2, w2a, w2b, scale):
     return HadaWeightCP.apply(t1, w1a, w1b, t2, w2a, w2b, scale)
 
 
-class LohaModule(nn.Module):
+class LohaModule(ModuleCustomSD):
     """
     Hadamard product Implementaion for Low Rank Adaptation
     """
 
     def __init__(
-        self, 
-        lora_name, 
-        org_module: nn.Module, 
-        multiplier=1.0, lora_dim=4, alpha=1, 
-        dropout=0., rank_dropout=0., module_dropout=0.,
+        self,
+        lora_name,
+        org_module: nn.Module,
+        multiplier=1.0,
+        lora_dim=4,
+        alpha=1,
+        dropout=0.0,
+        rank_dropout=0.0,
+        module_dropout=0.0,
         use_tucker=False,
         use_scalar=False,
+        rank_dropout_scale=False,
         **kwargs,
     ):
-        """ if alpha == 0 or None, alpha is rank (no scaling). """
+        """if alpha == 0 or None, alpha is rank (no scaling)."""
         super().__init__()
         self.lora_name = lora_name
         self.lora_dim = lora_dim
-        self.tucker=False
-        
+        self.tucker = False
+
         self.shape = org_module.weight.shape
-        if org_module.__class__.__name__ == 'Conv2d':
+        if org_module.__class__.__name__ == "Conv2d":
             in_dim = org_module.in_channels
             k_size = org_module.kernel_size
             out_dim = org_module.out_channels
-            self.tucker = use_tucker and k_size!=(1, 1)
+            self.tucker = use_tucker and k_size != (1, 1)
             if self.tucker:
                 shape = (out_dim, in_dim, *k_size)
             else:
-                shape = (out_dim, in_dim*k_size[0]*k_size[1])
+                shape = (out_dim, in_dim * k_size[0] * k_size[1])
             self.op = F.conv2d
             self.extra_args = {
                 "stride": org_module.stride,
                 "padding": org_module.padding,
                 "dilation": org_module.dilation,
-                "groups": org_module.groups
+                "groups": org_module.groups,
             }
         else:
             in_dim = org_module.in_features
@@ -126,33 +130,46 @@ def __init__(
             shape = (out_dim, in_dim)
             self.op = F.linear
             self.extra_args = {}
-        
+
         if self.tucker:
-            self.hada_t1 = nn.Parameter(torch.empty(lora_dim, lora_dim, shape[2], shape[3]))
-            self.hada_w1_a = nn.Parameter(torch.empty(lora_dim, shape[0])) # out_dim, 1-mode
-            self.hada_w1_b = nn.Parameter(torch.empty(lora_dim, shape[1])) # in_dim , 2-mode
-            
-            self.hada_t2 = nn.Parameter(torch.empty(lora_dim, lora_dim, shape[2], shape[3]))
-            self.hada_w2_a = nn.Parameter(torch.empty(lora_dim, shape[0])) # out_dim, 1-mode
-            self.hada_w2_b = nn.Parameter(torch.empty(lora_dim, shape[1])) # in_dim , 2-mode
+            self.hada_t1 = nn.Parameter(
+                torch.empty(lora_dim, lora_dim, shape[2], shape[3])
+            )
+            self.hada_w1_a = nn.Parameter(
+                torch.empty(lora_dim, shape[0])
+            )  # out_dim, 1-mode
+            self.hada_w1_b = nn.Parameter(
+                torch.empty(lora_dim, shape[1])
+            )  # in_dim , 2-mode
+
+            self.hada_t2 = nn.Parameter(
+                torch.empty(lora_dim, lora_dim, shape[2], shape[3])
+            )
+            self.hada_w2_a = nn.Parameter(
+                torch.empty(lora_dim, shape[0])
+            )  # out_dim, 1-mode
+            self.hada_w2_b = nn.Parameter(
+                torch.empty(lora_dim, shape[1])
+            )  # in_dim , 2-mode
         else:
             self.hada_w1_a = nn.Parameter(torch.empty(shape[0], lora_dim))
             self.hada_w1_b = nn.Parameter(torch.empty(lora_dim, shape[1]))
-            
+
             self.hada_w2_a = nn.Parameter(torch.empty(shape[0], lora_dim))
             self.hada_w2_b = nn.Parameter(torch.empty(lora_dim, shape[1]))
-        
+
         self.dropout = dropout
         if dropout:
             print("[WARN]LoHa/LoKr haven't implemented normal dropout yet.")
         self.rank_dropout = rank_dropout
+        self.rank_dropout_scale = rank_dropout_scale
         self.module_dropout = module_dropout
-        
+
         if type(alpha) == torch.Tensor:
             alpha = alpha.detach().float().numpy()  # without casting, bf16 causes error
         alpha = lora_dim if alpha is None or alpha == 0 else alpha
         self.scale = alpha / self.lora_dim
-        self.register_buffer('alpha', torch.tensor(alpha)) # 定数として扱える
+        self.register_buffer("alpha", torch.tensor(alpha))  # 定数として扱える
 
         if use_scalar:
             self.scalar = nn.Parameter(torch.tensor(0.0))
@@ -171,10 +188,10 @@ def __init__(
             torch.nn.init.constant_(self.hada_w2_a, 0)
 
         self.multiplier = multiplier
-        self.org_module = [org_module] # remove in applying
+        self.org_module = [org_module]  # remove in applying
         self.grad_ckpt = False
         self.register_load_state_dict_post_hook(self.load_weight_hook)
-    
+
     def load_weight_hook(self, *args, **kwargs):
         self.scalar = nn.Parameter(torch.ones_like(self.scalar))
 
@@ -184,64 +201,56 @@ def apply_to(self):
     def get_weight(self, orig_weight=None):
         if self.tucker:
             weight = make_weight_cp(
-                self.hada_t1, self.hada_w1_a, self.hada_w1_b,
-                self.hada_t1, self.hada_w2_a, self.hada_w2_b,
-                scale = torch.tensor(self.scale),
+                self.hada_t1,
+                self.hada_w1_a,
+                self.hada_w1_b,
+                self.hada_t2,
+                self.hada_w2_a,
+                self.hada_w2_b,
+                scale=torch.tensor(self.scale),
             )
         else:
             weight = make_weight(
-                self.hada_w1_a, self.hada_w1_b,
-                self.hada_w2_a, self.hada_w2_b,
-                scale = torch.tensor(self.scale),
+                self.hada_w1_a,
+                self.hada_w1_b,
+                self.hada_w2_a,
+                self.hada_w2_b,
+                scale=torch.tensor(self.scale),
             )
         if orig_weight is not None:
             weight = weight.reshape(orig_weight.shape)
         if self.training and self.rank_dropout:
-            drop = torch.rand(weight.size(0)) < self.rank_dropout
-            weight *= drop.view(-1, [1]*len(weight.shape[1:])).to(weight.device)
+            drop = (torch.rand(weight.size(0)) > self.rank_dropout).to(weight.dtype)
+            drop = drop.view(-1, *[1] * len(weight.shape[1:])).to(weight.device)
+            if self.rank_dropout_scale:
+                drop /= drop.mean()
+            weight *= drop
         return weight
-    
-    def state_dict(self, *args, destination=None, prefix='', keep_vars=False):
-        # TODO: Remove `args` and the parsing logic when BC allows.
-        if len(args) > 0:
-            if destination is None:
-                destination = args[0]
-            if len(args) > 1 and prefix == '':
-                prefix = args[1]
-            if len(args) > 2 and keep_vars is False:
-                keep_vars = args[2]
-            # DeprecationWarning is ignored by default
-
-        if destination is None:
-            destination = OrderedDict()
-            destination._metadata = OrderedDict()
-
-        local_metadata = dict(version=self._version)
-        if hasattr(destination, "_metadata"):
-            destination._metadata[prefix[:-1]] = local_metadata
-
-        destination[f'{prefix}alpha'] = self.alpha
-        destination[f'{prefix}hada_w1_a'] = self.hada_w1_a * self.scalar
-        destination[f'{prefix}hada_w1_b'] = self.hada_w1_b
-        destination[f'{prefix}hada_w2_a'] = self.hada_w2_a
-        destination[f'{prefix}hada_w2_b'] = self.hada_w2_b
+
+    def custom_state_dict(self):
+        destination = {}
+        destination["alpha"] = self.alpha
+        destination["hada_w1_a"] = self.hada_w1_a * self.scalar
+        destination["hada_w1_b"] = self.hada_w1_b
+        destination["hada_w2_a"] = self.hada_w2_a
+        destination["hada_w2_b"] = self.hada_w2_b
         if self.tucker:
-            destination[f'{prefix}hada_t1'] = self.hada_t1
-            destination[f'{prefix}hada_t2'] = self.hada_t2
+            destination["hada_t1"] = self.hada_t1
+            destination["hada_t2"] = self.hada_t2
         return destination
 
     @torch.no_grad()
     def apply_max_norm(self, max_norm, device=None):
-        orig_norm = (self.get_weight()*self.scalar).norm()
-        norm = torch.clamp(orig_norm, max_norm/2)
+        orig_norm = (self.get_weight() * self.scalar).norm()
+        norm = torch.clamp(orig_norm, max_norm / 2)
         desired = torch.clamp(norm, max=max_norm)
-        ratio = desired.cpu()/norm.cpu()
-        
+        ratio = desired.cpu() / norm.cpu()
+
         scaled = ratio.item() != 1.0
         if scaled:
             self.scalar *= ratio
-        
-        return scaled, orig_norm*ratio
+
+        return scaled, orig_norm * ratio
 
     @torch.enable_grad()
     def forward(self, x):
@@ -250,16 +259,15 @@ def forward(self, x):
                 return self.op(
                     x,
                     self.org_module[0].weight.data,
-                    None if self.org_module[0].bias is None else self.org_module[0].bias.data
+                    None
+                    if self.org_module[0].bias is None
+                    else self.org_module[0].bias.data,
                 )
         weight = (
             self.org_module[0].weight.data.to(x.device, dtype=self.hada_w1_a.dtype)
-            + self.get_weight(self.org_module[0].weight.data) * self.scalar * self.multiplier
+            + self.get_weight(self.org_module[0].weight.data)
+            * self.scalar
+            * self.multiplier
         )
         bias = None if self.org_module[0].bias is None else self.org_module[0].bias.data
-        return self.op(
-            x, 
-            weight.view(self.shape),
-            bias,
-            **self.extra_args
-        )
+        return self.op(x, weight.view(self.shape), bias, **self.extra_args)
diff --git a/lycoris/modules/lokr.py b/lycoris/modules/lokr.py
index 36f21d2..dbe07d7 100644
--- a/lycoris/modules/lokr.py
+++ b/lycoris/modules/lokr.py
@@ -1,24 +1,22 @@
 import math
-from collections import OrderedDict
 
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
-# 4, build custom backward function
-#  - 
+from .base import ModuleCustomSD
 
 
-def factorization(dimension: int, factor:int=-1) -> tuple[int, int]:
-    '''
+def factorization(dimension: int, factor: int = -1) -> tuple[int, int]:
+    """
     return a tuple of two value of input dimension decomposed by the number closest to factor
     second value is higher or equal than first value.
-    
+
     In LoRA with Kroneckor Product, first value is a value for weight scale.
     secon value is a value for weight.
-    
+
     Becuase of non-commutative property, A⊗B ≠ B⊗A. Meaning of two matrices is slightly different.
-    
+
     examples)
     factor
         -1               2                4               8               16               ...
@@ -28,8 +26,8 @@ def factorization(dimension: int, factor:int=-1) -> tuple[int, int]:
     360 -> 8, 45    360 -> 2, 180    360 -> 4, 90    360 -> 8, 45    360 -> 12, 30
     512 -> 16, 32   512 -> 2, 256    512 -> 4, 128   512 -> 8, 64    512 -> 16, 32
     1024 -> 32, 32  1024 -> 2, 512   1024 -> 4, 256  1024 -> 8, 128  1024 -> 16, 64
-    '''
-    
+    """
+
     if factor > 0 and (dimension % factor) == 0:
         m = factor
         n = dimension // factor
@@ -40,12 +38,12 @@ def factorization(dimension: int, factor:int=-1) -> tuple[int, int]:
         factor = dimension
     m, n = 1, dimension
     length = m + n
-    while m<n:
+    while m < n:
         new_m = m + 1
-        while dimension%new_m != 0:
+        while dimension % new_m != 0:
             new_m += 1
         new_n = dimension // new_m
-        if new_m + new_n > length or new_m>factor:
+        if new_m + new_n > length or new_m > factor:
             break
         else:
             m, n = new_m, new_n
@@ -55,7 +53,7 @@ def factorization(dimension: int, factor:int=-1) -> tuple[int, int]:
 
 
 def make_weight_cp(t, wa, wb):
-    rebuild2 = torch.einsum('i j k l, i p, j r -> p r k l', t, wa, wb) # [c, d, k1, k2]
+    rebuild2 = torch.einsum("i j k l, i p, j r -> p r k l", t, wa, wb)  # [c, d, k1, k2]
     return rebuild2
 
 
@@ -64,11 +62,11 @@ def make_kron(w1, w2, scale):
         w1 = w1.unsqueeze(2).unsqueeze(2)
     w2 = w2.contiguous()
     rebuild = torch.kron(w1, w2)
-    
-    return rebuild*scale
+
+    return rebuild * scale
 
 
-class LokrModule(nn.Module):
+class LokrModule(ModuleCustomSD):
     """
     modifed from kohya-ss/sd-scripts/networks/lora:LoRAModule
         and from KohakuBlueleaf/LyCORIS/lycoris:loha:LoHaModule
@@ -76,18 +74,23 @@ class LokrModule(nn.Module):
     """
 
     def __init__(
-        self, 
-        lora_name, org_module: nn.Module, 
-        multiplier=1.0, 
-        lora_dim=4, alpha=1, 
-        dropout=0., rank_dropout=0., module_dropout=0.,
+        self,
+        lora_name,
+        org_module: nn.Module,
+        multiplier=1.0,
+        lora_dim=4,
+        alpha=1,
+        dropout=0.0,
+        rank_dropout=0.0,
+        module_dropout=0.0,
         use_cp=False,
         use_scalar=False,
-        decompose_both = False,
-        factor:int=-1, # factorization factor
+        decompose_both=False,
+        factor: int = -1,  # factorization factor
+        rank_dropout_scale=False,
         **kwargs,
     ):
-        """ if alpha == 0 or None, alpha is rank (no scaling). """
+        """if alpha == 0 or None, alpha is rank (no scaling)."""
         super().__init__()
         factor = int(factor)
         self.lora_name = lora_name
@@ -97,58 +100,75 @@ def __init__(
         self.use_w2 = False
 
         self.shape = org_module.weight.shape
-        if org_module.__class__.__name__ == 'Conv2d':
+        if org_module.__class__.__name__ == "Conv2d":
             in_dim = org_module.in_channels
             k_size = org_module.kernel_size
             out_dim = org_module.out_channels
-            
+
             in_m, in_n = factorization(in_dim, factor)
             out_l, out_k = factorization(out_dim, factor)
-            shape = ((out_l, out_k), (in_m, in_n), *k_size) # ((a, b), (c, d), *k_size)
-            self.cp = use_cp and k_size!=(1, 1)
-            if decompose_both and lora_dim < max(shape[0][0], shape[1][0])/2:
+            shape = ((out_l, out_k), (in_m, in_n), *k_size)  # ((a, b), (c, d), *k_size)
+            self.cp = use_cp and k_size != (1, 1)
+            if decompose_both and lora_dim < max(shape[0][0], shape[1][0]) / 2:
                 self.lokr_w1_a = nn.Parameter(torch.empty(shape[0][0], lora_dim))
                 self.lokr_w1_b = nn.Parameter(torch.empty(lora_dim, shape[1][0]))
             else:
                 self.use_w1 = True
-                self.lokr_w1 = nn.Parameter(torch.empty(shape[0][0], shape[1][0]))  # a*c, 1-mode
-            
-            if lora_dim >= max(shape[0][1], shape[1][1])/2:
+                self.lokr_w1 = nn.Parameter(
+                    torch.empty(shape[0][0], shape[1][0])
+                )  # a*c, 1-mode
+
+            if lora_dim >= max(shape[0][1], shape[1][1]) / 2:
                 self.use_w2 = True
-                self.lokr_w2 = nn.Parameter(torch.empty(shape[0][1], shape[1][1], *k_size))
+                self.lokr_w2 = nn.Parameter(
+                    torch.empty(shape[0][1], shape[1][1], *k_size)
+                )
             elif self.cp:
-                self.lokr_t2 = nn.Parameter(torch.empty(lora_dim, lora_dim, shape[2], shape[3]))
-                self.lokr_w2_a = nn.Parameter(torch.empty(lora_dim, shape[0][1])) # b, 1-mode
-                self.lokr_w2_b = nn.Parameter(torch.empty(lora_dim, shape[1][1])) # d, 2-mode
-            else: # Conv2d not cp
+                self.lokr_t2 = nn.Parameter(
+                    torch.empty(lora_dim, lora_dim, shape[2], shape[3])
+                )
+                self.lokr_w2_a = nn.Parameter(
+                    torch.empty(lora_dim, shape[0][1])
+                )  # b, 1-mode
+                self.lokr_w2_b = nn.Parameter(
+                    torch.empty(lora_dim, shape[1][1])
+                )  # d, 2-mode
+            else:  # Conv2d not cp
                 # bigger part. weight and LoRA. [b, dim] x [dim, d*k1*k2]
                 self.lokr_w2_a = nn.Parameter(torch.empty(shape[0][1], lora_dim))
-                self.lokr_w2_b = nn.Parameter(torch.empty(lora_dim, shape[1][1]*shape[2]*shape[3]))
+                self.lokr_w2_b = nn.Parameter(
+                    torch.empty(lora_dim, shape[1][1] * shape[2] * shape[3])
+                )
                 # w1 ⊗ (w2_a x w2_b) = (a, b)⊗((c, dim)x(dim, d*k1*k2)) = (a, b)⊗(c, d*k1*k2) = (ac, bd*k1*k2)
-            
+
             self.op = F.conv2d
             self.extra_args = {
                 "stride": org_module.stride,
                 "padding": org_module.padding,
                 "dilation": org_module.dilation,
-                "groups": org_module.groups
+                "groups": org_module.groups,
             }
 
-        else: # Linear
+        else:  # Linear
             in_dim = org_module.in_features
             out_dim = org_module.out_features
-            
+
             in_m, in_n = factorization(in_dim, factor)
             out_l, out_k = factorization(out_dim, factor)
-            shape = ((out_l, out_k), (in_m, in_n)) # ((a, b), (c, d)), out_dim = a*c, in_dim = b*d
+            shape = (
+                (out_l, out_k),
+                (in_m, in_n),
+            )  # ((a, b), (c, d)), out_dim = a*c, in_dim = b*d
             # smaller part. weight scale
-            if decompose_both and lora_dim < max(shape[0][0], shape[1][0])/2:
+            if decompose_both and lora_dim < max(shape[0][0], shape[1][0]) / 2:
                 self.lokr_w1_a = nn.Parameter(torch.empty(shape[0][0], lora_dim))
                 self.lokr_w1_b = nn.Parameter(torch.empty(lora_dim, shape[1][0]))
             else:
                 self.use_w1 = True
-                self.lokr_w1 = nn.Parameter(torch.empty(shape[0][0], shape[1][0]))  # a*c, 1-mode
-            if lora_dim < max(shape[0][1], shape[1][1])/2:
+                self.lokr_w1 = nn.Parameter(
+                    torch.empty(shape[0][0], shape[1][0])
+                )  # a*c, 1-mode
+            if lora_dim < max(shape[0][1], shape[1][1]) / 2:
                 # bigger part. weight and LoRA. [b, dim] x [dim, d]
                 self.lokr_w2_a = nn.Parameter(torch.empty(shape[0][1], lora_dim))
                 self.lokr_w2_b = nn.Parameter(torch.empty(lora_dim, shape[1][1]))
@@ -159,27 +179,28 @@ def __init__(
 
             self.op = F.linear
             self.extra_args = {}
-        
+
         self.dropout = dropout
         if dropout:
             print("[WARN]LoHa/LoKr haven't implemented normal dropout yet.")
         self.rank_dropout = rank_dropout
+        self.rank_dropout_scale = rank_dropout_scale
         self.module_dropout = module_dropout
-        
+
         if isinstance(alpha, torch.Tensor):
             alpha = alpha.detach().float().numpy()  # without casting, bf16 causes error
         alpha = lora_dim if alpha is None or alpha == 0 else alpha
         if self.use_w2 and self.use_w1:
-            #use scale = 1
+            # use scale = 1
             alpha = lora_dim
         self.scale = alpha / self.lora_dim
-        self.register_buffer('alpha', torch.tensor(alpha)) # 定数として扱える
+        self.register_buffer("alpha", torch.tensor(alpha))  # 定数として扱える
 
         if use_scalar:
             self.scalar = nn.Parameter(torch.tensor(0.0))
         else:
             self.scalar = torch.tensor(1.0)
-        
+
         if self.use_w2:
             if use_scalar:
                 torch.nn.init.kaiming_uniform_(self.lokr_w2, a=math.sqrt(5))
@@ -193,7 +214,7 @@ def __init__(
                 torch.nn.init.kaiming_uniform_(self.lokr_w2_b, a=math.sqrt(5))
             else:
                 torch.nn.init.constant_(self.lokr_w2_b, 0)
-        
+
         if self.use_w1:
             torch.nn.init.kaiming_uniform_(self.lokr_w1, a=math.sqrt(5))
         else:
@@ -203,15 +224,19 @@ def __init__(
         self.multiplier = multiplier
         self.org_module = [org_module]
         weight = make_kron(
-            self.lokr_w1 if self.use_w1 else self.lokr_w1_a@self.lokr_w1_b,
-            (self.lokr_w2 if self.use_w2 
-             else make_weight_cp(self.lokr_t2, self.lokr_w2_a, self.lokr_w2_b) if self.cp 
-             else self.lokr_w2_a@self.lokr_w2_b),
-            torch.tensor(self.multiplier * self.scale)
+            self.lokr_w1 if self.use_w1 else self.lokr_w1_a @ self.lokr_w1_b,
+            (
+                self.lokr_w2
+                if self.use_w2
+                else make_weight_cp(self.lokr_t2, self.lokr_w2_a, self.lokr_w2_b)
+                if self.cp
+                else self.lokr_w2_a @ self.lokr_w2_b
+            ),
+            torch.tensor(self.multiplier * self.scale),
         )
         assert torch.sum(torch.isnan(weight)) == 0, "weight is nan"
         self.register_load_state_dict_post_hook(self.load_weight_hook)
-    
+
     def load_weight_hook(self, *args, **kwargs):
         self.scalar = nn.Parameter(torch.ones_like(self.scalar))
 
@@ -219,82 +244,72 @@ def load_weight_hook(self, *args, **kwargs):
     def apply_to(self):
         self.org_forward = self.org_module[0].forward
         self.org_module[0].forward = self.forward
-    
-    def get_weight(self, orig_weight = None):
+
+    def get_weight(self, orig_weight=None):
         weight = make_kron(
-            self.lokr_w1 if self.use_w1 else self.lokr_w1_a@self.lokr_w1_b,
-            (self.lokr_w2 if self.use_w2 
-             else make_weight_cp(self.lokr_t2, self.lokr_w2_a, self.lokr_w2_b) if self.cp 
-             else self.lokr_w2_a@self.lokr_w2_b),
-            torch.tensor(self.scale)
+            self.lokr_w1 if self.use_w1 else self.lokr_w1_a @ self.lokr_w1_b,
+            (
+                self.lokr_w2
+                if self.use_w2
+                else make_weight_cp(self.lokr_t2, self.lokr_w2_a, self.lokr_w2_b)
+                if self.cp
+                else self.lokr_w2_a @ self.lokr_w2_b
+            ),
+            torch.tensor(self.scale),
         )
         if orig_weight is not None:
             weight = weight.reshape(orig_weight.shape)
         if self.training and self.rank_dropout:
-            drop = torch.rand(weight.size(0)) < self.rank_dropout
-            weight *= drop.view(-1, [1]*len(weight.shape[1:])).to(weight.device)
+            drop = (torch.rand(weight.size(0)) > self.rank_dropout).to(weight.dtype)
+            drop = drop.view(-1, *[1] * len(weight.shape[1:])).to(weight.device)
+            if self.rank_dropout_scale:
+                drop /= drop.mean()
+            weight *= drop
         return weight
-    
-    def state_dict(self, *args, destination=None, prefix='', keep_vars=False):
-        # TODO: Remove `args` and the parsing logic when BC allows.
-        if len(args) > 0:
-            if destination is None:
-                destination = args[0]
-            if len(args) > 1 and prefix == '':
-                prefix = args[1]
-            if len(args) > 2 and keep_vars is False:
-                keep_vars = args[2]
-            # DeprecationWarning is ignored by default
-
-        if destination is None:
-            destination = OrderedDict()
-            destination._metadata = OrderedDict()
-
-        local_metadata = dict(version=self._version)
-        if hasattr(destination, "_metadata"):
-            destination._metadata[prefix[:-1]] = local_metadata
-
-        destination[f'{prefix}alpha'] = self.alpha
+
+    def custom_state_dict(self):
+        destination = {}
+        destination["alpha"] = self.alpha
         if self.use_w1:
-            destination[f'{prefix}lokr_w1'] = self.lokr_w1 * self.scalar
+            destination["lokr_w1"] = self.lokr_w1 * self.scalar
         else:
-            destination[f'{prefix}lokr_w1_a'] = self.lokr_w1_a * self.scalar
-            destination[f'{prefix}lokr_w1_b'] = self.lokr_w1_b
-        
+            destination["lokr_w1_a"] = self.lokr_w1_a * self.scalar
+            destination["lokr_w1_b"] = self.lokr_w1_b
+
         if self.use_w2:
-            destination[f'{prefix}lokr_w2'] = self.lokr_w2
+            destination["lokr_w2"] = self.lokr_w2
         else:
-            destination[f'{prefix}lokr_w2_a'] = self.lokr_w2_a
-            destination[f'{prefix}lokr_w2_b'] = self.lokr_w2_b
+            destination["lokr_w2_a"] = self.lokr_w2_a
+            destination["lokr_w2_b"] = self.lokr_w2_b
             if self.cp:
-                destination[f'{prefix}lokr_t2'] = self.lokr_t2
+                destination["lokr_t2"] = self.lokr_t2
         return destination
 
     @torch.no_grad()
     def apply_max_norm(self, max_norm, device=None):
         orig_norm = self.get_weight().norm()
-        norm = torch.clamp(orig_norm, max_norm/2)
+        norm = torch.clamp(orig_norm, max_norm / 2)
         desired = torch.clamp(norm, max=max_norm)
-        ratio = desired.cpu()/norm.cpu()
-        
+        ratio = desired.cpu() / norm.cpu()
+
         scaled = ratio.item() != 1.0
         if scaled:
-            modules = (4 - self.use_w1 - self.use_w2 + (not self.use_w2 and self.cp))
+            modules = 4 - self.use_w1 - self.use_w2 + (not self.use_w2 and self.cp)
             if self.use_w1:
-                self.lokr_w1 *= ratio**(1/modules)
+                self.lokr_w1 *= ratio ** (1 / modules)
             else:
-                self.lokr_w1_a *= ratio**(1/modules)
-                self.lokr_w1_b *= ratio**(1/modules)
-            
+                self.lokr_w1_a *= ratio ** (1 / modules)
+                self.lokr_w1_b *= ratio ** (1 / modules)
+
             if self.use_w2:
-                self.lokr_w2 *= ratio**(1/modules)
+                self.lokr_w2 *= ratio ** (1 / modules)
             else:
                 if self.cp:
-                    self.lokr_t2 *= ratio**(1/modules)
-                self.lokr_w2_a  *= ratio**(1/modules)
-                self.lokr_w2_b  *= ratio**(1/modules)
-        
-        return scaled, orig_norm*ratio
+                    self.lokr_t2 *= ratio ** (1 / modules)
+                self.lokr_w2_a *= ratio ** (1 / modules)
+                self.lokr_w2_b *= ratio ** (1 / modules)
+
+        return scaled, orig_norm * ratio
 
     def forward(self, x):
         if self.module_dropout and self.training:
@@ -302,16 +317,15 @@ def forward(self, x):
                 return self.op(
                     x,
                     self.org_module[0].weight.data,
-                    None if self.org_module[0].bias is None else self.org_module[0].bias.data
+                    None
+                    if self.org_module[0].bias is None
+                    else self.org_module[0].bias.data,
                 )
         weight = (
             self.org_module[0].weight.data.to(x.device, dtype=self.lokr_w1.dtype)
-            + self.get_weight(self.org_module[0].weight.data) * self.scalar * self.multiplier
+            + self.get_weight(self.org_module[0].weight.data)
+            * self.scalar
+            * self.multiplier
         )
         bias = None if self.org_module[0].bias is None else self.org_module[0].bias.data
-        return self.op(
-            x, 
-            weight.view(self.shape),
-            bias,
-            **self.extra_args
-        )
+        return self.op(x, weight.view(self.shape), bias, **self.extra_args)
diff --git a/lycoris/modules/norms.py b/lycoris/modules/norms.py
index 6745451..6b93df8 100644
--- a/lycoris/modules/norms.py
+++ b/lycoris/modules/norms.py
@@ -6,26 +6,35 @@
 import torch.nn as nn
 import torch.nn.functional as F
 
+from .base import ModuleCustomSD
 
-class NormModule(nn.Module):
+
+class NormModule(ModuleCustomSD):
     """
     modifed from kohya-ss/sd-scripts/networks/lora:LoRAModule
     """
 
     def __init__(
-        self, 
-        lora_name, org_module: nn.Module, 
-        multiplier=1.0, rank_dropout=0., module_dropout=0.,
+        self,
+        lora_name,
+        org_module: nn.Module,
+        multiplier=1.0,
+        rank_dropout=0.0,
+        module_dropout=0.0,
+        rank_dropout_scale=False,
         **kwargs,
     ):
-        """ if alpha == 0 or None, alpha is rank (no scaling). """
+        """if alpha == 0 or None, alpha is rank (no scaling)."""
         super().__init__()
         self.lora_name = lora_name
 
         if isinstance(org_module, nn.LayerNorm):
             self.op = F.layer_norm
             self.dim = org_module.normalized_shape[0]
-            self.kw_dict = {"normalized_shape": org_module.normalized_shape, "eps": org_module.eps}
+            self.kw_dict = {
+                "normalized_shape": org_module.normalized_shape,
+                "eps": org_module.eps,
+            }
         elif isinstance(org_module, nn.GroupNorm):
             self.op = F.group_norm
             self.group_num = org_module.num_groups
@@ -35,10 +44,11 @@ def __init__(
             raise NotImplementedError
         self.w_norm = nn.Parameter(torch.zeros(self.dim))
         self.b_norm = nn.Parameter(torch.zeros(self.dim))
-        
+
         self.rank_dropout = rank_dropout
+        self.rank_dropout_scale = rank_dropout_scale
         self.module_dropout = module_dropout
-        
+
         self.multiplier = multiplier
         self.org_module = [org_module]
 
@@ -46,12 +56,20 @@ def apply_to(self, **kwargs):
         self.org_forward = self.org_module[0].forward
         self.org_module[0].forward = self.forward
 
-    def make_weight(self, scale = 1, device=None):
+    def make_weight(self, scale=1, device=None):
         org_weight = self.org_module[0].weight.to(device, dtype=self.w_norm.dtype)
         org_bias = self.org_module[0].bias.to(device, dtype=self.b_norm.dtype)
+        if self.rank_dropout and self.training:
+            drop = (torch.rand(self.dim, device=device) < self.rank_dropout).to(
+                self.w_norm.device
+            )
+            if self.rank_dropout_scale:
+                drop /= drop.mean()
+        else:
+            drop = 1
         drop = (
-            torch.rand(self.dim, device=device) < self.rank_dropout 
-            if self.rank_dropout and self.training 
+            torch.rand(self.dim, device=device) < self.rank_dropout
+            if self.rank_dropout and self.training
             else 1
         )
         weight = self.w_norm.to(device) * drop * scale
@@ -63,7 +81,7 @@ def forward(self, x):
             if torch.rand(1) < self.module_dropout:
                 return self.org_forward(x)
         scale = self.multiplier
-        
+
         w, b = self.make_weight(scale, x.device)
         kw_dict = self.kw_dict | {"weight": w, "bias": b}
         return self.op(x, **kw_dict)
diff --git a/lycoris/utils/__init__.py b/lycoris/utils/__init__.py
index 87ff778..7e5ef33 100644
--- a/lycoris/utils/__init__.py
+++ b/lycoris/utils/__init__.py
@@ -11,6 +11,10 @@
 from tqdm import tqdm
 
 
+def str_bool(val):
+    return str(val).lower() != "false"
+
+
 def default(val, d):
     return val if val is not None else d
 
@@ -25,110 +29,120 @@ def make_sparse(t: torch.Tensor, sparsity=0.95):
 
 def extract_conv(
     weight: Union[torch.Tensor, nn.Parameter],
-    mode = 'fixed',
-    mode_param = 0,
-    device = 'cpu',
-    is_cp = False,
+    mode="fixed",
+    mode_param=0,
+    device="cpu",
+    is_cp=False,
 ) -> Tuple[nn.Parameter, nn.Parameter]:
     weight = weight.to(device)
     out_ch, in_ch, kernel_size, _ = weight.shape
-    
+
     U, S, Vh = linalg.svd(weight.reshape(out_ch, -1))
-    
-    if mode=='fixed':
+
+    if mode == "full":
+        return weight, "full"
+    elif mode == "fixed":
         lora_rank = mode_param
-    elif mode=='threshold':
-        assert mode_param>=0
-        lora_rank = torch.sum(S>mode_param)
-    elif mode=='ratio':
-        assert 1>=mode_param>=0
-        min_s = torch.max(S)*mode_param
-        lora_rank = torch.sum(S>min_s)
-    elif mode=='quantile' or mode=='percentile':
-        assert 1>=mode_param>=0
+    elif mode == "threshold":
+        assert mode_param >= 0
+        lora_rank = torch.sum(S > mode_param)
+    elif mode == "ratio":
+        assert 1 >= mode_param >= 0
+        min_s = torch.max(S) * mode_param
+        lora_rank = torch.sum(S > min_s)
+    elif mode == "quantile" or mode == "percentile":
+        assert 1 >= mode_param >= 0
         s_cum = torch.cumsum(S, dim=0)
         min_cum_sum = mode_param * torch.sum(S)
-        lora_rank = torch.sum(s_cum<min_cum_sum)
+        lora_rank = torch.sum(s_cum < min_cum_sum)
     else:
-        raise NotImplementedError('Extract mode should be "fixed", "threshold", "ratio" or "quantile"')
+        raise NotImplementedError(
+            'Extract mode should be "fixed", "threshold", "ratio" or "quantile"'
+        )
     lora_rank = max(1, lora_rank)
     lora_rank = min(out_ch, in_ch, lora_rank)
-    if lora_rank>=out_ch/2 and not is_cp:
-        return weight, 'full'
-    
+    if lora_rank >= out_ch / 2 and not is_cp:
+        return weight, "full"
+
     U = U[:, :lora_rank]
     S = S[:lora_rank]
     U = U @ torch.diag(S)
     Vh = Vh[:lora_rank, :]
-    
+
     diff = (weight - (U @ Vh).reshape(out_ch, in_ch, kernel_size, kernel_size)).detach()
     extract_weight_A = Vh.reshape(lora_rank, in_ch, kernel_size, kernel_size).detach()
     extract_weight_B = U.reshape(out_ch, lora_rank, 1, 1).detach()
     del U, S, Vh, weight
-    return (extract_weight_A, extract_weight_B, diff), 'low rank'
+    return (extract_weight_A, extract_weight_B, diff), "low rank"
 
 
 def extract_linear(
     weight: Union[torch.Tensor, nn.Parameter],
-    mode = 'fixed',
-    mode_param = 0,
-    device = 'cpu',
+    mode="fixed",
+    mode_param=0,
+    device="cpu",
 ) -> Tuple[nn.Parameter, nn.Parameter]:
     weight = weight.to(device)
     out_ch, in_ch = weight.shape
-    
+
     U, S, Vh = linalg.svd(weight)
-    
-    if mode=='fixed':
+
+    if mode == "full":
+        return weight, "full"
+    elif mode == "fixed":
         lora_rank = mode_param
-    elif mode=='threshold':
-        assert mode_param>=0
-        lora_rank = torch.sum(S>mode_param)
-    elif mode=='ratio':
-        assert 1>=mode_param>=0
-        min_s = torch.max(S)*mode_param
-        lora_rank = torch.sum(S>min_s)
-    elif mode=='quantile' or mode=='percentile':
-        assert 1>=mode_param>=0
+    elif mode == "threshold":
+        assert mode_param >= 0
+        lora_rank = torch.sum(S > mode_param)
+    elif mode == "ratio":
+        assert 1 >= mode_param >= 0
+        min_s = torch.max(S) * mode_param
+        lora_rank = torch.sum(S > min_s)
+    elif mode == "quantile" or mode == "percentile":
+        assert 1 >= mode_param >= 0
         s_cum = torch.cumsum(S, dim=0)
         min_cum_sum = mode_param * torch.sum(S)
-        lora_rank = torch.sum(s_cum<min_cum_sum)
+        lora_rank = torch.sum(s_cum < min_cum_sum)
     else:
-        raise NotImplementedError('Extract mode should be "fixed", "threshold", "ratio" or "quantile"')
+        raise NotImplementedError(
+            'Extract mode should be "fixed", "threshold", "ratio" or "quantile"'
+        )
     lora_rank = max(1, lora_rank)
     lora_rank = min(out_ch, in_ch, lora_rank)
-    if lora_rank>=out_ch/2:
-        return weight, 'full'
-    
+    if lora_rank >= out_ch / 2:
+        return weight, "full"
+
     U = U[:, :lora_rank]
     S = S[:lora_rank]
     U = U @ torch.diag(S)
     Vh = Vh[:lora_rank, :]
-    
+
     diff = (weight - U @ Vh).detach()
     extract_weight_A = Vh.reshape(lora_rank, in_ch).detach()
     extract_weight_B = U.reshape(out_ch, lora_rank).detach()
     del U, S, Vh, weight
-    return (extract_weight_A, extract_weight_B, diff), 'low rank'
+    return (extract_weight_A, extract_weight_B, diff), "low rank"
 
 
 def extract_diff(
-    base_model,
-    db_model,
-    mode = 'fixed',
-    linear_mode_param = 0,
-    conv_mode_param = 0,
-    extract_device = 'cpu',
-    use_bias = False,
-    sparsity = 0.98,
-    small_conv = True
+    base_tes,
+    db_tes,
+    base_unet,
+    db_unet,
+    mode="fixed",
+    linear_mode_param=0,
+    conv_mode_param=0,
+    extract_device="cpu",
+    use_bias=False,
+    sparsity=0.98,
+    small_conv=True,
 ):
     UNET_TARGET_REPLACE_MODULE = [
-        "Transformer2DModel", 
-        "Attention", 
-        "ResnetBlock2D", 
-        "Downsample2D", 
-        "Upsample2D"
+        "Transformer2DModel",
+        "Attention",
+        "ResnetBlock2D",
+        "Downsample2D",
+        "Upsample2D",
     ]
     UNET_TARGET_REPLACE_NAME = [
         "conv_in",
@@ -137,241 +151,298 @@ def extract_diff(
         "time_embedding.linear_2",
     ]
     TEXT_ENCODER_TARGET_REPLACE_MODULE = ["CLIPAttention", "CLIPMLP"]
-    LORA_PREFIX_UNET = 'lora_unet'
-    LORA_PREFIX_TEXT_ENCODER = 'lora_te'
+    LORA_PREFIX_UNET = "lora_unet"
+    LORA_PREFIX_TEXT_ENCODER = "lora_te"
+
     def make_state_dict(
-        prefix, 
+        prefix,
         root_module: torch.nn.Module,
         target_module: torch.nn.Module,
         target_replace_modules,
-        target_replace_names = []
+        target_replace_names=[],
     ):
         loras = {}
         temp = {}
         temp_name = {}
-        
+
         for name, module in root_module.named_modules():
             if module.__class__.__name__ in target_replace_modules:
                 temp[name] = {}
                 for child_name, child_module in module.named_modules():
-                    if child_module.__class__.__name__ not in {'Linear', 'Conv2d'}:
+                    if child_module.__class__.__name__ not in {"Linear", "Conv2d"}:
                         continue
-                    temp[name][child_name] = child_module.weight
+                    temp[name][child_name] = child_module
             elif name in target_replace_names:
-                temp_name[name] = module.weight
-        
+                temp_name[name] = module
+
         for name, module in tqdm(list(target_module.named_modules())):
             if name in temp:
                 weights = temp[name]
                 for child_name, child_module in module.named_modules():
-                    lora_name = prefix + '.' + name + '.' + child_name
-                    lora_name = lora_name.replace('.', '_')
+                    lora_name = prefix + "." + name + "." + child_name
+                    lora_name = lora_name.replace(".", "_")
                     layer = child_module.__class__.__name__
-                    if layer in {'Linear', 'Conv2d'}:
+                    if layer in {"Linear", "Conv2d"}:
                         root_weight = child_module.weight
-                        if torch.allclose(root_weight, weights[child_name]):
+                        if torch.allclose(root_weight, weights[child_name].weight):
                             continue
-                    
-                    if layer == 'Linear':
+
+                    if layer == "Linear":
                         weight, decompose_mode = extract_linear(
-                            (child_module.weight - weights[child_name]),
+                            (child_module.weight - weights[child_name].weight),
                             mode,
                             linear_mode_param,
-                            device = extract_device,
+                            device=extract_device,
                         )
-                        if decompose_mode == 'low rank':
+                        if decompose_mode == "low rank":
                             extract_a, extract_b, diff = weight
-                    elif layer == 'Conv2d':
-                        is_linear = (child_module.weight.shape[2] == 1
-                                     and child_module.weight.shape[3] == 1)
+                    elif layer == "Conv2d":
+                        is_linear = (
+                            child_module.weight.shape[2] == 1
+                            and child_module.weight.shape[3] == 1
+                        )
                         weight, decompose_mode = extract_conv(
-                            (child_module.weight - weights[child_name]), 
+                            (child_module.weight - weights[child_name].weight),
                             mode,
                             linear_mode_param if is_linear else conv_mode_param,
-                            device = extract_device,
+                            device=extract_device,
                         )
-                        if decompose_mode == 'low rank':
+                        if decompose_mode == "low rank":
                             extract_a, extract_b, diff = weight
-                        if small_conv and not is_linear and decompose_mode == 'low rank':
+                        if (
+                            small_conv
+                            and not is_linear
+                            and decompose_mode == "low rank"
+                        ):
                             dim = extract_a.size(0)
                             (extract_c, extract_a, _), _ = extract_conv(
-                                extract_a.transpose(0, 1), 
-                                'fixed', dim, 
-                                extract_device, True
+                                extract_a.transpose(0, 1),
+                                "fixed",
+                                dim,
+                                extract_device,
+                                True,
                             )
                             extract_a = extract_a.transpose(0, 1)
                             extract_c = extract_c.transpose(0, 1)
-                            loras[f'{lora_name}.lora_mid.weight'] = extract_c.detach().cpu().contiguous().half()
-                            diff = child_module.weight - torch.einsum(
-                                'i j k l, j r, p i -> p r k l', 
-                                extract_c, extract_a.flatten(1, -1), extract_b.flatten(1, -1)
-                            ).detach().cpu().contiguous()
+                            loras[f"{lora_name}.lora_mid.weight"] = (
+                                extract_c.detach().cpu().contiguous().half()
+                            )
+                            diff = (
+                                child_module.weight
+                                - torch.einsum(
+                                    "i j k l, j r, p i -> p r k l",
+                                    extract_c,
+                                    extract_a.flatten(1, -1),
+                                    extract_b.flatten(1, -1),
+                                )
+                                .detach()
+                                .cpu()
+                                .contiguous()
+                            )
                             del extract_c
                     else:
                         continue
-                    if decompose_mode == 'low rank':
-                        loras[f'{lora_name}.lora_down.weight'] = extract_a.detach().cpu().contiguous().half()
-                        loras[f'{lora_name}.lora_up.weight'] = extract_b.detach().cpu().contiguous().half()
-                        loras[f'{lora_name}.alpha'] = torch.Tensor([extract_a.shape[0]]).half()
+                    if decompose_mode == "low rank":
+                        loras[f"{lora_name}.lora_down.weight"] = (
+                            extract_a.detach().cpu().contiguous().half()
+                        )
+                        loras[f"{lora_name}.lora_up.weight"] = (
+                            extract_b.detach().cpu().contiguous().half()
+                        )
+                        loras[f"{lora_name}.alpha"] = torch.Tensor(
+                            [extract_a.shape[0]]
+                        ).half()
                         if use_bias:
                             diff = diff.detach().cpu().reshape(extract_b.size(0), -1)
-                            sparse_diff = make_sparse(diff, sparsity).to_sparse().coalesce()
-                            
+                            sparse_diff = (
+                                make_sparse(diff, sparsity).to_sparse().coalesce()
+                            )
+
                             indices = sparse_diff.indices().to(torch.int16)
                             values = sparse_diff.values().half()
-                            loras[f'{lora_name}.bias_indices'] = indices
-                            loras[f'{lora_name}.bias_values'] = values
-                            loras[f'{lora_name}.bias_size'] = torch.tensor(diff.shape).to(torch.int16)
+                            loras[f"{lora_name}.bias_indices"] = indices
+                            loras[f"{lora_name}.bias_values"] = values
+                            loras[f"{lora_name}.bias_size"] = torch.tensor(
+                                diff.shape
+                            ).to(torch.int16)
                         del extract_a, extract_b, diff
-                    elif decompose_mode == 'full':
-                        loras[f'{lora_name}.diff'] = weight.detach().cpu().contiguous().half()
+                    elif decompose_mode == "full":
+                        loras[f"{lora_name}.diff"] = (
+                            weight.detach().cpu().contiguous().half()
+                        )
+                        if weights[child_name].bias is not None:
+                            bias_diff = child_module.bias - weights[child_name].bias
+                            loras[f"{lora_name}.diff_b"] = (
+                                bias_diff.detach().cpu().contiguous().half()
+                            )
                     else:
                         raise NotImplementedError
             elif name in temp_name:
                 weights = temp_name[name]
-                lora_name = prefix + '.' + name
-                lora_name = lora_name.replace('.', '_')
+                lora_name = prefix + "." + name
+                lora_name = lora_name.replace(".", "_")
                 layer = module.__class__.__name__
-                
-                if layer in {'Linear', 'Conv2d'}:
+
+                if layer in {"Linear", "Conv2d"}:
                     root_weight = module.weight
-                    if torch.allclose(root_weight, weights):
+                    if torch.allclose(root_weight, weights.weight):
                         continue
-                
-                if layer == 'Linear':
+
+                if layer == "Linear":
                     weight, decompose_mode = extract_linear(
-                        (root_weight - weights),
+                        (root_weight - weights.weight),
                         mode,
                         linear_mode_param,
-                        device = extract_device,
+                        device=extract_device,
                     )
-                    if decompose_mode == 'low rank':
+                    if decompose_mode == "low rank":
                         extract_a, extract_b, diff = weight
-                elif layer == 'Conv2d':
-                    is_linear = (
-                        root_weight.shape[2] == 1
-                        and root_weight.shape[3] == 1
-                    )
+                elif layer == "Conv2d":
+                    is_linear = root_weight.shape[2] == 1 and root_weight.shape[3] == 1
                     weight, decompose_mode = extract_conv(
-                        (root_weight - weights), 
+                        (root_weight - weights.weight),
                         mode,
                         linear_mode_param if is_linear else conv_mode_param,
-                        device = extract_device,
+                        device=extract_device,
                     )
-                    if decompose_mode == 'low rank':
+                    if decompose_mode == "low rank":
                         extract_a, extract_b, diff = weight
-                    if small_conv and not is_linear and decompose_mode == 'low rank':
+                    if small_conv and not is_linear and decompose_mode == "low rank":
                         dim = extract_a.size(0)
                         (extract_c, extract_a, _), _ = extract_conv(
-                            extract_a.transpose(0, 1), 
-                            'fixed', dim, 
-                            extract_device, True
+                            extract_a.transpose(0, 1),
+                            "fixed",
+                            dim,
+                            extract_device,
+                            True,
                         )
                         extract_a = extract_a.transpose(0, 1)
                         extract_c = extract_c.transpose(0, 1)
-                        loras[f'{lora_name}.lora_mid.weight'] = extract_c.detach().cpu().contiguous().half()
-                        diff = root_weight - torch.einsum(
-                            'i j k l, j r, p i -> p r k l', 
-                            extract_c, extract_a.flatten(1, -1), extract_b.flatten(1, -1)
-                        ).detach().cpu().contiguous()
+                        loras[f"{lora_name}.lora_mid.weight"] = (
+                            extract_c.detach().cpu().contiguous().half()
+                        )
+                        diff = (
+                            root_weight
+                            - torch.einsum(
+                                "i j k l, j r, p i -> p r k l",
+                                extract_c,
+                                extract_a.flatten(1, -1),
+                                extract_b.flatten(1, -1),
+                            )
+                            .detach()
+                            .cpu()
+                            .contiguous()
+                        )
                         del extract_c
                 else:
                     continue
-                if decompose_mode == 'low rank':
-                    loras[f'{lora_name}.lora_down.weight'] = extract_a.detach().cpu().contiguous().half()
-                    loras[f'{lora_name}.lora_up.weight'] = extract_b.detach().cpu().contiguous().half()
-                    loras[f'{lora_name}.alpha'] = torch.Tensor([extract_a.shape[0]]).half()
+                if decompose_mode == "low rank":
+                    loras[f"{lora_name}.lora_down.weight"] = (
+                        extract_a.detach().cpu().contiguous().half()
+                    )
+                    loras[f"{lora_name}.lora_up.weight"] = (
+                        extract_b.detach().cpu().contiguous().half()
+                    )
+                    loras[f"{lora_name}.alpha"] = torch.Tensor(
+                        [extract_a.shape[0]]
+                    ).half()
                     if use_bias:
                         diff = diff.detach().cpu().reshape(extract_b.size(0), -1)
                         sparse_diff = make_sparse(diff, sparsity).to_sparse().coalesce()
-                        
+
                         indices = sparse_diff.indices().to(torch.int16)
                         values = sparse_diff.values().half()
-                        loras[f'{lora_name}.bias_indices'] = indices
-                        loras[f'{lora_name}.bias_values'] = values
-                        loras[f'{lora_name}.bias_size'] = torch.tensor(diff.shape).to(torch.int16)
+                        loras[f"{lora_name}.bias_indices"] = indices
+                        loras[f"{lora_name}.bias_values"] = values
+                        loras[f"{lora_name}.bias_size"] = torch.tensor(diff.shape).to(
+                            torch.int16
+                        )
                     del extract_a, extract_b, diff
-                elif decompose_mode == 'full':
-                    loras[f'{lora_name}.diff'] = weight.detach().cpu().contiguous().half()
+                elif decompose_mode == "full":
+                    loras[f"{lora_name}.diff"] = (
+                        weight.detach().cpu().contiguous().half()
+                    )
+                    if weights.bias is not None:
+                        bias_diff = module.bias - weights.bias
+                        loras[f"{lora_name}.diff_b"] = (
+                            bias_diff.detach().cpu().contiguous().half()
+                        )
                 else:
                     raise NotImplementedError
         return loras
-    
-    text_encoder_loras = make_state_dict(
-        LORA_PREFIX_TEXT_ENCODER, 
-        base_model[0], db_model[0], 
-        TEXT_ENCODER_TARGET_REPLACE_MODULE
-    )
-    
-    unet_loras = make_state_dict(
+
+    all_loras = {}
+    for te1, te2 in zip(base_tes, db_tes):
+        all_loras |= make_state_dict(
+            LORA_PREFIX_TEXT_ENCODER, te1, te2, TEXT_ENCODER_TARGET_REPLACE_MODULE
+        )
+
+    all_loras |= make_state_dict(
         LORA_PREFIX_UNET,
-        base_model[2], db_model[2], 
+        base_unet,
+        db_unet,
         UNET_TARGET_REPLACE_MODULE,
-        UNET_TARGET_REPLACE_NAME
+        UNET_TARGET_REPLACE_NAME,
     )
-    print(len(text_encoder_loras), len(unet_loras))
-    return text_encoder_loras|unet_loras
+    print(len(all_loras))
+    return all_loras
 
 
-def get_module(
-    lyco_state_dict: Dict,
-    lora_name
-):
-    if f'{lora_name}.lora_up.weight' in lyco_state_dict:
-        up = lyco_state_dict[f'{lora_name}.lora_up.weight']
-        down = lyco_state_dict[f'{lora_name}.lora_down.weight']
-        mid = lyco_state_dict.get(f'{lora_name}.lora_mid.weight', None)
-        alpha = lyco_state_dict.get(f'{lora_name}.alpha', None)
-        return 'locon', (up, down, mid, alpha)
-    elif f'{lora_name}.hada_w1_a' in lyco_state_dict:
-        w1a = lyco_state_dict[f'{lora_name}.hada_w1_a']
-        w1b = lyco_state_dict[f'{lora_name}.hada_w1_b']
-        w2a = lyco_state_dict[f'{lora_name}.hada_w2_a']
-        w2b = lyco_state_dict[f'{lora_name}.hada_w2_b']
-        t1 = lyco_state_dict.get(f'{lora_name}.hada_t1', None)
-        t2 = lyco_state_dict.get(f'{lora_name}.hada_t2', None)
-        alpha = lyco_state_dict.get(f'{lora_name}.alpha', None)
-        return 'hada', (w1a, w1b, w2a, w2b, t1, t2, alpha)
-    elif f'{lora_name}.weight' in lyco_state_dict:
-        weight = lyco_state_dict[f'{lora_name}.weight']
-        on_input = lyco_state_dict.get(f'{lora_name}.on_input', False)
-        return 'ia3', (weight, on_input)
-    elif (f'{lora_name}.lokr_w1' in lyco_state_dict
-          or f'{lora_name}.lokr_w1_a' in lyco_state_dict):
-        w1 = lyco_state_dict.get(f'{lora_name}.lokr_w1', None)
-        w1a = lyco_state_dict.get(f'{lora_name}.lokr_w1_a', None)
-        w1b = lyco_state_dict.get(f'{lora_name}.lokr_w1_b', None)
-        w2 = lyco_state_dict.get(f'{lora_name}.lokr_w2', None)
-        w2a = lyco_state_dict.get(f'{lora_name}.lokr_w2_a', None)
-        w2b = lyco_state_dict.get(f'{lora_name}.lokr_w2_b', None)
-        t1 = lyco_state_dict.get(f'{lora_name}.lokr_t1', None)
-        t2 = lyco_state_dict.get(f'{lora_name}.lokr_t2', None)
-        alpha = lyco_state_dict.get(f'{lora_name}.alpha', None)
-        return 'kron', (w1, w1a, w1b, w2, w2a, w2b, t1, t2, alpha)
-    elif f'{lora_name}.diff' in lyco_state_dict:
-        diff = lyco_state_dict[f'{lora_name}.diff']
-        diff_b = lyco_state_dict.get(f'{lora_name}.diff_b', None)
-        return 'full', (diff, diff_b)
-    elif f'{lora_name}.w_norm' in lyco_state_dict:
-        w_norm = lyco_state_dict[f'{lora_name}.w_norm']
-        b_norm = lyco_state_dict.get(f'{lora_name}.b_norm')
-        return 'norm', (w_norm, b_norm)
+def get_module(lyco_state_dict: Dict, lora_name):
+    if f"{lora_name}.lora_up.weight" in lyco_state_dict:
+        up = lyco_state_dict[f"{lora_name}.lora_up.weight"]
+        down = lyco_state_dict[f"{lora_name}.lora_down.weight"]
+        mid = lyco_state_dict.get(f"{lora_name}.lora_mid.weight", None)
+        alpha = lyco_state_dict.get(f"{lora_name}.alpha", None)
+        return "locon", (up, down, mid, alpha)
+    elif f"{lora_name}.hada_w1_a" in lyco_state_dict:
+        w1a = lyco_state_dict[f"{lora_name}.hada_w1_a"]
+        w1b = lyco_state_dict[f"{lora_name}.hada_w1_b"]
+        w2a = lyco_state_dict[f"{lora_name}.hada_w2_a"]
+        w2b = lyco_state_dict[f"{lora_name}.hada_w2_b"]
+        t1 = lyco_state_dict.get(f"{lora_name}.hada_t1", None)
+        t2 = lyco_state_dict.get(f"{lora_name}.hada_t2", None)
+        alpha = lyco_state_dict.get(f"{lora_name}.alpha", None)
+        return "hada", (w1a, w1b, w2a, w2b, t1, t2, alpha)
+    elif f"{lora_name}.weight" in lyco_state_dict:
+        weight = lyco_state_dict[f"{lora_name}.weight"]
+        on_input = lyco_state_dict.get(f"{lora_name}.on_input", False)
+        return "ia3", (weight, on_input)
+    elif (
+        f"{lora_name}.lokr_w1" in lyco_state_dict
+        or f"{lora_name}.lokr_w1_a" in lyco_state_dict
+    ):
+        w1 = lyco_state_dict.get(f"{lora_name}.lokr_w1", None)
+        w1a = lyco_state_dict.get(f"{lora_name}.lokr_w1_a", None)
+        w1b = lyco_state_dict.get(f"{lora_name}.lokr_w1_b", None)
+        w2 = lyco_state_dict.get(f"{lora_name}.lokr_w2", None)
+        w2a = lyco_state_dict.get(f"{lora_name}.lokr_w2_a", None)
+        w2b = lyco_state_dict.get(f"{lora_name}.lokr_w2_b", None)
+        t1 = lyco_state_dict.get(f"{lora_name}.lokr_t1", None)
+        t2 = lyco_state_dict.get(f"{lora_name}.lokr_t2", None)
+        alpha = lyco_state_dict.get(f"{lora_name}.alpha", None)
+        return "kron", (w1, w1a, w1b, w2, w2a, w2b, t1, t2, alpha)
+    elif f"{lora_name}.diff" in lyco_state_dict:
+        diff = lyco_state_dict[f"{lora_name}.diff"]
+        diff_b = lyco_state_dict.get(f"{lora_name}.diff_b", None)
+        return "full", (diff, diff_b)
+    elif f"{lora_name}.w_norm" in lyco_state_dict:
+        w_norm = lyco_state_dict[f"{lora_name}.w_norm"]
+        b_norm = lyco_state_dict.get(f"{lora_name}.b_norm")
+        return "norm", (w_norm, b_norm)
     else:
-        return 'None', ()
+        return "None", ()
 
 
-def cp_weight_from_conv(
-    up, down, mid
-):
+def cp_weight_from_conv(up, down, mid):
     up = up.reshape(up.size(0), up.size(1))
     down = down.reshape(down.size(0), down.size(1))
-    return torch.einsum('m n w h, i m, n j -> i j w h', mid, up, down)
+    return torch.einsum("m n w h, i m, n j -> i j w h", mid, up, down)
 
-def cp_weight(
-    wa, wb, t
-):
-    temp = torch.einsum('i j k l, j r -> i r k l', t, wb)
-    return torch.einsum('i j k l, i r -> r j k l', temp, wa)
+
+def cp_weight(wa, wb, t):
+    temp = torch.einsum("i j k l, j r -> i r k l", t, wb)
+    return torch.einsum("i j k l, i r -> r j k l", temp, wa)
 
 
 @torch.no_grad()
@@ -380,17 +451,17 @@ def rebuild_weight(module_type, params, orig_weight, orig_bias, scale=1):
         return orig_weight, orig_bias
     merged = orig_weight
     merged_bias = orig_bias
-    if module_type == 'locon':
+    if module_type == "locon":
         up, down, mid, alpha = params
         if alpha is not None:
-            scale *= alpha/up.size(1)
+            scale *= alpha / up.size(1)
         if mid is not None:
             rebuild = cp_weight_from_conv(up, down, mid)
         else:
-            rebuild = up.reshape(up.size(0),-1) @ down.reshape(down.size(0), -1)
+            rebuild = up.reshape(up.size(0), -1) @ down.reshape(down.size(0), -1)
         merged = orig_weight + rebuild.reshape(orig_weight.shape) * scale
         del up, down, mid, alpha, params, rebuild
-    elif module_type == 'hada':
+    elif module_type == "hada":
         w1a, w1b, w2a, w2b, t1, t2, alpha = params
         if alpha is not None:
             scale *= alpha / w1b.size(0)
@@ -405,13 +476,13 @@ def rebuild_weight(module_type, params, orig_weight, orig_bias, scale=1):
         rebuild = (rebuild1 * rebuild2).reshape(orig_weight.shape)
         merged = orig_weight + rebuild * scale
         del w1a, w1b, w2a, w2b, t1, t2, alpha, params, rebuild, rebuild1, rebuild2
-    elif module_type == 'ia3':
+    elif module_type == "ia3":
         weight, on_input = params
         if not on_input:
             weight = weight.reshape(-1, 1)
         merged = orig_weight + weight * orig_weight * scale
         del weight, on_input, params
-    elif module_type == 'kron':
+    elif module_type == "kron":
         w1, w1a, w1b, w2, w2a, w2b, t1, t2, alpha = params
         if alpha is not None and (w1b is not None or w2b is not None):
             scale *= alpha / (w1b.size(0) if w1b else w2b.size(0))
@@ -428,40 +499,38 @@ def rebuild_weight(module_type, params, orig_weight, orig_bias, scale=1):
         if len(w2.shape) == 4:
             w1 = w1.unsqueeze(2).unsqueeze(2)
         w2 = w2.contiguous()
-        rebuild = torch.kron(w1, w2).reshape(orig_weight.shape) 
-        merged = orig_weight + rebuild* scale
+        rebuild = torch.kron(w1, w2).reshape(orig_weight.shape)
+        merged = orig_weight + rebuild * scale
         del w1, w1a, w1b, w2, w2a, w2b, t1, t2, alpha, params, rebuild
-    elif module_type == 'full':
-        rebuild = params[0].reshape(orig_weight.shape) 
-        rebuild_b = params[1].reshape(orig_bias.shape) if orig_bias is not None else None
+    elif module_type == "full":
+        rebuild = params[0].reshape(orig_weight.shape)
+        rebuild_b = (
+            params[1].reshape(orig_bias.shape) if orig_bias is not None else None
+        )
         merged = orig_weight + rebuild * scale
         if rebuild_b is not None:
             merged_bias = orig_bias + rebuild_b * scale
         del params, rebuild, rebuild_b
-    elif module_type == 'norm':
+    elif module_type == "norm":
         rebuild = params[0].reshape(orig_weight.shape)
         rebuild_b = params[1].reshape(orig_bias.shape)
         merged = orig_weight + rebuild * scale
         merged_bias = orig_bias + rebuild_b * scale
     else:
-        raise NotImplementedError
-    
+        merged = orig_weight
+        merged_bias = orig_bias
+
     return merged, merged_bias
 
 
 @torch.no_grad()
-def merge(
-    base_model,
-    lyco_state_dict,
-    scale: float = 1.0,
-    device = 'cpu'
-):
+def merge(tes, unet, lyco_state_dict, scale: float = 1.0, device="cpu"):
     UNET_TARGET_REPLACE_MODULE = [
-        "Transformer2DModel", 
-        "Attention", 
-        "ResnetBlock2D", 
-        "Downsample2D", 
-        "Upsample2D"
+        "Transformer2DModel",
+        "Attention",
+        "ResnetBlock2D",
+        "Downsample2D",
+        "Upsample2D",
     ]
     UNET_TARGET_REPLACE_NAME = [
         "conv_in",
@@ -470,28 +539,34 @@ def merge(
         "time_embedding.linear_2",
     ]
     TEXT_ENCODER_TARGET_REPLACE_MODULE = ["CLIPAttention", "CLIPMLP"]
-    LORA_PREFIX_UNET = 'lora_unet'
-    LORA_PREFIX_TEXT_ENCODER = 'lora_te'
+    LORA_PREFIX_UNET = "lora_unet"
+    LORA_PREFIX_TEXT_ENCODER = "lora_te"
     merged = 0
+
     def merge_state_dict(
-        prefix, 
+        prefix,
         root_module: torch.nn.Module,
-        lyco_state_dict: Dict[str,torch.Tensor],
+        lyco_state_dict: Dict[str, torch.Tensor],
         target_replace_modules,
-        target_replace_names = []
+        target_replace_names=[],
     ):
         nonlocal merged
-        for name, module in tqdm(list(root_module.named_modules()), desc=f'Merging {prefix}'):
+        for name, module in tqdm(
+            list(root_module.named_modules()), desc=f"Merging {prefix}"
+        ):
             if module.__class__.__name__ in target_replace_modules:
                 for child_name, child_module in module.named_modules():
-                    if child_module.__class__.__name__ not in {'Linear', 'Conv2d'}:
+                    if child_module.__class__.__name__ not in {"Linear", "Conv2d"}:
                         continue
-                    lora_name = prefix + '.' + name + '.' + child_name
-                    lora_name = lora_name.replace('.', '_')
-                    
-                    result, result_b = rebuild_weight(*get_module(
-                        lyco_state_dict, lora_name
-                    ), getattr(child_module, 'weight'), getattr(child_module, 'bias', None), scale)
+                    lora_name = prefix + "." + name + "." + child_name
+                    lora_name = lora_name.replace(".", "_")
+
+                    result, result_b = rebuild_weight(
+                        *get_module(lyco_state_dict, lora_name),
+                        getattr(child_module, "weight"),
+                        getattr(child_module, "bias", None),
+                        scale,
+                    )
                     if result is not None:
                         merged += 1
                         child_module.requires_grad_(False)
@@ -499,37 +574,43 @@ def merge_state_dict(
                     if result_b is not None:
                         child_module.bias.copy_(result_b)
             elif name in target_replace_names:
-                lora_name = prefix + '.' + name
-                lora_name = lora_name.replace('.', '_')
-                    
-                result, result_b = rebuild_weight(*get_module(
-                    lyco_state_dict, lora_name
-                ), getattr(module, 'weight'), getattr(module, 'bias', None), scale)
+                lora_name = prefix + "." + name
+                lora_name = lora_name.replace(".", "_")
+
+                result, result_b = rebuild_weight(
+                    *get_module(lyco_state_dict, lora_name),
+                    getattr(module, "weight"),
+                    getattr(module, "bias", None),
+                    scale,
+                )
                 if result is not None:
                     merged += 1
                     module.requires_grad_(False)
                     module.weight.copy_(result)
                 if result_b is not None:
                     module.bias.copy_(result_b)
-    
-    for k, v in tqdm(list(lyco_state_dict.items()), desc='Converting Dtype and Device'):
-        if device=='cpu':
+
+    for k, v in tqdm(list(lyco_state_dict.items()), desc="Converting Dtype and Device"):
+        if device == "cpu":
             lyco_state_dict[k] = v.float().cpu()
         else:
-            lyco_state_dict[k] = v.to(device, dtype=base_model[0].parameters().__next__().dtype)
-    
-    merge_state_dict(
-        LORA_PREFIX_TEXT_ENCODER,
-        base_model[0].to(device),
-        lyco_state_dict,
-        TEXT_ENCODER_TARGET_REPLACE_MODULE,
-        UNET_TARGET_REPLACE_NAME
-    )
+            lyco_state_dict[k] = v.to(
+                device, dtype=tes[0].parameters().__next__().dtype
+            )
+
+    for te in tes:
+        merge_state_dict(
+            LORA_PREFIX_TEXT_ENCODER,
+            te.to(device),
+            lyco_state_dict,
+            TEXT_ENCODER_TARGET_REPLACE_MODULE,
+            UNET_TARGET_REPLACE_NAME,
+        )
     merge_state_dict(
         LORA_PREFIX_UNET,
-        base_model[2].to(device),
+        unet.to(device),
         lyco_state_dict,
         UNET_TARGET_REPLACE_MODULE,
-        UNET_TARGET_REPLACE_NAME
+        UNET_TARGET_REPLACE_NAME,
     )
-    print(f'{merged} Modules been merged')
+    print(f"{merged} Modules been merged")
diff --git a/lycoris/utils/logger.py b/lycoris/utils/logger.py
new file mode 100644
index 0000000..e3f51e5
--- /dev/null
+++ b/lycoris/utils/logger.py
@@ -0,0 +1,35 @@
+import logging
+import copy
+import sys
+
+
+class ColoredFormatter(logging.Formatter):
+    COLORS = {
+        "DEBUG": "\033[0;36m",  # CYAN
+        "INFO": "\033[0;32m",  # GREEN
+        "WARNING": "\033[0;33m",  # YELLOW
+        "ERROR": "\033[0;31m",  # RED
+        "CRITICAL": "\033[0;37;41m",  # WHITE ON RED
+        "RESET": "\033[0m",  # RESET COLOR
+    }
+
+    def format(self, record):
+        colored_record = copy.copy(record)
+        levelname = colored_record.levelname
+        seq = self.COLORS.get(levelname, self.COLORS["RESET"])
+        colored_record.levelname = f"{seq}{levelname}{self.COLORS['RESET']}"
+        return super().format(colored_record)
+
+
+# Create a new logger
+logger = logging.getLogger("LyCORIS")
+logger.propagate = False
+
+# Add handler if we don't have one.
+if not logger.handlers:
+    handler = logging.StreamHandler(sys.stdout)
+    handler.setFormatter(ColoredFormatter("[%(name)s]-%(levelname)s: %(message)s"))
+    logger.addHandler(handler)
+
+logger.setLevel(logging.DEBUG)
+logger.debug("Logger initialized.")
diff --git a/lycoris/utils/preset.py b/lycoris/utils/preset.py
index cd0e049..f6ec844 100644
--- a/lycoris/utils/preset.py
+++ b/lycoris/utils/preset.py
@@ -5,5 +5,5 @@ def read_preset(preset):
     try:
         return toml.load(preset)
     except Exception as e:
-        print('Error: cannot read preset file. ', e)
-        return None
\ No newline at end of file
+        print("Error: cannot read preset file. ", e)
+        return None
diff --git a/lycoris/utils/xformers_utils.py b/lycoris/utils/xformers_utils.py
index 14d1ba0..cc0bd85 100644
--- a/lycoris/utils/xformers_utils.py
+++ b/lycoris/utils/xformers_utils.py
@@ -6,6 +6,7 @@
 
 try:
     from xformers.ops import memory_efficient_attention
+
     XFORMERS_AVAIL = True
 except:
     memory_efficient_attention = None
diff --git a/lycoris/wrapper.py b/lycoris/wrapper.py
new file mode 100644
index 0000000..ce42f71
--- /dev/null
+++ b/lycoris/wrapper.py
@@ -0,0 +1,485 @@
+# General LyCORIS wrapper based on kohya-ss/sd-scripts' style
+
+import math
+import os
+import regex as re
+import sys
+
+sys.setrecursionlimit(10000)
+from typing import List
+from warnings import warn
+
+import torch
+
+from .utils import *
+from .modules.locon import LoConModule
+from .modules.loha import LohaModule
+from .modules.lokr import LokrModule
+from .modules.dylora import DyLoraModule
+from .modules.glora import GLoRAModule
+from .modules.norms import NormModule
+from .modules.full import FullModule
+from .modules.diag_oft import DiagOFTModule
+from .modules import make_module
+
+from .config import PRESET
+from .utils.preset import read_preset
+from .utils import get_module, str_bool
+
+
+network_module_dict = {
+    "lora": LoConModule,
+    "locon": LoConModule,
+    "loha": LohaModule,
+    "lokr": LokrModule,
+    "dylora": DyLoraModule,
+    "glora": GLoRAModule,
+    "full": FullModule,
+    "diag-oft": DiagOFTModule,
+}
+
+
+def create_lycoris(module, multiplier, linear_dim, linear_alpha, **kwargs):
+    if linear_dim is None:
+        linear_dim = 4  # default
+    conv_dim = int(kwargs.get("conv_dim", linear_dim) or linear_dim)
+    conv_alpha = float(kwargs.get("conv_alpha", linear_alpha) or linear_alpha)
+    dropout = float(kwargs.get("dropout", 0.0) or 0.0)
+    rank_dropout = float(kwargs.get("rank_dropout", 0.0) or 0.0)
+    module_dropout = float(kwargs.get("module_dropout", 0.0) or 0.0)
+    algo = (kwargs.get("algo", "lora") or "lora").lower()
+    use_tucker = str_bool(
+        not kwargs.get("disable_conv_cp", True)
+        or kwargs.get("use_conv_cp", False)
+        or kwargs.get("use_cp", False)
+        or kwargs.get("use_tucker", False)
+    )
+    if "disable_conv_cp" in kwargs or "use_cp" in kwargs or "use_conv_cp" in kwargs:
+        warn(
+            "disable_conv_cp and use_cp are deprecated. Please use use_tucker instead.",
+            stacklevel=2,
+        )
+    use_scalar = str_bool(kwargs.get("use_scalar", False))
+    block_size = int(kwargs.get("block_size", 4) or 4)
+    train_norm = str_bool(kwargs.get("train_norm", False))
+    constrain = float(kwargs.get("constrain", 0) or 0)
+    rescaled = str_bool(kwargs.get("rescaled", False))
+
+    if algo == "glora" and conv_dim > 0:
+        conv_dim = 0
+        print("Disable conv layer for GLoRA")
+
+    preset = kwargs.get("preset", "full")
+    if preset not in PRESET:
+        preset = read_preset(preset)
+    else:
+        preset = PRESET[preset]
+    assert preset is not None
+    LycorisNetwork.apply_preset(preset)
+
+    print(f"Using rank adaptation algo: {algo}")
+
+    if (
+        (algo == "loha")
+        and not kwargs.get("no_dim_warn", False)
+        and (linear_dim > 64 or conv_dim > 64)
+    ):
+        print("=" * 20 + "WARNING" + "=" * 20)
+        warning_type = {"loha": "Hadamard Product representation"}
+        warning_msg = (
+            "You are not supposed to use dim>64 (64*64 = 4096, it already has enough rank)\n"
+            f"in {warning_type[algo]}!\n"
+            "Please consider use lower dim or disable this warning with --network_args no_dim_warn=True\n"
+            f"If you just want to use high dim {algo}, please consider use lower lr."
+        )
+        warn(warning_msg, stacklevel=2)
+        print("=" * 20 + "WARNING" + "=" * 20)
+
+    network = LycorisNetwork(
+        module,
+        multiplier=multiplier,
+        lora_dim=linear_dim,
+        conv_lora_dim=conv_dim,
+        alpha=linear_alpha,
+        conv_alpha=conv_alpha,
+        dropout=dropout,
+        rank_dropout=rank_dropout,
+        module_dropout=module_dropout,
+        use_tucker=use_tucker,
+        use_scalar=use_scalar,
+        network_module=algo,
+        train_norm=train_norm,
+        decompose_both=kwargs.get("decompose_both", False),
+        factor=kwargs.get("factor", -1),
+        block_size=block_size,
+        constrain=constrain,
+        rescaled=rescaled,
+    )
+
+    if algo == "dylora":
+        # dylora didn't support scale weight norm yet
+        delattr(type(network), "apply_max_norm_regularization")
+
+    return network
+
+
+def create_lycoris_from_weights(multiplier, file, module, weights_sd=None, **kwargs):
+    if weights_sd is None:
+        if os.path.splitext(file)[1] == ".safetensors":
+            from safetensors.torch import load_file
+
+            weights_sd = load_file(file)
+        else:
+            weights_sd = torch.load(file, map_location="cpu")
+
+    # get dim/alpha mapping
+    loras = {}
+    for key in weights_sd:
+        if "." not in key:
+            continue
+
+        lora_name = key.split(".")[0]
+        loras[lora_name] = None
+
+    for name, modules in module.named_modules():
+        lora_name = f"{LycorisNetwork.LORA_PREFIX}_{name}".replace(".", "_")
+        if lora_name in loras:
+            loras[lora_name] = modules
+
+    network = LycorisNetwork(module)
+    network.loras = []
+
+    for lora_name, orig_modules in loras.items():
+        if orig_modules is None:
+            continue
+        lyco_type, params = get_module(weights_sd, lora_name)
+        module = make_module(lyco_type, params, lora_name, orig_modules)
+        if module is not None:
+            network.loras.append(module)
+
+    for lora in network.loras:
+        lora.multiplier = multiplier
+
+    return network, weights_sd
+
+
+class LycorisNetwork(torch.nn.Module):
+    ENABLE_CONV = True
+    TARGET_REPLACE_MODULE = []
+    TARGET_REPLACE_NAME = []
+    LORA_PREFIX = "lycoris"
+    MODULE_ALGO_MAP = {}
+    NAME_ALGO_MAP = {}
+
+    @classmethod
+    def apply_preset(cls, preset):
+        if "enable_conv" in preset:
+            cls.ENABLE_CONV = preset["enable_conv"]
+        if "target_module" in preset:
+            cls.TARGET_REPLACE_MODULE = preset["target_module"]
+        if "target_name" in preset:
+            cls.TARGET_REPLACE_NAME = preset["target_name"]
+        if "module_algo_map" in preset:
+            cls.MODULE_ALGO_MAP = preset["module_algo_map"]
+        if "name_algo_map" in preset:
+            cls.NAME_ALGO_MAP = preset["name_algo_map"]
+        return cls
+
+    def __init__(
+        self,
+        module: nn.Module,
+        multiplier=1.0,
+        lora_dim=4,
+        conv_lora_dim=4,
+        alpha=1,
+        conv_alpha=1,
+        use_tucker=False,
+        dropout=0,
+        rank_dropout=0,
+        module_dropout=0,
+        network_module: str = "locon",
+        norm_modules=NormModule,
+        train_norm=False,
+        **kwargs,
+    ) -> None:
+        super().__init__()
+        root_kwargs = kwargs
+        self.multiplier = multiplier
+        self.lora_dim = lora_dim
+
+        if not self.ENABLE_CONV:
+            conv_lora_dim = 0
+
+        self.conv_lora_dim = int(conv_lora_dim)
+        if self.conv_lora_dim and self.conv_lora_dim != self.lora_dim:
+            print("Apply different lora dim for conv layer")
+            print(f"Conv Dim: {conv_lora_dim}, Linear Dim: {lora_dim}")
+        elif self.conv_lora_dim == 0:
+            print("Disable conv layer")
+
+        self.alpha = alpha
+        self.conv_alpha = float(conv_alpha)
+        if self.conv_lora_dim and self.alpha != self.conv_alpha:
+            print("Apply different alpha value for conv layer")
+            print(f"Conv alpha: {conv_alpha}, Linear alpha: {alpha}")
+
+        if 1 >= dropout >= 0:
+            print(f"Use Dropout value: {dropout}")
+        self.dropout = dropout
+        self.rank_dropout = rank_dropout
+        self.module_dropout = module_dropout
+
+        self.use_tucker = use_tucker
+
+        def create_single_module(
+            lora_name: str,
+            module: torch.nn.Module,
+            algo_name,
+            dim=None,
+            alpha=None,
+            use_tucker=self.use_tucker,
+            **kwargs,
+        ):
+            for k, v in root_kwargs.items():
+                if k in kwargs:
+                    continue
+                kwargs[k] = v
+
+            if train_norm and "Norm" in module.__class__.__name__:
+                return norm_modules(
+                    lora_name,
+                    module,
+                    self.multiplier,
+                    self.rank_dropout,
+                    self.module_dropout,
+                    **kwargs,
+                )
+            lora = None
+            if module.__class__.__name__ == "Linear" and lora_dim > 0:
+                dim = dim or lora_dim
+                alpha = alpha or self.alpha
+            elif module.__class__.__name__ == "Conv2d":
+                k_size, *_ = module.kernel_size
+                if k_size == 1 and lora_dim > 0:
+                    dim = dim or lora_dim
+                    alpha = alpha or self.alpha
+                elif conv_lora_dim > 0 or dim:
+                    dim = dim or conv_lora_dim
+                    alpha = alpha or self.conv_alpha
+                else:
+                    return None
+            else:
+                return None
+            lora = network_module_dict[algo_name](
+                lora_name,
+                module,
+                self.multiplier,
+                dim,
+                alpha,
+                self.dropout,
+                self.rank_dropout,
+                self.module_dropout,
+                use_tucker,
+                **kwargs,
+            )
+            return lora
+
+        def create_modules_(
+            prefix: str,
+            root_module: torch.nn.Module,
+            algo,
+            configs={},
+        ):
+            loras = {}
+            lora_names = []
+            for name, module in root_module.named_modules():
+                module_name = module.__class__.__name__
+                if module_name in self.MODULE_ALGO_MAP:
+                    next_config = self.MODULE_ALGO_MAP[module_name]
+                    next_algo = next_config.get("algo", algo)
+                    new_loras, new_lora_names = create_modules_(
+                        f"{prefix}_{name}", module, next_algo, next_config
+                    )
+                    for lora_name, lora in zip(new_lora_names, new_loras):
+                        if lora_name not in loras:
+                            loras[lora_name] = lora
+                            lora_names.append(lora_name)
+                    continue
+                lora_name = prefix + "." + name
+                lora_name = lora_name.replace(".", "_")
+                if lora_name in loras:
+                    continue
+
+                lora = create_single_module(lora_name, module, algo, **configs)
+                if lora is not None:
+                    loras[lora_name] = lora
+                    lora_names.append(lora_name)
+            return [loras[lora_name] for lora_name in lora_names], lora_names
+
+        # create module instances
+        def create_modules(
+            prefix,
+            root_module: torch.nn.Module,
+            target_replace_modules,
+            target_replace_names=[],
+        ) -> List:
+            print("Create LyCORIS Module")
+            loras = []
+            next_config = {}
+            for name, module in root_module.named_modules():
+                module_name = module.__class__.__name__
+                if module_name in target_replace_modules:
+                    if module_name in self.MODULE_ALGO_MAP:
+                        next_config = self.MODULE_ALGO_MAP[module_name]
+                        algo = next_config.get("algo", network_module)
+                    else:
+                        algo = network_module
+                    loras.extend(
+                        create_modules_(f"{prefix}_{name}", module, algo, next_config)[
+                            0
+                        ]
+                    )
+                    next_config = {}
+                elif name in target_replace_names or any(
+                    re.match(t, name) for t in target_replace_names
+                ):
+                    if name in self.NAME_ALGO_MAP:
+                        next_config = self.NAME_ALGO_MAP[name]
+                        algo = next_config.get("algo", network_module)
+                    elif module_name in self.MODULE_ALGO_MAP:
+                        next_config = self.MODULE_ALGO_MAP[module_name]
+                        algo = next_config.get("algo", network_module)
+                    else:
+                        algo = network_module
+                    lora_name = prefix + "." + name
+                    lora_name = lora_name.replace(".", "_")
+                    lora = create_single_module(lora_name, module, algo, **next_config)
+                    next_config = {}
+                    if lora is not None:
+                        loras.append(lora)
+            return loras
+
+        self.loras = create_modules(
+            LycorisNetwork.LORA_PREFIX,
+            module,
+            LycorisNetwork.TARGET_REPLACE_MODULE,
+            LycorisNetwork.TARGET_REPLACE_NAME,
+        )
+        print(f"create LyCORIS: {len(self.loras)} modules.")
+
+        algo_table = {}
+        for lora in self.loras:
+            algo_table[lora.__class__.__name__] = (
+                algo_table.get(lora.__class__.__name__, 0) + 1
+            )
+        print(f"module type table: {algo_table}")
+
+        self.weights_sd = None
+
+        # assertion
+        names = set()
+        for lora in self.loras:
+            assert (
+                lora.lora_name not in names
+            ), f"duplicated lora name: {lora.lora_name}"
+            names.add(lora.lora_name)
+
+    def set_multiplier(self, multiplier):
+        self.multiplier = multiplier
+        for lora in self.loras:
+            lora.multiplier = self.multiplier
+
+    def load_weights(self, file):
+        if os.path.splitext(file)[1] == ".safetensors":
+            from safetensors.torch import load_file, safe_open
+
+            self.weights_sd = load_file(file)
+        else:
+            self.weights_sd = torch.load(file, map_location="cpu")
+        missing, unexpected = self.load_state_dict(self.weights_sd, strict=False)
+        state = {}
+        if missing:
+            state["missing keys"] = missing
+        if unexpected:
+            state["unexpected keys"] = unexpected
+        return state
+
+    def apply_to(self):
+        for lora in self.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 apply_max_norm_regularization(self, max_norm_value, device):
+        key_scaled = 0
+        norms = []
+        for model in self.loras:
+            if hasattr(model, "apply_max_norm"):
+                scaled, norm = model.apply_max_norm(max_norm_value, device)
+                norms.append(norm)
+                key_scaled += scaled
+
+        if key_scaled == 0:
+            return key_scaled, 0, 0
+
+        return key_scaled, sum(norms) / len(norms), max(norms)
+
+    def enable_gradient_checkpointing(self):
+        # not supported
+        def make_ckpt(module):
+            if isinstance(module, torch.nn.Module):
+                module.grad_ckpt = True
+
+        self.apply(make_ckpt)
+        pass
+
+    def prepare_optimizer_params(self, lr):
+        def enumerate_params(loras):
+            params = []
+            for lora in loras:
+                params.extend(lora.parameters())
+            return params
+
+        self.requires_grad_(True)
+        all_params = []
+
+        param_data = {"params": enumerate_params(self.loras)}
+        if lr is not None:
+            param_data["lr"] = lr
+        all_params.append(param_data)
+        return all_params
+
+    def prepare_grad_etc(self):
+        self.requires_grad_(True)
+
+    def on_epoch_start(self):
+        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
+
+            # Precalculate model hashes to save time on indexing
+            if metadata is None:
+                metadata = {}
+            save_file(state_dict, file, metadata)
+        else:
+            torch.save(state_dict, file)
diff --git a/setup.py b/setup.py
index 8809725..44d6449 100644
--- a/setup.py
+++ b/setup.py
@@ -1,18 +1,14 @@
 from setuptools import setup, find_packages
 
+
 setup(
-    name='lycoris_lora',
+    name="lycoris_lora",
     packages=find_packages(),
-    version='1.9.0',
-    url='https://github.com/KohakuBlueleaf/LyCORIS',
-    description='Lora beYond Conventional methods, Other Rank adaptation Implementations for Stable diffusion',
-    author='Shih-Ying Yeh(KohakuBlueLeaf), Yu-Guan Hsieh, Zhidong Gao',
-    author_email='apolloyeh0123@gmail.com',
+    version="2.0.0",
+    url="https://github.com/KohakuBlueleaf/LyCORIS",
+    description="Lora beYond Conventional methods, Other Rank adaptation Implementations for Stable diffusion",
+    author="Shih-Ying Yeh(KohakuBlueLeaf), Yu-Guan Hsieh, Zhidong Gao",
+    author_email="apolloyeh0123@gmail.com",
     zip_safe=False,
-    install_requires=[
-        'torch',
-        'safetensors',
-        'diffusers',
-        'transformers'
-    ],
+    install_requires=["torch", "safetensors", "diffusers", "transformers"],
 )
diff --git a/standalone_example.py b/standalone_example.py
new file mode 100644
index 0000000..f4e78b5
--- /dev/null
+++ b/standalone_example.py
@@ -0,0 +1,83 @@
+from itertools import chain
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.data as Data
+
+from torchvision.datasets import MNIST
+from torchvision.transforms import ToTensor, Lambda, Compose
+
+from lycoris import create_lycoris, LycorisNetwork
+
+
+class DemoNet(nn.Module):
+    """
+    A Simple Pytorch Module for demo
+    """
+    def __init__(self):
+        super().__init__()
+        self.test_1 = nn.Linear(784, 2048)
+        self.te_2st = nn.Linear(2048, 784)
+        self._3test = nn.Linear(784, 10)
+
+    def forward(self, x):
+        h = self.test_1(x)
+        h = F.mish(h)
+        h = self.te_2st(h)
+        h = x + h
+        h = self._3test(h)
+        return h
+
+
+# LyCORIS wrapper demo
+net = DemoNet()
+LycorisNetwork.apply_preset({"target_name": [".*te.*"]})
+lycoris_net1 = create_lycoris(net, 1.0, linear_dim=16, linear_alpha=2.0, algo="lokr")
+lycoris_net1.apply_to()
+
+LycorisNetwork.apply_preset({"target_name": [".*es.*"]})
+lycoris_net2 = create_lycoris(net, 1.0, linear_dim=16, linear_alpha=2.0, algo="lokr")
+lycoris_net2.apply_to()
+
+print(f"#Modules of net1: {len(lycoris_net1.loras)}")
+print(f"#Modules of net2: {len(lycoris_net2.loras)}")
+
+print("Total params:", sum(p.numel() for p in net.parameters()))
+print("Net1 Params:", sum(p.numel() for p in lycoris_net1.parameters()))
+print("Net2 Params:", sum(p.numel() for p in lycoris_net2.parameters()))
+
+
+# Training loop demo
+trns = Compose([ToTensor(), Lambda(lambda x: x.view(-1))])
+train_ds = MNIST(root="data", download=True, train=True, transform=trns)
+test_ds = MNIST(root="data", download=True, train=False, transform=trns)
+train_loader = Data.DataLoader(train_ds, batch_size=32, shuffle=True)
+test_loader = Data.DataLoader(test_ds, batch_size=32)
+
+optimizer = torch.optim.AdamW(
+    chain(lycoris_net1.parameters(), lycoris_net2.parameters()), lr=0.005
+)
+ema_loss = 0
+
+
+for i, (x, t) in enumerate(train_loader):
+    optimizer.zero_grad()
+    y = net(x)
+    loss = F.cross_entropy(y, t)
+    loss.backward()
+    optimizer.step()
+    ema_decay = min(0.999, i / 1000)
+    ema_loss = ema_decay * ema_loss + (1 - ema_decay) * loss.item()
+    if i % 100 == 0:
+        print(i, ema_loss)
+
+
+total_correct = 0
+total_num = 0
+for i, (x, t) in enumerate(test_loader):
+    y = net(x)
+    pred = y.argmax(dim=1)
+    total_correct += (pred == t).sum().item()
+    total_num += len(t)
+print(total_correct / total_num)
\ No newline at end of file
diff --git a/tools/batch_bundle_convert.py b/tools/batch_bundle_convert.py
new file mode 100644
index 0000000..f38a5d7
--- /dev/null
+++ b/tools/batch_bundle_convert.py
@@ -0,0 +1,290 @@
+import os
+import argparse
+from typing import List
+from collections import defaultdict
+
+import torch
+from torch import load, save
+from safetensors import safe_open
+from safetensors.torch import save_file
+
+
+def load_state_dict(file_path):
+    is_safetensors = file_path.rsplit(".", 1)[-1] == "safetensors"
+    if is_safetensors:
+        state_dict = {}
+        with safe_open(file_path, framework="pt", device="cpu") as f:
+            for key in f.keys():
+                state_dict[key] = f.get_tensor(key)
+    else:
+        state_dict = load(file_path)
+    return state_dict
+
+
+def save_state_dict(state, output_path):
+    if output_path.endswith(".safetensors"):
+        save_file(state, output_path)
+    else:
+        save(state, output_path)
+
+
+def pack_bundle(lora, emb_dict, verbose=False):
+    for emb, emb_sd in emb_dict.items():
+        for key, value in emb_sd.items():
+            if isinstance(value, dict):
+                for subkey, subvalue in value.items():
+                    lora[f"bundle_emb.{emb}.{key}.{subkey}"] = subvalue
+            elif isinstance(value, torch.Tensor):
+                lora[f"bundle_emb.{emb}.{key}"] = value
+    if verbose:
+        print("The following content has been added to lora")
+        for key, value in lora.items():
+            if key.startswith("bundle_emb"):
+                if isinstance(value, torch.Tensor):
+                    print(f" {key}: tensor of shape {value.shape}")
+                else:
+                    print(f" {key}: {value}")
+    return lora
+
+
+def unpack_bundle(lora, verbose, step="", emb_format=".pt"):
+    assert emb_format in [".pt", ".safetensors"]
+    if step != "":
+        step = "-" + str(step)
+    emb_dict = {}
+    bundle_keys = []
+    for lora_key, value in lora.items():
+        if lora_key.startswith("bundle_emb"):
+            bundle_keys.append(lora_key)
+            _, emb, *rest = lora_key.split(".")
+            emb = emb + step
+            if emb not in emb_dict:
+                emb_dict[emb] = {}
+            if len(rest) == 2:
+                key, subkey = rest
+                if emb_format == ".pt":
+                    if key not in emb_dict[emb]:
+                        emb_dict[emb][key] = {}
+                    emb_dict[emb][key][subkey] = value
+                else:
+                    emb_dict[emb][subkey] = value
+            elif len(rest) == 1:
+                key = rest[0]
+                emb_dict[emb][key] = value
+    for bundle_key in bundle_keys:
+        del lora[bundle_key]
+    if emb_format == ".pt":
+        for emb, emb_sd in emb_dict.items():
+            emb_sd["name"] = emb
+    if verbose:
+        print("The following embeddings have been loaded from bundle")
+        print_emb_information(emb_dict)
+    return lora, emb_dict
+
+
+def print_emb_information(emb_dict):
+    for emb, emb_sd in emb_dict.items():
+        print(emb)
+        for key, value in emb_sd.items():
+            if isinstance(value, dict):  # Check if the value is another dictionary
+                for subkey, subvalue in value.items():
+                    if isinstance(subvalue, torch.Tensor):
+                        print(f" {key}.{subkey}: tensor of shape" f" {subvalue.shape}")
+                    else:
+                        print(f"  {key}.{subkey}: {subvalue}")
+            elif isinstance(value, torch.Tensor):
+                print(f" {key}: tensor of shape {value.shape}")
+            else:
+                print(f" {key}: {value}")
+
+
+def extract_step(file_path):
+    filename = os.path.splitext(os.path.basename(file_path))[0]
+    step = filename.split("-")[-1].replace("step", "")
+    if step.isdigit():
+        name = "-".join(filename.split("-")[:-1])
+        return name, int(step)
+    else:
+        return name, ""
+
+
+def gather_files_from_list(
+    paths: List[str], extensions: List[str], recursive: bool
+) -> List[str]:
+    """
+    Gather files from given paths based on specific extensions.
+
+    :param paths: A list of paths which can be files or directories.
+    :type paths: List[str]
+    :param extensions: A list of file extensions to filter by.
+    :type extensions: List[str]
+    :param recursive: If True, search for files recursively in directories.
+    :type recursive: bool
+    :return: A list of file paths that match the given extensions.
+    :rtype: List[str]
+    """
+    files = []
+
+    def is_extension_valid(file: str) -> bool:
+        return any(file.endswith(ext) for ext in extensions)
+
+    def add_files_from_directory(directory: str):
+        for root, _, filenames in os.walk(directory):
+            for filename in filenames:
+                filepath = os.path.join(root, filename)
+                if is_extension_valid(filepath):
+                    files.append(filepath)
+            if not recursive:
+                break
+
+    for path in paths:
+        if os.path.isfile(path) and is_extension_valid(path):
+            files.append(path)
+        elif os.path.isdir(path):
+            add_files_from_directory(path)
+
+    return files
+
+
+def get_lora_embs_step_correspondance(lora_files: List[str], emb_files: List[str]):
+    """
+    Associate LoRA model files with embedding files based on their step count.
+
+    This function takes in lists of LoRA file paths and embedding file paths,
+    extracts their step counts, and associates them based on matching steps.
+    If a file's step count cannot be determined, we use the key 'none'.
+
+    :param lora_files: A list of file paths to LoRA model files.
+    :type lora_files: List[str]
+    :param emb_files: A list of file paths to embedding files.
+    :type emb_files: List[str]
+    :return: A dictionary where keys are step counts (or 'none') and values
+             are dictionaries containing 'lora' (path to the LoRA model) and
+             'embs' (a list of paths to associated embedding files).
+    :rtype: Dict[str, Dict[str, Union[str, List[str]]]]
+    """
+    lora_embs = defaultdict(lambda: {"lora": None, "embs": []})
+    for lora_path in lora_files:
+        _, step = extract_step(lora_path)
+        if step in lora_embs:
+            raise ValueError(
+                "Find two Lora files with the same" f" step count {step}, abort"
+            )
+        lora_embs[step]["lora"] = lora_path
+    for emb_path in emb_files:
+        _, step = extract_step(emb_path)
+        if step in lora_embs:
+            lora_embs[step]["embs"].append(emb_path)
+        else:
+            print(f"Warning: no corresponding lora found for {emb_path}")
+    return lora_embs
+
+
+def convert_lora_name(lora_path, dst_dir, to_bundle):
+    name, step = extract_step(lora_path)
+    if step != "":
+        step = "-" + str(step)
+    if to_bundle:
+        name = name + "-bundle"
+    elif name.endswith("-bundle"):
+        name = name[:-7]
+    lora_save_path = os.path.join(dst_dir, name + step + os.path.splitext(lora_path)[1])
+    return lora_save_path
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Tool for packing and unpacking LoRA and embeddings."
+    )
+    parser.add_argument(
+        "--lora_path",
+        default=[],
+        type=str,
+        nargs="+",
+        help="Paths to LoRA model files.",
+    )
+    parser.add_argument(
+        "--emb_path", default=[], type=str, nargs="+", help="Paths to embedding files."
+    )
+    parser.add_argument(
+        "--dst_dir",
+        default=None,
+        type=str,
+        help="Destination directory for output files.",
+    )
+    parser.add_argument(
+        "--from_bundle", action="store_true", help="Unpack from bundle."
+    )
+    parser.add_argument("--to_bundle", action="store_true", help="Pack to bundle.")
+    parser.add_argument(
+        "--lora_ext",
+        default=[".safetensors"],
+        type=str,
+        nargs="+",
+        help="Extensions for LoRA files.",
+    )
+    parser.add_argument(
+        "--emb_ext",
+        default=[".pt"],
+        type=str,
+        nargs="+",
+        help="Extensions for embedding files.",
+    )
+    parser.add_argument(
+        "--recursive",
+        action="store_true",
+        help="Recursively search for files in directories.",
+    )
+    parser.add_argument(
+        "--pack_all_embeddings",
+        action="store_true",
+        help=(
+            "Pack all embeddings to all LoRA files"
+            " instead of using step correspondance."
+        ),
+    )
+    parser.add_argument("--verbose", default=1, type=int, help="Verbosity level.")
+    args = parser.parse_args()
+
+    lora_paths = gather_files_from_list(args.lora_path, args.lora_ext, args.recursive)
+
+    if args.from_bundle:
+        dst_dir = "bundles_unpack" if args.dst_dir is None else args.dst_dir
+        os.makedirs(dst_dir, exist_ok=True)
+        for lora_path in lora_paths:
+            if args.verbose >= 1:
+                print(f"Unpacking {lora_path}")
+            lora = load_state_dict(lora_path)
+            _, step = extract_step(lora_path)
+            lora, emb_dict = unpack_bundle(
+                lora, args.verbose >= 2, step=step, emb_format=args.emb_ext[0]
+            )
+            lora_save_path = convert_lora_name(lora_path, dst_dir, to_bundle=False)
+            save_state_dict(lora, lora_save_path)
+            for emb, emb_sd in emb_dict.items():
+                emb_save_path = os.path.join(dst_dir, emb + args.emb_ext[0])
+                save_state_dict(emb_sd, emb_save_path)
+    elif args.to_bundle:
+        if args.emb_path == []:
+            args.emb_path = args.lora_path
+        emb_paths = gather_files_from_list(args.emb_path, args.emb_ext, args.recursive)
+        dst_dir = "bundles" if args.dst_dir is None else args.dst_dir
+        os.makedirs(dst_dir, exist_ok=True)
+        lora_embs_dict = {}
+        if args.pack_all_embeddings:
+            for i, lora_path in enumerate(lora_paths):
+                lora_embs_dict[i] = {"lora": lora_path, "embs": emb_paths}
+        else:
+            lora_embs_dict = get_lora_embs_step_correspondance(lora_paths, emb_paths)
+        for _, lora_embs_pair in lora_embs_dict.items():
+            lora_path = lora_embs_pair["lora"]
+            if args.verbose >= 1:
+                print(f"Packing {lora_path}")
+            lora = load_state_dict(lora_path)
+            emb_dict = {}
+            for emb_path in lora_embs_pair["embs"]:
+                name, _ = extract_step(emb_path)
+                emb_dict[name] = load_state_dict(emb_path)
+            bundle = pack_bundle(lora, emb_dict, args.verbose >= 2)
+            lora_save_path = convert_lora_name(lora_path, dst_dir, to_bundle=True)
+            save_state_dict(bundle, lora_save_path)
diff --git a/tools/batch_hcp_convert.py b/tools/batch_hcp_convert.py
new file mode 100644
index 0000000..9d659fe
--- /dev/null
+++ b/tools/batch_hcp_convert.py
@@ -0,0 +1,429 @@
+import os
+import math
+import argparse
+from typing import List
+from collections import defaultdict
+
+from hcpdiff.ckpt_manager import auto_manager
+
+
+class LoraConverter(object):
+    com_name_unet = [
+        "down_blocks",
+        "up_blocks",
+        "mid_block",
+        "transformer_blocks",
+        "to_q",
+        "to_k",
+        "to_v",
+        "to_out",
+        "proj_in",
+        "proj_out",
+        "input_blocks",
+        "middle_block",
+        "output_blocks",
+    ]
+    com_name_TE = ["self_attn", "q_proj", "v_proj", "k_proj", "out_proj", "text_model"]
+    prefix_unet = "lora_unet_"
+    prefix_TE = "lora_te_"
+
+    def __init__(self):
+        self.com_name_unet_tmp = [x.replace("_", "%") for x in self.com_name_unet]
+        self.com_name_TE_tmp = [x.replace("_", "%") for x in self.com_name_TE]
+
+    def convert_from_webui(
+        self, state, network_type="lora", auto_scale_alpha=False, sdxl=False
+    ):
+        assert network_type in ["lora", "plugin"]
+        if not sdxl:
+            sd_unet = self.convert_from_webui_(
+                state,
+                network_type=network_type,
+                prefix=self.prefix_unet,
+                com_name=self.com_name_unet,
+                com_name_tmp=self.com_name_unet_tmp,
+            )
+            sd_te = self.convert_from_webui_(
+                state,
+                network_type=network_type,
+                prefix=self.prefix_TE,
+                com_name=self.com_name_TE,
+                com_name_tmp=self.com_name_TE_tmp,
+            )
+        else:
+            sd_unet = self.convert_from_webui_xl_unet_(
+                state,
+                network_type=network_type,
+                prefix=self.prefix_unet,
+                com_name=self.com_name_unet,
+                com_name_tmp=self.com_name_unet_tmp,
+            )
+            sd_te = self.convert_from_webui_xl_te_(
+                state,
+                network_type=network_type,
+                prefix=self.prefix_TE_xl_clip_B,
+                com_name=self.com_name_TE,
+                com_name_tmp=self.com_name_TE_tmp,
+            )
+            sd_te2 = self.convert_from_webui_xl_te_(
+                state,
+                network_type=network_type,
+                prefix=self.prefix_TE_xl_clip_bigG,
+                com_name=self.com_name_TE,
+                com_name_tmp=self.com_name_TE_tmp,
+            )
+            sd_te.update(sd_te2)
+        if auto_scale_alpha and network_type == "lora":
+            sd_unet = self.alpha_scale_from_webui(sd_unet)
+            sd_te = self.alpha_scale_from_webui(sd_te)
+        return {network_type: sd_te}, {network_type: sd_unet}
+
+    def convert_to_webui(
+        self, sd_unet, sd_te, network_type="lora", auto_scale_alpha=False, sdxl=False
+    ):
+        assert network_type in ["lora", "plugin"]
+        sd_unet = self.convert_to_webui_(
+            sd_unet, network_type=network_type, prefix=self.prefix_unet
+        )
+        if sdxl:
+            sd_te = self.convert_to_webui_xl_(
+                sd_te, network_type=network_type, prefix=self.prefix_TE
+            )
+        else:
+            sd_te = self.convert_to_webui_(
+                sd_te, network_type=network_type, prefix=self.prefix_TE
+            )
+        sd_unet.update(sd_te)
+        if auto_scale_alpha and network_type == "lora":
+            sd_unet = self.alpha_scale_to_webui(sd_unet)
+        return sd_unet
+
+    def convert_from_webui_(self, state, network_type, prefix, com_name, com_name_tmp):
+        state = {k: v for k, v in state.items() if k.startswith(prefix)}
+        prefix_len = len(prefix)
+        sd_covert = {}
+        for k, v in state.items():
+            model_k, lora_k = k[prefix_len:].split(".", 1)
+            model_k = (
+                self.replace_all(model_k, com_name, com_name_tmp)
+                .replace("_", ".")
+                .replace("%", "_")
+            )
+            if lora_k == "alpha" or network_type == "plugin":
+                sd_covert[f"{model_k}.___.{lora_k}"] = v
+            else:
+                sd_covert[f"{model_k}.___.layer.{lora_k}"] = v
+        return sd_covert
+
+    def convert_to_webui_(self, state, network_type, prefix):
+        sd_covert = {}
+        for k, v in state.items():
+            if network_type == "plugin" or "alpha" in k or "scale" in k:
+                separator = ".___."
+            else:
+                separator = ".___.layer."
+            model_k, lora_k = k.split(separator, 1)
+            sd_covert[f"{prefix}{model_k.replace('.', '_')}.{lora_k}"] = v
+        return sd_covert
+
+    def convert_to_webui_xl_(self, state, network_type, prefix):
+        sd_convert = {}
+        for k, v in state.items():
+            if network_type == "plugin" or "alpha" in k or "scale" in k:
+                separator = ".___."
+            else:
+                separator = ".___.layer."
+            model_k, lora_k = k.split(separator, 1)
+            new_k = f"{prefix}{model_k.replace('.', '_')}.{lora_k}"
+            if "clip" in new_k:
+                new_k = (
+                    new_k.replace("_clip_B", "1")
+                    if "clip_B" in new_k
+                    else new_k.replace("_clip_bigG", "2")
+                )
+            sd_convert[new_k] = v
+        return sd_convert
+
+    def convert_from_webui_xl_te_(
+        self, state, network_type, prefix, com_name, com_name_tmp
+    ):
+        state = {k: v for k, v in state.items() if k.startswith(prefix)}
+        sd_covert = {}
+        prefix_len = len(prefix)
+
+        for k, v in state.items():
+            model_k, lora_k = k[prefix_len:].split(".", 1)
+            model_k = (
+                self.replace_all(model_k, com_name, com_name_tmp)
+                .replace("_", ".")
+                .replace("%", "_")
+            )
+            if prefix == "lora_te1_":
+                model_k = f"clip_B.{model_k}"
+            else:
+                model_k = f"clip_bigG.{model_k}"
+
+            if lora_k == "alpha" or network_type == "plugin":
+                sd_covert[f"{model_k}.___.{lora_k}"] = v
+            else:
+                sd_covert[f"{model_k}.___.layer.{lora_k}"] = v
+        return sd_covert
+
+    def convert_from_webui_xl_unet_(
+        self, state, network_type, prefix, com_name, com_name_tmp
+    ):
+        # Down:
+        # 4 -> 1, 0  4 = 1 + 3 * 1 + 0
+        # 5 -> 1, 1  5 = 1 + 3 * 1 + 1
+        # 7 -> 2, 0  7 = 1 + 3 * 2 + 0
+        # 8 -> 2, 1  8 = 1 + 3 * 2 + 1
+
+        # Up
+        # 0 -> 0, 0  0 = 0 * 3 + 0
+        # 1 -> 0, 1  1 = 0 * 3 + 1
+        # 2 -> 0, 2  2 = 0 * 3 + 2
+        # 3 -> 1, 0  3 = 1 * 3 + 0
+        # 4 -> 1, 1  4 = 1 * 3 + 1
+        # 5 -> 1, 2  5 = 1 * 3 + 2
+
+        down = {
+            "4": [1, 0],
+            "5": [1, 1],
+            "7": [2, 0],
+            "8": [2, 1],
+        }
+        up = {
+            "0": [0, 0],
+            "1": [0, 1],
+            "2": [0, 2],
+            "3": [1, 0],
+            "4": [1, 1],
+            "5": [1, 2],
+        }
+
+        import re
+
+        m = []
+
+        def match(key, regex_text):
+            regex = re.compile(regex_text)
+            r = re.match(regex, key)
+            if not r:
+                return False
+
+            m.clear()
+            m.extend(r.groups())
+            return True
+
+        state = {k: v for k, v in state.items() if k.startswith(prefix)}
+        sd_covert = {}
+        prefix_len = len(prefix)
+        for k, v in state.items():
+            model_k, lora_k = k[prefix_len:].split(".", 1)
+
+            model_k = (
+                self.replace_all(model_k, com_name, com_name_tmp)
+                .replace("_", ".")
+                .replace("%", "_")
+            )
+
+            if match(model_k, r"input_blocks.(\d+).1.(.+)"):
+                new_k = (
+                    f"down_blocks.{down[m[0]][0]}.attentions" f".{down[m[0]][1]}.{m[1]}"
+                )
+            elif match(model_k, r"middle_block.1.(.+)"):
+                new_k = f"mid_block.attentions.0.{m[0]}"
+                pass
+            elif match(model_k, r"output_blocks.(\d+).(\d+).(.+)"):
+                new_k = f"up_blocks.{up[m[0]][0]}.attentions" f".{up[m[0]][1]}.{m[2]}"
+            else:
+                raise NotImplementedError
+
+            if lora_k == "alpha" or network_type == "plugin":
+                sd_covert[f"{new_k}.___.{lora_k}"] = v
+            else:
+                sd_covert[f"{new_k}.___.layer.{lora_k}"] = v
+
+        return sd_covert
+
+    @staticmethod
+    def replace_all(data: str, srcs: List[str], dsts: List[str]):
+        for src, dst in zip(srcs, dsts):
+            data = data.replace(src, dst)
+        return data
+
+    @staticmethod
+    def alpha_scale_from_webui(state):
+        # Apply to "lora_down" and "lora_up" respectively to prevent overflow
+        for k, v in state.items():
+            if "lora_up" in k:
+                state[k] = v * math.sqrt(v.shape[1])
+            elif "lora_down" in k:
+                state[k] = v * math.sqrt(v.shape[0])
+        return state
+
+    @staticmethod
+    def alpha_scale_to_webui(state):
+        for k, v in state.items():
+            if "lora_up" in k:
+                state[k] = v * math.sqrt(v.shape[1])
+            elif "lora_down" in k:
+                state[k] = v * math.sqrt(v.shape[0])
+        return state
+
+
+def save_and_print_path(sd, path):
+    os.makedirs(args.dump_path, exist_ok=True)
+    ckpt_manager._save_ckpt(sd, save_path=path)
+    print("Saved to:", path)
+
+
+def get_unet_te_pairs(lora_path):
+    file_pairs = defaultdict(lambda: {"TE": None, "unet": None})
+    for filename in os.listdir(lora_path):
+        if filename.endswith(".safetensors"):
+            parts = os.path.splitext(filename)[0].split("-")
+            prefix, name = parts[0], "-".join(parts[1:])
+
+            if "text_encoder" in prefix:
+                file_pairs[name]["TE"] = os.path.join(args.lora_path, filename)
+            elif "unet" in prefix:
+                file_pairs[name]["unet"] = os.path.join(args.lora_path, filename)
+    return file_pairs
+
+
+def get_network_type(sd_unet, sd_te):
+    if "lora" in sd_unet.keys() and "lora" in sd_te.keys():
+        return "lora"
+    elif "plugin" in sd_unet.keys() and "plugin" in sd_te.keys():
+        return "plugin"
+    else:
+        return None
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Convert LoRA models.")
+    parser.add_argument(
+        "--lora_path",
+        required=True,
+        type=str,
+        help="Path to the LoRA or folder containing LoRA models.",
+    )
+    parser.add_argument(
+        "--lora_path_TE", type=str, help="Path to the HCP Text Encoder LoRA."
+    )
+    parser.add_argument(
+        "--dump_path",
+        required=True,
+        type=str,
+        help="Path to save the converted state dict.",
+    )
+    parser.add_argument(
+        "--from_webui", action="store_true", help="Convert from webui format."
+    )
+    parser.add_argument(
+        "--save_network_type", type=str, help="Specify the network type for conversion."
+    )
+    parser.add_argument(
+        "--to_webui", action="store_true", help="Convert to webui format."
+    )
+    parser.add_argument(
+        "--output_prefix", default="", type=str, help="Prefix for output filenames."
+    )
+    parser.add_argument(
+        "--auto_scale_alpha", action="store_true", help="Automatically scale alpha."
+    )
+    parser.add_argument("--sdxl", action="store_true", help="Enable SDXL conversion.")
+    args = parser.parse_args()
+
+    converter = LoraConverter()
+
+    if os.path.isdir(args.lora_path):
+        ckpt_manager = auto_manager(".safetensors")()
+        if args.from_webui:
+            for filename in os.listdir(args.lora_path):
+                if filename.endswith(".safetensors"):
+                    file_path = os.path.join(args.lora_path, filename)
+                    print(f"Converting {file_path}")
+                    state = ckpt_manager.load_ckpt(file_path)
+                    sd_te, sd_unet = converter.convert_from_webui(
+                        state,
+                        network_type=args.save_network_type,
+                        auto_scale_alpha=args.auto_scale_alpha,
+                        sdxl=args.sdxl,
+                    )
+
+                    TE_path = os.path.join(args.dump_path, "text_encoder-" + filename)
+                    unet_path = os.path.join(args.dump_path, "unet-" + filename)
+                    save_and_print_path(sd_te, TE_path)
+                    save_and_print_path(sd_unet, unet_path)
+
+        elif args.to_webui:
+            file_pairs = get_unet_te_pairs(args.lora_path)
+
+            for name, paths in file_pairs.items():
+                if paths["TE"] and paths["unet"]:
+                    sd_unet = ckpt_manager.load_ckpt(paths["unet"])
+                    sd_te = ckpt_manager.load_ckpt(paths["TE"])
+                    network_type = get_network_type(sd_unet, sd_te)
+                    if network_type is None:
+                        print("no saved lora/lycoris found, skip")
+                        continue
+                    print(
+                        f'Converting pair: {paths["TE"]} and {paths["unet"]}'
+                        f' with key "{network_type}"'
+                    )
+                    state = converter.convert_to_webui(
+                        sd_unet[network_type],
+                        sd_te[network_type],
+                        network_type=network_type,
+                        auto_scale_alpha=args.auto_scale_alpha,
+                        sdxl=args.sdxl,
+                    )
+
+                    output_path = os.path.join(
+                        args.dump_path, f"{args.output_prefix}-{name}.safetensors"
+                    )
+                    save_and_print_path(state, output_path)
+
+    else:
+        print("Converting LoRA model")
+        ckpt_manager = auto_manager(args.lora_path)()
+        if args.from_webui:
+            state = ckpt_manager.load_ckpt(args.lora_path)
+            sd_te, sd_unet = converter.convert_from_webui(
+                state,
+                network_type=args.save_network_type,
+                auto_scale_alpha=args.auto_scale_alpha,
+                sdxl=args.sdxl,
+            )
+            TE_path = os.path.join(
+                args.dump_path, "text_encoder-" + os.path.basename(args.lora_path)
+            )
+            unet_path = os.path.join(
+                args.dump_path, "unet-" + os.path.basename(args.lora_path)
+            )
+            save_and_print_path(sd_te, TE_path)
+            save_and_print_path(sd_unet, unet_path)
+        elif args.to_webui:
+            sd_unet = ckpt_manager.load_ckpt(args.lora_path)
+            sd_te = ckpt_manager.load_ckpt(args.lora_path_TE)
+            network_type = get_network_type(sd_unet, sd_te)
+            if network_type is None:
+                print("no saved lora/lycoris found, terminating")
+                exit(1)
+            print(f'Converting with key "{network_type}"')
+            state = converter.convert_to_webui(
+                sd_unet[network_type],
+                sd_te[network_type],
+                network_type=network_type,
+                auto_scale_alpha=args.auto_scale_alpha,
+                sdxl=args.sdxl,
+            )
+            lora_name = os.path.basename(args.lora_path)
+            if "-" in lora_name:
+                lora_name = "-".join(lora_name.split("-")[1:])
+            output_path = os.path.join(
+                args.dump_path, args.output_prefix + "-" + lora_name
+            )
+            save_and_print_path(state, output_path)
diff --git a/tools/extract_locon.py b/tools/extract_locon.py
index 8610ed8..28b25ee 100644
--- a/tools/extract_locon.py
+++ b/tools/extract_locon.py
@@ -1,4 +1,5 @@
 import os, sys
+
 sys.path.insert(0, os.getcwd())
 import argparse
 
@@ -6,87 +7,125 @@
 def get_args():
     parser = argparse.ArgumentParser()
     parser.add_argument(
-        "base_model", help="The model which use it to train the dreambooth model",
-        default='', type=str
+        "base_model",
+        help="The model which use it to train the dreambooth model",
+        default="",
+        type=str,
+    )
+    parser.add_argument(
+        "db_model",
+        help="the dreambooth model you want to extract the locon",
+        default="",
+        type=str,
     )
     parser.add_argument(
-        "db_model", help="the dreambooth model you want to extract the locon",
-        default='', type=str
+        "output_name", help="the output model", default="./out.pt", type=str
     )
     parser.add_argument(
-        "output_name", help="the output model",
-        default='./out.pt', type=str
+        "--is_v2",
+        help="Your base/db model is sd v2 or not",
+        default=False,
+        action="store_true",
     )
     parser.add_argument(
-        "--is_v2", help="Your base/db model is sd v2 or not",
-        default=False, action="store_true"
+        "--is_sdxl",
+        help="Your base/db model is sdxl or not",
+        default=False,
+        action="store_true",
     )
     parser.add_argument(
-        "--device", help="Which device you want to use to extract the locon",
-        default='cpu', type=str
+        "--device",
+        help="Which device you want to use to extract the locon",
+        default="cpu",
+        type=str,
     )
     parser.add_argument(
-        "--mode", 
+        "--mode",
         help=(
-            'extraction mode, can be "fixed", "threshold", "ratio", "quantile". '
+            'extraction mode, can be "full", "fixed", "threshold", "ratio", "quantile". '
             'If not "fixed", network_dim and conv_dim will be ignored'
         ),
-        default='fixed', type=str
+        default="fixed",
+        type=str,
     )
     parser.add_argument(
-        "--safetensors", help='use safetensors to save locon model',
-        default=False, action="store_true"
+        "--safetensors",
+        help="use safetensors to save locon model",
+        default=False,
+        action="store_true",
     )
     parser.add_argument(
-        "--linear_dim", help="network dim for linear layer in fixed mode",
-        default=1, type=int
+        "--linear_dim",
+        help="network dim for linear layer in fixed mode",
+        default=1,
+        type=int,
     )
     parser.add_argument(
-        "--conv_dim", help="network dim for conv layer in fixed mode",
-        default=1, type=int
+        "--conv_dim",
+        help="network dim for conv layer in fixed mode",
+        default=1,
+        type=int,
     )
     parser.add_argument(
-        "--linear_threshold", help="singular value threshold for linear layer in threshold mode",
-        default=0., type=float
+        "--linear_threshold",
+        help="singular value threshold for linear layer in threshold mode",
+        default=0.0,
+        type=float,
     )
     parser.add_argument(
-        "--conv_threshold", help="singular value threshold for conv layer in threshold mode",
-        default=0., type=float
+        "--conv_threshold",
+        help="singular value threshold for conv layer in threshold mode",
+        default=0.0,
+        type=float,
     )
     parser.add_argument(
-        "--linear_ratio", help="singular ratio for linear layer in ratio mode",
-        default=0., type=float
+        "--linear_ratio",
+        help="singular ratio for linear layer in ratio mode",
+        default=0.0,
+        type=float,
     )
     parser.add_argument(
-        "--conv_ratio", help="singular ratio for conv layer in ratio mode",
-        default=0., type=float
+        "--conv_ratio",
+        help="singular ratio for conv layer in ratio mode",
+        default=0.0,
+        type=float,
     )
     parser.add_argument(
-        "--linear_quantile", help="singular value quantile for linear layer quantile mode",
-        default=1., type=float
+        "--linear_quantile",
+        help="singular value quantile for linear layer quantile mode",
+        default=1.0,
+        type=float,
     )
     parser.add_argument(
-        "--conv_quantile", help="singular value quantile for conv layer quantile mode",
-        default=1., type=float
+        "--conv_quantile",
+        help="singular value quantile for conv layer quantile mode",
+        default=1.0,
+        type=float,
     )
     parser.add_argument(
-        "--use_sparse_bias", help="enable sparse bias",
-        default=False, action="store_true"
+        "--use_sparse_bias",
+        help="enable sparse bias",
+        default=False,
+        action="store_true",
     )
     parser.add_argument(
-        "--sparsity", help="sparsity for sparse bias",
-        default=0.98, type=float
+        "--sparsity", help="sparsity for sparse bias", default=0.98, type=float
     )
     parser.add_argument(
-        "--disable_cp", help="don't use cp decomposition",
-        default=False, action="store_true"
+        "--disable_cp",
+        help="don't use cp decomposition",
+        default=False,
+        action="store_true",
     )
     return parser.parse_args()
+
+
 ARGS = get_args()
 
 
 from lycoris.utils import extract_diff
 from lycoris.kohya.model_utils import load_models_from_stable_diffusion_checkpoint
+from lycoris.kohya.sdxl_model_util import load_models_from_sdxl_checkpoint
 
 import torch
 from safetensors.torch import save_file
@@ -94,36 +133,58 @@ def get_args():
 
 def main():
     args = ARGS
-    base = load_models_from_stable_diffusion_checkpoint(args.is_v2, args.base_model)
-    db = load_models_from_stable_diffusion_checkpoint(args.is_v2, args.db_model)
-    
+    if args.is_sdxl:
+        base = load_models_from_sdxl_checkpoint(None, args.base_model, args.device)
+        db = load_models_from_sdxl_checkpoint(None, args.db_model, args.device)
+    else:
+        base = load_models_from_stable_diffusion_checkpoint(args.is_v2, args.base_model)
+        db = load_models_from_stable_diffusion_checkpoint(args.is_v2, args.db_model)
+
     linear_mode_param = {
-        'fixed': args.linear_dim,
-        'threshold': args.linear_threshold,
-        'ratio': args.linear_ratio,
-        'quantile': args.linear_quantile,
+        "fixed": args.linear_dim,
+        "threshold": args.linear_threshold,
+        "ratio": args.linear_ratio,
+        "quantile": args.linear_quantile,
+        "full": None,
     }[args.mode]
     conv_mode_param = {
-        'fixed': args.conv_dim,
-        'threshold': args.conv_threshold,
-        'ratio': args.conv_ratio,
-        'quantile': args.conv_quantile,
+        "fixed": args.conv_dim,
+        "threshold": args.conv_threshold,
+        "ratio": args.conv_ratio,
+        "quantile": args.conv_quantile,
+        "full": None,
     }[args.mode]
-    
+
+    if args.is_sdxl:
+        db_tes = [db[0], db[1]]
+        db_unet = db[3]
+        base_tes = [base[0], base[1]]
+        base_unet = base[3]
+    else:
+        db_tes = [db[0]]
+        db_unet = db[2]
+        base_tes = [base[0]]
+        base_unet = base[2]
+
     state_dict = extract_diff(
-        base, db,
+        base_tes,
+        db_tes,
+        base_unet,
+        db_unet,
         args.mode,
-        linear_mode_param, conv_mode_param,
-        args.device, 
-        args.use_sparse_bias, args.sparsity,
-        not args.disable_cp
+        linear_mode_param,
+        conv_mode_param,
+        args.device,
+        args.use_sparse_bias,
+        args.sparsity,
+        not args.disable_cp,
     )
-    
+
     if args.safetensors:
         save_file(state_dict, args.output_name)
     else:
         torch.save(state_dict, args.output_name)
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     main()
diff --git a/tools/merge.py b/tools/merge.py
index 8662ddc..25f673d 100644
--- a/tools/merge.py
+++ b/tools/merge.py
@@ -1,4 +1,5 @@
 import os, sys
+
 sys.path.insert(0, os.getcwd())
 import argparse
 
@@ -6,42 +7,54 @@
 def get_args():
     parser = argparse.ArgumentParser()
     parser.add_argument(
-        "base_model", help="The model you want to merge with loha",
-        default='', type=str
+        "base_model", help="The model you want to merge with loha", default="", type=str
     )
     parser.add_argument(
-        "lycoris_model", help="the lyco model you want to merge into sd model",
-        default='', type=str
+        "lycoris_model",
+        help="the lyco model you want to merge into sd model",
+        default="",
+        type=str,
     )
     parser.add_argument(
-        "output_name", help="the output model",
-        default='./out.pt', type=str
+        "output_name", help="the output model", default="./out.pt", type=str
     )
     parser.add_argument(
-        "--is_v2", help="Your base model is sd v2 or not",
-        default=False, action="store_true"
+        "--is_v2",
+        help="Your base model is sd v2 or not",
+        default=False,
+        action="store_true",
     )
     parser.add_argument(
-        "--device", help="Which device you want to use to merge the weight",
-        default='cpu', type=str
+        "--is_sdxl",
+        help="Your base/db model is sdxl or not",
+        default=False,
+        action="store_true",
     )
     parser.add_argument(
-        "--dtype", help='dtype to save',
-        default='float', type=str
+        "--device",
+        help="Which device you want to use to merge the weight",
+        default="cpu",
+        type=str,
     )
+    parser.add_argument("--dtype", help="dtype to save", default="float", type=str)
     parser.add_argument(
-        "--weight", help='weight for the lyco model to merge',
-        default='1.0', type=float
+        "--weight", help="weight for the lyco model to merge", default="1.0", type=float
     )
     return parser.parse_args()
-ARGS = get_args()
+
+
+args = ARGS = get_args()
 
 
 from lycoris.utils import merge
 from lycoris.kohya.model_utils import (
     load_models_from_stable_diffusion_checkpoint,
     save_stable_diffusion_checkpoint,
-    load_file
+    load_file,
+)
+from lycoris.kohya.sdxl_model_util import (
+    load_models_from_sdxl_checkpoint,
+    save_stable_diffusion_checkpoint as save_sdxl_checkpoint,
 )
 
 import torch
@@ -49,38 +62,64 @@ def get_args():
 
 @torch.no_grad()
 def main():
-    base = load_models_from_stable_diffusion_checkpoint(ARGS.is_v2, ARGS.base_model)
-    if ARGS.lycoris_model.rsplit('.', 1)[-1] == 'safetensors':
+    if args.is_sdxl:
+        base = load_models_from_sdxl_checkpoint(
+            None, args.base_model, map_location=args.device
+        )
+    else:
+        base = load_models_from_stable_diffusion_checkpoint(args.is_v2, args.base_model)
+    if ARGS.lycoris_model.rsplit(".", 1)[-1] == "safetensors":
         lyco = load_file(ARGS.lycoris_model)
     else:
         lyco = torch.load(ARGS.lycoris_model)
-    
-    dtype_str = ARGS.dtype.replace('fp', 'float').replace('bf', 'bfloat')
+
+    dtype_str = ARGS.dtype.replace("fp", "float").replace("bf", "bfloat")
     dtype = {
-        'float': torch.float,
-        'float16': torch.float16,
-        'float32': torch.float32,
-        'float64': torch.float64,
-        'bfloat': torch.bfloat16,
-        'bfloat16': torch.bfloat16,
+        "float": torch.float,
+        "float16": torch.float16,
+        "float32": torch.float32,
+        "float64": torch.float64,
+        "bfloat": torch.bfloat16,
+        "bfloat16": torch.bfloat16,
     }.get(dtype_str, None)
     if dtype is None:
         raise ValueError(f'Cannot Find the dtype "{dtype}"')
-    
-    merge(
-        base,
-        lyco,
-        ARGS.weight,
-        ARGS.device
-    )
-    
-    save_stable_diffusion_checkpoint(
-        ARGS.is_v2, ARGS.output_name, 
-        base[0].cpu(), base[2].cpu(), 
-        None, 0, 0, dtype, 
-        base[1]
-    )
+
+    if args.is_sdxl:
+        base_tes = [base[0], base[1]]
+        base_unet = base[3]
+    else:
+        base_tes = [base[0]]
+        base_unet = base[2]
+
+    merge(base_tes, base_unet, lyco, ARGS.weight, ARGS.device)
+
+    if args.is_sdxl:
+        save_sdxl_checkpoint(
+            ARGS.output_name,
+            base[0].cpu(),
+            base[1].cpu(),
+            base[3].cpu(),
+            0,
+            0,
+            None,
+            base[2],
+            getattr(base[1], "logit_scale", None),
+            dtype,
+        )
+    else:
+        save_stable_diffusion_checkpoint(
+            ARGS.is_v2,
+            ARGS.output_name,
+            base[0].cpu(),
+            base[2].cpu(),
+            None,
+            0,
+            0,
+            dtype,
+            base[1],
+        )
 
 
-if __name__ == '__main__':
-    main()
\ No newline at end of file
+if __name__ == "__main__":
+    main()
diff --git a/tools/pack_bundle.py b/tools/pack_bundle.py
new file mode 100644
index 0000000..64ad572
--- /dev/null
+++ b/tools/pack_bundle.py
@@ -0,0 +1,94 @@
+# See: https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13568
+
+import os, sys
+
+sys.path.insert(0, os.getcwd())
+import argparse
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "lora_model",
+        help="The model you want to pack embeddings into it.",
+        default="",
+        type=str,
+    )
+    parser.add_argument(
+        "embeddings",
+        help="the embeddings you want to pack into the lora model.",
+        default="",
+        type=str,
+        nargs="+",
+    )
+    parser.add_argument(
+        "output", help="output file name (with path)", default="", type=str, nargs="*"
+    )
+    parser.add_argument(
+        "--safetensors",
+        help="use safetensors to save bundled file",
+        default=True,
+        action="store_true",
+    )
+    return parser.parse_args()
+
+
+ARGS = get_args()
+
+import torch
+from torch import load, save
+from safetensors import safe_open
+from safetensors.torch import save_file
+
+
+def load_state_dict(file_path):
+    is_safetensors = file_path.rsplit(".", 1)[-1] == "safetensors"
+    if is_safetensors:
+        state_dict = {}
+        with safe_open(file_path, framework="pt", device="cpu") as f:
+            for key in f.keys():
+                state_dict[key] = f.get_tensor(key)
+    else:
+        state_dict = load(file_path)
+    return state_dict
+
+
+def main():
+    args = ARGS
+    is_safetensors = args.lora_model.rsplit(".", 1)[-1] == "safetensors"
+
+    lora_sd = load_state_dict(args.lora_model)
+    embs_sd = {
+        os.path.splitext(os.path.basename(x))[0]: load_state_dict(x)
+        for x in args.embeddings
+    }
+
+    for emb, emb_sd in embs_sd.items():
+        for key, value in emb_sd.items():
+            if isinstance(value, dict):
+                for subkey, subvalue in value.items():
+                    lora_sd[f"bundle_emb.{emb}.{key}.{subkey}"] = subvalue
+            elif isinstance(value, torch.Tensor):
+                lora_sd[f"bundle_emb.{emb}.{key}"] = value
+
+    for k in lora_sd:
+        if k.startswith("bundle_emb"):
+            print(k)
+
+    if args.output:
+        output_name = args.output
+    else:
+        output_name = f"{os.path.splitext(args.lora_model)[0]}_bundle"
+        if args.safetensors:
+            output_name += ".safetensors"
+        else:
+            output_name += ".pt"
+
+    if is_safetensors:
+        save_file(lora_sd, output_name)
+    else:
+        save(lora_sd, output_name)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/train_hyperdream.py b/train_hyperdream.py
index 44bff7f..8c699b6 100644
--- a/train_hyperdream.py
+++ b/train_hyperdream.py
@@ -1,6 +1,6 @@
-'''
+"""
 Modified from kohya-ss/sd-scripts/train_network.py
-'''
+"""
 
 import importlib
 import argparse
@@ -44,10 +44,11 @@
 # from functools import partial
 # import torch.utils.checkpoint
 # torch.utils.checkpoint.checkpoint = partial(
-#     torch.utils.checkpoint.checkpoint, 
+#     torch.utils.checkpoint.checkpoint,
 #     use_reentrant=False
 # )
 
+
 def patch_getitem(self, index):
     bucket = self.bucket_manager.buckets[self.buckets_indices[index].bucket_index]
     bucket_batch_size = self.buckets_indices[index].bucket_batch_size
@@ -69,26 +70,36 @@ def patch_getitem(self, index):
         subset = self.image_to_subset[image_key]
         loss_weights.append(self.prior_loss_weight if image_info.is_reg else 1.0)
 
-        flipped = subset.flip_aug and random.random() < 0.5  # not flipped or flipped with 50% chance
+        flipped = (
+            subset.flip_aug and random.random() < 0.5
+        )  # not flipped or flipped with 50% chance
 
         latents = None
         image = None
         # image/latentsを処理する
         if image_info.latents is not None:  # cache_latents=Trueの場合
             original_size = image_info.latents_original_size
-            crop_left_top = image_info.latents_crop_left_top  # calc values later if flipped
+            crop_left_top = (
+                image_info.latents_crop_left_top
+            )  # calc values later if flipped
             if not flipped:
                 latents = image_info.latents
             else:
                 latents = image_info.latents_flipped
 
-        elif image_info.latents_npz is not None:  # FineTuningDatasetまたはcache_latents_to_disk=Trueの場合
-            latents, original_size, crop_left_top = self.load_latents_from_npz(image_info, flipped)
+        elif (
+            image_info.latents_npz is not None
+        ):  # FineTuningDatasetまたはcache_latents_to_disk=Trueの場合
+            latents, original_size, crop_left_top = self.load_latents_from_npz(
+                image_info, flipped
+            )
             latents = torch.FloatTensor(latents)
 
         if True:
             # 画像を読み込み、必要ならcropする
-            img, face_cx, face_cy, face_w, face_h = self.load_image_with_face_info(subset, image_info.absolute_path)
+            img, face_cx, face_cy, face_w, face_h = self.load_image_with_face_info(
+                subset, image_info.absolute_path
+            )
             im_h, im_w = img.shape[0:2]
 
             if self.enable_bucket:
@@ -97,7 +108,9 @@ def patch_getitem(self, index):
                 )
             else:
                 if face_cx > 0:  # 顔位置情報あり
-                    img = self.crop_target(subset, img, face_cx, face_cy, face_w, face_h)
+                    img = self.crop_target(
+                        subset, img, face_cx, face_cy, face_w, face_h
+                    )
                 elif im_h > self.height or im_w > self.width:
                     assert (
                         subset.random_crop
@@ -130,10 +143,17 @@ def patch_getitem(self, index):
         images.append(image)
         latents_list.append(latents)
 
-        target_size = (image.shape[2], image.shape[1]) if image is not None else (latents.shape[2] * 8, latents.shape[1] * 8)
+        target_size = (
+            (image.shape[2], image.shape[1])
+            if image is not None
+            else (latents.shape[2] * 8, latents.shape[1] * 8)
+        )
 
         if flipped:
-            crop_left_top = (original_size[0] - crop_left_top[0] - target_size[0], crop_left_top[1])
+            crop_left_top = (
+                original_size[0] - crop_left_top[0] - target_size[0],
+                crop_left_top[1],
+            )
 
         original_sizes_hw.append((original_size[1], original_size[0]))
         crop_top_lefts.append((crop_left_top[1], crop_left_top[0]))
@@ -145,7 +165,9 @@ def patch_getitem(self, index):
             caption_layer = []
             for layer in self.XTI_layers:
                 token_strings_from = " ".join(self.token_strings)
-                token_strings_to = " ".join([f"{x}_{layer}" for x in self.token_strings])
+                token_strings_to = " ".join(
+                    [f"{x}_{layer}" for x in self.token_strings]
+                )
                 caption_ = caption.replace(token_strings_from, token_strings_to)
                 caption_layer.append(caption_)
             captions.append(caption_layer)
@@ -161,7 +183,9 @@ def patch_getitem(self, index):
 
             if len(self.tokenizers) > 1:
                 if self.XTI_layers:
-                    token_caption2 = self.get_input_ids(caption_layer, self.tokenizers[1])
+                    token_caption2 = self.get_input_ids(
+                        caption_layer, self.tokenizers[1]
+                    )
                 else:
                     token_caption2 = self.get_input_ids(caption, self.tokenizers[1])
                 input_ids2_list.append(token_caption2)
@@ -171,13 +195,19 @@ def patch_getitem(self, index):
 
     if self.token_padding_disabled:
         # padding=True means pad in the batch
-        example["input_ids"] = self.tokenizer[0](captions, padding=True, truncation=True, return_tensors="pt").input_ids
+        example["input_ids"] = self.tokenizer[0](
+            captions, padding=True, truncation=True, return_tensors="pt"
+        ).input_ids
         if len(self.tokenizers) > 1:
             # following may not work in SDXL, keep the line for future update
-            example["input_ids2"] = self.tokenizer[1](captions, padding=True, truncation=True, return_tensors="pt").input_ids
+            example["input_ids2"] = self.tokenizer[1](
+                captions, padding=True, truncation=True, return_tensors="pt"
+            ).input_ids
     else:
         example["input_ids"] = torch.stack(input_ids_list)
-        example["input_ids2"] = torch.stack(input_ids2_list) if len(self.tokenizers) > 1 else None
+        example["input_ids2"] = (
+            torch.stack(input_ids2_list) if len(self.tokenizers) > 1 else None
+        )
 
     if images[0] is not None:
         images = torch.stack(images)
@@ -186,18 +216,27 @@ def patch_getitem(self, index):
         images = None
     example["images"] = images
 
-    example["latents"] = torch.stack(latents_list) if latents_list[0] is not None else None
+    example["latents"] = (
+        torch.stack(latents_list) if latents_list[0] is not None else None
+    )
     example["captions"] = captions
 
-    example["original_sizes_hw"] = torch.stack([torch.LongTensor(x) for x in original_sizes_hw])
-    example["crop_top_lefts"] = torch.stack([torch.LongTensor(x) for x in crop_top_lefts])
-    example["target_sizes_hw"] = torch.stack([torch.LongTensor(x) for x in target_sizes_hw])
+    example["original_sizes_hw"] = torch.stack(
+        [torch.LongTensor(x) for x in original_sizes_hw]
+    )
+    example["crop_top_lefts"] = torch.stack(
+        [torch.LongTensor(x) for x in crop_top_lefts]
+    )
+    example["target_sizes_hw"] = torch.stack(
+        [torch.LongTensor(x) for x in target_sizes_hw]
+    )
     example["flippeds"] = flippeds
 
     if self.debug_dataset:
         example["image_keys"] = bucket[image_index : image_index + self.batch_size]
     return example
 
+
 train_util.BaseDataset.__getitem__ = patch_getitem
 
 
@@ -207,7 +246,14 @@ def __init__(self):
 
     # TODO 他のスクリプトと共通化する
     def generate_step_logs(
-        self, args: argparse.Namespace, current_loss, avr_loss, lr_scheduler, keys_scaled=None, mean_norm=None, maximum_norm=None
+        self,
+        args: argparse.Namespace,
+        current_loss,
+        avr_loss,
+        lr_scheduler,
+        keys_scaled=None,
+        mean_norm=None,
+        maximum_norm=None,
     ):
         logs = {"loss/current": current_loss, "loss/average": avr_loss}
 
@@ -218,7 +264,9 @@ def generate_step_logs(
 
         lrs = lr_scheduler.get_last_lr()
 
-        if args.network_train_text_encoder_only or len(lrs) <= 2:  # not block lr (or single block)
+        if (
+            args.network_train_text_encoder_only or len(lrs) <= 2
+        ):  # not block lr (or single block)
             if args.network_train_unet_only:
                 logs["lr/unet"] = float(lrs[0])
             elif args.network_train_text_encoder_only:
@@ -228,10 +276,12 @@ def generate_step_logs(
                 logs["lr/unet"] = float(lrs[-1])  # may be same to textencoder
 
             if (
-                args.optimizer_type.lower().startswith("DAdapt".lower()) or args.optimizer_type.lower() == "Prodigy".lower()
+                args.optimizer_type.lower().startswith("DAdapt".lower())
+                or args.optimizer_type.lower() == "Prodigy".lower()
             ):  # tracking d*lr value of unet.
                 logs["lr/d*lr"] = (
-                    lr_scheduler.optimizers[-1].param_groups[0]["d"] * lr_scheduler.optimizers[-1].param_groups[0]["lr"]
+                    lr_scheduler.optimizers[-1].param_groups[0]["d"]
+                    * lr_scheduler.optimizers[-1].param_groups[0]["lr"]
                 )
         else:
             idx = 0
@@ -241,9 +291,13 @@ def generate_step_logs(
 
             for i in range(idx, len(lrs)):
                 logs[f"lr/group{i}"] = float(lrs[i])
-                if args.optimizer_type.lower().startswith("DAdapt".lower()) or args.optimizer_type.lower() == "Prodigy".lower():
+                if (
+                    args.optimizer_type.lower().startswith("DAdapt".lower())
+                    or args.optimizer_type.lower() == "Prodigy".lower()
+                ):
                     logs[f"lr/d*lr/group{i}"] = (
-                        lr_scheduler.optimizers[-1].param_groups[i]["d"] * lr_scheduler.optimizers[-1].param_groups[i]["lr"]
+                        lr_scheduler.optimizers[-1].param_groups[i]["d"]
+                        * lr_scheduler.optimizers[-1].param_groups[i]["lr"]
                     )
 
         return logs
@@ -252,8 +306,17 @@ def assert_extra_args(self, args, train_dataset_group):
         pass
 
     def load_target_model(self, args, weight_dtype, accelerator):
-        text_encoder, vae, unet, _ = train_util.load_target_model(args, weight_dtype, accelerator)
-        return model_util.get_model_version_str_for_sd1_sd2(args.v2, args.v_parameterization), text_encoder, vae, unet
+        text_encoder, vae, unet, _ = train_util.load_target_model(
+            args, weight_dtype, accelerator
+        )
+        return (
+            model_util.get_model_version_str_for_sd1_sd2(
+                args.v2, args.v_parameterization
+            ),
+            text_encoder,
+            vae,
+            unet,
+        )
 
     def load_tokenizer(self, args):
         tokenizer = train_util.load_tokenizer(args)
@@ -263,22 +326,65 @@ def is_text_encoder_outputs_cached(self, args):
         return False
 
     def cache_text_encoder_outputs_if_needed(
-        self, args, accelerator, unet, vae, tokenizers, text_encoders, data_loader, weight_dtype
+        self,
+        args,
+        accelerator,
+        unet,
+        vae,
+        tokenizers,
+        text_encoders,
+        data_loader,
+        weight_dtype,
     ):
         for t_enc in text_encoders:
             t_enc.to(accelerator.device)
 
-    def get_text_cond(self, args, accelerator, batch, tokenizers, text_encoders, weight_dtype):
+    def get_text_cond(
+        self, args, accelerator, batch, tokenizers, text_encoders, weight_dtype
+    ):
         input_ids = batch["input_ids"].to(accelerator.device)
-        encoder_hidden_states = train_util.get_hidden_states(args, input_ids, tokenizers[0], text_encoders[0], weight_dtype)
+        encoder_hidden_states = train_util.get_hidden_states(
+            args, input_ids, tokenizers[0], text_encoders[0], weight_dtype
+        )
         return encoder_hidden_states
 
-    def call_unet(self, args, accelerator, unet, noisy_latents, timesteps, text_conds, batch, weight_dtype):
+    def call_unet(
+        self,
+        args,
+        accelerator,
+        unet,
+        noisy_latents,
+        timesteps,
+        text_conds,
+        batch,
+        weight_dtype,
+    ):
         noise_pred = unet(noisy_latents, timesteps, text_conds).sample
         return noise_pred
 
-    def sample_images(self, accelerator, args, epoch, global_step, device, vae, tokenizer, text_encoder, unet):
-        train_util.sample_images(accelerator, args, epoch, global_step, device, vae, tokenizer, text_encoder, unet)
+    def sample_images(
+        self,
+        accelerator,
+        args,
+        epoch,
+        global_step,
+        device,
+        vae,
+        tokenizer,
+        text_encoder,
+        unet,
+    ):
+        train_util.sample_images(
+            accelerator,
+            args,
+            epoch,
+            global_step,
+            device,
+            vae,
+            tokenizer,
+            text_encoder,
+            unet,
+        )
 
     def train(self, args):
         session_id = random.randint(0, 2**32)
@@ -300,7 +406,9 @@ def train(self, args):
 
         # データセットを準備する
         if args.dataset_class is None:
-            blueprint_generator = BlueprintGenerator(ConfigSanitizer(True, True, False, True))
+            blueprint_generator = BlueprintGenerator(
+                ConfigSanitizer(True, True, False, True)
+            )
             if use_user_config:
                 print(f"Loading dataset config from {args.dataset_config}")
                 user_config = config_util.load_user_config(args.dataset_config)
@@ -338,16 +446,24 @@ def train(self, args):
                         ]
                     }
 
-            blueprint = blueprint_generator.generate(user_config, args, tokenizer=tokenizer)
-            train_dataset_group = config_util.generate_dataset_group_by_blueprint(blueprint.dataset_group)
+            blueprint = blueprint_generator.generate(
+                user_config, args, tokenizer=tokenizer
+            )
+            train_dataset_group = config_util.generate_dataset_group_by_blueprint(
+                blueprint.dataset_group
+            )
         else:
             # use arbitrary dataset class
             train_dataset_group = train_util.load_arbitrary_dataset(args, tokenizer)
 
         current_epoch = Value("i", 0)
         current_step = Value("i", 0)
-        ds_for_collater = train_dataset_group if args.max_data_loader_n_workers == 0 else None
-        collater = train_util.collater_class(current_epoch, current_step, ds_for_collater)
+        ds_for_collater = (
+            train_dataset_group if args.max_data_loader_n_workers == 0 else None
+        )
+        collater = train_util.collater_class(
+            current_epoch, current_step, ds_for_collater
+        )
 
         if args.debug_dataset:
             train_util.debug_dataset(train_dataset_group)
@@ -375,13 +491,19 @@ def train(self, args):
         vae_dtype = torch.float32 if args.no_half_vae else weight_dtype
 
         # モデルを読み込む
-        model_version, text_encoder, vae, unet = self.load_target_model(args, weight_dtype, accelerator)
+        model_version, text_encoder, vae, unet = self.load_target_model(
+            args, weight_dtype, accelerator
+        )
 
         # text_encoder is List[CLIPTextModel] or CLIPTextModel
-        text_encoders = text_encoder if isinstance(text_encoder, list) else [text_encoder]
+        text_encoders = (
+            text_encoder if isinstance(text_encoder, list) else [text_encoder]
+        )
 
         # モデルに xformers とか memory efficient attention を組み込む
-        train_util.replace_unet_modules(unet, args.mem_eff_attn, args.xformers, args.sdpa)
+        train_util.replace_unet_modules(
+            unet, args.mem_eff_attn, args.xformers, args.sdpa
+        )
         vae.set_use_memory_efficient_attention_xformers(args.xformers)
 
         # 差分追加学習のためにモデルを読み込む
@@ -393,7 +515,12 @@ def train(self, args):
             vae.requires_grad_(False)
             vae.eval()
             with torch.no_grad():
-                train_dataset_group.cache_latents(vae, args.vae_batch_size, args.cache_latents_to_disk, accelerator.is_main_process)
+                train_dataset_group.cache_latents(
+                    vae,
+                    args.vae_batch_size,
+                    args.cache_latents_to_disk,
+                    accelerator.is_main_process,
+                )
             vae.to("cpu")
             if torch.cuda.is_available():
                 torch.cuda.empty_cache()
@@ -403,7 +530,14 @@ def train(self, args):
 
         # 必要ならテキストエンコーダーの出力をキャッシュする: Text Encoderはcpuまたはgpuへ移される
         self.cache_text_encoder_outputs_if_needed(
-            args, accelerator, unet, vae, tokenizers, text_encoders, train_dataset_group, weight_dtype
+            args,
+            accelerator,
+            unet,
+            vae,
+            tokenizers,
+            text_encoders,
+            train_dataset_group,
+            weight_dtype,
         )
 
         # prepare network
@@ -426,19 +560,26 @@ def train(self, args):
 
         if hasattr(network, "prepare_network"):
             network.prepare_network(args)
-        if args.scale_weight_norms and not hasattr(network, "apply_max_norm_regularization"):
+        if args.scale_weight_norms and not hasattr(
+            network, "apply_max_norm_regularization"
+        ):
             print(
                 "warning: scale_weight_norms is specified but the network does not support it / scale_weight_normsが指定されていますが、ネットワークが対応していません"
             )
             args.scale_weight_norms = False
 
         train_unet = not args.network_train_text_encoder_only
-        train_text_encoder = not args.network_train_unet_only and not self.is_text_encoder_outputs_cached(args)
+        train_text_encoder = (
+            not args.network_train_unet_only
+            and not self.is_text_encoder_outputs_cached(args)
+        )
         network.apply_to(text_encoder, unet, train_text_encoder, train_unet)
 
         if args.network_weights is not None:
             info = network.load_weights(args.network_weights)
-            accelerator.print(f"load network weights from {args.network_weights}: {info}")
+            accelerator.print(
+                f"load network weights from {args.network_weights}: {info}"
+            )
 
         if args.gradient_checkpointing:
             unet.enable_gradient_checkpointing()
@@ -452,18 +593,26 @@ def train(self, args):
 
         # 後方互換性を確保するよ
         try:
-            trainable_params = network.prepare_optimizer_params(args.text_encoder_lr, args.unet_lr, args.learning_rate)
+            trainable_params = network.prepare_optimizer_params(
+                args.text_encoder_lr, args.unet_lr, args.learning_rate
+            )
         except TypeError:
             accelerator.print(
                 "Deprecated: use prepare_optimizer_params(text_encoder_lr, unet_lr, learning_rate) instead of prepare_optimizer_params(text_encoder_lr, unet_lr)"
             )
-            trainable_params = network.prepare_optimizer_params(args.text_encoder_lr, args.unet_lr)
+            trainable_params = network.prepare_optimizer_params(
+                args.text_encoder_lr, args.unet_lr
+            )
 
-        optimizer_name, optimizer_args, optimizer = train_util.get_optimizer(args, trainable_params)
+        optimizer_name, optimizer_args, optimizer = train_util.get_optimizer(
+            args, trainable_params
+        )
 
         # dataloaderを準備する
         # DataLoaderのプロセス数：0はメインプロセスになる
-        n_workers = min(args.max_data_loader_n_workers, os.cpu_count() - 1)  # cpu_count-1 ただし最大で指定された数まで
+        n_workers = min(
+            args.max_data_loader_n_workers, os.cpu_count() - 1
+        )  # cpu_count-1 ただし最大で指定された数まで
 
         train_dataloader = torch.utils.data.DataLoader(
             train_dataset_group,
@@ -477,7 +626,9 @@ def train(self, args):
         # 学習ステップ数を計算する
         if args.max_train_epochs is not None:
             args.max_train_steps = args.max_train_epochs * math.ceil(
-                len(train_dataloader) / accelerator.num_processes / args.gradient_accumulation_steps
+                len(train_dataloader)
+                / accelerator.num_processes
+                / args.gradient_accumulation_steps
             )
             accelerator.print(
                 f"override steps. steps for {args.max_train_epochs} epochs is / 指定エポックまでのステップ数: {args.max_train_steps}"
@@ -487,7 +638,9 @@ def train(self, args):
         train_dataset_group.set_max_train_steps(args.max_train_steps)
 
         # lr schedulerを用意する
-        lr_scheduler = train_util.get_scheduler_fix(args, optimizer, accelerator.num_processes)
+        lr_scheduler = train_util.get_scheduler_fix(
+            args, optimizer, accelerator.num_processes
+        )
 
         # 実験的機能：勾配も含めたfp16学習を行う　モデル全体をfp16にする
         if args.full_fp16:
@@ -507,29 +660,79 @@ def train(self, args):
         # TODO めちゃくちゃ冗長なのでコードを整理する
         if train_unet and train_text_encoder:
             if len(text_encoders) > 1:
-                unet, t_enc1, t_enc2, network, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
-                    unet, text_encoders[0], text_encoders[1], network, optimizer, train_dataloader, lr_scheduler
+                (
+                    unet,
+                    t_enc1,
+                    t_enc2,
+                    network,
+                    optimizer,
+                    train_dataloader,
+                    lr_scheduler,
+                ) = accelerator.prepare(
+                    unet,
+                    text_encoders[0],
+                    text_encoders[1],
+                    network,
+                    optimizer,
+                    train_dataloader,
+                    lr_scheduler,
                 )
                 text_encoder = text_encoders = [t_enc1, t_enc2]
                 del t_enc1, t_enc2
             else:
-                unet, text_encoder, network, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
-                    unet, text_encoder, network, optimizer, train_dataloader, lr_scheduler
+                (
+                    unet,
+                    text_encoder,
+                    network,
+                    optimizer,
+                    train_dataloader,
+                    lr_scheduler,
+                ) = accelerator.prepare(
+                    unet,
+                    text_encoder,
+                    network,
+                    optimizer,
+                    train_dataloader,
+                    lr_scheduler,
                 )
                 text_encoders = [text_encoder]
         elif train_unet:
-            unet, network, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
+            (
+                unet,
+                network,
+                optimizer,
+                train_dataloader,
+                lr_scheduler,
+            ) = accelerator.prepare(
                 unet, network, optimizer, train_dataloader, lr_scheduler
             )
         elif train_text_encoder:
             if len(text_encoders) > 1:
-                t_enc1, t_enc2, network, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
-                    text_encoders[0], text_encoders[1], network, optimizer, train_dataloader, lr_scheduler
+                (
+                    t_enc1,
+                    t_enc2,
+                    network,
+                    optimizer,
+                    train_dataloader,
+                    lr_scheduler,
+                ) = accelerator.prepare(
+                    text_encoders[0],
+                    text_encoders[1],
+                    network,
+                    optimizer,
+                    train_dataloader,
+                    lr_scheduler,
                 )
                 text_encoder = text_encoders = [t_enc1, t_enc2]
                 del t_enc1, t_enc2
             else:
-                text_encoder, network, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
+                (
+                    text_encoder,
+                    network,
+                    optimizer,
+                    train_dataloader,
+                    lr_scheduler,
+                ) = accelerator.prepare(
                     text_encoder, network, optimizer, train_dataloader, lr_scheduler
                 )
                 text_encoders = [text_encoder]
@@ -577,26 +780,44 @@ def train(self, args):
         train_util.resume_from_local_or_hf_if_specified(accelerator, args)
 
         # epoch数を計算する
-        num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
+        num_update_steps_per_epoch = math.ceil(
+            len(train_dataloader) / args.gradient_accumulation_steps
+        )
         num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)
         if (args.save_n_epoch_ratio is not None) and (args.save_n_epoch_ratio > 0):
-            args.save_every_n_epochs = math.floor(num_train_epochs / args.save_n_epoch_ratio) or 1
+            args.save_every_n_epochs = (
+                math.floor(num_train_epochs / args.save_n_epoch_ratio) or 1
+            )
 
         # 学習する
         # TODO: find a way to handle total batch size when there are multiple datasets
-        total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
+        total_batch_size = (
+            args.train_batch_size
+            * accelerator.num_processes
+            * args.gradient_accumulation_steps
+        )
 
         accelerator.print("running training / 学習開始")
-        accelerator.print(f"  num train images * repeats / 学習画像の数×繰り返し回数: {train_dataset_group.num_train_images}")
-        accelerator.print(f"  num reg images / 正則化画像の数: {train_dataset_group.num_reg_images}")
-        accelerator.print(f"  num batches per epoch / 1epochのバッチ数: {len(train_dataloader)}")
+        accelerator.print(
+            f"  num train images * repeats / 学習画像の数×繰り返し回数: {train_dataset_group.num_train_images}"
+        )
+        accelerator.print(
+            f"  num reg images / 正則化画像の数: {train_dataset_group.num_reg_images}"
+        )
+        accelerator.print(
+            f"  num batches per epoch / 1epochのバッチ数: {len(train_dataloader)}"
+        )
         accelerator.print(f"  num epochs / epoch数: {num_train_epochs}")
         accelerator.print(
             f"  batch size per device / バッチサイズ: {', '.join([str(d.batch_size) for d in train_dataset_group.datasets])}"
         )
         # accelerator.print(f"  total train batch size (with parallel & distributed & accumulation) / 総バッチサイズ（並列学習、勾配合計含む）: {total_batch_size}")
-        accelerator.print(f"  gradient accumulation steps / 勾配を合計するステップ数 = {args.gradient_accumulation_steps}")
-        accelerator.print(f"  total optimization steps / 学習ステップ数: {args.max_train_steps}")
+        accelerator.print(
+            f"  gradient accumulation steps / 勾配を合計するステップ数 = {args.gradient_accumulation_steps}"
+        )
+        accelerator.print(
+            f"  total optimization steps / 学習ステップ数: {args.max_train_steps}"
+        )
 
         # TODO refactor metadata creation and move to util
         metadata = {
@@ -634,7 +855,8 @@ def train(self, args):
             "ss_adaptive_noise_scale": args.adaptive_noise_scale,
             "ss_training_comment": args.training_comment,  # will not be updated after training
             "ss_sd_scripts_commit_hash": train_util.get_git_revision_hash(),
-            "ss_optimizer": optimizer_name + (f"({optimizer_args})" if len(optimizer_args) > 0 else ""),
+            "ss_optimizer": optimizer_name
+            + (f"({optimizer_args})" if len(optimizer_args) > 0 else ""),
             "ss_max_grad_norm": args.max_grad_norm,
             "ss_caption_dropout_rate": args.caption_dropout_rate,
             "ss_caption_dropout_every_n_epochs": args.caption_dropout_every_n_epochs,
@@ -742,7 +964,10 @@ def train(self, args):
             if use_dreambooth_method:
                 for subset in dataset.subsets:
                     info = reg_dataset_dirs_info if subset.is_reg else dataset_dirs_info
-                    info[os.path.basename(subset.image_dir)] = {"n_repeats": subset.num_repeats, "img_count": subset.img_count}
+                    info[os.path.basename(subset.image_dir)] = {
+                        "n_repeats": subset.num_repeats,
+                        "img_count": subset.img_count,
+                    }
             else:
                 for subset in dataset.subsets:
                     dataset_dirs_info[os.path.basename(subset.metadata_file)] = {
@@ -780,7 +1005,9 @@ def train(self, args):
             sd_model_name = args.pretrained_model_name_or_path
             if os.path.exists(sd_model_name):
                 metadata["ss_sd_model_hash"] = train_util.model_hash(sd_model_name)
-                metadata["ss_new_sd_model_hash"] = train_util.calculate_sha256(sd_model_name)
+                metadata["ss_new_sd_model_hash"] = train_util.calculate_sha256(
+                    sd_model_name
+                )
                 sd_model_name = os.path.basename(sd_model_name)
             metadata["ss_sd_model_name"] = sd_model_name
 
@@ -808,16 +1035,29 @@ def train(self, args):
             if key in metadata:
                 minimum_metadata[key] = metadata[key]
 
-        progress_bar = tqdm(range(args.max_train_steps), smoothing=0, disable=not accelerator.is_local_main_process, desc="steps")
+        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
+            beta_start=0.00085,
+            beta_end=0.012,
+            beta_schedule="scaled_linear",
+            num_train_timesteps=1000,
+            clip_sample=False,
         )
         prepare_scheduler_for_custom_training(noise_scheduler, accelerator.device)
 
         if accelerator.is_main_process:
-            accelerator.init_trackers("network_train" if args.log_tracker_name is None else args.log_tracker_name)
+            accelerator.init_trackers(
+                "network_train"
+                if args.log_tracker_name is None
+                else args.log_tracker_name
+            )
 
         loss_list = []
         loss_total = 0.0
@@ -830,7 +1070,9 @@ def train(self, args):
             on_step_start = lambda *args, **kwargs: None
 
         # function for saving/removing
-        def save_model(ckpt_name, unwrapped_nw, steps, epoch_no, force_sync_upload=False):
+        def save_model(
+            ckpt_name, unwrapped_nw, steps, epoch_no, force_sync_upload=False
+        ):
             os.makedirs(args.output_dir, exist_ok=True)
             ckpt_file = os.path.join(args.output_dir, ckpt_name)
 
@@ -839,9 +1081,18 @@ def save_model(ckpt_name, unwrapped_nw, steps, epoch_no, force_sync_upload=False
             metadata["ss_steps"] = str(steps)
             metadata["ss_epoch"] = str(epoch_no)
 
-            unwrapped_nw.save_weights(ckpt_file, save_dtype, minimum_metadata if args.no_metadata else metadata)
+            unwrapped_nw.save_weights(
+                ckpt_file,
+                save_dtype,
+                minimum_metadata if args.no_metadata else metadata,
+            )
             if args.huggingface_repo_id is not None:
-                huggingface_util.upload(args, ckpt_file, "/" + ckpt_name, force_sync_upload=force_sync_upload)
+                huggingface_util.upload(
+                    args,
+                    ckpt_file,
+                    "/" + ckpt_name,
+                    force_sync_upload=force_sync_upload,
+                )
 
         def remove_model(old_ckpt_name):
             old_ckpt_file = os.path.join(args.output_dir, old_ckpt_name)
@@ -874,12 +1125,20 @@ def remove_model(old_ckpt_name):
                             latents = batch["latents"].to(accelerator.device)
                         else:
                             # latentに変換
-                            latents = vae.encode(batch["images"].to(dtype=vae_dtype)).latent_dist.sample()
+                            latents = vae.encode(
+                                batch["images"].to(dtype=vae_dtype)
+                            ).latent_dist.sample()
 
                             # NaNが含まれていれば警告を表示し0に置き換える
                             if torch.any(torch.isnan(latents)):
-                                accelerator.print("NaN found in latents, replacing with zeros")
-                                latents = torch.where(torch.isnan(latents), torch.zeros_like(latents), latents)
+                                accelerator.print(
+                                    "NaN found in latents, replacing with zeros"
+                                )
+                                latents = torch.where(
+                                    torch.isnan(latents),
+                                    torch.zeros_like(latents),
+                                    latents,
+                                )
                         latents = latents * self.vae_scale_factor
                     b_size = latents.shape[0]
 
@@ -891,24 +1150,42 @@ def remove_model(old_ckpt_name):
                                 text_encoder,
                                 batch["captions"],
                                 accelerator.device,
-                                args.max_token_length // 75 if args.max_token_length else 1,
+                                args.max_token_length // 75
+                                if args.max_token_length
+                                else 1,
                                 clip_skip=args.clip_skip,
                             )
                         else:
                             text_encoder_conds = self.get_text_cond(
-                                args, accelerator, batch, tokenizers, text_encoders, weight_dtype
+                                args,
+                                accelerator,
+                                batch,
+                                tokenizers,
+                                text_encoders,
+                                weight_dtype,
                             )
 
                     # Sample noise, sample a random timestep for each image, and add noise to the latents,
                     # with noise offset and/or multires noise if specified
-                    noise, noisy_latents, timesteps = train_util.get_noise_noisy_latents_and_timesteps(
+                    (
+                        noise,
+                        noisy_latents,
+                        timesteps,
+                    ) = train_util.get_noise_noisy_latents_and_timesteps(
                         args, noise_scheduler, latents
                     )
 
                     # Predict the noise residual
                     with accelerator.autocast():
                         noise_pred = self.call_unet(
-                            args, accelerator, unet, noisy_latents, timesteps, text_encoder_conds, batch, weight_dtype
+                            args,
+                            accelerator,
+                            unet,
+                            noisy_latents,
+                            timesteps,
+                            text_encoder_conds,
+                            batch,
+                            weight_dtype,
                         )
 
                     if args.v_parameterization:
@@ -917,16 +1194,22 @@ def remove_model(old_ckpt_name):
                     else:
                         target = noise
 
-                    loss = torch.nn.functional.mse_loss(noise_pred.float(), target.float(), reduction="none")
+                    loss = torch.nn.functional.mse_loss(
+                        noise_pred.float(), target.float(), reduction="none"
+                    )
                     loss = loss.mean([1, 2, 3])
 
                     loss_weights = batch["loss_weights"]  # 各sampleごとのweight
                     loss = loss * loss_weights
 
                     if args.min_snr_gamma:
-                        loss = apply_snr_weight(loss, timesteps, noise_scheduler, args.min_snr_gamma)
+                        loss = apply_snr_weight(
+                            loss, timesteps, noise_scheduler, args.min_snr_gamma
+                        )
                     if args.scale_v_pred_loss_like_noise_pred:
-                        loss = scale_v_prediction_loss_like_noise_prediction(loss, timesteps, noise_scheduler)
+                        loss = scale_v_prediction_loss_like_noise_prediction(
+                            loss, timesteps, noise_scheduler
+                        )
 
                     loss = loss.mean()  # 平均なのでbatch_sizeで割る必要なし
 
@@ -934,16 +1217,23 @@ def remove_model(old_ckpt_name):
                     if accelerator.sync_gradients and args.max_grad_norm != 0.0:
                         params_to_clip = network.get_trainable_params()
                         accelerator.clip_grad_norm_(params_to_clip, args.max_grad_norm)
-                    
+
                     optimizer.step()
                     lr_scheduler.step()
                     optimizer.zero_grad(set_to_none=True)
 
                 if args.scale_weight_norms:
-                    keys_scaled, mean_norm, maximum_norm = network.apply_max_norm_regularization(
+                    (
+                        keys_scaled,
+                        mean_norm,
+                        maximum_norm,
+                    ) = network.apply_max_norm_regularization(
                         args.scale_weight_norms, accelerator.device
                     )
-                    max_mean_logs = {"Keys Scaled": keys_scaled, "Average key norm": mean_norm}
+                    max_mean_logs = {
+                        "Keys Scaled": keys_scaled,
+                        "Average key norm": mean_norm,
+                    }
                 else:
                     keys_scaled, mean_norm, maximum_norm = None, None, None
 
@@ -952,21 +1242,47 @@ def remove_model(old_ckpt_name):
                     progress_bar.update(1)
                     global_step += 1
 
-                    self.sample_images(accelerator, args, None, global_step, accelerator.device, vae, tokenizer, text_encoder, unet)
+                    self.sample_images(
+                        accelerator,
+                        args,
+                        None,
+                        global_step,
+                        accelerator.device,
+                        vae,
+                        tokenizer,
+                        text_encoder,
+                        unet,
+                    )
 
                     # 指定ステップごとにモデルを保存
-                    if args.save_every_n_steps is not None and global_step % args.save_every_n_steps == 0:
+                    if (
+                        args.save_every_n_steps is not None
+                        and global_step % args.save_every_n_steps == 0
+                    ):
                         accelerator.wait_for_everyone()
                         if accelerator.is_main_process:
-                            ckpt_name = train_util.get_step_ckpt_name(args, "." + args.save_model_as, global_step)
-                            save_model(ckpt_name, accelerator.unwrap_model(network), global_step, epoch)
+                            ckpt_name = train_util.get_step_ckpt_name(
+                                args, "." + args.save_model_as, global_step
+                            )
+                            save_model(
+                                ckpt_name,
+                                accelerator.unwrap_model(network),
+                                global_step,
+                                epoch,
+                            )
 
                             if args.save_state:
-                                train_util.save_and_remove_state_stepwise(args, accelerator, global_step)
+                                train_util.save_and_remove_state_stepwise(
+                                    args, accelerator, global_step
+                                )
 
-                            remove_step_no = train_util.get_remove_step_no(args, global_step)
+                            remove_step_no = train_util.get_remove_step_no(
+                                args, global_step
+                            )
                             if remove_step_no is not None:
-                                remove_ckpt_name = train_util.get_step_ckpt_name(args, "." + args.save_model_as, remove_step_no)
+                                remove_ckpt_name = train_util.get_step_ckpt_name(
+                                    args, "." + args.save_model_as, remove_step_no
+                                )
                                 remove_model(remove_ckpt_name)
 
                 current_loss = loss.detach().item()
@@ -984,7 +1300,15 @@ def remove_model(old_ckpt_name):
                     progress_bar.set_postfix(**{**max_mean_logs, **logs})
 
                 if args.logging_dir is not None:
-                    logs = self.generate_step_logs(args, current_loss, avr_loss, lr_scheduler, keys_scaled, mean_norm, maximum_norm)
+                    logs = self.generate_step_logs(
+                        args,
+                        current_loss,
+                        avr_loss,
+                        lr_scheduler,
+                        keys_scaled,
+                        mean_norm,
+                        maximum_norm,
+                    )
                     accelerator.log(logs, step=global_step)
 
                 if global_step >= args.max_train_steps:
@@ -998,20 +1322,43 @@ def remove_model(old_ckpt_name):
 
             # 指定エポックごとにモデルを保存
             if args.save_every_n_epochs is not None:
-                saving = (epoch + 1) % args.save_every_n_epochs == 0 and (epoch + 1) < num_train_epochs
+                saving = (epoch + 1) % args.save_every_n_epochs == 0 and (
+                    epoch + 1
+                ) < num_train_epochs
                 if is_main_process and saving:
-                    ckpt_name = train_util.get_epoch_ckpt_name(args, "." + args.save_model_as, epoch + 1)
-                    save_model(ckpt_name, accelerator.unwrap_model(network), global_step, epoch + 1)
+                    ckpt_name = train_util.get_epoch_ckpt_name(
+                        args, "." + args.save_model_as, epoch + 1
+                    )
+                    save_model(
+                        ckpt_name,
+                        accelerator.unwrap_model(network),
+                        global_step,
+                        epoch + 1,
+                    )
 
                     remove_epoch_no = train_util.get_remove_epoch_no(args, epoch + 1)
                     if remove_epoch_no is not None:
-                        remove_ckpt_name = train_util.get_epoch_ckpt_name(args, "." + args.save_model_as, remove_epoch_no)
+                        remove_ckpt_name = train_util.get_epoch_ckpt_name(
+                            args, "." + args.save_model_as, remove_epoch_no
+                        )
                         remove_model(remove_ckpt_name)
 
                     if args.save_state:
-                        train_util.save_and_remove_state_on_epoch_end(args, accelerator, epoch + 1)
+                        train_util.save_and_remove_state_on_epoch_end(
+                            args, accelerator, epoch + 1
+                        )
 
-            self.sample_images(accelerator, args, epoch + 1, global_step, accelerator.device, vae, tokenizer, text_encoder, unet)
+            self.sample_images(
+                accelerator,
+                args,
+                epoch + 1,
+                global_step,
+                accelerator.device,
+                vae,
+                tokenizer,
+                text_encoder,
+                unet,
+            )
 
             # end of epoch
 
@@ -1028,7 +1375,13 @@ def remove_model(old_ckpt_name):
 
         if is_main_process:
             ckpt_name = train_util.get_last_ckpt_name(args, "." + args.save_model_as)
-            save_model(ckpt_name, network, global_step, num_train_epochs, force_sync_upload=True)
+            save_model(
+                ckpt_name,
+                network,
+                global_step,
+                num_train_epochs,
+                force_sync_upload=True,
+            )
 
             print("model saved.")
 
@@ -1043,7 +1396,11 @@ def setup_parser() -> argparse.ArgumentParser:
     config_util.add_config_arguments(parser)
     custom_train_functions.add_custom_train_arguments(parser)
 
-    parser.add_argument("--no_metadata", action="store_true", help="do not save metadata in output model / メタデータを出力先モデルに保存しない")
+    parser.add_argument(
+        "--no_metadata",
+        action="store_true",
+        help="do not save metadata in output model / メタデータを出力先モデルに保存しない",
+    )
     parser.add_argument(
         "--save_model_as",
         type=str,
@@ -1052,13 +1409,36 @@ def setup_parser() -> argparse.ArgumentParser:
         help="format to save the model (default is .safetensors) / モデル保存時の形式（デフォルトはsafetensors）",
     )
 
-    parser.add_argument("--unet_lr", type=float, default=None, help="learning rate for U-Net / U-Netの学習率")
-    parser.add_argument("--text_encoder_lr", type=float, default=None, help="learning rate for Text Encoder / Text Encoderの学習率")
+    parser.add_argument(
+        "--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) / モジュールの次元数（ネットワークにより定義は異なります）"
+        "--network_weights",
+        type=str,
+        default=None,
+        help="pretrained weights for network / 学習するネットワークの初期重み",
+    )
+    parser.add_argument(
+        "--network_module",
+        type=str,
+        default=None,
+        help="network module to train / 学習対象のネットワークのモジュール",
+    )
+    parser.add_argument(
+        "--network_dim",
+        type=int,
+        default=None,
+        help="network dimensions (depends on each network) / モジュールの次元数（ネットワークにより定義は異なります）",
     )
     parser.add_argument(
         "--network_alpha",
@@ -1073,14 +1453,27 @@ def setup_parser() -> argparse.ArgumentParser:
         help="Drops neurons out of training every step (0 or None is default behavior (no dropout), 1 would drop all neurons) / 訓練時に毎ステップでニューロンをdropする（0またはNoneはdropoutなし、1は全ニューロンをdropout）",
     )
     parser.add_argument(
-        "--network_args", type=str, default=None, nargs="*", help="additional argmuments for network (key=value) / ネットワークへの追加の引数"
+        "--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_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関連部分のみ学習する"
+        "--network_train_text_encoder_only",
+        action="store_true",
+        help="only training Text Encoder part / Text Encoder関連部分のみ学習する",
     )
     parser.add_argument(
-        "--training_comment", type=str, default=None, help="arbitrary comment string stored in metadata / メタデータに記録する任意のコメント文字列"
+        "--training_comment",
+        type=str,
+        default=None,
+        help="arbitrary comment string stored in metadata / メタデータに記録する任意のコメント文字列",
     )
     parser.add_argument(
         "--dim_from_weights",


Test patch:
diff --git a/experiments/locon_extract_test.py b/experiments/locon_extract_test.py
index 24324fa..454ca5b 100644
--- a/experiments/locon_extract_test.py
+++ b/experiments/locon_extract_test.py
@@ -1,4 +1,5 @@
 import os, sys
+
 sys.path.insert(0, os.getcwd())
 import torch
 import torch.nn as nn
@@ -29,6 +30,6 @@
 test_out_lora = convB(convA(test_x))
 test_out_orig = conv_orig(test_x)
 
-print('MSE Loss: ', F.mse_loss(test_out_orig, test_out_lora))
-print('L1 Loss : ', F.l1_loss(test_out_orig, test_out_lora))
-print('Distance: ', torch.dist(test_out_orig, test_out_lora))
\ No newline at end of file
+print("MSE Loss: ", F.mse_loss(test_out_orig, test_out_lora))
+print("L1 Loss : ", F.l1_loss(test_out_orig, test_out_lora))
+print("Distance: ", torch.dist(test_out_orig, test_out_lora))
diff --git a/experiments/locon_merge_test.py b/experiments/locon_merge_test.py
index 8edd565..64d9b48 100644
--- a/experiments/locon_merge_test.py
+++ b/experiments/locon_merge_test.py
@@ -24,6 +24,6 @@
 test_out_lora = convB(convA(test_x))
 test_out_orig = conv_orig(test_x)
 
-print('MSE Loss: ', F.mse_loss(test_out_orig, test_out_lora))
-print('L1 Loss : ', F.l1_loss(test_out_orig, test_out_lora))
-print('Distance: ', torch.dist(test_out_orig, test_out_lora))
\ No newline at end of file
+print("MSE Loss: ", F.mse_loss(test_out_orig, test_out_lora))
+print("L1 Loss : ", F.l1_loss(test_out_orig, test_out_lora))
+print("Distance: ", torch.dist(test_out_orig, test_out_lora))
diff --git a/experiments/singular_value_test.py b/experiments/singular_value_test.py
index a6795e6..28cc6a8 100644
--- a/experiments/singular_value_test.py
+++ b/experiments/singular_value_test.py
@@ -8,22 +8,19 @@
 from collections import Counter
 
 
-def check_diff_singular_value(
-    a: torch.TensorType,
-    b: torch.TensorType
-):
+def check_diff_singular_value(a: torch.TensorType, b: torch.TensorType):
     assert a.shape == b.shape
-    
-    diff = a-b
+
+    diff = a - b
     U, S, Vh = torch.linalg.svd(diff)
-    
+
     return S
 
 
 def load_models(a, b):
-    model_a = torch.load(a)['state_dict']
-    model_b = torch.load(b)['state_dict']
-    
+    model_a = torch.load(a)["state_dict"]
+    model_b = torch.load(b)["state_dict"]
+
     linear_singular = []
     linear_layer = {}
     conv_singular = []
@@ -31,52 +28,57 @@ def load_models(a, b):
     for k, va, vb in ((k, model_a[k], model_b[k]) for k in model_a if k in model_b):
         if va.shape != vb.shape:
             continue
-        if torch.allclose(va.float(), vb.float()): continue
-        if 'first_stage' in k: continue
-        
+        if torch.allclose(va.float(), vb.float()):
+            continue
+        if "first_stage" in k:
+            continue
+
         if len(va.shape) == 2:
             out, in_ = va.shape
-            if out>10000 or in_>10000:
+            if out > 10000 or in_ > 10000:
                 continue
             linear_singular.append(
-                check_diff_singular_value(
-                    va.float().cuda(), 
-                    vb.float().cuda()
-                ).cpu()
+                check_diff_singular_value(va.float().cuda(), vb.float().cuda()).cpu()
             )
-            linear_layer[k] = torch.sum(linear_singular[-1]>0.05).item()
+            linear_layer[k] = torch.sum(linear_singular[-1] > 0.05).item()
         elif len(va.shape) == 4:
             out_ch, in_ch, kw, kh = va.shape
             conv_singular.append(
                 check_diff_singular_value(
                     va.reshape(out_ch, -1).float().cuda(),
-                    vb.reshape(out_ch, -1).float().cuda()
+                    vb.reshape(out_ch, -1).float().cuda(),
                 ).cpu()
             )
-            conv_layer[k] = torch.sum(conv_singular[-1]>0.1).item()
+            conv_layer[k] = torch.sum(conv_singular[-1] > 0.1).item()
         else:
             continue
-        
+
         print(k)
-    
-    linear_singular = torch.stack([i[:128] for i in linear_singular if i.shape[0]>100])
-    rank_counter = Counter([torch.sum(i>0.05).item() for i in linear_singular if i.shape[0]>100])
+
+    linear_singular = torch.stack(
+        [i[:128] for i in linear_singular if i.shape[0] > 100]
+    )
+    rank_counter = Counter(
+        [torch.sum(i > 0.05).item() for i in linear_singular if i.shape[0] > 100]
+    )
     linear_singular_rank = [rank_counter[i] for i in range(128)]
-    
-    conv_singular = torch.stack([i[:128] for i in conv_singular if i.shape[0]>100])
-    rank_counter = Counter([torch.sum(i>0.1).item() for i in conv_singular if i.shape[0]>100])
+
+    conv_singular = torch.stack([i[:128] for i in conv_singular if i.shape[0] > 100])
+    rank_counter = Counter(
+        [torch.sum(i > 0.1).item() for i in conv_singular if i.shape[0] > 100]
+    )
     conv_singular_rank = [rank_counter[i] for i in range(128)]
-    
+
     lin_std, lin_mean = torch.std_mean(linear_singular, dim=0)
     lin_max = torch.max(linear_singular, dim=0).values
     lin_min = torch.min(linear_singular, dim=0).values
     conv_std, conv_mean = torch.std_mean(conv_singular, dim=0)
     conv_max = torch.max(conv_singular, dim=0).values
     conv_min = torch.min(conv_singular, dim=0).values
-    
+
     fig, ax = plt.subplots()
     fig, ax_conv = plt.subplots()
-    
+
     x = np.arange(1, 129)
     # ax.plot(x, np.array(lin_mean))
     # ax.fill_between(x, np.array(lin_mean-lin_std), np.array(lin_mean+lin_std), alpha=0.5)
@@ -84,14 +86,14 @@ def load_models(a, b):
     # ax_conv.plot(x, np.array(conv_mean))
     # ax_conv.fill_between(x, np.array(conv_mean-conv_std), np.array(conv_mean+conv_std), alpha=0.5)
     # ax_conv.fill_between(x, conv_min.numpy(), conv_max.numpy(), alpha=0.2)
-    
+
     # ax.plot(list(linear_layer.keys()), list(linear_layer.values()))
     ax_conv.plot(list(conv_layer.keys()), list(conv_layer.values()))
     plt.show()
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     load_models(
-        './test_model/umamusume0224-continue-ep00-gs123386.ckpt',
-        './test_model/ACertainty.ckpt',
-    )
\ No newline at end of file
+        "./test_model/umamusume0224-continue-ep00-gs123386.ckpt",
+        "./test_model/ACertainty.ckpt",
+    )
diff --git a/experiments/sparse_bias_test.py b/experiments/sparse_bias_test.py
index d15ac21..ab89c19 100644
--- a/experiments/sparse_bias_test.py
+++ b/experiments/sparse_bias_test.py
@@ -1,4 +1,5 @@
 import os, sys
+
 sys.path.insert(0, os.getcwd())
 import torch
 import torch.nn as nn
@@ -9,12 +10,10 @@
 
 
 def make_sparse(t, sparsity=0.95):
-    sparse_t = t.masked_fill(
-        torch.abs(t) < torch.quantile(torch.abs(t), sparsity), 
-        0
-    )
+    sparse_t = t.masked_fill(torch.abs(t) < torch.quantile(torch.abs(t), sparsity), 0)
     return sparse_t
 
+
 seed_everything(0)
 
 KERNEL_SIZE = 3
@@ -27,45 +26,44 @@ def make_sparse(t, sparsity=0.95):
 
 conv_orig = nn.Conv2d(IN_CH, OUT_CH, 3, 1, 1, bias=False).cuda()
 conv_orig.weight = nn.Parameter(
-    (torch.randn(OUT_CH, LORA_RANK*2) 
-     @ torch.randn(LORA_RANK*2, IN_CH*9) 
-     * 0.05
-     ).reshape(conv_orig.weight.shape).cuda()
+    (torch.randn(OUT_CH, LORA_RANK * 2) @ torch.randn(LORA_RANK * 2, IN_CH * 9) * 0.05)
+    .reshape(conv_orig.weight.shape)
+    .cuda()
 )
 conv_rebd = nn.Conv2d(IN_CH, OUT_CH, 3, 1, 1, bias=False).cuda()
 
-extract_a, extract_b, diff = extract_conv(conv_orig.weight, 'fixed', LORA_RANK, 'cuda')
+extract_a, extract_b, diff = extract_conv(conv_orig.weight, "fixed", LORA_RANK, "cuda")
 out_weight = extract_b.reshape(OUT_CH, -1) @ extract_a.reshape(LORA_RANK, -1)
 out_weight = out_weight.reshape(conv_orig.weight.shape)
 
 
 conv_rebd.weight = nn.Parameter(out_weight)
 print()
-print('without sparse bias')
-print('MSE Loss: ', F.mse_loss(conv_orig.weight, conv_rebd.weight))
-print('L1 Loss : ', F.l1_loss(conv_orig.weight, conv_rebd.weight))
-print('Distance: ', torch.dist(conv_orig.weight, conv_rebd.weight))
+print("without sparse bias")
+print("MSE Loss: ", F.mse_loss(conv_orig.weight, conv_rebd.weight))
+print("L1 Loss : ", F.l1_loss(conv_orig.weight, conv_rebd.weight))
+print("Distance: ", torch.dist(conv_orig.weight, conv_rebd.weight))
 
 
 conv_rebd.weight = nn.Parameter(out_weight + make_sparse(diff, 0.99))
 print()
-print('with sparse bias, sparsity 99% ')
-print('MSE Loss: ', F.mse_loss(conv_orig.weight, conv_rebd.weight))
-print('L1 Loss : ', F.l1_loss(conv_orig.weight, conv_rebd.weight))
-print('Distance: ', torch.dist(conv_orig.weight, conv_rebd.weight))
+print("with sparse bias, sparsity 99% ")
+print("MSE Loss: ", F.mse_loss(conv_orig.weight, conv_rebd.weight))
+print("L1 Loss : ", F.l1_loss(conv_orig.weight, conv_rebd.weight))
+print("Distance: ", torch.dist(conv_orig.weight, conv_rebd.weight))
 
 
 conv_rebd.weight = nn.Parameter(out_weight + make_sparse(diff, 0.98))
 print()
-print('with sparse bias, sparsity 98% ')
-print('MSE Loss: ', F.mse_loss(conv_orig.weight, conv_rebd.weight))
-print('L1 Loss : ', F.l1_loss(conv_orig.weight, conv_rebd.weight))
-print('Distance: ', torch.dist(conv_orig.weight, conv_rebd.weight))
+print("with sparse bias, sparsity 98% ")
+print("MSE Loss: ", F.mse_loss(conv_orig.weight, conv_rebd.weight))
+print("L1 Loss : ", F.l1_loss(conv_orig.weight, conv_rebd.weight))
+print("Distance: ", torch.dist(conv_orig.weight, conv_rebd.weight))
 
 
 conv_rebd.weight = nn.Parameter(out_weight + make_sparse(diff, 0.95))
 print()
-print('with sparse bias, sparsity 95% ')
-print('MSE Loss: ', F.mse_loss(conv_orig.weight, conv_rebd.weight))
-print('L1 Loss : ', F.l1_loss(conv_orig.weight, conv_rebd.weight))
-print('Distance: ', torch.dist(conv_orig.weight, conv_rebd.weight))
\ No newline at end of file
+print("with sparse bias, sparsity 95% ")
+print("MSE Loss: ", F.mse_loss(conv_orig.weight, conv_rebd.weight))
+print("L1 Loss : ", F.l1_loss(conv_orig.weight, conv_rebd.weight))
+print("Distance: ", torch.dist(conv_orig.weight, conv_rebd.weight))
diff --git a/experiments/time_test.py b/experiments/time_test.py
index 083cb72..ad3b94d 100644
--- a/experiments/time_test.py
+++ b/experiments/time_test.py
@@ -5,58 +5,58 @@
 
 
 def ein(t, x1, x2, grad):
-    #inference
-    rebuild = torch.einsum('i j k l, j r, i p -> p r k l', t, x1, x2)
+    # inference
+    rebuild = torch.einsum("i j k l, j r, i p -> p r k l", t, x1, x2)
     # temp = torch.einsum('i j k l, j r -> i r k l', t, x1)
     # rebuild = torch.einsum('i j k l, i r -> r j k l', temp, x2)
-    
-    #backward
-    temp = torch.einsum('i j k l, j r -> i r k l', t, x1)
-    rebuild = torch.einsum('i j k l, i r -> r j k l', temp, x2)
-    
-    grad_w = rebuild*grad
-    grad_x2 = torch.einsum('r j k l, i j k l -> r i', temp, grad_w)
-    grad_temp = torch.einsum('i j k l, i r -> r j k l', grad_w, x2.T)
-    
-    grad_x1 = torch.einsum('i r k l, i j k l -> r j', t, grad_temp)
-    grad_t = torch.einsum('i j k l, j r -> i r k l', grad_temp, x1.T)
-    
+
+    # backward
+    temp = torch.einsum("i j k l, j r -> i r k l", t, x1)
+    rebuild = torch.einsum("i j k l, i r -> r j k l", temp, x2)
+
+    grad_w = rebuild * grad
+    grad_x2 = torch.einsum("r j k l, i j k l -> r i", temp, grad_w)
+    grad_temp = torch.einsum("i j k l, i r -> r j k l", grad_w, x2.T)
+
+    grad_x1 = torch.einsum("i r k l, i j k l -> r j", t, grad_temp)
+    grad_t = torch.einsum("i j k l, j r -> i r k l", grad_temp, x1.T)
+
     # assert t.shape == grad_t.shape
     # assert x1.shape == grad_x1.shape
     # assert x2.shape == grad_x2.shape
 
 
 def mat(t, x1, x2, grad):
-    #inference
+    # inference
     temp = (t.transpose(1, 3) @ x1).transpose(1, 3)
     rebuild = (temp.transpose(0, 3) @ x2).transpose(0, 3)
-    
-    #backward
+
+    # backward
     temp = (t.transpose(1, 3) @ x1).transpose(1, 3)
     rebuild = (temp.transpose(0, 3) @ x2).transpose(0, 3)
-    
-    grad_w = rebuild*grad
-    grad_x2 = torch.einsum('r j k l, i j k l -> r i', temp, grad_w)
+
+    grad_w = rebuild * grad
+    grad_x2 = torch.einsum("r j k l, i j k l -> r i", temp, grad_w)
     grad_temp = (grad_w.transpose(0, 3) @ x2.T).transpose(0, 3)
-    
-    grad_x1 = torch.einsum('i r k l, i j k l -> r j', t, grad_temp)
+
+    grad_x1 = torch.einsum("i r k l, i j k l -> r j", t, grad_temp)
     grad_t = (grad_temp.transpose(1, 3) @ x1.T).transpose(1, 3)
-    
+
     # assert t.shape == grad_t.shape
     # assert x1.shape == grad_x1.shape
     # assert x2.shape == grad_x2.shape
 
 
 def fold(w1a, w1b, grad):
-    #inference
-    rebuild = (w1a@w1b).reshape(OUT_DIM, IN_DIM, K_SIZE, K_SIZE)
-    
-    #backward
-    rebuild = (w1a@w1b)
+    # inference
+    rebuild = (w1a @ w1b).reshape(OUT_DIM, IN_DIM, K_SIZE, K_SIZE)
+
+    # backward
+    rebuild = w1a @ w1b
     grad_w1 = rebuild * grad.reshape(rebuild.shape)
     grad_w1a = grad_w1 @ w1b.T
     grad_w1b = w1a.T @ grad_w1
-    
+
     # assert w1a.shape == grad_w1a.shape
     # assert w1b.shape == grad_w1b.shape
 
@@ -71,7 +71,7 @@ def fold(w1a, w1b, grad):
 x2 = torch.randn(RANK, OUT_DIM).cuda()
 
 w1a = torch.randn(OUT_DIM, RANK).cuda()
-w1b = torch.randn(RANK, IN_DIM*K_SIZE**2).cuda()
+w1b = torch.randn(RANK, IN_DIM * K_SIZE**2).cuda()
 
 grad = torch.randn(OUT_DIM, IN_DIM, K_SIZE, K_SIZE).cuda()
 
@@ -79,10 +79,14 @@ def fold(w1a, w1b, grad):
 mat(t, x1, x2, grad)
 fold(w1a, w1b, grad)
 
-print(f'Proposition 1 Params: {w1a.size().numel() + w1b.size().numel()}')
-print(f'Proposition 3 Params: {t.size().numel() + x1.size().numel() + x2.size().numel()}')
+print(f"Proposition 1 Params: {w1a.size().numel() + w1b.size().numel()}")
+print(
+    f"Proposition 3 Params: {t.size().numel() + x1.size().numel() + x2.size().numel()}"
+)
 
 NUM = 5000
-print('CP einsum', timeit('ein(t, x1, x2, grad)', globals=globals(), number=NUM)/NUM)
-print('CP matmul', timeit('mat(t, x1, x2, grad)', globals=globals(), number=NUM)/NUM)
-print('Pro1 matmul', timeit('fold(w1a, w1b, grad)', globals=globals(), number=NUM)/NUM)
\ No newline at end of file
+print("CP einsum", timeit("ein(t, x1, x2, grad)", globals=globals(), number=NUM) / NUM)
+print("CP matmul", timeit("mat(t, x1, x2, grad)", globals=globals(), number=NUM) / NUM)
+print(
+    "Pro1 matmul", timeit("fold(w1a, w1b, grad)", globals=globals(), number=NUM) / NUM
+)