Repository: KohakuBlueleaf/LyCORIS (language: python)

Issue:
Fix division by zero weight_norm causes NaN loss in DoRA
In apply_weight_decompose, weight_norm can equal 0 which causes NaN weight during training.

We can prevent the division by zero by adding a very small value to weight_norm.


Discussion / hints:

Fix patch:
diff --git a/Change.md b/Change.md
index 4349bcf..5196fd2 100644
--- a/Change.md
+++ b/Change.md
@@ -1,5 +1,62 @@
 # Change Log
 
+## 2024/06/xx update to 3.0.0 - Brand New Functional API, Parametrize API and Module API
+
+### The reasons of 3.0.0
+
+We reconstruct the whole library with new Class definition and brand new Functional API system.
+
+We also removed lot of redundant/unused modules.
+
+Since the whole library are changed significantly. We decide to call it 3.0.0 as a new major version.
+
+### Major Changes
+
+* New Module API
+* Add Parametrize API
+* Add Functional API
+  * LoCon/LoHa/LoKr/Diag-OFT/BOFT only.
+* Remove optional deps from install_requires
+* Remove lot of redundant/deprecated modules
+* Better testing
+* HunYuan DiT Support ([PR](https://github.com/kohya-ss/sd-scripts/pull/1378) in kohya-ss/sd-scripts)
+
+### Full change log
+
+#### New Features
+
+* LyCORIS now have consistent API for different algorithm like `bypass_forward_diff` or `get_diff_weight` method. Developers of other project can utilize these API to do more tricks or integrate LyCORIS into their framework more easily.
+* LyCORIS now have parametrize API which utilize `torch.nn.utils.parametrize.register_parametrization` to directly patch individual parameters. Which can be useful for MHA layer or other tricky modules.
+  * Currently only support 2~5D tensors. And LyCORIS will pretend these weights are weight of Linear/Conv1,2,3D then send it into LyCORIS modules
+  * More native implementation or more detailed control will be added in the future.
+* LyCORIS now have functional API. Developers who prefer functional more than Module things can utilize this feature.
+  * Functional API also allow developers who don't want to introduce new dependencies. Just copy-paste the source code and utilizing it. (with Apache-2 License, directly copy-paste is totally allowed)
+* Add support for Conv1d and Conv3d module on LoCon/LoHa/LoKr/Full/OFT/BOFT/GLoRA (not All algo in LyCORIS support them, you may receive error when apply unsopported algo), support inherited module (for example: `LoRACompatibleConv` or `LoRACompatibleLinear` from [`huggingface/diffusers`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/lora.py))
+* HunYuan DiT support.
+
+#### Improvements, Fixes, Slight Changes
+
+* Drop dependencies related to kohya-ss/sd-scripts:
+  * We now take kohya-ss/sd-scripts as optional dependency
+  * Which means `transformers`, `diffusers` and anything related to kohya are all optional deps now.
+* The definition of dropout and rank_dropout in each algorithm are changed. Since some concept of original rank_dropout in the lora of kohya-ss/sd-script is hard to applied to other algorithm. We can only design the dropout for each module seperatedly.
+* `apply_max_norm` issue are all fixed.
+* DyLoRA, (IA)^3, GLoRA are all rewritten and support Linear/Conv1,2,3d.
+* (IA)^3, GLoRA, Diag-OFT, BOFT are supported in `create_lycoris_from_weights`
+  * `lycoris.kohya.create_network_from_weights` also support them as well.
+* Fix wrong implementation of BOFT.
+* `create_lycoris_from_weights` and `create_network_from_weights` now have correct logging infos.
+* `get_module` and `make_module` are moved into modules' API.
+
+#### Deprecation
+
+* HCP modules are dropped. We will wait until HCP have better wrapper API.
+* HyperNetwork-related modules like `hypernet/`, `attention.py`, `lilora.py` are removed.
+* Uncompleted GLoKr are removed.
+* code copied from kohya-ss/sd-scripts are removed. The original sd-scripts repo is now an optional dependency.
+
+---
+
 ## 2024/03/15 update to 2.2.0 - QLyCORIS and DoRA
 
 #### New Algo
diff --git a/README.md b/README.md
index 78875b4..a8aac19 100644
--- a/README.md
+++ b/README.md
@@ -1,6 +1,15 @@
+
+![pypi](https://img.shields.io/pypi/v/lycoris-lora.svg)
+![versions](https://img.shields.io/pypi/pyversions/lycoris-lora.svg)
+![PyPI - License](https://img.shields.io/pypi/l/lycoris-lora)
+![downloads](https://img.shields.io/pypi/dm/lycoris-lora)
+![commits](https://img.shields.io/github/commit-activity/m/KohakuBlueleaf/LyCORIS/dev)
+![Discord](https://img.shields.io/discord/1082218577395986452)
+
+
 # LyCORIS - Lora beYond Conventional methods, Other Rank adaptation Implementations for Stable diffusion. (ICLR'24)
 
-![banner image](docs/images/banner.png)
+![banner image](docs/images/banner2.png)
 
 A project that implements different parameter-efficient fine-tuning algorithms for Stable Diffusion.
 
@@ -20,21 +29,21 @@ Please check [List of Implemented Algorithms](docs/Algo-List.md) and [Guidelines
 
 A simple comparison of some of these methods are provided below (to be taken with a grain of salt)
 
-|                       | Full | LoRA | LoHa | LoKr low factor | LoKr high factor $^+$ |
-| --------------------- | ---- | ---- | ---- | --------------- | ---------------- |
-| Fidelity              | ★   | ●   | ▲   | ◉              | ▲               |
-| Flexibility $^*$     | ★   | ●   | ◉   | ▲              | ● $^†$        |
-| Diversity             | ▲   | ◉   | ★   | ●              | ★               |
-| Size                  | ▲   | ●   | ●   | ●              | ★               |
-| Training Speed Linear | ★   | ●   | ●   | ★              | ★               |
-| Training Speed Conv   | ●   | ★   | ▲   | ●              | ●               |
+|                       | Full | LoRA | LoHa | LoKr low factor | LoKr high factor$^+$ |
+| --------------------- | ---- | ---- | ---- | --------------- | ---------------------- |
+| Fidelity              | ★   | ●   | ▲   | ◉              | ▲                     |
+| Flexibility$^*$     | ★   | ●   | ◉   | ▲              | ●$^†$              |
+| Diversity             | ▲   | ◉   | ★   | ●              | ★                     |
+| Size                  | ▲   | ●   | ●   | ●              | ★                     |
+| Training Speed Linear | ★   | ●   | ●   | ★              | ★                     |
+| Training Speed Conv   | ●   | ★   | ▲   | ●              | ●                     |
 
 ★ > ◉ > ● > ▲
 [> means better and smaller size is better]
 
-$^+$ Usually we take `factor <= 0.5 * sqrt(dim)` as low factor and `factor >= sqrt(dim` as high factor. For example, factor<=8 for SD1.x/SD2.x/SDXL can be seen as low factor, and, factor>=16 can be seen as high factor. <br>
-$^*$ Flexibility means anything related to generating images not similar to those in the training set, and combination of multiple concepts, whether they are trained together or not <br>
-$^†$ It may become more difficult to switch base model or combine multiple concepts in this situation <br>
+$^+$ Usually we take `factor <= 0.5 * sqrt(dim)` as low factor and `factor >= sqrt(dim` as high factor. For example, factor<=8 for SD1.x/SD2.x/SDXL can be seen as low factor, and, factor>=16 can be seen as high factor. `<br>`
+$^*$ Flexibility means anything related to generating images not similar to those in the training set, and combination of multiple concepts, whether they are trained together or not `<br>`
+$^†$ It may become more difficult to switch base model or combine multiple concepts in this situation `<br>`
 
 **The actual performance may vary depending on the datasets, tasks, and hyperparameters used. It is recommended to experiment with different settings to achieve optimal results.**
 
@@ -66,7 +75,7 @@ However, newer model types may not always be supported. If you encounter this is
 There are three different ways to train LyCORIS models.
 
 - 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 [Naifu-Diffusion](https://github.com/Mikubill/naifu)
 - 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
@@ -74,7 +83,7 @@ In any case, please install this package in the corresponding virtual environmen
 - through pip
 
   ```bash
-  pip install lycoris_lora
+  pip install lycoris-lora
   ```
 - or from source
 
@@ -111,6 +120,8 @@ You can use this package's kohya module to run kohya's training script to train
 
 #### HCP-Diffusion
 
+**The support for HCP-Diffusion has been dropped on LyCORIS3.0.0, we will wait until HCP side finish the implementation of new wrapper**
+
 You can use this package's hcp module to run HCP-Diffusion's training script to train lycoris module for SD models
 
 ```bash
@@ -156,6 +167,13 @@ 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.
 
+#### Other method
+
+After LyCORIS3.0.0, Parametrize API and Functional API have been added, which provide more different ways on utilizing LyCORIS library.
+
+Check API reference for more informations.
+You can also take the [test suites](test/) as a kind of examples.
+
 #### Bitsandbytes support
 
 See [bnb_example.py](bnb_example.py) for example. Basically as same as standalone wrapper.
@@ -243,44 +261,61 @@ See [docs/Conversion-scripts.md](docs/Conversion-scripts.md) for more informatio
 
 For full log, please see [Change.md](Change.md)
 
-## 2024/03/15 update to 2.2.0 - QLyCORIS and DoRA
+### 2024/06/xx update to 3.0.0 - Brand New Functional API, Parametrize API and Module API
 
-#### New Algo
+#### The reasons of 3.0.0
 
-* DoRA
-  * Ref: [DoRA: Weight-Decomposed Low-Rank Adaptation](https://github.com/KohakuBlueleaf/LyCORIS)
-* Weight decompose for LoHa and LoKr. (A.K.A DoHa/DoKr)
-  * DoRA/DoHa/DoKr will require smaller Learning rate!
+We reconstruct the whole library with new Class definition and brand new Functional API system.
 
-#### New Features
+We also removed lot of redundant/unused modules.
 
-* Support "bypass" (a.k.a. adapter) mode for LoHa/LoKr/OFT/BOFT
-  * LoHa will require 2xFLOPs since we rebuild full diff weight and then do one more forward.
-  * LoKr, OFT, BOFT should be more efficient than LoHa in bypass mode.
-* Support [bnb 8bit/4bit Linear layer](https://github.com/TimDettmers/bitsandbytes) (a.k.a. QLyCORIS) with LoHa/LoKr/OFT/BOFT.
-  * This will force module to enable bypass mode.
+Since the whole library are changed significantly. We decide to call it 3.0.0 as a new major version.
 
-#### Fixes, slight changes
+#### Major Changes
 
-* Refine some details about code quality. Based on the report from GitRoll. (Thx you gitroll!)
-* Remove redundant calculation in BOFT
-* rank_dropout has been removed from OFT/BOFT temporarily untill we ensure how to apply it.
-* Fix bugs in lokr when `lokr_w1_a` not exist.
-* Fix bugs in conversion scritps.
+* New Module API
+* Add Parametrize API
+* Add Functional API
+  * LoCon/LoHa/LoKr/Diag-OFT/BOFT only.
+* Remove optional deps from install_requires
+* Remove lot of redundant/deprecated modules
+* Better testing
+* HunYuan DiT Support ([PR](https://github.com/kohya-ss/sd-scripts/pull/1378) in kohya-ss/sd-scripts)
+
+#### Full change log
+
+##### New Features
+
+* LyCORIS now have consistent API for different algorithm like `bypass_forward_diff` or `get_diff_weight` method. Developers of other project can utilize these API to do more tricks or integrate LyCORIS into their framework more easily.
+* LyCORIS now have parametrize API which utilize `torch.nn.utils.parametrize.register_parametrization` to directly patch individual parameters. Which can be useful for MHA layer or other tricky modules.
+  * Currently only support 2~5D tensors. And LyCORIS will pretend these weights are weight of Linear/Conv1,2,3D then send it into LyCORIS modules
+  * More native implementation or more detailed control will be added in the future.
+* LyCORIS now have functional API. Developers who prefer functional more than Module things can utilize this feature.
+  * Functional API also allow developers who don't want to introduce new dependencies. Just copy-paste the source code and utilizing it. (with Apache-2 License, directly copy-paste is totally allowed)
+* Add support for Conv1d and Conv3d module on LoCon/LoHa/LoKr/Full/OFT/BOFT/GLoRA (not All algo in LyCORIS support them, you may receive error when apply unsopported algo), support inherited module (for example: `LoRACompatibleConv` or `LoRACompatibleLinear` from [`huggingface/diffusers`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/lora.py))
+* HunYuan DiT support.
+
+##### Improvements, Fixes, Slight Changes
+
+* Drop dependencies related to kohya-ss/sd-scripts:
+  * We now take kohya-ss/sd-scripts as optional dependency
+  * Which means `transformers`, `diffusers` and anything related to kohya are all optional deps now.
+* The definition of dropout and rank_dropout in each algorithm are changed. Since some concept of original rank_dropout in the lora of kohya-ss/sd-script is hard to applied to other algorithm. We can only design the dropout for each module seperatedly.
+* `apply_max_norm` issue are all fixed.
+* DyLoRA, (IA)^3, GLoRA are all rewritten and support Linear/Conv1,2,3d.
+* (IA)^3, GLoRA, Diag-OFT, BOFT are supported in `create_lycoris_from_weights`
+  * `lycoris.kohya.create_network_from_weights` also support them as well.
+* Fix wrong implementation of BOFT.
+* `create_lycoris_from_weights` and `create_network_from_weights` now have correct logging infos.
+* `get_module` and `make_module` are moved into modules' API.
 
 ## Todo list
 
-- [X] Module and Document for using LyCORIS in any other model, Not only SD.
-- [X] Proposition3 in [FedPara](https://arxiv.org/abs/2108.06098)
-  * also need custom backward to save the vram
-- [ ] Low rank + sparse representation
-  - [X] For extraction
-  - [ ] For training
-- [ ] Support more operation, not only linear and conv2d.
-- [X] Configure varying ranks or dimensions for specific modules as needed.
 - [ ] Automatically selecting an algorithm based on the specific rank requirement.
-- [ ] Explore other low-rank representations or parameter-efficient methods to fine-tune either the entire model or specific parts of it.
 - [ ] More experiments for different task, not only diffusion models.
+  - [X] LoKr and LoHa have been proven to be useful for Large Language Model.
+- [ ] Explore other low-rank representations or parameter-efficient methods to fine-tune either the entire model or specific parts of it.
+- [ ] Documentation for whole library.
 
 ## Citation
 
diff --git a/docs/API.md b/docs/API.md
new file mode 100644
index 0000000..c9ab8f1
--- /dev/null
+++ b/docs/API.md
@@ -0,0 +1,73 @@
+# API Reference [WIP]
+
+## Module
+
+### Class: `LycorisBaseModule`:
+
+* classmethod
+
+  * `parametrize`
+  * `algo_check`
+  * `extract_state_dict`
+  * `make_module_from_state_dict`
+* property
+
+  * `dtype`
+  * `device`
+  * `org_weight`
+* methods
+
+  * `apply_to`
+  * `restore`
+  * `merge_to`
+  * `get_diff_weight`
+  * `get_merged_weight`
+  * `apply_max_norm`
+  * `bypass_forward_diff`
+  * `bypass_forward`
+  * `parametrize_forward`
+  * `forward`
+
+#### Subclasses
+
+* `LoConModule`
+* `LohaModule`
+* `LokrModule`
+* `DyLoraModule`
+* `GLoRAModule`
+* `NormModule`
+* `FullModule`
+* `DiagOFTModule`
+
+### Functions
+
+* `get_module`: determine the algorithm and extract corresponding weights from state dict.
+* `make_module`: based on given algorithm and weights to construct modules.
+
+## Functional
+
+For each modules, we have 3 basic methods:
+
+* `weight_gen`: Generate weights for corresponding algorithm
+* `weight_diff`: calculate $\Delta W$
+* `bypass_forward_diff`: calculate $\Delta W X$ 
+
+There are some other utilities:
+
+* `factorization`: $fact(p, factor) = (m, n)$ 
+  * where $m \times n = p$, $m < n$, $m<=factor$ and $m, n \in \mathbb{N}$
+  * This method have been used in LoKr and Diag-OFT.
+* `power2factorization`: $p2fact(p, factor) = (m, n)$ 
+  * where $m \times n = p$, $m < n$, $m<=factor$, $m=2k$, $n=2^p$ and $m, n, p, k \in \mathbb{N}$
+  * This method have been used in BOFT.
+* `tucker_weight` and `tucker_weight_from_conv`: Reconstruct tucker decomposed weight from tensors or conv modules.
+
+## Others
+
+### wrapper
+* `LycorisNetwork`: the wrapper class to patch any pytorch modules to apply LyCORIS algorithms.
+* `create_lycoris`: see example
+* `create_lycoris_from_weights`: see example
+
+### kohya
+* the specialized wrapper for kohya-ss/sd-scripts.
\ No newline at end of file
diff --git a/docs/Algo-List.md b/docs/Algo-List.md
index b27c7a7..6bae257 100644
--- a/docs/Algo-List.md
+++ b/docs/Algo-List.md
@@ -86,7 +86,7 @@ However, newer methods may only be available in the latest release / the dev bra
 * 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
+* Set `constraint` 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)
 
 ### BOFT
@@ -94,6 +94,6 @@ However, newer methods may only be available in the latest release / the dev bra
 * Triggered by `algo=boft`
 * An advanced version of Diag-OFT which use butterfly operation to get full orthogonol matrix.
 * With some extreme hyperparameter settings, it can become diag-oft or oft. Its capbility is between this 2 algorithm.
-* constrain and rescaled also supported.
+* constraint and rescaled also supported.
 * TODO
 * Ref: [Parameter-Efficient Orthogonal Finetuning via Butterfly Factorization](https://arxiv.org/abs/2311.06243)
diff --git a/docs/Network-Args.md b/docs/Network-Args.md
index 33afeea..8ee1d5f 100644
--- a/docs/Network-Args.md
+++ b/docs/Network-Args.md
@@ -91,5 +91,5 @@ Arguments to put in `network_args` for kohya sd scripts
 
 ### Constrained OFT
 
-- Enabled with `constrain=FLOAT`
+- Enabled with `constraint=FLOAT`
 - Valid for Diag-OFT
diff --git a/docs/images/banner2.png b/docs/images/banner2.png
new file mode 100644
index 0000000..6e26465
Binary files /dev/null and b/docs/images/banner2.png differ
diff --git a/bnb_example.py b/example/bnb_example.py
similarity index 100%
rename from bnb_example.py
rename to example/bnb_example.py
diff --git a/standalone_example.py b/example/standalone_example.py
similarity index 100%
rename from standalone_example.py
rename to example/standalone_example.py
diff --git a/experiments/boft_factorize.py b/experiments/boft_factorize.py
index 209479c..760a3d9 100644
--- a/experiments/boft_factorize.py
+++ b/experiments/boft_factorize.py
@@ -29,7 +29,7 @@ def butterfly_factor(dimension: int, factor: int = -1) -> tuple[int, int]:
 
     if n == 0:
         raise ValueError(
-            f"It is impossible to decompose {dimension} with factor {factor} under BOFT constrains."
+            f"It is impossible to decompose {dimension} with factor {factor} under BOFT constraints."
         )
 
     log_butterfly_factorize(dimension, factor, (dimension // n, n))
diff --git a/experiments/concept.py b/experiments/concept.py
deleted file mode 100644
index b02357d..0000000
--- a/experiments/concept.py
+++ /dev/null
@@ -1,278 +0,0 @@
-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
-        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
deleted file mode 100644
index 289fa31..0000000
--- a/experiments/concept_neuron.py
+++ /dev/null
@@ -1,263 +0,0 @@
-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 *
-from lycoris.kohya_model_utils import *
-from lycoris.locon import LoConModule
-from lycoris.loha import LohaModule
-
-
-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.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 - 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",
-        "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, **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 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,
-        )
-        print(f"create Concept 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/module_count.py b/experiments/module_count.py
index ea93b02..a4133a6 100644
--- a/experiments/module_count.py
+++ b/experiments/module_count.py
@@ -2,7 +2,7 @@
 import torch.nn as nn
 import torch.nn.functional as F
 
-from lycoris.kohya.model_utils import load_models_from_stable_diffusion_checkpoint
+from library.model_util import load_models_from_stable_diffusion_checkpoint
 import lycoris.kohya as network_module
 
 
diff --git a/lycoris/__init__.py b/lycoris/__init__.py
index a33d22b..af78faa 100644
--- a/lycoris/__init__.py
+++ b/lycoris/__init__.py
@@ -1,5 +1,8 @@
+try:
+    from . import kohya
+except:
+    pass
 from . import (
-    kohya,
     modules,
     utils,
 )
diff --git a/lycoris/config.py b/lycoris/config.py
index fc22a65..b011458 100644
--- a/lycoris/config.py
+++ b/lycoris/config.py
@@ -6,6 +6,7 @@
             "ResnetBlock2D",
             "Downsample2D",
             "Upsample2D",
+            "HunYuanDiTBlock",
         ],
         "unet_target_name": [
             "conv_in",
@@ -13,7 +14,12 @@
             "time_embedding.linear_1",
             "time_embedding.linear_2",
         ],
-        "text_encoder_target_module": ["CLIPAttention", "CLIPMLP"],
+        "text_encoder_target_module": [
+            "CLIPAttention",
+            "CLIPMLP",
+            "MT5Block",
+            "BertLayer",
+        ],
         "text_encoder_target_name": [],
     },
     "full-lin": {
@@ -21,30 +27,47 @@
         "unet_target_module": [
             "Transformer2DModel",
             "ResnetBlock2D",
+            "HunYuanDiTBlock",
         ],
         "unet_target_name": [
             "time_embedding.linear_1",
             "time_embedding.linear_2",
         ],
-        "text_encoder_target_module": ["CLIPAttention", "CLIPMLP"],
+        "text_encoder_target_module": [
+            "CLIPAttention",
+            "CLIPMLP",
+            "MT5Block",
+            "BertLayer",
+        ],
         "text_encoder_target_name": [],
     },
     "attn-mlp": {
         "enable_conv": False,
         "unet_target_module": [
             "Transformer2DModel",
+            "HunYuanDiTBlock",
         ],
         "unet_target_name": [],
-        "text_encoder_target_module": ["CLIPAttention", "CLIPMLP"],
+        "text_encoder_target_module": [
+            "CLIPAttention",
+            "CLIPMLP",
+            "MT5Block",
+            "BertLayer",
+        ],
         "text_encoder_target_name": [],
     },
     "attn-only": {
         "enable_conv": False,
         "unet_target_module": [
             "CrossAttention",
+            "SelfAttention",
         ],
         "unet_target_name": [],
-        "text_encoder_target_module": ["CLIPAttention"],
+        "text_encoder_target_module": [
+            "CLIPAttention",
+            "BertAttention",
+            "MT5LayerSelfAttention",
+        ],
         "text_encoder_target_name": [],
     },
     "unet-only": {
@@ -54,6 +77,7 @@
             "ResnetBlock2D",
             "Downsample2D",
             "Upsample2D",
+            "HunYuanDiTBlock",
         ],
         "unet_target_name": [
             "conv_in",
@@ -68,6 +92,7 @@
         "enable_conv": False,
         "unet_target_module": [
             "Transformer2DModel",
+            "HunYuanDiTBlock",
         ],
         "unet_target_name": [],
         "text_encoder_target_module": [],
@@ -83,4 +108,15 @@
         "text_encoder_target_module": [],
         "text_encoder_target_name": [],
     },
+    "ia3": {
+        "enable_conv": False,
+        "unet_target_module": [],
+        "unet_target_name": ["to_k", "to_v", "ff.net.2"],
+        "text_encoder_target_module": [],
+        "text_encoder_target_name": ["k_proj", "v_proj", "mlp.fc2"],
+        "name_algo_map": {
+            "mlp.fc2": {"train_on_input": True},
+            "ff.net.2": {"train_on_input": True},
+        },
+    },
 }
diff --git a/lycoris/functional/__init__.py b/lycoris/functional/__init__.py
new file mode 100644
index 0000000..a489a16
--- /dev/null
+++ b/lycoris/functional/__init__.py
@@ -0,0 +1,9 @@
+from .general import (
+    rebuild_tucker,
+    factorization,
+    power2factorization,
+    FUNC_LIST,
+    tucker_weight,
+    tucker_weight_from_conv,
+    apply_dora_scale,
+)
diff --git a/lycoris/functional/boft.py b/lycoris/functional/boft.py
new file mode 100644
index 0000000..2d95397
--- /dev/null
+++ b/lycoris/functional/boft.py
@@ -0,0 +1,99 @@
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from einops import rearrange
+
+from .general import power2factorization, FUNC_LIST
+from .diag_oft import get_r
+
+
+def weight_gen(org_weight, max_block_size, boft_m=-1, rescale=False):
+    """### boft_weight_gen
+
+    Args:
+        org_weight (torch.Tensor): the weight tensor
+        max_block_size (int): max block size
+        rescale (bool, optional): whether to rescale the weight. Defaults to False.
+
+    Returns:
+        torch.Tensor: oft_blocks[, rescale_weight]
+    """
+    out_dim, *rest = org_weight.shape
+    block_size, block_num = power2factorization(out_dim, max_block_size)
+    max_boft_m = sum(int(i) for i in f"{block_num-1:b}") + 1
+    if boft_m == -1:
+        boft_m = max_boft_m
+    boft_m = min(boft_m, max_boft_m)
+    oft_blocks = torch.zeros(boft_m, block_num, block_size, block_size)
+    if rescale is not None:
+        return oft_blocks, torch.ones(out_dim, *[1] * len(rest))
+    else:
+        return oft_blocks, None
+
+
+def diff_weight(org_weight, oft_blocks, rescale=None, constraint=None):
+    """### boft_diff_weight
+
+    Args:
+        TODO
+
+    Returns:
+        torch.Tensor: ΔW
+    """
+    m, num, b, _ = oft_blocks.shape
+    r_b = b // 2
+    I = torch.eye(b, device=oft_blocks.device)
+    r = get_r(oft_blocks, I, constraint)
+    inp = org = org_weight.to(dtype=r.dtype)
+
+    for i in range(m):
+        bi = r[i]  # b_num, b_size, b_size
+        inp = rearrange(inp, "(c g k) ... -> (c k g) ...", g=2, k=2**i * r_b)
+        inp = rearrange(inp, "(d b) ... -> d b ...", b=b)
+        inp = torch.einsum("b i j, b j ... -> b i ...", bi, inp)
+        inp = rearrange(inp, "d b ... -> (d b) ...")
+        inp = rearrange(inp, "(c k g) ... -> (c g k) ...", g=2, k=2**i * r_b)
+
+    if rescale is not None:
+        inp = inp * rescale
+
+    return inp - org
+
+
+def bypass_forward_diff(
+    org_out, oft_blocks, rescale=None, constraint=None, need_transpose=False
+):
+    """### boft_bypass_forward_diff
+
+    Args:
+        TODO
+
+    Returns:
+        torch.Tensor: output tensor
+    """
+    m, num, b, _ = oft_blocks.shape
+    r_b = b // 2
+    I = torch.eye(b, device=oft_blocks.device)
+    r = get_r(oft_blocks, I, constraint)
+    inp = org = org_out.to(dtype=r.dtype)
+    if need_transpose:
+        inp = org = inp.transpose(1, -1)
+
+    for i in range(m):
+        bi = r[i]  # b_num, b_size, b_size
+        inp = rearrange(inp, "... (c g k) ->... (c k g)", g=2, k=2**i * r_b)
+        inp = rearrange(inp, "... (d b) -> ... d b", b=b)
+        inp = torch.einsum("b i j, ... b j -> ... b i", bi, inp)
+        inp = rearrange(inp, "... d b -> ... (d b)")
+        inp = rearrange(inp, "... (c k g) -> ... (c g k)", g=2, k=2**i * r_b)
+
+    if rescale is not None:
+        inp = inp * rescale.transpose(0, -1)
+
+    inp = inp - org
+    if need_transpose:
+        inp = inp.transpose(1, -1)
+    return inp
diff --git a/lycoris/functional/diag_oft.py b/lycoris/functional/diag_oft.py
new file mode 100644
index 0000000..bfda8b9
--- /dev/null
+++ b/lycoris/functional/diag_oft.py
@@ -0,0 +1,103 @@
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .general import factorization, FUNC_LIST
+
+
+def get_r(oft_blocks, I=None, constraint=0):
+    if I is None:
+        I = torch.eye(oft_blocks.shape[-1], device=oft_blocks.device)
+    if I.ndim < oft_blocks.ndim:
+        for _ in range(oft_blocks.ndim - I.ndim):
+            I = I.unsqueeze(0)
+    # for Q = -Q^T
+    q = oft_blocks - oft_blocks.transpose(-1, -2)
+    normed_q = q
+    if constraint is not None and constraint > 0:
+        q_norm = torch.norm(q) + 1e-8
+        if q_norm > constraint:
+            normed_q = q * constraint / q_norm
+    # use float() to prevent unsupported type
+    r = (I + normed_q) @ (I - normed_q).float().inverse()
+    return r
+
+
+def weight_gen(org_weight, max_block_size, rescale=False):
+    """### weight_gen
+
+    Args:
+        org_weight (torch.Tensor): the weight tensor
+        max_block_size (int): max block size
+        rescale (bool, optional): whether to rescale the weight. Defaults to False.
+
+    Returns:
+        torch.Tensor: oft_blocks[, rescale_weight]
+    """
+    out_dim, *rest = org_weight.shape
+    block_size, block_num = factorization(out_dim, max_block_size)
+    oft_blocks = torch.zeros(block_num, block_size, block_size)
+    if rescale:
+        return oft_blocks, torch.ones(out_dim, *[1] * len(rest))
+    else:
+        return oft_blocks, None
+
+
+def diff_weight(org_weight, oft_blocks, rescale=None, constraint=None):
+    """### diff_weight
+
+    Args:
+        TODO
+
+    Returns:
+        torch.Tensor: ΔW
+    """
+    I = torch.eye(oft_blocks.shape[1], device=oft_blocks.device)
+    r = get_r(oft_blocks, I, constraint)
+
+    block_num, block_size, _ = oft_blocks.shape
+    _, *shape = org_weight.shape
+    org_weight = org_weight.to(dtype=r.dtype)
+    org_weight = org_weight.view(block_num, block_size, *shape)
+    # 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 - I,
+        org_weight,
+    ).view(-1, *shape)
+    if rescale is not None:
+        weight = rescale * weight
+        weight = weight + (rescale - 1) * org_weight
+    return weight
+
+
+def bypass_forward_diff(
+    org_out, oft_blocks, rescale=None, constraint=None, need_transpose=False
+):
+    """### bypass_forward_diff
+
+    Args:
+        TODO
+
+    Returns:
+        torch.Tensor: output tensor
+    """
+    block_num, block_size, _ = oft_blocks.shape
+    I = torch.eye(block_size, device=oft_blocks.device)
+    r = get_r(oft_blocks, I, constraint)
+    if need_transpose:
+        org_out = org_out.transpose(1, -1)
+    org_out = org_out.to(dtype=r.dtype)
+    *shape, _ = org_out.shape
+    oft_out = torch.einsum(
+        "k n m, ... k n -> ... k m", r - I, org_out.view(*shape, block_num, block_size)
+    )
+    out = oft_out.view(*shape, -1)
+    if rescale is not None:
+        out = rescale.transpose(-1, 0) * out
+        out = out + (rescale - 1).transpose(-1, 0) * org_out
+    if need_transpose:
+        out = out.transpose(1, -1)
+    return out
diff --git a/lycoris/functional/general.py b/lycoris/functional/general.py
new file mode 100644
index 0000000..001e38a
--- /dev/null
+++ b/lycoris/functional/general.py
@@ -0,0 +1,106 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+FUNC_LIST = [None, None, F.linear, F.conv1d, F.conv2d, F.conv3d]
+
+
+def rebuild_tucker(t, wa, wb):
+    rebuild2 = torch.einsum("i j ..., i p, j r -> p r ...", t, wa, wb)
+    return rebuild2
+
+
+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.
+    second value is a value for weight.
+
+    Because of non-commutative property, A⊗B ≠ B⊗A. Meaning of two matrices is slightly different.
+
+    examples)
+    factor
+        -1               2                4               8               16               ...
+    127 -> 1, 127   127 -> 1, 127    127 -> 1, 127   127 -> 1, 127   127 -> 1, 127
+    128 -> 8, 16    128 -> 2, 64     128 -> 4, 32    128 -> 8, 16    128 -> 8, 16
+    250 -> 10, 25   250 -> 2, 125    250 -> 2, 125   250 -> 5, 50    250 -> 10, 25
+    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
+        if m > n:
+            n, m = m, n
+        return m, n
+    if factor < 0:
+        factor = dimension
+    m, n = 1, dimension
+    length = m + n
+    while m < n:
+        new_m = m + 1
+        while dimension % new_m != 0:
+            new_m += 1
+        new_n = dimension // new_m
+        if new_m + new_n > length or new_m > factor:
+            break
+        else:
+            m, n = new_m, new_n
+    if m > n:
+        n, m = m, n
+    return m, n
+
+
+def power2factorization(dimension: int, factor: int = -1) -> tuple[int, int]:
+    """
+    m = 2k
+    n = 2**p
+    m*n = dim
+    """
+    if factor == -1:
+        factor = dimension
+
+    # Find the first solution and check if it is even doable
+    m = n = 0
+    while m <= factor:
+        m += 2
+        while dimension % m != 0 and m < dimension:
+            m += 2
+        if m > factor:
+            break
+        if sum(int(i) for i in f"{dimension//m:b}") == 1:
+            n = dimension // m
+
+    return dimension // n, n
+
+
+def tucker_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 ..., i m, n j -> i j ...", mid, up, down)
+
+
+def tucker_weight(wa, wb, t):
+    temp = torch.einsum("i j ..., j r -> i r ...", t, wb)
+    return torch.einsum("i j ..., i r -> r j ...", temp, wa)
+
+
+def apply_dora_scale(org_weight, rebuild, dora_scale, scale):
+    dora_norm_dims = org_weight.dim() - 1
+    weight = org_weight + rebuild
+    weight = weight.to(dora_scale.dtype)
+    weight_norm = (
+        weight.transpose(0, 1)
+        .reshape(weight.shape[1], -1)
+        .norm(dim=1, keepdim=True)
+        .reshape(weight.shape[1], *[1] * dora_norm_dims)
+        .transpose(0, 1)
+    )
+    merged_scale1 = weight / weight_norm * dora_scale
+    diff_weight = merged_scale1 - org_weight
+    return org_weight + diff_weight * scale
diff --git a/lycoris/functional/locon.py b/lycoris/functional/locon.py
new file mode 100644
index 0000000..e1aad22
--- /dev/null
+++ b/lycoris/functional/locon.py
@@ -0,0 +1,87 @@
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .general import rebuild_tucker, FUNC_LIST
+
+
+def weight_gen(org_weight, rank, tucker=True):
+    """### weight_gen
+
+    Args:
+        org_weight (torch.Tensor): the weight tensor
+        rank (int): low rank
+
+    Returns:
+        torch.Tensor: down, up[, mid]
+    """
+    out_dim, in_dim, *k = org_weight.shape
+    if k and tucker:
+        down = torch.empty(rank, in_dim, *(1 for _ in k))
+        up = torch.empty(out_dim, rank, *(1 for _ in k))
+        mid = torch.empty(rank, rank, *k)
+        nn.init.kaiming_uniform_(down, a=math.sqrt(5))
+        nn.init.constant_(up, 0)
+        nn.init.kaiming_uniform_(mid, a=math.sqrt(5))
+        return down, up, mid
+    else:
+        down = torch.empty(rank, in_dim)
+        up = torch.empty(out_dim, rank)
+        nn.init.kaiming_uniform_(down, a=math.sqrt(5))
+        nn.init.constant_(up, 0)
+        return down, up
+
+
+def diff_weight(d, u, m=None, gamma=1.0):
+    """### diff_weight
+
+    Get ΔW = BA, where BA is low rank decomposition
+
+    Args:
+        d (torch.Tensor): weight of down proj linear/conv layer
+        u (torch.Tensor): weight of up proj linear/conv layer
+        m (torch.Tensor, optional): middle weight of tucker decomposition
+        gamma (float, optional): scale factor, normally alpha/rank here
+
+    Returns:
+        torch.Tensor: ΔW
+    """
+    R, I, *k = d.shape
+    O, R, *_ = u.shape
+    u = u * gamma
+
+    if m is None:
+        result = u.reshape(-1, u.size(1)) @ d.reshape(d.size(0), -1)
+    else:
+        R, R, *k = m.shape
+        u = u.reshape(u.size(0), -1).transpose(0, 1)
+        d = d.reshape(d.size(0), -1)
+        result = rebuild_tucker(m, u, d)
+    return result.reshape(O, I, *k)
+
+
+def bypass_forward_diff(x, d, u, m=None, gamma=1.0, extra_args={}):
+    """### bypass_forward_diff
+
+    Args:
+        x (torch.Tensor): input tensor
+        d (torch.Tensor): weight of down proj linear/conv layer
+        u (torch.Tensor): weight of up proj linear/conv layer
+        m (torch.Tensor, optional): middle weight of tucker decomposition
+        gamma (float, optional): scale factor, normally alpha/rank here
+        extra_args (dict, optional): extra args for forward func, \
+            e.g. padding, stride for Conv1/2/3d
+
+    Returns:
+        torch.Tensor: output tensor
+    """
+    if m is not None:
+        down = FUNC_LIST[d.dim()](x, d)
+        mid = FUNC_LIST[d.dim()](down, m, **extra_args)
+        up = FUNC_LIST[d.dim()](mid, u)
+    else:
+        down = FUNC_LIST[d.dim()](x, d, **extra_args)
+        up = FUNC_LIST[d.dim()](down, u)
+    return up * gamma
diff --git a/lycoris/functional/loha.py b/lycoris/functional/loha.py
new file mode 100644
index 0000000..e32f427
--- /dev/null
+++ b/lycoris/functional/loha.py
@@ -0,0 +1,168 @@
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .general import FUNC_LIST
+
+
+class HadaWeight(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, w1d, w1u, w2d, w2u, scale=torch.tensor(1)):
+        ctx.save_for_backward(w1d, w1u, w2d, w2u, scale)
+        diff_weight = ((w1u @ w1d) * (w2u @ w2d)) * scale
+        return diff_weight
+
+    @staticmethod
+    def backward(ctx, grad_out):
+        (w1d, w1u, w2d, w2u, scale) = ctx.saved_tensors
+        grad_out = grad_out * scale
+        temp = grad_out * (w2u @ w2d)
+        grad_w1u = temp @ w1d.T
+        grad_w1d = w1u.T @ temp
+
+        temp = grad_out * (w1u @ w1d)
+        grad_w2u = temp @ w2d.T
+        grad_w2d = w2u.T @ temp
+
+        del temp
+        return grad_w1d, grad_w1u, grad_w2d, grad_w2u, None
+
+
+class HadaWeightTucker(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, t1, w1d, w1u, t2, w2d, w2u, scale=torch.tensor(1)):
+        ctx.save_for_backward(t1, w1d, w1u, t2, w2d, w2u, scale)
+
+        rebuild1 = torch.einsum("i j ..., j r, i p -> p r ...", t1, w1d, w1u)
+        rebuild2 = torch.einsum("i j ..., j r, i p -> p r ...", t2, w2d, w2u)
+
+        return rebuild1 * rebuild2 * scale
+
+    @staticmethod
+    def backward(ctx, grad_out):
+        (t1, w1d, w1u, t2, w2d, w2u, scale) = ctx.saved_tensors
+        grad_out = grad_out * scale
+
+        temp = torch.einsum("i j ..., j r -> i r ...", t2, w2d)
+        rebuild = torch.einsum("i j ..., i r -> r j ...", temp, w2u)
+
+        grad_w = rebuild * grad_out
+        del rebuild
+
+        grad_w1u = torch.einsum("r j ..., i j ... -> r i", temp, grad_w)
+        grad_temp = torch.einsum("i j ..., i r -> r j ...", grad_w, w1u.T)
+        del grad_w, temp
+
+        grad_w1d = torch.einsum("i r ..., i j ... -> r j", t1, grad_temp)
+        grad_t1 = torch.einsum("i j ..., j r -> i r ...", grad_temp, w1d.T)
+        del grad_temp
+
+        temp = torch.einsum("i j ..., j r -> i r ...", t1, w1d)
+        rebuild = torch.einsum("i j ..., i r -> r j ...", temp, w1u)
+
+        grad_w = rebuild * grad_out
+        del rebuild
+
+        grad_w2u = torch.einsum("r j ..., i j ... -> r i", temp, grad_w)
+        grad_temp = torch.einsum("i j ..., i r -> r j ...", grad_w, w2u.T)
+        del grad_w, temp
+
+        grad_w2d = torch.einsum("i r ..., i j ... -> r j", t2, grad_temp)
+        grad_t2 = torch.einsum("i j ..., j r -> i r ...", grad_temp, w2d.T)
+        del grad_temp
+        return grad_t1, grad_w1d, grad_w1u, grad_t2, grad_w2d, grad_w2u, None
+
+
+def make_weight(w1d, w1u, w2d, w2u, scale):
+    return HadaWeight.apply(w1d, w1u, w2d, w2u, scale)
+
+
+def make_weight_tucker(t1, w1d, w1u, t2, w2d, w2u, scale):
+    return HadaWeightTucker.apply(t1, w1d, w1u, t2, w2d, w2u, scale)
+
+
+def weight_gen(org_weight, rank, tucker=True):
+    """### weight_gen
+
+    Args:
+        org_weight (torch.Tensor): the weight tensor
+        rank (int): low rank
+
+    Returns:
+        torch.Tensor: w1d, w2d, w1u, w2u[, t1, t2]
+    """
+    out_dim, in_dim, *k = org_weight.shape
+    if k and tucker:
+        w1d = torch.empty(rank, in_dim)
+        w1u = torch.empty(rank, out_dim)
+        t1 = torch.empty(rank, rank, *k)
+        w2d = torch.empty(rank, in_dim)
+        w2u = torch.empty(rank, out_dim)
+        t2 = torch.empty(rank, rank, *k)
+        nn.init.normal_(t1, std=0.1)
+        nn.init.normal_(t2, std=0.1)
+    else:
+        w1d = torch.empty(rank, in_dim)
+        w1u = torch.empty(out_dim, rank)
+        w2d = torch.empty(rank, in_dim)
+        w2u = torch.empty(out_dim, rank)
+        t1 = t2 = None
+    nn.init.normal_(w1d, std=1)
+    nn.init.constant_(w1u, 0)
+    nn.init.normal_(w2d, std=1)
+    nn.init.normal_(w2u, std=0.1)
+    return w1d, w1u, w2d, w2u, t1, t2
+
+
+def diff_weight(w1d, w1u, w2d, w2u, t1=None, t2=None, gamma=1.0):
+    """### diff_weight
+
+    Get ΔW = BA, where BA is low rank decomposition
+
+    Args:
+        w1d, w2d (torch.Tensor): weight of down proj linear/conv layer
+        w1u, w2u (torch.Tensor): weight of up proj linear/conv layer
+        t1, t2 (torch.Tensor, optional): middle weight of tucker decomposition
+        gamma (float, optional): scale factor, normally alpha/rank here
+
+    Returns:
+        torch.Tensor: ΔW
+    """
+    if t1 is not None and t2 is not None:
+        R, I = w1d.shape
+        R, O = w1u.shape
+        R, R, *k = t1.shape
+        result = make_weight_tucker(t1, w1d, w1u, t2, w2d, w2u, gamma)
+    else:
+        R, I, *k = w1d.shape
+        O, R, *_ = w1u.shape
+        w1d = w1d.reshape(w1d.size(0), -1)
+        w1u = w1u.reshape(-1, w1u.size(1))
+        w2d = w2d.reshape(w2d.size(0), -1)
+        w2u = w1u.reshape(-1, w2u.size(1))
+        result = make_weight(w1d, w1u, w2d, w2u, gamma)
+
+    result = result.reshape(O, I, *k)
+    return result
+
+
+def bypass_forward_diff(
+    x, w1d, w1u, w2d, w2u, t1=None, t2=None, gamma=1.0, extra_args={}
+):
+    """### bypass_forward_diff
+
+    Args:
+        x (torch.Tensor): input tensor
+        w1d, w2d (torch.Tensor): weight of up proj linear/conv layer
+        w1u, w2u (torch.Tensor): weight of down proj linear/conv layer
+        gamma (float, optional): scale factor, normally alpha/rank here
+        extra_args (dict, optional): extra args for forward func, \
+            e.g. padding, stride for Conv1/2/3d
+
+    Returns:
+        torch.Tensor: output tensor
+    """
+    diff_w = diff_weight(w1d, w1u, w2d, w2u, t1, t2, gamma)
+    return FUNC_LIST[w1d.dim() if t1 is None else t1.dim()](x, diff_w, **extra_args)
diff --git a/lycoris/functional/lokr.py b/lycoris/functional/lokr.py
new file mode 100644
index 0000000..e252553
--- /dev/null
+++ b/lycoris/functional/lokr.py
@@ -0,0 +1,240 @@
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .general import rebuild_tucker, FUNC_LIST
+from .general import factorization
+
+
+def make_kron(w1, w2, scale):
+    for _ in range(w2.dim() - w1.dim()):
+        w1 = w1.unsqueeze(-1)
+    w2 = w2.contiguous()
+    rebuild = torch.kron(w1, w2)
+
+    return rebuild * scale
+
+
+def weight_gen(
+    org_weight,
+    rank,
+    tucker=True,
+    factor=-1,
+    decompose_both=False,
+    full_matrix=False,
+    unbalanced_factorization=False,
+):
+    """### weight_gen
+
+    Args:
+        org_weight (torch.Tensor): the weight tensor
+        rank (int): low rank
+
+    Returns:
+        torch.Tensor | None: w1, w1a, w1b, w2, w2a, w2b, t2
+    """
+    out_dim, in_dim, *k = org_weight.shape
+    w1 = w1a = w1b = None
+    w2 = w2a = w2b = None
+    t2 = None
+    use_w1 = use_w2 = False
+
+    if k:
+        k_size = k
+        shape = (out_dim, in_dim, *k_size)
+
+        in_m, in_n = factorization(in_dim, factor)
+        out_l, out_k = factorization(out_dim, factor)
+        if unbalanced_factorization:
+            out_l, out_k = out_k, out_l
+        shape = ((out_l, out_k), (in_m, in_n), *k_size)  # ((a, b), (c, d), *k_size)
+        tucker = tucker and any(i != 1 for i in k_size)
+        if (
+            decompose_both
+            and rank < max(shape[0][0], shape[1][0]) / 2
+            and not full_matrix
+        ):
+            w1a = torch.empty(shape[0][0], rank)
+            w1b = torch.empty(rank, shape[1][0])
+        else:
+            use_w1 = True
+            w1 = torch.empty(shape[0][0], shape[1][0])  # a*c, 1-mode
+
+        if rank >= max(shape[0][1], shape[1][1]) / 2 or full_matrix:
+            use_w2 = True
+            w2 = torch.empty(shape[0][1], shape[1][1], *k_size)
+        elif tucker:
+            t2 = torch.empty(rank, rank, *shape[2:])
+            w2a = torch.empty(rank, shape[0][1])  # b, 1-mode
+            w2b = torch.empty(rank, shape[1][1])  # d, 2-mode
+        else:  # Conv2d not tucker
+            # bigger part. weight and LoRA. [b, dim] x [dim, d*k1*k2]
+            w2a = torch.empty(shape[0][1], rank)
+            w2b = torch.empty(rank, shape[1][1], *shape[2:])
+            # w1 ⊗ (w2a x w2b) = (a, b)⊗((c, dim)x(dim, d*k1*k2)) = (a, b)⊗(c, d*k1*k2) = (ac, bd*k1*k2)
+    else:  # Linear
+        shape = (out_dim, in_dim)
+
+        in_m, in_n = factorization(in_dim, factor)
+        out_l, out_k = factorization(out_dim, factor)
+        if unbalanced_factorization:
+            out_l, out_k = out_k, out_l
+        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 rank < max(shape[0][0], shape[1][0]) / 2:
+            w1a = torch.empty(shape[0][0], rank)
+            w1b = torch.empty(rank, shape[1][0])
+        else:
+            use_w1 = True
+            w1 = torch.empty(shape[0][0], shape[1][0])  # a*c, 1-mode
+        if rank < max(shape[0][1], shape[1][1]) / 2:
+            # bigger part. weight and LoRA. [b, dim] x [dim, d]
+            w2a = torch.empty(shape[0][1], rank)
+            w2b = torch.empty(rank, shape[1][1])
+            # w1 ⊗ (w2a x w2b) = (a, b)⊗((c, dim)x(dim, d)) = (a, b)⊗(c, d) = (ac, bd)
+        else:
+            use_w2 = True
+            w2 = torch.empty(shape[0][1], shape[1][1])
+
+    if use_w2:
+        torch.nn.init.constant_(w2, 1)
+    else:
+        if tucker:
+            torch.nn.init.kaiming_uniform_(t2, a=math.sqrt(5))
+        torch.nn.init.kaiming_uniform_(w2a, a=math.sqrt(5))
+        torch.nn.init.constant_(w2b, 1)
+
+    if use_w1:
+        torch.nn.init.kaiming_uniform_(w1, a=math.sqrt(5))
+    else:
+        torch.nn.init.kaiming_uniform_(w1a, a=math.sqrt(5))
+        torch.nn.init.kaiming_uniform_(w1b, a=math.sqrt(5))
+
+    return w1, w1a, w1b, w2, w2a, w2b, t2
+
+
+def diff_weight(w1, w1a, w1b, w2, w2a, w2b, t, gamma=1.0):
+    """### diff_weight
+
+    Args:
+        gamma (float, optional): scale factor, normally alpha/rank here
+
+    Returns:
+        torch.Tensor: ΔW
+    """
+    if w1a is not None:
+        rank = w1a.shape[1]
+    elif w2a is not None:
+        rank = w2a.shape[1]
+    else:
+        rank = gamma
+    scale = gamma / rank
+    if w1 is None:
+        w1 = w1a @ w1b
+    if w2 is None:
+        if t is None:
+            r, o, *k = w2b.shape
+            w2 = w2a @ w2b.view(r, -1)
+            w2 = w2.view(-1, o, *k)
+        else:
+            w2 = rebuild_tucker(t, w2a, w2b)
+    return make_kron(w1, w2, scale)
+
+
+def bypass_forward_diff(h, w1, w1a, w1b, w2, w2a, w2b, t, gamma=1.0, extra_args={}):
+    """### bypass_forward_diff
+
+    Args:
+        gamma (float, optional): scale factor, normally alpha/rank here
+        extra_args (dict, optional): extra args for forward func, \
+            e.g. padding, stride for Conv1/2/3d
+
+    Returns:
+        torch.Tensor: output tensor
+    """
+    use_w1 = w1 is not None
+    use_w2 = w2 is not None
+    tucker = t is not None
+    dim = t.dim() if tucker else w2.dim() if w2 is not None else w2b.dim()
+    rank = w1b.size(0) if not use_w1 else w2b.size(0) if not use_w2 else gamma
+    scale = gamma / rank
+    is_conv = dim > 2
+    op = FUNC_LIST[dim]
+
+    if is_conv:
+        kw_dict = extra_args
+    else:
+        kw_dict = {}
+
+    if use_w2:
+        ba = w2
+    else:
+        a = w2b
+        b = w2a
+
+        if t is not None:
+            a = a.view(*a.shape, *[1] * (dim - 2))
+            b = b.view(*b.shape, *[1] * (dim - 2))
+        elif is_conv:
+            b = b.view(*b.shape, *[1] * (dim - 2))
+
+    if use_w1:
+        c = w1
+    else:
+        c = w1a @ w1b
+    uq = c.size(1)
+
+    if is_conv:
+        # (b, uq), vq, ...
+        B, _, *rest = h.shape
+        h_in_group = h.reshape(B * uq, -1, *rest)
+    else:
+        # b, ..., uq, vq
+        h_in_group = h.reshape(*h.shape[:-1], uq, -1)
+
+    if use_w2:
+        hb = op(h_in_group, ba, **kw_dict)
+    else:
+        if is_conv:
+            if tucker:
+                ha = op(h_in_group, a)
+                ht = op(ha, t, **kw_dict)
+                hb = op(ht, b)
+            else:
+                ha = op(h_in_group, a, **kw_dict)
+                hb = op(ha, b)
+        else:
+            ha = op(h_in_group, a, **kw_dict)
+            hb = op(ha, b)
+
+    if is_conv:
+        # (b, uq), vp, ..., f
+        # -> b, uq, vp, ..., f
+        # -> b, f, vp, ..., uq
+        hb = hb.view(B, -1, *hb.shape[1:])
+        h_cross_group = hb.transpose(1, -1)
+    else:
+        # b, ..., uq, vq
+        # -> b, ..., vq, uq
+        h_cross_group = hb.transpose(-1, -2)
+
+    hc = F.linear(h_cross_group, c)
+    if is_conv:
+        # b, f, vp, ..., up
+        # -> b, up, vp, ... ,f
+        # -> b, c, ..., f
+        hc = hc.transpose(1, -1)
+        h = hc.reshape(B, -1, *hc.shape[3:])
+    else:
+        # b, ..., vp, up
+        # -> b, ..., up, vp
+        # -> b, ..., c
+        hc = hc.transpose(-1, -2)
+        h = hc.reshape(*hc.shape[:-2], -1)
+
+    return h * scale
diff --git a/lycoris/hcp/__init__.py b/lycoris/hcp/__init__.py
deleted file mode 100644
index 9bd7dbf..0000000
--- a/lycoris/hcp/__init__.py
+++ /dev/null
@@ -1,4 +0,0 @@
-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
deleted file mode 100644
index 937d536..0000000
--- a/lycoris/hcp/base.py
+++ /dev/null
@@ -1,153 +0,0 @@
-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
deleted file mode 100644
index 7992fe6..0000000
--- a/lycoris/hcp/diag_oft.py
+++ /dev/null
@@ -1,62 +0,0 @@
-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, log_oft_factorize
-
-
-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(out_dim))
-        log_oft_factorize(
-            dim=out_dim,
-            factor=self.dim,
-            num=self.block_num,
-            bdim=self.block_size,
-        )
-
-    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) ...")
-        if self.rescaled:
-            weight = self.rescale.to(weight) * weight
-        # 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
deleted file mode 100644
index 891760e..0000000
--- a/lycoris/hcp/loha.py
+++ /dev/null
@@ -1,85 +0,0 @@
-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_tucker
-
-
-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_tucker(
-                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
deleted file mode 100644
index 03bcef8..0000000
--- a/lycoris/hcp/lokr.py
+++ /dev/null
@@ -1,131 +0,0 @@
-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_tucker, 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_tucker(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),
-        )
-        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_tucker(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.py
similarity index 59%
rename from lycoris/kohya/__init__.py
rename to lycoris/kohya.py
index b65ddb5..267768e 100644
--- a/lycoris/kohya/__init__.py
+++ b/lycoris/kohya.py
@@ -1,52 +1,40 @@
-# network module for kohya
-# reference:
-# https://github.com/microsoft/LoRA/blob/main/loralib/layers.py
-# https://github.com/cloneofsimo/lora/blob/master/lora_diffusion/lora.py
-# https://github.com/kohya-ss/sd-scripts/blob/main/networks/lora.py
-
 import os
 import re
+import logging
 from typing import List
 
 import torch
 
-from .utils import *
-from ..wrapper import LycorisNetwork
-from ..modules.locon import LoConModule
-from ..modules.loha import LohaModule
-from ..modules.ia3 import IA3Module
-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.boft import ButterflyOFTModule
-from ..modules import make_module
-
-from ..config import PRESET
-from ..utils.preset import read_preset
-from ..utils import get_module, str_bool
-from ..logging import logger
-
-
-network_module_dict = {
-    "lora": LoConModule,
-    "locon": LoConModule,
-    "loha": LohaModule,
-    "ia3": IA3Module,
-    "lokr": LokrModule,
-    "dylora": DyLoraModule,
-    "glora": GLoRAModule,
-    "full": FullModule,
-    "diag-oft": DiagOFTModule,
-    "boft": ButterflyOFTModule,
-}
+from .utils import precalculate_safetensors_hashes
+from .wrapper import LycorisNetwork, network_module_dict, deprecated_arg_dict
+from .modules.locon import LoConModule
+from .modules.loha import LohaModule
+from .modules.ia3 import IA3Module
+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.boft import ButterflyOFTModule
+from .modules import make_module, get_module
+
+from .config import PRESET
+from .utils.preset import read_preset
+from .utils import str_bool
+from .logging import logger
 
 
 def create_network(
     multiplier, network_dim, network_alpha, vae, text_encoder, unet, **kwargs
 ):
+    for key, value in kwargs.items():
+        if key in deprecated_arg_dict:
+            logger.warning(
+                f"{key} is deprecated. Please use {deprecated_arg_dict[key]} instead.",
+                stacklevel=2,
+            )
+            kwargs[deprecated_arg_dict[key]] = value
     if network_dim is None:
         network_dim = 4  # default
     conv_dim = int(kwargs.get("conv_dim", network_dim) or network_dim)
@@ -61,19 +49,19 @@ def create_network(
         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:
-        logger.warning(
-            "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)
+    block_size = int(kwargs.get("block_size", None) or 4)
     train_norm = str_bool(kwargs.get("train_norm", False))
-    constrain = float(kwargs.get("constrain", 0) or 0)
+    constraint = float(kwargs.get("constraint", None) or 0)
     rescaled = str_bool(kwargs.get("rescaled", False))
     weight_decompose = str_bool(kwargs.get("dora_wd", False))
     full_matrix = str_bool(kwargs.get("full_matrix", False))
     bypass_mode = str_bool(kwargs.get("bypass_mode", False))
+    rs_lora = str_bool(kwargs.get("rs_lora", False))
+    unbalanced_factorization = str_bool(kwargs.get("unbalanced_factorization", False))
+
+    if unbalanced_factorization:
+        logger.info("Unbalanced factorization for LoKr is enabled")
 
     if bypass_mode:
         logger.info("Bypass mode is enabled")
@@ -84,68 +72,45 @@ def create_network(
     if full_matrix:
         logger.info("Full matrix mode for LoKr is enabled")
 
-    if algo == "glora" and conv_dim > 0:
-        conv_dim = 0
-        logger.info("Disable conv layer for GLoRA")
-
-    preset = kwargs.get("preset", "full")
-    if preset not in PRESET:
-        preset = read_preset(preset)
+    preset_str = kwargs.get("preset", "full")
+    if preset_str not in PRESET:
+        preset = read_preset(preset_str)
     else:
-        preset = PRESET[preset]
+        preset = PRESET[preset_str]
     assert preset is not None
     LycorisNetworkKohya.apply_preset(preset)
 
     logger.info(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)
-    ):
-        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.
-        """
-        logger.warning(warning_msg, stacklevel=2)
-
-    if algo == "ia3":
-        network = IA3Network(
-            text_encoder,
-            unet,
-            multiplier=multiplier,
-        )
-    else:
-        network = LycorisNetworkKohya(
-            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,
-            weight_decompose=weight_decompose,
-            full_matrix=full_matrix,
-            bypass_mode=bypass_mode,
-        )
+    if algo == "ia3" and preset_str != "ia3":
+        logger.warning("It is recommended to use preset ia3 for IA^3 algorithm")
 
-    if algo == "dylora":
-        # dylora didn't support scale weight norm yet
-        delattr(type(network), "apply_max_norm_regularization")
+    network = LycorisNetworkKohya(
+        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,
+        constraint=constraint,
+        rescaled=rescaled,
+        weight_decompose=weight_decompose,
+        full_matrix=full_matrix,
+        bypass_mode=bypass_mode,
+        rs_lora=rs_lora,
+        unbalanced_factorization=unbalanced_factorization,
+    )
 
     return network
 
@@ -203,10 +168,14 @@ def create_network_from_weights(
             if lora_name in te_loras:
                 te_loras[lora_name] = modules
 
+    original_level = logger.level
+    logger.setLevel(logging.ERROR)
     network = LycorisNetworkKohya(text_encoder, unet)
     network.unet_loras = []
     network.text_encoder_loras = []
+    logger.setLevel(original_level)
 
+    logger.info("Loading UNet Modules from state dict...")
     for lora_name, orig_modules in unet_loras.items():
         if orig_modules is None:
             continue
@@ -214,7 +183,9 @@ def create_network_from_weights(
         module = make_module(lyco_type, params, lora_name, orig_modules)
         if module is not None:
             network.unet_loras.append(module)
+    logger.info(f"{len(network.unet_loras)} Modules Loaded")
 
+    logger.info("Loading TE Modules from state dict...")
     for lora_name, orig_modules in te_loras.items():
         if orig_modules is None:
             continue
@@ -222,6 +193,7 @@ def create_network_from_weights(
         module = make_module(lyco_type, params, lora_name, orig_modules)
         if module is not None:
             network.text_encoder_loras.append(module)
+    logger.info(f"{len(network.text_encoder_loras)} Modules Loaded")
 
     for lora in network.unet_loras + network.text_encoder_loras:
         lora.multiplier = multiplier
@@ -241,6 +213,7 @@ class LycorisNetworkKohya(LycorisNetwork):
         "ResnetBlock2D",
         "Downsample2D",
         "Upsample2D",
+        "HunYuanDiTBlock",
     ]
     UNET_TARGET_REPLACE_NAME = [
         "conv_in",
@@ -248,7 +221,12 @@ class LycorisNetworkKohya(LycorisNetwork):
         "time_embedding.linear_1",
         "time_embedding.linear_2",
     ]
-    TEXT_ENCODER_TARGET_REPLACE_MODULE = ["CLIPAttention", "CLIPMLP"]
+    TEXT_ENCODER_TARGET_REPLACE_MODULE = [
+        "CLIPAttention",
+        "CLIPMLP",
+        "MT5Block",
+        "BertLayer",
+    ]
     TEXT_ENCODER_TARGET_REPLACE_NAME = []
     LORA_PREFIX_UNET = "lora_unet"
     LORA_PREFIX_TEXT_ENCODER = "lora_te"
@@ -349,7 +327,9 @@ def create_single_module(
             if isinstance(module, torch.nn.Linear) and lora_dim > 0:
                 dim = dim or lora_dim
                 alpha = alpha or self.alpha
-            elif isinstance(module, torch.nn.Conv2d):
+            elif isinstance(
+                module, (torch.nn.Conv1d, torch.nn.Conv2d, torch.nn.Conv3d)
+            ):
                 k_size, *_ = module.kernel_size
                 if k_size == 1 and lora_dim > 0:
                     dim = dim or lora_dim
@@ -384,10 +364,8 @@ def create_modules_(
             loras = {}
             lora_names = []
             for name, module in root_module.named_modules():
-                if module is root_module:
-                    continue
                 module_name = module.__class__.__name__
-                if module_name in self.MODULE_ALGO_MAP:
+                if module_name in self.MODULE_ALGO_MAP and module is not root_module:
                     next_config = self.MODULE_ALGO_MAP[module_name]
                     next_algo = next_config.get("algo", algo)
                     new_loras, new_lora_names = create_modules_(
@@ -398,7 +376,10 @@ def create_modules_(
                             loras[lora_name] = lora
                             lora_names.append(lora_name)
                     continue
-                lora_name = prefix + "." + name
+                if name:
+                    lora_name = prefix + "." + name
+                else:
+                    lora_name = prefix
                 lora_name = lora_name.replace(".", "_")
                 if lora_name in loras:
                     continue
@@ -614,245 +595,8 @@ def save_weights(self, file, dtype, metadata):
             # 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)
-
-
-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"]
-    TRAIN_INPUT = ["mlp.fc2", "ff.net.2"]
-    LORA_PREFIX_UNET = "lora_unet"
-    LORA_PREFIX_TEXT_ENCODER = "lora_te"
-
-    def __init__(
-        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,
-            target_replace_modules,
-            target_replace_names=[],
-            target_train_input=[],
-        ) -> List[IA3Module]:
-            logger.info("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 = IA3Module(
-                                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 = IA3Module(
-                            lora_name,
-                            module,
-                            self.multiplier,
-                            name in target_train_input,
-                            **kwargs,
-                        )
-                        loras.append(lora)
-            return loras
-
-        if isinstance(text_encoder, list):
-            text_encoders = text_encoder
-            use_index = True
-        else:
-            text_encoders = [text_encoder]
-            use_index = False
-
-        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,
-                )
-            )
-        logger.info(
-            f"create LyCORIS for Text Encoder: {len(self.text_encoder_loras)} modules."
-        )
-
-        self.unet_loras = create_modules(
-            IA3Network.LORA_PREFIX_UNET,
-            unet,
-            IA3Network.UNET_TARGET_REPLACE_MODULE,
-            IA3Network.UNET_TARGET_REPLACE_NAME,
-            IA3Network.TRAIN_INPUT,
-        )
-        logger.info(f"create LyCORIS for U-Net: {len(self.unet_loras)} modules.")
-
-        self.weights_sd = None
-
-        # assertion
-        names = set()
-        for lora in self.text_encoder_loras + self.unet_loras:
-            assert (
-                lora.lora_name not in names
-            ), f"duplicated lora name: {lora.lora_name}"
-            names.add(lora.lora_name)
-
-    def 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(LycorisNetworkKohya.LORA_PREFIX_TEXT_ENCODER):
-                    weights_has_text_encoder = True
-                elif key.startswith(LycorisNetworkKohya.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:
-            logger.info("enable LyCORIS for text encoder")
-        else:
-            self.text_encoder_loras = []
-
-        if apply_unet:
-            logger.info("enable LyCORIS 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)
-            logger.info(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, learning_rate):
-        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
-            )
+            model_hash = precalculate_safetensors_hashes(state_dict)
             metadata["sshs_model_hash"] = model_hash
-            metadata["sshs_legacy_hash"] = legacy_hash
 
             save_file(state_dict, file, metadata)
         else:
diff --git a/lycoris/kohya/model_utils.py b/lycoris/kohya/model_utils.py
deleted file mode 100644
index c8e2485..0000000
--- a/lycoris/kohya/model_utils.py
+++ /dev/null
@@ -1,1633 +0,0 @@
-# https://github.com/kohya-ss/sd-scripts/blob/b78c0e2a69e52ce6c79abc6c8c82d1a9cabcf05c/library/model_util.py
-# v1: split from train_db_fixed.py.
-# v2: support safetensors
-
-import math
-import os
-import torch
-import diffusers
-from transformers import CLIPTextModel, CLIPTokenizer, CLIPTextConfig, logging
-from diffusers import (
-    AutoencoderKL,
-    DDIMScheduler,
-    StableDiffusionPipeline,
-)  # , UNet2DConditionModel
-from safetensors.torch import load_file, save_file
-from .original_unet import UNet2DConditionModel
-
-# DiffUsers版StableDiffusionのモデルパラメータ
-NUM_TRAIN_TIMESTEPS = 1000
-BETA_START = 0.00085
-BETA_END = 0.0120
-
-UNET_PARAMS_MODEL_CHANNELS = 320
-UNET_PARAMS_CHANNEL_MULT = [1, 2, 4, 4]
-UNET_PARAMS_ATTENTION_RESOLUTIONS = [4, 2, 1]
-UNET_PARAMS_IMAGE_SIZE = 64  # fixed from old invalid value `32`
-UNET_PARAMS_IN_CHANNELS = 4
-UNET_PARAMS_OUT_CHANNELS = 4
-UNET_PARAMS_NUM_RES_BLOCKS = 2
-UNET_PARAMS_CONTEXT_DIM = 768
-UNET_PARAMS_NUM_HEADS = 8
-# UNET_PARAMS_USE_LINEAR_PROJECTION = False
-
-VAE_PARAMS_Z_CHANNELS = 4
-VAE_PARAMS_RESOLUTION = 256
-VAE_PARAMS_IN_CHANNELS = 3
-VAE_PARAMS_OUT_CH = 3
-VAE_PARAMS_CH = 128
-VAE_PARAMS_CH_MULT = [1, 2, 4, 4]
-VAE_PARAMS_NUM_RES_BLOCKS = 2
-
-# V2
-V2_UNET_PARAMS_ATTENTION_HEAD_DIM = [5, 10, 20, 20]
-V2_UNET_PARAMS_CONTEXT_DIM = 1024
-# V2_UNET_PARAMS_USE_LINEAR_PROJECTION = True
-
-# Diffusersの設定を読み込むための参照モデル
-DIFFUSERS_REF_MODEL_ID_V1 = "runwayml/stable-diffusion-v1-5"
-DIFFUSERS_REF_MODEL_ID_V2 = "stabilityai/stable-diffusion-2-1"
-
-
-# region StableDiffusion->Diffusersの変換コード
-# convert_original_stable_diffusion_to_diffusers をコピーして修正している（ASL 2.0）
-
-
-def shave_segments(path, n_shave_prefix_segments=1):
-    """
-    Removes segments. Positive values shave the first segments, negative shave the last segments.
-    """
-    if n_shave_prefix_segments >= 0:
-        return ".".join(path.split(".")[n_shave_prefix_segments:])
-    else:
-        return ".".join(path.split(".")[:n_shave_prefix_segments])
-
-
-def renew_resnet_paths(old_list, n_shave_prefix_segments=0):
-    """
-    Updates paths inside resnets to the new naming scheme (local renaming)
-    """
-    mapping = []
-    for old_item in old_list:
-        new_item = old_item.replace("in_layers.0", "norm1")
-        new_item = new_item.replace("in_layers.2", "conv1")
-
-        new_item = new_item.replace("out_layers.0", "norm2")
-        new_item = new_item.replace("out_layers.3", "conv2")
-
-        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
-        )
-
-        mapping.append({"old": old_item, "new": new_item})
-
-    return mapping
-
-
-def renew_vae_resnet_paths(old_list, n_shave_prefix_segments=0):
-    """
-    Updates paths inside resnets to the new naming scheme (local renaming)
-    """
-    mapping = []
-    for old_item in old_list:
-        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
-        )
-
-        mapping.append({"old": old_item, "new": new_item})
-
-    return mapping
-
-
-def renew_attention_paths(old_list, n_shave_prefix_segments=0):
-    """
-    Updates paths inside attentions to the new naming scheme (local renaming)
-    """
-    mapping = []
-    for old_item in old_list:
-        new_item = old_item
-
-        #         new_item = new_item.replace('norm.weight', 'group_norm.weight')
-        #         new_item = new_item.replace('norm.bias', 'group_norm.bias')
-
-        #         new_item = new_item.replace('proj_out.weight', 'proj_attn.weight')
-        #         new_item = new_item.replace('proj_out.bias', 'proj_attn.bias')
-
-        #         new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments)
-
-        mapping.append({"old": old_item, "new": new_item})
-
-    return mapping
-
-
-def renew_vae_attention_paths(old_list, n_shave_prefix_segments=0):
-    """
-    Updates paths inside attentions to the new naming scheme (local renaming)
-    """
-    mapping = []
-    for old_item in old_list:
-        new_item = old_item
-
-        new_item = new_item.replace("norm.weight", "group_norm.weight")
-        new_item = new_item.replace("norm.bias", "group_norm.bias")
-
-        if diffusers.__version__ < "0.17.0":
-            new_item = new_item.replace("q.weight", "query.weight")
-            new_item = new_item.replace("q.bias", "query.bias")
-
-            new_item = new_item.replace("k.weight", "key.weight")
-            new_item = new_item.replace("k.bias", "key.bias")
-
-            new_item = new_item.replace("v.weight", "value.weight")
-            new_item = new_item.replace("v.bias", "value.bias")
-
-            new_item = new_item.replace("proj_out.weight", "proj_attn.weight")
-            new_item = new_item.replace("proj_out.bias", "proj_attn.bias")
-        else:
-            new_item = new_item.replace("q.weight", "to_q.weight")
-            new_item = new_item.replace("q.bias", "to_q.bias")
-
-            new_item = new_item.replace("k.weight", "to_k.weight")
-            new_item = new_item.replace("k.bias", "to_k.bias")
-
-            new_item = new_item.replace("v.weight", "to_v.weight")
-            new_item = new_item.replace("v.bias", "to_v.bias")
-
-            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
-        )
-
-        mapping.append({"old": old_item, "new": new_item})
-
-    return mapping
-
-
-def assign_to_checkpoint(
-    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
-    to them. It splits attention layers, and takes into account additional replacements
-    that may arise.
-
-    Assigns the weights to the new checkpoint.
-    """
-    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:
-        for path, path_map in attention_paths_to_split.items():
-            old_tensor = old_checkpoint[path]
-            channels = old_tensor.shape[0] // 3
-
-            target_shape = (-1, channels) if len(old_tensor.shape) == 3 else (-1)
-
-            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:]
-            )
-            query, key, value = old_tensor.split(channels // num_heads, dim=1)
-
-            checkpoint[path_map["query"]] = query.reshape(target_shape)
-            checkpoint[path_map["key"]] = key.reshape(target_shape)
-            checkpoint[path_map["value"]] = value.reshape(target_shape)
-
-    for path in paths:
-        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
-        ):
-            continue
-
-        # Global renaming happens here
-        new_path = new_path.replace("middle_block.0", "mid_block.resnets.0")
-        new_path = new_path.replace("middle_block.1", "mid_block.attentions.0")
-        new_path = new_path.replace("middle_block.2", "mid_block.resnets.1")
-
-        if additional_replacements is not None:
-            for replacement in additional_replacements:
-                new_path = new_path.replace(replacement["old"], replacement["new"])
-
-        # proj_attn.weight has to be converted from conv 1D to linear
-        reshaping = False
-        if diffusers.__version__ < "0.17.0":
-            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
-            ):
-                reshaping = True
-
-        if reshaping:
-            checkpoint[new_path] = old_checkpoint[path["old"]][:, :, 0, 0]
-        else:
-            checkpoint[new_path] = old_checkpoint[path["old"]]
-
-
-def conv_attn_to_linear(checkpoint):
-    keys = list(checkpoint.keys())
-    attn_keys = ["query.weight", "key.weight", "value.weight"]
-    for key in keys:
-        if ".".join(key.split(".")[-2:]) in attn_keys:
-            if checkpoint[key].ndim > 2:
-                checkpoint[key] = checkpoint[key][:, :, 0, 0]
-        elif "proj_attn.weight" in key:
-            if checkpoint[key].ndim > 2:
-                checkpoint[key] = checkpoint[key][:, :, 0]
-
-
-def linear_transformer_to_conv(checkpoint):
-    keys = list(checkpoint.keys())
-    tf_keys = ["proj_in.weight", "proj_out.weight"]
-    for key in keys:
-        if ".".join(key.split(".")[-2:]) in tf_keys:
-            if checkpoint[key].ndim == 2:
-                checkpoint[key] = checkpoint[key].unsqueeze(2).unsqueeze(2)
-
-
-def convert_ldm_unet_checkpoint(v2, checkpoint, config):
-    """
-    Takes a state dict and a config, and returns a converted checkpoint.
-    """
-
-    # extract state_dict for UNet
-    unet_state_dict = {}
-    unet_key = "model.diffusion_model."
-    keys = list(checkpoint.keys())
-    for key in keys:
-        if key.startswith(unet_key):
-            unet_state_dict[key.replace(unet_key, "")] = checkpoint.pop(key)
-
-    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["conv_in.weight"] = unet_state_dict["input_blocks.0.0.weight"]
-    new_checkpoint["conv_in.bias"] = unet_state_dict["input_blocks.0.0.bias"]
-
-    new_checkpoint["conv_norm_out.weight"] = unet_state_dict["out.0.weight"]
-    new_checkpoint["conv_norm_out.bias"] = unet_state_dict["out.0.bias"]
-    new_checkpoint["conv_out.weight"] = unet_state_dict["out.2.weight"]
-    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
-        }
-    )
-    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
-        }
-    )
-    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
-        }
-    )
-    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
-        ]
-        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")
-            )
-
-        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,
-        )
-
-        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,
-            )
-
-    resnet_0 = middle_blocks[0]
-    attentions = middle_blocks[1]
-    resnet_1 = middle_blocks[2]
-
-    resnet_0_paths = renew_resnet_paths(resnet_0)
-    assign_to_checkpoint(resnet_0_paths, new_checkpoint, unet_state_dict, config=config)
-
-    resnet_1_paths = renew_resnet_paths(resnet_1)
-    assign_to_checkpoint(resnet_1_paths, new_checkpoint, unet_state_dict, config=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,
-    )
-
-    for i in range(num_output_blocks):
-        block_id = i // (config["layers_per_block"] + 1)
-        layer_in_block_id = i % (config["layers_per_block"] + 1)
-        output_block_layers = [shave_segments(name, 2) for name in output_blocks[i]]
-        output_block_list = {}
-
-        for layer in output_block_layers:
-            layer_id, layer_name = layer.split(".")[0], shave_segments(layer, 1)
-            if layer_id in output_block_list:
-                output_block_list[layer_id].append(layer_name)
-            else:
-                output_block_list[layer_id] = [layer_name]
-
-        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
-            ]
-
-            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,
-            )
-
-            # オリジナル：
-            # if ["conv.weight", "conv.bias"] in output_block_list.values():
-            #   index = list(output_block_list.values()).index(["conv.weight", "conv.bias"])
-
-            # biasとweightの順番に依存しないようにする：もっといいやり方がありそうだが
-            for l in output_block_list.values():
-                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"]
-                )
-
-                # Clear attentions as they have been attributed above.
-                if len(attentions) == 2:
-                    attentions = []
-
-            if len(attentions):
-                paths = renew_attention_paths(attentions)
-                meta_path = {
-                    "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,
-                )
-        else:
-            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_checkpoint[new_path] = unet_state_dict[old_path]
-
-    # SDのv2では1*1のconv2dがlinearに変わっている
-    # 誤って Diffusers 側を conv2d のままにしてしまったので、変換必要
-    if v2 and not config.get("use_linear_projection", False):
-        linear_transformer_to_conv(new_checkpoint)
-
-    return new_checkpoint
-
-
-def convert_ldm_vae_checkpoint(checkpoint, config):
-    # extract state dict for VAE
-    vae_state_dict = {}
-    vae_key = "first_stage_model."
-    keys = list(checkpoint.keys())
-    for key in keys:
-        if key.startswith(vae_key):
-            vae_state_dict[key.replace(vae_key, "")] = checkpoint.get(key)
-    # if len(vae_state_dict) == 0:
-    #   # 渡されたcheckpointは.ckptから読み込んだcheckpointではなくvaeのstate_dict
-    #   vae_state_dict = checkpoint
-
-    new_checkpoint = {}
-
-    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.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["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.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["quant_conv.weight"] = vae_state_dict["quant_conv.weight"]
-    new_checkpoint["quant_conv.bias"] = vae_state_dict["quant_conv.bias"]
-    new_checkpoint["post_quant_conv.weight"] = vae_state_dict["post_quant_conv.weight"]
-    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)
-    }
-
-    # 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)
-    }
-
-    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
-        ]
-
-        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")
-            )
-
-        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,
-        )
-
-    mid_resnets = [key for key in vae_state_dict if "encoder.mid.block" in key]
-    num_mid_res_blocks = 2
-    for i in range(1, num_mid_res_blocks + 1):
-        resnets = [key for key in mid_resnets if f"encoder.mid.block_{i}" in key]
-
-        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,
-        )
-
-    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,
-    )
-    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
-        ]
-
-        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"]
-            )
-
-        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,
-        )
-
-    mid_resnets = [key for key in vae_state_dict if "decoder.mid.block" in key]
-    num_mid_res_blocks = 2
-    for i in range(1, num_mid_res_blocks + 1):
-        resnets = [key for key in mid_resnets if f"decoder.mid.block_{i}" in key]
-
-        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,
-        )
-
-    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,
-    )
-    conv_attn_to_linear(new_checkpoint)
-    return new_checkpoint
-
-
-def create_unet_diffusers_config(v2, use_linear_projection_in_v2=False):
-    """
-    Creates a config for the diffusers based on the config of the LDM model.
-    """
-    # unet_params = original_config.model.params.unet_config.params
-
-    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"
-        )
-        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"
-        )
-        up_block_types.append(block_type)
-        resolution //= 2
-
-    config = dict(
-        sample_size=UNET_PARAMS_IMAGE_SIZE,
-        in_channels=UNET_PARAMS_IN_CHANNELS,
-        out_channels=UNET_PARAMS_OUT_CHANNELS,
-        down_block_types=tuple(down_block_types),
-        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
-        ),
-        # 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:
-        config["use_linear_projection"] = True
-
-    return config
-
-
-def create_vae_diffusers_config():
-    """
-    Creates a config for the diffusers based on the config of the LDM model.
-    """
-    # vae_params = original_config.model.params.first_stage_config.params.ddconfig
-    # _ = original_config.model.params.first_stage_config.params.embed_dim
-    block_out_channels = [VAE_PARAMS_CH * mult for mult in VAE_PARAMS_CH_MULT]
-    down_block_types = ["DownEncoderBlock2D"] * len(block_out_channels)
-    up_block_types = ["UpDecoderBlock2D"] * len(block_out_channels)
-
-    config = dict(
-        sample_size=VAE_PARAMS_RESOLUTION,
-        in_channels=VAE_PARAMS_IN_CHANNELS,
-        out_channels=VAE_PARAMS_OUT_CH,
-        down_block_types=tuple(down_block_types),
-        up_block_types=tuple(up_block_types),
-        block_out_channels=tuple(block_out_channels),
-        latent_channels=VAE_PARAMS_Z_CHANNELS,
-        layers_per_block=VAE_PARAMS_NUM_RES_BLOCKS,
-    )
-    return config
-
-
-def convert_ldm_clip_checkpoint_v1(checkpoint):
-    keys = list(checkpoint.keys())
-    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
-            ]
-    return text_model_dict
-
-
-def convert_ldm_clip_checkpoint_v2(checkpoint, max_length):
-    # 嫌になるくらい違うぞ！
-    def convert_key(key):
-        if not key.startswith("cond_stage_model"):
-            return None
-
-        # common conversion
-        key = key.replace("cond_stage_model.model.transformer.", "text_model.encoder.")
-        key = key.replace("cond_stage_model.model.", "text_model.")
-
-        if "resblocks" in key:
-            # resblocks conversion
-            key = key.replace(".resblocks.", ".layers.")
-            if ".ln_" in key:
-                key = key.replace(".ln_", ".layer_norm")
-            elif ".mlp." in key:
-                key = key.replace(".c_fc.", ".fc1.")
-                key = key.replace(".c_proj.", ".fc2.")
-            elif ".attn.out_proj" in key:
-                key = key.replace(".attn.out_proj.", ".self_attn.out_proj.")
-            elif ".attn.in_proj" in key:
-                key = None  # 特殊なので後で処理する
-            else:
-                raise ValueError(f"unexpected key in SD: {key}")
-        elif ".positional_embedding" in key:
-            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"
-            )
-        elif ".ln_final" in key:
-            key = key.replace(".ln_final", ".final_layer_norm")
-        return key
-
-    keys = list(checkpoint.keys())
-    new_sd = {}
-    for key in keys:
-        # remove resblocks 23
-        if ".resblocks.23." in key:
-            continue
-        new_key = convert_key(key)
-        if new_key is None:
-            continue
-        new_sd[new_key] = checkpoint[key]
-
-    # attnの変換
-    for key in keys:
-        if ".resblocks.23." in key:
-            continue
-        if ".resblocks" in key and ".attn.in_proj_" in 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_pfx.replace("_weight", "")
-            key_pfx = key_pfx.replace("_bias", "")
-            key_pfx = key_pfx.replace(".attn.in_proj", ".self_attn.")
-            new_sd[key_pfx + "q_proj" + key_suffix] = values[0]
-            new_sd[key_pfx + "k_proj" + key_suffix] = values[1]
-            new_sd[key_pfx + "v_proj" + key_suffix] = values[2]
-
-    # rename or add position_ids
-    ANOTHER_POSITION_IDS_KEY = "text_model.encoder.text_model.embeddings.position_ids"
-    if ANOTHER_POSITION_IDS_KEY in new_sd:
-        # waifu diffusion v1.4
-        position_ids = new_sd[ANOTHER_POSITION_IDS_KEY]
-        del new_sd[ANOTHER_POSITION_IDS_KEY]
-    else:
-        position_ids = torch.Tensor([list(range(max_length))]).to(torch.int64)
-
-    new_sd["text_model.embeddings.position_ids"] = position_ids
-    return new_sd
-
-
-# endregion
-
-
-# region Diffusers->StableDiffusion の変換コード
-# convert_diffusers_to_original_stable_diffusion をコピーして修正している（ASL 2.0）
-
-
-def conv_transformer_to_linear(checkpoint):
-    keys = list(checkpoint.keys())
-    tf_keys = ["proj_in.weight", "proj_out.weight"]
-    for key in keys:
-        if ".".join(key.split(".")[-2:]) in tf_keys:
-            if checkpoint[key].ndim > 2:
-                checkpoint[key] = checkpoint[key][:, :, 0, 0]
-
-
-def convert_unet_state_dict_to_sd(v2, unet_state_dict):
-    unet_conversion_map = [
-        # (stable-diffusion, HF Diffusers)
-        ("time_embed.0.weight", "time_embedding.linear_1.weight"),
-        ("time_embed.0.bias", "time_embedding.linear_1.bias"),
-        ("time_embed.2.weight", "time_embedding.linear_2.weight"),
-        ("time_embed.2.bias", "time_embedding.linear_2.bias"),
-        ("input_blocks.0.0.weight", "conv_in.weight"),
-        ("input_blocks.0.0.bias", "conv_in.bias"),
-        ("out.0.weight", "conv_norm_out.weight"),
-        ("out.0.bias", "conv_norm_out.bias"),
-        ("out.2.weight", "conv_out.weight"),
-        ("out.2.bias", "conv_out.bias"),
-    ]
-
-    unet_conversion_map_resnet = [
-        # (stable-diffusion, HF Diffusers)
-        ("in_layers.0", "norm1"),
-        ("in_layers.2", "conv1"),
-        ("out_layers.0", "norm2"),
-        ("out_layers.3", "conv2"),
-        ("emb_layers.1", "time_emb_proj"),
-        ("skip_connection", "conv_shortcut"),
-    ]
-
-    unet_conversion_map_layer = []
-    for i in range(4):
-        # loop over downblocks/upblocks
-
-        for j in range(2):
-            # loop over resnets/attentions for downblocks
-            hf_down_res_prefix = f"down_blocks.{i}.resnets.{j}."
-            sd_down_res_prefix = f"input_blocks.{3*i + j + 1}.0."
-            unet_conversion_map_layer.append((sd_down_res_prefix, hf_down_res_prefix))
-
-            if i < 3:
-                # 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)
-                )
-
-        for j in range(3):
-            # loop over resnets/attentions for upblocks
-            hf_up_res_prefix = f"up_blocks.{i}.resnets.{j}."
-            sd_up_res_prefix = f"output_blocks.{3*i + j}.0."
-            unet_conversion_map_layer.append((sd_up_res_prefix, hf_up_res_prefix))
-
-            if i > 0:
-                # no attention layers in up_blocks.0
-                hf_up_atn_prefix = f"up_blocks.{i}.attentions.{j}."
-                sd_up_atn_prefix = f"output_blocks.{3*i + j}.1."
-                unet_conversion_map_layer.append((sd_up_atn_prefix, hf_up_atn_prefix))
-
-        if i < 3:
-            # 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)
-            )
-
-            # no upsample in up_blocks.3
-            hf_upsample_prefix = f"up_blocks.{i}.upsamplers.0."
-            sd_upsample_prefix = f"output_blocks.{3*i + 2}.{1 if i == 0 else 2}."
-            unet_conversion_map_layer.append((sd_upsample_prefix, hf_upsample_prefix))
-
-    hf_mid_atn_prefix = "mid_block.attentions.0."
-    sd_mid_atn_prefix = "middle_block.1."
-    unet_conversion_map_layer.append((sd_mid_atn_prefix, hf_mid_atn_prefix))
-
-    for j in range(2):
-        hf_mid_res_prefix = f"mid_block.resnets.{j}."
-        sd_mid_res_prefix = f"middle_block.{2*j}."
-        unet_conversion_map_layer.append((sd_mid_res_prefix, hf_mid_res_prefix))
-
-    # buyer beware: this is a *brittle* function,
-    # and correct output requires that all of these pieces interact in
-    # the exact order in which I have arranged them.
-    mapping = {k: k for k in unet_state_dict.keys()}
-    for sd_name, hf_name in unet_conversion_map:
-        mapping[hf_name] = sd_name
-    for k, v in mapping.items():
-        if "resnets" in k:
-            for sd_part, hf_part in unet_conversion_map_resnet:
-                v = v.replace(hf_part, sd_part)
-            mapping[k] = v
-    for k, v in mapping.items():
-        for sd_part, hf_part in unet_conversion_map_layer:
-            v = v.replace(hf_part, sd_part)
-        mapping[k] = v
-    new_state_dict = {v: unet_state_dict[k] for k, v in mapping.items()}
-
-    if v2:
-        conv_transformer_to_linear(new_state_dict)
-
-    return new_state_dict
-
-
-def controlnet_conversion_map():
-    unet_conversion_map = [
-        ("time_embed.0.weight", "time_embedding.linear_1.weight"),
-        ("time_embed.0.bias", "time_embedding.linear_1.bias"),
-        ("time_embed.2.weight", "time_embedding.linear_2.weight"),
-        ("time_embed.2.bias", "time_embedding.linear_2.bias"),
-        ("input_blocks.0.0.weight", "conv_in.weight"),
-        ("input_blocks.0.0.bias", "conv_in.bias"),
-        ("middle_block_out.0.weight", "controlnet_mid_block.weight"),
-        ("middle_block_out.0.bias", "controlnet_mid_block.bias"),
-    ]
-
-    unet_conversion_map_resnet = [
-        ("in_layers.0", "norm1"),
-        ("in_layers.2", "conv1"),
-        ("out_layers.0", "norm2"),
-        ("out_layers.3", "conv2"),
-        ("emb_layers.1", "time_emb_proj"),
-        ("skip_connection", "conv_shortcut"),
-    ]
-
-    unet_conversion_map_layer = []
-    for i in range(4):
-        for j in range(2):
-            hf_down_res_prefix = f"down_blocks.{i}.resnets.{j}."
-            sd_down_res_prefix = f"input_blocks.{3*i + j + 1}.0."
-            unet_conversion_map_layer.append((sd_down_res_prefix, hf_down_res_prefix))
-
-            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)
-                )
-
-        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)
-            )
-
-    hf_mid_atn_prefix = "mid_block.attentions.0."
-    sd_mid_atn_prefix = "middle_block.1."
-    unet_conversion_map_layer.append((sd_mid_atn_prefix, hf_mid_atn_prefix))
-
-    for j in range(2):
-        hf_mid_res_prefix = f"mid_block.resnets.{j}."
-        sd_mid_res_prefix = f"middle_block.{2*j}."
-        unet_conversion_map_layer.append((sd_mid_res_prefix, hf_mid_res_prefix))
-
-    controlnet_cond_embedding_names = (
-        ["conv_in"] + [f"blocks.{i}" for i in range(6)] + ["conv_out"]
-    )
-    for i, hf_prefix in enumerate(controlnet_cond_embedding_names):
-        hf_prefix = f"controlnet_cond_embedding.{hf_prefix}."
-        sd_prefix = f"input_hint_block.{i*2}."
-        unet_conversion_map_layer.append((sd_prefix, hf_prefix))
-
-    for i in range(12):
-        hf_prefix = f"controlnet_down_blocks.{i}."
-        sd_prefix = f"zero_convs.{i}.0."
-        unet_conversion_map_layer.append((sd_prefix, hf_prefix))
-
-    return unet_conversion_map, unet_conversion_map_resnet, unet_conversion_map_layer
-
-
-def convert_controlnet_state_dict_to_sd(controlnet_state_dict):
-    (
-        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:
-        mapping[diffusers_name] = sd_name
-    for k, v in mapping.items():
-        if "resnets" in k:
-            for sd_part, diffusers_part in unet_conversion_map_resnet:
-                v = v.replace(diffusers_part, sd_part)
-            mapping[k] = v
-    for k, v in mapping.items():
-        for sd_part, diffusers_part in unet_conversion_map_layer:
-            v = v.replace(diffusers_part, sd_part)
-        mapping[k] = v
-    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()
-
-    mapping = {k: k for k in controlnet_state_dict.keys()}
-    for sd_name, diffusers_name in unet_conversion_map:
-        mapping[sd_name] = diffusers_name
-    for k, v in mapping.items():
-        for sd_part, diffusers_part in unet_conversion_map_layer:
-            v = v.replace(sd_part, diffusers_part)
-        mapping[k] = v
-    for k, v in mapping.items():
-        if "resnets" in v:
-            for sd_part, diffusers_part in unet_conversion_map_resnet:
-                v = v.replace(sd_part, diffusers_part)
-            mapping[k] = v
-    new_state_dict = {v: controlnet_state_dict[k] for k, v in mapping.items()}
-    return new_state_dict
-
-
-# ================#
-# VAE Conversion #
-# ================#
-
-
-def reshape_weight_for_sd(w):
-    # convert HF linear weights to SD conv2d weights
-    return w.reshape(*w.shape, 1, 1)
-
-
-def convert_vae_state_dict(vae_state_dict):
-    vae_conversion_map = [
-        # (stable-diffusion, HF Diffusers)
-        ("nin_shortcut", "conv_shortcut"),
-        ("norm_out", "conv_norm_out"),
-        ("mid.attn_1.", "mid_block.attentions.0."),
-    ]
-
-    for i in range(4):
-        # down_blocks have two resnets
-        for j in range(2):
-            hf_down_prefix = f"encoder.down_blocks.{i}.resnets.{j}."
-            sd_down_prefix = f"encoder.down.{i}.block.{j}."
-            vae_conversion_map.append((sd_down_prefix, hf_down_prefix))
-
-        if i < 3:
-            hf_downsample_prefix = f"down_blocks.{i}.downsamplers.0."
-            sd_downsample_prefix = f"down.{i}.downsample."
-            vae_conversion_map.append((sd_downsample_prefix, hf_downsample_prefix))
-
-            hf_upsample_prefix = f"up_blocks.{i}.upsamplers.0."
-            sd_upsample_prefix = f"up.{3-i}.upsample."
-            vae_conversion_map.append((sd_upsample_prefix, hf_upsample_prefix))
-
-        # up_blocks have three resnets
-        # also, up blocks in hf are numbered in reverse from sd
-        for j in range(3):
-            hf_up_prefix = f"decoder.up_blocks.{i}.resnets.{j}."
-            sd_up_prefix = f"decoder.up.{3-i}.block.{j}."
-            vae_conversion_map.append((sd_up_prefix, hf_up_prefix))
-
-    # this part accounts for mid blocks in both the encoder and the decoder
-    for i in range(2):
-        hf_mid_res_prefix = f"mid_block.resnets.{i}."
-        sd_mid_res_prefix = f"mid.block_{i+1}."
-        vae_conversion_map.append((sd_mid_res_prefix, hf_mid_res_prefix))
-
-    if diffusers.__version__ < "0.17.0":
-        vae_conversion_map_attn = [
-            # (stable-diffusion, HF Diffusers)
-            ("norm.", "group_norm."),
-            ("q.", "query."),
-            ("k.", "key."),
-            ("v.", "value."),
-            ("proj_out.", "proj_attn."),
-        ]
-    else:
-        vae_conversion_map_attn = [
-            # (stable-diffusion, HF Diffusers)
-            ("norm.", "group_norm."),
-            ("q.", "to_q."),
-            ("k.", "to_k."),
-            ("v.", "to_v."),
-            ("proj_out.", "to_out.0."),
-        ]
-
-    mapping = {k: k for k in vae_state_dict.keys()}
-    for k, v in mapping.items():
-        for sd_part, hf_part in vae_conversion_map:
-            v = v.replace(hf_part, sd_part)
-        mapping[k] = v
-    for k, v in mapping.items():
-        if "attentions" in k:
-            for sd_part, hf_part in vae_conversion_map_attn:
-                v = v.replace(hf_part, sd_part)
-            mapping[k] = v
-    new_state_dict = {v: vae_state_dict[k] for k, v in mapping.items()}
-    weights_to_convert = ["q", "k", "v", "proj_out"]
-    for k, v in new_state_dict.items():
-        for weight_name in weights_to_convert:
-            if f"mid.attn_1.{weight_name}.weight" in k:
-                # print(f"Reshaping {k} for SD format: shape {v.shape} -> {v.shape} x 1 x 1")
-                new_state_dict[k] = reshape_weight_for_sd(v)
-
-    return new_state_dict
-
-
-# endregion
-
-# region 自作のモデル読み書きなど
-
-
-def is_safetensors(path):
-    return os.path.splitext(path)[1].lower() == ".safetensors"
-
-
-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.",
-        ),
-    ]
-
-    if is_safetensors(ckpt_path):
-        checkpoint = None
-        state_dict = load_file(ckpt_path)  # , device) # may causes error
-    else:
-        checkpoint = torch.load(ckpt_path, map_location=device)
-        if "state_dict" in checkpoint:
-            state_dict = checkpoint["state_dict"]
-        else:
-            state_dict = checkpoint
-            checkpoint = None
-
-    key_reps = []
-    for rep_from, rep_to in TEXT_ENCODER_KEY_REPLACEMENTS:
-        for key in state_dict.keys():
-            if key.startswith(rep_from):
-                new_key = rep_to + key[len(rep_from) :]
-                key_reps.append((key, new_key))
-
-    for key, new_key in key_reps:
-        state_dict[new_key] = state_dict[key]
-        del state_dict[key]
-
-    return checkpoint, state_dict
-
-
-# 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
-):
-    _, state_dict = load_checkpoint_with_text_encoder_conversion(ckpt_path, device)
-
-    # Convert the UNet2DConditionModel model.
-    unet_config = create_unet_diffusers_config(v2, unet_use_linear_projection_in_v2)
-    converted_unet_checkpoint = convert_ldm_unet_checkpoint(v2, state_dict, unet_config)
-
-    unet = UNet2DConditionModel(**unet_config).to(device)
-    info = unet.load_state_dict(converted_unet_checkpoint)
-    print("loading u-net:", info)
-
-    # Convert the VAE model.
-    vae_config = create_vae_diffusers_config()
-    converted_vae_checkpoint = convert_ldm_vae_checkpoint(state_dict, vae_config)
-
-    vae = AutoencoderKL(**vae_config).to(device)
-    info = vae.load_state_dict(converted_vae_checkpoint)
-    print("loading vae:", info)
-
-    # convert text_model
-    if v2:
-        converted_text_encoder_checkpoint = convert_ldm_clip_checkpoint_v2(
-            state_dict, 77
-        )
-        cfg = CLIPTextConfig(
-            vocab_size=49408,
-            hidden_size=1024,
-            intermediate_size=4096,
-            num_hidden_layers=23,
-            num_attention_heads=16,
-            max_position_embeddings=77,
-            hidden_act="gelu",
-            layer_norm_eps=1e-05,
-            dropout=0.0,
-            attention_dropout=0.0,
-            initializer_range=0.02,
-            initializer_factor=1.0,
-            pad_token_id=1,
-            bos_token_id=0,
-            eos_token_id=2,
-            model_type="clip_text_model",
-            projection_dim=512,
-            torch_dtype="float32",
-            transformers_version="4.25.0.dev0",
-        )
-        text_model = CLIPTextModel._from_config(cfg)
-        info = text_model.load_state_dict(
-            converted_text_encoder_checkpoint, strict=False
-        )
-    else:
-        converted_text_encoder_checkpoint = convert_ldm_clip_checkpoint_v1(state_dict)
-
-        # logging.set_verbosity_error()  # don't show annoying warning
-        # text_model = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14").to(device)
-        # logging.set_verbosity_warning()
-        # print(f"config: {text_model.config}")
-        cfg = CLIPTextConfig(
-            vocab_size=49408,
-            hidden_size=768,
-            intermediate_size=3072,
-            num_hidden_layers=12,
-            num_attention_heads=12,
-            max_position_embeddings=77,
-            hidden_act="quick_gelu",
-            layer_norm_eps=1e-05,
-            dropout=0.0,
-            attention_dropout=0.0,
-            initializer_range=0.02,
-            initializer_factor=1.0,
-            pad_token_id=1,
-            bos_token_id=0,
-            eos_token_id=2,
-            model_type="clip_text_model",
-            projection_dim=768,
-            torch_dtype="float32",
-        )
-        text_model = CLIPTextModel._from_config(cfg)
-        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"
-    if v2:
-        version_str += "_v2"
-    else:
-        version_str += "_v1"
-    if 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の除去
-        if ".position_ids" in key:
-            return None
-
-        # common
-        key = key.replace("text_model.encoder.", "transformer.")
-        key = key.replace("text_model.", "")
-        if "layers" in key:
-            # resblocks conversion
-            key = key.replace(".layers.", ".resblocks.")
-            if ".layer_norm" in key:
-                key = key.replace(".layer_norm", ".ln_")
-            elif ".mlp." in key:
-                key = key.replace(".fc1.", ".c_fc.")
-                key = key.replace(".fc2.", ".c_proj.")
-            elif ".self_attn.out_proj" in key:
-                key = key.replace(".self_attn.out_proj.", ".attn.out_proj.")
-            elif ".self_attn." in key:
-                key = None  # 特殊なので後で処理する
-            else:
-                raise ValueError(f"unexpected key in DiffUsers model: {key}")
-        elif ".position_embedding" in key:
-            key = key.replace(
-                "embeddings.position_embedding.weight", "positional_embedding"
-            )
-        elif ".token_embedding" in key:
-            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
-
-    keys = list(checkpoint.keys())
-    new_sd = {}
-    for key in keys:
-        new_key = convert_key(key)
-        if new_key is None:
-            continue
-        new_sd[new_key] = checkpoint[key]
-
-    # attnの変換
-    for key in keys:
-        if "layers" in key and "q_proj" in key:
-            # 三つを結合
-            key_q = key
-            key_k = key.replace("q_proj", "k_proj")
-            key_v = key.replace("q_proj", "v_proj")
-
-            value_q = checkpoint[key_q]
-            value_k = checkpoint[key_k]
-            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 = new_key.replace(".self_attn.q_proj.", ".attn.in_proj_")
-            new_sd[new_key] = value
-
-    # 最後の層などを捏造するか
-    if make_dummy_weights:
-        print("make dummy weights for resblock.23, text_projection and logit scale.")
-        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の保存で落ちる
-
-        # Diffusersに含まれない重みを作っておく
-        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,
-):
-    if ckpt_path is not None:
-        # epoch/stepを参照する。またVAEがメモリ上にないときなど、もう一度VAEを含めて読み込む
-        checkpoint, state_dict = load_checkpoint_with_text_encoder_conversion(ckpt_path)
-        if checkpoint is None:  # safetensors または state_dictのckpt
-            checkpoint = {}
-            strict = False
-        else:
-            strict = True
-        if "state_dict" in state_dict:
-            del state_dict["state_dict"]
-    else:
-        # 新しく作る
-        assert (
-            vae is not None
-        ), "VAE is required to save a checkpoint without a given checkpoint"
-        checkpoint = {}
-        state_dict = {}
-        strict = False
-
-    def update_sd(prefix, sd):
-        for k, v in sd.items():
-            key = prefix + k
-            assert not strict or key in state_dict, f"Illegal key in save SD: {key}"
-            if save_dtype is not None:
-                v = v.detach().clone().to("cpu").to(save_dtype)
-            state_dict[key] = v
-
-    # Convert the UNet model
-    unet_state_dict = convert_unet_state_dict_to_sd(v2, unet.state_dict())
-    update_sd("model.diffusion_model.", unet_state_dict)
-
-    # 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
-        )
-        update_sd("cond_stage_model.model.", text_enc_dict)
-    else:
-        text_enc_dict = text_encoder.state_dict()
-        update_sd("cond_stage_model.transformer.", text_enc_dict)
-
-    # Convert the VAE
-    if vae is not None:
-        vae_dict = convert_vae_state_dict(vae.state_dict())
-        update_sd("first_stage_model.", vae_dict)
-
-    # Put together new checkpoint
-    key_count = len(state_dict.keys())
-    new_ckpt = {"state_dict": state_dict}
-
-    # epoch and global_step are sometimes not int
-    try:
-        if "epoch" in checkpoint:
-            epochs += checkpoint["epoch"]
-        if "global_step" in checkpoint:
-            steps += checkpoint["global_step"]
-    except:
-        pass
-
-    new_ckpt["epoch"] = epochs
-    new_ckpt["global_step"] = steps
-
-    if is_safetensors(output_file):
-        # TODO Tensor以外のdictの値を削除したほうがいいか
-        save_file(state_dict, output_file)
-    else:
-        torch.save(new_ckpt, output_file)
-
-    return key_count
-
-
-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:
-            pretrained_model_name_or_path = DIFFUSERS_REF_MODEL_ID_V2
-        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"
-    )
-    if vae is None:
-        vae = AutoencoderKL.from_pretrained(
-            pretrained_model_name_or_path, subfolder="vae"
-        )
-
-    pipeline = StableDiffusionPipeline(
-        unet=unet,
-        text_encoder=text_encoder,
-        vae=vae,
-        scheduler=scheduler,
-        tokenizer=tokenizer,
-        safety_checker=None,
-        feature_extractor=None,
-        requires_safety_checker=None,
-    )
-    pipeline.save_pretrained(output_dir, safe_serialization=use_safetensors)
-
-
-VAE_PREFIX = "first_stage_model."
-
-
-def load_vae(vae_id, dtype):
-    print(f"load VAE: {vae_id}")
-    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
-            )
-        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
-            )
-        return vae
-
-    # local
-    vae_config = create_vae_diffusers_config()
-
-    if vae_id.endswith(".bin"):
-        # SD 1.5 VAE on Huggingface
-        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_sd = vae_model["state_dict"] if "state_dict" in vae_model else vae_model
-
-        # vae only or full model
-        full_model = False
-        for vae_key in vae_sd:
-            if vae_key.startswith(VAE_PREFIX):
-                full_model = True
-                break
-        if not full_model:
-            sd = {}
-            for key, value in vae_sd.items():
-                sd[VAE_PREFIX + key] = value
-            vae_sd = sd
-            del sd
-
-        # Convert the VAE model.
-        converted_vae_checkpoint = convert_ldm_vae_checkpoint(vae_sd, vae_config)
-
-    vae = AutoencoderKL(**vae_config)
-    vae.load_state_dict(converted_vae_checkpoint)
-    return vae
-
-
-# endregion
-
-
-def make_bucket_resolutions(max_reso, min_size=256, max_size=1024, divisible=64):
-    max_width, max_height = max_reso
-    max_area = (max_width // divisible) * (max_height // divisible)
-
-    resos = set()
-
-    size = int(math.sqrt(max_area)) * divisible
-    resos.add((size, size))
-
-    size = min_size
-    while size <= max_size:
-        width = size
-        height = min(max_size, (max_area // (width // divisible)) * divisible)
-        resos.add((width, height))
-        resos.add((height, width))
-
-        # # make additional resos
-        # if width >= height and width - divisible >= min_size:
-        #   resos.add((width - divisible, height))
-        #   resos.add((height, width - divisible))
-        # if height >= width and height - divisible >= min_size:
-        #   resos.add((width, height - divisible))
-        #   resos.add((height - divisible, width))
-
-        size += divisible
-
-    resos = list(resos)
-    resos.sort()
-    return resos
-
-
-if __name__ == "__main__":
-    resos = make_bucket_resolutions((512, 768))
-    print(len(resos))
-    print(resos)
-    aspect_ratios = [w / h for w, h in resos]
-    print(aspect_ratios)
-
-    ars = set()
-    for ar in aspect_ratios:
-        if ar in ars:
-            print("error! duplicate ar:", ar)
-        ars.add(ar)
diff --git a/lycoris/kohya/original_unet.py b/lycoris/kohya/original_unet.py
deleted file mode 100644
index 2d1ff36..0000000
--- a/lycoris/kohya/original_unet.py
+++ /dev/null
@@ -1,1760 +0,0 @@
-# https://github.com/kohya-ss/sd-scripts/blob/b78c0e2a69e52ce6c79abc6c8c82d1a9cabcf05c/library/original_unet.py
-# Diffusers 0.10.2からStable Diffusionに必要な部分だけを持ってくる
-# 条件分岐等で不要な部分は削除している
-# コードの多くはDiffusersからコピーしている
-# 制約として、モデルのstate_dictがDiffusers 0.10.2のものと同じ形式である必要がある
-
-# Copy from Diffusers 0.10.2 for Stable Diffusion. Most of the code is copied from Diffusers.
-# Unnecessary parts are deleted by condition branching.
-# As a constraint, the state_dict of the model must be in the same format as that of Diffusers 0.10.2
-
-"""
-v1.5とv2.1の相違点は
-- attention_head_dimがintかlist[int]か
-- cross_attention_dimが768か1024か
-- use_linear_projection: trueがない（=False, 1.5）かあるか
-- upcast_attentionがFalse(1.5)かTrue(2.1)か
-- （以下は多分無視していい）
-- sample_sizeが64か96か
-- dual_cross_attentionがあるかないか
-- num_class_embedsがあるかないか
-- only_cross_attentionがあるかないか
-
-v1.5
-{
-  "_class_name": "UNet2DConditionModel",
-  "_diffusers_version": "0.6.0",
-  "act_fn": "silu",
-  "attention_head_dim": 8,
-  "block_out_channels": [
-    320,
-    640,
-    1280,
-    1280
-  ],
-  "center_input_sample": false,
-  "cross_attention_dim": 768,
-  "down_block_types": [
-    "CrossAttnDownBlock2D",
-    "CrossAttnDownBlock2D",
-    "CrossAttnDownBlock2D",
-    "DownBlock2D"
-  ],
-  "downsample_padding": 1,
-  "flip_sin_to_cos": true,
-  "freq_shift": 0,
-  "in_channels": 4,
-  "layers_per_block": 2,
-  "mid_block_scale_factor": 1,
-  "norm_eps": 1e-05,
-  "norm_num_groups": 32,
-  "out_channels": 4,
-  "sample_size": 64,
-  "up_block_types": [
-    "UpBlock2D",
-    "CrossAttnUpBlock2D",
-    "CrossAttnUpBlock2D",
-    "CrossAttnUpBlock2D"
-  ]
-}
-
-v2.1
-{
-  "_class_name": "UNet2DConditionModel",
-  "_diffusers_version": "0.10.0.dev0",
-  "act_fn": "silu",
-  "attention_head_dim": [
-    5,
-    10,
-    20,
-    20
-  ],
-  "block_out_channels": [
-    320,
-    640,
-    1280,
-    1280
-  ],
-  "center_input_sample": false,
-  "cross_attention_dim": 1024,
-  "down_block_types": [
-    "CrossAttnDownBlock2D",
-    "CrossAttnDownBlock2D",
-    "CrossAttnDownBlock2D",
-    "DownBlock2D"
-  ],
-  "downsample_padding": 1,
-  "dual_cross_attention": false,
-  "flip_sin_to_cos": true,
-  "freq_shift": 0,
-  "in_channels": 4,
-  "layers_per_block": 2,
-  "mid_block_scale_factor": 1,
-  "norm_eps": 1e-05,
-  "norm_num_groups": 32,
-  "num_class_embeds": null,
-  "only_cross_attention": false,
-  "out_channels": 4,
-  "sample_size": 96,
-  "up_block_types": [
-    "UpBlock2D",
-    "CrossAttnUpBlock2D",
-    "CrossAttnUpBlock2D",
-    "CrossAttnUpBlock2D"
-  ],
-  "use_linear_projection": true,
-  "upcast_attention": true
-}
-"""
-
-import math
-from types import SimpleNamespace
-from typing import Dict, Optional, Tuple, Union
-import torch
-from torch import nn
-from torch.nn import functional as F
-from einops import rearrange
-
-BLOCK_OUT_CHANNELS: Tuple[int] = (320, 640, 1280, 1280)
-TIMESTEP_INPUT_DIM = BLOCK_OUT_CHANNELS[0]
-TIME_EMBED_DIM = BLOCK_OUT_CHANNELS[0] * 4
-IN_CHANNELS: int = 4
-OUT_CHANNELS: int = 4
-LAYERS_PER_BLOCK: int = 2
-LAYERS_PER_BLOCK_UP: int = LAYERS_PER_BLOCK + 1
-TIME_EMBED_FLIP_SIN_TO_COS: bool = True
-TIME_EMBED_FREQ_SHIFT: int = 0
-NORM_GROUPS: int = 32
-NORM_EPS: float = 1e-5
-TRANSFORMER_NORM_NUM_GROUPS = 32
-
-DOWN_BLOCK_TYPES = [
-    "CrossAttnDownBlock2D",
-    "CrossAttnDownBlock2D",
-    "CrossAttnDownBlock2D",
-    "DownBlock2D",
-]
-UP_BLOCK_TYPES = [
-    "UpBlock2D",
-    "CrossAttnUpBlock2D",
-    "CrossAttnUpBlock2D",
-    "CrossAttnUpBlock2D",
-]
-
-
-# region memory effcient attention
-
-# FlashAttentionを使うCrossAttention
-# based on https://github.com/lucidrains/memory-efficient-attention-pytorch/blob/main/memory_efficient_attention_pytorch/flash_attention.py
-# LICENSE MIT https://github.com/lucidrains/memory-efficient-attention-pytorch/blob/main/LICENSE
-
-# constants
-
-EPSILON = 1e-6
-
-# helper functions
-
-
-def exists(val):
-    return val is not None
-
-
-def default(val, d):
-    return val if exists(val) else d
-
-
-# flash attention forwards and backwards
-
-# https://arxiv.org/abs/2205.14135
-
-
-class FlashAttentionFunction(torch.autograd.Function):
-    @staticmethod
-    @torch.no_grad()
-    def forward(ctx, q, k, v, mask, causal, q_bucket_size, k_bucket_size):
-        """Algorithm 2 in the paper"""
-
-        device = q.device
-        dtype = q.dtype
-        max_neg_value = -torch.finfo(q.dtype).max
-        qk_len_diff = max(k.shape[-2] - q.shape[-2], 0)
-
-        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
-        )
-
-        scale = q.shape[-1] ** -0.5
-
-        if not exists(mask):
-            mask = (None,) * math.ceil(q.shape[-2] / q_bucket_size)
-        else:
-            mask = rearrange(mask, "b n -> b 1 1 n")
-            mask = mask.split(q_bucket_size, dim=-1)
-
-        row_splits = zip(
-            q.split(q_bucket_size, dim=-2),
-            o.split(q_bucket_size, dim=-2),
-            mask,
-            all_row_sums.split(q_bucket_size, dim=-2),
-            all_row_maxes.split(q_bucket_size, dim=-2),
-        )
-
-        for ind, (qc, oc, row_mask, row_sums, row_maxes) in enumerate(row_splits):
-            q_start_index = ind * q_bucket_size - qk_len_diff
-
-            col_splits = zip(
-                k.split(k_bucket_size, dim=-2),
-                v.split(k_bucket_size, dim=-2),
-            )
-
-            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
-                )
-
-                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)
-                    attn_weights.masked_fill_(causal_mask, max_neg_value)
-
-                block_row_maxes = attn_weights.amax(dim=-1, keepdims=True)
-                attn_weights -= block_row_maxes
-                exp_weights = torch.exp(attn_weights)
-
-                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
-                )
-
-                new_row_maxes = torch.maximum(block_row_maxes, row_maxes)
-
-                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
-                )
-
-                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)
-
-        ctx.args = (causal, scale, mask, q_bucket_size, k_bucket_size)
-        ctx.save_for_backward(q, k, v, o, all_row_sums, all_row_maxes)
-
-        return o
-
-    @staticmethod
-    @torch.no_grad()
-    def backward(ctx, do):
-        """Algorithm 4 in the paper"""
-
-        causal, scale, mask, q_bucket_size, k_bucket_size = ctx.args
-        q, k, v, o, l, m = ctx.saved_tensors
-
-        device = q.device
-
-        max_neg_value = -torch.finfo(q.dtype).max
-        qk_len_diff = max(k.shape[-2] - q.shape[-2], 0)
-
-        dq = torch.zeros_like(q)
-        dk = torch.zeros_like(k)
-        dv = torch.zeros_like(v)
-
-        row_splits = zip(
-            q.split(q_bucket_size, dim=-2),
-            o.split(q_bucket_size, dim=-2),
-            do.split(q_bucket_size, dim=-2),
-            mask,
-            l.split(q_bucket_size, dim=-2),
-            m.split(q_bucket_size, dim=-2),
-            dq.split(q_bucket_size, dim=-2),
-        )
-
-        for ind, (qc, oc, doc, row_mask, lc, mc, dqc) in enumerate(row_splits):
-            q_start_index = ind * q_bucket_size - qk_len_diff
-
-            col_splits = zip(
-                k.split(k_bucket_size, dim=-2),
-                v.split(k_bucket_size, dim=-2),
-                dk.split(k_bucket_size, dim=-2),
-                dv.split(k_bucket_size, dim=-2),
-            )
-
-            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
-                )
-
-                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)
-                    attn_weights.masked_fill_(causal_mask, max_neg_value)
-
-                exp_attn_weights = torch.exp(attn_weights - mc)
-
-                if exists(row_mask):
-                    exp_attn_weights.masked_fill_(~row_mask, 0.0)
-
-                p = exp_attn_weights / lc
-
-                dv_chunk = torch.einsum("... i j, ... i d -> ... j d", p, doc)
-                dp = torch.einsum("... i d, ... j d -> ... i j", doc, vc)
-
-                D = (doc * oc).sum(dim=-1, keepdims=True)
-                ds = p * scale * (dp - D)
-
-                dq_chunk = torch.einsum("... i j, ... j d -> ... i d", ds, kc)
-                dk_chunk = torch.einsum("... i j, ... i d -> ... j d", ds, qc)
-
-                dqc.add_(dq_chunk)
-                dkc.add_(dk_chunk)
-                dvc.add_(dv_chunk)
-
-        return dq, dk, dv, None, None, None, None
-
-
-# endregion
-
-
-def get_parameter_dtype(parameter: torch.nn.Module):
-    return next(parameter.parameters()).dtype
-
-
-def get_parameter_device(parameter: torch.nn.Module):
-    return next(parameter.parameters()).device
-
-
-def get_timestep_embedding(
-    timesteps: torch.Tensor,
-    embedding_dim: int,
-    flip_sin_to_cos: bool = False,
-    downscale_freq_shift: float = 1,
-    scale: float = 1,
-    max_period: int = 10000,
-):
-    """
-    This matches the implementation in Denoising Diffusion Probabilistic Models: Create sinusoidal timestep embeddings.
-
-    :param timesteps: a 1-D Tensor of N indices, one per batch element.
-                      These may be fractional.
-    :param embedding_dim: the dimension of the output. :param max_period: controls the minimum frequency of the
-    embeddings. :return: an [N x dim] Tensor of positional embeddings.
-    """
-    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 = exponent / (half_dim - downscale_freq_shift)
-
-    emb = torch.exp(exponent)
-    emb = timesteps[:, None].float() * emb[None, :]
-
-    # scale embeddings
-    emb = scale * emb
-
-    # concat sine and cosine embeddings
-    emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)
-
-    # flip sine and cosine embeddings
-    if flip_sin_to_cos:
-        emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1)
-
-    # zero pad
-    if embedding_dim % 2 == 1:
-        emb = torch.nn.functional.pad(emb, (0, 1, 0, 0))
-    return emb
-
-
-class SampleOutput:
-    def __init__(self, sample):
-        self.sample = sample
-
-
-class TimestepEmbedding(nn.Module):
-    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)
-        self.act = None
-        if act_fn == "silu":
-            self.act = nn.SiLU()
-        elif act_fn == "mish":
-            self.act = nn.Mish()
-
-        if out_dim is not None:
-            time_embed_dim_out = out_dim
-        else:
-            time_embed_dim_out = time_embed_dim
-        self.linear_2 = nn.Linear(time_embed_dim, time_embed_dim_out)
-
-    def forward(self, sample):
-        sample = self.linear_1(sample)
-
-        if self.act is not None:
-            sample = self.act(sample)
-
-        sample = self.linear_2(sample)
-        return sample
-
-
-class Timesteps(nn.Module):
-    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
-        self.downscale_freq_shift = downscale_freq_shift
-
-    def forward(self, timesteps):
-        t_emb = get_timestep_embedding(
-            timesteps,
-            self.num_channels,
-            flip_sin_to_cos=self.flip_sin_to_cos,
-            downscale_freq_shift=self.downscale_freq_shift,
-        )
-        return t_emb
-
-
-class ResnetBlock2D(nn.Module):
-    def __init__(
-        self,
-        in_channels,
-        out_channels,
-    ):
-        super().__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.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
-        )
-
-        # if non_linearity == "swish":
-        self.nonlinearity = lambda x: F.silu(x)
-
-        self.use_in_shortcut = self.in_channels != self.out_channels
-
-        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
-            )
-
-    def forward(self, input_tensor, temb):
-        hidden_states = input_tensor
-
-        hidden_states = self.norm1(hidden_states)
-        hidden_states = self.nonlinearity(hidden_states)
-
-        hidden_states = self.conv1(hidden_states)
-
-        temb = self.time_emb_proj(self.nonlinearity(temb))[:, :, None, None]
-        hidden_states = hidden_states + temb
-
-        hidden_states = self.norm2(hidden_states)
-        hidden_states = self.nonlinearity(hidden_states)
-
-        hidden_states = self.conv2(hidden_states)
-
-        if self.conv_shortcut is not None:
-            input_tensor = self.conv_shortcut(input_tensor)
-
-        output_tensor = input_tensor + hidden_states
-
-        return output_tensor
-
-
-class DownBlock2D(nn.Module):
-    def __init__(
-        self,
-        in_channels: int,
-        out_channels: int,
-        add_downsample=True,
-    ):
-        super().__init__()
-
-        self.has_cross_attention = False
-        resnets = []
-
-        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,
-                )
-            )
-        self.resnets = nn.ModuleList(resnets)
-
-        if add_downsample:
-            self.downsamplers = [Downsample2D(out_channels, out_channels=out_channels)]
-        else:
-            self.downsamplers = None
-
-        self.gradient_checkpointing = False
-
-    def set_use_memory_efficient_attention(self, xformers, mem_eff):
-        pass
-
-    def set_use_sdpa(self, sdpa):
-        pass
-
-    def forward(self, hidden_states, temb=None):
-        output_states = ()
-
-        for resnet in self.resnets:
-            if self.training and self.gradient_checkpointing:
-
-                def create_custom_forward(module):
-                    def custom_forward(*inputs):
-                        return module(*inputs)
-
-                    return custom_forward
-
-                hidden_states = torch.utils.checkpoint.checkpoint(
-                    create_custom_forward(resnet), hidden_states, temb
-                )
-            else:
-                hidden_states = resnet(hidden_states, temb)
-
-            output_states += (hidden_states,)
-
-        if self.downsamplers is not None:
-            for downsampler in self.downsamplers:
-                hidden_states = downsampler(hidden_states)
-
-            output_states += (hidden_states,)
-
-        return hidden_states, output_states
-
-
-class Downsample2D(nn.Module):
-    def __init__(self, channels, out_channels):
-        super().__init__()
-
-        self.channels = channels
-        self.out_channels = out_channels
-
-        self.conv = nn.Conv2d(self.channels, self.out_channels, 3, stride=2, padding=1)
-
-    def forward(self, hidden_states):
-        assert hidden_states.shape[1] == self.channels
-        hidden_states = self.conv(hidden_states)
-
-        return hidden_states
-
-
-class CrossAttention(nn.Module):
-    def __init__(
-        self,
-        query_dim: int,
-        cross_attention_dim: Optional[int] = None,
-        heads: int = 8,
-        dim_head: int = 64,
-        upcast_attention: bool = False,
-    ):
-        super().__init__()
-        inner_dim = dim_head * heads
-        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
-        self.heads = heads
-
-        self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
-        self.to_k = nn.Linear(cross_attention_dim, inner_dim, bias=False)
-        self.to_v = nn.Linear(cross_attention_dim, inner_dim, bias=False)
-
-        self.to_out = nn.ModuleList([])
-        self.to_out.append(nn.Linear(inner_dim, query_dim))
-        # no dropout here
-
-        self.use_memory_efficient_attention_xformers = False
-        self.use_memory_efficient_attention_mem_eff = False
-        self.use_sdpa = False
-
-    def set_use_memory_efficient_attention(self, xformers, mem_eff):
-        self.use_memory_efficient_attention_xformers = xformers
-        self.use_memory_efficient_attention_mem_eff = mem_eff
-
-    def set_use_sdpa(self, sdpa):
-        self.use_sdpa = sdpa
-
-    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
-        )
-        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
-        )
-        return tensor
-
-    def forward(self, hidden_states, context=None, mask=None):
-        if self.use_memory_efficient_attention_xformers:
-            return self.forward_memory_efficient_xformers(hidden_states, context, mask)
-        if self.use_memory_efficient_attention_mem_eff:
-            return self.forward_memory_efficient_mem_eff(hidden_states, context, mask)
-        if self.use_sdpa:
-            return self.forward_sdpa(hidden_states, context, mask)
-
-        query = self.to_q(hidden_states)
-        context = context if context is not None else hidden_states
-        key = self.to_k(context)
-        value = self.to_v(context)
-
-        query = self.reshape_heads_to_batch_dim(query)
-        key = self.reshape_heads_to_batch_dim(key)
-        value = self.reshape_heads_to_batch_dim(value)
-
-        hidden_states = self._attention(query, key, value)
-
-        # linear proj
-        hidden_states = self.to_out[0](hidden_states)
-        # hidden_states = self.to_out[1](hidden_states)     # no dropout
-        return hidden_states
-
-    def _attention(self, query, key, value):
-        if self.upcast_attention:
-            query = query.float()
-            key = key.float()
-
-        attention_scores = torch.baddbmm(
-            torch.empty(
-                query.shape[0],
-                query.shape[1],
-                key.shape[1],
-                dtype=query.dtype,
-                device=query.device,
-            ),
-            query,
-            key.transpose(-1, -2),
-            beta=0,
-            alpha=self.scale,
-        )
-        attention_probs = attention_scores.softmax(dim=-1)
-
-        # cast back to the original dtype
-        attention_probs = attention_probs.to(value.dtype)
-
-        # compute attention output
-        hidden_states = torch.bmm(attention_probs, value)
-
-        # reshape hidden_states
-        hidden_states = self.reshape_batch_dim_to_heads(hidden_states)
-        return hidden_states
-
-    # TODO support Hypernetworks
-    def forward_memory_efficient_xformers(self, x, context=None, mask=None):
-        import xformers.ops
-
-        h = self.heads
-        q_in = self.to_q(x)
-        context = context if context is not None else x
-        context = context.to(x.dtype)
-        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)
-        )
-        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 = rearrange(out, "b n h d -> b n (h d)", h=h)
-
-        out = self.to_out[0](out)
-        return out
-
-    def forward_memory_efficient_mem_eff(self, x, context=None, mask=None):
-        flash_func = FlashAttentionFunction
-
-        q_bucket_size = 512
-        k_bucket_size = 1024
-
-        h = self.heads
-        q = self.to_q(x)
-        context = context if context is not None else x
-        context = context.to(x.dtype)
-        k = self.to_k(context)
-        v = self.to_v(context)
-        del context, x
-
-        q, k, v = map(lambda t: rearrange(t, "b n (h d) -> b h n d", h=h), (q, k, v))
-
-        out = flash_func.apply(q, k, v, mask, False, q_bucket_size, k_bucket_size)
-
-        out = rearrange(out, "b h n d -> b n (h d)")
-
-        out = self.to_out[0](out)
-        return out
-
-    def forward_sdpa(self, x, context=None, mask=None):
-        h = self.heads
-        q_in = self.to_q(x)
-        context = context if context is not None else x
-        context = context.to(x.dtype)
-        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)
-        )
-        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 = rearrange(out, "b h n d -> b n (h d)", h=h)
-
-        out = self.to_out[0](out)
-        return out
-
-
-# feedforward
-class GEGLU(nn.Module):
-    r"""
-    A variant of the gated linear unit activation function from https://arxiv.org/abs/2002.05202.
-
-    Parameters:
-        dim_in (`int`): The number of channels in the input.
-        dim_out (`int`): The number of channels in the output.
-    """
-
-    def __init__(self, dim_in: int, dim_out: int):
-        super().__init__()
-        self.proj = nn.Linear(dim_in, dim_out * 2)
-
-    def gelu(self, gate):
-        if gate.device.type != "mps":
-            return F.gelu(gate)
-        # mps: gelu is not implemented for float16
-        return F.gelu(gate.to(dtype=torch.float32)).to(dtype=gate.dtype)
-
-    def forward(self, hidden_states):
-        hidden_states, gate = self.proj(hidden_states).chunk(2, dim=-1)
-        return hidden_states * self.gelu(gate)
-
-
-class FeedForward(nn.Module):
-    def __init__(
-        self,
-        dim: int,
-    ):
-        super().__init__()
-        inner_dim = int(dim * 4)  # mult is always 4
-
-        self.net = nn.ModuleList([])
-        # project in
-        self.net.append(GEGLU(dim, inner_dim))
-        # project dropout
-        self.net.append(nn.Identity())  # nn.Dropout(0)) # dummy for dropout with 0
-        # project out
-        self.net.append(nn.Linear(inner_dim, dim))
-
-    def forward(self, hidden_states):
-        for module in self.net:
-            hidden_states = module(hidden_states)
-        return 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,
-    ):
-        super().__init__()
-
-        # 1. Self-Attn
-        self.attn1 = CrossAttention(
-            query_dim=dim,
-            cross_attention_dim=None,
-            heads=num_attention_heads,
-            dim_head=attention_head_dim,
-            upcast_attention=upcast_attention,
-        )
-        self.ff = FeedForward(dim)
-
-        # 2. Cross-Attn
-        self.attn2 = CrossAttention(
-            query_dim=dim,
-            cross_attention_dim=cross_attention_dim,
-            heads=num_attention_heads,
-            dim_head=attention_head_dim,
-            upcast_attention=upcast_attention,
-        )
-
-        self.norm1 = nn.LayerNorm(dim)
-        self.norm2 = nn.LayerNorm(dim)
-
-        # 3. Feed-forward
-        self.norm3 = nn.LayerNorm(dim)
-
-    def set_use_memory_efficient_attention(self, xformers: bool, mem_eff: bool):
-        self.attn1.set_use_memory_efficient_attention(xformers, mem_eff)
-        self.attn2.set_use_memory_efficient_attention(xformers, mem_eff)
-
-    def set_use_sdpa(self, sdpa: bool):
-        self.attn1.set_use_sdpa(sdpa)
-        self.attn2.set_use_sdpa(sdpa)
-
-    def forward(self, hidden_states, context=None, timestep=None):
-        # 1. Self-Attention
-        norm_hidden_states = self.norm1(hidden_states)
-
-        hidden_states = self.attn1(norm_hidden_states) + hidden_states
-
-        # 2. Cross-Attention
-        norm_hidden_states = self.norm2(hidden_states)
-        hidden_states = self.attn2(norm_hidden_states, context=context) + hidden_states
-
-        # 3. Feed-forward
-        hidden_states = self.ff(self.norm3(hidden_states)) + hidden_states
-
-        return hidden_states
-
-
-class Transformer2DModel(nn.Module):
-    def __init__(
-        self,
-        num_attention_heads: int = 16,
-        attention_head_dim: int = 88,
-        in_channels: Optional[int] = None,
-        cross_attention_dim: Optional[int] = None,
-        use_linear_projection: bool = False,
-        upcast_attention: bool = False,
-    ):
-        super().__init__()
-        self.in_channels = in_channels
-        self.num_attention_heads = num_attention_heads
-        self.attention_head_dim = attention_head_dim
-        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,
-        )
-
-        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.transformer_blocks = nn.ModuleList(
-            [
-                BasicTransformerBlock(
-                    inner_dim,
-                    num_attention_heads,
-                    attention_head_dim,
-                    cross_attention_dim=cross_attention_dim,
-                    upcast_attention=upcast_attention,
-                )
-            ]
-        )
-
-        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
-            )
-
-    def set_use_memory_efficient_attention(self, xformers, mem_eff):
-        for transformer in self.transformer_blocks:
-            transformer.set_use_memory_efficient_attention(xformers, mem_eff)
-
-    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,
-    ):
-        # 1. Input
-        batch, _, height, weight = hidden_states.shape
-        residual = hidden_states
-
-        hidden_states = self.norm(hidden_states)
-        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
-            )
-        else:
-            inner_dim = hidden_states.shape[1]
-            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
-            )
-
-        # 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 = 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()
-            )
-
-        output = hidden_states + residual
-
-        if not return_dict:
-            return (output,)
-
-        return SampleOutput(sample=output)
-
-
-class CrossAttnDownBlock2D(nn.Module):
-    def __init__(
-        self,
-        in_channels: int,
-        out_channels: int,
-        add_downsample=True,
-        cross_attention_dim=1280,
-        attn_num_head_channels=1,
-        use_linear_projection=False,
-        upcast_attention=False,
-    ):
-        super().__init__()
-        self.has_cross_attention = True
-        resnets = []
-        attentions = []
-
-        self.attn_num_head_channels = attn_num_head_channels
-
-        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)
-            )
-            attentions.append(
-                Transformer2DModel(
-                    attn_num_head_channels,
-                    out_channels // attn_num_head_channels,
-                    in_channels=out_channels,
-                    cross_attention_dim=cross_attention_dim,
-                    use_linear_projection=use_linear_projection,
-                    upcast_attention=upcast_attention,
-                )
-            )
-        self.attentions = nn.ModuleList(attentions)
-        self.resnets = nn.ModuleList(resnets)
-
-        if add_downsample:
-            self.downsamplers = nn.ModuleList(
-                [Downsample2D(out_channels, out_channels)]
-            )
-        else:
-            self.downsamplers = None
-
-        self.gradient_checkpointing = False
-
-    def set_use_memory_efficient_attention(self, xformers, mem_eff):
-        for attn in self.attentions:
-            attn.set_use_memory_efficient_attention(xformers, mem_eff)
-
-    def set_use_sdpa(self, sdpa):
-        for attn in self.attentions:
-            attn.set_use_sdpa(sdpa)
-
-    def forward(self, hidden_states, temb=None, encoder_hidden_states=None):
-        output_states = ()
-
-        for resnet, attn in zip(self.resnets, self.attentions):
-            if self.training and self.gradient_checkpointing:
-
-                def create_custom_forward(module, return_dict=None):
-                    def custom_forward(*inputs):
-                        if return_dict is not None:
-                            return module(*inputs, return_dict=return_dict)
-                        else:
-                            return module(*inputs)
-
-                    return custom_forward
-
-                hidden_states = torch.utils.checkpoint.checkpoint(
-                    create_custom_forward(resnet), hidden_states, temb
-                )
-                hidden_states = torch.utils.checkpoint.checkpoint(
-                    create_custom_forward(attn, return_dict=False),
-                    hidden_states,
-                    encoder_hidden_states,
-                )[0]
-            else:
-                hidden_states = resnet(hidden_states, temb)
-                hidden_states = attn(
-                    hidden_states, encoder_hidden_states=encoder_hidden_states
-                ).sample
-
-            output_states += (hidden_states,)
-
-        if self.downsamplers is not None:
-            for downsampler in self.downsamplers:
-                hidden_states = downsampler(hidden_states)
-
-            output_states += (hidden_states,)
-
-        return hidden_states, output_states
-
-
-class UNetMidBlock2DCrossAttn(nn.Module):
-    def __init__(
-        self,
-        in_channels: int,
-        attn_num_head_channels=1,
-        cross_attention_dim=1280,
-        use_linear_projection=False,
-    ):
-        super().__init__()
-
-        self.has_cross_attention = True
-        self.attn_num_head_channels = attn_num_head_channels
-
-        # Middle block has two resnets and one attention
-        resnets = [
-            ResnetBlock2D(
-                in_channels=in_channels,
-                out_channels=in_channels,
-            ),
-            ResnetBlock2D(
-                in_channels=in_channels,
-                out_channels=in_channels,
-            ),
-        ]
-        attentions = [
-            Transformer2DModel(
-                attn_num_head_channels,
-                in_channels // attn_num_head_channels,
-                in_channels=in_channels,
-                cross_attention_dim=cross_attention_dim,
-                use_linear_projection=use_linear_projection,
-            )
-        ]
-
-        self.attentions = nn.ModuleList(attentions)
-        self.resnets = nn.ModuleList(resnets)
-
-        self.gradient_checkpointing = False
-
-    def set_use_memory_efficient_attention(self, xformers, mem_eff):
-        for attn in self.attentions:
-            attn.set_use_memory_efficient_attention(xformers, mem_eff)
-
-    def set_use_sdpa(self, sdpa):
-        for attn in self.attentions:
-            attn.set_use_sdpa(sdpa)
-
-    def forward(self, hidden_states, temb=None, encoder_hidden_states=None):
-        for i, resnet in enumerate(self.resnets):
-            attn = None if i == 0 else self.attentions[i - 1]
-
-            if self.training and self.gradient_checkpointing:
-
-                def create_custom_forward(module, return_dict=None):
-                    def custom_forward(*inputs):
-                        if return_dict is not None:
-                            return module(*inputs, return_dict=return_dict)
-                        else:
-                            return module(*inputs)
-
-                    return custom_forward
-
-                if attn is not None:
-                    hidden_states = torch.utils.checkpoint.checkpoint(
-                        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
-                )
-            else:
-                if attn is not None:
-                    hidden_states = attn(hidden_states, encoder_hidden_states).sample
-                hidden_states = resnet(hidden_states, temb)
-
-        return hidden_states
-
-
-class Upsample2D(nn.Module):
-    def __init__(self, channels, out_channels):
-        super().__init__()
-        self.channels = channels
-        self.out_channels = out_channels
-        self.conv = nn.Conv2d(self.channels, self.out_channels, 3, padding=1)
-
-    def forward(self, hidden_states, output_size):
-        assert hidden_states.shape[1] == self.channels
-
-        # Cast to float32 to as 'upsample_nearest2d_out_frame' op does not support bfloat16
-        # TODO(Suraj): Remove this cast once the issue is fixed in PyTorch
-        # https://github.com/pytorch/pytorch/issues/86679
-        dtype = hidden_states.dtype
-        if dtype == torch.bfloat16:
-            hidden_states = hidden_states.to(torch.float32)
-
-        # upsample_nearest_nhwc fails with large batch sizes. see https://github.com/huggingface/diffusers/issues/984
-        if hidden_states.shape[0] >= 64:
-            hidden_states = hidden_states.contiguous()
-
-        # 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"
-            )
-        else:
-            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:
-            hidden_states = hidden_states.to(dtype)
-
-        hidden_states = self.conv(hidden_states)
-
-        return hidden_states
-
-
-class UpBlock2D(nn.Module):
-    def __init__(
-        self,
-        in_channels: int,
-        prev_output_channel: int,
-        out_channels: int,
-        add_upsample=True,
-    ):
-        super().__init__()
-
-        self.has_cross_attention = False
-        resnets = []
-
-        for i in range(LAYERS_PER_BLOCK_UP):
-            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(
-                ResnetBlock2D(
-                    in_channels=resnet_in_channels + res_skip_channels,
-                    out_channels=out_channels,
-                )
-            )
-
-        self.resnets = nn.ModuleList(resnets)
-
-        if add_upsample:
-            self.upsamplers = nn.ModuleList([Upsample2D(out_channels, out_channels)])
-        else:
-            self.upsamplers = None
-
-        self.gradient_checkpointing = False
-
-    def set_use_memory_efficient_attention(self, xformers, mem_eff):
-        pass
-
-    def set_use_sdpa(self, sdpa):
-        pass
-
-    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]
-            res_hidden_states_tuple = res_hidden_states_tuple[:-1]
-            hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)
-
-            if self.training and self.gradient_checkpointing:
-
-                def create_custom_forward(module):
-                    def custom_forward(*inputs):
-                        return module(*inputs)
-
-                    return custom_forward
-
-                hidden_states = torch.utils.checkpoint.checkpoint(
-                    create_custom_forward(resnet), hidden_states, temb
-                )
-            else:
-                hidden_states = resnet(hidden_states, temb)
-
-        if self.upsamplers is not None:
-            for upsampler in self.upsamplers:
-                hidden_states = upsampler(hidden_states, upsample_size)
-
-        return hidden_states
-
-
-class CrossAttnUpBlock2D(nn.Module):
-    def __init__(
-        self,
-        in_channels: int,
-        out_channels: int,
-        prev_output_channel: int,
-        attn_num_head_channels=1,
-        cross_attention_dim=1280,
-        add_upsample=True,
-        use_linear_projection=False,
-        upcast_attention=False,
-    ):
-        super().__init__()
-        resnets = []
-        attentions = []
-
-        self.has_cross_attention = True
-        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
-            )
-            resnet_in_channels = prev_output_channel if i == 0 else out_channels
-
-            resnets.append(
-                ResnetBlock2D(
-                    in_channels=resnet_in_channels + res_skip_channels,
-                    out_channels=out_channels,
-                )
-            )
-            attentions.append(
-                Transformer2DModel(
-                    attn_num_head_channels,
-                    out_channels // attn_num_head_channels,
-                    in_channels=out_channels,
-                    cross_attention_dim=cross_attention_dim,
-                    use_linear_projection=use_linear_projection,
-                    upcast_attention=upcast_attention,
-                )
-            )
-
-        self.attentions = nn.ModuleList(attentions)
-        self.resnets = nn.ModuleList(resnets)
-
-        if add_upsample:
-            self.upsamplers = nn.ModuleList([Upsample2D(out_channels, out_channels)])
-        else:
-            self.upsamplers = None
-
-        self.gradient_checkpointing = False
-
-    def set_use_memory_efficient_attention(self, xformers, mem_eff):
-        for attn in self.attentions:
-            attn.set_use_memory_efficient_attention(xformers, mem_eff)
-
-    def set_use_sdpa(self, spda):
-        for attn in self.attentions:
-            attn.set_use_sdpa(spda)
-
-    def forward(
-        self,
-        hidden_states,
-        res_hidden_states_tuple,
-        temb=None,
-        encoder_hidden_states=None,
-        upsample_size=None,
-    ):
-        for resnet, attn in zip(self.resnets, self.attentions):
-            # pop res hidden states
-            res_hidden_states = res_hidden_states_tuple[-1]
-            res_hidden_states_tuple = res_hidden_states_tuple[:-1]
-            hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)
-
-            if self.training and self.gradient_checkpointing:
-
-                def create_custom_forward(module, return_dict=None):
-                    def custom_forward(*inputs):
-                        if return_dict is not None:
-                            return module(*inputs, return_dict=return_dict)
-                        else:
-                            return module(*inputs)
-
-                    return custom_forward
-
-                hidden_states = torch.utils.checkpoint.checkpoint(
-                    create_custom_forward(resnet), hidden_states, temb
-                )
-                hidden_states = torch.utils.checkpoint.checkpoint(
-                    create_custom_forward(attn, return_dict=False),
-                    hidden_states,
-                    encoder_hidden_states,
-                )[0]
-            else:
-                hidden_states = resnet(hidden_states, temb)
-                hidden_states = attn(
-                    hidden_states, encoder_hidden_states=encoder_hidden_states
-                ).sample
-
-        if self.upsamplers is not None:
-            for upsampler in self.upsamplers:
-                hidden_states = upsampler(hidden_states, upsample_size)
-
-        return hidden_states
-
-
-def get_down_block(
-    down_block_type,
-    in_channels,
-    out_channels,
-    add_downsample,
-    attn_num_head_channels,
-    cross_attention_dim,
-    use_linear_projection,
-    upcast_attention,
-):
-    if down_block_type == "DownBlock2D":
-        return DownBlock2D(
-            in_channels=in_channels,
-            out_channels=out_channels,
-            add_downsample=add_downsample,
-        )
-    elif down_block_type == "CrossAttnDownBlock2D":
-        return CrossAttnDownBlock2D(
-            in_channels=in_channels,
-            out_channels=out_channels,
-            add_downsample=add_downsample,
-            cross_attention_dim=cross_attention_dim,
-            attn_num_head_channels=attn_num_head_channels,
-            use_linear_projection=use_linear_projection,
-            upcast_attention=upcast_attention,
-        )
-
-
-def get_up_block(
-    up_block_type,
-    in_channels,
-    out_channels,
-    prev_output_channel,
-    add_upsample,
-    attn_num_head_channels,
-    cross_attention_dim=None,
-    use_linear_projection=False,
-    upcast_attention=False,
-):
-    if up_block_type == "UpBlock2D":
-        return UpBlock2D(
-            in_channels=in_channels,
-            prev_output_channel=prev_output_channel,
-            out_channels=out_channels,
-            add_upsample=add_upsample,
-        )
-    elif up_block_type == "CrossAttnUpBlock2D":
-        return CrossAttnUpBlock2D(
-            in_channels=in_channels,
-            out_channels=out_channels,
-            prev_output_channel=prev_output_channel,
-            attn_num_head_channels=attn_num_head_channels,
-            cross_attention_dim=cross_attention_dim,
-            add_upsample=add_upsample,
-            use_linear_projection=use_linear_projection,
-            upcast_attention=upcast_attention,
-        )
-
-
-class UNet2DConditionModel(nn.Module):
-    _supports_gradient_checkpointing = True
-
-    def __init__(
-        self,
-        sample_size: Optional[int] = None,
-        attention_head_dim: Union[int, Tuple[int]] = 8,
-        cross_attention_dim: int = 1280,
-        use_linear_projection: bool = False,
-        upcast_attention: bool = False,
-        **kwargs,
-    ):
-        super().__init__()
-        assert sample_size is not None, "sample_size must be specified"
-        print(
-            f"UNet2DConditionModel: {sample_size}, {attention_head_dim}, {cross_attention_dim}, {use_linear_projection}, {upcast_attention}"
-        )
-
-        # 外部からの参照用に定義しておく
-        self.in_channels = IN_CHANNELS
-        self.out_channels = OUT_CHANNELS
-
-        self.sample_size = sample_size
-        self.prepare_config()
-
-        # state_dictの書式が変わるのでmoduleの持ち方は変えられない
-
-        # input
-        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_embedding = TimestepEmbedding(TIMESTEP_INPUT_DIM, TIME_EMBED_DIM)
-
-        self.down_blocks = nn.ModuleList([])
-        self.mid_block = None
-        self.up_blocks = nn.ModuleList([])
-
-        if isinstance(attention_head_dim, int):
-            attention_head_dim = (attention_head_dim,) * 4
-
-        # down
-        output_channel = BLOCK_OUT_CHANNELS[0]
-        for i, down_block_type in enumerate(DOWN_BLOCK_TYPES):
-            input_channel = output_channel
-            output_channel = BLOCK_OUT_CHANNELS[i]
-            is_final_block = i == len(BLOCK_OUT_CHANNELS) - 1
-
-            down_block = get_down_block(
-                down_block_type,
-                in_channels=input_channel,
-                out_channels=output_channel,
-                add_downsample=not is_final_block,
-                attn_num_head_channels=attention_head_dim[i],
-                cross_attention_dim=cross_attention_dim,
-                use_linear_projection=use_linear_projection,
-                upcast_attention=upcast_attention,
-            )
-            self.down_blocks.append(down_block)
-
-        # mid
-        self.mid_block = UNetMidBlock2DCrossAttn(
-            in_channels=BLOCK_OUT_CHANNELS[-1],
-            attn_num_head_channels=attention_head_dim[-1],
-            cross_attention_dim=cross_attention_dim,
-            use_linear_projection=use_linear_projection,
-        )
-
-        # count how many layers upsample the images
-        self.num_upsamplers = 0
-
-        # up
-        reversed_block_out_channels = list(reversed(BLOCK_OUT_CHANNELS))
-        reversed_attention_head_dim = list(reversed(attention_head_dim))
-        output_channel = reversed_block_out_channels[0]
-        for i, up_block_type in enumerate(UP_BLOCK_TYPES):
-            is_final_block = i == len(BLOCK_OUT_CHANNELS) - 1
-
-            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)
-            ]
-
-            # add upsample block for all BUT final layer
-            if not is_final_block:
-                add_upsample = True
-                self.num_upsamplers += 1
-            else:
-                add_upsample = False
-
-            up_block = get_up_block(
-                up_block_type,
-                in_channels=input_channel,
-                out_channels=output_channel,
-                prev_output_channel=prev_output_channel,
-                add_upsample=add_upsample,
-                attn_num_head_channels=reversed_attention_head_dim[i],
-                cross_attention_dim=cross_attention_dim,
-                use_linear_projection=use_linear_projection,
-                upcast_attention=upcast_attention,
-            )
-            self.up_blocks.append(up_block)
-            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_act = nn.SiLU()
-        self.conv_out = nn.Conv2d(
-            BLOCK_OUT_CHANNELS[0], OUT_CHANNELS, kernel_size=3, padding=1
-        )
-
-    # region diffusers compatibility
-    def prepare_config(self):
-        self.config = SimpleNamespace()
-
-    @property
-    def dtype(self) -> torch.dtype:
-        # `torch.dtype`: The dtype of the module (assuming that all the module parameters have the same dtype).
-        return get_parameter_dtype(self)
-
-    @property
-    def device(self) -> torch.device:
-        # `torch.device`: The device on which the module is (assuming that all the module parameters are on the same device).
-        return get_parameter_device(self)
-
-    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()
-        )
-
-    def enable_gradient_checkpointing(self):
-        self.set_gradient_checkpointing(value=True)
-
-    def disable_gradient_checkpointing(self):
-        self.set_gradient_checkpointing(value=False)
-
-    def set_use_memory_efficient_attention(self, xformers: bool, mem_eff: bool) -> None:
-        modules = self.down_blocks + [self.mid_block] + self.up_blocks
-        for module in modules:
-            module.set_use_memory_efficient_attention(xformers, mem_eff)
-
-    def set_use_sdpa(self, sdpa: bool) -> None:
-        modules = self.down_blocks + [self.mid_block] + self.up_blocks
-        for module in modules:
-            module.set_use_sdpa(sdpa)
-
-    def set_gradient_checkpointing(self, value=False):
-        modules = self.down_blocks + [self.mid_block] + self.up_blocks
-        for module in modules:
-            print(module.__class__.__name__, module.gradient_checkpointing, "->", value)
-            module.gradient_checkpointing = value
-
-    # endregion
-
-    def forward(
-        self,
-        sample: torch.FloatTensor,
-        timestep: Union[torch.Tensor, float, int],
-        encoder_hidden_states: torch.Tensor,
-        class_labels: Optional[torch.Tensor] = None,
-        return_dict: bool = True,
-        down_block_additional_residuals: Optional[Tuple[torch.Tensor]] = None,
-        mid_block_additional_residual: Optional[torch.Tensor] = None,
-    ) -> Union[Dict, Tuple]:
-        r"""
-        Args:
-            sample (`torch.FloatTensor`): (batch, channel, height, width) noisy inputs tensor
-            timestep (`torch.FloatTensor` or `float` or `int`): (batch) timesteps
-            encoder_hidden_states (`torch.FloatTensor`): (batch, sequence_length, feature_dim) encoder hidden states
-            return_dict (`bool`, *optional*, defaults to `True`):
-                Whether or not to return a dict instead of a plain tuple.
-
-        Returns:
-            `SampleOutput` or `tuple`:
-            `SampleOutput` if `return_dict` is True, otherwise a `tuple`. When returning a tuple, the first element is the sample tensor.
-        """
-        # By default samples have to be AT least a multiple of the overall upsampling factor.
-        # The overall upsampling factor is equal to 2 ** (# num of upsampling layears).
-        # However, the upsampling interpolation output size can be forced to fit any upsampling size
-        # on the fly if necessary.
-        # デフォルトではサンプルは「2^アップサンプルの数」、つまり64の倍数である必要がある
-        # ただそれ以外のサイズにも対応できるように、必要ならアップサンプルのサイズを変更する
-        # 多分画質が悪くなるので、64で割り切れるようにしておくのが良い
-        default_overall_up_factor = 2**self.num_upsamplers
-
-        # upsample size should be forwarded when sample is not a multiple of `default_overall_up_factor`
-        # 64で割り切れないときはupsamplerにサイズを伝える
-        forward_upsample_size = False
-        upsample_size = None
-
-        if any(s % default_overall_up_factor != 0 for s in sample.shape[-2:]):
-            # logger.info("Forward upsample size to force interpolation output size.")
-            forward_upsample_size = True
-
-        # 1. time
-        timesteps = timestep
-        timesteps = self.handle_unusual_timesteps(sample, timesteps)  # 変な時だけ処理
-
-        t_emb = self.time_proj(timesteps)
-
-        # timesteps does not contain any weights and will always return f32 tensors
-        # but time_embedding might actually be running in fp16. so we need to cast here.
-        # there might be better ways to encapsulate this.
-        # timestepsは重みを含まないので常にfloat32のテンソルを返す
-        # しかしtime_embeddingはfp16で動いているかもしれないので、ここでキャストする必要がある
-        # time_projでキャストしておけばいいんじゃね？
-        t_emb = t_emb.to(dtype=self.dtype)
-        emb = self.time_embedding(t_emb)
-
-        # 2. pre-process
-        sample = self.conv_in(sample)
-
-        # 3. down
-        down_block_res_samples = (sample,)
-        for downsample_block in self.down_blocks:
-            # downblockはforwardで必ずencoder_hidden_statesを受け取るようにしても良さそうだけど、
-            # まあこちらのほうがわかりやすいかもしれない
-            if downsample_block.has_cross_attention:
-                sample, res_samples = downsample_block(
-                    hidden_states=sample,
-                    temb=emb,
-                    encoder_hidden_states=encoder_hidden_states,
-                )
-            else:
-                sample, res_samples = downsample_block(hidden_states=sample, temb=emb)
-
-            down_block_res_samples += res_samples
-
-        # skip connectionにControlNetの出力を追加する
-        if down_block_additional_residuals is not None:
-            down_block_res_samples = list(down_block_res_samples)
-            for i in range(len(down_block_res_samples)):
-                down_block_res_samples[i] += down_block_additional_residuals[i]
-            down_block_res_samples = tuple(down_block_res_samples)
-
-        # 4. mid
-        sample = self.mid_block(
-            sample, emb, encoder_hidden_states=encoder_hidden_states
-        )
-
-        # ControlNetの出力を追加する
-        if mid_block_additional_residual is not None:
-            sample += mid_block_additional_residual
-
-        # 5. up
-        for i, upsample_block in enumerate(self.up_blocks):
-            is_final_block = i == len(self.up_blocks) - 1
-
-            res_samples = down_block_res_samples[-len(upsample_block.resnets) :]
-            down_block_res_samples = down_block_res_samples[
-                : -len(upsample_block.resnets)
-            ]  # skip connection
-
-            # if we have not reached the final block and need to forward the upsample size, we do it here
-            # 前述のように最後のブロック以外ではupsample_sizeを伝える
-            if not is_final_block and forward_upsample_size:
-                upsample_size = down_block_res_samples[-1].shape[2:]
-
-            if upsample_block.has_cross_attention:
-                sample = upsample_block(
-                    hidden_states=sample,
-                    temb=emb,
-                    res_hidden_states_tuple=res_samples,
-                    encoder_hidden_states=encoder_hidden_states,
-                    upsample_size=upsample_size,
-                )
-            else:
-                sample = upsample_block(
-                    hidden_states=sample,
-                    temb=emb,
-                    res_hidden_states_tuple=res_samples,
-                    upsample_size=upsample_size,
-                )
-
-        # 6. post-process
-        sample = self.conv_norm_out(sample)
-        sample = self.conv_act(sample)
-        sample = self.conv_out(sample)
-
-        if not return_dict:
-            return (sample,)
-
-        return SampleOutput(sample=sample)
-
-    def handle_unusual_timesteps(self, sample, timesteps):
-        r"""
-        timestampsがTensorでない場合、Tensorに変換する。またOnnx/Core MLと互換性のあるようにbatchサイズまでbroadcastする。
-        """
-        if not torch.is_tensor(timesteps):
-            # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can
-            # This would be a good case for the `match` statement (Python 3.10+)
-            is_mps = sample.device.type == "mps"
-            if isinstance(timesteps, float):
-                dtype = torch.float32 if is_mps else torch.float64
-            else:
-                dtype = torch.int32 if is_mps else torch.int64
-            timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device)
-        elif len(timesteps.shape) == 0:
-            timesteps = timesteps[None].to(sample.device)
-
-        # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
-        timesteps = timesteps.expand(sample.shape[0])
-
-        return timesteps
diff --git a/lycoris/kohya/sdxl_model_util.py b/lycoris/kohya/sdxl_model_util.py
deleted file mode 100644
index 666ddaa..0000000
--- a/lycoris/kohya/sdxl_model_util.py
+++ /dev/null
@@ -1,581 +0,0 @@
-import torch
-from safetensors.torch import load_file, save_file
-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
-
-
-VAE_SCALE_FACTOR = 0.13025
-MODEL_VERSION_SDXL_BASE_V0_9 = "sdxl_base_v0-9"
-
-# Diffusersの設定を読み込むための参照モデル
-DIFFUSERS_REF_MODEL_ID_SDXL = (
-    "stabilityai/stable-diffusion-xl-base-0.9"  # アクセス権が必要
-)
-
-DIFFUSERS_SDXL_UNET_CONFIG = {
-    "act_fn": "silu",
-    "addition_embed_type": "text_time",
-    "addition_embed_type_num_heads": 64,
-    "addition_time_embed_dim": 256,
-    "attention_head_dim": [5, 10, 20],
-    "block_out_channels": [320, 640, 1280],
-    "center_input_sample": False,
-    "class_embed_type": None,
-    "class_embeddings_concat": False,
-    "conv_in_kernel": 3,
-    "conv_out_kernel": 3,
-    "cross_attention_dim": 2048,
-    "cross_attention_norm": None,
-    "down_block_types": ["DownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D"],
-    "downsample_padding": 1,
-    "dual_cross_attention": False,
-    "encoder_hid_dim": None,
-    "encoder_hid_dim_type": None,
-    "flip_sin_to_cos": True,
-    "freq_shift": 0,
-    "in_channels": 4,
-    "layers_per_block": 2,
-    "mid_block_only_cross_attention": None,
-    "mid_block_scale_factor": 1,
-    "mid_block_type": "UNetMidBlock2DCrossAttn",
-    "norm_eps": 1e-05,
-    "norm_num_groups": 32,
-    "num_attention_heads": None,
-    "num_class_embeds": None,
-    "only_cross_attention": False,
-    "out_channels": 4,
-    "projection_class_embeddings_input_dim": 2816,
-    "resnet_out_scale_factor": 1.0,
-    "resnet_skip_time_act": False,
-    "resnet_time_scale_shift": "default",
-    "sample_size": 128,
-    "time_cond_proj_dim": None,
-    "time_embedding_act_fn": None,
-    "time_embedding_dim": None,
-    "time_embedding_type": "positional",
-    "timestep_post_act": None,
-    "transformer_layers_per_block": [1, 2, 10],
-    "up_block_types": ["CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "UpBlock2D"],
-    "upcast_attention": False,
-    "use_linear_projection": True,
-}
-
-
-def convert_sdxl_text_encoder_2_checkpoint(checkpoint, max_length):
-    SDXL_KEY_PREFIX = "conditioner.embedders.1.model."
-
-    # SD2のと、基本的には同じ。logit_scaleを後で使うので、それを追加で返す
-    # logit_scaleはcheckpointの保存時に使用する
-    def convert_key(key):
-        # common conversion
-        key = key.replace(SDXL_KEY_PREFIX + "transformer.", "text_model.encoder.")
-        key = key.replace(SDXL_KEY_PREFIX, "text_model.")
-
-        if "resblocks" in key:
-            # resblocks conversion
-            key = key.replace(".resblocks.", ".layers.")
-            if ".ln_" in key:
-                key = key.replace(".ln_", ".layer_norm")
-            elif ".mlp." in key:
-                key = key.replace(".c_fc.", ".fc1.")
-                key = key.replace(".c_proj.", ".fc2.")
-            elif ".attn.out_proj" in key:
-                key = key.replace(".attn.out_proj.", ".self_attn.out_proj.")
-            elif ".attn.in_proj" in key:
-                key = None  # 特殊なので後で処理する
-            else:
-                raise ValueError(f"unexpected key in SD: {key}")
-        elif ".positional_embedding" in key:
-            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"
-            )
-        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
-        elif ".embeddings.position_ids" in key:
-            key = None  # remove this key: make position_ids by ourselves
-        return key
-
-    keys = list(checkpoint.keys())
-    new_sd = {}
-    for key in keys:
-        new_key = convert_key(key)
-        if new_key is None:
-            continue
-        new_sd[new_key] = checkpoint[key]
-
-    # attnの変換
-    for key in keys:
-        if ".resblocks" in key and ".attn.in_proj_" in 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_pfx.replace("_weight", "")
-            key_pfx = key_pfx.replace("_bias", "")
-            key_pfx = key_pfx.replace(".attn.in_proj", ".self_attn.")
-            new_sd[key_pfx + "q_proj" + key_suffix] = values[0]
-            new_sd[key_pfx + "k_proj" + key_suffix] = values[1]
-            new_sd[key_pfx + "v_proj" + key_suffix] = values[2]
-
-    # original SD にはないので、position_idsを追加
-    position_ids = torch.Tensor([list(range(max_length))]).to(torch.int64)
-    new_sd["text_model.embeddings.position_ids"] = position_ids
-
-    # logit_scale はDiffusersには含まれないが、保存時に戻したいので別途返す
-    logit_scale = checkpoint.get(SDXL_KEY_PREFIX + "logit_scale", None)
-
-    return new_sd, logit_scale
-
-
-def load_models_from_sdxl_checkpoint(model_version, ckpt_path, map_location):
-    # model_version is reserved for future use
-
-    # Load the state dict
-    if model_util.is_safetensors(ckpt_path):
-        checkpoint = None
-        state_dict = load_file(ckpt_path, device=map_location)
-        epoch = None
-        global_step = None
-    else:
-        checkpoint = torch.load(ckpt_path, map_location=map_location)
-        if "state_dict" in checkpoint:
-            state_dict = checkpoint["state_dict"]
-            epoch = checkpoint.get("epoch", 0)
-            global_step = checkpoint.get("global_step", 0)
-        else:
-            state_dict = checkpoint
-            epoch = 0
-            global_step = 0
-        checkpoint = None
-
-    # U-Net
-    print("building U-Net")
-    unet = sdxl_original_unet.SdxlUNet2DConditionModel()
-
-    print("loading U-Net from checkpoint")
-    unet_sd = {}
-    for k in list(state_dict.keys()):
-        if k.startswith("model.diffusion_model."):
-            unet_sd[k.replace("model.diffusion_model.", "")] = state_dict.pop(k)
-    info = unet.load_state_dict(unet_sd)
-    print("U-Net: ", info)
-    del unet_sd
-
-    # Text Encoders
-    print("building text encoders")
-
-    # Text Encoder 1 is same to Stability AI's SDXL
-    text_model1_cfg = CLIPTextConfig(
-        vocab_size=49408,
-        hidden_size=768,
-        intermediate_size=3072,
-        num_hidden_layers=12,
-        num_attention_heads=12,
-        max_position_embeddings=77,
-        hidden_act="quick_gelu",
-        layer_norm_eps=1e-05,
-        dropout=0.0,
-        attention_dropout=0.0,
-        initializer_range=0.02,
-        initializer_factor=1.0,
-        pad_token_id=1,
-        bos_token_id=0,
-        eos_token_id=2,
-        model_type="clip_text_model",
-        projection_dim=768,
-        # torch_dtype="float32",
-        # transformers_version="4.25.0.dev0",
-    )
-    text_model1 = CLIPTextModel._from_config(text_model1_cfg)
-
-    # Text Encoder 2 is different from Stability AI's SDXL. SDXL uses open clip, but we use the model from HuggingFace.
-    # Note: Tokenizer from HuggingFace is different from SDXL. We must use open clip's tokenizer.
-    text_model2_cfg = CLIPTextConfig(
-        vocab_size=49408,
-        hidden_size=1280,
-        intermediate_size=5120,
-        num_hidden_layers=32,
-        num_attention_heads=20,
-        max_position_embeddings=77,
-        hidden_act="gelu",
-        layer_norm_eps=1e-05,
-        dropout=0.0,
-        attention_dropout=0.0,
-        initializer_range=0.02,
-        initializer_factor=1.0,
-        pad_token_id=1,
-        bos_token_id=0,
-        eos_token_id=2,
-        model_type="clip_text_model",
-        projection_dim=1280,
-        # torch_dtype="float32",
-        # transformers_version="4.25.0.dev0",
-    )
-    text_model2 = CLIPTextModelWithProjection(text_model2_cfg)
-
-    print("loading text encoders from checkpoint")
-    te1_sd = {}
-    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)
-            )
-        elif k.startswith("conditioner.embedders.1.model."):
-            te2_sd[k] = state_dict.pop(k)
-
-    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, strict=False)
-    print("text encoder 2:", info2)
-
-    # prepare vae
-    print("building VAE")
-    vae_config = model_util.create_vae_diffusers_config()
-    vae = AutoencoderKL(**vae_config)  # .to(device)
-
-    print("loading VAE from checkpoint")
-    converted_vae_checkpoint = model_util.convert_ldm_vae_checkpoint(
-        state_dict, vae_config
-    )
-    info = vae.load_state_dict(converted_vae_checkpoint)
-    print("VAE:", info)
-
-    ckpt_info = (epoch, global_step) if epoch is not None else None
-    return text_model1, text_model2, vae, unet, logit_scale, ckpt_info
-
-
-def make_unet_conversion_map():
-    unet_conversion_map_layer = []
-
-    for i in range(3):  # num_blocks is 3 in sdxl
-        # loop over downblocks/upblocks
-        for j in range(2):
-            # loop over resnets/attentions for downblocks
-            hf_down_res_prefix = f"down_blocks.{i}.resnets.{j}."
-            sd_down_res_prefix = f"input_blocks.{3*i + j + 1}.0."
-            unet_conversion_map_layer.append((sd_down_res_prefix, hf_down_res_prefix))
-
-            if i < 3:
-                # 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)
-                )
-
-        for j in range(3):
-            # loop over resnets/attentions for upblocks
-            hf_up_res_prefix = f"up_blocks.{i}.resnets.{j}."
-            sd_up_res_prefix = f"output_blocks.{3*i + j}.0."
-            unet_conversion_map_layer.append((sd_up_res_prefix, hf_up_res_prefix))
-
-            # if i > 0: commentout for sdxl
-            # no attention layers in up_blocks.0
-            hf_up_atn_prefix = f"up_blocks.{i}.attentions.{j}."
-            sd_up_atn_prefix = f"output_blocks.{3*i + j}.1."
-            unet_conversion_map_layer.append((sd_up_atn_prefix, hf_up_atn_prefix))
-
-        if i < 3:
-            # 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)
-            )
-
-            # no upsample in up_blocks.3
-            hf_upsample_prefix = f"up_blocks.{i}.upsamplers.0."
-            sd_upsample_prefix = f"output_blocks.{3*i + 2}.{2}."  # change for sdxl
-            unet_conversion_map_layer.append((sd_upsample_prefix, hf_upsample_prefix))
-
-    hf_mid_atn_prefix = "mid_block.attentions.0."
-    sd_mid_atn_prefix = "middle_block.1."
-    unet_conversion_map_layer.append((sd_mid_atn_prefix, hf_mid_atn_prefix))
-
-    for j in range(2):
-        hf_mid_res_prefix = f"mid_block.resnets.{j}."
-        sd_mid_res_prefix = f"middle_block.{2*j}."
-        unet_conversion_map_layer.append((sd_mid_res_prefix, hf_mid_res_prefix))
-
-    unet_conversion_map_resnet = [
-        # (stable-diffusion, HF Diffusers)
-        ("in_layers.0.", "norm1."),
-        ("in_layers.2.", "conv1."),
-        ("out_layers.0.", "norm2."),
-        ("out_layers.3.", "conv2."),
-        ("emb_layers.1.", "time_emb_proj."),
-        ("skip_connection.", "conv_shortcut."),
-    ]
-
-    unet_conversion_map = []
-    for sd, hf in unet_conversion_map_layer:
-        if "resnets" in hf:
-            for sd_res, hf_res in unet_conversion_map_resnet:
-                unet_conversion_map.append((sd + sd_res, hf + hf_res))
-        else:
-            unet_conversion_map.append((sd, hf))
-
-    for j in range(2):
-        hf_time_embed_prefix = f"time_embedding.linear_{j+1}."
-        sd_time_embed_prefix = f"time_embed.{j*2}."
-        unet_conversion_map.append((sd_time_embed_prefix, hf_time_embed_prefix))
-
-    for j in range(2):
-        hf_label_embed_prefix = f"add_embedding.linear_{j+1}."
-        sd_label_embed_prefix = f"label_emb.0.{j*2}."
-        unet_conversion_map.append((sd_label_embed_prefix, hf_label_embed_prefix))
-
-    unet_conversion_map.append(("input_blocks.0.0.", "conv_in."))
-    unet_conversion_map.append(("out.0.", "conv_norm_out."))
-    unet_conversion_map.append(("out.2.", "conv_out."))
-
-    return unet_conversion_map
-
-
-def convert_diffusers_unet_state_dict_to_sdxl(du_sd):
-    unet_conversion_map = make_unet_conversion_map()
-
-    conversion_map = {hf: sd for sd, hf in unet_conversion_map}
-    return convert_unet_state_dict(du_sd, conversion_map)
-
-
-def convert_unet_state_dict(src_sd, conversion_map):
-    converted_sd = {}
-    for src_key, value in src_sd.items():
-        # さすがに全部回すのは時間がかかるので右から要素を削りつつprefixを探す
-        src_key_fragments = src_key.split(".")[:-1]  # remove weight/bias
-        while len(src_key_fragments) > 0:
-            src_key_prefix = ".".join(src_key_fragments) + "."
-            if src_key_prefix in conversion_map:
-                converted_prefix = conversion_map[src_key_prefix]
-                converted_key = converted_prefix + src_key[len(src_key_prefix) :]
-                converted_sd[converted_key] = value
-                break
-            src_key_fragments.pop(-1)
-        assert len(src_key_fragments) > 0, f"key {src_key} not found in conversion map"
-
-    return converted_sd
-
-
-def convert_sdxl_unet_state_dict_to_diffusers(sd):
-    unet_conversion_map = make_unet_conversion_map()
-
-    conversion_dict = {sd: hf for sd, hf in unet_conversion_map}
-    return convert_unet_state_dict(sd, conversion_dict)
-
-
-def convert_text_encoder_2_state_dict_to_sdxl(checkpoint, logit_scale):
-    def convert_key(key):
-        # position_idsの除去
-        if ".position_ids" in key:
-            return None
-
-        # common
-        key = key.replace("text_model.encoder.", "transformer.")
-        key = key.replace("text_model.", "")
-        if "layers" in key:
-            # resblocks conversion
-            key = key.replace(".layers.", ".resblocks.")
-            if ".layer_norm" in key:
-                key = key.replace(".layer_norm", ".ln_")
-            elif ".mlp." in key:
-                key = key.replace(".fc1.", ".c_fc.")
-                key = key.replace(".fc2.", ".c_proj.")
-            elif ".self_attn.out_proj" in key:
-                key = key.replace(".self_attn.out_proj.", ".attn.out_proj.")
-            elif ".self_attn." in key:
-                key = None  # 特殊なので後で処理する
-            else:
-                raise ValueError(f"unexpected key in DiffUsers model: {key}")
-        elif ".position_embedding" in key:
-            key = key.replace(
-                "embeddings.position_embedding.weight", "positional_embedding"
-            )
-        elif ".token_embedding" in key:
-            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:
-            key = key.replace("final_layer_norm", "ln_final")
-        return key
-
-    keys = list(checkpoint.keys())
-    new_sd = {}
-    for key in keys:
-        new_key = convert_key(key)
-        if new_key is None:
-            continue
-        new_sd[new_key] = checkpoint[key]
-
-    # attnの変換
-    for key in keys:
-        if "layers" in key and "q_proj" in key:
-            # 三つを結合
-            key_q = key
-            key_k = key.replace("q_proj", "k_proj")
-            key_v = key.replace("q_proj", "v_proj")
-
-            value_q = checkpoint[key_q]
-            value_k = checkpoint[key_k]
-            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 = new_key.replace(".self_attn.q_proj.", ".attn.in_proj_")
-            new_sd[new_key] = value
-
-    if logit_scale is not None:
-        new_sd["logit_scale"] = logit_scale
-
-    return new_sd
-
-
-def save_stable_diffusion_checkpoint(
-    output_file,
-    text_encoder1,
-    text_encoder2,
-    unet,
-    epochs,
-    steps,
-    ckpt_info,
-    vae,
-    logit_scale,
-    save_dtype=None,
-):
-    state_dict = {}
-
-    def update_sd(prefix, sd):
-        for k, v in sd.items():
-            key = prefix + k
-            if save_dtype is not None:
-                v = v.detach().clone().to("cpu").to(save_dtype)
-            state_dict[key] = v
-
-    # Convert the UNet model
-    update_sd("model.diffusion_model.", unet.state_dict())
-
-    # 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
-    )
-    update_sd("conditioner.embedders.1.model.", text_enc2_dict)
-
-    # Convert the VAE
-    vae_dict = model_util.convert_vae_state_dict(vae.state_dict())
-    update_sd("first_stage_model.", vae_dict)
-
-    # Put together new checkpoint
-    key_count = len(state_dict.keys())
-    new_ckpt = {"state_dict": state_dict}
-
-    # epoch and global_step are sometimes not int
-    if ckpt_info is not None:
-        epochs += ckpt_info[0]
-        steps += ckpt_info[1]
-
-    new_ckpt["epoch"] = epochs
-    new_ckpt["global_step"] = steps
-
-    if model_util.is_safetensors(output_file):
-        save_file(state_dict, output_file)
-    else:
-        torch.save(new_ckpt, output_file)
-
-    return key_count
-
-
-def save_diffusers_checkpoint(
-    output_dir,
-    text_encoder1,
-    text_encoder2,
-    unet,
-    pretrained_model_name_or_path,
-    vae=None,
-    use_safetensors=False,
-    save_dtype=None,
-):
-    from diffusers import StableDiffusionXLPipeline
-
-    # convert U-Net
-    unet_sd = unet.state_dict()
-    du_unet_sd = convert_sdxl_unet_state_dict_to_diffusers(unet_sd)
-
-    diffusers_unet = UNet2DConditionModel(**DIFFUSERS_SDXL_UNET_CONFIG)
-    if save_dtype is not None:
-        diffusers_unet.to(save_dtype)
-    diffusers_unet.load_state_dict(du_unet_sd)
-
-    # create pipeline to save
-    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"
-    )
-    if vae is None:
-        vae = AutoencoderKL.from_pretrained(
-            pretrained_model_name_or_path, subfolder="vae"
-        )
-
-    # prevent local path from being saved
-    def remove_name_or_path(model):
-        if hasattr(model, "config"):
-            model.config._name_or_path = None
-            model.config._name_or_path = None
-
-    remove_name_or_path(diffusers_unet)
-    remove_name_or_path(text_encoder1)
-    remove_name_or_path(text_encoder2)
-    remove_name_or_path(scheduler)
-    remove_name_or_path(tokenizer1)
-    remove_name_or_path(tokenizer2)
-    remove_name_or_path(vae)
-
-    pipeline = StableDiffusionXLPipeline(
-        unet=diffusers_unet,
-        text_encoder=text_encoder1,
-        text_encoder_2=text_encoder2,
-        vae=vae,
-        scheduler=scheduler,
-        tokenizer=tokenizer1,
-        tokenizer_2=tokenizer2,
-    )
-    if save_dtype is not None:
-        pipeline.to(None, save_dtype)
-    pipeline.save_pretrained(output_dir, safe_serialization=use_safetensors)
diff --git a/lycoris/kohya/sdxl_original_unet.py b/lycoris/kohya/sdxl_original_unet.py
deleted file mode 100644
index 9beeb73..0000000
--- a/lycoris/kohya/sdxl_original_unet.py
+++ /dev/null
@@ -1,1243 +0,0 @@
-# Diffusersのコードをベースとした sd_xl_baseのU-Net
-# state dictの形式をSDXLに合わせてある
-
-"""
-      target: sgm.modules.diffusionmodules.openaimodel.UNetModel
-      params:
-        adm_in_channels: 2816
-        num_classes: sequential
-        use_checkpoint: True
-        in_channels: 4
-        out_channels: 4
-        model_channels: 320
-        attention_resolutions: [4, 2]
-        num_res_blocks: 2
-        channel_mult: [1, 2, 4]
-        num_head_channels: 64
-        use_spatial_transformer: True
-        use_linear_in_transformer: True
-        transformer_depth: [1, 2, 10]  # note: the first is unused (due to attn_res starting at 2) 32, 16, 8 --> 64, 32, 16
-        context_dim: 2048
-        spatial_transformer_attn_type: softmax-xformers
-        legacy: False
-"""
-
-import math
-from types import SimpleNamespace
-from typing import Optional
-import torch
-import torch.utils.checkpoint
-from torch import nn
-from torch.nn import functional as F
-from einops import rearrange
-
-
-IN_CHANNELS: int = 4
-OUT_CHANNELS: int = 4
-ADM_IN_CHANNELS: int = 2816
-CONTEXT_DIM: int = 2048
-MODEL_CHANNELS: int = 320
-TIME_EMBED_DIM = 320 * 4
-
-
-# region memory effcient attention
-
-# FlashAttentionを使うCrossAttention
-# based on https://github.com/lucidrains/memory-efficient-attention-pytorch/blob/main/memory_efficient_attention_pytorch/flash_attention.py
-# LICENSE MIT https://github.com/lucidrains/memory-efficient-attention-pytorch/blob/main/LICENSE
-
-# constants
-
-EPSILON = 1e-6
-
-# helper functions
-
-
-def exists(val):
-    return val is not None
-
-
-def default(val, d):
-    return val if exists(val) else d
-
-
-# flash attention forwards and backwards
-
-# https://arxiv.org/abs/2205.14135
-
-
-class FlashAttentionFunction(torch.autograd.Function):
-    @staticmethod
-    @torch.no_grad()
-    def forward(ctx, q, k, v, mask, causal, q_bucket_size, k_bucket_size):
-        """Algorithm 2 in the paper"""
-
-        device = q.device
-        dtype = q.dtype
-        max_neg_value = -torch.finfo(q.dtype).max
-        qk_len_diff = max(k.shape[-2] - q.shape[-2], 0)
-
-        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
-        )
-
-        scale = q.shape[-1] ** -0.5
-
-        if not exists(mask):
-            mask = (None,) * math.ceil(q.shape[-2] / q_bucket_size)
-        else:
-            mask = rearrange(mask, "b n -> b 1 1 n")
-            mask = mask.split(q_bucket_size, dim=-1)
-
-        row_splits = zip(
-            q.split(q_bucket_size, dim=-2),
-            o.split(q_bucket_size, dim=-2),
-            mask,
-            all_row_sums.split(q_bucket_size, dim=-2),
-            all_row_maxes.split(q_bucket_size, dim=-2),
-        )
-
-        for ind, (qc, oc, row_mask, row_sums, row_maxes) in enumerate(row_splits):
-            q_start_index = ind * q_bucket_size - qk_len_diff
-
-            col_splits = zip(
-                k.split(k_bucket_size, dim=-2),
-                v.split(k_bucket_size, dim=-2),
-            )
-
-            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
-                )
-
-                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)
-                    attn_weights.masked_fill_(causal_mask, max_neg_value)
-
-                block_row_maxes = attn_weights.amax(dim=-1, keepdims=True)
-                attn_weights -= block_row_maxes
-                exp_weights = torch.exp(attn_weights)
-
-                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
-                )
-
-                new_row_maxes = torch.maximum(block_row_maxes, row_maxes)
-
-                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
-                )
-
-                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)
-
-        ctx.args = (causal, scale, mask, q_bucket_size, k_bucket_size)
-        ctx.save_for_backward(q, k, v, o, all_row_sums, all_row_maxes)
-
-        return o
-
-    @staticmethod
-    @torch.no_grad()
-    def backward(ctx, do):
-        """Algorithm 4 in the paper"""
-
-        causal, scale, mask, q_bucket_size, k_bucket_size = ctx.args
-        q, k, v, o, l, m = ctx.saved_tensors
-
-        device = q.device
-
-        max_neg_value = -torch.finfo(q.dtype).max
-        qk_len_diff = max(k.shape[-2] - q.shape[-2], 0)
-
-        dq = torch.zeros_like(q)
-        dk = torch.zeros_like(k)
-        dv = torch.zeros_like(v)
-
-        row_splits = zip(
-            q.split(q_bucket_size, dim=-2),
-            o.split(q_bucket_size, dim=-2),
-            do.split(q_bucket_size, dim=-2),
-            mask,
-            l.split(q_bucket_size, dim=-2),
-            m.split(q_bucket_size, dim=-2),
-            dq.split(q_bucket_size, dim=-2),
-        )
-
-        for ind, (qc, oc, doc, row_mask, lc, mc, dqc) in enumerate(row_splits):
-            q_start_index = ind * q_bucket_size - qk_len_diff
-
-            col_splits = zip(
-                k.split(k_bucket_size, dim=-2),
-                v.split(k_bucket_size, dim=-2),
-                dk.split(k_bucket_size, dim=-2),
-                dv.split(k_bucket_size, dim=-2),
-            )
-
-            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
-                )
-
-                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)
-                    attn_weights.masked_fill_(causal_mask, max_neg_value)
-
-                exp_attn_weights = torch.exp(attn_weights - mc)
-
-                if exists(row_mask):
-                    exp_attn_weights.masked_fill_(~row_mask, 0.0)
-
-                p = exp_attn_weights / lc
-
-                dv_chunk = torch.einsum("... i j, ... i d -> ... j d", p, doc)
-                dp = torch.einsum("... i d, ... j d -> ... i j", doc, vc)
-
-                D = (doc * oc).sum(dim=-1, keepdims=True)
-                ds = p * scale * (dp - D)
-
-                dq_chunk = torch.einsum("... i j, ... j d -> ... i d", ds, kc)
-                dk_chunk = torch.einsum("... i j, ... i d -> ... j d", ds, qc)
-
-                dqc.add_(dq_chunk)
-                dkc.add_(dk_chunk)
-                dvc.add_(dv_chunk)
-
-        return dq, dk, dv, None, None, None, None
-
-
-# endregion
-
-
-def get_parameter_dtype(parameter: torch.nn.Module):
-    return next(parameter.parameters()).dtype
-
-
-def get_parameter_device(parameter: torch.nn.Module):
-    return next(parameter.parameters()).device
-
-
-def get_timestep_embedding(
-    timesteps: torch.Tensor,
-    embedding_dim: int,
-    downscale_freq_shift: float = 1,
-    scale: float = 1,
-    max_period: int = 10000,
-):
-    """
-    This matches the implementation in Denoising Diffusion Probabilistic Models: Create sinusoidal timestep embeddings.
-
-    :param timesteps: a 1-D Tensor of N indices, one per batch element.
-                      These may be fractional.
-    :param embedding_dim: the dimension of the output. :param max_period: controls the minimum frequency of the
-    embeddings. :return: an [N x dim] Tensor of positional embeddings.
-    """
-    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 = exponent / (half_dim - downscale_freq_shift)
-
-    emb = torch.exp(exponent)
-    emb = timesteps[:, None].float() * emb[None, :]
-
-    # scale embeddings
-    emb = scale * emb
-
-    # concat sine and cosine embeddings: flipped from Diffusers original ver because always flip_sin_to_cos=True
-    emb = torch.cat([torch.cos(emb), torch.sin(emb)], dim=-1)
-
-    # zero pad
-    if embedding_dim % 2 == 1:
-        emb = torch.nn.functional.pad(emb, (0, 1, 0, 0))
-    return emb
-
-
-class GroupNorm32(nn.GroupNorm):
-    def forward(self, x):
-        if self.weight.dtype != torch.float32:
-            return super().forward(x)
-        return super().forward(x.float()).type(x.dtype)
-
-
-class ResnetBlock2D(nn.Module):
-    def __init__(
-        self,
-        in_channels,
-        out_channels,
-    ):
-        super().__init__()
-        self.in_channels = in_channels
-        self.out_channels = out_channels
-
-        self.in_layers = nn.Sequential(
-            GroupNorm32(32, in_channels),
-            nn.SiLU(),
-            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.out_layers = nn.Sequential(
-            GroupNorm32(32, out_channels),
-            nn.SiLU(),
-            nn.Identity(),  # to make state_dict compatible with original model
-            nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1),
-        )
-
-        if in_channels != out_channels:
-            self.skip_connection = nn.Conv2d(
-                in_channels, out_channels, kernel_size=1, stride=1, padding=0
-            )
-        else:
-            self.skip_connection = nn.Identity()
-
-        self.gradient_checkpointing = False
-
-    def forward_body(self, x, emb):
-        h = self.in_layers(x)
-        emb_out = self.emb_layers(emb).type(h.dtype)
-        h = h + emb_out[:, :, None, None]
-        h = self.out_layers(h)
-        x = self.skip_connection(x)
-        return x + h
-
-    def forward(self, x, emb):
-        if self.training and self.gradient_checkpointing:
-            # print("ResnetBlock2D: gradient_checkpointing")
-
-            def create_custom_forward(func):
-                def custom_forward(*inputs):
-                    return func(*inputs)
-
-                return custom_forward
-
-            x = torch.utils.checkpoint.checkpoint(
-                create_custom_forward(self.forward_body), x, emb
-            )
-        else:
-            x = self.forward_body(x, emb)
-
-        return x
-
-
-class Downsample2D(nn.Module):
-    def __init__(self, channels, out_channels):
-        super().__init__()
-
-        self.channels = channels
-        self.out_channels = out_channels
-
-        self.op = nn.Conv2d(self.channels, self.out_channels, 3, stride=2, padding=1)
-
-        self.gradient_checkpointing = False
-
-    def forward_body(self, hidden_states):
-        assert hidden_states.shape[1] == self.channels
-        hidden_states = self.op(hidden_states)
-
-        return hidden_states
-
-    def forward(self, hidden_states):
-        if self.training and self.gradient_checkpointing:
-            # print("Downsample2D: gradient_checkpointing")
-
-            def create_custom_forward(func):
-                def custom_forward(*inputs):
-                    return func(*inputs)
-
-                return custom_forward
-
-            hidden_states = torch.utils.checkpoint.checkpoint(
-                create_custom_forward(self.forward_body), hidden_states
-            )
-        else:
-            hidden_states = self.forward_body(hidden_states)
-
-        return hidden_states
-
-
-class CrossAttention(nn.Module):
-    def __init__(
-        self,
-        query_dim: int,
-        cross_attention_dim: Optional[int] = None,
-        heads: int = 8,
-        dim_head: int = 64,
-        upcast_attention: bool = False,
-    ):
-        super().__init__()
-        inner_dim = dim_head * heads
-        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
-        self.heads = heads
-
-        self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
-        self.to_k = nn.Linear(cross_attention_dim, inner_dim, bias=False)
-        self.to_v = nn.Linear(cross_attention_dim, inner_dim, bias=False)
-
-        self.to_out = nn.ModuleList([])
-        self.to_out.append(nn.Linear(inner_dim, query_dim))
-        # no dropout here
-
-        self.use_memory_efficient_attention_xformers = False
-        self.use_memory_efficient_attention_mem_eff = False
-        self.use_sdpa = False
-
-    def set_use_memory_efficient_attention(self, xformers, mem_eff):
-        self.use_memory_efficient_attention_xformers = xformers
-        self.use_memory_efficient_attention_mem_eff = mem_eff
-
-    def set_use_sdpa(self, sdpa):
-        self.use_sdpa = sdpa
-
-    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
-        )
-        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
-        )
-        return tensor
-
-    def forward(self, hidden_states, context=None, mask=None):
-        if self.use_memory_efficient_attention_xformers:
-            return self.forward_memory_efficient_xformers(hidden_states, context, mask)
-        if self.use_memory_efficient_attention_mem_eff:
-            return self.forward_memory_efficient_mem_eff(hidden_states, context, mask)
-        if self.use_sdpa:
-            return self.forward_sdpa(hidden_states, context, mask)
-
-        query = self.to_q(hidden_states)
-        context = context if context is not None else hidden_states
-        key = self.to_k(context)
-        value = self.to_v(context)
-
-        query = self.reshape_heads_to_batch_dim(query)
-        key = self.reshape_heads_to_batch_dim(key)
-        value = self.reshape_heads_to_batch_dim(value)
-
-        hidden_states = self._attention(query, key, value)
-
-        # linear proj
-        hidden_states = self.to_out[0](hidden_states)
-        # hidden_states = self.to_out[1](hidden_states)     # no dropout
-        return hidden_states
-
-    def _attention(self, query, key, value):
-        if self.upcast_attention:
-            query = query.float()
-            key = key.float()
-
-        attention_scores = torch.baddbmm(
-            torch.empty(
-                query.shape[0],
-                query.shape[1],
-                key.shape[1],
-                dtype=query.dtype,
-                device=query.device,
-            ),
-            query,
-            key.transpose(-1, -2),
-            beta=0,
-            alpha=self.scale,
-        )
-        attention_probs = attention_scores.softmax(dim=-1)
-
-        # cast back to the original dtype
-        attention_probs = attention_probs.to(value.dtype)
-
-        # compute attention output
-        hidden_states = torch.bmm(attention_probs, value)
-
-        # reshape hidden_states
-        hidden_states = self.reshape_batch_dim_to_heads(hidden_states)
-        return hidden_states
-
-    # TODO support Hypernetworks
-    def forward_memory_efficient_xformers(self, x, context=None, mask=None):
-        import xformers.ops
-
-        h = self.heads
-        q_in = self.to_q(x)
-        context = context if context is not None else x
-        context = context.to(x.dtype)
-        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)
-        )
-        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
-        )  # 最適なのを選んでくれる
-        del q, k, v
-
-        out = rearrange(out, "b n h d -> b n (h d)", h=h)
-
-        out = self.to_out[0](out)
-        return out
-
-    def forward_memory_efficient_mem_eff(self, x, context=None, mask=None):
-        flash_func = FlashAttentionFunction
-
-        q_bucket_size = 512
-        k_bucket_size = 1024
-
-        h = self.heads
-        q = self.to_q(x)
-        context = context if context is not None else x
-        context = context.to(x.dtype)
-        k = self.to_k(context)
-        v = self.to_v(context)
-        del context, x
-
-        q, k, v = map(lambda t: rearrange(t, "b n (h d) -> b h n d", h=h), (q, k, v))
-
-        out = flash_func.apply(q, k, v, mask, False, q_bucket_size, k_bucket_size)
-
-        out = rearrange(out, "b h n d -> b n (h d)")
-
-        out = self.to_out[0](out)
-        return out
-
-    def forward_sdpa(self, x, context=None, mask=None):
-        h = self.heads
-        q_in = self.to_q(x)
-        context = context if context is not None else x
-        context = context.to(x.dtype)
-        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)
-        )
-        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 = rearrange(out, "b h n d -> b n (h d)", h=h)
-
-        out = self.to_out[0](out)
-        return out
-
-
-# feedforward
-class GEGLU(nn.Module):
-    r"""
-    A variant of the gated linear unit activation function from https://arxiv.org/abs/2002.05202.
-
-    Parameters:
-        dim_in (`int`): The number of channels in the input.
-        dim_out (`int`): The number of channels in the output.
-    """
-
-    def __init__(self, dim_in: int, dim_out: int):
-        super().__init__()
-        self.proj = nn.Linear(dim_in, dim_out * 2)
-
-    def gelu(self, gate):
-        if gate.device.type != "mps":
-            return F.gelu(gate)
-        # mps: gelu is not implemented for float16
-        return F.gelu(gate.to(dtype=torch.float32)).to(dtype=gate.dtype)
-
-    def forward(self, hidden_states):
-        hidden_states, gate = self.proj(hidden_states).chunk(2, dim=-1)
-        return hidden_states * self.gelu(gate)
-
-
-class FeedForward(nn.Module):
-    def __init__(
-        self,
-        dim: int,
-    ):
-        super().__init__()
-        inner_dim = int(dim * 4)  # mult is always 4
-
-        self.net = nn.ModuleList([])
-        # project in
-        self.net.append(GEGLU(dim, inner_dim))
-        # project dropout
-        self.net.append(nn.Identity())  # nn.Dropout(0)) # dummy for dropout with 0
-        # project out
-        self.net.append(nn.Linear(inner_dim, dim))
-
-    def forward(self, hidden_states):
-        for module in self.net:
-            hidden_states = module(hidden_states)
-        return 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,
-    ):
-        super().__init__()
-
-        self.gradient_checkpointing = False
-
-        # 1. Self-Attn
-        self.attn1 = CrossAttention(
-            query_dim=dim,
-            cross_attention_dim=None,
-            heads=num_attention_heads,
-            dim_head=attention_head_dim,
-            upcast_attention=upcast_attention,
-        )
-        self.ff = FeedForward(dim)
-
-        # 2. Cross-Attn
-        self.attn2 = CrossAttention(
-            query_dim=dim,
-            cross_attention_dim=cross_attention_dim,
-            heads=num_attention_heads,
-            dim_head=attention_head_dim,
-            upcast_attention=upcast_attention,
-        )
-
-        self.norm1 = nn.LayerNorm(dim)
-        self.norm2 = nn.LayerNorm(dim)
-
-        # 3. Feed-forward
-        self.norm3 = nn.LayerNorm(dim)
-
-    def set_use_memory_efficient_attention(self, xformers: bool, mem_eff: bool):
-        self.attn1.set_use_memory_efficient_attention(xformers, mem_eff)
-        self.attn2.set_use_memory_efficient_attention(xformers, mem_eff)
-
-    def set_use_sdpa(self, sdpa: bool):
-        self.attn1.set_use_sdpa(sdpa)
-        self.attn2.set_use_sdpa(sdpa)
-
-    def forward_body(self, hidden_states, context=None, timestep=None):
-        # 1. Self-Attention
-        norm_hidden_states = self.norm1(hidden_states)
-
-        hidden_states = self.attn1(norm_hidden_states) + hidden_states
-
-        # 2. Cross-Attention
-        norm_hidden_states = self.norm2(hidden_states)
-        hidden_states = self.attn2(norm_hidden_states, context=context) + hidden_states
-
-        # 3. Feed-forward
-        hidden_states = self.ff(self.norm3(hidden_states)) + hidden_states
-
-        return hidden_states
-
-    def forward(self, hidden_states, context=None, timestep=None):
-        if self.training and self.gradient_checkpointing:
-            # print("BasicTransformerBlock: checkpointing")
-
-            def create_custom_forward(func):
-                def custom_forward(*inputs):
-                    return func(*inputs)
-
-                return custom_forward
-
-            output = torch.utils.checkpoint.checkpoint(
-                create_custom_forward(self.forward_body),
-                hidden_states,
-                context,
-                timestep,
-            )
-        else:
-            output = self.forward_body(hidden_states, context, timestep)
-
-        return output
-
-
-class Transformer2DModel(nn.Module):
-    def __init__(
-        self,
-        num_attention_heads: int = 16,
-        attention_head_dim: int = 88,
-        in_channels: Optional[int] = None,
-        cross_attention_dim: Optional[int] = None,
-        use_linear_projection: bool = False,
-        upcast_attention: bool = False,
-        num_transformer_layers: int = 1,
-    ):
-        super().__init__()
-        self.in_channels = in_channels
-        self.num_attention_heads = num_attention_heads
-        self.attention_head_dim = attention_head_dim
-        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 = 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
-            )
-
-        blocks = []
-        for _ in range(num_transformer_layers):
-            blocks.append(
-                BasicTransformerBlock(
-                    inner_dim,
-                    num_attention_heads,
-                    attention_head_dim,
-                    cross_attention_dim=cross_attention_dim,
-                    upcast_attention=upcast_attention,
-                )
-            )
-
-        self.transformer_blocks = nn.ModuleList(blocks)
-
-        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.gradient_checkpointing = False
-
-    def set_use_memory_efficient_attention(self, xformers, mem_eff):
-        for transformer in self.transformer_blocks:
-            transformer.set_use_memory_efficient_attention(xformers, mem_eff)
-
-    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):
-        # 1. Input
-        batch, _, height, weight = hidden_states.shape
-        residual = hidden_states
-
-        hidden_states = self.norm(hidden_states)
-        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
-            )
-        else:
-            inner_dim = hidden_states.shape[1]
-            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
-            )
-
-        # 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 = 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()
-            )
-
-        output = hidden_states + residual
-
-        return output
-
-
-class Upsample2D(nn.Module):
-    def __init__(self, channels, out_channels):
-        super().__init__()
-        self.channels = channels
-        self.out_channels = out_channels
-        self.conv = nn.Conv2d(self.channels, self.out_channels, 3, padding=1)
-
-        self.gradient_checkpointing = False
-
-    def forward_body(self, hidden_states, output_size=None):
-        assert hidden_states.shape[1] == self.channels
-
-        # Cast to float32 to as 'upsample_nearest2d_out_frame' op does not support bfloat16
-        # TODO(Suraj): Remove this cast once the issue is fixed in PyTorch
-        # https://github.com/pytorch/pytorch/issues/86679
-        dtype = hidden_states.dtype
-        if dtype == torch.bfloat16:
-            hidden_states = hidden_states.to(torch.float32)
-
-        # upsample_nearest_nhwc fails with large batch sizes. see https://github.com/huggingface/diffusers/issues/984
-        if hidden_states.shape[0] >= 64:
-            hidden_states = hidden_states.contiguous()
-
-        # 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"
-            )
-        else:
-            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:
-            hidden_states = hidden_states.to(dtype)
-
-        hidden_states = self.conv(hidden_states)
-
-        return hidden_states
-
-    def forward(self, hidden_states, output_size=None):
-        if self.training and self.gradient_checkpointing:
-            # print("Upsample2D: gradient_checkpointing")
-
-            def create_custom_forward(func):
-                def custom_forward(*inputs):
-                    return func(*inputs)
-
-                return custom_forward
-
-            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)
-
-        return hidden_states
-
-
-class SdxlUNet2DConditionModel(nn.Module):
-    _supports_gradient_checkpointing = True
-
-    def __init__(
-        self,
-        **kwargs,
-    ):
-        super().__init__()
-
-        self.in_channels = IN_CHANNELS
-        self.out_channels = OUT_CHANNELS
-        self.model_channels = MODEL_CHANNELS
-        self.time_embed_dim = TIME_EMBED_DIM
-        self.adm_in_channels = ADM_IN_CHANNELS
-
-        self.gradient_checkpointing = False
-        # self.sample_size = sample_size
-
-        # time embedding
-        self.time_embed = nn.Sequential(
-            nn.Linear(self.model_channels, self.time_embed_dim),
-            nn.SiLU(),
-            nn.Linear(self.time_embed_dim, self.time_embed_dim),
-        )
-
-        # label embedding
-        self.label_emb = nn.Sequential(
-            nn.Sequential(
-                nn.Linear(self.adm_in_channels, self.time_embed_dim),
-                nn.SiLU(),
-                nn.Linear(self.time_embed_dim, self.time_embed_dim),
-            )
-        )
-
-        # input
-        self.input_blocks = nn.ModuleList(
-            [
-                nn.Sequential(
-                    nn.Conv2d(
-                        self.in_channels,
-                        self.model_channels,
-                        kernel_size=3,
-                        padding=(1, 1),
-                    ),
-                )
-            ]
-        )
-
-        # level 0
-        for i in range(2):
-            layers = [
-                ResnetBlock2D(
-                    in_channels=1 * self.model_channels,
-                    out_channels=1 * self.model_channels,
-                ),
-            ]
-            self.input_blocks.append(nn.ModuleList(layers))
-
-        self.input_blocks.append(
-            nn.Sequential(
-                Downsample2D(
-                    channels=1 * self.model_channels,
-                    out_channels=1 * self.model_channels,
-                ),
-            )
-        )
-
-        # level 1
-        for i in range(2):
-            layers = [
-                ResnetBlock2D(
-                    in_channels=(1 if i == 0 else 2) * self.model_channels,
-                    out_channels=2 * self.model_channels,
-                ),
-                Transformer2DModel(
-                    num_attention_heads=2 * self.model_channels // 64,
-                    attention_head_dim=64,
-                    in_channels=2 * self.model_channels,
-                    num_transformer_layers=2,
-                    use_linear_projection=True,
-                    cross_attention_dim=2048,
-                ),
-            ]
-            self.input_blocks.append(nn.ModuleList(layers))
-
-        self.input_blocks.append(
-            nn.Sequential(
-                Downsample2D(
-                    channels=2 * self.model_channels,
-                    out_channels=2 * self.model_channels,
-                ),
-            )
-        )
-
-        # level 2
-        for i in range(2):
-            layers = [
-                ResnetBlock2D(
-                    in_channels=(2 if i == 0 else 4) * self.model_channels,
-                    out_channels=4 * self.model_channels,
-                ),
-                Transformer2DModel(
-                    num_attention_heads=4 * self.model_channels // 64,
-                    attention_head_dim=64,
-                    in_channels=4 * self.model_channels,
-                    num_transformer_layers=10,
-                    use_linear_projection=True,
-                    cross_attention_dim=2048,
-                ),
-            ]
-            self.input_blocks.append(nn.ModuleList(layers))
-
-        # mid
-        self.middle_block = nn.ModuleList(
-            [
-                ResnetBlock2D(
-                    in_channels=4 * self.model_channels,
-                    out_channels=4 * self.model_channels,
-                ),
-                Transformer2DModel(
-                    num_attention_heads=4 * self.model_channels // 64,
-                    attention_head_dim=64,
-                    in_channels=4 * self.model_channels,
-                    num_transformer_layers=10,
-                    use_linear_projection=True,
-                    cross_attention_dim=2048,
-                ),
-                ResnetBlock2D(
-                    in_channels=4 * self.model_channels,
-                    out_channels=4 * self.model_channels,
-                ),
-            ]
-        )
-
-        # output
-        self.output_blocks = nn.ModuleList([])
-
-        # level 2
-        for i in range(3):
-            layers = [
-                ResnetBlock2D(
-                    in_channels=4 * self.model_channels
-                    + (4 if i <= 1 else 2) * self.model_channels,
-                    out_channels=4 * self.model_channels,
-                ),
-                Transformer2DModel(
-                    num_attention_heads=4 * self.model_channels // 64,
-                    attention_head_dim=64,
-                    in_channels=4 * self.model_channels,
-                    num_transformer_layers=10,
-                    use_linear_projection=True,
-                    cross_attention_dim=2048,
-                ),
-            ]
-            if i == 2:
-                layers.append(
-                    Upsample2D(
-                        channels=4 * self.model_channels,
-                        out_channels=4 * self.model_channels,
-                    )
-                )
-
-            self.output_blocks.append(nn.ModuleList(layers))
-
-        # level 1
-        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,
-                    out_channels=2 * self.model_channels,
-                ),
-                Transformer2DModel(
-                    num_attention_heads=2 * self.model_channels // 64,
-                    attention_head_dim=64,
-                    in_channels=2 * self.model_channels,
-                    num_transformer_layers=2,
-                    use_linear_projection=True,
-                    cross_attention_dim=2048,
-                ),
-            ]
-            if i == 2:
-                layers.append(
-                    Upsample2D(
-                        channels=2 * self.model_channels,
-                        out_channels=2 * self.model_channels,
-                    )
-                )
-
-            self.output_blocks.append(nn.ModuleList(layers))
-
-        # level 0
-        for i in range(3):
-            layers = [
-                ResnetBlock2D(
-                    in_channels=1 * self.model_channels
-                    + (2 if i == 0 else 1) * self.model_channels,
-                    out_channels=1 * self.model_channels,
-                ),
-            ]
-
-            self.output_blocks.append(nn.ModuleList(layers))
-
-        # output
-        self.out = nn.ModuleList(
-            [
-                GroupNorm32(32, self.model_channels),
-                nn.SiLU(),
-                nn.Conv2d(self.model_channels, self.out_channels, 3, padding=1),
-            ]
-        )
-
-    # region diffusers compatibility
-    def prepare_config(self):
-        self.config = SimpleNamespace()
-
-    @property
-    def dtype(self) -> torch.dtype:
-        # `torch.dtype`: The dtype of the module (assuming that all the module parameters have the same dtype).
-        return get_parameter_dtype(self)
-
-    @property
-    def device(self) -> torch.device:
-        # `torch.device`: The device on which the module is (assuming that all the module parameters are on the same device).
-        return get_parameter_device(self)
-
-    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()
-        )
-
-    def enable_gradient_checkpointing(self):
-        self.gradient_checkpointing = True
-        self.set_gradient_checkpointing(value=True)
-
-    def disable_gradient_checkpointing(self):
-        self.gradient_checkpointing = False
-        self.set_gradient_checkpointing(value=False)
-
-    def set_use_memory_efficient_attention(self, xformers: bool, mem_eff: bool) -> None:
-        blocks = self.input_blocks + [self.middle_block] + self.output_blocks
-        for block in blocks:
-            for module in block:
-                if hasattr(module, "set_use_memory_efficient_attention"):
-                    # print(module.__class__.__name__)
-                    module.set_use_memory_efficient_attention(xformers, mem_eff)
-
-    def set_use_sdpa(self, sdpa: bool) -> None:
-        blocks = self.input_blocks + [self.middle_block] + self.output_blocks
-        for block in blocks:
-            for module in block:
-                if hasattr(module, "set_use_sdpa"):
-                    module.set_use_sdpa(sdpa)
-
-    def set_gradient_checkpointing(self, value=False):
-        blocks = self.input_blocks + [self.middle_block] + self.output_blocks
-        for block in blocks:
-            for module in block.modules():
-                if hasattr(module, "gradient_checkpointing"):
-                    # print(module.__class__.__name__, module.gradient_checkpointing, "->", value)
-                    module.gradient_checkpointing = value
-
-    # endregion
-
-    def forward(self, x, timesteps=None, context=None, y=None, **kwargs):
-        # broadcast timesteps to batch dimension
-        timesteps = timesteps.expand(x.shape[0])
-
-        hs = []
-        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.dtype == y.dtype, f"dtype mismatch: {x.dtype} != {y.dtype}"
-        # assert x.dtype == self.dtype
-        emb = emb + self.label_emb(y)
-
-        def call_module(module, h, emb, context):
-            x = h
-            for layer in module:
-                # print(layer.__class__.__name__, x.dtype, emb.dtype, context.dtype if context is not None else None)
-                if isinstance(layer, ResnetBlock2D):
-                    x = layer(x, emb)
-                elif isinstance(layer, Transformer2DModel):
-                    x = layer(x, context)
-                else:
-                    x = layer(x)
-            return x
-
-        # h = x.type(self.dtype)
-        h = x
-        for module in self.input_blocks:
-            h = call_module(module, h, emb, context)
-            hs.append(h)
-
-        h = call_module(self.middle_block, h, emb, context)
-
-        for module in self.output_blocks:
-            h = torch.cat([h, hs.pop()], dim=1)
-            h = call_module(module, h, emb, context)
-
-        h = h.type(x.dtype)
-        h = call_module(self.out, h, emb, context)
-
-        return h
-
-
-if __name__ == "__main__":
-    import time
-
-    print("create unet")
-    unet = SdxlUNet2DConditionModel()
-
-    unet.to("cuda")
-    unet.set_use_memory_efficient_attention(True, False)
-    unet.set_gradient_checkpointing(True)
-    unet.train()
-
-    # 使用メモリ量確認用の疑似学習ループ
-    print("preparing optimizer")
-
-    # optimizer = torch.optim.SGD(unet.parameters(), lr=1e-3, nesterov=True, momentum=0.9) # not working
-
-    # import bitsandbytes
-    # optimizer = bitsandbytes.adam.Adam8bit(unet.parameters(), lr=1e-3)        # not working
-    # optimizer = bitsandbytes.optim.RMSprop8bit(unet.parameters(), lr=1e-3)  # working at 23.5 GB with torch2
-    # optimizer=bitsandbytes.optim.Adagrad8bit(unet.parameters(), lr=1e-3)  # working at 23.5 GB with torch2
-
-    import transformers
-
-    optimizer = transformers.optimization.Adafactor(
-        unet.parameters(), relative_step=True
-    )  # working at 22.2GB with torch2
-
-    scaler = torch.cuda.amp.GradScaler(enabled=True)
-
-    print("start training")
-    steps = 10
-    batch_size = 1
-
-    for step in range(steps):
-        print(f"step {step}")
-        if step == 1:
-            time_start = time.perf_counter()
-
-        x = torch.randn(batch_size, 4, 128, 128).cuda()  # 1024x1024
-        t = torch.randint(low=0, high=10, size=(batch_size,), device="cuda")
-        ctx = torch.randn(batch_size, 77, 2048).cuda()
-        y = torch.randn(batch_size, ADM_IN_CHANNELS).cuda()
-
-        with torch.cuda.amp.autocast(enabled=True):
-            output = unet(x, t, ctx, y)
-            target = torch.randn_like(output)
-            loss = torch.nn.functional.mse_loss(output, target)
-
-        scaler.scale(loss).backward()
-        scaler.step(optimizer)
-        scaler.update()
-        optimizer.zero_grad(set_to_none=True)
-
-    time_end = time.perf_counter()
-    print(f"elapsed time: {time_end - time_start} [sec] for last {steps - 1} steps")
diff --git a/lycoris/kohya/utils.py b/lycoris/kohya/utils.py
deleted file mode 100644
index 892b93f..0000000
--- a/lycoris/kohya/utils.py
+++ /dev/null
@@ -1,48 +0,0 @@
-# part of https://github.com/kohya-ss/sd-scripts/blob/main/library/train_util.py
-
-import hashlib
-import safetensors
-from io import BytesIO
-
-
-def addnet_hash_legacy(b):
-    """Old model hash used by sd-webui-additional-networks for .safetensors format files"""
-    m = hashlib.sha256()
-
-    b.seek(0x100000)
-    m.update(b.read(0x10000))
-    return m.hexdigest()[0:8]
-
-
-def addnet_hash_safetensors(b):
-    """New model hash used by sd-webui-additional-networks for .safetensors format files"""
-    hash_sha256 = hashlib.sha256()
-    blksize = 1024 * 1024
-
-    b.seek(0)
-    header = b.read(8)
-    n = int.from_bytes(header, "little")
-
-    offset = n + 8
-    b.seek(offset)
-    for chunk in iter(lambda: b.read(blksize), b""):
-        hash_sha256.update(chunk)
-
-    return hash_sha256.hexdigest()
-
-
-def precalculate_safetensors_hashes(tensors, metadata):
-    """Precalculate the model hashes needed by sd-webui-additional-networks to
-    save time on indexing the model later."""
-
-    # Because writing user metadata to the file can change the result of
-    # sd_models.model_hash(), only retain the training metadata for purposes of
-    # calculating the hash, as they are meant to be immutable
-    metadata = {k: v for k, v in metadata.items() if k.startswith("ss_")}
-
-    bytes = safetensors.torch.save(tensors, metadata)
-    b = BytesIO(bytes)
-
-    model_hash = addnet_hash_safetensors(b)
-    legacy_hash = addnet_hash_legacy(b)
-    return model_hash, legacy_hash
diff --git a/lycoris/logging.py b/lycoris/logging.py
index 94c0ba6..6c51ccc 100644
--- a/lycoris/logging.py
+++ b/lycoris/logging.py
@@ -1,6 +1,7 @@
 import sys
 import copy
 import logging
+from functools import cache
 
 
 class ColoredFormatter(logging.Formatter):
@@ -34,3 +35,18 @@ def format(self, record):
         )
     )
     logger.addHandler(handler)
+
+
+@cache
+def info_once(msg):
+    logger.info(msg)
+
+
+@cache
+def warning_once(msg):
+    logger.warning(msg)
+
+
+@cache
+def error_once(msg):
+    logger.error(msg)
diff --git a/lycoris/modules/__init__.py b/lycoris/modules/__init__.py
index 85f5ec9..9d0c8bd 100644
--- a/lycoris/modules/__init__.py
+++ b/lycoris/modules/__init__.py
@@ -1,124 +1,46 @@
 import torch
+import torch.nn as nn
 
-from .full import FullModule
-from .norms import NormModule
+from .base import LycorisBaseModule
 from .locon import LoConModule
 from .loha import LohaModule
-from .lokr import LokrModule, factorization
+from .lokr import LokrModule
+from .full import FullModule
+from .norms import NormModule
+from .diag_oft import DiagOFTModule
+from .boft import ButterflyOFTModule
+from .glora import GLoRAModule
+from .dylora import DyLoraModule
 from .ia3 import IA3Module
 
+from ..functional.general import factorization
 
-@torch.no_grad()
-def make_module(lyco_type, params, lora_name, orig_module):
-    module = None
-    if lyco_type == "locon":
-        up, down, mid, alpha = params
-        module = LoConModule(
-            lora_name,
-            orig_module,
-            1,
-            down.size(0),
-            float(alpha),
-            use_tucker=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":
-        w1a, w1b, w2a, w2b, t1, t2, alpha = params
-        module = LohaModule(
-            lora_name,
-            orig_module,
-            1,
-            w1b.size(0),
-            float(alpha),
-            use_tucker=t1 is not None,
-        )
-        module.hada_w1_a.copy_(w1a)
-        module.hada_w1_b.copy_(w1b)
-        module.hada_w2_a.copy_(w2a)
-        module.hada_w2_b.copy_(w2b)
-        if t1 is not None:
-            module.hada_t1.copy_(t1)
-            module.hada_t2.copy_(t2)
-    elif lyco_type == "kron":
-        w1, w1a, w1b, w2, w2a, w2b, _, t2, alpha = params
-        if w1a is not None:
-            lora_dim = w1a.size(1)
-        elif w2a is not None:
-            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]
+MODULE_LIST = [
+    LoConModule,
+    LohaModule,
+    IA3Module,
+    LokrModule,
+    FullModule,
+    NormModule,
+    DiagOFTModule,
+    ButterflyOFTModule,
+    GLoRAModule,
+    DyLoraModule,
+]
 
-        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]
-        ):
-            factor = -1
-        else:
-            factor = max(w1.shape) if w1 is not None else max(w1a.size(0), w1b.size(1))
+def get_module(lyco_state_dict, lora_name):
+    for module in MODULE_LIST:
+        if module.algo_check(lyco_state_dict, lora_name):
+            return module, tuple(module.extract_state_dict(lyco_state_dict, lora_name))
+    return None, None
 
-        module = LokrModule(
-            lora_name,
-            orig_module,
-            1,
-            lora_dim,
-            float(alpha),
-            use_tucker=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)
-        else:
-            module.lokr_w1_a.copy_(w1a)
-            module.lokr_w1_b.copy_(w1b)
-        if w2 is not None:
-            module.lokr_w2.copy_(w2)
-        else:
-            module.lokr_w2_a.copy_(w2a)
-            module.lokr_w2_b.copy_(w2b)
-        if t2 is not None:
-            module.lokr_t2.copy_(t2)
-    elif lyco_type == "norm":
-        w_norm, b_norm = params
-        module = NormModule(
-            lora_name,
-            orig_module,
-            1,
-        )
-        module.w_norm.copy_(w_norm)
-        if b_norm is not None:
-            module.b_norm.copy_(b_norm)
-    elif lyco_type == "full":
-        diff, diff_b = params
-        module = FullModule(
-            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":
-        pass
-    else:
-        return None
 
+@torch.no_grad()
+def make_module(lyco_type: LycorisBaseModule, params, lora_name, orig_module):
+    try:
+        module = lyco_type.make_module_from_state_dict(lora_name, orig_module, *params)
+    except NotImplementedError:
+        module = None
     return module
diff --git a/lycoris/modules/attention.py b/lycoris/modules/attention.py
deleted file mode 100644
index c8f340d..0000000
--- a/lycoris/modules/attention.py
+++ /dev/null
@@ -1,259 +0,0 @@
-import math
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-from einops import rearrange, repeat
-
-from lycoris.utils import default
-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)
-
-
-class FeedForward(nn.Module):
-    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))
-        # nn.init.constant_(self.net[-1].weight, 0)
-        # nn.init.constant_(self.net[-1].bias, 0)
-
-    def forward(self, x):
-        return self.net(x)
-
-
-def vanilla_attention(q, k, v, mask, scale=None):
-    if scale is 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 (...)")
-        max_neg_value = -torch.finfo(scores.dtype).max
-        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)",
-        lambda x: (1, x, 1, 1),
-    ),
-    "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)",
-        lambda x: (1, x, 1, 1),
-    ),
-}
-ATTN_FUNCTION = {
-    "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,
-        context_ch=None,
-        heads=8,
-        head_ch=64,
-        self_cross=False,
-        double_attn=False,
-        single_kv_head=False,
-        attn_backend="torch",
-        cosine_attn=False,
-        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:
-            qk_head_ch = head_ch
-        q_ch = heads * qk_head_ch
-        k_ch = (1 if single_kv_head else heads) * qk_head_ch
-        v_ch = (1 if single_kv_head else heads) * head_ch
-        inner_ch = heads * head_ch
-        assert inner_ch == in_ch
-        use_context = context_ch is not None
-        context_ch = default(context_ch, in_ch)
-
-        if attn_backend == "xformers":
-            assert XFORMERS_AVAIL
-        if attn_backend == "torch":
-            assert torch.version.__version__ >= "2.0.0"
-
-        self.heads = heads
-        self.self_cross = self_cross
-        self.double_attn = double_attn
-        self.single_kv_head = single_kv_head
-        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))
-            )
-        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)
-            self.v = nn.Linear(in_ch, v_ch, bias=False)
-            self.ck = nn.Linear(context_ch, k_ch, bias=False)
-            self.cv = nn.Linear(context_ch, v_ch, bias=False)
-        else:
-            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)
-        else:
-            self.out = nn.Linear(inner_ch, in_ch)
-
-    def forward(self, x: torch.Tensor, context=None, mask=None):
-        # Input Projection
-        heads = self.heads
-        q = self.q(x)
-        ck = cv = None
-        if self.self_cross:
-            k = self.k(x)
-            v = self.v(x)
-            if context is not None:
-                ck = self.ck(context)
-                cv = self.cv(context)
-                if not self.double_attn:
-                    k = torch.concat([k, ck], dim=1)
-                    v = torch.concat([v, cv], dim=1)
-        else:
-            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 = rearrange(out, self.memory_layout[1], h=heads)
-        if self.double_attn:
-            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)
-
-
-class TransformerBlock(nn.Module):
-    def __init__(
-        self,
-        dim,
-        n_heads,
-        d_head,
-        context_dim=None,
-        gated_ff=True,
-        self_cross=False,
-        single_kv_head=False,
-        attn_backend="torch",
-        cosine_attn=False,
-        qk_head_ch=-1,
-        disable_self_attn=False,
-        single_attn=False,
-    ):
-        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,
-        )
-        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,
-            )
-        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
-        )
-        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
diff --git a/lycoris/modules/base.py b/lycoris/modules/base.py
index f3d87c2..f82467a 100644
--- a/lycoris/modules/base.py
+++ b/lycoris/modules/base.py
@@ -1,5 +1,11 @@
 from collections import OrderedDict
+
+import torch
 import torch.nn as nn
+import torch.nn.functional as F
+import torch.nn.utils.parametrize as parametrize
+
+from ..utils.bnb import QuantLinears, log_bypass
 
 
 class ModuleCustomSD(nn.Module):
@@ -53,3 +59,244 @@ def state_dict(self, *args, destination=None, prefix="", keep_vars=False):
             return super().state_dict(
                 *args, destination=destination, prefix=prefix, keep_vars=keep_vars
             )
+
+
+class LycorisBaseModule(ModuleCustomSD):
+    name: str
+    dtype_tensor: torch.Tensor
+    support_module = {}
+    weight_list = []
+    weight_list_det = []
+
+    def __init__(
+        self,
+        lora_name,
+        org_module: nn.Module,
+        multiplier=1.0,
+        dropout=0.0,
+        rank_dropout=0.0,
+        module_dropout=0.0,
+        rank_dropout_scale=False,
+        bypass_mode=False,
+        **kwargs,
+    ):
+        """if alpha == 0 or None, alpha is rank (no scaling)."""
+        super().__init__()
+        self.lora_name = lora_name
+        self.not_supported = False
+
+        self.module = type(org_module)
+        if isinstance(org_module, nn.Linear):
+            self.module_type = "linear"
+            self.shape = (org_module.out_features, org_module.in_features)
+            self.op = F.linear
+            self.dim = org_module.out_features
+            self.kw_dict = {}
+        elif isinstance(org_module, nn.Conv1d):
+            self.module_type = "conv1d"
+            self.shape = (
+                org_module.out_channels,
+                org_module.in_channels,
+                *org_module.kernel_size,
+            )
+            self.op = F.conv1d
+            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,
+            }
+        elif isinstance(org_module, nn.Conv2d):
+            self.module_type = "conv2d"
+            self.shape = (
+                org_module.out_channels,
+                org_module.in_channels,
+                *org_module.kernel_size,
+            )
+            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,
+            }
+        elif isinstance(org_module, nn.Conv3d):
+            self.module_type = "conv3d"
+            self.shape = (
+                org_module.out_channels,
+                org_module.in_channels,
+                *org_module.kernel_size,
+            )
+            self.op = F.conv3d
+            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,
+            }
+        elif isinstance(org_module, nn.LayerNorm):
+            self.module_type = "layernorm"
+            self.shape = tuple(org_module.normalized_shape)
+            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,
+            }
+        elif isinstance(org_module, nn.GroupNorm):
+            self.module_type = "groupnorm"
+            self.shape = (org_module.num_channels,)
+            self.op = F.group_norm
+            self.group_num = org_module.num_groups
+            self.dim = org_module.num_channels
+            self.kw_dict = {"num_groups": org_module.num_groups, "eps": org_module.eps}
+        else:
+            self.not_supported = True
+            self.module_type = "unknown"
+
+        self.register_buffer("dtype_tensor", torch.tensor(0.0), persistent=False)
+
+        self.is_bnb = False
+        if isinstance(org_module, QuantLinears):
+            if not bypass_mode:
+                log_bypass()
+            self.is_bnb = True
+            bypass_mode = True
+        self.bypass_mode = bypass_mode
+        self.dropout = dropout
+        self.rank_dropout = rank_dropout
+        self.rank_dropout_scale = rank_dropout_scale
+        self.module_dropout = module_dropout
+
+        ## Dropout things
+        # Since LoKr/LoHa/OFT/BOFT are hard to follow the rank_dropout definition from kohya
+        # We redefine the dropout procedure here.
+        # g(x) = WX + drop(Brank_drop(AX)) for LoCon(lora), bypass
+        # g(x) = WX + drop(ΔWX) for any algo except LoCon(lora), bypass
+        # g(x) = (W + Brank_drop(A))X for LoCon(lora), rebuid
+        # g(x) = (W + rank_drop(ΔW))X for any algo except LoCon(lora), rebuild
+        self.drop = nn.Identity() if dropout == 0.0 else nn.Dropout(dropout)
+        self.rank_drop = (
+            nn.Identity() if rank_dropout == 0.0 else nn.Dropout(rank_dropout)
+        )
+
+        self.multiplier = multiplier
+        self.org_forward = org_module.forward
+        self.org_module = [org_module]
+
+    @classmethod
+    def parametrize(cls, org_module, attr, *args, **kwargs):
+        from .full import FullModule
+
+        if cls is FullModule:
+            raise RuntimeError("FullModule cannot be used for parametrize.")
+        target_param = getattr(org_module, attr)
+        kwargs["bypass_mode"] = False
+        if target_param.dim() == 2:
+            proxy_module = nn.Linear(
+                target_param.shape[0], target_param.shape[1], bias=False
+            )
+            proxy_module.weight = target_param
+        elif target_param.dim() > 2:
+            module_type = [
+                None,
+                None,
+                None,
+                nn.Conv1d,
+                nn.Conv2d,
+                nn.Conv3d,
+                None,
+                None,
+            ][target_param.dim()]
+            proxy_module = module_type(
+                target_param.shape[0],
+                target_param.shape[1],
+                *target_param.shape[2:],
+                bias=False,
+            )
+            proxy_module.weight = target_param
+        module_obj = cls("", proxy_module, *args, **kwargs)
+        module_obj.forward = module_obj.parametrize_forward
+        module_obj.to(target_param)
+        parametrize.register_parametrization(org_module, attr, module_obj)
+        return module_obj
+
+    @classmethod
+    def algo_check(cls, state_dict, lora_name):
+        return any(f"{lora_name}.{k}" in state_dict for k in cls.weight_list_det)
+
+    @classmethod
+    def extract_state_dict(cls, state_dict, lora_name):
+        return [state_dict.get(f"{lora_name}.{k}", None) for k in cls.weight_list]
+
+    @classmethod
+    def make_module_from_state_dict(cls, lora_name, orig_module, *weights):
+        raise NotImplementedError
+
+    @property
+    def dtype(self):
+        return self.dtype_tensor.dtype
+
+    @property
+    def device(self):
+        return self.dtype_tensor.device
+
+    @property
+    def org_weight(self):
+        return self.org_module[0].weight
+
+    @org_weight.setter
+    def org_weight(self, value):
+        self.org_module[0].weight.data.copy_(value)
+
+    def apply_to(self, **kwargs):
+        if self.not_supported:
+            return
+        self.org_forward = self.org_module[0].forward
+        self.org_module[0].forward = self.forward
+
+    def restore(self):
+        if self.not_supported:
+            return
+        self.org_module[0].forward = self.org_forward
+
+    def merge_to(self, multiplier=1.0):
+        if self.not_supported:
+            return
+        weight, bias = self.get_merged_weight(
+            multiplier, self.org_weight.shape, self.org_weight.device
+        )
+        self.org_weight = weight.to(self.org_weight)
+        if bias is not None:
+            bias = bias.to(self.org_weight)
+            if self.org_module[0].bias is not None:
+                self.org_module[0].bias.data.copy_(bias)
+            else:
+                self.org_module[0].bias = nn.Parameter(bias)
+
+    def get_diff_weight(self, multiplier=1.0, shape=None, device=None):
+        raise NotImplementedError
+
+    def get_merged_weight(self, multiplier=1.0, shape=None, device=None):
+        raise NotImplementedError
+
+    @torch.no_grad()
+    def apply_max_norm(self, max_norm, device=None):
+        return None, None
+
+    def bypass_forward_diff(self, x, scale=1):
+        raise NotImplementedError
+
+    def bypass_forward(self, x, scale=1):
+        raise NotImplementedError
+
+    def parametrize_forward(self, x: torch.Tensor, *args, **kwargs):
+        return self.get_merged_weight(
+            multiplier=self.multiplier, shape=x.shape, device=x.device
+        )[0].to(x.dtype)
+
+    def forward(self, *args, **kwargs):
+        raise NotImplementedError
diff --git a/lycoris/modules/boft.py b/lycoris/modules/boft.py
index 759e3e1..e8c3353 100644
--- a/lycoris/modules/boft.py
+++ b/lycoris/modules/boft.py
@@ -7,9 +7,10 @@
 
 from einops import rearrange
 
-from .base import ModuleCustomSD
+from .base import LycorisBaseModule
+from ..functional import power2factorization
 from ..logging import logger
-from ..utils.bnb import LinearNF4, QuantLinears, log_bypass
+from ..utils.bnb import LinearNF4
 
 
 @cache
@@ -22,36 +23,31 @@ def log_butterfly_factorize(dim, factor, result):
 
 
 def butterfly_factor(dimension: int, factor: int = -1) -> tuple[int, int]:
-    """
-    m = 2k
-    n = 2**p
-    m*n = dim
-    """
-
-    # Find the first solution and check if it is even doable
-    m = n = 0
-    while m <= factor:
-        m += 2
-        while dimension % m != 0 and m < dimension:
-            m += 2
-        if m > factor:
-            break
-        if sum(int(i) for i in f"{dimension//m:b}") == 1:
-            n = dimension // m
+    m, n = power2factorization(dimension, factor)
 
     if n == 0:
         raise ValueError(
-            f"It is impossible to decompose {dimension} with factor {factor} under BOFT constrains."
+            f"It is impossible to decompose {dimension} with factor {factor} under BOFT constraints."
         )
 
-    log_butterfly_factorize(dimension, factor, (dimension // n, n))
-    return dimension // n, n
-
-
-class ButterflyOFTModule(ModuleCustomSD):
-    """
-    modifed from kohya-ss/sd-scripts/networks/lora:LoRAModule
-    """
+    log_butterfly_factorize(dimension, factor, (m, n))
+    return m, n
+
+
+class ButterflyOFTModule(LycorisBaseModule):
+    name = "boft"
+    support_module = {
+        "linear",
+        "conv1d",
+        "conv2d",
+        "conv3d",
+    }
+    weight_list = [
+        "oft_blocks",
+        "rescale",
+        "alpha",
+    ]
+    weight_list_det = ["oft_blocks"]
 
     def __init__(
         self,
@@ -66,36 +62,23 @@ def __init__(
         use_tucker=False,
         use_scalar=False,
         rank_dropout_scale=False,
-        constrain=0,
+        constraint=0,
         rescaled=False,
         bypass_mode=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
-
-        if isinstance(org_module, QuantLinears):
-            if not bypass_mode:
-                log_bypass()
-            bypass_mode = True
-        self.bypass_mode = bypass_mode
+        super().__init__(
+            lora_name,
+            org_module,
+            multiplier,
+            dropout,
+            rank_dropout,
+            module_dropout,
+            rank_dropout_scale,
+            bypass_mode,
+        )
+        if self.module_type not in self.support_module:
+            raise ValueError(f"{self.module_type} is not supported in BOFT algo.")
 
         out_dim = self.dim
         b, m_exp = butterfly_factor(out_dim, lora_dim)
@@ -106,8 +89,8 @@ def __init__(
         self.boft_m = sum(int(i) for i in f"{m_exp-1:b}") + 1
         # block_num > block_size
         self.rescaled = rescaled
-        self.constrain = constrain * out_dim
-        self.register_buffer("alpha", torch.tensor(constrain))
+        self.constraint = constraint * out_dim
+        self.register_buffer("alpha", torch.tensor(constraint))
         self.oft_blocks = nn.Parameter(
             torch.zeros(self.boft_m, self.block_num, self.block_size, self.block_size)
         )
@@ -116,28 +99,32 @@ def __init__(
                 torch.ones(out_dim, *(1 for _ in range(org_module.weight.dim() - 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]
-        self.org_forward = self.org_module[0].forward
+    @classmethod
+    def algo_check(cls, state_dict, lora_name):
+        if f"{lora_name}.oft_blocks" in state_dict:
+            oft_blocks = state_dict[f"{lora_name}.oft_blocks"]
+            if oft_blocks.ndim == 4:
+                return True
+        return False
+
+    @classmethod
+    def make_module_from_state_dict(
+        cls, lora_name, orig_module, oft_blocks, rescale, alpha
+    ):
+        m, n, s, _ = oft_blocks.shape
+        module = cls(
+            lora_name,
+            orig_module,
+            1,
+            lora_dim=n,
+            alpha=float(alpha),
+            rescale=rescale is not None,
+        )
+        return 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 restore(self):
-        self.org_module[0].forward = self.org_forward
-
-    def merge_to(self, multiplier=1.0):
-        weight = self.make_weight(scale=multiplier)
-        self.org_module[0].weight.data.copy_(weight)
+        return torch.eye(self.block_size, device=self.device)
 
     def get_r(self):
         I = self.I
@@ -145,35 +132,50 @@ def get_r(self):
         q = self.oft_blocks - self.oft_blocks.transpose(-1, -2)
         normed_q = q
         # Diag OFT style constrain
-        if self.constrain > 0:
+        if self.constraint > 0:
             q_norm = torch.norm(q) + 1e-8
-            if q_norm > self.constrain:
-                normed_q = q * self.constrain / q_norm
+            if q_norm > self.constraint:
+                normed_q = q * self.constraint / q_norm
         # use float() to prevent unsupported type
         r = (I + normed_q) @ (I - normed_q).float().inverse()
         return r
 
-    def make_weight(self, scale=1, device=None):
+    def make_weight(self, scale=1, device=None, diff=False):
         m = self.boft_m
         b = self.boft_b
         r_b = b // 2
         r = self.get_r()
-        inp = self.org_module[0].weight.to(device, dtype=r.dtype)
+        inp = org = self.org_weight.to(device, dtype=r.dtype)
 
         for i in range(m):
             bi = r[i]  # b_num, b_size, b_size
-            if i == 0:
-                # Apply multiplier/scale and rescale into first weight
+            if scale != 1:
                 bi = bi * scale + (1 - scale) * self.I
-                if self.rescaled:
-                    bi = bi * self.rescale
             inp = rearrange(inp, "(c g k) ... -> (c k g) ...", g=2, k=2**i * r_b)
             inp = rearrange(inp, "(d b) ... -> d b ...", b=b)
             inp = torch.einsum("b i j, b j ... -> b i ...", bi, inp)
             inp = rearrange(inp, "d b ... -> (d b) ...")
             inp = rearrange(inp, "(c k g) ... -> (c g k) ...", g=2, k=2**i * r_b)
 
-        return inp
+        if self.rescaled:
+            inp = inp * self.rescale
+
+        if diff:
+            inp = inp - org
+
+        return inp.to(self.oft_blocks.dtype)
+
+    def get_diff_weight(self, multiplier=1, shape=None, device=None):
+        diff = self.make_weight(scale=multiplier, device=device, diff=True)
+        if shape is not None:
+            diff = diff.view(shape)
+        return diff, None
+
+    def get_merged_weight(self, multiplier=1, shape=None, device=None):
+        diff = self.make_weight(scale=multiplier, device=device)
+        if shape is not None:
+            diff = diff.view(shape)
+        return diff, None
 
     @torch.no_grad()
     def apply_max_norm(self, max_norm, device=None):
@@ -188,29 +190,42 @@ def apply_max_norm(self, max_norm, device=None):
 
         return scaled, orig_norm * ratio
 
-    def bypass_forward(self, x, scale=1):
+    def _bypass_forward(self, x, scale=1, diff=False):
         m = self.boft_m
         b = self.boft_b
         r_b = b // 2
         r = self.get_r()
-        inp = self.org_forward(x)
+        inp = org = self.org_forward(x)
+        if self.op in {F.conv2d, F.conv1d, F.conv3d}:
+            inp = inp.transpose(1, -1)
 
         for i in range(m):
             bi = r[i]  # b_num, b_size, b_size
-            if i == 0:
-                # Apply multiplier/scale and rescale into first weight
+            if scale != 1:
                 bi = bi * scale + (1 - scale) * self.I
-                if self.rescaled:
-                    bi = bi * self.rescale
             inp = rearrange(inp, "... (c g k) ->... (c k g)", g=2, k=2**i * r_b)
             inp = rearrange(inp, "... (d b) -> ... d b", b=b)
             inp = torch.einsum("b i j, ... b j -> ... b i", bi, inp)
             inp = rearrange(inp, "... d b -> ... (d b)")
             inp = rearrange(inp, "... (c k g) -> ... (c g k)", g=2, k=2**i * r_b)
 
+        if self.rescaled:
+            inp = inp * self.rescale.transpose(0, -1)
+
+        if self.op in {F.conv2d, F.conv1d, F.conv3d}:
+            inp = inp.transpose(1, -1)
+
+        if diff:
+            inp = inp - org
         return inp
 
-    def forward(self, x):
+    def bypass_forward_diff(self, x, scale=1):
+        return self._bypass_forward(x, scale, diff=True)
+
+    def bypass_forward(self, x, scale=1):
+        return self._bypass_forward(x, scale, diff=False)
+
+    def forward(self, x, *args, **kwargs):
         if self.module_dropout and self.training:
             if torch.rand(1) < self.module_dropout:
                 return self.org_forward(x)
@@ -222,33 +237,3 @@ def forward(self, x):
             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)
-
-
-if __name__ == "__main__":
-    base = nn.Linear(128, 128).cuda()
-    lokr = ButterflyOFTModule("test", base, 1, 4, 1, weight_decompose=True).cuda()
-    print(lokr)
-    test_input = torch.randn(1, 128).cuda()
-    test_output = lokr(test_input)
-    print(test_output.shape)
-    print(F.mse_loss(test_output, lokr.bypass_forward(test_input)))
-
-    base_4bit = LinearNF4(128, 128)
-    base_4bit.load_state_dict(base.state_dict())
-    base_4bit.cuda()
-    qlocon = ButterflyOFTModule(
-        "test", base_4bit, 1, 4, 1, weight_decompose=False
-    ).cuda()
-    print(qlocon)
-    test_input = torch.randn(1, 128).cuda()
-    test_output = qlocon(test_input)
-    print(test_output.shape)
-
-    base = nn.Conv2d(128, 128, 3, 1, 1)
-    lokr = ButterflyOFTModule(
-        "test", base, 1, 4, 1, weight_decompose=True, use_tucker=True
-    )
-    print(lokr)
-    test_input = torch.randn(1, 128, 16, 16)
-    test_output = lokr(test_input)
-    print(test_output.shape)
diff --git a/lycoris/modules/diag_oft.py b/lycoris/modules/diag_oft.py
index 1fb8360..0b05547 100644
--- a/lycoris/modules/diag_oft.py
+++ b/lycoris/modules/diag_oft.py
@@ -4,12 +4,10 @@
 import torch.nn as nn
 import torch.nn.functional as F
 
-from einops import rearrange
-
-from .base import ModuleCustomSD
-from .lokr import factorization
+from .base import LycorisBaseModule
+from ..functional import factorization
 from ..logging import logger
-from ..utils.bnb import LinearNF4, QuantLinears, log_bypass
+from ..utils.bnb import LinearNF4
 
 
 @cache
@@ -21,10 +19,20 @@ def log_oft_factorize(dim, factor, num, bdim):
     )
 
 
-class DiagOFTModule(ModuleCustomSD):
-    """
-    modifed from kohya-ss/sd-scripts/networks/lora:LoRAModule
-    """
+class DiagOFTModule(LycorisBaseModule):
+    name = "diag-oft"
+    support_module = {
+        "linear",
+        "conv1d",
+        "conv2d",
+        "conv3d",
+    }
+    weight_list = [
+        "oft_blocks",
+        "rescale",
+        "alpha",
+    ]
+    weight_list_det = ["oft_blocks"]
 
     def __init__(
         self,
@@ -39,43 +47,30 @@ def __init__(
         use_tucker=False,
         use_scalar=False,
         rank_dropout_scale=False,
-        constrain=0,
+        constraint=0,
         rescaled=False,
         bypass_mode=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
-
-        if isinstance(org_module, QuantLinears):
-            if not bypass_mode:
-                log_bypass()
-            bypass_mode = True
-        self.bypass_mode = bypass_mode
+        super().__init__(
+            lora_name,
+            org_module,
+            multiplier,
+            dropout,
+            rank_dropout,
+            module_dropout,
+            rank_dropout_scale,
+            bypass_mode,
+        )
+        if self.module_type not in self.support_module:
+            raise ValueError(f"{self.module_type} is not supported in Diag-OFT algo.")
 
         out_dim = self.dim
         self.block_size, self.block_num = factorization(out_dim, lora_dim)
         # block_num > block_size
         self.rescaled = rescaled
-        self.constrain = constrain * out_dim
-        self.register_buffer("alpha", torch.tensor(constrain))
+        self.constraint = constraint * out_dim
+        self.register_buffer("alpha", torch.tensor(constraint))
         self.oft_blocks = nn.Parameter(
             torch.zeros(self.block_num, self.block_size, self.block_size)
         )
@@ -91,61 +86,74 @@ def __init__(
             bdim=self.block_size,
         )
 
-        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]
-        self.org_forward = self.org_module[0].forward
+    @classmethod
+    def algo_check(cls, state_dict, lora_name):
+        if f"{lora_name}.oft_blocks" in state_dict:
+            oft_blocks = state_dict[f"{lora_name}.oft_blocks"]
+            if oft_blocks.ndim == 3:
+                return True
+        return False
+
+    @classmethod
+    def make_module_from_state_dict(
+        cls, lora_name, orig_module, oft_blocks, rescale, alpha
+    ):
+        n, s, _ = oft_blocks.shape
+        module = cls(
+            lora_name,
+            orig_module,
+            1,
+            lora_dim=s,
+            alpha=float(alpha),
+            rescale=rescale is not None,
+        )
+        return 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 restore(self):
-        self.org_module[0].forward = self.org_forward
-
-    def merge_to(self, multiplier=1.0):
-        weight = self.make_weight(scale=multiplier)
-        self.org_module[0].weight.data.copy_(weight)
+        return torch.eye(self.block_size, device=self.device)
 
     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:
+        if self.constraint > 0:
             q_norm = torch.norm(q) + 1e-8
-            if q_norm > self.constrain:
-                normed_q = q * self.constrain / q_norm
+            if q_norm > self.constraint:
+                normed_q = q * self.constraint / q_norm
         # use float() to prevent unsupported type
         r = (I + normed_q) @ (I - normed_q).float().inverse()
         return r
 
-    def make_weight(self, scale=1, device=None):
+    def make_weight(self, scale=1, device=None, diff=False):
         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,
-        )
+        _, *shape = self.org_weight.shape
+        org_weight = self.org_weight.to(device, dtype=r.dtype)
+        org_weight = org_weight.view(self.block_num, self.block_size, *shape)
         # 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,
+            self.rank_drop(r * scale) - scale * self.I + (0 if diff else self.I),
             org_weight,
-        )
-        weight = rearrange(weight, "k m ... -> (k m) ...")
+        ).view(-1, *shape)
         if self.rescaled:
             weight = self.rescale * weight
-        return weight
+            if diff:
+                weight = weight + (self.rescale - 1) * org_weight
+        return weight.to(self.oft_blocks.dtype)
+
+    def get_diff_weight(self, multiplier=1, shape=None, device=None):
+        diff = self.make_weight(scale=multiplier, device=device, diff=True)
+        if shape is not None:
+            diff = diff.view(shape)
+        return diff, None
+
+    def get_merged_weight(self, multiplier=1, shape=None, device=None):
+        diff = self.make_weight(scale=multiplier, device=device)
+        if shape is not None:
+            diff = diff.view(shape)
+        return diff, None
 
     @torch.no_grad()
     def apply_max_norm(self, max_norm, device=None):
@@ -160,26 +168,40 @@ def apply_max_norm(self, max_norm, device=None):
 
         return scaled, orig_norm * ratio
 
-    def bypass_forward(self, x, scale=1):
+    def _bypass_forward(self, x, scale=1, diff=False):
         r = self.get_r()
         org_out = self.org_forward(x)
-        if self.op == F.conv2d:
+        if self.op in {F.conv2d, F.conv1d, F.conv3d}:
             org_out = org_out.transpose(1, -1)
-        org_out = rearrange(
-            org_out,
-            "... (k n) ->... k n",
-            k=self.block_num,
-            n=self.block_size,
-        )
+        *shape, _ = org_out.shape
+        org_out = org_out.view(*shape, self.block_num, self.block_size)
+        mask = neg_mask = 1
+        if self.dropout != 0 and self.training:
+            mask = torch.ones_like(org_out)
+            mask = self.drop(mask)
+            neg_mask = torch.max(mask) - mask
         oft_out = torch.einsum(
-            "k n m, ... k n -> ... k m", r * scale + (1 - scale) * self.I, org_out
+            "k n m, ... k n -> ... k m",
+            r * scale * mask + (1 - scale) * self.I * neg_mask,
+            org_out,
         )
-        out = rearrange(oft_out, "... k m -> ... (k m)")
-        if self.op == F.conv2d:
+        if diff:
+            out = out - org_out
+        out = oft_out.view(*shape, -1)
+        if self.rescaled:
+            out = self.rescale.transpose(-1, 0) * out
+            out = out + (self.rescale.transpose(-1, 0) - 1) * org_out
+        if self.op in {F.conv2d, F.conv1d, F.conv3d}:
             out = out.transpose(1, -1)
         return out
 
-    def forward(self, x):
+    def bypass_forward_diff(self, x, scale=1):
+        return self._bypass_forward(x, scale, diff=True)
+
+    def bypass_forward(self, x, scale=1):
+        return self._bypass_forward(x, scale, diff=False)
+
+    def forward(self, x: torch.Tensor, *args, **kwargs):
         if self.module_dropout and self.training:
             if torch.rand(1) < self.module_dropout:
                 return self.org_forward(x)
@@ -191,28 +213,3 @@ def forward(self, x):
             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)
-
-
-if __name__ == "__main__":
-    base = nn.Linear(128, 128).cuda()
-    lokr = DiagOFTModule("test", base, 1, 4, 1, weight_decompose=True).cuda()
-    print(lokr)
-    test_input = torch.randn(1, 128).cuda()
-    test_output = lokr(test_input)
-    print(test_output.shape)
-
-    base_4bit = LinearNF4(128, 128)
-    base_4bit.load_state_dict(base.state_dict())
-    base_4bit.cuda()
-    qlocon = DiagOFTModule("test", base_4bit, 1, 4, 1, weight_decompose=False).cuda()
-    print(qlocon)
-    test_input = torch.randn(1, 128).cuda()
-    test_output = qlocon(test_input)
-    print(test_output.shape)
-
-    base = nn.Conv2d(128, 128, 3, 1, 1)
-    lokr = DiagOFTModule("test", base, 1, 4, 1, weight_decompose=True, use_tucker=True)
-    print(lokr)
-    test_input = torch.randn(1, 128, 16, 16)
-    test_output = lokr(test_input)
-    print(test_output.shape)
diff --git a/lycoris/modules/dylora.py b/lycoris/modules/dylora.py
index 7d2a5f0..f785c25 100644
--- a/lycoris/modules/dylora.py
+++ b/lycoris/modules/dylora.py
@@ -3,15 +3,19 @@
 
 import torch
 import torch.nn as nn
-import torch.nn.functional as F
 
-from .base import ModuleCustomSD
+from .base import LycorisBaseModule
+from ..utils import product
+from ..utils.bnb import LinearNF4
 
 
-class DyLoraModule(ModuleCustomSD):
-    """
-    Hadamard product Implementaion for Dynamic Low Rank adaptation
-    """
+class DyLoraModule(LycorisBaseModule):
+    support_module = {
+        "linear",
+        "conv1d",
+        "conv2d",
+        "conv3d",
+    }
 
     def __init__(
         self,
@@ -24,48 +28,45 @@ def __init__(
         rank_dropout=0.0,
         module_dropout=0.0,
         use_tucker=False,
-        block_size=1,
+        block_size=4,
+        use_scalar=False,
+        rank_dropout_scale=False,
+        weight_decompose=False,
+        bypass_mode=False,
+        rs_lora=False,
+        train_on_input=False,
         **kwargs,
     ):
         """if alpha == 0 or None, alpha is rank (no scaling)."""
-        super().__init__()
-        self.lora_name = lora_name
-        self.lora_dim = lora_dim
+        super().__init__(
+            lora_name,
+            org_module,
+            multiplier,
+            dropout,
+            rank_dropout,
+            module_dropout,
+            rank_dropout_scale,
+            bypass_mode,
+        )
+        if self.module_type not in self.support_module:
+            raise ValueError(f"{self.module_type} is not supported in IA^3 algo.")
         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":
-            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])
-            self.op = F.conv2d
-            self.extra_args = {
-                "stride": org_module.stride,
-                "padding": org_module.padding,
-                "dilation": org_module.dilation,
-                "groups": org_module.groups,
-            }
-        else:
-            in_dim = org_module.in_features
-            out_dim = org_module.out_features
-            shape = (out_dim, in_dim)
-            self.op = F.linear
-            self.extra_args = {}
+        shape = (
+            self.shape[0],
+            product(self.shape[1:]),
+        )
 
         self.lora_dim = lora_dim
         self.up_list = nn.ParameterList(
-            [torch.empty(shape[0], 1) for i in range(lora_dim)]
+            [torch.empty(shape[0], self.block_size) for i in range(self.block_count)]
         )
-
         self.down_list = nn.ParameterList(
-            [torch.empty(1, shape[1]) for i in range(lora_dim)]
+            [torch.empty(self.block_size, shape[1]) for i in range(self.block_count)]
         )
 
-        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
@@ -73,17 +74,13 @@ def __init__(
         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_forward = self.org_module[0].forward
-        self.grad_ckpt = False
-        self.state_dict()
+    def load_state_dict(self, state_dict, strict: bool = True, assign: bool = False):
+        return
 
     def custom_state_dict(self):
         destination = {}
@@ -92,41 +89,69 @@ def custom_state_dict(self):
             torch.concat(list(self.up_list), dim=1)
         )
         destination["lora_down.weight"] = nn.Parameter(
-            torch.concat(list(self.down_list))
+            torch.concat(list(self.down_list)).reshape(
+                self.lora_dim, -1, *self.shape[2:]
+            )
         )
         return destination
 
-    def apply_to(self):
-        self.org_module[0].forward = self.forward
-
-    def restore(self):
-        self.org_module[0].forward = self.org_forward
-
-    def merge_to(self, multiplier=1.0):
-        up = torch.concat(list(self.up_list), dim=1)
-        down = torch.concat(list(self.down_list))
-        diff_weight = up @ down * self.scale * multiplier / up.size(1)
-        self.org_module[0].weight.data.add_(diff_weight)
-
-    @torch.enable_grad()
-    def forward(self, x):
-        b = random.randint(0, self.block_count - 1)
-
+    def get_weight(self, rank):
+        b = math.ceil(rank / self.block_size)
         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]) + list(self.down_list[b : (b + 1)])
         )
         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]),
+            list(i.data for i in self.up_list[:b]) + list(self.up_list[b : (b + 1)]),
             dim=1,
         )
+        return down, up, self.alpha / (b + 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),
-            bias,
-            **self.extra_args,
-        )
+    def get_random_rank_weight(self):
+        b = random.randint(0, self.block_count - 1)
+        return self.get_weight(b * self.block_size)
+
+    def get_diff_weight(self, multiplier=1, shape=None, device=None, rank=None):
+        if rank is None:
+            down, up, scale = self.get_random_rank_weight()
+        else:
+            down, up, scale = self.get_random_rank_weight()
+        w = up @ (down * (scale * multiplier))
+        if device is not None:
+            w = w.to(device)
+        if shape is not None:
+            w = w.view(shape)
+        else:
+            w = w.view(self.shape)
+        return w, None
+
+    def get_merged_weight(self, multiplier=1, shape=None, device=None, rank=None):
+        diff, _ = self.get_diff_weight(multiplier, shape, device, rank)
+        return diff + self.org_weight, None
+
+    def bypass_forward_diff(self, x, scale=1, rank=None):
+        if rank is None:
+            down, up, gamma = self.get_random_rank_weight()
+        else:
+            down, up, gamma = self.get_random_rank_weight()
+        down = down.view(self.lora_dim, -1, *self.shape[2:])
+        up = up.view(-1, self.lora_dim, *(1 for _ in self.shape[2:]))
+        scale = scale * gamma
+        return self.op(self.op(x, down, **self.kw_dict), up)
+
+    def bypass_forward(self, x, scale=1, rank=None):
+        return self.org_forward(x) + self.bypass_forward_diff(x, scale, rank)
+
+    def forward(self, x, *args, **kwargs):
+        if self.module_dropout and self.training:
+            if torch.rand(1) < self.module_dropout:
+                return self.org_forward(x)
+        if self.bypass_mode:
+            return self.bypass_forward(x, self.multiplier)
+        else:
+            weight = self.get_merged_weight(multiplier=self.multiplier)[0]
+            bias = (
+                None
+                if self.org_module[0].bias is None
+                else self.org_module[0].bias.data
+            )
+            return self.op(x, weight, bias, **self.kw_dict)
diff --git a/lycoris/modules/full.py b/lycoris/modules/full.py
index 164c303..896d5f8 100644
--- a/lycoris/modules/full.py
+++ b/lycoris/modules/full.py
@@ -1,11 +1,20 @@
 import torch
 import torch.nn as nn
-import torch.nn.functional as F
 
-from .base import ModuleCustomSD
+from .base import LycorisBaseModule
 
 
-class FullModule(ModuleCustomSD):
+class FullModule(LycorisBaseModule):
+    name = "full"
+    support_module = {
+        "linear",
+        "conv1d",
+        "conv2d",
+        "conv3d",
+    }
+    weight_list = ["diff", "diff_b"]
+    weight_list_det = ["diff"]
+
     def __init__(
         self,
         lora_name,
@@ -19,43 +28,64 @@ def __init__(
         use_tucker=False,
         use_scalar=False,
         rank_dropout_scale=False,
+        bypass_mode=False,
         **kwargs,
     ):
-        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
+        super().__init__(
+            lora_name,
+            org_module,
+            multiplier,
+            dropout,
+            rank_dropout,
+            module_dropout,
+            rank_dropout_scale,
+            bypass_mode,
+        )
+        if self.module_type not in self.support_module:
+            raise ValueError(f"{self.module_type} is not supported in Full algo.")
+
+        if self.is_bnb:
+            raise ValueError(
+                "Quant Linear is not supported and meaningless in Full algo."
+            )
+
+        if self.bypass_mode:
+            raise ValueError("bypass mode is not supported in Full algo.")
+
         self.weight = nn.Parameter(torch.zeros_like(org_module.weight))
         if org_module.bias is not None:
             self.bias = nn.Parameter(torch.zeros_like(org_module.bias))
         else:
             self.bias = None
 
-        self.rank_dropout = rank_dropout
-        self.module_dropout = module_dropout
+    @classmethod
+    def make_module_from_state_dict(cls, lora_name, orig_module, diff, diff_b):
+        module = cls(
+            lora_name,
+            orig_module,
+            1,
+        )
+        module.weight.copy_(diff)
+        if diff_b is not None:
+            if orig_module.bias is not None:
+                module.bias.copy_(diff_b)
+            else:
+                module.bias = nn.Parameter(diff_b)
+        return module
+
+    @property
+    def org_weight(self):
+        return self._org_weight[0]
 
-        self.multiplier = multiplier
-        self.org_module = [org_module]
+    @org_weight.setter
+    def org_weight(self, value):
+        self.org_module[0].weight.data.copy_(value)
 
     def apply_to(self, **kwargs):
         self.org_forward = self.org_module[0].forward
         self.org_module[0].forward = self.forward
         self.weight.data.add_(self.org_module[0].weight.data)
-        self.org_weight = [self.org_module[0].weight.data.cpu().clone()]
+        self._org_weight = [self.org_module[0].weight.data.cpu().clone()]
         delattr(self.org_module[0], "weight")
         if self.org_module[0].bias is not None:
             self.bias.data.add_(self.org_module[0].bias.data)
@@ -64,14 +94,14 @@ def apply_to(self, **kwargs):
         else:
             self.org_bias = None
 
-    def merge_to(self, multiplier=1.0):
-        weight, bias = self.make_weight(scale=multiplier)
-        self.org_module[0].weight.data.copy_(weight)
-        if bias is not None:
-            self.org_module[0].bias.data.copy_(bias)
+    def restore(self):
+        self.org_module[0].forward = self.org_forward
+        self.org_module[0].weight = nn.Parameter(self._org_weight[0])
+        if self.org_bias is not None:
+            self.org_module[0].bias = nn.Parameter(self.org_bias[0])
 
     def custom_state_dict(self):
-        sd = {"diff": self.weight.data.cpu() - self.org_weight[0]}
+        sd = {"diff": self.weight.data.cpu() - self._org_weight[0]}
         if self.bias is not None:
             sd["diff_b"] = self.bias.data.cpu() - self.org_bias[0]
         return sd
@@ -86,39 +116,55 @@ def load_weight_prehook(
         unexpected_keys,
         error_msgs,
     ):
-        diff_weight = state_dict["diff"]
-        self.weight.data.add_(diff_weight)
-        if "diff_b" in state_dict:
-            diff_bias = state_dict["diff_b"]
-            self.bias.data.add_(diff_bias)
+        diff_weight = state_dict.pop(f"{prefix}diff")
+        state_dict[f"{prefix}weight"] = diff_weight + self.weight.data.to(diff_weight)
+        if f"{prefix}diff_b" in state_dict:
+            diff_bias = state_dict.pop(f"{prefix}diff_b")
+            state_dict[f"{prefix}bias"] = diff_bias + self.bias.data.to(diff_bias)
 
-    def make_weight(self, scale=1, device=None, original=False):
-        if original:
-            weight = self.org_weight[0].to(device, dtype=self.weight.dtype)
-            if self.org_bias is not None:
-                bias = self.org_bias[0].to(device, dtype=self.bias.dtype)
-            else:
-                bias = None
-            return weight, bias
+    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
             else 1
         )
         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
+            diff_w, diff_b = self.get_diff_weight(scale, device=device)
+            weight = self.org_module[0].weight + diff_w * drop
+            bias = self.org_module[0].bias + diff_b * drop
         else:
             weight = self.weight
             bias = self.bias
         return weight, bias
 
-    def forward(self, x: torch.Tensor):
+    def get_diff_weight(self, multiplier=1, shape=None, device=None):
+        org_weight = self.org_module[0].weight.to(device, dtype=self.weight.dtype)
+        diff = self.weight.to(device) - org_weight
+        diff_b = None
+        if shape:
+            diff = diff.view(shape)
+        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
+        if device is not None:
+            diff = diff.to(device)
+            if self.bias is not None:
+                diff_b = diff_b.to(device)
+        if multiplier != 1:
+            diff = diff * multiplier
+            if diff_b is not None:
+                diff_b = diff_b * multiplier
+        return diff * multiplier, diff_b
+
+    def get_merged_weight(self, multiplier=1, shape=None, device=None):
+        weight, bias = self.make_weight(multiplier, device)
+        if shape is not None:
+            weight = weight.view(shape)
+            if bias is not None:
+                bias = bias.view(shape[0])
+        return weight, bias
+
+    def forward(self, x: torch.Tensor, *args, **kwargs):
         if (
             self.module_dropout
             and self.training
@@ -127,7 +173,9 @@ def forward(self, x: torch.Tensor):
             original = True
         else:
             original = False
+        if original:
+            return self.org_forward(x)
         scale = self.multiplier
-        weight, bias = self.make_weight(scale, x.device, original=original)
+        weight, bias = self.make_weight(scale, x.device)
         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
deleted file mode 100644
index d03f343..0000000
--- a/lycoris/modules/glokr.py
+++ /dev/null
@@ -1,300 +0,0 @@
-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]:
-    """
-    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               ...
-    127 -> 127, 1   127 -> 127, 1    127 -> 127, 1   127 -> 127, 1   127 -> 127, 1
-    128 -> 16, 8    128 -> 64, 2     128 -> 32, 4    128 -> 16, 8    128 -> 16, 8
-    250 -> 25, 10   250 -> 125, 2    250 -> 125, 2   250 -> 50, 5   250 -> 25, 10
-    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
-        return m, n
-    if factor == -1:
-        factor = dimension
-    m, n = 1, dimension
-    length = m + n
-    while m < n:
-        new_m = m + 1
-        while dimension % new_m != 0:
-            new_m += 1
-        new_n = dimension // new_m
-        if new_m + new_n > length or new_m > factor:
-            break
-        else:
-            m, n = new_m, new_n
-    if m > n:
-        n, m = m, n
-    return m, n
-
-
-def make_weight_tucker(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]
-    return rebuild2
-
-
-def make_kron(w1, w2, scale):
-    if len(w2.shape) == 4:
-        w1 = w1.unsqueeze(2).unsqueeze(2)
-    w2 = w2.contiguous()
-    rebuild = torch.kron(w1, w2)
-
-    return rebuild * scale
-
-
-class LokrModule(ModuleCustomSD):
-    """
-    modifed from kohya-ss/sd-scripts/networks/lora:LoRAModule
-        and from KohakuBlueleaf/LyCORIS/lycoris:loha:LoHaModule
-        and from KohakuBlueleaf/LyCORIS/lycoris:locon:LoconModule
-    """
-
-    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,
-        decompose_both=False,
-        factor: int = -1,  # factorization factor
-        rank_dropout_scale=False,
-        **kwargs,
-    ):
-        """if alpha == 0 or None, alpha is rank (no scaling)."""
-        super().__init__()
-        factor = int(factor)
-        self.lora_name = lora_name
-        self.lora_dim = lora_dim
-        self.tucker = False
-        self.use_w1 = False
-        self.use_w2 = False
-
-        self.shape = org_module.weight.shape
-        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:
-                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.use_w2 = True
-                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
-                # 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])
-                )
-                # 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,
-            }
-
-        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
-
-            # smaller part. weight scale
-            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:
-                # 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]))
-                # 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(shape[0][1], shape[1][1]))
-
-            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
-            alpha = lora_dim
-        self.scale = alpha / self.lora_dim
-        self.register_buffer("alpha", torch.tensor(alpha))  # 定数として扱える
-
-        if self.use_w2:
-            torch.nn.init.constant_(self.lokr_w2, 0)
-        else:
-            if self.tucker:
-                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:
-            torch.nn.init.normal_(self.lokr_w1_a, std=1)
-            torch.nn.init.normal_(self.lokr_w1_b, std=0.1)
-
-        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_tucker(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"
-
-    # Same as locon.py
-    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):
-        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_tucker(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)
-            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)
-        desired = torch.clamp(norm, max=max_norm)
-        ratio = desired.cpu() / norm.cpu()
-
-        scaled = norm != desired
-        if scaled:
-            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)
-            else:
-                self.lokr_w1_a *= ratio ** (1 / modules)
-                self.lokr_w1_b *= ratio ** (1 / modules)
-
-            if self.use_w2:
-                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
-
-    def forward(self, x):
-        if self.module_dropout and self.training:
-            if torch.rand(1) < self.module_dropout:
-                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
-                    ),
-                )
-        weight = (
-            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)
diff --git a/lycoris/modules/glora.py b/lycoris/modules/glora.py
index 0b3994e..620af85 100644
--- a/lycoris/modules/glora.py
+++ b/lycoris/modules/glora.py
@@ -2,14 +2,30 @@
 
 import torch
 import torch.nn as nn
+import torch.nn.functional as F
 
-from .base import ModuleCustomSD
+from .base import LycorisBaseModule
+from ..functional import tucker_weight_from_conv
+from ..utils.bnb import LinearNF4
 
 
-class GLoRAModule(ModuleCustomSD):
-    """
-    modifed from kohya-ss/sd-scripts/networks/lora:LoRAModule
-    """
+class GLoRAModule(LycorisBaseModule):
+    name = "glora"
+    support_module = {
+        "linear",
+        "conv1d",
+        "conv2d",
+        "conv3d",
+    }
+    weight_list = [
+        "a1.weight",
+        "a2.weight",
+        "b1.weight",
+        "b2.weight",
+        "bm.weight",
+        "alpha",
+    ]
+    weight_list_det = ["a1.weight"]
 
     def __init__(
         self,
@@ -21,97 +37,191 @@ def __init__(
         dropout=0.0,
         rank_dropout=0.0,
         module_dropout=0.0,
+        use_tucker=False,
+        use_scalar=False,
         rank_dropout_scale=False,
-        *args,
+        weight_decompose=False,
+        bypass_mode=False,
+        rs_lora=False,
         **kwargs,
     ):
-        """if alpha == 0 or None, alpha is rank (no scaling)."""
-        super().__init__()
-        self.lora_name = lora_name
+        """
+        f(x) = WX + WAX + BX, where A and B are low-rank matrices
+        bypass_forward(x) = W(X+A(X)) + B(X)
+        bypass_forward_diff(x) = W(A(X)) + B(X)
+        get_merged_weight() = W + WA + B
+        get_diff_weight() = WA + B
+        """
+        super().__init__(
+            lora_name,
+            org_module,
+            multiplier,
+            dropout,
+            rank_dropout,
+            module_dropout,
+            rank_dropout_scale,
+            bypass_mode,
+        )
+        if self.module_type not in self.support_module:
+            raise ValueError(f"{self.module_type} is not supported in GLoRA algo.")
         self.lora_dim = lora_dim
         self.tucker = False
+        self.rs_lora = rs_lora
 
-        if isinstance(org_module, nn.Conv2d):
-            assert org_module.kernel_size == (1, 1)
+        if self.module_type.startswith("conv"):
+            self.isconv = True
+            # For general LoCon
             in_dim = org_module.in_channels
+            k_size = org_module.kernel_size
+            stride = org_module.stride
+            padding = org_module.padding
             out_dim = org_module.out_channels
-            self.a1 = nn.Conv2d(in_dim, lora_dim, (1, 1), bias=False)
-            self.a2 = nn.Conv2d(lora_dim, in_dim, (1, 1), bias=False)
-            self.b1 = nn.Conv2d(in_dim, lora_dim, (1, 1), bias=False)
-            self.b2 = nn.Conv2d(lora_dim, out_dim, (1, 1), bias=False)
-        elif isinstance(org_module, nn.Linear):
-            in_dim = org_module.in_features
-            out_dim = org_module.out_features
-            self.a1 = nn.Linear(in_dim, lora_dim, bias=False)
-            self.a2 = nn.Linear(lora_dim, in_dim, bias=False)
-            self.b1 = nn.Linear(in_dim, lora_dim, bias=False)
-            self.b2 = nn.Linear(lora_dim, out_dim, bias=False)
-        else:
-            raise NotImplementedError
-        self.shape = org_module.weight.shape
+            self.tucker = use_tucker and all(i == 1 for i in k_size)
+            self.down_op = self.op
+            self.up_op = self.op
+
+            # A
+            self.a2 = self.module(in_dim, lora_dim, 1, bias=False)
+            self.a1 = self.module(lora_dim, in_dim, 1, bias=False)
 
-        if dropout:
-            self.dropout = nn.Dropout(dropout)
+            # B
+            if use_tucker and any(i != 1 for i in k_size):
+                self.b2 = self.module(in_dim, lora_dim, 1, bias=False)
+                self.bm = self.module(
+                    lora_dim, lora_dim, k_size, stride, padding, bias=False
+                )
+                self.tucker = True
+            else:
+                self.b2 = self.module(
+                    in_dim, lora_dim, k_size, stride, padding, bias=False
+                )
+            self.b1 = self.module(lora_dim, out_dim, 1, bias=False)
         else:
-            self.dropout = nn.Identity()
-        self.rank_dropout = rank_dropout
-        self.rank_dropout_scale = rank_dropout_scale
-        self.module_dropout = module_dropout
+            self.isconv = False
+            self.down_op = F.linear
+            self.up_op = F.linear
+            in_dim = org_module.in_features
+            out_dim = org_module.out_features
+            self.a2 = nn.Linear(in_dim, lora_dim, bias=False)
+            self.a1 = nn.Linear(lora_dim, in_dim, bias=False)
+            self.b2 = nn.Linear(in_dim, lora_dim, bias=False)
+            self.b1 = nn.Linear(lora_dim, out_dim, bias=False)
 
         if type(alpha) == torch.Tensor:
             alpha = alpha.detach().float().numpy()  # without casting, bf16 causes error
         alpha = lora_dim if alpha is None or alpha == 0 else alpha
-        self.scale = alpha / self.lora_dim
+
+        r_factor = lora_dim
+        if self.rs_lora:
+            r_factor = math.sqrt(r_factor)
+
+        self.scale = alpha / r_factor
+
         self.register_buffer("alpha", torch.tensor(alpha))  # 定数として扱える
 
+        if use_scalar:
+            self.scalar = nn.Parameter(torch.tensor(0.0))
+        else:
+            self.register_buffer("scalar", torch.tensor(1.0), persistent=False)
+
         # same as microsoft's
         torch.nn.init.kaiming_uniform_(self.a1.weight, a=math.sqrt(5))
-        torch.nn.init.zeros_(self.a2.weight)
         torch.nn.init.kaiming_uniform_(self.b1.weight, a=math.sqrt(5))
-        torch.nn.init.zeros_(self.b2.weight)
-
-        self.multiplier = multiplier
-        self.org_module = [org_module]
-        self.org_forward = self.org_module[0].forward
+        if use_scalar:
+            torch.nn.init.kaiming_uniform_(self.a2.weight, a=math.sqrt(5))
+            torch.nn.init.kaiming_uniform_(self.b2.weight, a=math.sqrt(5))
+        else:
+            torch.nn.init.zeros_(self.a2.weight)
+            torch.nn.init.zeros_(self.b2.weight)
 
-    def apply_to(self):
-        self.org_forward = self.org_module[0].forward
-        self.org_module[0].forward = self.forward
+    @classmethod
+    def make_module_from_state_dict(
+        cls, lora_name, orig_module, a1, a2, b1, b2, bm, alpha
+    ):
+        module = cls(
+            lora_name,
+            orig_module,
+            1,
+            a2.size(0),
+            float(alpha),
+            use_tucker=bm is not None,
+        )
+        module.a1.weight.data.copy_(a1)
+        module.a2.weight.data.copy_(a2)
+        module.b1.weight.data.copy_(b1)
+        module.b2.weight.data.copy_(b2)
+        if bm is not None:
+            module.bm.weight.data.copy_(bm)
+        return module
 
-    def restore(self):
-        self.org_module[0].forward = self.org_forward
+    def custom_state_dict(self):
+        destination = {}
+        destination["alpha"] = self.alpha
+        destination["a1.weight"] = self.a1.weight
+        destination["a2.weight"] = self.a2.weight * self.scalar
+        destination["b1.weight"] = self.b1.weight
+        destination["b2.weight"] = self.b2.weight * self.scalar
+        if self.tucker:
+            destination["bm.weight"] = self.bm.weight
+        return destination
 
-    def merge_to(self, multiplier=1.0):
-        weight = self.make_weight() * self.scale * multiplier
-        self.org_module[0].weight.data.add_(weight)
+    def load_weight_hook(self, module: nn.Module, incompatible_keys):
+        missing_keys = incompatible_keys.missing_keys
+        for key in missing_keys:
+            if "scalar" in key:
+                del missing_keys[missing_keys.index(key)]
+        if isinstance(self.scalar, nn.Parameter):
+            self.scalar.data.copy_(torch.ones_like(self.scalar))
+        elif getattr(self, "scalar", None) is not None:
+            self.scalar.copy_(torch.ones_like(self.scalar))
+        else:
+            self.register_buffer(
+                "scalar", torch.ones_like(self.scalar), persistent=False
+            )
 
     def make_weight(self, device=None):
         wa1 = self.a1.weight.view(self.a1.weight.size(0), -1)
         wa2 = self.a2.weight.view(self.a2.weight.size(0), -1)
-        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)
+        orig = self.org_weight
 
-    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
+        if self.tucker:
+            wb = tucker_weight_from_conv(self.b1.weight, self.b2.weight, self.bm.weight)
+        else:
+            wb1 = self.b1.weight.view(self.b1.weight.size(0), -1)
+            wb2 = self.b2.weight.view(self.b2.weight.size(0), -1)
+            wb = wb1 @ wb2
+            wb = wb.view(*orig.shape)
+        if orig.dim() > 2:
+            w_wa1 = torch.einsum("o i ..., i j -> o j ...", orig, wa1)
+            w_wa2 = torch.einsum("o i ..., i j -> o j ...", w_wa1, wa2)
+        else:
+            w_wa2 = (orig @ wa1) @ wa2
+        return (wb + w_wa2) * self.scale * self.scalar
 
-        ax_mid = self.a1(x) * scale
-        bx_mid = self.b1(x) * scale
+    def get_diff_weight(self, multiplier=1.0, shape=None, device=None):
+        weight = self.make_weight(device) * multiplier
+        if shape is not None:
+            weight = weight.view(shape)
+        return weight, None
+
+    def get_merged_weight(self, multiplier=1, shape=None, device=None):
+        diff_w, _ = self.get_diff_weight(multiplier, shape, device)
+        return self.org_weight + diff_w, None
+
+    def _bypass_forward(self, x, scale=1, diff=False):
+        scale = self.scale * scale
+        ax_mid = self.a2(x) * scale
+        bx_mid = self.b2(x) * scale
 
         if self.rank_dropout and self.training:
             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_a /= drop_a.mean()
                 drop_b /= drop_b.mean()
             if (dims := len(x.shape)) == 4:
                 drop_a = drop_a.view(1, -1, 1, 1)
@@ -121,8 +231,33 @@ def forward(self, x):
                 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
+            self.org_forward(
+                (0 if diff else x) + self.drop(self.a1(ax_mid)) * self.scale
+            )
+            + self.drop(self.b1(bx_mid)) * self.scale
         )
+
+    def bypass_forward_diff(self, x, scale=1):
+        return self._bypass_forward(x, scale=scale, diff=True)
+
+    def bypass_forward(self, x, scale=1):
+        return self._bypass_forward(x, scale=scale, diff=False)
+
+    def forward(self, x, *args, **kwargs):
+        if self.module_dropout and self.training:
+            if torch.rand(1) < self.module_dropout:
+                return self.org_forward(x)
+        if self.bypass_mode:
+            return self.bypass_forward(x, self.multiplier)
+        else:
+            weight = (
+                self.org_module[0].weight.data.to(self.dtype)
+                + self.get_diff_weight(multiplier=self.multiplier)[0]
+            )
+            bias = (
+                None
+                if self.org_module[0].bias is None
+                else self.org_module[0].bias.data
+            )
+            return self.op(x, weight, bias, **self.kw_dict)
diff --git a/lycoris/modules/hypernet/__init__.py b/lycoris/modules/hypernet/__init__.py
deleted file mode 100644
index 243c210..0000000
--- a/lycoris/modules/hypernet/__init__.py
+++ /dev/null
@@ -1,1 +0,0 @@
-from .generater import WeightDecoder, ImgWeightGenerator, TextWeightGenerator
diff --git a/lycoris/modules/hypernet/generater.py b/lycoris/modules/hypernet/generater.py
deleted file mode 100644
index b0c24bd..0000000
--- a/lycoris/modules/hypernet/generater.py
+++ /dev/null
@@ -1,223 +0,0 @@
-from typing import *
-import numpy as np
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-from torchvision.transforms.functional import resize
-
-from timm import create_model
-from einops import rearrange
-
-from lycoris.modules.attention import TransformerBlock
-
-
-def _get_sinusoid_encoding_table(n_position, d_hid):
-    """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)
-        ]
-
-    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
-
-    return torch.FloatTensor(sinusoid_table).unsqueeze(0)
-
-
-class WeightDecoder(nn.Module):
-    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)
-        )
-
-        # calc heads for mem-eff or flash_attn
-        heads = 1
-        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,
-            )
-            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)
-        )
-        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()
-        )
-        h = weight + pos_emb
-        for decoder in self.decoder_model:
-            h = decoder(h, context=features)
-        weight = weight + self.delta_proj(h)
-        return weight
-
-
-class ImgWeightGenerator(nn.Module):
-    def __init__(
-        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
-        decoder_blocks: int = 4,
-        sample_iters: int = 1,
-    ):
-        super(ImgWeightGenerator, 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)
-        )
-
-        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):
-            test_output = test_output[-1]
-        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:
-            with torch.no_grad():
-                features = self.encoder_model.forward_features(ref_img)
-        else:
-            features = self.encoder_model.forward_features(ref_img + ensure_grad)
-        if isinstance(features, list):
-            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)
-
-        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
-
-
-class TextWeightGenerator(nn.Module):
-    def __init__(
-        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
-        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)
-        )
-
-        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):
-            test_output = test_output[-1]
-        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():
-                features = self.encoder_model(caption)
-        else:
-            features = self.encoder_model(caption)
-        if isinstance(features, list):
-            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)
-
-        if weight is None:
-            weight = torch.zeros(
-                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
diff --git a/lycoris/modules/hypernet/text_encoder.py b/lycoris/modules/hypernet/text_encoder.py
deleted file mode 100644
index 10555d0..0000000
--- a/lycoris/modules/hypernet/text_encoder.py
+++ /dev/null
@@ -1,266 +0,0 @@
-"""
-https://github.com/Stability-AI/stablediffusion/blob/main/ldm/modules/encoders/modules.py
-"""
-
-import torch
-import torch.nn as nn
-from torch.utils.checkpoint import checkpoint
-
-from transformers import T5Tokenizer, T5EncoderModel, CLIPTokenizer, CLIPTextModel
-
-import open_clip
-
-
-class AbstractEncoder(nn.Module):
-    def __init__(self):
-        super().__init__()
-
-    def encode(self, *args, **kwargs):
-        raise NotImplementedError
-
-
-class IdentityEncoder(AbstractEncoder):
-    def encode(self, x):
-        return x
-
-
-class ClassEmbedder(nn.Module):
-    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)
-        self.n_classes = n_classes
-        self.ucg_rate = ucg_rate
-
-    def forward(self, batch, key=None, disable_dropout=False):
-        if key is None:
-            key = self.key
-        # this is for use in crossattn
-        c = batch[key][:, None]
-        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 = torch.ones((bs,), device=device) * uc_class
-        uc = {self.key: uc}
-        return uc
-
-
-def disabled_train(self, mode=True):
-    """Overwrite model.train with this function to make sure train/eval mode
-    does not change anymore."""
-    return self
-
-
-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
-        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?
-        if freeze:
-            self.freeze()
-
-    def freeze(self):
-        self.transformer = self.transformer.eval()
-        # 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",
-        )
-        tokens = batch_encoding["input_ids"].to(self.device)
-        outputs = self.transformer(input_ids=tokens)
-
-        z = outputs.last_hidden_state
-        return z
-
-    def encode(self, text):
-        return 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
-        super().__init__()
-        assert layer in self.LAYERS
-        self.tokenizer = CLIPTokenizer.from_pretrained(version)
-        self.transformer = CLIPTextModel.from_pretrained(version)
-        self.device = device
-        self.max_length = max_length
-        if freeze:
-            self.freeze()
-        self.layer = layer
-        self.layer_idx = layer_idx
-        if layer == "hidden":
-            assert layer_idx is not None
-            assert 0 <= abs(layer_idx) <= 12
-
-    def freeze(self):
-        self.transformer = self.transformer.eval()
-        # 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",
-        )
-        tokens = batch_encoding["input_ids"].to(self.device)
-        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)
-        else:
-            z = outputs.hidden_states[self.layer_idx]
-        return z
-
-    def encode(self, text):
-        return self(text)
-
-
-class FrozenOpenCLIPEmbedder(AbstractEncoder):
-    """
-    Uses the OpenCLIP transformer encoder for text
-    """
-
-    LAYERS = [
-        # "pooled",
-        "last",
-        "penultimate",
-    ]
-
-    def __init__(
-        self,
-        arch="ViT-L-14",
-        version="datacomp_xl_s13b_b90k",
-        max_length=77,
-        freeze=True,
-        layer="last",
-    ):
-        super().__init__()
-        assert layer in self.LAYERS
-        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()
-        self.layer = layer
-        if self.layer == "last":
-            self.layer_idx = 0
-        elif self.layer == "penultimate":
-            self.layer_idx = 1
-        else:
-            raise NotImplementedError()
-
-    def freeze(self):
-        self.model = self.model.eval()
-        for param in self.parameters():
-            param.requires_grad = False
-
-    def forward(self, text):
-        tokens = open_clip.tokenize(text)
-        z = self.encode_with_transformer(
-            tokens.to(self.model.token_embedding.weight.device)
-        )
-        return z
-
-    def encode_with_transformer(self, text):
-        x = self.model.token_embedding(text)  # [batch_size, n_ctx, d_model]
-        x = x + self.model.positional_embedding
-        x = x.permute(1, 0, 2)  # NLD -> LND
-        x = self.text_transformer_forward(x, attn_mask=self.model.attn_mask)
-        x = x.permute(1, 0, 2)  # LND -> NLD
-        x = self.model.ln_final(x)
-        return x
-
-    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()
-            ):
-                x = checkpoint(r, x, attn_mask)
-            else:
-                x = r(x, attn_mask=attn_mask)
-        return x
-
-    def encode(self, text):
-        return 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,
-    ):
-        super().__init__()
-        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):
-        return self(text)
-
-    def forward(self, text):
-        clip_z = self.clip_encoder.encode(text)
-        t5_z = self.t5_encoder.encode(text)
-        return [clip_z, t5_z]
-
-
-if __name__ == "__main__":
-    clip = FrozenCLIPEmbedder(layer="last_with_pooled").cuda()
-    ctx, embed = clip("test")
-    print(ctx.shape, embed.shape)
diff --git a/lycoris/modules/ia3.py b/lycoris/modules/ia3.py
index 643a986..6f82e07 100644
--- a/lycoris/modules/ia3.py
+++ b/lycoris/modules/ia3.py
@@ -1,36 +1,65 @@
 import torch
 import torch.nn as nn
 
-from .base import ModuleCustomSD
+from .base import LycorisBaseModule
+from ..utils.bnb import LinearNF4
 
 
-class IA3Module(ModuleCustomSD):
-    """
-    Hadamard product Implementaion for Low Rank Adaptation
-    """
+class IA3Module(LycorisBaseModule):
+    name = "ia3"
+    support_module = {
+        "linear",
+        "conv1d",
+        "conv2d",
+        "conv3d",
+    }
+    weight_list = ["weight", "on_input"]
+    weight_list_det = ["on_input"]
 
     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,
+        weight_decompose=False,
+        bypass_mode=False,
+        rs_lora=False,
         train_on_input=False,
-        **kwargs
+        **kwargs,
     ):
         """if alpha == 0 or None, alpha is rank (no scaling)."""
-        super().__init__()
-        self.lora_name = lora_name
-        self.tucker = False
+        super().__init__(
+            lora_name,
+            org_module,
+            multiplier,
+            dropout,
+            rank_dropout,
+            module_dropout,
+            rank_dropout_scale,
+            bypass_mode,
+        )
+        if self.module_type not in self.support_module:
+            raise ValueError(f"{self.module_type} is not supported in IA^3 algo.")
 
-        self.shape = org_module.weight.shape
-        if org_module.__class__.__name__ == "Conv2d":
+        if self.module_type.startswith("conv"):
+            self.isconv = True
             in_dim = org_module.in_channels
             out_dim = org_module.out_channels
             if train_on_input:
                 train_dim = in_dim
             else:
                 train_dim = out_dim
-            self.weight = nn.Parameter(torch.empty(1, train_dim, 1, 1))
+            self.weight = nn.Parameter(
+                torch.empty(1, train_dim, *(1 for _ in self.shape[2:]))
+            )
         else:
             in_dim = org_module.in_features
             out_dim = org_module.out_features
@@ -43,25 +72,72 @@ def __init__(
 
         # Need more experiences on init method
         torch.nn.init.constant_(self.weight, 0)
-
-        self.multiplier = multiplier
-        self.org_forward = None
-        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)))
 
+    @classmethod
+    def make_module_from_state_dict(cls, lora_name, orig_module, weight):
+        module = cls(
+            lora_name,
+            orig_module,
+            1,
+        )
+        module.weight.data.copy_(weight)
+        return module
+
     def apply_to(self):
         self.org_forward = self.org_module[0].forward
         self.org_module[0].forward = self.forward
 
-    @torch.enable_grad()
-    def forward(self, x):
+    def make_weight(self, multiplier=1, shape=None, device=None, diff=False):
+        weight = self.weight * multiplier + int(not diff)
         if self.train_input:
-            x = x * (1 + self.weight * self.multiplier)
+            diff = self.org_weight * weight
+        else:
+            diff = self.org_weight.transpose(0, 1) * weight
+            diff = diff.transpose(0, 1)
+        if shape is not None:
+            diff = diff.view(shape)
+        if device is not None:
+            diff = diff.to(device)
+        return diff
+
+    def get_diff_weight(self, multiplier=1, shape=None, device=None):
+        diff = self.make_weight(
+            multiplier=multiplier, shape=shape, device=device, diff=True
+        )
+        return diff, None
+
+    def get_merged_weight(self, multiplier=1, shape=None, device=None):
+        diff = self.make_weight(multiplier=multiplier, shape=shape, device=device)
+        return diff, None
+
+    def _bypass_forward(self, x, scale=1, diff=False):
+        weight = self.weight * scale + int(not diff)
+        if self.train_input:
+            x = x * weight
         out = self.org_forward(x)
-        dtype = out.dtype
         if not self.train_input:
-            out = out * (1 + self.weight * self.multiplier)
-            out = out.to(dtype)
+            out = out * weight
         return out
+
+    def bypass_forward_diff(self, x, scale=1):
+        return self._bypass_forward(x, scale, diff=True)
+
+    def bypass_forward(self, x, scale=1):
+        return self._bypass_forward(x, scale, diff=False)
+
+    def forward(self, x, *args, **kwargs):
+        if self.module_dropout and self.training:
+            if torch.rand(1) < self.module_dropout:
+                return self.org_forward(x)
+        if self.bypass_mode:
+            return self.bypass_forward(x, self.multiplier)
+        else:
+            weight = self.get_merged_weight(multiplier=self.multiplier)[0]
+            bias = (
+                None
+                if self.org_module[0].bias is None
+                else self.org_module[0].bias.data
+            )
+            return self.op(x, weight, bias, **self.kw_dict)
diff --git a/lycoris/modules/lilora.py b/lycoris/modules/lilora.py
deleted file mode 100644
index 14e9d83..0000000
--- a/lycoris/modules/lilora.py
+++ /dev/null
@@ -1,326 +0,0 @@
-import math
-from functools import cache
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-from .base import ModuleCustomSD
-from ..logging import logger
-
-
-@cache
-def log_wd():
-    return logger.warning(
-        "Using weight_decompose=True with LoRA (DoRA) will ignore network_dropout."
-        "Only rank dropout and module dropout will be applied"
-    )
-
-
-class LiLoConModule(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,
-        weight_decompose=False,
-        bypass_mode=False,
-        **kwargs,
-    ):
-        """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
-        assert not (
-            bypass_mode and weight_decompose
-        ), "bypass_mode and dora_wd cannot be used together"
-
-        if isinstance(org_module, nn.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.op = F.conv2d
-            self.extra_args = {
-                "stride": org_module.stride,
-                "padding": org_module.padding,
-                "dilation": org_module.dilation,
-                "groups": org_module.groups,
-            }
-        else:
-            in_dim = org_module.in_features
-            out_dim = org_module.out_features
-            self.op = F.linear
-            self.extra_args = {}
-
-        if isinstance(org_module, nn.Conv2d):
-            self.isconv = True
-            # For general LoCon
-            in_dim = org_module.in_channels
-            k_size = org_module.kernel_size
-            stride = org_module.stride
-            padding = org_module.padding
-            out_dim = org_module.out_channels
-            self.down_op = F.conv2d
-            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.tucker = True
-            else:
-                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
-            self.down_op = F.linear
-            self.up_op = F.linear
-            in_dim = org_module.in_features
-            out_dim = org_module.out_features
-            self.lora_down = nn.Linear(in_dim, lora_dim, bias=False)
-            self.lora_up = nn.Linear(lora_dim, out_dim, bias=False)
-        else:
-            raise NotImplementedError
-        self.shape = org_module.weight.shape
-
-        self.wd = weight_decompose
-        if self.wd:
-            org_weight: nn.Parameter = org_module.weight
-            self.dora_mean_dim = tuple(i for i in range(org_weight.dim()) if i != 1)
-            self.dora_scale = nn.Parameter(
-                torch.mean(org_weight, dim=self.dora_mean_dim, keepdim=True)
-            )
-
-        if dropout:
-            self.dropout = nn.Dropout(dropout)
-            if self.wd:
-                log_wd()
-        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))  # 定数として扱える
-
-        if use_scalar:
-            self.scalar = nn.Parameter(torch.tensor(0.0))
-        else:
-            self.scalar: torch.FloatTensor
-            self.register_buffer("scalar", torch.tensor(1.0), persistent=False)
-        # same as microsoft's
-        torch.nn.init.kaiming_uniform_(self.lora_down.weight, a=math.sqrt(5))
-        if use_scalar:
-            torch.nn.init.kaiming_uniform_(self.lora_up.weight, a=math.sqrt(5))
-        else:
-            torch.nn.init.constant_(self.lora_up.weight, 0)
-        if self.tucker:
-            torch.nn.init.kaiming_uniform_(self.lora_mid.weight, a=math.sqrt(5))
-
-        self.multiplier = multiplier
-        self.org_module = [org_module]
-        self.org_forward = self.org_module[0].forward
-
-    def load_weight_hook(self, module: nn.Module, incompatible_keys):
-        missing_keys = incompatible_keys.missing_keys
-        for key in missing_keys:
-            if "scalar" in key:
-                del missing_keys[missing_keys.index(key)]
-        self.scalar = nn.Parameter(torch.ones_like(self.scalar))
-
-    def apply_to(self, is_hypernet=False, **kwargs):
-        self.org_forward = self.org_module[0].forward
-        if is_hypernet:
-            self.org_module[0].forward = self.hypernet_forward
-        else:
-            self.org_module[0].forward = self.forward
-
-    def restore(self):
-        self.org_module[0].forward = self.org_forward
-
-    def merge_to(self, multiplier=1.0):
-        weight = self.make_weight() * self.scale * multiplier
-        self.org_module[0].weight.data.add_(weight)
-
-    def make_weight(self, device=None):
-        wa = self.lora_up.weight.to(device)
-        wb = self.lora_down.weight.to(device)
-        if self.tucker:
-            t = self.lora_mid.weight.to(device)
-            wa = wa.squeeze(-1, -2)
-            wb = wb.squeeze(-1, -2)
-            weight = torch.einsum("i j k l, p i, j r -> p r k l", t, wa, wb)
-        else:
-            weight = wa.view(wa.size(0), -1) @ wb.view(wb.size(0), -1)
-
-        weight = weight.view(self.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)
-            if self.rank_dropout_scale:
-                drop /= drop.mean()
-            weight *= drop
-
-        return weight * self.scalar.to(device)
-
-    def apply_weight_decompose(self, weight):
-        return (
-            weight / weight.mean(dim=self.dora_mean_dim, keepdim=True) * self.dora_scale
-        )
-
-    def custom_state_dict(self):
-        destination = {}
-        if self.wd:
-            destination["dora_scale"] = self.dora_scale
-        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["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)
-        desired = torch.clamp(norm, max=max_norm)
-        ratio = desired.cpu() / norm.cpu()
-
-        scaled = norm != desired
-        if scaled:
-            self.scalar *= 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)
-
-    def make_lightweight(self, down, up, seed=None, down_aux=None, up_aux=None):
-        if down.dim() == 3:
-            down = down.reshape(down.size(0), self.lora_dim, -1)
-            up = up.reshape(up.size(0), -1, self.lora_dim)
-        else:
-            down = down.reshape(self.lora_dim, -1)
-            up = up.reshape(-1, self.lora_dim)
-        # print(up.shape)
-        if seed is None:
-            assert down_aux is not None and up_aux is not None
-        else:
-            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
-                )
-                nn.init.orthogonal_(down_aux)
-                nn.init.orthogonal_(up_aux)
-                # print(up_aux.shape)
-            torch.set_rng_state(rng_state)
-        if down.dim() == 3 and down.size(0) == 1:
-            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
-        return (down @ down_aux), (up_aux @ up)
-
-    def apply_lightweight(self, down, up, seed=None, down_aux=None, up_aux=None):
-        down_weight, up_weight = self.make_lightweight(down, up, seed, down_aux, up_aux)
-        self.lora_down.weight.data = down_weight
-        self.lora_up.weight.data = up_weight
-        return down_weight, up_weight
-
-    def hypernet_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
-
-        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"
-                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)
-            else:
-                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
-            ).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)
-            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"
-                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)
-            else:
-                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:
-            if self.isconv:
-                weight = up_weight.unsqueeze(-1).unsqueeze(-1)
-            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
diff --git a/lycoris/modules/locon.py b/lycoris/modules/locon.py
index bd82ae7..a8581a3 100644
--- a/lycoris/modules/locon.py
+++ b/lycoris/modules/locon.py
@@ -5,9 +5,10 @@
 import torch.nn as nn
 import torch.nn.functional as F
 
-from .base import ModuleCustomSD
+from .base import LycorisBaseModule
+from ..functional.general import rebuild_tucker
 from ..logging import logger
-from ..utils.bnb import LinearNF4, QuantLinears, log_bypass
+from ..utils.bnb import LinearNF4
 
 
 @cache
@@ -18,10 +19,22 @@ def log_wd():
     )
 
 
-class LoConModule(ModuleCustomSD):
-    """
-    modifed from kohya-ss/sd-scripts/networks/lora:LoRAModule
-    """
+class LoConModule(LycorisBaseModule):
+    name = "locon"
+    support_module = {
+        "linear",
+        "conv1d",
+        "conv2d",
+        "conv3d",
+    }
+    weight_list = [
+        "lora_up.weight",
+        "lora_down.weight",
+        "lora_mid.weight",
+        "alpha",
+        "dora_scale",
+    ]
+    weight_list_det = ["lora_up.weight"]
 
     def __init__(
         self,
@@ -38,42 +51,27 @@ def __init__(
         rank_dropout_scale=False,
         weight_decompose=False,
         bypass_mode=False,
+        rs_lora=False,
         **kwargs,
     ):
         """if alpha == 0 or None, alpha is rank (no scaling)."""
-        super().__init__()
-        self.lora_name = lora_name
+        super().__init__(
+            lora_name,
+            org_module,
+            multiplier,
+            dropout,
+            rank_dropout,
+            module_dropout,
+            rank_dropout_scale,
+            bypass_mode,
+        )
+        if self.module_type not in self.support_module:
+            raise ValueError(f"{self.module_type} is not supported in LoRA/LoCon algo.")
         self.lora_dim = lora_dim
         self.tucker = False
+        self.rs_lora = rs_lora
 
-        if isinstance(org_module, nn.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.op = F.conv2d
-            self.extra_args = {
-                "stride": org_module.stride,
-                "padding": org_module.padding,
-                "dilation": org_module.dilation,
-                "groups": org_module.groups,
-            }
-        else:
-            in_dim = org_module.in_features
-            out_dim = org_module.out_features
-            self.op = F.linear
-            self.extra_args = {}
-
-        if isinstance(org_module, QuantLinears):
-            if not bypass_mode:
-                log_bypass()
-            bypass_mode = True
-        self.bypass_mode = bypass_mode
-        assert not (
-            bypass_mode and weight_decompose
-        ), "bypass_mode and dora_wd cannot be used together"
-
-        if isinstance(org_module, nn.Conv2d):
+        if self.module_type.startswith("conv"):
             self.isconv = True
             # For general LoCon
             in_dim = org_module.in_channels
@@ -81,19 +79,20 @@ def __init__(
             stride = org_module.stride
             padding = org_module.padding
             out_dim = org_module.out_channels
-            self.down_op = F.conv2d
-            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(
+            self.tucker = use_tucker and any(i != 1 for i in k_size)
+            self.down_op = self.op
+            self.up_op = self.op
+            if use_tucker and any(i != 1 for i in k_size):
+                self.lora_down = self.module(in_dim, lora_dim, 1, bias=False)
+                self.lora_mid = self.module(
                     lora_dim, lora_dim, k_size, stride, padding, bias=False
                 )
                 self.tucker = True
             else:
-                self.lora_down = nn.Conv2d(
+                self.lora_down = self.module(
                     in_dim, lora_dim, k_size, stride, padding, bias=False
                 )
-            self.lora_up = nn.Conv2d(lora_dim, out_dim, (1, 1), bias=False)
+            self.lora_up = self.module(lora_dim, out_dim, 1, bias=False)
         elif isinstance(org_module, nn.Linear):
             self.isconv = False
             self.down_op = F.linear
@@ -104,11 +103,10 @@ def __init__(
             self.lora_up = nn.Linear(lora_dim, out_dim, bias=False)
         else:
             raise NotImplementedError
-        self.shape = org_module.weight.shape
 
         self.wd = weight_decompose
         if self.wd:
-            org_weight: nn.Parameter = org_module.weight
+            org_weight = org_module.weight.cpu().clone().float()
             self.dora_norm_dims = org_weight.dim() - 1
             self.dora_scale = nn.Parameter(
                 torch.norm(
@@ -126,14 +124,17 @@ def __init__(
                 log_wd()
         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
+
+        r_factor = lora_dim
+        if self.rs_lora:
+            r_factor = math.sqrt(r_factor)
+
+        self.scale = alpha / r_factor
+
         self.register_buffer("alpha", torch.tensor(alpha))  # 定数として扱える
 
         if use_scalar:
@@ -149,62 +150,93 @@ def __init__(
         if self.tucker:
             torch.nn.init.kaiming_uniform_(self.lora_mid.weight, a=math.sqrt(5))
 
-        self.multiplier = multiplier
-        self.org_module = [org_module]
-        self.org_forward = self.org_module[0].forward
+    @classmethod
+    def make_module_from_state_dict(
+        cls, lora_name, orig_module, up, down, mid, alpha, dora_scale
+    ):
+        module = cls(
+            lora_name,
+            orig_module,
+            1,
+            down.size(0),
+            float(alpha),
+            use_tucker=mid is not None,
+            weight_decompose=dora_scale 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)
+        if dora_scale is not None:
+            module.dora_scale.copy_(dora_scale)
+        return module
 
     def load_weight_hook(self, module: nn.Module, incompatible_keys):
         missing_keys = incompatible_keys.missing_keys
         for key in missing_keys:
             if "scalar" in key:
                 del missing_keys[missing_keys.index(key)]
-        self.scalar = nn.Parameter(torch.ones_like(self.scalar))
-
-    def apply_to(self, is_hypernet=False, **kwargs):
-        self.org_forward = self.org_module[0].forward
-        if is_hypernet:
-            self.org_module[0].forward = self.hypernet_forward
+        if isinstance(self.scalar, nn.Parameter):
+            self.scalar.data.copy_(torch.ones_like(self.scalar))
+        elif getattr(self, "scalar", None) is not None:
+            self.scalar.copy_(torch.ones_like(self.scalar))
         else:
-            self.org_module[0].forward = self.forward
-
-    def restore(self):
-        self.org_module[0].forward = self.org_forward
-
-    def merge_to(self, multiplier=1.0):
-        weight = self.make_weight() * self.scale * multiplier
-        self.org_module[0].weight.data.add_(weight)
+            self.register_buffer(
+                "scalar", torch.ones_like(self.scalar), persistent=False
+            )
 
     def make_weight(self, device=None):
         wa = self.lora_up.weight.to(device)
         wb = self.lora_down.weight.to(device)
         if self.tucker:
-            t = self.lora_mid.weight.to(device)
-            wa = wa.squeeze(-1, -2)
-            wb = wb.squeeze(-1, -2)
-            weight = torch.einsum("i j k l, p i, j r -> p r k l", t, wa, wb)
+            t = self.lora_mid.weight
+            wa = wa.view(wa.size(0), -1).transpose(0, 1)
+            wb = wb.view(wb.size(0), -1)
+            weight = rebuild_tucker(t, wa, wb)
         else:
             weight = wa.view(wa.size(0), -1) @ wb.view(wb.size(0), -1)
 
         weight = weight.view(self.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 = (torch.rand(weight.size(0), device=device) > self.rank_dropout).to(
+                weight.dtype
+            )
+            drop = drop.view(-1, *[1] * len(weight.shape[1:]))
             if self.rank_dropout_scale:
                 drop /= drop.mean()
             weight *= drop
 
         return weight * self.scalar.to(device)
 
+    def get_diff_weight(self, multiplier=1, shape=None, device=None):
+        scale = self.scale * multiplier
+        diff = self.make_weight(device=device) * scale
+        if shape is not None:
+            diff = diff.view(shape)
+        if device is not None:
+            diff = diff.to(device)
+        return diff, None
+
+    def get_merged_weight(self, multiplier=1, shape=None, device=None):
+        merged = (
+            self.org_module[0].weight.data
+            + self.get_diff_weight(multiplier=multiplier, shape=shape, device=device)[0]
+        )
+        if self.wd:
+            merged = self.apply_weight_decompose(merged)
+        return merged, None
+
     def apply_weight_decompose(self, weight):
+        weight = weight.to(self.dora_scale.dtype)
         weight_norm = (
             weight.transpose(0, 1)
             .reshape(weight.shape[1], -1)
             .norm(dim=1, keepdim=True)
             .reshape(weight.shape[1], *[1] * self.dora_norm_dims)
             .transpose(0, 1)
-        )
+        ) + torch.finfo(weight.dtype).eps
 
-        return weight * (self.dora_scale / weight_norm)
+        return weight * (self.dora_scale.to(weight.device) / weight_norm)
 
     def custom_state_dict(self):
         destination = {}
@@ -230,17 +262,40 @@ def apply_max_norm(self, max_norm, device=None):
 
         return scaled, orig_norm * ratio
 
+    def bypass_forward_diff(self, x, scale=1):
+        if self.tucker:
+            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
+            ).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)
+            mid = mid * drop
+
+        return self.dropout(self.lora_up(mid) * self.scalar * self.scale * scale)
+
+    def bypass_forward(self, x, scale=1):
+        return self.org_forward(x) + self.bypass_forward_diff(x, scale=scale)
+
     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
 
-        dtype = next(self.parameters()).dtype
+        dtype = self.dtype
         if not self.bypass_mode:
             weight = (
-                self.org_module[0].weight.data.to(device=x.device, dtype=dtype)
-                + self.make_weight(x.device).to(device=x.device, dtype=dtype) * scale
+                self.org_module[0].weight.data.to(dtype)
+                + self.make_weight(x.device).to(dtype) * scale
             )
             if self.wd:
                 weight = self.apply_weight_decompose(weight)
@@ -249,50 +304,6 @@ def forward(self, x):
                 if self.org_module[0].bias is None
                 else self.org_module[0].bias.data
             )
-            return self.op(x, weight, bias, **self.extra_args)
+            return self.op(x, weight, bias, **self.kw_dict)
         else:
-            if self.tucker:
-                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
-                ).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)
-                mid = mid * drop
-
-            return self.org_forward(x) + self.dropout(
-                self.lora_up(mid) * self.scalar * scale
-            )
-
-
-if __name__ == "__main__":
-    base = nn.Linear(128, 128).cuda()
-    lokr = LoConModule("test", base, 1, 4, 1, weight_decompose=True).cuda()
-    print(lokr)
-    test_input = torch.randn(1, 128).cuda()
-    test_output = lokr(test_input)
-    print(test_output.shape)
-
-    base_4bit = LinearNF4(128, 128)
-    base_4bit.load_state_dict(base.state_dict())
-    base_4bit.cuda()
-    qlocon = LoConModule("test", base_4bit, 1, 4, 1, weight_decompose=False).cuda()
-    print(qlocon)
-    test_input = torch.randn(1, 128).cuda()
-    test_output = qlocon(test_input)
-    print(test_output.shape)
-
-    base = nn.Conv2d(128, 128, 3, 1, 1)
-    lokr = LoConModule("test", base, 1, 4, 1, weight_decompose=True, use_tucker=True)
-    print(lokr)
-    test_input = torch.randn(1, 128, 16, 16)
-    test_output = lokr(test_input)
-    print(test_output.shape)
+            return self.bypass_forward(x, scale=self.multiplier)
diff --git a/lycoris/modules/loha.py b/lycoris/modules/loha.py
index 316f039..fb0047c 100644
--- a/lycoris/modules/loha.py
+++ b/lycoris/modules/loha.py
@@ -1,91 +1,32 @@
+import math
+
 import torch
 import torch.nn as nn
-import torch.nn.functional as F
-
-from .base import ModuleCustomSD
-from ..utils.bnb import LinearNF4, QuantLinears, log_bypass
-
-
-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
-        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)
-        grad_w1a = temp @ w1b.T
-        grad_w1b = w1a.T @ temp
-
-        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
-
-
-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
-
-    @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
-        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)
-        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)
-        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
-        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)
-        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)
-        del grad_temp
-        return grad_t1, grad_w1a, grad_w1b, grad_t2, grad_w2a, grad_w2b, None
-
-
-def make_weight(w1a, w1b, w2a, w2b, scale):
-    return HadaWeight.apply(w1a, w1b, w2a, w2b, scale)
-
-
-def make_weight_tucker(t1, w1a, w1b, t2, w2a, w2b, scale):
-    return HadaWeightCP.apply(t1, w1a, w1b, t2, w2a, w2b, scale)
-
-
-class LohaModule(ModuleCustomSD):
-    """
-    Hadamard product Implementaion for Low Rank Adaptation
-    """
+from .base import LycorisBaseModule
+from ..functional.loha import diff_weight as loha_diff_weight
+from ..utils.bnb import LinearNF4
+
+
+class LohaModule(LycorisBaseModule):
+    name = "loha"
+    support_module = {
+        "linear",
+        "conv1d",
+        "conv2d",
+        "conv3d",
+    }
+    weight_list = [
+        "hada_w1_a",
+        "hada_w1_b",
+        "hada_w2_a",
+        "hada_w2_b",
+        "hada_t1",
+        "hada_t2",
+        "alpha",
+        "dora_scale",
+    ]
+    weight_list_det = ["hada_w1_a"]
 
     def __init__(
         self,
@@ -102,79 +43,64 @@ def __init__(
         rank_dropout_scale=False,
         weight_decompose=False,
         bypass_mode=False,
+        rs_lora=False,
         **kwargs,
     ):
-        """if alpha == 0 or None, alpha is rank (no scaling)."""
-        super().__init__()
+        super().__init__(
+            lora_name,
+            org_module,
+            multiplier,
+            dropout,
+            rank_dropout,
+            module_dropout,
+            rank_dropout_scale,
+            bypass_mode,
+        )
+        if self.module_type not in self.support_module:
+            raise ValueError(f"{self.module_type} is not supported in LoHa algo.")
         self.lora_name = lora_name
         self.lora_dim = lora_dim
         self.tucker = False
+        self.rs_lora = rs_lora
 
-        self.shape = org_module.weight.shape
-        if org_module.__class__.__name__ == "Conv2d":
+        w_shape = self.shape
+        if self.module_type.startswith("conv"):
             in_dim = org_module.in_channels
             k_size = org_module.kernel_size
             out_dim = org_module.out_channels
             self.shape = (out_dim, in_dim, *k_size)
-            self.tucker = use_tucker and k_size != (1, 1)
+            self.tucker = use_tucker and any(i != 1 for i in k_size)
             if self.tucker:
-                shape = (out_dim, in_dim, *k_size)
+                w_shape = (out_dim, in_dim, *k_size)
             else:
-                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,
-            }
-        else:
-            in_dim = org_module.in_features
-            out_dim = org_module.out_features
-            shape = (out_dim, in_dim)
-            self.shape = shape
-            self.op = F.linear
-            self.extra_args = {}
-
-        if isinstance(org_module, QuantLinears):
-            if not bypass_mode:
-                log_bypass()
-            bypass_mode = True
-        self.bypass_mode = bypass_mode
-        assert not (
-            bypass_mode and weight_decompose
-        ), "bypass_mode and dora_wd cannot be used together"
+                w_shape = (out_dim, in_dim * torch.tensor(k_size).prod().item())
 
         if self.tucker:
-            self.hada_t1 = nn.Parameter(
-                torch.empty(lora_dim, lora_dim, shape[2], shape[3])
-            )
+            self.hada_t1 = nn.Parameter(torch.empty(lora_dim, lora_dim, *w_shape[2:]))
             self.hada_w1_a = nn.Parameter(
-                torch.empty(lora_dim, shape[0])
+                torch.empty(lora_dim, w_shape[0])
             )  # out_dim, 1-mode
             self.hada_w1_b = nn.Parameter(
-                torch.empty(lora_dim, shape[1])
+                torch.empty(lora_dim, w_shape[1])
             )  # in_dim , 2-mode
 
-            self.hada_t2 = nn.Parameter(
-                torch.empty(lora_dim, lora_dim, shape[2], shape[3])
-            )
+            self.hada_t2 = nn.Parameter(torch.empty(lora_dim, lora_dim, *w_shape[2:]))
             self.hada_w2_a = nn.Parameter(
-                torch.empty(lora_dim, shape[0])
+                torch.empty(lora_dim, w_shape[0])
             )  # out_dim, 1-mode
             self.hada_w2_b = nn.Parameter(
-                torch.empty(lora_dim, shape[1])
+                torch.empty(lora_dim, w_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_w1_a = nn.Parameter(torch.empty(w_shape[0], lora_dim))
+            self.hada_w1_b = nn.Parameter(torch.empty(lora_dim, w_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.hada_w2_a = nn.Parameter(torch.empty(w_shape[0], lora_dim))
+            self.hada_w2_b = nn.Parameter(torch.empty(lora_dim, w_shape[1]))
 
         self.wd = weight_decompose
         if self.wd:
-            org_weight: nn.Parameter = org_module.weight
+            org_weight = org_module.weight.cpu().clone().float()
             self.dora_norm_dims = org_weight.dim() - 1
             self.dora_scale = nn.Parameter(
                 torch.norm(
@@ -186,17 +112,19 @@ def __init__(
                 .transpose(1, 0)
             ).float()
 
-        self.dropout = dropout
-        if dropout:
+        if self.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
+
+        r_factor = lora_dim
+        if self.rs_lora:
+            r_factor = math.sqrt(r_factor)
+
+        self.scale = alpha / r_factor
+
         self.register_buffer("alpha", torch.tensor(alpha))  # 定数として扱える
 
         if use_scalar:
@@ -215,46 +143,64 @@ def __init__(
         else:
             torch.nn.init.constant_(self.hada_w2_a, 0)
 
-        self.multiplier = multiplier
-        self.org_module = [org_module]  # remove in applying
-        self.org_forward = self.org_module[0].forward
-        self.grad_ckpt = False
+    @classmethod
+    def make_module_from_state_dict(
+        cls, lora_name, orig_module, w1a, w1b, w2a, w2b, t1, t2, alpha, dora_scale
+    ):
+        module = cls(
+            lora_name,
+            orig_module,
+            1,
+            w1b.size(0),
+            float(alpha),
+            use_tucker=t1 is not None,
+            weight_decompose=dora_scale is not None,
+        )
+        module.hada_w1_a.copy_(w1a)
+        module.hada_w1_b.copy_(w1b)
+        module.hada_w2_a.copy_(w2a)
+        module.hada_w2_b.copy_(w2b)
+        if t1 is not None:
+            module.hada_t1.copy_(t1)
+            module.hada_t2.copy_(t2)
+        if dora_scale is not None:
+            module.dora_scale.copy_(dora_scale)
+        return module
 
     def load_weight_hook(self, module: nn.Module, incompatible_keys):
         missing_keys = incompatible_keys.missing_keys
         for key in missing_keys:
             if "scalar" in key:
                 del missing_keys[missing_keys.index(key)]
-        self.scalar = nn.Parameter(torch.ones_like(self.scalar))
-
-    def apply_to(self):
-        self.org_module[0].forward = self.forward
-
-    def restore(self):
-        self.org_module[0].forward = self.org_forward
-
-    def merge_to(self, multiplier=1.0):
-        weight = self.get_weight(self.shape)
-        self.org_module[0].weight.data.add_(weight * multiplier)
+        if isinstance(self.scalar, nn.Parameter):
+            self.scalar.data.copy_(torch.ones_like(self.scalar))
+        elif getattr(self, "scalar", None) is not None:
+            self.scalar.copy_(torch.ones_like(self.scalar))
+        else:
+            self.register_buffer(
+                "scalar", torch.ones_like(self.scalar), persistent=False
+            )
 
     def get_weight(self, shape):
         if self.tucker:
-            weight = make_weight_tucker(
-                self.hada_t1,
-                self.hada_w1_a,
+            weight = loha_diff_weight(
                 self.hada_w1_b,
-                self.hada_t2,
-                self.hada_w2_a,
+                self.hada_w1_a,
                 self.hada_w2_b,
-                scale=torch.tensor(self.scale),
+                self.hada_w2_a,
+                self.hada_t1,
+                self.hada_t2,
+                gamma=torch.tensor(self.scale),
             )
         else:
-            weight = make_weight(
-                self.hada_w1_a,
+            weight = loha_diff_weight(
                 self.hada_w1_b,
-                self.hada_w2_a,
+                self.hada_w1_a,
                 self.hada_w2_b,
-                scale=torch.tensor(self.scale),
+                self.hada_w2_a,
+                None,
+                None,
+                gamma=torch.tensor(self.scale),
             )
         if shape is not None:
             weight = weight.reshape(shape)
@@ -266,14 +212,31 @@ def get_weight(self, shape):
             weight *= drop
         return weight
 
+    def get_diff_weight(self, multiplier=1, shape=None, device=None):
+        scale = self.scale * multiplier
+        diff = self.get_weight(shape) * scale
+        if device is not None:
+            diff = diff.to(device)
+        return diff, None
+
+    def get_merged_weight(self, multiplier=1, shape=None, device=None):
+        merged = (
+            self.org_module[0].weight.data
+            + self.get_diff_weight(multiplier=multiplier, shape=shape, device=device)[0]
+        )
+        if self.wd:
+            merged = self.apply_weight_decompose(merged)
+        return merged, None
+
     def apply_weight_decompose(self, weight):
+        weight = weight.to(self.dora_scale.dtype)
         weight_norm = (
             weight.transpose(0, 1)
             .reshape(weight.shape[1], -1)
             .norm(dim=1, keepdim=True)
             .reshape(weight.shape[1], *[1] * self.dora_norm_dims)
             .transpose(0, 1)
-        )
+        ) + torch.finfo(weight.dtype).eps
 
         return weight * (self.dora_scale / weight_norm)
 
@@ -304,8 +267,14 @@ def apply_max_norm(self, max_norm, device=None):
 
         return scaled, orig_norm * ratio
 
-    @torch.enable_grad()
-    def forward(self, x):
+    def bypass_forward_diff(self, x, scale=1):
+        diff_weight = self.get_weight(self.shape) * self.scalar * scale
+        return self.drop(self.op(x, diff_weight, **self.kw_dict))
+
+    def bypass_forward(self, x, scale=1):
+        return self.org_forward(x) + self.bypass_forward_diff(x, scale=scale)
+
+    def forward(self, x: torch.Tensor, *args, **kwargs):
         if self.module_dropout and self.training:
             if torch.rand(1) < self.module_dropout:
                 return self.op(
@@ -318,13 +287,10 @@ def forward(self, x):
                     ),
                 )
         if self.bypass_mode:
-            diff_weight = self.get_weight(self.shape) * self.scalar * self.multiplier
-            return self.org_forward(x) + self.op(
-                x, diff_weight.view(self.shape), **self.extra_args
-            )
+            return self.bypass_forward(x, scale=self.multiplier)
         else:
             weight = (
-                self.org_module[0].weight.data.to(x.device, dtype=self.hada_w1_a.dtype)
+                self.org_module[0].weight.data.to(self.dtype)
                 + self.get_weight(self.shape) * self.scalar * self.multiplier
             )
             bias = (
@@ -335,29 +301,4 @@ def forward(self, x):
 
             if self.wd:
                 weight = self.apply_weight_decompose(weight)
-            return self.op(x, weight.view(self.shape), bias, **self.extra_args)
-
-
-if __name__ == "__main__":
-    base = nn.Linear(128, 128).cuda()
-    lokr = LohaModule("test", base, 1, 4, 1, weight_decompose=True).cuda()
-    print(lokr)
-    test_input = torch.randn(1, 128).cuda()
-    test_output = lokr(test_input)
-    print(test_output.shape)
-
-    base_4bit = LinearNF4(128, 128)
-    base_4bit.load_state_dict(base.state_dict())
-    base_4bit.cuda()
-    qlocon = LohaModule("test", base_4bit, 1, 4, 1, weight_decompose=False).cuda()
-    print(qlocon)
-    test_input = torch.randn(1, 128).cuda()
-    test_output = qlocon(test_input)
-    print(test_output.shape)
-
-    base = nn.Conv2d(128, 128, 3, 1, 1)
-    lokr = LohaModule("test", base, 1, 4, 1, weight_decompose=True, use_tucker=True)
-    print(lokr)
-    test_input = torch.randn(1, 128, 16, 16)
-    test_output = lokr(test_input)
-    print(test_output.shape)
+            return self.op(x, weight, bias, **self.kw_dict)
diff --git a/lycoris/modules/lokr.py b/lycoris/modules/lokr.py
index 8613f68..ebe6bbf 100644
--- a/lycoris/modules/lokr.py
+++ b/lycoris/modules/lokr.py
@@ -7,9 +7,11 @@
 
 from einops import rearrange
 
-from .base import ModuleCustomSD
+from .base import LycorisBaseModule
+from ..functional import factorization, rebuild_tucker
+from ..functional.lokr import make_kron
 from ..logging import logger
-from ..utils.bnb import LinearNF4, QuantLinears, log_bypass
+from ..utils.bnb import LinearNF4
 
 
 @cache
@@ -21,71 +23,27 @@ def logging_force_full_matrix(lora_dim, dim, factor):
     )
 
 
-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.
-    second value is a value for weight.
-
-    Because of non-commutative property, A⊗B ≠ B⊗A. Meaning of two matrices is slightly different.
-
-    examples)
-    factor
-        -1               2                4               8               16               ...
-    127 -> 1, 127   127 -> 1, 127    127 -> 1, 127   127 -> 1, 127   127 -> 1, 127
-    128 -> 8, 16    128 -> 2, 64     128 -> 4, 32    128 -> 8, 16    128 -> 8, 16
-    250 -> 10, 25   250 -> 2, 125    250 -> 2, 125   250 -> 5, 50    250 -> 10, 25
-    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
-        if m > n:
-            n, m = m, n
-        return m, n
-    if factor < 0:
-        factor = dimension
-    m, n = 1, dimension
-    length = m + n
-    while m < n:
-        new_m = m + 1
-        while dimension % new_m != 0:
-            new_m += 1
-        new_n = dimension // new_m
-        if new_m + new_n > length or new_m > factor:
-            break
-        else:
-            m, n = new_m, new_n
-    if m > n:
-        n, m = m, n
-    return m, n
-
-
-def make_weight_tucker(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]
-    return rebuild2
-
-
-def make_kron(w1, w2, scale):
-    if len(w2.shape) == 4:
-        w1 = w1.unsqueeze(2).unsqueeze(2)
-    w2 = w2.contiguous()
-    rebuild = torch.kron(w1, w2)
-
-    return rebuild * scale
-
-
-class LokrModule(ModuleCustomSD):
-    """
-    modifed from kohya-ss/sd-scripts/networks/lora:LoRAModule
-        and from KohakuBlueleaf/LyCORIS/lycoris:loha:LoHaModule
-        and from KohakuBlueleaf/LyCORIS/lycoris:locon:LoconModule
-    """
+class LokrModule(LycorisBaseModule):
+    name = "kron"
+    support_module = {
+        "linear",
+        "conv1d",
+        "conv2d",
+        "conv3d",
+    }
+    weight_list = [
+        "lokr_w1",
+        "lokr_w1_a",
+        "lokr_w1_b",
+        "lokr_w2",
+        "lokr_w2_a",
+        "lokr_w2_b",
+        "lokr_t1",
+        "lokr_t2",
+        "alpha",
+        "dora_scale",
+    ]
+    weight_list_det = ["lokr_w1", "lokr_w1_a"]
 
     def __init__(
         self,
@@ -105,20 +63,32 @@ def __init__(
         weight_decompose=False,
         full_matrix=False,
         bypass_mode=False,
+        rs_lora=False,
+        unbalanced_factorization=False,
         **kwargs,
     ):
-        """if alpha == 0 or None, alpha is rank (no scaling)."""
-        super().__init__()
+        super().__init__(
+            lora_name,
+            org_module,
+            multiplier,
+            dropout,
+            rank_dropout,
+            module_dropout,
+            rank_dropout_scale,
+            bypass_mode,
+        )
+        if self.module_type not in self.support_module:
+            raise ValueError(f"{self.module_type} is not supported in LoKr algo.")
+
         factor = int(factor)
-        self.lora_name = lora_name
         self.lora_dim = lora_dim
         self.tucker = False
         self.use_w1 = False
         self.use_w2 = False
         self.full_matrix = full_matrix
+        self.rs_lora = rs_lora
 
-        self.shape = org_module.weight.shape
-        if org_module.__class__.__name__ == "Conv2d":
+        if self.module_type.startswith("conv"):
             in_dim = org_module.in_channels
             k_size = org_module.kernel_size
             out_dim = org_module.out_channels
@@ -126,8 +96,10 @@ def __init__(
 
             in_m, in_n = factorization(in_dim, factor)
             out_l, out_k = factorization(out_dim, factor)
+            if unbalanced_factorization:
+                out_l, out_k = out_k, out_l
             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)
+            self.tucker = use_tucker and any(i != 1 for i in k_size)
             if (
                 decompose_both
                 and lora_dim < max(shape[0][0], shape[1][0]) / 2
@@ -149,9 +121,7 @@ def __init__(
                     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_t2 = nn.Parameter(torch.empty(lora_dim, lora_dim, *shape[2:]))
                 self.lokr_w2_a = nn.Parameter(
                     torch.empty(lora_dim, shape[0][1])
                 )  # b, 1-mode
@@ -162,18 +132,11 @@ def __init__(
                 # 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])
+                    torch.empty(
+                        lora_dim, shape[1][1] * torch.tensor(shape[2:]).prod().item()
+                    )
                 )
                 # 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,
-            }
-
         else:  # Linear
             in_dim = org_module.in_features
             out_dim = org_module.out_features
@@ -181,12 +144,18 @@ def __init__(
 
             in_m, in_n = factorization(in_dim, factor)
             out_l, out_k = factorization(out_dim, factor)
+            if unbalanced_factorization:
+                out_l, out_k = out_k, out_l
             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
+                and not self.full_matrix
+            ):
                 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:
@@ -194,30 +163,20 @@ def __init__(
                 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 and not self.full_matrix:
                 # 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]))
                 # w1 ⊗ (w2_a x w2_b) = (a, b)⊗((c, dim)x(dim, d)) = (a, b)⊗(c, d) = (ac, bd)
             else:
+                if not self.full_matrix:
+                    logging_force_full_matrix(lora_dim, max(in_dim, out_dim), factor)
                 self.use_w2 = True
                 self.lokr_w2 = nn.Parameter(torch.empty(shape[0][1], shape[1][1]))
 
-            self.op = F.linear
-            self.extra_args = {}
-
-        if isinstance(org_module, QuantLinears):
-            if not bypass_mode:
-                log_bypass()
-            bypass_mode = True
-        self.bypass_mode = bypass_mode
-        assert not (
-            bypass_mode and weight_decompose
-        ), "bypass_mode and dora_wd cannot be used together"
-
         self.wd = weight_decompose
         if self.wd:
-            org_weight: nn.Parameter = org_module.weight
+            org_weight = org_module.weight.cpu().clone().float()
             self.dora_norm_dims = org_weight.dim() - 1
             self.dora_scale = nn.Parameter(
                 torch.norm(
@@ -242,7 +201,13 @@ def __init__(
         if self.use_w2 and self.use_w1:
             # use scale = 1
             alpha = lora_dim
-        self.scale = alpha / self.lora_dim
+
+        r_factor = lora_dim
+        if self.rs_lora:
+            r_factor = math.sqrt(r_factor)
+
+        self.scale = alpha / r_factor
+
         self.register_buffer("alpha", torch.tensor(alpha))  # 定数として扱える
 
         if use_scalar:
@@ -270,42 +235,95 @@ def __init__(
             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))
 
-        self.multiplier = multiplier
-        self.org_module = [org_module]
-        self.org_forward = self.org_module[0].forward
-        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_tucker(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),
+    @classmethod
+    def make_module_from_state_dict(
+        cls,
+        lora_name,
+        orig_module,
+        w1,
+        w1a,
+        w1b,
+        w2,
+        w2a,
+        w2b,
+        _,
+        t2,
+        alpha,
+        dora_scale,
+    ):
+        full_matrix = False
+        if w1a is not None:
+            lora_dim = w1a.size(1)
+        elif w2a is not None:
+            lora_dim = w2a.size(1)
+        else:
+            full_matrix = True
+            lora_dim = 1
+
+        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]
+        ):
+            factor = -1
+        else:
+            factor = max(w1.shape) if w1 is not None else max(w1a.size(0), w1b.size(1))
+
+        module = cls(
+            lora_name,
+            orig_module,
+            1,
+            lora_dim,
+            float(alpha),
+            use_tucker=t2 is not None,
+            decompose_both=w1 is None and w2 is None,
+            factor=factor,
+            weight_decompose=dora_scale is not None,
+            full_matrix=full_matrix,
         )
-        assert torch.sum(torch.isnan(weight)) == 0, "weight is nan"
+        if w1 is not None:
+            module.lokr_w1.copy_(w1)
+        else:
+            module.lokr_w1_a.copy_(w1a)
+            module.lokr_w1_b.copy_(w1b)
+        if w2 is not None:
+            module.lokr_w2.copy_(w2)
+        else:
+            module.lokr_w2_a.copy_(w2a)
+            module.lokr_w2_b.copy_(w2b)
+        if t2 is not None:
+            module.lokr_t2.copy_(t2)
+        if dora_scale is not None:
+            module.dora_scale.copy_(dora_scale)
+        return module
 
     def load_weight_hook(self, module: nn.Module, incompatible_keys):
         missing_keys = incompatible_keys.missing_keys
         for key in missing_keys:
             if "scalar" in key:
                 del missing_keys[missing_keys.index(key)]
-        self.scalar = nn.Parameter(torch.ones_like(self.scalar))
-
-    # Same as locon.py
-    def apply_to(self):
-        self.org_forward = self.org_module[0].forward
-        self.org_module[0].forward = self.forward
-
-    def restore(self):
-        self.org_module[0].forward = self.org_forward
-
-    def merge_to(self, multiplier=1.0):
-        weight = self.get_weight(self.shape)
-        self.org_module[0].weight.data.add_(weight * multiplier)
+        if isinstance(self.scalar, nn.Parameter):
+            self.scalar.data.copy_(torch.ones_like(self.scalar))
+        elif getattr(self, "scalar", None) is not None:
+            self.scalar.copy_(torch.ones_like(self.scalar))
+        else:
+            self.register_buffer(
+                "scalar", torch.ones_like(self.scalar), persistent=False
+            )
 
     def get_weight(self, shape):
         weight = make_kron(
@@ -314,31 +332,49 @@ def get_weight(self, shape):
                 self.lokr_w2
                 if self.use_w2
                 else (
-                    make_weight_tucker(self.lokr_t2, self.lokr_w2_a, self.lokr_w2_b)
+                    rebuild_tucker(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),
         )
+        dtype = weight.dtype
         if shape is not None:
-            weight = weight.reshape(shape)
+            weight = weight.view(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 = (torch.rand(weight.size(0)) > self.rank_dropout).to(dtype)
+            drop = drop.view(-1, *[1] * len(weight.shape[1:]))
             if self.rank_dropout_scale:
                 drop /= drop.mean()
             weight *= drop
         return weight
 
+    def get_diff_weight(self, multiplier=1, shape=None, device=None):
+        scale = self.scale * multiplier
+        diff = self.get_weight(shape) * scale
+        if device is not None:
+            diff = diff.to(device)
+        return diff, None
+
+    def get_merged_weight(self, multiplier=1, shape=None, device=None):
+        merged = (
+            self.org_module[0].weight.data
+            + self.get_diff_weight(multiplier=multiplier, shape=shape, device=device)[0]
+        )
+        if self.wd:
+            merged = self.apply_weight_decompose(merged)
+        return merged, None
+
     def apply_weight_decompose(self, weight):
+        weight = weight.to(self.dora_scale.dtype)
         weight_norm = (
             weight.transpose(0, 1)
             .reshape(weight.shape[1], -1)
             .norm(dim=1, keepdim=True)
             .reshape(weight.shape[1], *[1] * self.dora_norm_dims)
             .transpose(0, 1)
-        )
+        ) + torch.finfo(weight.dtype).eps
 
         return weight * (self.dora_scale / weight_norm)
 
@@ -388,63 +424,90 @@ def apply_max_norm(self, max_norm, device=None):
 
         return scaled, orig_norm * ratio
 
-    def bypass_forward(self, h):
-        if isinstance(self.org_module[0], nn.Conv2d):
-            h = h.transpose(1, -1)
+    def bypass_forward_diff(self, h, scale=1):
+        is_conv = self.module_type.startswith("conv")
         if self.use_w2:
             ba = self.lokr_w2
         else:
             a = self.lokr_w2_b
             b = self.lokr_w2_a
+
+            if self.tucker:
+                t = self.lokr_t2
+                a = a.view(*a.shape, *[1] * (len(t.shape) - 2))
+                b = b.view(*b.shape, *[1] * (len(t.shape) - 2))
+            elif is_conv:
+                a = a.view(*a.shape, *self.shape[2:])
+                b = b.view(*b.shape, *[1] * (len(self.shape) - 2))
+
         if self.use_w1:
             c = self.lokr_w1
         else:
             c = self.lokr_w1_a @ self.lokr_w1_b
         uq = c.size(1)
 
+        if is_conv:
+            # (b, uq), vq, ...
+            b, _, *rest = h.shape
+            h_in_group = h.reshape(b * uq, -1, *rest)
+        else:
+            # b, ..., uq, vq
+            h_in_group = h.reshape(*h.shape[:-1], uq, -1)
+
         if self.use_w2:
-            vq = ba.size(1)
-            h_in_group = rearrange(h, "b ... (uq vq) -> b ... uq vq", uq=uq, vq=vq)
-            hb = F.linear(h_in_group, ba)
+            hb = self.op(h_in_group, ba, **self.kw_dict)
         else:
-            vq = a.size(1)
-            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)
+            if is_conv:
+                if self.tucker:
+                    ha = self.op(h_in_group, a)
+                    ht = self.op(ha, t, **self.kw_dict)
+                    hb = self.op(ht, b)
+                else:
+                    ha = self.op(h_in_group, a, **self.kw_dict)
+                    hb = self.op(ha, b)
+            else:
+                ha = self.op(h_in_group, a, **self.kw_dict)
+                hb = self.op(ha, b)
+
+        if is_conv:
+            # (b, uq), vp, ..., f
+            # -> b, uq, vp, ..., f
+            # -> b, f, vp, ..., uq
+            hb = hb.view(b, -1, *hb.shape[1:])
+            h_cross_group = hb.transpose(1, -1)
+        else:
+            # b, ..., uq, vq
+            # -> b, ..., vq, uq
+            h_cross_group = hb.transpose(-1, -2)
 
-        h_cross_group = hb.transpose(-1, -2)
         hc = F.linear(h_cross_group, c)
+        if is_conv:
+            # b, f, vp, ..., up
+            # -> b, up, vp, ... ,f
+            # -> b, c, ..., f
+            hc = hc.transpose(1, -1)
+            h = hc.reshape(b, -1, *hc.shape[3:])
+        else:
+            # b, ..., vp, up
+            # -> b, ..., up, vp
+            # -> b, ..., c
+            hc = hc.transpose(-1, -2)
+            h = hc.reshape(*hc.shape[:-2], -1)
 
-        h = rearrange(hc, "b ... vp up -> b ... (up vp)")
-        if isinstance(self.org_module[0], nn.Conv2d):
-            h = h.transpose(1, -1)
+        return self.drop(h * scale)
 
-        return h
+    def bypass_forward(self, x, scale=1):
+        return self.org_forward(x) + self.bypass_forward_diff(x, scale=scale)
 
-    def forward(self, x):
+    def forward(self, x: torch.Tensor, *args, **kwargs):
         if self.module_dropout and self.training:
             if torch.rand(1) < self.module_dropout:
-                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
-                    ),
-                )
+                return self.org_forward(x)
         if self.bypass_mode:
-            if len(self.shape) > 2 and self.shape[2] > 1:
-                return self.org_forward(x) + self.op(
-                    x, self.get_weight(self.shape), **self.extra_args
-                )
-            else:
-                return self.org_forward(x) + self.bypass_forward(x)
+            return self.bypass_forward(x, self.multiplier)
         else:
             weight = (
-                self.org_module[0].weight.data.to(
-                    x.device, dtype=next(self.parameters()).dtype
-                )
+                self.org_module[0].weight.data.to(self.dtype)
                 + self.get_weight(self.shape) * self.scalar * self.multiplier
             )
             bias = (
@@ -455,40 +518,48 @@ def forward(self, x):
 
             if self.wd:
                 weight = self.apply_weight_decompose(weight)
-            return self.op(x, weight.view(self.shape), bias, **self.extra_args)
+            return self.op(x, weight, bias, **self.kw_dict)
 
 
 if __name__ == "__main__":
-    base = nn.Linear(128, 128).cuda()
-    lokr = LokrModule("test", base, 1, 4, 1, weight_decompose=True, factor=8).cuda()
-    print(lokr)
-    test_input = torch.randn(1, 77, 128).cuda()
-    test_output = lokr(test_input)
-    print(test_output.shape)
-
-    # opt = torch.optim.AdamW(lokr.parameters(), lr=1e-2)
-    # for _ in range(100):
-    #     x = torch.randn(128, 128).cuda()
-    #     t = x / 10
-    #     y = lokr(x)
-    #     loss = F.mse_loss(y, t)
-    #     loss.backward()
-    #     opt.step()
-    #     opt.zero_grad()
-    #     print(loss.item())
-
-    base_4bit = LinearNF4(128, 128)
-    base_4bit.load_state_dict(base.state_dict())
-    base_4bit.cuda()
-    qlocon = LokrModule("test", base_4bit, 1, 4, 1, weight_decompose=False).cuda()
-    print(qlocon)
-    test_output2 = qlocon(test_input)
-    print(test_output2.shape)
-    print(F.mse_loss(test_output, test_output2))
-
     base = nn.Conv2d(128, 128, 3, 1, 1)
-    lokr = LokrModule("test", base, 1, 4, 1, weight_decompose=True, use_tucker=True)
-    print(lokr)
+    net = LokrModule(
+        "",
+        base,
+        multiplier=1,
+        lora_dim=4,
+        alpha=1,
+        weight_decompose=False,
+        use_tucker=False,
+        use_scalar=False,
+        decompose_both=True,
+    )
+    net.apply_to()
+    sd = net.state_dict()
+    for key in sd:
+        if key != "alpha":
+            sd[key] = torch.randn_like(sd[key])
+    net.load_state_dict(sd)
+
     test_input = torch.randn(1, 128, 16, 16)
-    test_output = lokr(test_input)
+    test_output = net(test_input)
     print(test_output.shape)
+
+    net2 = LokrModule(
+        "",
+        base,
+        multiplier=1,
+        lora_dim=4,
+        alpha=1,
+        weight_decompose=False,
+        use_tucker=False,
+        use_scalar=False,
+        bypass_mode=True,
+        decompose_both=True,
+    )
+    net2.apply_to()
+    net2.load_state_dict(sd)
+    print(net2)
+
+    test_output2 = net(test_input)
+    print(F.mse_loss(test_output, test_output2))
diff --git a/lycoris/modules/norms.py b/lycoris/modules/norms.py
index bd45dee..3a855d5 100644
--- a/lycoris/modules/norms.py
+++ b/lycoris/modules/norms.py
@@ -1,14 +1,18 @@
 import torch
 import torch.nn as nn
-import torch.nn.functional as F
 
-from .base import ModuleCustomSD
+from .base import LycorisBaseModule
+from ..logging import warning_once
 
 
-class NormModule(ModuleCustomSD):
-    """
-    modifed from kohya-ss/sd-scripts/networks/lora:LoRAModule
-    """
+class NormModule(LycorisBaseModule):
+    name = "norm"
+    support_module = {
+        "layernorm",
+        "groupnorm",
+    }
+    weight_list = ["w_norm", "b_norm"]
+    weight_list_det = ["w_norm"]
 
     def __init__(
         self,
@@ -21,49 +25,52 @@ def __init__(
         **kwargs,
     ):
         """if alpha == 0 or None, alpha is rank (no scaling)."""
-        super().__init__()
-        self.lora_name = lora_name
+        super().__init__(
+            lora_name=lora_name,
+            org_module=org_module,
+            multiplier=multiplier,
+            rank_dropout=rank_dropout,
+            module_dropout=module_dropout,
+            rank_dropout_scale=rank_dropout_scale,
+            **kwargs,
+        )
+        if self.module_type == "unknown":
+            if not hasattr(org_module, "weight") or not hasattr(org_module, "_norm"):
+                warning_once(f"{type(org_module)} is not supported in Norm algo.")
+                return
+            else:
+                self.dim = org_module.weight.numel()
+                self.not_supported = False
+        elif self.module_type not in self.support_module:
+            warning_once(f"{self.module_type} is not supported in Norm algo.")
+            return
 
-        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,
-            }
-        elif isinstance(org_module, nn.GroupNorm):
-            self.op = F.group_norm
-            self.group_num = org_module.num_groups
-            self.dim = org_module.num_channels
-            self.kw_dict = {"num_groups": org_module.num_groups, "eps": org_module.eps}
-        else:
-            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]
-        self.org_forward = self.org_module[0].forward
-
-    def apply_to(self, **kwargs):
-        self.org_forward = self.org_module[0].forward
-        self.org_module[0].forward = self.forward
-
-    def restore(self):
-        self.org_module[0].forward = self.org_forward
+        if hasattr(org_module, "bias"):
+            self.b_norm = nn.Parameter(torch.zeros(self.dim))
+        if hasattr(org_module, "_norm"):
+            self.org_norm = org_module._norm
+        else:
+            self.org_norm = None
 
-    def merge_to(self, multiplier=1.0):
-        weight, bias = self.make_weight(scale=multiplier)
-        self.org_module[0].weight.data.copy_(weight)
-        self.org_module[0].bias.data.copy_(bias)
+    @classmethod
+    def make_module_from_state_dict(cls, lora_name, orig_module, w_norm, b_norm):
+        module = cls(
+            lora_name,
+            orig_module,
+            1,
+        )
+        module.w_norm.copy_(w_norm)
+        if b_norm is not None:
+            module.b_norm.copy_(b_norm)
+        return module
 
     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 hasattr(self.org_module[0], "bias"):
+            org_bias = self.org_module[0].bias.to(device, dtype=self.b_norm.dtype)
+        else:
+            org_bias = None
         if self.rank_dropout and self.training:
             drop = (torch.rand(self.dim, device=device) < self.rank_dropout).to(
                 self.w_norm.device
@@ -78,15 +85,75 @@ def make_weight(self, scale=1, device=None):
             else 1
         )
         weight = self.w_norm.to(device) * drop * scale
-        bias = self.b_norm.to(device) * drop * scale
-        return org_weight + weight, org_bias + bias
+        if org_bias is not None:
+            bias = self.b_norm.to(device) * drop * scale
+        return org_weight + weight, org_bias + bias if org_bias is not None else None
+
+    def get_diff_weight(self, multiplier=1, shape=None, device=None):
+        if self.not_supported:
+            return 0, 0
+        w = self.w_norm * multiplier
+        if device is not None:
+            w = w.to(device)
+        if shape is not None:
+            w = w.view(shape)
+        if self.b_norm is not None:
+            b = self.b_norm * multiplier
+            if device is not None:
+                b = b.to(device)
+            if shape is not None:
+                b = b.view(shape)
+        else:
+            b = None
+        return w, b
+
+    def get_merged_weight(self, multiplier=1, shape=None, device=None):
+        if self.not_supported:
+            return None, None
+        diff_w, diff_b = self.get_diff_weight(multiplier, shape, device)
+        org_w = self.org_module[0].weight.to(device, dtype=self.w_norm.dtype)
+        weight = org_w + diff_w
+        if diff_b is not None:
+            org_b = self.org_module[0].bias.to(device, dtype=self.b_norm.dtype)
+            bias = org_b + diff_b
+        else:
+            bias = None
+        return weight, bias
 
     def forward(self, x):
-        if self.module_dropout and self.training:
-            if torch.rand(1) < self.module_dropout:
-                return self.org_forward(x)
+        if self.not_supported or (
+            self.module_dropout
+            and self.training
+            and torch.rand(1) < self.module_dropout
+        ):
+            return self.org_forward(x)
         scale = self.multiplier
 
         w, b = self.make_weight(scale, x.device)
+        if self.org_norm is not None:
+            normed = self.org_norm(x)
+            scaled = normed * w
+            if b is not None:
+                scaled += b
+            return scaled
+
         kw_dict = self.kw_dict | {"weight": w, "bias": b}
         return self.op(x, **kw_dict)
+
+
+if __name__ == "__main__":
+    base = nn.LayerNorm(128).cuda()
+    norm = NormModule("test", base, 1).cuda()
+    print(norm)
+    test_input = torch.randn(1, 128).cuda()
+    test_output = norm(test_input)
+    torch.sum(test_output).backward()
+    print(test_output.shape)
+
+    base = nn.GroupNorm(4, 128).cuda()
+    norm = NormModule("test", base, 1).cuda()
+    print(norm)
+    test_input = torch.randn(1, 128, 3, 3).cuda()
+    test_output = norm(test_input)
+    torch.sum(test_output).backward()
+    print(test_output.shape)
diff --git a/lycoris/utils/__init__.py b/lycoris/utils/__init__.py
index 18aac08..f1acf0e 100644
--- a/lycoris/utils/__init__.py
+++ b/lycoris/utils/__init__.py
@@ -1,4 +1,6 @@
 import re
+import hashlib
+from io import BytesIO
 from typing import Dict, Tuple, Union
 
 import numpy as np
@@ -6,10 +8,37 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-
 import torch.linalg as linalg
 
+import safetensors.torch
+
 from tqdm import tqdm
+from .general import *
+
+
+def load_bytes_in_safetensors(tensors):
+    bytes = safetensors.torch.save(tensors)
+    b = BytesIO(bytes)
+
+    b.seek(0)
+    header = b.read(8)
+    n = int.from_bytes(header, "little")
+
+    offset = n + 8
+    b.seek(offset)
+
+    return b.read()
+
+
+def precalculate_safetensors_hashes(state_dict):
+    # calculate each tensor one by one to reduce memory usage
+    hash_sha256 = hashlib.sha256()
+    for tensor in state_dict.values():
+        single_tensor_sd = {"tensor": tensor}
+        bytes_for_tensor = load_bytes_in_safetensors(single_tensor_sd)
+        hash_sha256.update(bytes_for_tensor)
+
+    return f"0x{hash_sha256.hexdigest()}"
 
 
 def str_bool(val):
@@ -388,179 +417,10 @@ def match(match_list, regex_text):
     return key
 
 
-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)
-        dora_scale = lyco_state_dict.get(f"{lora_name}.dora_scale", None)
-        return "locon", (up, down, mid, alpha, dora_scale)
-    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)
-        dora_scale = lyco_state_dict.get(f"{lora_name}.dora_scale", None)
-        return "hada", (w1a, w1b, w2a, w2b, t1, t2, alpha, dora_scale)
-    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)
-        dora_scale = lyco_state_dict.get(f"{lora_name}.dora_scale", None)
-        return "kron", (w1, w1a, w1b, w2, w2a, w2b, t1, t2, alpha, dora_scale)
-    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", None)
-        return "norm", (w_norm, b_norm)
-    else:
-        return "None", ()
-
-
-def tucker_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)
-
-
-def tucker_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 apply_dora_scale(org_weight, rebuild, dora_scale, scale):
-    dora_mean_dim = tuple(i for i in range(org_weight.dim()) if i != 1)
-    weight = org_weight + rebuild
-    merged_scale1 = weight / weight.mean(dim=dora_mean_dim, keepdim=True) * dora_scale
-    diff_weight = merged_scale1 - org_weight
-    return org_weight + diff_weight * scale
-
-
-@torch.no_grad()
-def rebuild_weight(module_type, params, orig_weight, orig_bias, scale=1):
-    if orig_weight is None:
-        return None, None
-    merged = orig_weight
-    merged_bias = orig_bias
-    if module_type == "locon":
-        up, down, mid, alpha, dora_scale = params
-        if alpha is not None:
-            scale *= alpha / up.size(1)
-        if mid is not None:
-            rebuild = tucker_weight_from_conv(up, down, mid)
-        else:
-            rebuild = up.reshape(up.size(0), -1) @ down.reshape(down.size(0), -1)
-        rebuild = rebuild.reshape(orig_weight.shape)
-        if dora_scale is None:
-            merged = orig_weight + rebuild * scale
-        else:
-            merged = apply_dora_scale(orig_weight, rebuild, dora_scale, scale)
-        del up, down, mid, alpha, params, rebuild
-    elif module_type == "hada":
-        w1a, w1b, w2a, w2b, t1, t2, alpha, dora_scale = params
-        if alpha is not None:
-            scale *= alpha / w1b.size(0)
-        if t1 is not None:
-            rebuild1 = tucker_weight(w1a, w1b, t1)
-        else:
-            rebuild1 = w1a @ w1b
-        if t2 is not None:
-            rebuild2 = tucker_weight(w2a, w2b, t2)
-        else:
-            rebuild2 = w2a @ w2b
-        rebuild = (rebuild1 * rebuild2).reshape(orig_weight.shape)
-        if dora_scale is None:
-            merged = orig_weight + rebuild * scale
-        else:
-            merged = apply_dora_scale(orig_weight, rebuild, dora_scale, scale)
-        del w1a, w1b, w2a, w2b, t1, t2, alpha, params, rebuild, rebuild1, rebuild2
-    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":
-        w1, w1a, w1b, w2, w2a, w2b, t1, t2, alpha, dora_scale = 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))
-        if w1a is not None and w1b is not None:
-            if t1 is not None:
-                w1 = tucker_weight(w1a, w1b, t1)
-            else:
-                w1 = w1a @ w1b
-        if w2a is not None and w2b is not None:
-            if t2 is not None:
-                w2 = tucker_weight(w2a, w2b, t2)
-            else:
-                w2 = w2a @ w2b
-        if len(w2.shape) == 4:
-            w1 = w1.unsqueeze(2).unsqueeze(2)
-        w2 = w2.contiguous()
-        rebuild = torch.kron(w1, w2).reshape(orig_weight.shape)
-        if dora_scale is None:
-            merged = orig_weight + rebuild * scale
-        else:
-            merged = apply_dora_scale(orig_weight, rebuild, dora_scale, 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
-        )
-        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":
-        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:
-        return None, None
-
-    return merged, merged_bias
-
-
 @torch.no_grad()
 def merge(tes, unet, lyco_state_dict, scale: float = 1.0, device="cpu"):
-    UNET_TARGET_REPLACE_MODULE = [
-        "Linear",
-        "Conv2d",
-        "LayerNorm",
-        "GroupNorm",
-        "GroupNorm32",
-    ]
-    TEXT_ENCODER_TARGET_REPLACE_MODULE = [
-        "Embedding",
-        "Linear",
-        "Conv2d",
-        "LayerNorm",
-        "GroupNorm",
-        "GroupNorm32",
-        "Embedding",
-    ]
+    from ..modules import make_module, get_module
+
     LORA_PREFIX_UNET = "lora_unet"
     LORA_PREFIX_TEXT_ENCODER = "lora_te"
     merged = 0
@@ -569,29 +429,25 @@ def merge_state_dict(
         prefix,
         root_module: torch.nn.Module,
         lyco_state_dict: Dict[str, torch.Tensor],
-        target_replace_modules,
     ):
         nonlocal merged
         for child_name, child_module in tqdm(
             list(root_module.named_modules()), desc=f"Merging {prefix}"
         ):
-            if child_module.__class__.__name__ in target_replace_modules:
-                lora_name = prefix + "." + 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:
-                    key_dict.pop(lora_name)
-                    merged += 1
-                    child_module.requires_grad_(False)
-                    child_module.weight.copy_(result)
-                if result_b is not None:
-                    child_module.bias.copy_(result_b)
+            lora_name = prefix + "." + child_name
+            lora_name = lora_name.replace(".", "_")
+
+            lyco_type, params = get_module(lyco_state_dict, lora_name)
+            if lyco_type is None:
+                continue
+            module = make_module(lyco_type, params, lora_name, child_module)
+            if module is None:
+                continue
+            module.to(device)
+            module.merge_to(scale)
+            key_dict.pop(convert_diffusers_name_to_compvis(lora_name), None)
+            key_dict.pop(lora_name, None)
+            merged += 1
 
     key_dict = {}
     for k, v in tqdm(list(lyco_state_dict.items()), desc="Converting Dtype and Device"):
@@ -604,12 +460,7 @@ def merge_state_dict(
             k = f"{convert_key}.{weight_key}"
         else:
             key_dict[module] = key_dict.get(module, []) + [k]
-        if device == "cpu":
-            lyco_state_dict[k] = v.float().cpu()
-        else:
-            lyco_state_dict[k] = v.to(
-                device, dtype=tes[0].parameters().__next__().dtype
-            )
+        lyco_state_dict[k] = v.float().cpu()
 
     for idx, te in enumerate(tes):
         if len(tes) > 1:
@@ -618,15 +469,15 @@ def merge_state_dict(
             prefix = LORA_PREFIX_TEXT_ENCODER
         merge_state_dict(
             prefix,
-            te.to(device),
+            te,
             lyco_state_dict,
-            TEXT_ENCODER_TARGET_REPLACE_MODULE,
         )
+        torch.cuda.empty_cache()
     merge_state_dict(
         LORA_PREFIX_UNET,
-        unet.to(device),
+        unet,
         lyco_state_dict,
-        UNET_TARGET_REPLACE_MODULE,
     )
+    torch.cuda.empty_cache()
     print(f"Unused state dict key: {key_dict}")
     print(f"{merged} Modules been merged")
diff --git a/lycoris/utils/bnb.py b/lycoris/utils/bnb.py
index c68a3f4..8f82afa 100644
--- a/lycoris/utils/bnb.py
+++ b/lycoris/utils/bnb.py
@@ -2,7 +2,10 @@
 
 try:
     from bitsandbytes.nn import LinearNF4, Linear8bitLt, LinearFP4
+
+    SUPPORT_BNB = True
 except:
+    SUPPORT_BNB = False
     import torch.nn as nn
 
     class LinearNF4(nn.Linear):
diff --git a/lycoris/utils/general.py b/lycoris/utils/general.py
new file mode 100644
index 0000000..ccfaaa1
--- /dev/null
+++ b/lycoris/utils/general.py
@@ -0,0 +1,5 @@
+def product(xs: list[int | float]):
+    res = 1
+    for x in xs:
+        res *= x
+    return res
diff --git a/lycoris/wrapper.py b/lycoris/wrapper.py
index b7996b2..4758ae3 100644
--- a/lycoris/wrapper.py
+++ b/lycoris/wrapper.py
@@ -1,11 +1,8 @@
 # General LyCORIS wrapper based on kohya-ss/sd-scripts' style
-
-import math
 import os
 import regex as re
-import sys
+import logging
 
-sys.setrecursionlimit(10000)
 from typing import List
 
 import torch
@@ -20,11 +17,11 @@
 from .modules.full import FullModule
 from .modules.diag_oft import DiagOFTModule
 from .modules.boft import ButterflyOFTModule
-from .modules import make_module
+from .modules import get_module, make_module
 
 from .config import PRESET
 from .utils.preset import read_preset
-from .utils import get_module, str_bool
+from .utils import str_bool
 from .logging import logger
 
 
@@ -39,9 +36,22 @@
     "diag-oft": DiagOFTModule,
     "boft": ButterflyOFTModule,
 }
+deprecated_arg_dict = {
+    "disable_conv_cp": "use_tucker",
+    "use_cp": "use_tucker",
+    "use_conv_cp": "use_tucker",
+    "constrain": "constraint",
+}
 
 
 def create_lycoris(module, multiplier, linear_dim, linear_alpha, **kwargs):
+    for key, value in kwargs.items():
+        if key in deprecated_arg_dict:
+            logger.warning(
+                f"{key} is deprecated. Please use {deprecated_arg_dict[key]} instead.",
+                stacklevel=2,
+            )
+            kwargs[deprecated_arg_dict[key]] = value
     if linear_dim is None:
         linear_dim = 4  # default
     conv_dim = int(kwargs.get("conv_dim", linear_dim) or linear_dim)
@@ -56,19 +66,18 @@ def create_lycoris(module, multiplier, linear_dim, linear_alpha, **kwargs):
         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:
-        logger.warning(
-            "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)
+    constraint = float(kwargs.get("constraint", 0) or 0)
     rescaled = str_bool(kwargs.get("rescaled", False))
     weight_decompose = str_bool(kwargs.get("dora_wd", False))
     full_matrix = str_bool(kwargs.get("full_matrix", False))
     bypass_mode = str_bool(kwargs.get("bypass_mode", False))
+    unbalanced_factorization = str_bool(kwargs.get("unbalanced_factorization", False))
+
+    if unbalanced_factorization:
+        logger.info("Unbalanced factorization for LoKr is enabled")
 
     if bypass_mode:
         logger.info("Bypass mode is enabled")
@@ -79,10 +88,6 @@ def create_lycoris(module, multiplier, linear_dim, linear_alpha, **kwargs):
     if full_matrix:
         logger.info("Full matrix mode for LoKr is enabled")
 
-    if algo == "glora" and conv_dim > 0:
-        conv_dim = 0
-        logger.info("Disable conv layer for GLoRA")
-
     preset = kwargs.get("preset", "full")
     if preset not in PRESET:
         preset = read_preset(preset)
@@ -93,20 +98,6 @@ def create_lycoris(module, multiplier, linear_dim, linear_alpha, **kwargs):
 
     logger.info(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)
-    ):
-        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."
-        )
-        logger.warning(warning_msg, stacklevel=2)
-
     network = LycorisNetwork(
         module,
         multiplier=multiplier,
@@ -124,17 +115,14 @@ def create_lycoris(module, multiplier, linear_dim, linear_alpha, **kwargs):
         decompose_both=kwargs.get("decompose_both", False),
         factor=kwargs.get("factor", -1),
         block_size=block_size,
-        constrain=constrain,
+        constraint=constraint,
         rescaled=rescaled,
         weight_decompose=weight_decompose,
         full_matrix=full_matrix,
         bypass_mode=bypass_mode,
+        unbalanced_factorization=unbalanced_factorization,
     )
 
-    if algo == "dylora":
-        # dylora didn't support scale weight norm yet
-        delattr(type(network), "apply_max_norm_regularization")
-
     return network
 
 
@@ -161,10 +149,14 @@ def create_lycoris_from_weights(multiplier, file, module, weights_sd=None, **kwa
         if lora_name in loras:
             loras[lora_name] = modules
 
+    original_level = logger.level
+    logger.setLevel(logging.ERROR)
     network = LycorisNetwork(module, init_only=True)
     network.multiplier = multiplier
     network.loras = []
+    logger.setLevel(original_level)
 
+    logger.info("Loading Modules from state dict...")
     for lora_name, orig_modules in loras.items():
         if orig_modules is None:
             continue
@@ -175,6 +167,7 @@ def create_lycoris_from_weights(multiplier, file, module, weights_sd=None, **kwa
             network.algo_table[module.__class__.__name__] = (
                 network.algo_table.get(module.__class__.__name__, 0) + 1
             )
+    logger.info(f"{len(network.loras)} Modules Loaded")
 
     for lora in network.loras:
         lora.multiplier = multiplier
@@ -184,7 +177,14 @@ def create_lycoris_from_weights(multiplier, file, module, weights_sd=None, **kwa
 
 class LycorisNetwork(torch.nn.Module):
     ENABLE_CONV = True
-    TARGET_REPLACE_MODULE = []
+    TARGET_REPLACE_MODULE = [
+        "Linear",
+        "Conv1d",
+        "Conv2d",
+        "Conv3d",
+        "GroupNorm",
+        "LayerNorm",
+    ]
     TARGET_REPLACE_NAME = []
     LORA_PREFIX = "lycoris"
     MODULE_ALGO_MAP = {}
@@ -224,6 +224,7 @@ def __init__(
     ) -> None:
         super().__init__()
         root_kwargs = kwargs
+        self.weights_sd = None
         if init_only:
             self.multiplier = 1
             self.lora_dim = 0
@@ -291,7 +292,9 @@ def create_single_module(
             if isinstance(module, torch.nn.Linear) and lora_dim > 0:
                 dim = dim or lora_dim
                 alpha = alpha or self.alpha
-            elif isinstance(module, torch.nn.Conv2d):
+            elif isinstance(
+                module, (torch.nn.Conv1d, torch.nn.Conv2d, torch.nn.Conv3d)
+            ):
                 k_size, *_ = module.kernel_size
                 if k_size == 1 and lora_dim > 0:
                     dim = dim or lora_dim
@@ -327,7 +330,7 @@ def create_modules_(
             lora_names = []
             for name, module in root_module.named_modules():
                 module_name = module.__class__.__name__
-                if module_name in self.MODULE_ALGO_MAP:
+                if module_name in self.MODULE_ALGO_MAP and module is not root_module:
                     next_config = self.MODULE_ALGO_MAP[module_name]
                     next_algo = next_config.get("algo", algo)
                     new_loras, new_lora_names = create_modules_(
@@ -338,7 +341,10 @@ def create_modules_(
                             loras[lora_name] = lora
                             lora_names.append(lora_name)
                     continue
-                lora_name = prefix + "." + name
+                if name:
+                    lora_name = prefix + "." + name
+                else:
+                    lora_name = prefix
                 lora_name = lora_name.replace(".", "_")
                 if lora_name in loras:
                     continue
@@ -407,8 +413,6 @@ def create_modules(
             )
         logger.info(f"module type table: {algo_table}")
 
-        self.weights_sd = None
-
         # assertion
         names = set()
         for lora in self.loras:
@@ -461,11 +465,12 @@ def merge_to(self, weight=1.0):
     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
+        for module in self.loras:
+            scaled, norm = module.apply_max_norm(max_norm_value, device)
+            if scaled is None:
+                continue
+            norms.append(norm)
+            key_scaled += scaled
 
         if key_scaled == 0:
             return key_scaled, 0, 0
diff --git a/requirements-dev.txt b/requirements-dev.txt
new file mode 100644
index 0000000..827aa3c
--- /dev/null
+++ b/requirements-dev.txt
@@ -0,0 +1,2 @@
+coverage
+parameterized
\ No newline at end of file
diff --git a/requirements-kohya.txt b/requirements-kohya.txt
new file mode 100644
index 0000000..2c23500
--- /dev/null
+++ b/requirements-kohya.txt
@@ -0,0 +1,41 @@
+accelerate==0.25.0
+transformers==4.36.2
+diffusers[torch]==0.25.0
+ftfy==6.1.1
+# albumentations==1.3.0
+opencv-python==4.7.0.68
+einops==0.7.0
+pytorch-lightning==1.9.0
+bitsandbytes==0.43.0
+prodigyopt==1.0
+lion-pytorch==0.0.6
+tensorboard
+safetensors==0.4.2
+# gradio==3.16.2
+altair==4.2.2
+easygui==0.98.3
+toml==0.10.2
+voluptuous==0.13.1
+huggingface-hub==0.20.1
+# for Image utils
+imagesize==1.4.1
+# for BLIP captioning
+# requests==2.28.2
+# timm==0.6.12
+# fairscale==0.4.13
+# for WD14 captioning (tensorflow)
+# tensorflow==2.10.1
+# for WD14 captioning (onnx)
+# onnx==1.15.0
+# onnxruntime-gpu==1.17.1
+# onnxruntime==1.17.1
+# for cuda 12.1(default 11.8)
+# onnxruntime-gpu --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-12/pypi/simple/
+
+# this is for onnx: 
+# protobuf==3.20.3
+# open clip for SDXL
+# open-clip-torch==2.20.0
+# For logging
+rich==13.7.0
+# for kohya_ss library
\ No newline at end of file
diff --git a/requirements.txt b/requirements.txt
index 4ae5621..ad8ac5b 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,10 +1,3 @@
 torch
-safetensors
-diffusers[torch]
-transformers
 toml
-einops
-git+https://github.com/kohya-ss/sd-scripts.git
-pillow>=10.3.0 # not directly required, pinned by Snyk to avoid a vulnerability
-requests>=2.32.0 # not directly required, pinned by Snyk to avoid a vulnerability
-urllib3>=2.2.2 # not directly required, pinned by Snyk to avoid a vulnerability
\ No newline at end of file
+einops
\ No newline at end of file
diff --git a/setup.py b/setup.py
index 689a56a..9b983f4 100644
--- a/setup.py
+++ b/setup.py
@@ -4,19 +4,26 @@
 setup(
     name="lycoris_lora",
     packages=find_packages(),
-    version="2.2.0.post3",
+    version="3.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",
-        "einops",
-        "toml",
-    ],
+    install_requires=["torch", "einops", "toml", "tqdm"],
     python_requires=">=3.10",
+    classifiers=[
+        "Development Status :: 5 - Production/Stable",
+        "Environment :: GPU :: NVIDIA CUDA :: 11.8",
+        "Environment :: GPU :: NVIDIA CUDA :: 12",
+        "Environment :: GPU :: NVIDIA CUDA :: 12 :: 12.1",
+        "Intended Audience :: Developers",
+        "Intended Audience :: End Users/Desktop",
+        "License :: OSI Approved :: Apache Software License",
+        "Programming Language :: Python",
+        "Programming Language :: Python :: 3.10",
+        "Programming Language :: Python :: 3.11",
+        "Programming Language :: Python :: 3.12",
+        "Operating System :: OS Independent",
+    ],
 )
diff --git a/tools/extract_locon.py b/tools/extract_locon.py
index dc1144a..cd58761 100644
--- a/tools/extract_locon.py
+++ b/tools/extract_locon.py
@@ -124,8 +124,8 @@ def 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
+from library.model_util import load_models_from_stable_diffusion_checkpoint
+from library.sdxl_model_util import load_models_from_sdxl_checkpoint
 
 import torch
 from safetensors.torch import save_file
diff --git a/tools/merge.py b/tools/merge.py
index 25f673d..2ba8d44 100644
--- a/tools/merge.py
+++ b/tools/merge.py
@@ -47,12 +47,12 @@ def get_args():
 
 
 from lycoris.utils import merge
-from lycoris.kohya.model_utils import (
+from library.model_util import (
     load_models_from_stable_diffusion_checkpoint,
     save_stable_diffusion_checkpoint,
     load_file,
 )
-from lycoris.kohya.sdxl_model_util import (
+from library.sdxl_model_util import (
     load_models_from_sdxl_checkpoint,
     save_stable_diffusion_checkpoint as save_sdxl_checkpoint,
 )
@@ -64,7 +64,7 @@ def get_args():
 def main():
     if args.is_sdxl:
         base = load_models_from_sdxl_checkpoint(
-            None, args.base_model, map_location=args.device
+            None, args.base_model, map_location="cpu"
         )
     else:
         base = load_models_from_stable_diffusion_checkpoint(args.is_v2, args.base_model)
@@ -105,6 +105,7 @@ def main():
             None,
             base[2],
             getattr(base[1], "logit_scale", None),
+            None,
             dtype,
         )
     else:
@@ -116,6 +117,7 @@ def main():
             None,
             0,
             0,
+            None,
             dtype,
             base[1],
         )
diff --git a/tools/sdxl_emb.py b/tools/sdxl_emb.py
index e270921..f163864 100644
--- a/tools/sdxl_emb.py
+++ b/tools/sdxl_emb.py
@@ -30,7 +30,7 @@ def get_args():
 args = ARGS = get_args()
 
 
-from lycoris.kohya.sdxl_model_util import (
+from library.sdxl_model_util import (
     load_models_from_sdxl_checkpoint,
 )
 


Test patch:
diff --git a/experiments/dora_extract_test.py b/experiments/dora_extract_test.py
new file mode 100644
index 0000000..3a493ee
--- /dev/null
+++ b/experiments/dora_extract_test.py
@@ -0,0 +1,37 @@
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+
+
+def input_norm(x):
+    return x.transpose(0, 1).norm(dim=1, keepdim=True).transpose(0, 1)
+
+
+RANK = 1
+
+
+w1 = torch.randn(128, 144) * math.sqrt(2 / 144)
+w2 = torch.randn(128, 144) * math.sqrt(2 / 144)
+diff_w = w2 - w1
+u, s, vh = torch.linalg.svd(diff_w)
+low_rank_diff = u[:, :RANK] @ torch.diag(s[:RANK]) @ vh[:RANK, :]
+reconstruction1 = w1 + low_rank_diff
+
+w1_norm = input_norm(w1)
+w2_norm = input_norm(w2)
+
+normed_w2 = w2 / w2_norm
+normed_w1 = w1 / w1_norm
+normed_diff_w = normed_w2 - normed_w1
+
+u, s, vh = torch.linalg.svd(normed_diff_w)
+low_rank_normed_diff = u[:, :RANK] @ torch.diag(s[:RANK]) @ vh[:RANK, :]
+normed_reconstruction = normed_w1 + low_rank_normed_diff
+reconstruction2 = normed_reconstruction / input_norm(normed_reconstruction) * w2_norm
+
+
+print(F.mse_loss(reconstruction1, w2))
+print(F.mse_loss(reconstruction2, w2))
diff --git a/test/__init__.py b/test/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/test/functional.py b/test/functional.py
new file mode 100644
index 0000000..fc27e07
--- /dev/null
+++ b/test/functional.py
@@ -0,0 +1,87 @@
+import unittest
+from itertools import product
+from parameterized import parameterized
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from lycoris.functional import locon, loha, lokr, diag_oft, boft
+
+
+EPS_DTYPE = {
+    torch.float32: 5e-6,
+    torch.float16: 5e-5,
+    torch.bfloat16: 5e-4,
+}
+
+
+modules = [locon, loha, lokr, diag_oft, boft]
+base_module_and_input_adn_weight = [
+    lambda dim: (F.linear, torch.randn(dim, dim), torch.randn(1, dim)),
+    lambda dim: (F.conv1d, torch.randn(dim, dim, 3), torch.randn(1, dim, 16)),
+    lambda dim: (F.conv2d, torch.randn(dim, dim, 3, 3), torch.randn(1, dim, 16, 16)),
+    lambda dim: (
+        F.conv3d,
+        torch.randn(dim, dim, 3, 3, 3),
+        torch.randn(1, dim, 16, 16, 16),
+    ),
+]
+device_and_dtype = [
+    (torch.device("cpu"), torch.float32),
+]
+
+if torch.cuda.is_available():
+    device_and_dtype.append((torch.device("cuda"), torch.float32))
+    device_and_dtype.append((torch.device("cuda"), torch.float16))
+    device_and_dtype.append((torch.device("cuda"), torch.bfloat16))
+
+if torch.backends.mps.is_available():
+    device_and_dtype.append((torch.device("mps"), torch.float32))
+
+
+patch_forward_param_list = list(
+    product(
+        modules,
+        base_module_and_input_adn_weight,
+        device_and_dtype,
+    )
+)
+
+
+class LycorisFunctionalTests(unittest.TestCase):
+    @parameterized.expand(patch_forward_param_list)
+    def test_lycoris_functional(self, module, base, device_dtype):
+        func, test_weight, test_input = base(16)
+        device, dtype = device_dtype
+        print(
+            f"{module.__name__: <27}",
+            f"{func.__name__: <7}",
+            f"device={str(device): <5}",
+            f"dtype={str(dtype): <15}",
+            sep="||",
+        )
+
+        w = test_weight.to(device, dtype)
+        x = test_input.to(device, dtype)
+        y = func(x, w)
+
+        params = list(module.weight_gen(w, 4))
+        for idx, param in enumerate(params):
+            if param is not None:
+                param = param.to(device, dtype)
+                params[idx] = param + torch.randn_like(param) * 0.01
+
+        if module in {boft, diag_oft}:
+            diff_w = module.diff_weight(w, *params)
+            diff_y = module.bypass_forward_diff(y, *params, need_transpose=w.ndim > 2)
+        else:
+            diff_w = module.diff_weight(*params)
+            diff_y = module.bypass_forward_diff(x, *params)
+
+        diff_y_from_diff_w = func(x, diff_w.to(x))
+        self.assertTrue(
+            F.mse_loss(diff_y, diff_y_from_diff_w).item() < EPS_DTYPE[dtype],
+            f"Error: {module.__name__} {base.__name__} {device} {dtype} ||"
+            f"{F.mse_loss(diff_y, diff_y_from_diff_w).item()}",
+        )
diff --git a/test/module.py b/test/module.py
new file mode 100644
index 0000000..a75e755
--- /dev/null
+++ b/test/module.py
@@ -0,0 +1,181 @@
+import unittest
+from itertools import product
+from parameterized import parameterized
+
+import torch
+import torch.nn as nn
+
+from lycoris.modules import (
+    LycorisBaseModule,
+    LoConModule,
+    LohaModule,
+    LokrModule,
+    FullModule,
+    DiagOFTModule,
+    ButterflyOFTModule,
+    GLoRAModule,
+    DyLoraModule,
+    IA3Module,
+)
+
+
+modules: list[LycorisBaseModule] = [
+    LoConModule,
+    LohaModule,
+    LokrModule,
+    FullModule,
+    DiagOFTModule,
+    ButterflyOFTModule,
+    GLoRAModule,
+    DyLoraModule,
+    IA3Module,
+]
+base_module_and_input = [
+    lambda dim: (nn.Linear(dim, dim), torch.randn(1, dim)),
+    lambda dim: (nn.Conv1d(dim, dim, 3, 1, 1), torch.randn(1, dim, 16)),
+    lambda dim: (nn.Conv2d(dim, dim, (3, 3), 1, 1), torch.randn(1, dim, 16, 16)),
+    lambda dim: (nn.Conv3d(dim, dim, (3, 3, 3), 1, 1), torch.randn(1, dim, 16, 16, 16)),
+]
+device_and_dtype = [
+    (torch.device("cpu"), torch.float32),
+]
+weight_decompose = [False, True]
+use_tucker = [False, True]
+use_scalar = [False, True]
+
+if torch.cuda.is_available():
+    device_and_dtype.append((torch.device("cuda"), torch.float32))
+    device_and_dtype.append((torch.device("cuda"), torch.float16))
+    device_and_dtype.append((torch.device("cuda"), torch.bfloat16))
+
+if torch.backends.mps.is_available():
+    device_and_dtype.append((torch.device("mps"), torch.float32))
+
+
+patch_forward_param_list = list(
+    product(
+        modules,
+        base_module_and_input,
+        device_and_dtype,
+        weight_decompose,
+        use_tucker,
+        use_scalar,
+    )
+)
+
+
+class LycorisModuleTests(unittest.TestCase):
+    @parameterized.expand(patch_forward_param_list)
+    def test_lycoris_modules(self, module, base, device_dtype, wd, tucker, scalar):
+        base, test_input = base(16)
+        device, dtype = device_dtype
+        print(
+            f"{module.__name__: <18}",
+            f"{base.__class__.__name__: <7}",
+            f"device={str(device): <5}",
+            f"dtype={str(dtype): <15}",
+            f"wd={str(wd): <6}",
+            f"tucker={str(tucker): <6}",
+            f"scalar={str(scalar): <6}",
+            sep="|| ",
+        )
+        base = base.to(device, dtype)
+        test_input = test_input.to(device, dtype)
+        net: LycorisBaseModule = module(
+            "test",
+            base,
+            multiplier=1,
+            lora_dim=4,
+            alpha=1,
+            weight_decompose=wd,
+            use_tucker=tucker,
+            use_scalar=scalar,
+        ).to(device, dtype)
+        net.apply_to()
+
+        with torch.autocast("cuda", dtype=dtype):
+            test_output = base(test_input)
+        torch.sum(test_output).backward()
+        net.apply_max_norm(1.0)
+        state_dict = net.state_dict()
+        net.load_state_dict(state_dict)
+        net.restore()
+        net.merge_to()
+
+    @parameterized.expand(patch_forward_param_list)
+    def test_lycoris_modules_bypass_mode(
+        self, module, base, device_dtype, wd, tucker, scalar
+    ):
+        base, test_input = base(16)
+        if module == FullModule:
+            # Full module not support bypass forward
+            return
+        device, dtype = device_dtype
+        print(
+            f"{module.__name__: <18}",
+            f"{base.__class__.__name__: <7}",
+            f"device={str(device): <5}",
+            f"dtype={str(dtype): <15}",
+            f"wd={str(wd): <6}",
+            f"tucker={str(tucker): <6}",
+            f"scalar={str(scalar): <6}",
+            sep="|| ",
+        )
+        base = base.to(device, dtype)
+        test_input = test_input.to(device, dtype)
+        net: LycorisBaseModule = module(
+            "test",
+            base,
+            multiplier=1,
+            lora_dim=4,
+            alpha=1,
+            weight_decompose=wd,
+            use_tucker=tucker,
+            use_scalar=scalar,
+            bypass_mode=True,
+        ).to(device, dtype)
+        net.apply_to()
+
+        with torch.autocast("cuda", dtype=dtype):
+            test_output = base(test_input)
+        torch.sum(test_output).backward()
+        state_dict = net.state_dict()
+        net.load_state_dict(state_dict)
+
+    @parameterized.expand(patch_forward_param_list)
+    def test_lycoris_modules_parametrize(
+        self, module, base, device_dtype, wd, tucker, scalar
+    ):
+        base, test_input = base(16)
+        if module == FullModule:
+            # Full module not support bypass forward
+            return
+        device, dtype = device_dtype
+        print(
+            f"{module.__name__: <18}",
+            f"{base.__class__.__name__: <7}",
+            f"device={str(device): <5}",
+            f"dtype={str(dtype): <15}",
+            f"wd={str(wd): <6}",
+            f"tucker={str(tucker): <6}",
+            f"scalar={str(scalar): <6}",
+            sep="|| ",
+        )
+        base = base.to(device, dtype)
+        test_input = test_input.to(device, dtype)
+        net = module.parametrize(
+            base,
+            "weight",
+            1,
+            4,
+            1,
+            weight_decompose=wd,
+            use_tucker=tucker,
+            use_scalar=scalar,
+        ).to(device, dtype)
+
+        with torch.autocast("cuda", dtype=dtype):
+            test_output = base(test_input)
+        torch.sum(test_output).backward()
+        state_dict = net.state_dict()
+        net.load_state_dict(state_dict)
diff --git a/test/wrapper.py b/test/wrapper.py
new file mode 100644
index 0000000..92e5aa3
--- /dev/null
+++ b/test/wrapper.py
@@ -0,0 +1,128 @@
+import unittest
+from itertools import product
+from parameterized import parameterized
+
+import torch
+import torch.nn as nn
+
+from lycoris import create_lycoris, create_lycoris_from_weights, LycorisNetwork
+
+
+class TestNetwork(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.conv3d = nn.Conv3d(dim, dim, (3, 3, 3), 1, 1)
+        self.conv2d = nn.Conv2d(dim, dim, (3, 3), 1, 1)
+        self.group_norm = nn.GroupNorm(4, dim)
+        self.conv1d = nn.Conv1d(dim, dim, 3, 1, 1)
+        self.layer_norm = nn.LayerNorm(dim)
+        self.linear = nn.Linear(dim, dim)
+
+    def forward(self, x):
+        h = self.conv3d(x)
+        h = h.flatten(2, 3)
+        h = self.group_norm(h)
+        h = self.conv2d(h)
+        h = h.flatten(2, 3)
+        h = self.conv1d(h)
+        h = h.transpose(-1, -2)
+        h = self.layer_norm(h)
+        h = self.linear(h)
+        return h
+
+
+algos: list[str] = [
+    "lora",
+    "loha",
+    "lokr",
+    "full",
+    "diag-oft",
+    "boft",
+    "glora",
+    # "ia3",
+]
+device_and_dtype = [
+    (torch.device("cpu"), torch.float32),
+]
+weight_decompose = [
+    False,
+    True,
+]
+use_tucker = [
+    False,
+    True,
+]
+use_scalar = [
+    False,
+    True,
+]
+
+if torch.cuda.is_available():
+    device_and_dtype.append((torch.device("cuda"), torch.float32))
+    device_and_dtype.append((torch.device("cuda"), torch.float16))
+    device_and_dtype.append((torch.device("cuda"), torch.bfloat16))
+
+if torch.backends.mps.is_available():
+    device_and_dtype.append((torch.device("mps"), torch.float32))
+
+
+patch_forward_param_list = list(
+    product(
+        algos,
+        device_and_dtype,
+        weight_decompose,
+        use_tucker,
+        use_scalar,
+    )
+)
+
+
+class LycorisWrapperTests(unittest.TestCase):
+    @parameterized.expand(patch_forward_param_list)
+    def test_lycoris_wrapper(self, algo, device_dtype, wd, tucker, scalar):
+        device, dtype = device_dtype
+        print(
+            f"{algo: <18}",
+            f"device={str(device): <5}",
+            f"dtype={str(dtype): <15}",
+            f"wd={str(wd): <6}",
+            f"tucker={str(tucker): <6}",
+            f"scalar={str(scalar): <6}",
+            sep="|| ",
+        )
+        test_net = TestNetwork(16).to(device, dtype)
+        test_lycoris: LycorisNetwork = create_lycoris(
+            test_net,
+            1,
+            algo=algo,
+            linear_dim=4,
+            linear_alpha=2.0,
+            conv_dim=4,
+            conv_alpha=2.0,
+            dropout=0.0,
+            rank_dropout=0.0,
+            weight_decompose=wd,
+            use_tucker=tucker,
+            use_scalar=scalar,
+            train_norm=True,
+        )
+        test_lycoris.apply_to()
+        test_lycoris.to(device, dtype)
+
+        test_input = torch.randn(1, 16, 8, 8, 8).to(device, dtype)
+        test_output = test_net(test_input)
+        test_lycoris.restore()
+
+        state_dict = test_lycoris.state_dict()
+        test_lycoris_from_weights: LycorisNetwork
+        test_lycoris_from_weights, _ = create_lycoris_from_weights(
+            1, None, test_net, state_dict
+        )
+        test_lycoris_from_weights.apply_to()
+        test_lycoris_from_weights.to(device, dtype)
+        test_output_from_weights = test_net(test_input)
+
+        test_lycoris_from_weights.load_state_dict(test_lycoris.state_dict())
+
+        self.assertTrue(len(test_lycoris.loras) == len(test_lycoris_from_weights.loras))
+        self.assertTrue(torch.allclose(test_output, test_output_from_weights))
diff --git a/unit-test.py b/unit-test.py
new file mode 100644
index 0000000..a5b55e4
--- /dev/null
+++ b/unit-test.py
@@ -0,0 +1,30 @@
+import unittest
+import logging
+import coverage
+
+cov = coverage.Coverage()
+cov.start()
+
+from lycoris.logging import logger
+
+logger.setLevel(logging.ERROR)
+
+from test.module import LycorisModuleTests
+from test.wrapper import LycorisWrapperTests
+from test.functional import LycorisFunctionalTests
+
+
+TESTS = [LycorisModuleTests, LycorisWrapperTests, LycorisFunctionalTests]
+
+
+if __name__ == "__main__":
+    test_loader = unittest.TestLoader()
+    runner = unittest.TextTestRunner(verbosity=0)
+    for test in TESTS:
+        suite = test_loader.loadTestsFromTestCase(test)
+        result = runner.run(suite)
+
+    cov.stop()
+    cov.save()
+    cov.report()
+    cov.html_report(directory="coverage_report")