lycoris-lora 3.0.0.dev8__tar.gz → 3.0.0.dev9__tar.gz

{lycoris_lora-3.0.0.dev8/lycoris_lora.egg-info → lycoris_lora-3.0.0.dev9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: lycoris_lora
-Version: 3.0.0.dev8
+Version: 3.0.0.dev9
 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.0.dev8 → lycoris_lora-3.0.0.dev9}/README.md

@@ -74,7 +74,7 @@ In any case, please install this package in the corresponding virtual environmen
 - through pip
 
   ```bash
-  pip install lycoris_lora
+  pip install lycoris-lora
   ```
 - or from source
 
@@ -111,6 +111,8 @@ You can use this package's kohya module to run kohya's training script to train
 
 #### HCP-Diffusion
 
+**The support for HCP-Diffusion has been dropped on LyCORIS3.0.0, we will wait until HCP side finish the implementation of new wrapper**
+
 You can use this package's hcp module to run HCP-Diffusion's training script to train lycoris module for SD models
 
 ```bash
@@ -156,6 +158,11 @@ forward_with_lyco = your_model(x)
 
 You can check my [HakuPhi](https://github.com/KohakuBlueleaf/HakuPhi) project to see how I utilize LyCORIS to finetune the Phi-1.5 models.
 
+#### Other method
+After LyCORIS3.0.0, Parametrize API and Functional API have been added, which provide more different ways on utilizing LyCORIS library.
+
+Check [xxx] for more informations.
+
 #### Bitsandbytes support
 
 See [bnb_example.py](bnb_example.py) for example. Basically as same as standalone wrapper.

lycoris_lora-3.0.0.dev9/lycoris/functional/lokr.py

@@ -0,0 +1,259 @@
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .general import rebuild_tucker, FUNC_LIST
+from .general import factorization
+
+
+def make_kron(w1, w2, scale):
+    for _ in range(w2.dim() - w1.dim()):
+        w1 = w1.unsqueeze(-1)
+    w2 = w2.contiguous()
+    rebuild = torch.kron(w1, w2)
+
+    return rebuild * scale
+
+
+def weight_gen(
+    org_weight,
+    rank,
+    tucker=True,
+    factor=-1,
+    decompose_both=False,
+    full_matrix=False,
+    unbalanced_factorization=False,
+):
+    """### weight_gen
+
+    Args:
+        org_weight (torch.Tensor): the weight tensor
+        rank (int): low rank
+
+    Returns:
+        torch.Tensor | None: w1, w1a, w1b, w2, w2a, w2b, t2
+    """
+    out_dim, in_dim, *k = org_weight.shape
+    w1 = w1a = w1b = None
+    w2 = w2a = w2b = None
+    t2 = None
+    use_w1 = use_w2 = False
+
+    if k:
+        k_size = k
+        shape = (out_dim, in_dim, *k_size)
+
+        in_m, in_n = factorization(in_dim, factor)
+        out_l, out_k = factorization(out_dim, factor)
+        if unbalanced_factorization:
+            out_l, out_k = out_k, out_l
+        shape = ((out_l, out_k), (in_m, in_n), *k_size)  # ((a, b), (c, d), *k_size)
+        tucker = tucker and any(i != 1 for i in k_size)
+        if (
+            decompose_both
+            and rank < max(shape[0][0], shape[1][0]) / 2
+            and not full_matrix
+        ):
+            w1a = torch.empty(shape[0][0], rank)
+            w1b = torch.empty(rank, shape[1][0])
+        else:
+            use_w1 = True
+            w1 = torch.empty(shape[0][0], shape[1][0])  # a*c, 1-mode
+
+        if rank >= max(shape[0][1], shape[1][1]) / 2 or full_matrix:
+            use_w2 = True
+            w2 = torch.empty(shape[0][1], shape[1][1], *k_size)
+        elif tucker:
+            t2 = torch.empty(rank, rank, *shape[2:])
+            w2a = torch.empty(rank, shape[0][1])  # b, 1-mode
+            w2b = torch.empty(rank, shape[1][1])  # d, 2-mode
+        else:  # Conv2d not tucker
+            # bigger part. weight and LoRA. [b, dim] x [dim, d*k1*k2]
+            w2a = torch.empty(shape[0][1], rank)
+            w2b = torch.empty(rank, shape[1][1], *shape[2:])
+            # w1 ⊗ (w2a x w2b) = (a, b)⊗((c, dim)x(dim, d*k1*k2)) = (a, b)⊗(c, d*k1*k2) = (ac, bd*k1*k2)
+    else:  # Linear
+        shape = (out_dim, in_dim)
+
+        in_m, in_n = factorization(in_dim, factor)
+        out_l, out_k = factorization(out_dim, factor)
+        if unbalanced_factorization:
+            out_l, out_k = out_k, out_l
+        shape = (
+            (out_l, out_k),
+            (in_m, in_n),
+        )  # ((a, b), (c, d)), out_dim = a*c, in_dim = b*d
+        # smaller part. weight scale
+        if decompose_both and rank < max(shape[0][0], shape[1][0]) / 2:
+            w1a = torch.empty(shape[0][0], rank)
+            w1b = torch.empty(rank, shape[1][0])
+        else:
+            use_w1 = True
+            w1 = torch.empty(shape[0][0], shape[1][0])  # a*c, 1-mode
+        if rank < max(shape[0][1], shape[1][1]) / 2:
+            # bigger part. weight and LoRA. [b, dim] x [dim, d]
+            w2a = torch.empty(shape[0][1], rank)
+            w2b = torch.empty(rank, shape[1][1])
+            # w1 ⊗ (w2a x w2b) = (a, b)⊗((c, dim)x(dim, d)) = (a, b)⊗(c, d) = (ac, bd)
+        else:
+            use_w2 = True
+            w2 = torch.empty(shape[0][1], shape[1][1])
+
+    if use_w2:
+        torch.nn.init.constant_(w2, 1)
+    else:
+        if tucker:
+            torch.nn.init.kaiming_uniform_(t2, a=math.sqrt(5))
+        torch.nn.init.kaiming_uniform_(w2a, a=math.sqrt(5))
+        torch.nn.init.constant_(w2b, 1)
+
+    if use_w1:
+        torch.nn.init.kaiming_uniform_(w1, a=math.sqrt(5))
+    else:
+        torch.nn.init.kaiming_uniform_(w1a, a=math.sqrt(5))
+        torch.nn.init.kaiming_uniform_(w1b, a=math.sqrt(5))
+
+    return w1, w1a, w1b, w2, w2a, w2b, t2
+
+
+def lokr_diff_weight(w1, w1a, w1b, w2, w2a, w2b, t, gamma=1.0):
+    """### lokr_diff_weight
+
+    Args:
+        gamma (float, optional): scale factor, normally alpha/rank here
+
+    Returns:
+        torch.Tensor: ΔW
+    """
+    if w1a is not None:
+        rank = w1a.shape[1]
+    elif w2a is not None:
+        rank = w2a.shape[1]
+    else:
+        rank = gamma
+    scale = gamma / rank
+    if w1 is None:
+        w1 = w1a @ w1b
+    if w2 is None:
+        if t is None:
+            r, o, *k = w2b.shape
+            w2 = w2a @ w2b.view(r, -1)
+            w2 = w2.view(-1, o, *k)
+        else:
+            w2 = rebuild_tucker(t, w2a, w2b)
+    return make_kron(w1, w2, scale)
+
+
+def lokr_bypass_forward_diff(
+    h, w1, w1a, w1b, w2, w2a, w2b, t, gamma=1.0, extra_args={}
+):
+    """### lokr_bypass_forward_diff
+
+    Args:
+        gamma (float, optional): scale factor, normally alpha/rank here
+        extra_args (dict, optional): extra args for forward func, \
+            e.g. padding, stride for Conv1/2/3d
+
+    Returns:
+        torch.Tensor: output tensor
+    """
+    use_w1 = w1 is not None
+    use_w2 = w2 is not None
+    tucker = t is not None
+    dim = t.dim() if tucker else w2.dim() if w2 is not None else w2b.dim()
+    print(dim)
+    rank = w1b.size(0) if not use_w1 else w2b.size(0) if not use_w2 else gamma
+    scale = gamma / rank
+    is_conv = dim > 2
+    op = FUNC_LIST[dim]
+
+    if is_conv:
+        kw_dict = extra_args
+    else:
+        kw_dict = {}
+
+    if use_w2:
+        ba = w2
+    else:
+        a = w2b
+        b = w2a
+
+        if t is not None:
+            a = a.view(*a.shape, *[1] * (dim - 2))
+            b = b.view(*b.shape, *[1] * (dim - 2))
+        elif is_conv:
+            b = b.view(*b.shape, *[1] * (dim - 2))
+
+    if use_w1:
+        c = w1
+    else:
+        c = w1a @ w1b
+    uq = c.size(1)
+
+    if is_conv:
+        # (b, uq), vq, ...
+        B, _, *rest = h.shape
+        h_in_group = h.reshape(B * uq, -1, *rest)
+    else:
+        # b, ..., uq, vq
+        h_in_group = h.reshape(*h.shape[:-1], uq, -1)
+
+    if use_w2:
+        hb = op(h_in_group, ba, **kw_dict)
+    else:
+        if is_conv:
+            if tucker:
+                ha = op(h_in_group, a)
+                ht = op(ha, t, **kw_dict)
+                hb = op(ht, b)
+            else:
+                ha = op(h_in_group, a, **kw_dict)
+                hb = op(ha, b)
+        else:
+            ha = op(h_in_group, a, **kw_dict)
+            hb = op(ha, b)
+
+    if is_conv:
+        # (b, uq), vp, ..., f
+        # -> b, uq, vp, ..., f
+        # -> b, f, vp, ..., uq
+        hb = hb.view(B, -1, *hb.shape[1:])
+        h_cross_group = hb.transpose(1, -1)
+    else:
+        # b, ..., uq, vq
+        # -> b, ..., vq, uq
+        h_cross_group = hb.transpose(-1, -2)
+
+    hc = F.linear(h_cross_group, c)
+    if is_conv:
+        # b, f, vp, ..., up
+        # -> b, up, vp, ... ,f
+        # -> b, c, ..., f
+        hc = hc.transpose(1, -1)
+        h = hc.reshape(B, -1, *hc.shape[3:])
+    else:
+        # b, ..., vp, up
+        # -> b, ..., up, vp
+        # -> b, ..., c
+        hc = hc.transpose(-1, -2)
+        h = hc.reshape(*hc.shape[:-2], -1)
+
+    return h * scale
+
+
+if __name__ == "__main__":
+    w = torch.randn(32, 32, 3, 3, 3)
+    w1, w1a, w1b, w2, w2a, w2b, t = weight_gen(w, 2, tucker=False)
+    extra_args = {"padding": 1}
+
+    x = torch.randn(1, 32, 8, 8, 8)
+    y = FUNC_LIST[w.dim()](x, w, **extra_args)
+    diff_w = lokr_diff_weight(w1, w1a, w1b, w2, w2a, w2b, t, 1)
+    diff_y_w = FUNC_LIST[w.dim()](x, diff_w, **extra_args)
+    diff_y = lokr_bypass_forward_diff(x, w1, w1a, w1b, w2, w2a, w2b, t, 1, extra_args)
+
+    print(F.mse_loss(y, y + diff_y))
+    print(F.mse_loss(w, w + diff_w))
+    print(F.mse_loss(diff_y, diff_y_w))

{lycoris_lora-3.0.0.dev8 → lycoris_lora-3.0.0.dev9}/lycoris/modules/base.py

@@ -184,6 +184,10 @@ class LycorisBaseModule(ModuleCustomSD):
 
     @classmethod
     def parametrize(cls, org_module, attr, *args, **kwargs):
+        from .full import FullModule
+
+        if cls is FullModule:
+            raise RuntimeError("FullModule cannot be used for parametrize.")
         target_param = getattr(org_module, attr)
         kwargs["bypass_mode"] = False
         if target_param.dim() == 2:
@@ -206,11 +210,12 @@ class LycorisBaseModule(ModuleCustomSD):
                 target_param.shape[0],
                 target_param.shape[1],
                 *target_param.shape[2:],
+                bias=False,
             )
             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)
+        module_obj.to(target_param)
         parametrize.register_parametrization(org_module, attr, module_obj)
         return module_obj
 

{lycoris_lora-3.0.0.dev8 → lycoris_lora-3.0.0.dev9}/lycoris/modules/boft.py

@@ -3,7 +3,6 @@ from math import log2
 
 import torch
 import torch.nn as nn
-import torch.nn.functional as F
 
 from einops import rearrange
 
@@ -197,46 +196,3 @@ class ButterflyOFTModule(LycorisBaseModule):
             w = self.make_weight(scale, x.device)
             kw_dict = self.kw_dict | {"weight": w, "bias": self.org_module[0].bias}
             return self.op(x, **kw_dict)
-
-
-if __name__ == "__main__":
-    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.dev8 → lycoris_lora-3.0.0.dev9}/lycoris/modules/diag_oft.py

@@ -4,8 +4,6 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
-from einops import rearrange
-
 from .base import LycorisBaseModule
 from ..functional import factorization
 from ..logging import logger
@@ -182,46 +180,3 @@ class DiagOFTModule(LycorisBaseModule):
             w = self.make_weight(scale, x.device)
             kw_dict = self.kw_dict | {"weight": w, "bias": self.org_module[0].bias}
             return self.op(x, **kw_dict)
-
-
-if __name__ == "__main__":
-    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.dev8 → lycoris_lora-3.0.0.dev9}/lycoris/modules/dylora.py

@@ -3,7 +3,6 @@ import random
 
 import torch
 import torch.nn as nn
-import torch.nn.functional as F
 
 from .base import LycorisBaseModule
 from ..utils import product
@@ -156,46 +155,3 @@ class DyLoraModule(LycorisBaseModule):
                 else self.org_module[0].bias.data
             )
             return self.op(x, weight, bias, **self.kw_dict)
-
-
-if __name__ == "__main__":
-    device = torch.device("cuda")
-    module = DyLoraModule
-    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.dev8 → lycoris_lora-3.0.0.dev9}/lycoris/modules/full.py

@@ -1,4 +1,3 @@
-from typing import Mapping
 import torch
 import torch.nn as nn
 
@@ -77,6 +76,12 @@ class FullModule(LycorisBaseModule):
         else:
             self.org_bias = None
 
+    def restore(self):
+        self.org_module[0].forward = self.org_forward
+        self.org_module[0].weight = nn.Parameter(self._org_weight[0])
+        if self.org_bias is not None:
+            self.org_module[0].bias = nn.Parameter(self.org_bias[0])
+
     def load_state_dict(self, state_dict, strict: bool = True, assign: bool = False):
         return
 
@@ -141,7 +146,7 @@ class FullModule(LycorisBaseModule):
         if shape is not None:
             weight = weight.view(shape)
             if bias is not None:
-                bias = bias.view(shape[0], 1)
+                bias = bias.view(shape[0])
         return weight, bias
 
     def forward(self, x: torch.Tensor, *args, **kwargs):
@@ -159,36 +164,3 @@ class FullModule(LycorisBaseModule):
         weight, bias = self.make_weight(scale, x.device)
         kw_dict = self.kw_dict | {"weight": weight, "bias": bias}
         return self.op(x, **kw_dict)
-
-
-if __name__ == "__main__":
-    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.dev8 → lycoris_lora-3.0.0.dev9}/lycoris/modules/glora.py

@@ -202,46 +202,3 @@ class GLoRAModule(LycorisBaseModule):
                 else self.org_module[0].bias.data
             )
             return self.op(x, weight, bias, **self.kw_dict)
-
-
-if __name__ == "__main__":
-    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.dev8 → lycoris_lora-3.0.0.dev9}/lycoris/modules/ia3.py

@@ -128,46 +128,3 @@ class IA3Module(LycorisBaseModule):
                 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.dev8 → lycoris_lora-3.0.0.dev9}/lycoris/modules/locon.py

@@ -74,7 +74,7 @@ class LoConModule(LycorisBaseModule):
             self.down_op = self.op
             self.up_op = self.op
             if use_tucker and any(i != 1 for i in k_size):
-                self.lora_down = self.module(in_dim, lora_dim, (1, 1), bias=False)
+                self.lora_down = self.module(in_dim, lora_dim, 1, bias=False)
                 self.lora_mid = self.module(
                     lora_dim, lora_dim, k_size, stride, padding, bias=False
                 )
@@ -83,7 +83,7 @@ class LoConModule(LycorisBaseModule):
                 self.lora_down = self.module(
                     in_dim, lora_dim, k_size, stride, padding, bias=False
                 )
-            self.lora_up = self.module(lora_dim, out_dim, (1, 1), bias=False)
+            self.lora_up = self.module(lora_dim, out_dim, 1, bias=False)
         elif isinstance(org_module, nn.Linear):
             self.isconv = False
             self.down_op = F.linear
@@ -277,46 +277,3 @@ class LoConModule(LycorisBaseModule):
             return self.op(x, weight, bias, **self.kw_dict)
         else:
             return self.bypass_forward(x, scale=self.multiplier)
-
-
-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 = 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.dev8 → lycoris_lora-3.0.0.dev9}/lycoris/modules/loha.py

@@ -2,7 +2,6 @@ import math
 
 import torch
 import torch.nn as nn
-import torch.nn.functional as F
 
 from .base import LycorisBaseModule
 from ..functional.loha import loha_diff_weight
@@ -267,46 +266,3 @@ class LohaModule(LycorisBaseModule):
             if self.wd:
                 weight = self.apply_weight_decompose(weight)
             return self.op(x, weight, bias, **self.kw_dict)
-
-
-if __name__ == "__main__":
-    device = torch.device("cuda")
-    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.dev8 → lycoris_lora-3.0.0.dev9}/lycoris/modules/lokr.py

@@ -9,6 +9,7 @@ from einops import rearrange
 
 from .base import LycorisBaseModule
 from ..functional import factorization, rebuild_tucker
+from ..functional.lokr import make_kron
 from ..logging import logger
 from ..utils.bnb import LinearNF4
 
@@ -22,15 +23,6 @@ def logging_force_full_matrix(lora_dim, dim, factor):
     )
 
 
-def make_kron(w1, w2, scale):
-    if len(w2.shape) == 4:
-        w1 = w1.unsqueeze(2).unsqueeze(2)
-    w2 = w2.contiguous()
-    rebuild = torch.kron(w1, w2)
-
-    return rebuild * scale
-
-
 class LokrModule(LycorisBaseModule):
     support_module = {
         "linear",
@@ -115,9 +107,7 @@ class LokrModule(LycorisBaseModule):
                     torch.empty(shape[0][1], shape[1][1], *k_size)
                 )
             elif self.tucker:
-                self.lokr_t2 = nn.Parameter(
-                    torch.empty(lora_dim, lora_dim, shape[2], shape[3])
-                )
+                self.lokr_t2 = nn.Parameter(torch.empty(lora_dim, lora_dim, *shape[2:]))
                 self.lokr_w2_a = nn.Parameter(
                     torch.empty(lora_dim, shape[0][1])
                 )  # b, 1-mode
@@ -339,40 +329,74 @@ class LokrModule(LycorisBaseModule):
         return scaled, orig_norm * ratio
 
     def bypass_forward_diff(self, h, scale=1):
-        if len(self.shape) > 2 and self.shape[2] > 1:
-            return (
-                self.op(h, self.get_weight(self.shape), **self.kw_dict)
-                * self.multiplier
-            )
-        if self.module_type.startswith("conv"):
-            h = h.transpose(1, -1)
+        is_conv = self.module_type.startswith("conv")
         if self.use_w2:
             ba = self.lokr_w2
         else:
             a = self.lokr_w2_b
             b = self.lokr_w2_a
+
+            if self.tucker:
+                t = self.lokr_t2
+                a = a.view(*a.shape, *[1] * (len(t.shape) - 2))
+                b = b.view(*b.shape, *[1] * (len(t.shape) - 2))
+            elif is_conv:
+                a = a.view(*a.shape, *self.shape[2:])
+                b = b.view(*b.shape, *[1] * (len(self.shape) - 2))
+
         if self.use_w1:
             c = self.lokr_w1
         else:
             c = self.lokr_w1_a @ self.lokr_w1_b
         uq = c.size(1)
 
+        if is_conv:
+            # (b, uq), vq, ...
+            b, _, *rest = h.shape
+            h_in_group = h.reshape(b * uq, -1, *rest)
+        else:
+            # b, ..., uq, vq
+            h_in_group = h.reshape(*h.shape[:-1], uq, -1)
+
         if self.use_w2:
-            vq = ba.size(1)
-            h_in_group = rearrange(h, "b ... (uq vq) -> b ... uq vq", uq=uq, vq=vq)
-            hb = F.linear(h_in_group, ba)
+            hb = self.op(h_in_group, ba, **self.kw_dict)
+        else:
+            if is_conv:
+                if self.tucker:
+                    ha = self.op(h_in_group, a)
+                    ht = self.op(ha, t, **self.kw_dict)
+                    hb = self.op(ht, b)
+                else:
+                    ha = self.op(h_in_group, a, **self.kw_dict)
+                    hb = self.op(ha, b)
+            else:
+                ha = self.op(h_in_group, a, **self.kw_dict)
+                hb = self.op(ha, b)
+
+        if is_conv:
+            # (b, uq), vp, ..., f
+            # -> b, uq, vp, ..., f
+            # -> b, f, vp, ..., uq
+            hb = hb.view(b, -1, *hb.shape[1:])
+            h_cross_group = hb.transpose(1, -1)
         else:
-            vq = a.size(1)
-            h_in_group = rearrange(h, "b ... (uq vq) -> b ... uq vq", uq=uq, vq=vq)
-            ha = F.linear(h_in_group, a)
-            hb = F.linear(ha, b)
+            # b, ..., uq, vq
+            # -> b, ..., vq, uq
+            h_cross_group = hb.transpose(-1, -2)
 
-        h_cross_group = hb.transpose(-1, -2)
         hc = F.linear(h_cross_group, c)
-
-        h = rearrange(hc, "b ... vp up -> b ... (up vp)")
-        if self.module_type.startswith("conv"):
-            h = h.transpose(1, -1)
+        if is_conv:
+            # b, f, vp, ..., up
+            # -> b, up, vp, ... ,f
+            # -> b, c, ..., f
+            hc = hc.transpose(1, -1)
+            h = hc.reshape(b, -1, *hc.shape[3:])
+        else:
+            # b, ..., vp, up
+            # -> b, ..., up, vp
+            # -> b, ..., c
+            hc = hc.transpose(-1, -2)
+            h = hc.reshape(*hc.shape[:-2], -1)
 
         return self.drop(h * scale)
 
@@ -402,43 +426,44 @@ class LokrModule(LycorisBaseModule):
 
 
 if __name__ == "__main__":
-    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)
+    base = nn.Conv2d(128, 128, 3, 1, 1)
+    net = LokrModule(
+        "",
+        base,
+        multiplier=1,
+        lora_dim=4,
+        alpha=1,
+        weight_decompose=False,
+        use_tucker=False,
+        use_scalar=False,
+        decompose_both=True,
+    )
+    net.apply_to()
+    sd = net.state_dict()
+    for key in sd:
+        if key != "alpha":
+            sd[key] = torch.randn_like(sd[key])
+    net.load_state_dict(sd)
+
+    test_input = torch.randn(1, 128, 16, 16)
+    test_output = net(test_input)
+    print(test_output.shape)
+
+    net2 = LokrModule(
+        "",
+        base,
+        multiplier=1,
+        lora_dim=4,
+        alpha=1,
+        weight_decompose=False,
+        use_tucker=False,
+        use_scalar=False,
+        bypass_mode=True,
+        decompose_both=True,
+    )
+    net2.apply_to()
+    net2.load_state_dict(sd)
+    print(net2)
+
+    test_output2 = net(test_input)
+    print(F.mse_loss(test_output, test_output2))

{lycoris_lora-3.0.0.dev8 → lycoris_lora-3.0.0.dev9}/lycoris/utils/bnb.py

@@ -2,7 +2,10 @@ from functools import cache
 
 try:
     from bitsandbytes.nn import LinearNF4, Linear8bitLt, LinearFP4
+
+    SUPPORT_BNB = True
 except:
+    SUPPORT_BNB = False
     import torch.nn as nn
 
     class LinearNF4(nn.Linear):

{lycoris_lora-3.0.0.dev8 → lycoris_lora-3.0.0.dev9/lycoris_lora.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: lycoris_lora
-Version: 3.0.0.dev8
+Version: 3.0.0.dev9
 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.0.dev8 → lycoris_lora-3.0.0.dev9}/lycoris_lora.egg-info/SOURCES.txt

@@ -10,6 +10,7 @@ lycoris/functional/__init__.py
 lycoris/functional/general.py
 lycoris/functional/locon.py
 lycoris/functional/loha.py
+lycoris/functional/lokr.py
 lycoris/modules/__init__.py
 lycoris/modules/base.py
 lycoris/modules/boft.py

{lycoris_lora-3.0.0.dev8 → lycoris_lora-3.0.0.dev9}/setup.py

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