lycoris-lora 2.1.0.dev3__tar.gz → 2.1.0.dev5__tar.gz

{lycoris_lora-2.1.0.dev3 → lycoris_lora-2.1.0.dev5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: lycoris_lora
-Version: 2.1.0.dev3
+Version: 2.1.0.dev5
 Summary: Lora beYond Conventional methods, Other Rank adaptation Implementations for Stable diffusion
 Home-page: https://github.com/KohakuBlueleaf/LyCORIS
 Author: Shih-Ying Yeh(KohakuBlueLeaf), Yu-Guan Hsieh, Zhidong Gao

{lycoris_lora-2.1.0.dev3 → lycoris_lora-2.1.0.dev5}/lycoris/kohya/__init__.py

@@ -15,6 +15,7 @@ import torch
 import torch.utils.checkpoint as checkpoint
 
 from .utils import *
+from ..wrapper import LycorisNetwork
 from ..modules.locon import LoConModule
 from ..modules.loha import LohaModule
 from ..modules.ia3 import IA3Module
@@ -186,9 +187,7 @@ def create_network_from_weights(
         else:
             prefix = LycorisNetworkKohya.LORA_PREFIX_TEXT_ENCODER
         for name, modules in te.named_modules():
-            lora_name = f"{prefix}_{name}".replace(
-                ".", "_"
-            )
+            lora_name = f"{prefix}_{name}".replace(".", "_")
             if lora_name in te_loras:
                 te_loras[lora_name] = modules
 
@@ -279,7 +278,7 @@ def create_hypernetwork(
     )
 
 
-class LycorisNetworkKohya(torch.nn.Module):
+class LycorisNetworkKohya(LycorisNetwork):
     """
     LoRA + LoCon
     """
@@ -343,7 +342,7 @@ class LycorisNetworkKohya(torch.nn.Module):
         train_norm=False,
         **kwargs,
     ) -> None:
-        super().__init__()
+        torch.nn.Module.__init__(self)
         root_kwargs = kwargs
         self.multiplier = multiplier
         self.lora_dim = lora_dim
@@ -548,19 +547,15 @@ class LycorisNetworkKohya(torch.nn.Module):
 
         self.weights_sd = None
 
+        self.loras = self.text_encoder_loras + self.unet_loras
         # assertion
         names = set()
-        for lora in self.text_encoder_loras + self.unet_loras:
+        for lora in self.loras:
             assert (
                 lora.lora_name not in names
             ), f"duplicated lora name: {lora.lora_name}"
             names.add(lora.lora_name)
 
-    def set_multiplier(self, multiplier):
-        self.multiplier = multiplier
-        for lora in self.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
@@ -591,7 +586,9 @@ class LycorisNetworkKohya(torch.nn.Module):
         else:
             self.unet_loras = []
 
-        for lora in self.text_encoder_loras + self.unet_loras:
+        self.loras = self.text_encoder_loras + self.unet_loras
+
+        for lora in self.loras:
             lora.apply_to()
             self.add_module(lora.lora_name, lora)
 
@@ -600,34 +597,27 @@ class LycorisNetworkKohya(torch.nn.Module):
             info = self.load_state_dict(self.weights_sd, False)
             print(f"weights are loaded: {info}")
 
-    def apply_max_norm_regularization(self, max_norm_value, device):
-        key_scaled = 0
-        norms = []
-        for model in self.unet_loras:
-            if hasattr(model, "apply_max_norm"):
-                scaled, norm = model.apply_max_norm(max_norm_value, device)
-                norms.append(norm)
-                key_scaled += scaled
-
-        for model in self.text_encoder_loras:
-            if hasattr(model, "apply_max_norm"):
-                scaled, norm = model.apply_max_norm(max_norm_value, device)
-                norms.append(norm)
-                key_scaled += scaled
+    # TODO refactor to common function with apply_to
+    def merge_to(self, text_encoder, unet, weights_sd, dtype, device):
+        apply_text_encoder = apply_unet = False
+        for key in weights_sd.keys():
+            if key.startswith(LycorisNetworkKohya.LORA_PREFIX_TEXT_ENCODER):
+                apply_text_encoder = True
+            elif key.startswith(LycorisNetworkKohya.LORA_PREFIX_UNET):
+                apply_unet = True
 
-        if key_scaled == 0:
-            return key_scaled, 0, 0
-
-        return key_scaled, sum(norms) / len(norms), max(norms)
+        if apply_text_encoder:
+            print("enable LoRA for text encoder")
+        else:
+            self.text_encoder_loras = []
 
-    def enable_gradient_checkpointing(self):
-        # not supported
-        def make_ckpt(module):
-            if isinstance(module, torch.nn.Module):
-                module.grad_ckpt = True
+        if apply_unet:
+            print("enable LoRA for U-Net")
+        else:
+            self.unet_loras = []
 
-        self.apply(make_ckpt)
-        pass
+        self.loras = self.text_encoder_loras + self.unet_loras
+        super().merge_to(1)
 
     def prepare_optimizer_params(self, text_encoder_lr, unet_lr, learning_rate):
         def enumerate_params(loras):
@@ -653,15 +643,6 @@ class LycorisNetworkKohya(torch.nn.Module):
 
         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

{lycoris_lora-2.1.0.dev3 → lycoris_lora-2.1.0.dev5}/lycoris/modules/__init__.py

@@ -99,7 +99,8 @@ def make_module(lyco_type, params, lora_name, orig_module):
             1,
         )
         module.w_norm.copy_(w_norm)
-        module.b_norm.copy_(b_norm)
+        if b_norm is not None:
+            module.b_norm.copy_(b_norm)
     elif lyco_type == "full":
         diff, diff_b = params
         module = FullModule(

{lycoris_lora-2.1.0.dev3 → lycoris_lora-2.1.0.dev5}/lycoris/modules/base.py

@@ -1,10 +1,31 @@
 from collections import OrderedDict
+from typing import Any, Mapping
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
 
 class ModuleCustomSD(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self._register_load_state_dict_pre_hook(self.load_weight_prehook)
+        self.register_load_state_dict_post_hook(self.load_weight_hook)
+
+    def load_weight_prehook(
+        self,
+        state_dict,
+        prefix,
+        local_metadata,
+        strict,
+        missing_keys,
+        unexpected_keys,
+        error_msgs,
+    ):
+        pass
+
+    def load_weight_hook(self, module, incompatible_keys):
+        pass
+
     def custom_state_dict(self):
         return None
 

{lycoris_lora-2.1.0.dev3 → lycoris_lora-2.1.0.dev5}/lycoris/modules/dylora.py

@@ -103,9 +103,7 @@ class DyLoraModule(ModuleCustomSD):
         self.org_module[0].forward = self.org_forward
 
     def merge_to(self, multiplier=1.0):
-        up = torch.concat(
-            list(self.up_list), dim=1
-        )
+        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)

{lycoris_lora-2.1.0.dev3 → lycoris_lora-2.1.0.dev5}/lycoris/modules/full.py

@@ -1,6 +1,7 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+from torch.nn.modules import Module
 
 from .base import ModuleCustomSD
 
@@ -54,11 +55,11 @@ class FullModule(ModuleCustomSD):
     def apply_to(self, **kwargs):
         self.org_forward = self.org_module[0].forward
         self.org_module[0].forward = self.forward
-        self.weight.data.copy_(self.org_module[0].weight.data)
+        self.weight.data.add_(self.org_module[0].weight.data)
         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.copy_(self.org_module[0].bias.data)
+            self.bias.data.add_(self.org_module[0].bias.data)
             self.org_bias = [self.org_module[0].bias.data.cpu().clone()]
             delattr(self.org_module[0], "bias")
         else:
@@ -76,6 +77,22 @@ class FullModule(ModuleCustomSD):
             sd["diff_b"] = self.bias.data.cpu() - self.org_bias[0]
         return sd
 
+    def load_weight_prehook(
+        self,
+        state_dict,
+        prefix,
+        local_metadata,
+        strict,
+        missing_keys,
+        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)
+
     def make_weight(self, scale=1, device=None, original=False):
         if original:
             weight = self.org_weight[0].to(device, dtype=self.weight.dtype)

{lycoris_lora-2.1.0.dev3 → lycoris_lora-2.1.0.dev5}/lycoris/modules/loha.py

@@ -191,9 +191,11 @@ class LohaModule(ModuleCustomSD):
         self.org_module = [org_module]  # remove in applying
         self.org_forward = self.org_module[0].forward
         self.grad_ckpt = False
-        self.register_load_state_dict_post_hook(self.load_weight_hook)
 
-    def load_weight_hook(self, *args, **kwargs):
+    def load_weight_hook(self, module: nn.Module, incompatible_keys):
+        unexpected_keys = incompatible_keys.unexpected_keys
+        if "scalar" in unexpected_keys:
+            del unexpected_keys[unexpected_keys.index("scalar")]
         self.scalar = nn.Parameter(torch.ones_like(self.scalar))
 
     def apply_to(self):

{lycoris_lora-2.1.0.dev3 → lycoris_lora-2.1.0.dev5}/lycoris/utils/__init__.py

@@ -67,7 +67,7 @@ def extract_conv(
 
     U = U[:, :lora_rank]
     S = S[:lora_rank]
-    U = U @ torch.diag(S)
+    U = U @ torch.diag(S).to(device)
     Vh = Vh[:lora_rank, :]
 
     diff = (weight - (U @ Vh).reshape(out_ch, in_ch, kernel_size, kernel_size)).detach()
@@ -115,7 +115,7 @@ def extract_linear(
 
     U = U[:, :lora_rank]
     S = S[:lora_rank]
-    U = U @ torch.diag(S)
+    U = U @ torch.diag(S).to(device)
     Vh = Vh[:lora_rank, :]
 
     diff = (weight - U @ Vh).detach()
@@ -125,6 +125,7 @@ def extract_linear(
     return (extract_weight_A, extract_weight_B, diff), "low rank"
 
 
+@torch.no_grad()
 def extract_diff(
     base_tes,
     db_tes,
@@ -139,19 +140,20 @@ def extract_diff(
     small_conv=True,
 ):
     UNET_TARGET_REPLACE_MODULE = [
-        "Transformer2DModel",
-        "Attention",
-        "ResnetBlock2D",
-        "Downsample2D",
-        "Upsample2D",
+        "Linear",
+        "Conv2d",
+        "LayerNorm",
+        "GroupNorm",
+        "GroupNorm32",
     ]
-    UNET_TARGET_REPLACE_NAME = [
-        "conv_in",
-        "conv_out",
-        "time_embedding.linear_1",
-        "time_embedding.linear_2",
+    TEXT_ENCODER_TARGET_REPLACE_MODULE = [
+        "Embedding",
+        "Linear",
+        "Conv2d",
+        "LayerNorm",
+        "GroupNorm",
+        "GroupNorm32",
     ]
-    TEXT_ENCODER_TARGET_REPLACE_MODULE = ["CLIPAttention", "CLIPMLP"]
     LORA_PREFIX_UNET = "lora_unet"
     LORA_PREFIX_TEXT_ENCODER = "lora_te"
 
@@ -160,236 +162,162 @@ def extract_diff(
         root_module: torch.nn.Module,
         target_module: torch.nn.Module,
         target_replace_modules,
-        target_replace_names=[],
     ):
         loras = {}
         temp = {}
-        temp_name = {}
 
         for name, module in root_module.named_modules():
             if module.__class__.__name__ in target_replace_modules:
-                temp[name] = {}
-                for child_name, child_module in module.named_modules():
-                    if child_module.__class__.__name__ not in {"Linear", "Conv2d"}:
-                        continue
-                    temp[name][child_name] = child_module
-            elif name in target_replace_names:
-                temp_name[name] = module
-
-        for name, module in tqdm(list(target_module.named_modules())):
-            if name in temp:
-                weights = temp[name]
-                for child_name, child_module in module.named_modules():
-                    lora_name = prefix + "." + name + "." + child_name
-                    lora_name = lora_name.replace(".", "_")
-                    layer = child_module.__class__.__name__
-                    if layer in {"Linear", "Conv2d"}:
-                        root_weight = child_module.weight
-                        if torch.allclose(root_weight, weights[child_name].weight):
-                            continue
-
-                    if layer == "Linear":
-                        weight, decompose_mode = extract_linear(
-                            (child_module.weight - weights[child_name].weight),
-                            mode,
-                            linear_mode_param,
-                            device=extract_device,
-                        )
-                        if decompose_mode == "low rank":
-                            extract_a, extract_b, diff = weight
-                    elif layer == "Conv2d":
-                        is_linear = (
-                            child_module.weight.shape[2] == 1
-                            and child_module.weight.shape[3] == 1
-                        )
-                        weight, decompose_mode = extract_conv(
-                            (child_module.weight - weights[child_name].weight),
-                            mode,
-                            linear_mode_param if is_linear else conv_mode_param,
-                            device=extract_device,
-                        )
-                        if decompose_mode == "low rank":
-                            extract_a, extract_b, diff = weight
-                        if (
-                            small_conv
-                            and not is_linear
-                            and decompose_mode == "low rank"
-                        ):
-                            dim = extract_a.size(0)
-                            (extract_c, extract_a, _), _ = extract_conv(
-                                extract_a.transpose(0, 1),
-                                "fixed",
-                                dim,
-                                extract_device,
-                                True,
-                            )
-                            extract_a = extract_a.transpose(0, 1)
-                            extract_c = extract_c.transpose(0, 1)
-                            loras[f"{lora_name}.lora_mid.weight"] = (
-                                extract_c.detach().cpu().contiguous().half()
-                            )
-                            diff = (
-                                child_module.weight
-                                - torch.einsum(
-                                    "i j k l, j r, p i -> p r k l",
-                                    extract_c,
-                                    extract_a.flatten(1, -1),
-                                    extract_b.flatten(1, -1),
-                                )
-                                .detach()
-                                .cpu()
-                                .contiguous()
-                            )
-                            del extract_c
-                    else:
-                        continue
-                    if decompose_mode == "low rank":
-                        loras[f"{lora_name}.lora_down.weight"] = (
-                            extract_a.detach().cpu().contiguous().half()
-                        )
-                        loras[f"{lora_name}.lora_up.weight"] = (
-                            extract_b.detach().cpu().contiguous().half()
-                        )
-                        loras[f"{lora_name}.alpha"] = torch.Tensor(
-                            [extract_a.shape[0]]
-                        ).half()
-                        if use_bias:
-                            diff = diff.detach().cpu().reshape(extract_b.size(0), -1)
-                            sparse_diff = (
-                                make_sparse(diff, sparsity).to_sparse().coalesce()
-                            )
-
-                            indices = sparse_diff.indices().to(torch.int16)
-                            values = sparse_diff.values().half()
-                            loras[f"{lora_name}.bias_indices"] = indices
-                            loras[f"{lora_name}.bias_values"] = values
-                            loras[f"{lora_name}.bias_size"] = torch.tensor(
-                                diff.shape
-                            ).to(torch.int16)
-                        del extract_a, extract_b, diff
-                    elif decompose_mode == "full":
-                        loras[f"{lora_name}.diff"] = (
-                            weight.detach().cpu().contiguous().half()
-                        )
-                        if weights[child_name].bias is not None:
-                            bias_diff = child_module.bias - weights[child_name].bias
-                            loras[f"{lora_name}.diff_b"] = (
-                                bias_diff.detach().cpu().contiguous().half()
-                            )
-                    else:
-                        raise NotImplementedError
-            elif name in temp_name:
-                weights = temp_name[name]
-                lora_name = prefix + "." + name
-                lora_name = lora_name.replace(".", "_")
-                layer = module.__class__.__name__
-
-                if layer in {"Linear", "Conv2d"}:
-                    root_weight = module.weight
-                    if torch.allclose(root_weight, weights.weight):
-                        continue
-
-                if layer == "Linear":
-                    weight, decompose_mode = extract_linear(
-                        (root_weight - weights.weight),
-                        mode,
-                        linear_mode_param,
-                        device=extract_device,
-                    )
-                    if decompose_mode == "low rank":
-                        extract_a, extract_b, diff = weight
-                elif layer == "Conv2d":
-                    is_linear = root_weight.shape[2] == 1 and root_weight.shape[3] == 1
-                    weight, decompose_mode = extract_conv(
-                        (root_weight - weights.weight),
-                        mode,
-                        linear_mode_param if is_linear else conv_mode_param,
-                        device=extract_device,
-                    )
-                    if decompose_mode == "low rank":
-                        extract_a, extract_b, diff = weight
-                    if small_conv and not is_linear and decompose_mode == "low rank":
-                        dim = extract_a.size(0)
-                        (extract_c, extract_a, _), _ = extract_conv(
-                            extract_a.transpose(0, 1),
-                            "fixed",
-                            dim,
-                            extract_device,
-                            True,
-                        )
-                        extract_a = extract_a.transpose(0, 1)
-                        extract_c = extract_c.transpose(0, 1)
-                        loras[f"{lora_name}.lora_mid.weight"] = (
-                            extract_c.detach().cpu().contiguous().half()
-                        )
-                        diff = (
-                            root_weight
-                            - torch.einsum(
-                                "i j k l, j r, p i -> p r k l",
-                                extract_c,
-                                extract_a.flatten(1, -1),
-                                extract_b.flatten(1, -1),
-                            )
-                            .detach()
-                            .cpu()
-                            .contiguous()
-                        )
-                        del extract_c
-                else:
+                temp[name] = module
+
+        for name, module in tqdm(
+            list((n, m) for n, m in target_module.named_modules() if n in temp)
+        ):
+            weights = temp[name]
+            lora_name = prefix + "." + name
+            lora_name = lora_name.replace(".", "_")
+            layer = module.__class__.__name__
+
+            if layer in {
+                "Linear",
+                "Conv2d",
+                "LayerNorm",
+                "GroupNorm",
+                "GroupNorm32",
+            }:
+                root_weight = module.weight
+                if torch.allclose(root_weight, weights.weight):
                     continue
+            else:
+                continue
+            module = module.to(extract_device)
+            weights = weights.to(extract_device)
+
+            if mode == "full":
+                decompose_mode = "full"
+            elif layer == "Linear":
+                weight, decompose_mode = extract_linear(
+                    (root_weight - weights.weight),
+                    mode,
+                    linear_mode_param,
+                    device=extract_device,
+                )
                 if decompose_mode == "low rank":
-                    loras[f"{lora_name}.lora_down.weight"] = (
-                        extract_a.detach().cpu().contiguous().half()
+                    extract_a, extract_b, diff = weight
+            elif layer == "Conv2d":
+                is_linear = root_weight.shape[2] == 1 and root_weight.shape[3] == 1
+                weight, decompose_mode = extract_conv(
+                    (root_weight - weights.weight),
+                    mode,
+                    linear_mode_param if is_linear else conv_mode_param,
+                    device=extract_device,
+                )
+                if decompose_mode == "low rank":
+                    extract_a, extract_b, diff = weight
+                if small_conv and not is_linear and decompose_mode == "low rank":
+                    dim = extract_a.size(0)
+                    (extract_c, extract_a, _), _ = extract_conv(
+                        extract_a.transpose(0, 1),
+                        "fixed",
+                        dim,
+                        extract_device,
+                        True,
                     )
-                    loras[f"{lora_name}.lora_up.weight"] = (
-                        extract_b.detach().cpu().contiguous().half()
+                    extract_a = extract_a.transpose(0, 1)
+                    extract_c = extract_c.transpose(0, 1)
+                    loras[f"{lora_name}.lora_mid.weight"] = (
+                        extract_c.detach().cpu().contiguous().half()
                     )
-                    loras[f"{lora_name}.alpha"] = torch.Tensor(
-                        [extract_a.shape[0]]
-                    ).half()
-                    if use_bias:
-                        diff = diff.detach().cpu().reshape(extract_b.size(0), -1)
-                        sparse_diff = make_sparse(diff, sparsity).to_sparse().coalesce()
-
-                        indices = sparse_diff.indices().to(torch.int16)
-                        values = sparse_diff.values().half()
-                        loras[f"{lora_name}.bias_indices"] = indices
-                        loras[f"{lora_name}.bias_values"] = values
-                        loras[f"{lora_name}.bias_size"] = torch.tensor(diff.shape).to(
-                            torch.int16
+                    diff = (
+                        root_weight
+                        - torch.einsum(
+                            "i j k l, j r, p i -> p r k l",
+                            extract_c,
+                            extract_a.flatten(1, -1),
+                            extract_b.flatten(1, -1),
                         )
-                    del extract_a, extract_b, diff
-                elif decompose_mode == "full":
-                    loras[f"{lora_name}.diff"] = (
-                        weight.detach().cpu().contiguous().half()
+                        .detach()
+                        .cpu()
+                        .contiguous()
                     )
-                    if weights.bias is not None:
-                        bias_diff = module.bias - weights.bias
-                        loras[f"{lora_name}.diff_b"] = (
-                            bias_diff.detach().cpu().contiguous().half()
-                        )
+                    del extract_c
+            else:
+                module = module.to("cpu")
+                weights = weights.to("cpu")
+                continue
+
+            if decompose_mode == "low rank":
+                loras[f"{lora_name}.lora_down.weight"] = (
+                    extract_a.detach().cpu().contiguous().half()
+                )
+                loras[f"{lora_name}.lora_up.weight"] = (
+                    extract_b.detach().cpu().contiguous().half()
+                )
+                loras[f"{lora_name}.alpha"] = torch.Tensor([extract_a.shape[0]]).half()
+                if use_bias:
+                    diff = diff.detach().cpu().reshape(extract_b.size(0), -1)
+                    sparse_diff = make_sparse(diff, sparsity).to_sparse().coalesce()
+
+                    indices = sparse_diff.indices().to(torch.int16)
+                    values = sparse_diff.values().half()
+                    loras[f"{lora_name}.bias_indices"] = indices
+                    loras[f"{lora_name}.bias_values"] = values
+                    loras[f"{lora_name}.bias_size"] = torch.tensor(diff.shape).to(
+                        torch.int16
+                    )
+                del extract_a, extract_b, diff
+            elif decompose_mode == "full":
+                if "Norm" in layer:
+                    w_key = "w_norm"
+                    b_key = "b_norm"
                 else:
-                    raise NotImplementedError
+                    w_key = "diff"
+                    b_key = "diff_b"
+                weight_diff = module.weight - weights.weight
+                loras[f"{lora_name}.{w_key}"] = (
+                    weight_diff.detach().cpu().contiguous().half()
+                )
+                if weights.bias is not None:
+                    bias_diff = module.bias - weights.bias
+                    loras[f"{lora_name}.{b_key}"] = (
+                        bias_diff.detach().cpu().contiguous().half()
+                    )
+            else:
+                raise NotImplementedError
+            module = module.to("cpu")
+            weights = weights.to("cpu")
         return loras
 
     all_loras = {}
+
+    all_loras |= make_state_dict(
+        LORA_PREFIX_UNET,
+        base_unet,
+        db_unet,
+        UNET_TARGET_REPLACE_MODULE,
+    )
+    del base_unet, db_unet
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
     for idx, (te1, te2) in enumerate(zip(base_tes, db_tes)):
         if len(base_tes) > 1:
             prefix = f"{LORA_PREFIX_TEXT_ENCODER}{idx+1}"
         else:
             prefix = LORA_PREFIX_TEXT_ENCODER
         all_loras |= make_state_dict(
-            prefix, te1, te2, TEXT_ENCODER_TARGET_REPLACE_MODULE
+            prefix,
+            te1,
+            te2,
+            TEXT_ENCODER_TARGET_REPLACE_MODULE,
         )
+        del te1, te2
 
-    all_loras |= make_state_dict(
-        LORA_PREFIX_UNET,
-        base_unet,
-        db_unet,
-        UNET_TARGET_REPLACE_MODULE,
-        UNET_TARGET_REPLACE_NAME,
-    )
-    print(len(all_loras))
+    all_lora_name = set()
+    for k in all_loras:
+        lora_name, weight = k.rsplit(".", 1)
+        all_lora_name.add(lora_name)
+    print(len(all_lora_name))
     return all_loras
 
 
@@ -497,7 +425,7 @@ def get_module(lyco_state_dict: Dict, lora_name):
         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")
+        b_norm = lyco_state_dict.get(f"{lora_name}.b_norm", None)
         return "norm", (w_norm, b_norm)
     else:
         return "None", ()
@@ -594,19 +522,20 @@ def rebuild_weight(module_type, params, orig_weight, orig_bias, scale=1):
 @torch.no_grad()
 def merge(tes, unet, lyco_state_dict, scale: float = 1.0, device="cpu"):
     UNET_TARGET_REPLACE_MODULE = [
-        "Transformer2DModel",
-        "Attention",
-        "ResnetBlock2D",
-        "Downsample2D",
-        "Upsample2D",
+        "Linear",
+        "Conv2d",
+        "LayerNorm",
+        "GroupNorm",
+        "GroupNorm32",
     ]
-    UNET_TARGET_REPLACE_NAME = [
-        "conv_in",
-        "conv_out",
-        "time_embedding.linear_1",
-        "time_embedding.linear_2",
+    TEXT_ENCODER_TARGET_REPLACE_MODULE = [
+        "Embedding",
+        "Linear",
+        "Conv2d",
+        "LayerNorm",
+        "GroupNorm",
+        "GroupNorm32",
     ]
-    TEXT_ENCODER_TARGET_REPLACE_MODULE = ["CLIPAttention", "CLIPMLP"]
     LORA_PREFIX_UNET = "lora_unet"
     LORA_PREFIX_TEXT_ENCODER = "lora_te"
     merged = 0
@@ -616,55 +545,28 @@ def merge(tes, unet, lyco_state_dict, scale: float = 1.0, device="cpu"):
         root_module: torch.nn.Module,
         lyco_state_dict: Dict[str, torch.Tensor],
         target_replace_modules,
-        target_replace_names=[],
     ):
         nonlocal merged
-        for name, module in tqdm(
+        for child_name, child_module in tqdm(
             list(root_module.named_modules()), desc=f"Merging {prefix}"
         ):
-            if module.__class__.__name__ in target_replace_modules:
-                for child_name, child_module in module.named_modules():
-                    if child_module.__class__.__name__ not in {
-                        "Linear",
-                        "Conv2d",
-                        "GroupNorm",
-                        "GroupNorm32",
-                        "LayerNorm",
-                    }:
-                        continue
-                    lora_name = prefix + "." + name + "." + child_name
-                    lora_name = lora_name.replace(".", "_")
-
-                    result, result_b = rebuild_weight(
-                        *get_module(lyco_state_dict, lora_name),
-                        getattr(child_module, "weight"),
-                        getattr(child_module, "bias", None),
-                        scale,
-                    )
-                    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)
-            elif name in target_replace_names:
-                lora_name = prefix + "." + name
+            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(module, "weight"),
-                    getattr(module, "bias", None),
+                    getattr(child_module, "weight"),
+                    getattr(child_module, "bias", None),
                     scale,
                 )
                 if result is not None:
                     key_dict.pop(lora_name)
                     merged += 1
-                    module.requires_grad_(False)
-                    module.weight.copy_(result)
+                    child_module.requires_grad_(False)
+                    child_module.weight.copy_(result)
                 if result_b is not None:
-                    module.bias.copy_(result_b)
+                    child_module.bias.copy_(result_b)
 
     key_dict = {}
     for k, v in tqdm(list(lyco_state_dict.items()), desc="Converting Dtype and Device"):
@@ -673,10 +575,10 @@ def merge(tes, unet, lyco_state_dict, scale: float = 1.0, device="cpu"):
         if convert_key != module and len(tes) > 1:
             # kohya's format for sdxl is as same as SGM, not diffusers
             del lyco_state_dict[k]
-            key_dict[convert_key] = k
+            key_dict[convert_key] = key_dict.get(convert_key, []) + [k]
             k = f"{convert_key}.{weight_key}"
         else:
-            key_dict[module] = k
+            key_dict[module] = key_dict.get(module, []) + [k]
         if device == "cpu":
             lyco_state_dict[k] = v.float().cpu()
         else:
@@ -685,7 +587,7 @@ def merge(tes, unet, lyco_state_dict, scale: float = 1.0, device="cpu"):
             )
 
     for idx, te in enumerate(tes):
-        if len(tes)>1:
+        if len(tes) > 1:
             prefix = LORA_PREFIX_TEXT_ENCODER + str(idx + 1)
         else:
             prefix = LORA_PREFIX_TEXT_ENCODER
@@ -694,14 +596,12 @@ def merge(tes, unet, lyco_state_dict, scale: float = 1.0, device="cpu"):
             te.to(device),
             lyco_state_dict,
             TEXT_ENCODER_TARGET_REPLACE_MODULE,
-            UNET_TARGET_REPLACE_NAME,
         )
     merge_state_dict(
         LORA_PREFIX_UNET,
         unet.to(device),
         lyco_state_dict,
         UNET_TARGET_REPLACE_MODULE,
-        UNET_TARGET_REPLACE_NAME,
     )
     print(f"Unused state dict key: {key_dict}")
     print(f"{merged} Modules been merged")

{lycoris_lora-2.1.0.dev3 → lycoris_lora-2.1.0.dev5}/lycoris/wrapper.py

@@ -432,6 +432,9 @@ class LycorisNetwork(torch.nn.Module):
             info = self.load_state_dict(self.weights_sd, False)
             print(f"weights are loaded: {info}")
 
+    def is_mergeable(self):
+        return True
+
     def restore(self):
         for lora in self.loras:
             lora.restore()
@@ -479,13 +482,13 @@ class LycorisNetwork(torch.nn.Module):
         all_params.append(param_data)
         return all_params
 
-    def prepare_grad_etc(self):
+    def prepare_grad_etc(self, *args):
         self.requires_grad_(True)
 
-    def on_epoch_start(self):
+    def on_epoch_start(self, *args):
         self.train()
 
-    def get_trainable_params(self):
+    def get_trainable_params(self, *args):
         return self.parameters()
 
     def save_weights(self, file, dtype, metadata):

{lycoris_lora-2.1.0.dev3 → lycoris_lora-2.1.0.dev5}/lycoris_lora.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: lycoris-lora
-Version: 2.1.0.dev3
+Version: 2.1.0.dev5
 Summary: Lora beYond Conventional methods, Other Rank adaptation Implementations for Stable diffusion
 Home-page: https://github.com/KohakuBlueleaf/LyCORIS
 Author: Shih-Ying Yeh(KohakuBlueLeaf), Yu-Guan Hsieh, Zhidong Gao

{lycoris_lora-2.1.0.dev3 → lycoris_lora-2.1.0.dev5}/setup.py

@@ -4,7 +4,7 @@ from setuptools import setup, find_packages
 setup(
     name="lycoris_lora",
     packages=find_packages(),
-    version="2.1.0.dev3",
+    version="2.1.0.dev5",
     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",