lycoris-lora 2.2.0.dev6__tar.gz → 2.2.0.dev8__tar.gz

{lycoris_lora-2.2.0.dev6/lycoris_lora.egg-info → lycoris_lora-2.2.0.dev8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: lycoris_lora
-Version: 2.2.0.dev6
+Version: 2.2.0.dev8
 Summary: Lora beYond Conventional methods, Other Rank adaptation Implementations for Stable diffusion
 Home-page: https://github.com/KohakuBlueleaf/LyCORIS
 Author: Shih-Ying Yeh(KohakuBlueLeaf), Yu-Guan Hsieh, Zhidong Gao

lycoris_lora-2.2.0.dev6/lycoris/modules/locon.py → lycoris_lora-2.2.0.dev8/lycoris/modules/lilora.py

@@ -8,7 +8,6 @@ import torch.nn.functional as F
 from .base import ModuleCustomSD
 from ..logging import logger
 
-
 @cache
 def log_wd():
     return logger.warning(
@@ -17,7 +16,7 @@ def log_wd():
     )
 
 
-class LoConModule(ModuleCustomSD):
+class LiLoConModule(ModuleCustomSD):
     """
     modifed from kohya-ss/sd-scripts/networks/lora:LoRAModule
     """
@@ -36,6 +35,7 @@ class LoConModule(ModuleCustomSD):
         use_scalar=False,
         rank_dropout_scale=False,
         weight_decompose=False,
+        bypass_mode=False,
         **kwargs,
     ):
         """if alpha == 0 or None, alpha is rank (no scaling)."""
@@ -43,8 +43,9 @@ class LoConModule(ModuleCustomSD):
         self.lora_name = lora_name
         self.lora_dim = lora_dim
         self.tucker = False
+        assert not (bypass_mode and weight_decompose), "bypass_mode and dora_wd cannot be used together"
 
-        if org_module.__class__.__name__ == "Conv2d":
+        if isinstance(org_module, nn.Conv2d):
             in_dim = org_module.in_channels
             k_size = org_module.kernel_size
             out_dim = org_module.out_channels
@@ -122,7 +123,8 @@ class LoConModule(ModuleCustomSD):
         if use_scalar:
             self.scalar = nn.Parameter(torch.tensor(0.0))
         else:
-            self.scalar = torch.tensor(1.0)
+            self.scalar: torch.FloatTensor
+            self.register_buffer("scalar", torch.tensor(1.0), persistent=False)
         # same as microsoft's
         torch.nn.init.kaiming_uniform_(self.lora_down.weight, a=math.sqrt(5))
         if use_scalar:
@@ -176,7 +178,7 @@ class LoConModule(ModuleCustomSD):
                 drop /= drop.mean()
             weight *= drop
 
-        return weight * self.scalar
+        return weight * self.scalar.to(device)
 
     def apply_weight_decompose(self, weight):
         return (
@@ -319,60 +321,3 @@ class LoConModule(ModuleCustomSD):
         org_out = self.org_forward(x)
         # print(x.shape, org_out.shape, up.shape)
         return org_out + self.dropout(up) * scale
-
-    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.wd:
-            weight = (
-                self.org_module[0].weight.data.to(self.scalar)
-                + self.make_weight(x.device).to(self.scalar) * scale
-            )
-            if self.wd:
-                weight = self.apply_weight_decompose(weight)
-            bias = (
-                None
-                if self.org_module[0].bias is None
-                else self.org_module[0].bias.data
-            )
-            return self.op(x, weight, bias, **self.extra_args)
-        else:
-            if self.tucker:
-                mid = self.lora_mid(self.lora_down(x))
-            else:
-                mid = self.lora_down(x)
-
-            if self.rank_dropout and self.training:
-                drop = (
-                    torch.rand(self.lora_dim, device=mid.device) > self.rank_dropout
-                ).to(mid.dtype)
-                if self.rank_dropout_scale:
-                    drop /= drop.mean()
-                if (dims := len(x.shape)) == 4:
-                    drop = drop.view(1, -1, 1, 1)
-                else:
-                    drop = drop.view(*[1] * (dims - 1), -1)
-                mid = mid * drop
-
-            return self.org_forward(x) + self.dropout(
-                self.lora_up(mid) * self.scalar * scale
-            )
-
-
-if __name__ == "__main__":
-    base = nn.Linear(128, 128)
-    lokr = LoConModule("test", base, 1, 4, 1, weight_decompose=True)
-    print(lokr)
-    test_input = torch.randn(1, 128)
-    test_output = lokr(test_input)
-    print(test_output.shape)
-
-    base = nn.Conv2d(128, 128, 3, 1, 1)
-    lokr = LoConModule("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)
-    print(test_output.shape)

lycoris_lora-2.2.0.dev8/lycoris/modules/locon.py

@@ -0,0 +1,288 @@
+import math
+from functools import cache
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .base import ModuleCustomSD
+from ..logging import logger
+from ..utils.bnb import (
+    LinearNF4,
+    QuantLinears,
+    log_bypass
+)
+
+@cache
+def log_wd():
+    return logger.warning(
+        "Using weight_decompose=True with LoRA (DoRA) will ignore network_dropout."
+        "Only rank dropout and module dropout will be applied"
+    )
+
+
+class LoConModule(ModuleCustomSD):
+    """
+    modifed from kohya-ss/sd-scripts/networks/lora:LoRAModule
+    """
+
+    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,
+        **kwargs,
+    ):
+        """if alpha == 0 or None, alpha is rank (no scaling)."""
+        super().__init__()
+        self.lora_name = lora_name
+        self.lora_dim = lora_dim
+        self.tucker = False
+
+        if isinstance(org_module, nn.Conv2d):
+            in_dim = org_module.in_channels
+            k_size = org_module.kernel_size
+            out_dim = org_module.out_channels
+            self.tucker = use_tucker and k_size != (1, 1)
+            self.op = F.conv2d
+            self.extra_args = {
+                "stride": org_module.stride,
+                "padding": org_module.padding,
+                "dilation": org_module.dilation,
+                "groups": org_module.groups,
+            }
+        else:
+            in_dim = org_module.in_features
+            out_dim = org_module.out_features
+            self.op = F.linear
+            self.extra_args = {}
+
+        if isinstance(org_module, QuantLinears):
+            if not bypass_mode:
+                log_bypass()
+            bypass_mode = True
+        self.bypass_mode = bypass_mode
+        assert not (bypass_mode and weight_decompose), "bypass_mode and dora_wd cannot be used together"
+
+        if isinstance(org_module, nn.Conv2d):
+            self.isconv = True
+            # For general LoCon
+            in_dim = org_module.in_channels
+            k_size = org_module.kernel_size
+            stride = org_module.stride
+            padding = org_module.padding
+            out_dim = org_module.out_channels
+            self.down_op = F.conv2d
+            self.up_op = F.conv2d
+            if use_tucker and k_size != (1, 1):
+                self.lora_down = nn.Conv2d(in_dim, lora_dim, (1, 1), bias=False)
+                self.lora_mid = nn.Conv2d(
+                    lora_dim, lora_dim, k_size, stride, padding, bias=False
+                )
+                self.tucker = True
+            else:
+                self.lora_down = nn.Conv2d(
+                    in_dim, lora_dim, k_size, stride, padding, bias=False
+                )
+            self.lora_up = nn.Conv2d(lora_dim, out_dim, (1, 1), bias=False)
+        elif isinstance(org_module, nn.Linear):
+            self.isconv = False
+            self.down_op = F.linear
+            self.up_op = F.linear
+            in_dim = org_module.in_features
+            out_dim = org_module.out_features
+            self.lora_down = nn.Linear(in_dim, lora_dim, bias=False)
+            self.lora_up = nn.Linear(lora_dim, out_dim, bias=False)
+        else:
+            raise NotImplementedError
+        self.shape = org_module.weight.shape
+
+        self.wd = weight_decompose
+        if self.wd:
+            org_weight: nn.Parameter = org_module.weight
+            self.dora_mean_dim = tuple(i for i in range(org_weight.dim()) if i != 1)
+            self.dora_scale = nn.Parameter(
+                torch.mean(org_weight, dim=self.dora_mean_dim, keepdim=True)
+            )
+
+        if dropout:
+            self.dropout = nn.Dropout(dropout)
+            if self.wd:
+                log_wd()
+        else:
+            self.dropout = nn.Identity()
+        self.rank_dropout = rank_dropout
+        self.rank_dropout_scale = rank_dropout_scale
+        self.module_dropout = module_dropout
+
+        if type(alpha) == torch.Tensor:
+            alpha = alpha.detach().float().numpy()  # without casting, bf16 causes error
+        alpha = lora_dim if alpha is None or alpha == 0 else alpha
+        self.scale = alpha / self.lora_dim
+        self.register_buffer("alpha", torch.tensor(alpha))  # 定数として扱える
+
+        if use_scalar:
+            self.scalar = nn.Parameter(torch.tensor(0.0))
+        else:
+            self.scalar: torch.FloatTensor
+            self.register_buffer("scalar", torch.tensor(1.0), persistent=False)
+        # same as microsoft's
+        torch.nn.init.kaiming_uniform_(self.lora_down.weight, a=math.sqrt(5))
+        if use_scalar:
+            torch.nn.init.kaiming_uniform_(self.lora_up.weight, a=math.sqrt(5))
+        else:
+            torch.nn.init.constant_(self.lora_up.weight, 0)
+        if self.tucker:
+            torch.nn.init.kaiming_uniform_(self.lora_mid.weight, a=math.sqrt(5))
+
+        self.multiplier = multiplier
+        self.org_module = [org_module]
+        self.org_forward = self.org_module[0].forward
+
+    def load_weight_hook(self, module: nn.Module, incompatible_keys):
+        missing_keys = incompatible_keys.missing_keys
+        for key in missing_keys:
+            if "scalar" in key:
+                del missing_keys[missing_keys.index(key)]
+        self.scalar = nn.Parameter(torch.ones_like(self.scalar))
+
+    def apply_to(self, is_hypernet=False, **kwargs):
+        self.org_forward = self.org_module[0].forward
+        if is_hypernet:
+            self.org_module[0].forward = self.hypernet_forward
+        else:
+            self.org_module[0].forward = self.forward
+
+    def restore(self):
+        self.org_module[0].forward = self.org_forward
+
+    def merge_to(self, multiplier=1.0):
+        weight = self.make_weight() * self.scale * multiplier
+        self.org_module[0].weight.data.add_(weight)
+
+    def make_weight(self, device=None):
+        wa = self.lora_up.weight.to(device)
+        wb = self.lora_down.weight.to(device)
+        if self.tucker:
+            t = self.lora_mid.weight.to(device)
+            wa = wa.squeeze(-1, -2)
+            wb = wb.squeeze(-1, -2)
+            weight = torch.einsum("i j k l, p i, j r -> p r k l", t, wa, wb)
+        else:
+            weight = wa.view(wa.size(0), -1) @ wb.view(wb.size(0), -1)
+
+        weight = weight.view(self.shape)
+        if self.training and self.rank_dropout:
+            drop = (torch.rand(weight.size(0)) > self.rank_dropout).to(weight.dtype)
+            drop = drop.view(-1, *[1] * len(weight.shape[1:])).to(weight.device)
+            if self.rank_dropout_scale:
+                drop /= drop.mean()
+            weight *= drop
+
+        return weight * self.scalar.to(device)
+
+    def apply_weight_decompose(self, weight):
+        return (
+            weight / weight.mean(dim=self.dora_mean_dim, keepdim=True) * self.dora_scale
+        )
+
+    def custom_state_dict(self):
+        destination = {}
+        if self.wd:
+            destination["dora_scale"] = self.dora_scale
+        destination["alpha"] = self.alpha
+        destination["lora_up.weight"] = self.lora_up.weight * self.scalar
+        destination["lora_down.weight"] = self.lora_down.weight
+        if self.tucker:
+            destination["lora_mid.weight"] = self.lora_mid.weight
+        return destination
+
+    @torch.no_grad()
+    def apply_max_norm(self, max_norm, device=None):
+        orig_norm = self.make_weight(device).norm() * self.scale
+        norm = torch.clamp(orig_norm, max_norm / 2)
+        desired = torch.clamp(norm, max=max_norm)
+        ratio = desired.cpu() / norm.cpu()
+
+        scaled = norm != desired
+        if scaled:
+            self.scalar *= ratio
+
+        return scaled, orig_norm * ratio
+
+    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
+
+        dtype_device = next(self.parameters())
+        if not self.bypass_mode:
+            weight = (
+                self.org_module[0].weight.data.to(dtype_device)
+                + self.make_weight(x.device).to(dtype_device) * scale
+            )
+            if self.wd:
+                weight = self.apply_weight_decompose(weight)
+            bias = (
+                None
+                if self.org_module[0].bias is None
+                else self.org_module[0].bias.data
+            )
+            return self.op(x, weight, bias, **self.extra_args)
+        else:
+            if self.tucker:
+                mid = self.lora_mid(self.lora_down(x))
+            else:
+                mid = self.lora_down(x)
+
+            if self.rank_dropout and self.training:
+                drop = (
+                    torch.rand(self.lora_dim, device=mid.device) > self.rank_dropout
+                ).to(mid.dtype)
+                if self.rank_dropout_scale:
+                    drop /= drop.mean()
+                if (dims := len(x.shape)) == 4:
+                    drop = drop.view(1, -1, 1, 1)
+                else:
+                    drop = drop.view(*[1] * (dims - 1), -1)
+                mid = mid * drop
+
+            return self.org_forward(x) + self.dropout(
+                self.lora_up(mid) * self.scalar * scale
+            )
+
+
+if __name__ == "__main__":
+    base = nn.Linear(128, 128).cuda()
+    lokr = LoConModule("test", base, 1, 4, 1, weight_decompose=True).cuda()
+    print(lokr)
+    test_input = torch.randn(1, 128).cuda()
+    test_output = lokr(test_input)
+    print(test_output.shape)
+
+    base_4bit = LinearNF4(128, 128)
+    base_4bit.load_state_dict(base.state_dict())
+    base_4bit.cuda()
+    qlocon = LoConModule("test", base_4bit, 1, 4, 1, weight_decompose=False).cuda()
+    print(qlocon)
+    test_input = torch.randn(1, 128).cuda()
+    test_output = qlocon(test_input)
+    print(test_output.shape)
+
+    base = nn.Conv2d(128, 128, 3, 1, 1)
+    lokr = LoConModule("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)
+    print(test_output.shape)

{lycoris_lora-2.2.0.dev6 → lycoris_lora-2.2.0.dev8}/lycoris/modules/loha.py

@@ -3,6 +3,11 @@ import torch.nn as nn
 import torch.nn.functional as F
 
 from .base import ModuleCustomSD
+from ..utils.bnb import (
+    LinearNF4,
+    QuantLinears,
+    log_bypass
+)
 
 
 class HadaWeight(torch.autograd.Function):
@@ -100,6 +105,7 @@ class LohaModule(ModuleCustomSD):
         use_scalar=False,
         rank_dropout_scale=False,
         weight_decompose=False,
+        bypass_mode=False,
         **kwargs,
     ):
         """if alpha == 0 or None, alpha is rank (no scaling)."""
@@ -113,6 +119,7 @@ class LohaModule(ModuleCustomSD):
             in_dim = org_module.in_channels
             k_size = org_module.kernel_size
             out_dim = org_module.out_channels
+            self.shape = (out_dim, in_dim, *k_size)
             self.tucker = use_tucker and k_size != (1, 1)
             if self.tucker:
                 shape = (out_dim, in_dim, *k_size)
@@ -129,9 +136,17 @@ class LohaModule(ModuleCustomSD):
             in_dim = org_module.in_features
             out_dim = org_module.out_features
             shape = (out_dim, in_dim)
+            self.shape = shape
             self.op = F.linear
             self.extra_args = {}
 
+        if isinstance(org_module, QuantLinears):
+            if not bypass_mode:
+                log_bypass()
+            bypass_mode = True
+        self.bypass_mode = bypass_mode
+        assert not (bypass_mode and weight_decompose), "bypass_mode and dora_wd cannot be used together"
+
         if self.tucker:
             self.hada_t1 = nn.Parameter(
                 torch.empty(lora_dim, lora_dim, shape[2], shape[3])
@@ -215,10 +230,10 @@ class LohaModule(ModuleCustomSD):
         self.org_module[0].forward = self.org_forward
 
     def merge_to(self, multiplier=1.0):
-        weight = self.get_weight(self.org_module[0].weight)
+        weight = self.get_weight(self.shape)
         self.org_module[0].weight.data.add_(weight * multiplier)
 
-    def get_weight(self, orig_weight=None):
+    def get_weight(self, shape):
         if self.tucker:
             weight = make_weight_tucker(
                 self.hada_t1,
@@ -237,8 +252,8 @@ class LohaModule(ModuleCustomSD):
                 self.hada_w2_b,
                 scale=torch.tensor(self.scale),
             )
-        if orig_weight is not None:
-            weight = weight.reshape(orig_weight.shape)
+        if shape is not None:
+            weight = weight.reshape(shape)
         if self.training and self.rank_dropout:
             drop = (torch.rand(weight.size(0)) > self.rank_dropout).to(weight.dtype)
             drop = drop.view(-1, *[1] * len(weight.shape[1:])).to(weight.device)
@@ -292,14 +307,47 @@ class LohaModule(ModuleCustomSD):
                         else self.org_module[0].bias.data
                     ),
                 )
-        weight = (
-            self.org_module[0].weight.data.to(x.device, dtype=self.hada_w1_a.dtype)
-            + self.get_weight(self.org_module[0].weight.data)
-            * self.scalar
-            * self.multiplier
-        )
-        bias = None if self.org_module[0].bias is None else self.org_module[0].bias.data
-
-        if self.wd:
-            weight = self.apply_weight_decompose(weight)
-        return self.op(x, weight.view(self.shape), bias, **self.extra_args)
+        if self.bypass_mode:
+            diff_weight = (
+                self.get_weight(self.shape)
+                * self.scalar
+                * self.multiplier
+            )
+            return self.org_forward(x) + self.op(x, diff_weight.view(self.shape), **self.extra_args)
+        else:
+            weight = (
+                self.org_module[0].weight.data.to(x.device, dtype=self.hada_w1_a.dtype)
+                + self.get_weight(self.shape)
+                * self.scalar
+                * self.multiplier
+            )
+            bias = None if self.org_module[0].bias is None else self.org_module[0].bias.data
+
+            if self.wd:
+                weight = self.apply_weight_decompose(weight)
+            return self.op(x, weight.view(self.shape), bias, **self.extra_args)
+
+
+if __name__ == "__main__":
+    base = nn.Linear(128, 128).cuda()
+    lokr = LohaModule("test", base, 1, 4, 1, weight_decompose=True).cuda()
+    print(lokr)
+    test_input = torch.randn(1, 128).cuda()
+    test_output = lokr(test_input)
+    print(test_output.shape)
+
+    base_4bit = LinearNF4(128, 128)
+    base_4bit.load_state_dict(base.state_dict())
+    base_4bit.cuda()
+    qlocon = LohaModule("test", base_4bit, 1, 4, 1, weight_decompose=False).cuda()
+    print(qlocon)
+    test_input = torch.randn(1, 128).cuda()
+    test_output = qlocon(test_input)
+    print(test_output.shape)
+
+    base = nn.Conv2d(128, 128, 3, 1, 1)
+    lokr = LohaModule("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)
+    print(test_output.shape)

{lycoris_lora-2.2.0.dev6 → lycoris_lora-2.2.0.dev8}/lycoris/modules/lokr.py

@@ -7,6 +7,11 @@ import torch.nn.functional as F
 
 from .base import ModuleCustomSD
 from ..logging import logger
+from ..utils.bnb import (
+    LinearNF4,
+    QuantLinears,
+    log_bypass
+)
 
 
 @cache
@@ -101,6 +106,7 @@ class LokrModule(ModuleCustomSD):
         rank_dropout_scale=False,
         weight_decompose=False,
         full_matrix=False,
+        bypass_mode=False,
         **kwargs,
     ):
         """if alpha == 0 or None, alpha is rank (no scaling)."""
@@ -118,6 +124,7 @@ class LokrModule(ModuleCustomSD):
             in_dim = org_module.in_channels
             k_size = org_module.kernel_size
             out_dim = org_module.out_channels
+            self.shape = (out_dim, in_dim, *k_size)
 
             in_m, in_n = factorization(in_dim, factor)
             out_l, out_k = factorization(out_dim, factor)
@@ -172,6 +179,7 @@ class LokrModule(ModuleCustomSD):
         else:  # Linear
             in_dim = org_module.in_features
             out_dim = org_module.out_features
+            self.shape = (out_dim, in_dim)
 
             in_m, in_n = factorization(in_dim, factor)
             out_l, out_k = factorization(out_dim, factor)
@@ -200,6 +208,13 @@ class LokrModule(ModuleCustomSD):
             self.op = F.linear
             self.extra_args = {}
 
+        if isinstance(org_module, QuantLinears):
+            if not bypass_mode:
+                log_bypass()
+            bypass_mode = True
+        self.bypass_mode = bypass_mode
+        assert not (bypass_mode and weight_decompose), "bypass_mode and dora_wd cannot be used together"
+
         self.wd = weight_decompose
         if self.wd:
             org_weight: nn.Parameter = org_module.weight
@@ -283,10 +298,10 @@ class LokrModule(ModuleCustomSD):
         self.org_module[0].forward = self.org_forward
 
     def merge_to(self, multiplier=1.0):
-        weight = self.get_weight(self.org_module[0].weight)
+        weight = self.get_weight(self.shape)
         self.org_module[0].weight.data.add_(weight * multiplier)
 
-    def get_weight(self, orig_weight=None):
+    def get_weight(self, shape):
         weight = make_kron(
             self.lokr_w1 if self.use_w1 else self.lokr_w1_a @ self.lokr_w1_b,
             (
@@ -300,8 +315,8 @@ class LokrModule(ModuleCustomSD):
             ),
             torch.tensor(self.scale),
         )
-        if orig_weight is not None:
-            weight = weight.reshape(orig_weight.shape)
+        if shape is not None:
+            weight = weight.reshape(shape)
         if self.training and self.rank_dropout:
             drop = (torch.rand(weight.size(0)) > self.rank_dropout).to(weight.dtype)
             drop = drop.view(-1, *[1] * len(weight.shape[1:])).to(weight.device)
@@ -373,23 +388,47 @@ class LokrModule(ModuleCustomSD):
                         else self.org_module[0].bias.data
                     ),
                 )
-        weight = (
-            self.org_module[0].weight.data.to(x.device, dtype=self.lokr_w1.dtype)
-            + self.get_weight(self.org_module[0].weight.data)
-            * self.scalar
-            * self.multiplier
-        )
-        bias = None if self.org_module[0].bias is None else self.org_module[0].bias.data
+        if self.bypass_mode:
+            diff_weight = (
+                self.get_weight(self.shape)
+                * self.scalar
+                * self.multiplier
+            )
+            return self.org_forward(x) + self.op(x, diff_weight.view(self.shape), **self.extra_args)
+        else:
+            weight = (
+                self.org_module[0].weight.data.to(x.device, dtype=next(self.parameters()).dtype)
+                + self.get_weight(self.shape)
+                * self.scalar
+                * self.multiplier
+            )
+            bias = None if self.org_module[0].bias is None else self.org_module[0].bias.data
 
-        if self.wd:
-            weight = self.apply_weight_decompose(weight)
-        return self.op(x, weight.view(self.shape), bias, **self.extra_args)
+            if self.wd:
+                weight = self.apply_weight_decompose(weight)
+            return self.op(x, weight.view(self.shape), bias, **self.extra_args)
 
 
 if __name__ == "__main__":
-    base = nn.Linear(128, 128)
-    lokr = LokrModule("test", base, 1, 4, 1, weight_decompose=True)
+    base = nn.Linear(128, 128).cuda()
+    lokr = LokrModule("test", base, 1, 4, 1, weight_decompose=True).cuda()
+    print(lokr)
+    test_input = torch.randn(1, 128).cuda()
+    test_output = lokr(test_input)
+    print(test_output.shape)
+
+    base_4bit = LinearNF4(128, 128)
+    base_4bit.load_state_dict(base.state_dict())
+    base_4bit.cuda()
+    qlocon = LokrModule("test", base_4bit, 1, 4, 1, weight_decompose=False).cuda()
+    print(qlocon)
+    test_input = torch.randn(1, 128).cuda()
+    test_output = qlocon(test_input)
+    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)
+    test_input = torch.randn(1, 128, 16, 16)
     test_output = lokr(test_input)
     print(test_output.shape)

lycoris_lora-2.2.0.dev8/lycoris/utils/bnb.py

@@ -0,0 +1,30 @@
+from functools import cache
+
+try:
+    from bitsandbytes.nn import (
+        LinearNF4, Linear8bitLt, LinearFP4
+    )
+except:
+    import torch.nn as nn
+    class LinearNF4(nn.Linear):
+        pass
+    class Linear8bitLt(nn.Linear):
+        pass
+    class LinearFP4(nn.Linear):
+        pass
+
+from ..logging import logger
+
+
+QuantLinears = (
+    Linear8bitLt,
+    LinearFP4,
+    LinearNF4
+)
+
+
+@cache
+def log_bypass():
+    return logger.warning(
+        "Using bnb with LyCORIS will enable force-bypass mode."
+    )

{lycoris_lora-2.2.0.dev6 → lycoris_lora-2.2.0.dev8/lycoris_lora.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: lycoris_lora
-Version: 2.2.0.dev6
+Version: 2.2.0.dev8
 Summary: Lora beYond Conventional methods, Other Rank adaptation Implementations for Stable diffusion
 Home-page: https://github.com/KohakuBlueleaf/LyCORIS
 Author: Shih-Ying Yeh(KohakuBlueLeaf), Yu-Guan Hsieh, Zhidong Gao

{lycoris_lora-2.2.0.dev6 → lycoris_lora-2.2.0.dev8}/lycoris_lora.egg-info/SOURCES.txt

@@ -27,6 +27,7 @@ lycoris/modules/full.py
 lycoris/modules/glokr.py
 lycoris/modules/glora.py
 lycoris/modules/ia3.py
+lycoris/modules/lilora.py
 lycoris/modules/locon.py
 lycoris/modules/loha.py
 lycoris/modules/lokr.py
@@ -35,6 +36,7 @@ lycoris/modules/hypernet/__init__.py
 lycoris/modules/hypernet/generater.py
 lycoris/modules/hypernet/text_encoder.py
 lycoris/utils/__init__.py
+lycoris/utils/bnb.py
 lycoris/utils/logger.py
 lycoris/utils/preset.py
 lycoris/utils/xformers_utils.py

{lycoris_lora-2.2.0.dev6 → lycoris_lora-2.2.0.dev8}/setup.py

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