lycoris-lora 3.0.0.dev2__tar.gz → 3.0.0.dev4__tar.gz

{lycoris_lora-3.0.0.dev2/lycoris_lora.egg-info → lycoris_lora-3.0.0.dev4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: lycoris_lora
-Version: 3.0.0.dev2
+Version: 3.0.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
@@ -8,8 +8,5 @@ Author-email: apolloyeh0123@gmail.com
 Requires-Python: >=3.10
 License-File: LICENSE.md
 Requires-Dist: torch
-Requires-Dist: safetensors
-Requires-Dist: diffusers
-Requires-Dist: transformers
 Requires-Dist: einops
 Requires-Dist: toml

{lycoris_lora-3.0.0.dev2 → lycoris_lora-3.0.0.dev4}/lycoris/functional/general.py

@@ -59,6 +59,8 @@ def power2factorization(dimension: int, factor: int = -1) -> tuple[int, int]:
     n = 2**p
     m*n = dim
     """
+    if factor == -1:
+        factor = dimension
 
     # Find the first solution and check if it is even doable
     m = n = 0

{lycoris_lora-3.0.0.dev2 → lycoris_lora-3.0.0.dev4}/lycoris/modules/base.py

@@ -3,6 +3,7 @@ from collections import OrderedDict
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+import torch.nn.utils.parametrize as parametrize
 
 from ..utils.bnb import QuantLinears, log_bypass
 
@@ -181,6 +182,38 @@ class LycorisBaseModule(ModuleCustomSD):
         self.org_forward = org_module.forward
         self.org_module = [org_module]
 
+    @classmethod
+    def parametrize(cls, org_module, attr, *args, **kwargs):
+        target_param = getattr(org_module, attr)
+        kwargs["bypass_mode"] = False
+        if target_param.dim() == 2:
+            proxy_module = nn.Linear(
+                target_param.shape[0], target_param.shape[1], bias=False
+            )
+            proxy_module.weight = target_param
+        elif target_param.dim() > 2:
+            module_type = [
+                None,
+                None,
+                None,
+                nn.Conv1d,
+                nn.Conv2d,
+                nn.Conv3d,
+                None,
+                None,
+            ][target_param.dim()]
+            proxy_module = module_type(
+                target_param.shape[0],
+                target_param.shape[1],
+                *target_param.shape[2:],
+            )
+            proxy_module.weight = target_param
+        module_obj = cls("", proxy_module, *args, **kwargs)
+        module_obj.forward = module_obj.parametrize_forward
+        module_obj.to(target_param.device)
+        parametrize.register_parametrization(org_module, attr, module_obj)
+        return module_obj
+
     @property
     def dtype(self):
         return self.dtype_tensor.dtype
@@ -231,5 +264,10 @@ class LycorisBaseModule(ModuleCustomSD):
     def bypass_forward(self, x, scale=1):
         raise NotImplementedError
 
+    def parametrize_forward(self, x: torch.Tensor, *args, **kwargs):
+        return self.get_merged_weight(
+            multiplier=self.multiplier, shape=x.shape, device=x.device
+        )[0].to(x.dtype)
+
     def forward(self, *args, **kwargs):
         raise NotImplementedError

{lycoris_lora-3.0.0.dev2 → lycoris_lora-3.0.0.dev4}/lycoris/modules/boft.py

@@ -136,7 +136,7 @@ class ButterflyOFTModule(LycorisBaseModule):
         diff = self.make_weight(scale=multiplier, device=device, diff=True)
         if shape is not None:
             diff = diff.view(shape)
-        return diff
+        return diff, None
 
     def get_merged_weight(self, multiplier=1, shape=None, device=None):
         diff = self.make_weight(scale=multiplier, device=device)
@@ -200,33 +200,43 @@ class ButterflyOFTModule(LycorisBaseModule):
 
 
 if __name__ == "__main__":
-    base = nn.Linear(128, 128).cuda()
-    boft = ButterflyOFTModule(
-        "test", base, 1, 4, 1, weight_decompose=True, factor=8
-    ).cuda()
-    print(boft)
-    test_input = torch.randn(1, 77, 128).cuda()
-    test_output = boft(test_input)
-    torch.sum(test_output).backward()
-    print(test_output.shape)
-
-    base_4bit = LinearNF4(128, 128)
-    base_4bit.load_state_dict(base.state_dict())
-    base_4bit.cuda()
-    qboft = ButterflyOFTModule(
-        "test", base_4bit, 1, 4, 1, weight_decompose=False
-    ).cuda()
-    print(qboft)
-    test_output = qboft(test_input)
-    torch.sum(test_output).backward()
-    print(test_output.shape)
-
-    base = nn.Conv2d(128, 128, 3, 1, 1)
-    boft = ButterflyOFTModule(
-        "test", base, 1, 4, 1, weight_decompose=True, use_tucker=True
-    )
-    print(boft)
-    test_input = torch.randn(1, 128, 16, 16)
-    test_output = boft(test_input)
-    torch.sum(test_output).backward()
-    print(test_output.shape)
+    device = torch.device("cuda")
+    module = ButterflyOFTModule
+    with torch.autocast("cuda" if torch.cuda.is_available() else "cpu"):
+        base = nn.Linear(128, 128).to(device).half()
+        net = module("test", base, 1, 4, 1, weight_decompose=True).to(device)
+        print(net)
+        test_input = torch.randn(1, 128).to(device).half()
+        test_output = net(test_input)
+        torch.sum(test_output).backward()
+        print(test_output.shape)
+
+        base_4bit = LinearNF4(128, 128, device="cuda")
+        base_4bit.load_state_dict(base.state_dict())
+        base_4bit.to(device)
+        qnet = module("test", base_4bit, 1, 4, 1, weight_decompose=False).to(device)
+        print(qnet)
+        test_input = torch.randn(1, 128).to(device).half()
+        test_output = qnet(test_input)
+        torch.sum(test_output).backward()
+        print(test_output.shape)
+
+        base = nn.Conv2d(128, 128, 3, 1, 1).to(device).half()
+        net = module("test", base, 1, 4, 1, weight_decompose=True, use_tucker=True).to(
+            device
+        )
+        print(net)
+        test_input = torch.randn(1, 128, 16, 16).to(device).half()
+        test_output = net(test_input)
+        torch.sum(test_output).backward()
+        print(test_output.shape)
+
+        base = nn.Conv2d(128, 128, 3, 1, 1).to(device).half()
+        net = module.parametrize(
+            base, "weight", 1, 4, 1, weight_decompose=True, use_tucker=True
+        ).to(device)
+        print(base)
+        test_input = torch.randn(1, 128, 16, 16).to(device).half()
+        test_output = base(test_input)
+        torch.sum(test_output).backward()
+        print(test_output.shape)

{lycoris_lora-3.0.0.dev2 → lycoris_lora-3.0.0.dev4}/lycoris/modules/diag_oft.py

@@ -119,7 +119,7 @@ class DiagOFTModule(LycorisBaseModule):
         diff = self.make_weight(scale=multiplier, device=device, diff=True)
         if shape is not None:
             diff = diff.view(shape)
-        return diff
+        return diff, None
 
     def get_merged_weight(self, multiplier=1, shape=None, device=None):
         diff = self.make_weight(scale=multiplier, device=device)
@@ -185,27 +185,43 @@ class DiagOFTModule(LycorisBaseModule):
 
 
 if __name__ == "__main__":
-    base = nn.Linear(128, 128).cuda()
-    doft = DiagOFTModule("test", base, 1, 4, 1, weight_decompose=True, factor=8).cuda()
-    print(doft)
-    test_input = torch.randn(1, 77, 128).cuda()
-    test_output = doft(test_input)
-    torch.sum(test_output).backward()
-    print(test_output.shape)
-
-    base_4bit = LinearNF4(128, 128)
-    base_4bit.load_state_dict(base.state_dict())
-    base_4bit.cuda()
-    qdoft = DiagOFTModule("test", base_4bit, 1, 4, 1, weight_decompose=False).cuda()
-    print(qdoft)
-    test_output = qdoft(test_input)
-    torch.sum(test_output).backward()
-    print(test_output.shape)
-
-    base = nn.Conv2d(128, 128, 3, 1, 1)
-    doft = DiagOFTModule("test", base, 1, 4, 1, weight_decompose=True, use_tucker=True)
-    print(doft)
-    test_input = torch.randn(1, 128, 16, 16)
-    test_output = doft(test_input)
-    torch.sum(test_output).backward()
-    print(test_output.shape)
+    device = torch.device("cuda")
+    module = DiagOFTModule
+    with torch.autocast("cuda" if torch.cuda.is_available() else "cpu"):
+        base = nn.Linear(128, 128).to(device).half()
+        net = module("test", base, 1, 4, 1, weight_decompose=True).to(device)
+        print(net)
+        test_input = torch.randn(1, 128).to(device).half()
+        test_output = net(test_input)
+        torch.sum(test_output).backward()
+        print(test_output.shape)
+
+        base_4bit = LinearNF4(128, 128, device="cuda")
+        base_4bit.load_state_dict(base.state_dict())
+        base_4bit.to(device)
+        qnet = module("test", base_4bit, 1, 4, 1, weight_decompose=False).to(device)
+        print(qnet)
+        test_input = torch.randn(1, 128).to(device).half()
+        test_output = qnet(test_input)
+        torch.sum(test_output).backward()
+        print(test_output.shape)
+
+        base = nn.Conv2d(128, 128, 3, 1, 1).to(device).half()
+        net = module("test", base, 1, 4, 1, weight_decompose=True, use_tucker=True).to(
+            device
+        )
+        print(net)
+        test_input = torch.randn(1, 128, 16, 16).to(device).half()
+        test_output = net(test_input)
+        torch.sum(test_output).backward()
+        print(test_output.shape)
+
+        base = nn.Conv2d(128, 128, 3, 1, 1).to(device).half()
+        net = module.parametrize(
+            base, "weight", 1, 4, 1, weight_decompose=True, use_tucker=True
+        ).to(device)
+        print(base)
+        test_input = torch.randn(1, 128, 16, 16).to(device).half()
+        test_output = base(test_input)
+        torch.sum(test_output).backward()
+        print(test_output.shape)

{lycoris_lora-3.0.0.dev2 → lycoris_lora-3.0.0.dev4}/lycoris/modules/full.py

@@ -158,18 +158,33 @@ class FullModule(LycorisBaseModule):
 
 
 if __name__ == "__main__":
-    base = nn.Linear(128, 128).cuda()
-    full = FullModule("test", base, 1, 4, 1, weight_decompose=True, factor=8).cuda()
-    print(full)
-    test_input = torch.randn(1, 77, 128).cuda()
-    test_output = full(test_input)
-    torch.sum(test_output).backward()
-    print(test_output.shape)
-
-    base = nn.Conv2d(128, 128, 3, 1, 1)
-    full = FullModule("test", base, 1, 4, 1, weight_decompose=True, use_tucker=True)
-    print(full)
-    test_input = torch.randn(1, 128, 16, 16)
-    test_output = full(test_input)
-    torch.sum(test_output).backward()
-    print(test_output.shape)
+    device = torch.device("cuda")
+    module = FullModule
+    with torch.autocast("cuda" if torch.cuda.is_available() else "cpu"):
+        base = nn.Linear(128, 128).to(device).half()
+        net = module("test", base, 1, 4, 1, weight_decompose=True).to(device)
+        print(net)
+        test_input = torch.randn(1, 128).to(device).half()
+        test_output = net(test_input)
+        torch.sum(test_output).backward()
+        print(test_output.shape)
+
+        base = nn.Conv2d(128, 128, 3, 1, 1).to(device).half()
+        net = module("test", base, 1, 4, 1, weight_decompose=True, use_tucker=True).to(
+            device
+        )
+        print(net)
+        test_input = torch.randn(1, 128, 16, 16).to(device).half()
+        test_output = net(test_input)
+        torch.sum(test_output).backward()
+        print(test_output.shape)
+
+        base = nn.Conv2d(128, 128, 3, 1, 1).to(device).half()
+        net = module.parametrize(
+            base, "weight", 1, 4, 1, weight_decompose=True, use_tucker=True
+        ).to(device)
+        print(base)
+        test_input = torch.randn(1, 128, 16, 16).to(device).half()
+        test_output = base(test_input)
+        torch.sum(test_output).backward()
+        print(test_output.shape)

{lycoris_lora-3.0.0.dev2 → lycoris_lora-3.0.0.dev4}/lycoris/modules/glora.py

@@ -5,6 +5,7 @@ import torch.nn as nn
 import torch.nn.functional as F
 
 from .base import LycorisBaseModule
+from ..utils import tucker_weight_from_conv
 from ..utils.bnb import LinearNF4
 
 
@@ -52,7 +53,7 @@ class GLoRAModule(LycorisBaseModule):
             bypass_mode,
         )
         if self.module_type not in self.support_module:
-            raise ValueError(f"{self.module_type} is not supported in LoRA/LoCon algo.")
+            raise ValueError(f"{self.module_type} is not supported in GLoRA algo.")
         self.lora_dim = lora_dim
         self.tucker = False
         self.rs_lora = rs_lora
@@ -68,17 +69,12 @@ class GLoRAModule(LycorisBaseModule):
             self.tucker = use_tucker and k_size != (1, 1)
             self.down_op = self.op
             self.up_op = self.op
-            if use_tucker and k_size != (1, 1):
-                self.a2 = self.module(in_dim, lora_dim, (1, 1), bias=False)
-                self.am = self.module(
-                    lora_dim, lora_dim, k_size, stride, padding, bias=False
-                )
-                self.tucker = True
-            else:
-                self.a2 = self.module(
-                    in_dim, lora_dim, k_size, stride, padding, bias=False
-                )
+
+            # A
+            self.a2 = self.module(in_dim, lora_dim, (1, 1), bias=False)
             self.a1 = self.module(lora_dim, in_dim, (1, 1), bias=False)
+
+            # B
             if use_tucker and k_size != (1, 1):
                 self.b2 = self.module(in_dim, lora_dim, (1, 1), bias=False)
                 self.bm = self.module(
@@ -90,7 +86,7 @@ class GLoRAModule(LycorisBaseModule):
                     in_dim, lora_dim, k_size, stride, padding, bias=False
                 )
             self.b1 = self.module(lora_dim, out_dim, (1, 1), bias=False)
-        elif isinstance(org_module, nn.Linear):
+        else:
             self.isconv = False
             self.down_op = F.linear
             self.up_op = F.linear
@@ -100,8 +96,6 @@ class GLoRAModule(LycorisBaseModule):
             self.a1 = nn.Linear(lora_dim, in_dim, bias=False)
             self.b2 = nn.Linear(in_dim, lora_dim, bias=False)
             self.b1 = nn.Linear(lora_dim, out_dim, bias=False)
-        else:
-            raise NotImplementedError
 
         if type(alpha) == torch.Tensor:
             alpha = alpha.detach().float().numpy()  # without casting, bf16 causes error
@@ -129,19 +123,35 @@ class GLoRAModule(LycorisBaseModule):
     def make_weight(self, device=None):
         wa1 = self.a1.weight.view(self.a1.weight.size(0), -1)
         wa2 = self.a2.weight.view(self.a2.weight.size(0), -1)
-        wb1 = self.b1.weight.view(self.b1.weight.size(0), -1)
-        wb2 = self.b2.weight.view(self.b2.weight.size(0), -1)
-        orig = self.org_weight.view(self.org_module[0].weight.size(0), -1)
-        return (wb2 @ wb1) + ((orig @ wa2) @ wa1)
 
-    def forward(self, x):
-        if self.module_dropout and self.training:
-            if torch.rand(1) < self.module_dropout:
-                return self.org_forward(x)
-        scale = self.scale * self.multiplier
+        if self.tucker:
+            wb = tucker_weight_from_conv(self.b1.weight, self.b2.weight, self.bm.weight)
+        else:
+            wb1 = self.b1.weight.view(self.b1.weight.size(0), -1)
+            wb2 = self.b2.weight.view(self.b2.weight.size(0), -1)
+            wb = wb1 @ wb2
+        orig = self.org_weight
+        if orig.dim() > 2:
+            w_wa1 = torch.einsum("o i ..., i j -> o j ...", orig, wa1)
+            w_wa2 = torch.einsum("o i ..., i j -> o j ...", w_wa1, wa2)
+        else:
+            w_wa2 = (orig @ wa1) @ wa2
+        return (wb + w_wa2) * self.scale
+
+    def get_diff_weight(self, multiplier=1.0, shape=None, device=None):
+        weight = self.make_weight(device) * multiplier
+        if shape is not None:
+            weight = weight.view(shape)
+        return weight, None
+
+    def get_merged_weight(self, multiplier=1, shape=None, device=None):
+        diff_w, _ = self.get_diff_weight(multiplier, shape, device)
+        return self.org_weight + diff_w, None
 
-        ax_mid = self.a1(x) * scale
-        bx_mid = self.b1(x) * scale
+    def _bypass_forward(self, x, scale=1, diff=False):
+        scale = self.scale * scale
+        ax_mid = self.a2(x) * scale
+        bx_mid = self.b2(x) * scale
 
         if self.rank_dropout and self.training:
             drop_a = (
@@ -161,36 +171,76 @@ class GLoRAModule(LycorisBaseModule):
                 drop_b = drop_b.view(*[1] * (dims - 1), -1)
             ax_mid = ax_mid * drop_a
             bx_mid = bx_mid * drop_b
-
         return (
-            self.org_forward(x + self.dropout(self.a2(ax_mid)) * self.scale)
-            + self.dropout(self.b2(bx_mid)) * self.scale
+            self.org_forward(
+                (0 if diff else x) + self.drop(self.a1(ax_mid)) * self.scale
+            )
+            + self.drop(self.b1(bx_mid)) * self.scale
         )
 
+    def bypass_forward_diff(self, x, scale=1):
+        return self._bypass_forward(x, scale=scale, diff=True)
+
+    def bypass_forward(self, x, scale=1):
+        return self._bypass_forward(x, scale=scale, diff=False)
+
+    def forward(self, x, *args, **kwargs):
+        if self.module_dropout and self.training:
+            if torch.rand(1) < self.module_dropout:
+                return self.org_forward(x)
+        if self.bypass_mode:
+            return self.bypass_forward(x, self.multiplier)
+        else:
+            weight = (
+                self.org_module[0].weight.data.to(self.dtype)
+                + self.get_diff_weight(multiplier=self.multiplier)[0]
+            )
+            bias = (
+                None
+                if self.org_module[0].bias is None
+                else self.org_module[0].bias.data
+            )
+            return self.op(x, weight, bias, **self.kw_dict)
+
 
 if __name__ == "__main__":
-    base = nn.Linear(128, 128).cuda()
-    glora = GLoRAModule("test", base, 1, 4, 1, weight_decompose=True).cuda()
-    print(glora)
-    test_input = torch.randn(1, 128).cuda()
-    test_output = glora(test_input)
-    torch.sum(test_output).backward()
-    print(test_output.shape)
-
-    base_4bit = LinearNF4(128, 128)
-    base_4bit.load_state_dict(base.state_dict())
-    base_4bit.cuda()
-    qglora = GLoRAModule("test", base_4bit, 1, 4, 1, weight_decompose=False).cuda()
-    print(qglora)
-    test_input = torch.randn(1, 128).cuda()
-    test_output = qglora(test_input)
-    torch.sum(test_output).backward()
-    print(test_output.shape)
-
-    base = nn.Conv2d(128, 128, 3, 1, 1)
-    glora = GLoRAModule("test", base, 1, 4, 1, weight_decompose=True, use_tucker=True)
-    print(glora)
-    test_input = torch.randn(1, 128, 16, 16)
-    test_output = glora(test_input)
-    torch.sum(test_output).backward()
-    print(test_output.shape)
+    device = torch.device("cuda")
+    module = GLoRAModule
+    with torch.autocast("cuda" if torch.cuda.is_available() else "cpu"):
+        base = nn.Linear(128, 128).to(device).half()
+        net = module("test", base, 1, 4, 1, weight_decompose=True).to(device)
+        print(net)
+        test_input = torch.randn(1, 128).to(device).half()
+        test_output = net(test_input)
+        torch.sum(test_output).backward()
+        print(test_output.shape)
+
+        base_4bit = LinearNF4(128, 128, device="cuda")
+        base_4bit.load_state_dict(base.state_dict())
+        base_4bit.to(device)
+        qnet = module("test", base_4bit, 1, 4, 1, weight_decompose=False).to(device)
+        print(qnet)
+        test_input = torch.randn(1, 128).to(device).half()
+        test_output = qnet(test_input)
+        torch.sum(test_output).backward()
+        print(test_output.shape)
+
+        base = nn.Conv2d(128, 128, 3, 1, 1).to(device).half()
+        net = module("test", base, 1, 4, 1, weight_decompose=True, use_tucker=True).to(
+            device
+        )
+        print(net)
+        test_input = torch.randn(1, 128, 16, 16).to(device).half()
+        test_output = net(test_input)
+        torch.sum(test_output).backward()
+        print(test_output.shape)
+
+        base = nn.Conv2d(128, 128, 3, 1, 1).to(device).half()
+        net = module.parametrize(
+            base, "weight", 1, 4, 1, weight_decompose=True, use_tucker=True
+        ).to(device)
+        print(base)
+        test_input = torch.randn(1, 128, 16, 16).to(device).half()
+        test_output = base(test_input)
+        torch.sum(test_output).backward()
+        print(test_output.shape)

lycoris_lora-3.0.0.dev4/lycoris/modules/ia3.py

@@ -0,0 +1,179 @@
+import torch
+import torch.nn as nn
+
+from .base import LycorisBaseModule
+from ..utils.bnb import LinearNF4
+
+
+class IA3Module(LycorisBaseModule):
+    support_module = {
+        "linear",
+        "conv1d",
+        "conv2d",
+        "conv3d",
+    }
+
+    def __init__(
+        self,
+        lora_name,
+        org_module: nn.Module,
+        multiplier=1.0,
+        lora_dim=4,
+        alpha=1,
+        dropout=0.0,
+        rank_dropout=0.0,
+        module_dropout=0.0,
+        use_tucker=False,
+        use_scalar=False,
+        rank_dropout_scale=False,
+        weight_decompose=False,
+        bypass_mode=False,
+        rs_lora=False,
+        train_on_input=False,
+        **kwargs,
+    ):
+        """if alpha == 0 or None, alpha is rank (no scaling)."""
+        super().__init__(
+            lora_name,
+            org_module,
+            multiplier,
+            dropout,
+            rank_dropout,
+            module_dropout,
+            rank_dropout_scale,
+            bypass_mode,
+        )
+        if self.module_type not in self.support_module:
+            raise ValueError(f"{self.module_type} is not supported in IA^3 algo.")
+        self.lora_dim = lora_dim
+        self.tucker = False
+        self.rs_lora = rs_lora
+
+        self.shape = org_module.weight.shape
+        if self.module_type.startswith("conv"):
+            self.isconv = True
+            in_dim = org_module.in_channels
+            out_dim = org_module.out_channels
+            if train_on_input:
+                train_dim = in_dim
+            else:
+                train_dim = out_dim
+            self.weight = nn.Parameter(
+                torch.empty(1, train_dim, *(1 for _ in self.shape[2:]))
+            )
+        else:
+            in_dim = org_module.in_features
+            out_dim = org_module.out_features
+            if train_on_input:
+                train_dim = in_dim
+            else:
+                train_dim = out_dim
+
+            self.weight = nn.Parameter(torch.empty(train_dim))
+
+        # Need more experiences on init method
+        torch.nn.init.constant_(self.weight, 0)
+
+        self.multiplier = multiplier
+        self.train_input = train_on_input
+        self.register_buffer("on_input", torch.tensor(int(train_on_input)))
+
+    def apply_to(self):
+        self.org_forward = self.org_module[0].forward
+        self.org_module[0].forward = self.forward
+
+    def make_weight(self, multiplier=1, shape=None, device=None, diff=False):
+        weight = self.weight * multiplier + int(not diff)
+        if self.train_input:
+            diff = self.org_weight * weight
+        else:
+            diff = self.org_weight.transpose(0, 1) * weight
+            diff = diff.transpose(0, 1)
+        if shape is not None:
+            diff = diff.view(shape)
+        if device is not None:
+            diff = diff.to(device)
+        return diff
+
+    def get_diff_weight(self, multiplier=1, shape=None, device=None):
+        diff = self.make_weight(
+            multiplier=multiplier, shape=shape, device=device, diff=True
+        )
+        return diff, None
+
+    def get_merged_weight(self, multiplier=1, shape=None, device=None):
+        diff = self.make_weight(multiplier=multiplier, shape=shape, device=device)
+        return diff, None
+
+    def _bypass_forward(self, x, scale=1, diff=False):
+        weight = self.weight * scale + int(not diff)
+        if self.train_input:
+            x = x * weight
+        out = self.org_forward(x)
+        if not self.train_input:
+            out = out * weight
+        return out
+
+    def bypass_forward_diff(self, x, scale=1):
+        return self._bypass_forward(x, scale, diff=True)
+
+    def bypass_forward(self, x, scale=1):
+        return self._bypass_forward(x, scale, diff=False)
+
+    def forward(self, x, *args, **kwargs):
+        if self.module_dropout and self.training:
+            if torch.rand(1) < self.module_dropout:
+                return self.org_forward(x)
+        if self.bypass_mode:
+            return self.bypass_forward(x, self.multiplier)
+        else:
+            weight = self.get_merged_weight(multiplier=self.multiplier)[0]
+            bias = (
+                None
+                if self.org_module[0].bias is None
+                else self.org_module[0].bias.data
+            )
+            return self.op(x, weight, bias, **self.kw_dict)
+
+
+if __name__ == "__main__":
+    device = torch.device("cuda")
+    module = IA3Module
+    with torch.autocast("cuda" if torch.cuda.is_available() else "cpu"):
+        base = nn.Linear(128, 128).to(device).half()
+        net = module("test", base, 1, 4, 1, weight_decompose=True).to(device)
+        print(net)
+        test_input = torch.randn(1, 128).to(device).half()
+        test_output = net(test_input)
+        torch.sum(test_output).backward()
+        print(test_output.shape)
+
+        base_4bit = LinearNF4(128, 128, device="cuda")
+        base_4bit.load_state_dict(base.state_dict())
+        base_4bit.to(device)
+        qnet = module("test", base_4bit, 1, 4, 1, weight_decompose=False).to(device)
+        print(qnet)
+        test_input = torch.randn(1, 128).to(device).half()
+        test_output = qnet(test_input)
+        torch.sum(test_output).backward()
+        print(test_output.shape)
+
+        base = nn.Conv2d(128, 128, 3, 1, 1).to(device).half()
+        net = module("test", base, 1, 4, 1, weight_decompose=True, use_tucker=True).to(
+            device
+        )
+        print(net)
+        test_input = torch.randn(1, 128, 16, 16).to(device).half()
+        test_output = net(test_input)
+        torch.sum(test_output).backward()
+        print(test_output.shape)
+
+        base = nn.Conv2d(128, 128, 3, 1, 1).to(device).half()
+        net = module.parametrize(
+            base, "weight", 1, 4, 1, weight_decompose=True, use_tucker=True
+        ).to(device)
+        print(base)
+        test_input = torch.randn(1, 128, 16, 16).to(device).half()
+        test_output = base(test_input)
+        torch.sum(test_output).backward()
+        print(test_output.shape)

{lycoris_lora-3.0.0.dev2 → lycoris_lora-3.0.0.dev4}/lycoris/modules/locon.py

@@ -183,11 +183,12 @@ class LoConModule(LycorisBaseModule):
             diff = diff.view(shape)
         if device is not None:
             diff = diff.to(device)
-        return diff
+        return diff, None
 
     def get_merged_weight(self, multiplier=1, shape=None, device=None):
-        merged = self.org_module[0].weight.data + self.get_diff_weight(
-            multiplier=multiplier, shape=shape, device=device
+        merged = (
+            self.org_module[0].weight.data
+            + self.get_diff_weight(multiplier=multiplier, shape=shape, device=device)[0]
         )
         if self.wd:
             merged = self.apply_weight_decompose(merged)
@@ -278,9 +279,10 @@ class LoConModule(LycorisBaseModule):
 
 if __name__ == "__main__":
     device = torch.device("cuda")
+    module = LoConModule
     with torch.autocast("cuda" if torch.cuda.is_available() else "cpu"):
         base = nn.Linear(128, 128).to(device).half()
-        net = LoConModule("test", base, 1, 4, 1, weight_decompose=True).to(device)
+        net = module("test", base, 1, 4, 1, weight_decompose=True).to(device)
         print(net)
         test_input = torch.randn(1, 128).to(device).half()
         test_output = net(test_input)
@@ -290,9 +292,7 @@ if __name__ == "__main__":
         base_4bit = LinearNF4(128, 128, device="cuda")
         base_4bit.load_state_dict(base.state_dict())
         base_4bit.to(device)
-        qnet = LoConModule("test", base_4bit, 1, 4, 1, weight_decompose=False).to(
-            device
-        )
+        qnet = module("test", base_4bit, 1, 4, 1, weight_decompose=False).to(device)
         print(qnet)
         test_input = torch.randn(1, 128).to(device).half()
         test_output = qnet(test_input)
@@ -300,11 +300,21 @@ if __name__ == "__main__":
         print(test_output.shape)
 
         base = nn.Conv2d(128, 128, 3, 1, 1).to(device).half()
-        net = LoConModule(
-            "test", base, 1, 4, 1, weight_decompose=True, use_tucker=True
+        net = module("test", base, 1, 4, 1, weight_decompose=True, use_tucker=True).to(
+            device
         )
         print(net)
         test_input = torch.randn(1, 128, 16, 16).to(device).half()
         test_output = net(test_input)
         torch.sum(test_output).backward()
         print(test_output.shape)
+
+        base = nn.Conv2d(128, 128, 3, 1, 1).to(device).half()
+        net = module.parametrize(
+            base, "weight", 1, 4, 1, weight_decompose=True, use_tucker=True
+        )
+        print(base)
+        test_input = torch.randn(1, 128, 16, 16).to(device).half()
+        test_output = base(test_input)
+        torch.sum(test_output).backward()
+        print(test_output.shape)

{lycoris_lora-3.0.0.dev2 → lycoris_lora-3.0.0.dev4}/lycoris/modules/loha.py

@@ -181,15 +181,16 @@ class LohaModule(LycorisBaseModule):
         diff = self.get_weight(shape) * scale
         if device is not None:
             diff = diff.to(device)
-        return diff
+        return diff, None
 
     def get_merged_weight(self, multiplier=1, shape=None, device=None):
-        merged = self.org_module[0].weight.data + self.get_diff_weight(
-            multiplier=multiplier, shape=shape, device=device
+        merged = (
+            self.org_module[0].weight.data
+            + self.get_diff_weight(multiplier=multiplier, shape=shape, device=device)[0]
         )
         if self.wd:
             merged = self.apply_weight_decompose(merged)
-        return merged
+        return merged, None
 
     def apply_weight_decompose(self, weight):
         weight = weight.to(self.dora_scale.dtype)
@@ -264,33 +265,47 @@ class LohaModule(LycorisBaseModule):
 
             if self.wd:
                 weight = self.apply_weight_decompose(weight)
-            return self.op(x, weight.view(self.shape), bias, **self.kw_dict)
+            return self.op(x, weight, bias, **self.kw_dict)
 
 
 if __name__ == "__main__":
     device = torch.device("cuda")
-    base = nn.Linear(128, 128).to(device)
-    loha = LohaModule("test", base, 1, 4, 1, weight_decompose=True).to(device)
-    print(loha)
-    test_input = torch.randn(1, 128).to(device)
-    test_output = loha(test_input)
-    torch.sum(test_output).backward()
-    print(test_output.shape)
-
-    base_4bit = LinearNF4(128, 128)
-    base_4bit.load_state_dict(base.state_dict())
-    base_4bit.to(device)
-    qloha = LohaModule("test", base_4bit, 1, 4, 1, weight_decompose=False).to(device)
-    print(qloha)
-    test_input = torch.randn(1, 128).to(device)
-    test_output = qloha(test_input)
-    torch.sum(test_output).backward()
-    print(test_output.shape)
-
-    base = nn.Conv2d(128, 128, 3, 1, 1)
-    loha = LohaModule("test", base, 1, 4, 1, weight_decompose=True, use_tucker=True)
-    print(loha)
-    test_input = torch.randn(1, 128, 16, 16)
-    test_output = loha(test_input)
-    torch.sum(test_output).backward()
-    print(test_output.shape)
+    module = LohaModule
+    with torch.autocast("cuda" if torch.cuda.is_available() else "cpu"):
+        base = nn.Linear(128, 128).to(device).half()
+        net = module("test", base, 1, 4, 1, weight_decompose=True).to(device)
+        print(net)
+        test_input = torch.randn(1, 128).to(device).half()
+        test_output = net(test_input)
+        torch.sum(test_output).backward()
+        print(test_output.shape)
+
+        base_4bit = LinearNF4(128, 128, device="cuda")
+        base_4bit.load_state_dict(base.state_dict())
+        base_4bit.to(device)
+        qnet = module("test", base_4bit, 1, 4, 1, weight_decompose=False).to(device)
+        print(qnet)
+        test_input = torch.randn(1, 128).to(device).half()
+        test_output = qnet(test_input)
+        torch.sum(test_output).backward()
+        print(test_output.shape)
+
+        base = nn.Conv2d(128, 128, 3, 1, 1).to(device).half()
+        net = module("test", base, 1, 4, 1, weight_decompose=True, use_tucker=True).to(
+            device
+        )
+        print(net)
+        test_input = torch.randn(1, 128, 16, 16).to(device).half()
+        test_output = net(test_input)
+        torch.sum(test_output).backward()
+        print(test_output.shape)
+
+        base = nn.Conv2d(128, 128, 3, 1, 1).to(device).half()
+        net = module.parametrize(
+            base, "weight", 1, 4, 1, weight_decompose=True, use_tucker=True
+        )
+        print(base)
+        test_input = torch.randn(1, 128, 16, 16).to(device).half()
+        test_output = base(test_input)
+        torch.sum(test_output).backward()
+        print(test_output.shape)

{lycoris_lora-3.0.0.dev2 → lycoris_lora-3.0.0.dev4}/lycoris/modules/lokr.py

@@ -72,7 +72,7 @@ class LokrModule(LycorisBaseModule):
             bypass_mode,
         )
         if self.module_type not in self.support_module:
-            raise ValueError(f"{self.module_type} is not supported in LoRA/LoCon algo.")
+            raise ValueError(f"{self.module_type} is not supported in LoKr algo.")
 
         factor = int(factor)
         self.lora_dim = lora_dim
@@ -262,15 +262,21 @@ class LokrModule(LycorisBaseModule):
             weight *= drop
         return weight
 
-    def get_merged_weight(self, multiplier=1, shape=None, device=None):
+    def get_diff_weight(self, multiplier=1, shape=None, device=None):
         scale = self.scale * multiplier
         diff = self.get_weight(shape) * scale
         if device is not None:
             diff = diff.to(device)
-        merged = self.org_module[0].weight.data + diff
+        return diff, None
+
+    def get_merged_weight(self, multiplier=1, shape=None, device=None):
+        merged = (
+            self.org_module[0].weight.data
+            + self.get_diff_weight(multiplier=multiplier, shape=shape, device=device)[0]
+        )
         if self.wd:
             merged = self.apply_weight_decompose(merged)
-        return merged
+        return merged, None
 
     def apply_weight_decompose(self, weight):
         weight = weight.to(self.dora_scale.dtype)
@@ -390,31 +396,47 @@ class LokrModule(LycorisBaseModule):
 
             if self.wd:
                 weight = self.apply_weight_decompose(weight)
-            return self.op(x, weight.view(self.shape), bias, **self.kw_dict)
+            return self.op(x, weight, bias, **self.kw_dict)
 
 
 if __name__ == "__main__":
-    base = nn.Linear(128, 128).cuda()
-    lokr = LokrModule("test", base, 1, 4, 1, weight_decompose=True, factor=8).cuda()
-    print(lokr)
-    test_input = torch.randn(1, 77, 128).cuda()
-    test_output = lokr(test_input)
-    torch.sum(test_output).backward()
-    print(test_output.shape)
-
-    base_4bit = LinearNF4(128, 128)
-    base_4bit.load_state_dict(base.state_dict())
-    base_4bit.cuda()
-    qlokr = LokrModule("test", base_4bit, 1, 4, 1, weight_decompose=False).cuda()
-    print(qlokr)
-    test_output = qlokr(test_input)
-    torch.sum(test_output).backward()
-    print(test_output.shape)
-
-    base = nn.Conv2d(128, 128, 3, 1, 1)
-    lokr = LokrModule("test", base, 1, 4, 1, weight_decompose=True, use_tucker=True)
-    print(lokr)
-    test_input = torch.randn(1, 128, 16, 16)
-    test_output = lokr(test_input)
-    torch.sum(test_output).backward()
-    print(test_output.shape)
+    device = torch.device("cuda")
+    module = LokrModule
+    with torch.autocast("cuda" if torch.cuda.is_available() else "cpu"):
+        base = nn.Linear(128, 128).to(device).half()
+        net = module("test", base, 1, 4, 1, weight_decompose=True).to(device)
+        print(net)
+        test_input = torch.randn(1, 128).to(device).half()
+        test_output = net(test_input)
+        torch.sum(test_output).backward()
+        print(test_output.shape)
+
+        base_4bit = LinearNF4(128, 128, device="cuda")
+        base_4bit.load_state_dict(base.state_dict())
+        base_4bit.to(device)
+        qnet = module("test", base_4bit, 1, 4, 1, weight_decompose=False).to(device)
+        print(qnet)
+        test_input = torch.randn(1, 128).to(device).half()
+        test_output = qnet(test_input)
+        torch.sum(test_output).backward()
+        print(test_output.shape)
+
+        base = nn.Conv2d(128, 128, 3, 1, 1).to(device).half()
+        net = module("test", base, 1, 4, 1, weight_decompose=True, use_tucker=True).to(
+            device
+        )
+        print(net)
+        test_input = torch.randn(1, 128, 16, 16).to(device).half()
+        test_output = net(test_input)
+        torch.sum(test_output).backward()
+        print(test_output.shape)
+
+        base = nn.Conv2d(128, 128, 3, 1, 1).to(device).half()
+        net = module.parametrize(
+            base, "weight", 1, 4, 1, weight_decompose=True, use_tucker=True
+        )
+        print(base)
+        test_input = torch.randn(1, 128, 16, 16).to(device).half()
+        test_output = base(test_input)
+        torch.sum(test_output).backward()
+        print(test_output.shape)

{lycoris_lora-3.0.0.dev2 → lycoris_lora-3.0.0.dev4/lycoris_lora.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: lycoris_lora
-Version: 3.0.0.dev2
+Version: 3.0.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
@@ -8,8 +8,5 @@ Author-email: apolloyeh0123@gmail.com
 Requires-Python: >=3.10
 License-File: LICENSE.md
 Requires-Dist: torch
-Requires-Dist: safetensors
-Requires-Dist: diffusers
-Requires-Dist: transformers
 Requires-Dist: einops
 Requires-Dist: toml

lycoris_lora-3.0.0.dev4/lycoris_lora.egg-info/requires.txt

@@ -0,0 +1,3 @@
+torch
+einops
+toml

{lycoris_lora-3.0.0.dev2 → lycoris_lora-3.0.0.dev4}/setup.py

@@ -4,7 +4,7 @@ from setuptools import setup, find_packages
 setup(
     name="lycoris_lora",
     packages=find_packages(),
-    version="3.0.0.dev2",
+    version="3.0.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",
@@ -12,9 +12,6 @@ setup(
     zip_safe=False,
     install_requires=[
         "torch",
-        "safetensors",
-        "diffusers",
-        "transformers",
         "einops",
         "toml",
     ],

lycoris_lora-3.0.0.dev2/lycoris/modules/ia3.py

@@ -1,67 +0,0 @@
-import torch
-import torch.nn as nn
-
-from .base import ModuleCustomSD
-
-
-class IA3Module(ModuleCustomSD):
-    """
-    Hadamard product Implementaion for Low Rank Adaptation
-    """
-
-    def __init__(
-        self,
-        lora_name,
-        org_module: nn.Module,
-        multiplier=1.0,
-        train_on_input=False,
-        **kwargs
-    ):
-        """if alpha == 0 or None, alpha is rank (no scaling)."""
-        super().__init__()
-        self.lora_name = lora_name
-        self.tucker = False
-
-        self.shape = org_module.weight.shape
-        if org_module.__class__.__name__ == "Conv2d":
-            in_dim = org_module.in_channels
-            out_dim = org_module.out_channels
-            if train_on_input:
-                train_dim = in_dim
-            else:
-                train_dim = out_dim
-            self.weight = nn.Parameter(torch.empty(1, train_dim, 1, 1))
-        else:
-            in_dim = org_module.in_features
-            out_dim = org_module.out_features
-            if train_on_input:
-                train_dim = in_dim
-            else:
-                train_dim = out_dim
-
-            self.weight = nn.Parameter(torch.empty(train_dim))
-
-        # Need more experiences on init method
-        torch.nn.init.constant_(self.weight, 0)
-
-        self.multiplier = multiplier
-        self.org_forward = None
-        self.org_module = [org_module]  # remove in applying
-        self.grad_ckpt = False
-        self.train_input = train_on_input
-        self.register_buffer("on_input", torch.tensor(int(train_on_input)))
-
-    def apply_to(self):
-        self.org_forward = self.org_module[0].forward
-        self.org_module[0].forward = self.forward
-
-    @torch.enable_grad()
-    def forward(self, x):
-        if self.train_input:
-            x = x * (1 + self.weight * self.multiplier)
-        out = self.org_forward(x)
-        dtype = out.dtype
-        if not self.train_input:
-            out = out * (1 + self.weight * self.multiplier)
-            out = out.to(dtype)
-        return out

lycoris_lora-3.0.0.dev2/lycoris_lora.egg-info/requires.txt

@@ -1,6 +0,0 @@
-torch
-safetensors
-diffusers
-transformers
-einops
-toml