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

{lycoris_lora-3.0.0.dev3/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.dev3
+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

{lycoris_lora-3.0.0.dev3 → 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
 
-        ax_mid = self.a1(x) * scale
-        bx_mid = self.b1(x) * scale
+    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
+
+    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,12 +171,37 @@ 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__":
     device = torch.device("cuda")

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.dev3 → lycoris_lora-3.0.0.dev4}/lycoris/modules/loha.py

@@ -265,7 +265,7 @@ 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__":

{lycoris_lora-3.0.0.dev3 → 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
@@ -396,7 +396,7 @@ 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__":

{lycoris_lora-3.0.0.dev3 → 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.dev3
+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

{lycoris_lora-3.0.0.dev3 → 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.dev3",
+    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",

lycoris_lora-3.0.0.dev3/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