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

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

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

@@ -1,7 +1,7 @@
 import os
 import re
 import logging
-from typing import List
+from typing import Any, List
 
 import torch
 
@@ -404,7 +404,9 @@ class LycorisNetworkKohya(LycorisNetwork):
             next_config = {}
             for name, module in root_module.named_modules():
                 module_name = module.__class__.__name__
-                if module_name in target_replace_modules:
+                if module_name in target_replace_modules and not any(
+                    re.match(t, name) for t in target_replace_names
+                ):
                     if module_name in self.MODULE_ALGO_MAP:
                         next_config = self.MODULE_ALGO_MAP[module_name]
                         algo = next_config.get("algo", network_module)
@@ -419,8 +421,9 @@ class LycorisNetworkKohya(LycorisNetwork):
                 elif name in target_replace_names or any(
                     re.match(t, name) for t in target_replace_names
                 ):
-                    if name in self.NAME_ALGO_MAP:
-                        next_config = self.NAME_ALGO_MAP[name]
+                    conf_from_name = self.find_conf_for_name(name)
+                    if conf_from_name is not None:
+                        next_config = conf_from_name
                         algo = next_config.get("algo", network_module)
                     elif module_name in self.MODULE_ALGO_MAP:
                         next_config = self.MODULE_ALGO_MAP[module_name]
@@ -492,6 +495,19 @@ class LycorisNetworkKohya(LycorisNetwork):
             ), f"duplicated lora name: {lora.lora_name}"
             names.add(lora.lora_name)
 
+    def find_conf_for_name(
+        self,
+        name: str,
+    ) -> dict[str, Any]:
+        if name in self.NAME_ALGO_MAP.keys():
+            return self.NAME_ALGO_MAP[name]
+
+        for key, value in self.NAME_ALGO_MAP.items():
+            if re.match(key, name):
+                return value
+
+        return None
+
     def load_weights(self, file):
         if os.path.splitext(file)[1] == ".safetensors":
             from safetensors.torch import load_file, safe_open

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

@@ -1,9 +1,9 @@
 # General LyCORIS wrapper based on kohya-ss/sd-scripts' style
 import os
-import regex as re
+import re
 import logging
 
-from typing import List
+from typing import Any, List
 
 import torch
 import torch.nn as nn
@@ -334,7 +334,7 @@ class LycorisNetwork(torch.nn.Module):
                     next_config = self.MODULE_ALGO_MAP[module_name]
                     next_algo = next_config.get("algo", algo)
                     new_loras, new_lora_names = create_modules_(
-                        f"{prefix}_{name}", module, next_algo, next_config
+                        f"{prefix}_{name}" if name else prefix, module, next_algo, next_config
                     )
                     for lora_name, lora in zip(new_lora_names, new_loras):
                         if lora_name not in loras:
@@ -367,7 +367,9 @@ class LycorisNetwork(torch.nn.Module):
             next_config = {}
             for name, module in root_module.named_modules():
                 module_name = module.__class__.__name__
-                if module_name in target_replace_modules:
+                if module_name in target_replace_modules and not any(
+                    re.match(t, name) for t in target_replace_names
+                ):
                     if module_name in self.MODULE_ALGO_MAP:
                         next_config = self.MODULE_ALGO_MAP[module_name]
                         algo = next_config.get("algo", network_module)
@@ -382,8 +384,9 @@ class LycorisNetwork(torch.nn.Module):
                 elif name in target_replace_names or any(
                     re.match(t, name) for t in target_replace_names
                 ):
-                    if name in self.NAME_ALGO_MAP:
-                        next_config = self.NAME_ALGO_MAP[name]
+                    conf_from_name = self.find_conf_for_name(name)
+                    if conf_from_name is not None:
+                        next_config = conf_from_name
                         algo = next_config.get("algo", network_module)
                     elif module_name in self.MODULE_ALGO_MAP:
                         next_config = self.MODULE_ALGO_MAP[module_name]
@@ -401,8 +404,14 @@ class LycorisNetwork(torch.nn.Module):
         self.loras = create_modules(
             LycorisNetwork.LORA_PREFIX,
             module,
-            LycorisNetwork.TARGET_REPLACE_MODULE,
-            LycorisNetwork.TARGET_REPLACE_NAME,
+            list(set([
+                *LycorisNetwork.TARGET_REPLACE_MODULE,
+                *LycorisNetwork.MODULE_ALGO_MAP.keys(),
+            ])),
+            list(set([
+                *LycorisNetwork.TARGET_REPLACE_NAME,
+                *LycorisNetwork.NAME_ALGO_MAP.keys(),
+            ])),
         )
         logger.info(f"create LyCORIS: {len(self.loras)} modules.")
 
@@ -421,6 +430,19 @@ class LycorisNetwork(torch.nn.Module):
             ), f"duplicated lora name: {lora.lora_name}"
             names.add(lora.lora_name)
 
+    def find_conf_for_name(
+        self,
+        name: str,
+    ) -> dict[str, Any]:
+        if name in self.NAME_ALGO_MAP.keys():
+            return self.NAME_ALGO_MAP[name]
+
+        for key, value in self.NAME_ALGO_MAP.items():
+            if re.match(key, name):
+                return value
+
+        return None
+
     def set_multiplier(self, multiplier):
         self.multiplier = multiplier
         for lora in self.loras:

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: lycoris_lora
-Version: 3.0.1.dev3
+Version: 3.0.1.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-3.0.1.dev3 → lycoris_lora-3.0.1.dev5}/setup.py

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

lycoris_lora-3.0.1.dev5/test/wrapper.py

@@ -0,0 +1,290 @@
+import unittest
+import re
+
+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 FunLinearNetwork(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.layer_norm = nn.LayerNorm(dim)
+        self.linear = nn.Linear(dim, dim)
+        self.act = nn.ReLU()
+
+    def forward(self, x):
+        x = self.layer_norm(x)
+        x = self.linear(x)
+        x = self.act(x)
+        return x
+
+
+class FunConvNetwork(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.layer_norm = nn.LayerNorm(dim)
+        self.conv = nn.Conv1d(dim, dim, 3, 1, 1)
+        self.act = nn.ReLU()
+
+    def forward(self, x):
+        x = self.layer_norm(x)
+        x = self.conv(x)
+        x = self.act(x)
+        return x
+
+
+class NamedSequential(nn.Module):
+    def __init__(self, name, *layers):
+        super().__init__()
+        self.name = name
+        self.layers = nn.ModuleList(layers)
+    
+    def forward(self, x):
+        for layer in self.layers:
+            x = layer(x)
+        return x
+
+
+class TestNetwork2(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.linear_layers = NamedSequential(
+            "linear_layers",
+            FunLinearNetwork(dim),
+            FunLinearNetwork(dim),
+            FunLinearNetwork(dim),
+        )
+        self.conv_layers = NamedSequential(
+            "conv_layers",
+            FunConvNetwork(dim),
+            FunConvNetwork(dim),
+            FunConvNetwork(dim),
+        )
+
+    def forward(self, x):
+        x = self.linear_layers(x)
+        x = self.conv_layers(x)
+        return x
+
+    def named_modules(self, memo=None, prefix=''):
+        # Call the parent class named_modules method
+        for layer_name, layer in super().named_modules(memo, prefix):
+            yield layer_name, layer
+
+
+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))
+
+    def test_lycoris_wrapper_regex_named_modules(
+        self,
+        device_dtype=(torch.device("cpu"), torch.float32),
+    ):
+        device, dtype = device_dtype
+
+        # Define the preset configuration with regex
+        preset = {
+            "name_algo_map": {
+                "linear_layers.layers.[0-1]..*": {
+                    "algo": "lokr",
+                    "factor": 4,
+                    "linear_dim":1000000,
+                    "linear_alpha":1,
+                    "full_matrix": True,
+                },
+                "linear_layers.layers.2..*": {
+                    "algo": "lora", "dim": 8,
+                },
+                "conv_layers.layers.0..*": {
+                   "algo": "locon", "dim": 8,
+                },
+                "conv_layers.layers.[1-2]..*": {
+                    "algo": "glora", "dim": 128,
+                }
+            }
+        }
+
+        # Apply the preset to the LycorisNetwork class
+        LycorisNetwork.apply_preset(preset)
+
+        # Create the test network and the Lycoris network wrapper
+        test_net = TestNetwork2(16).to(device, dtype)
+        test_lycoris: LycorisNetwork = create_lycoris(
+            test_net,
+            multiplier=1,
+            linear_dim=16,
+            linear_alpha=1,
+            train_norm=True,
+        )
+        test_lycoris.apply_to()
+        test_lycoris.to(device, dtype)
+
+        # Assertions to verify the correct modules were created and configured
+        assert len(test_lycoris.loras) == 12
+
+        lora_names = sorted([ lora.lora_name for lora in test_lycoris.loras ])
+        expected = sorted([
+            "lycoris_linear_layers_layers_0_layer_norm",
+            "lycoris_linear_layers_layers_0_linear",
+            "lycoris_linear_layers_layers_1_layer_norm",
+            "lycoris_linear_layers_layers_1_linear",
+            "lycoris_linear_layers_layers_2_layer_norm",
+            "lycoris_linear_layers_layers_2_linear",
+            "lycoris_conv_layers_layers_0_layer_norm",
+            "lycoris_conv_layers_layers_0_conv",
+            "lycoris_conv_layers_layers_1_layer_norm",
+            "lycoris_conv_layers_layers_1_conv",
+            "lycoris_conv_layers_layers_2_layer_norm",
+            "lycoris_conv_layers_layers_2_conv",
+        ])
+        self.assertEqual(lora_names, expected)
+
+        for lora in test_lycoris.loras:
+            if "lycoris_linear_layers_layers_0" in lora.lora_name or \
+                "lycoris_linear_layers_layers_1" in lora.lora_name:
+                self.assertEqual(lora.dim, 16)
+                if 'norm' in lora.lora_name:
+                    self.assertEqual(lora.__class__.__name__, "NormModule")
+                else:
+                    self.assertEqual(lora.__class__.__name__, "LokrModule")
+                    self.assertEqual(lora.multiplier, 1)
+                    self.assertEqual(lora.full_matrix, True)
+            elif "lycoris_linear_layers_layers_2" in lora.lora_name:
+                self.assertEqual(lora.dim, 16)
+                if 'norm' in lora.lora_name:
+                    self.assertEqual(lora.__class__.__name__, "NormModule")
+                else:
+                    self.assertEqual(lora.__class__.__name__, "LoConModule")
+            elif "lycoris_conv_layers_layers_0" in lora.lora_name:
+                self.assertEqual(lora.dim, 16)
+                if 'norm' in lora.lora_name:
+                    self.assertEqual(lora.__class__.__name__, "NormModule")
+                else:
+                    self.assertEqual(lora.__class__.__name__, "LoConModule")
+            elif "lycoris_linear_layers_layers_1" in lora.lora_name or \
+                "lycoris_linear_layers_layers_2" in lora.lora_name:
+                self.assertEqual(lora.dim, 128)
+                if 'norm' in lora.lora_name:
+                    self.assertEqual(lora.__class__.__name__, "NormModule")
+                else:
+                    self.assertEqual(lora.__class__.__name__, "GLoRAModule")

lycoris_lora-3.0.1.dev3/test/wrapper.py

@@ -1,128 +0,0 @@
-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))