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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: lycoris_lora
-Version: 2.1.0.dev4
+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.dev4 → lycoris_lora-2.1.0.dev5}/lycoris/kohya/__init__.py

@@ -586,7 +586,9 @@ class LycorisNetworkKohya(LycorisNetwork):
         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)
 

{lycoris_lora-2.1.0.dev4 → 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.dev4 → 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.dev4 → 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.dev4 → 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.dev4 → lycoris_lora-2.1.0.dev5}/lycoris/utils/__init__.py

@@ -147,6 +147,7 @@ def extract_diff(
         "GroupNorm32",
     ]
     TEXT_ENCODER_TARGET_REPLACE_MODULE = [
+        "Embedding",
         "Linear",
         "Conv2d",
         "LayerNorm",
@@ -312,7 +313,11 @@ def extract_diff(
         )
         del te1, te2
 
-    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
 
 
@@ -420,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", ()
@@ -517,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
@@ -539,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"):
@@ -596,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:
@@ -617,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.dev4 → 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.dev4
+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.dev4 → 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.dev4",
+    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",