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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: lycoris_lora
-Version: 2.1.0.dev3
+Version: 2.1.0.dev4
 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.dev4}/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
@@ -600,34 +595,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
+    # 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
 
-        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
-
-        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 +641,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.dev4}/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.dev4}/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,19 @@ 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 = [
+        "Linear",
+        "Conv2d",
+        "LayerNorm",
+        "GroupNorm",
+        "GroupNorm32",
     ]
-    TEXT_ENCODER_TARGET_REPLACE_MODULE = ["CLIPAttention", "CLIPMLP"]
     LORA_PREFIX_UNET = "lora_unet"
     LORA_PREFIX_TEXT_ENCODER = "lora_te"
 
@@ -160,235 +161,157 @@ 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
+                temp[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:
+        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":
+                    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":
-                    loras[f"{lora_name}.lora_down.weight"] = (
-                        extract_a.detach().cpu().contiguous().half()
+                    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))
     return all_loras
 
@@ -685,7 +608,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

{lycoris_lora-2.1.0.dev3 → lycoris_lora-2.1.0.dev4}/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.dev4}/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.dev4
 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.dev4}/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.dev4",
     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",