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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: lycoris_lora
-Version: 3.0.0.dev3
+Version: 3.0.0.dev5
 Summary: Lora beYond Conventional methods, Other Rank adaptation Implementations for Stable diffusion
 Home-page: https://github.com/KohakuBlueleaf/LyCORIS
 Author: Shih-Ying Yeh(KohakuBlueLeaf), Yu-Guan Hsieh, Zhidong Gao
@@ -10,3 +10,4 @@ License-File: LICENSE.md
 Requires-Dist: torch
 Requires-Dist: einops
 Requires-Dist: toml
+Requires-Dist: tqdm

{lycoris_lora-3.0.0.dev3 → lycoris_lora-3.0.0.dev5}/lycoris/functional/general.py

@@ -3,6 +3,9 @@ import torch.nn as nn
 import torch.nn.functional as F
 
 
+FUNC_LIST = [None, None, F.linear, F.conv1d, F.conv2d, F.conv3d]
+
+
 def rebuild_tucker(t, wa, wb):
     rebuild2 = torch.einsum("i j ..., i p, j r -> p r ...", t, wa, wb)
     return rebuild2

{lycoris_lora-3.0.0.dev3 → lycoris_lora-3.0.0.dev5}/lycoris/functional/locon.py

@@ -4,21 +4,18 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
-from .general import rebuild_tucker
+from .general import rebuild_tucker, FUNC_LIST
 
 
 def lora_weight_gen(org_weight, rank, tucker=True):
     """### lora_weight_gen
 
     Args:
-        out_dim (int): output dimension
-        in_dim (int): input dimension
+        org_weight (torch.Tensor): the weight tensor
         rank (int): low rank
 
     Returns:
-        torch.Tensor: down proj weight
-        torch.Tensor: up proj weight
-        torch.Tensor(optional): mid proj weight
+        torch.Tensor: down, up[, mid]
     """
     out_dim, in_dim, *k = org_weight.shape
     if k and tucker:
@@ -45,14 +42,12 @@ def lora_diff_weight(d, u, m=None, gamma=1.0):
     Args:
         d (torch.Tensor): weight of down proj linear/conv layer
         u (torch.Tensor): weight of up proj linear/conv layer
-        m (torch.Tensor, optional): weight of mid proj linear/conv layer, \
-            for tucker deomposition
+        m (torch.Tensor, optional): middle weight of tucker decomposition
         gamma (float, optional): scale factor, normally alpha/rank here
 
     Returns:
         torch.Tensor: ΔW
     """
-    assert d.size(0) == u.size(1)
     R, I, *k = d.shape
     O, R, *_ = u.shape
     u = u * gamma
@@ -67,9 +62,6 @@ def lora_diff_weight(d, u, m=None, gamma=1.0):
     return result.reshape(O, I, *k)
 
 
-FUNC_LIST = [None, None, F.linear, F.conv1d, F.conv2d, F.conv3d]
-
-
 def lora_bypass_forward_diff(x, d, u, m=None, gamma=1.0, extra_args={}):
     """### lora_bypass_forward_diff
 
@@ -77,9 +69,10 @@ def lora_bypass_forward_diff(x, d, u, m=None, gamma=1.0, extra_args={}):
         x (torch.Tensor): input tensor
         d (torch.Tensor): weight of down proj linear/conv layer
         u (torch.Tensor): weight of up proj linear/conv layer
-        m (torch.Tensor, optional): weight of mid proj linear/conv layer, \
-            for tucker deomposition
+        m (torch.Tensor, optional): middle weight of tucker decomposition
         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

{lycoris_lora-3.0.0.dev3 → lycoris_lora-3.0.0.dev5}/lycoris/functional/loha.py

@@ -1,7 +1,11 @@
+import math
+
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
+from .general import FUNC_LIST
+
 
 class HadaWeight(torch.autograd.Function):
     @staticmethod
@@ -79,6 +83,43 @@ def make_weight_tucker(t1, w1d, w1u, t2, w2d, w2u, scale):
     return HadaWeightTucker.apply(t1, w1d, w1u, t2, w2d, w2u, scale)
 
 
+def loha_weight_gen(org_weight, rank, tucker=True):
+    """### loha_weight_gen
+
+    Args:
+        org_weight (torch.Tensor): the weight tensor
+        rank (int): low rank
+
+    Returns:
+        torch.Tensor: w1d, w2d, w1u, w2u[, t1, t2]
+    """
+    out_dim, in_dim, *k = org_weight.shape
+    if k and tucker:
+        w1d = torch.empty(rank, in_dim)
+        w1u = torch.empty(rank, out_dim)
+        t1 = torch.empty(rank, rank, *k)
+        w2d = torch.empty(rank, in_dim)
+        w2u = torch.empty(rank, out_dim)
+        t2 = torch.empty(rank, rank, *k)
+        nn.init.normal_(w1d, std=1)
+        nn.init.constant_(w1u, 0)
+        nn.init.normal_(w2d, std=1)
+        nn.init.normal_(w2u, std=0.1)
+        nn.init.normal_(t1, std=0.1)
+        nn.init.normal_(t2, std=0.1)
+        return w1d, w2d, w1u, w2u, t1, t2
+    else:
+        w1d = torch.empty(rank, in_dim)
+        w1u = torch.empty(rank, out_dim)
+        w2d = torch.empty(rank, in_dim)
+        w2u = torch.empty(rank, out_dim)
+        nn.init.normal_(w1d, std=1)
+        nn.init.constant_(w1u, 0)
+        nn.init.normal_(w2d, std=1)
+        nn.init.normal_(w2u, std=0.1)
+        return w1d, w2d, w1u, w2u
+
+
 def loha_diff_weight(w1d, w1u, w2d, w2u, t1=None, t2=None, gamma=1.0):
     """### loha_diff_weight
 
@@ -87,23 +128,18 @@ def loha_diff_weight(w1d, w1u, w2d, w2u, t1=None, t2=None, gamma=1.0):
     Args:
         w1d, w2d (torch.Tensor): weight of down proj linear/conv layer
         w1u, w2u (torch.Tensor): weight of up proj linear/conv layer
+        t1, t2 (torch.Tensor, optional): middle weight of tucker decomposition
         gamma (float, optional): scale factor, normally alpha/rank here
 
     Returns:
         torch.Tensor: ΔW
     """
     if t1 is not None and t2 is not None:
-        assert w1d.size(0) == t1.size(1)
-        assert w1u.size(0) == t1.size(0)
-        assert w2d.size(0) == t2.size(1)
-        assert w2u.size(0) == t2.size(0)
         R, I = w1d.shape
         R, O = w1u.shape
         R, R, *k = t1.shape
         result = make_weight_tucker(t1, w1d, w1u, t2, w2d, w2u, gamma)
     else:
-        assert w1d.size(0) == w1u.size(1)
-        assert w2d.size(0) == w2u.size(1)
         R, I, *k = w1d.shape
         O, R, *_ = w1u.shape
         w1d = w1d.reshape(w1d.size(0), -1)
@@ -116,9 +152,6 @@ def loha_diff_weight(w1d, w1u, w2d, w2u, t1=None, t2=None, gamma=1.0):
     return result
 
 
-FUNC_LIST = [None, None, F.linear, F.conv1d, F.conv2d, F.conv3d]
-
-
 def loha_bypass_forward_diff(
     x, w1d, w1u, w2d, w2u, t1=None, t2=None, gamma=1.0, extra_args={}
 ):
@@ -129,6 +162,8 @@ def loha_bypass_forward_diff(
         w1d, w2d (torch.Tensor): weight of up proj linear/conv layer
         w1u, w2u (torch.Tensor): weight of down proj linear/conv layer
         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
@@ -139,15 +174,16 @@ def loha_bypass_forward_diff(
 
 if __name__ == "__main__":
     w = torch.randn(128, 128, 3, 3, 3)
-    a = torch.randn(16, 128) * 0.01
-    b = torch.randn(16, 128) * 0.01
-    m = torch.randn(16, 16, 3, 3, 3) * 0.01
+    w1d, w2d, w1u, w2u, t1, t2 = loha_weight_gen(
+        w, 4, tucker=True
+    )
+    w1u += 0.01
     extra_args = {"padding": 1}
 
     x = torch.randn(1, 128, 8, 8, 8)
     y = FUNC_LIST[w.dim()](x, w, **extra_args)
-    diff_w = loha_diff_weight(a, b, a, b, m, m, 1)
-    diff_y = loha_bypass_forward_diff(x, a, b, a, b, m, m, 1, extra_args)
+    diff_w = loha_diff_weight(w1d, w2d, w1u, w2u, t1, t2, 1)
+    diff_y = loha_bypass_forward_diff(x, w1d, w2d, w1u, w2u, t1, t2, 1, extra_args)
 
     print(F.mse_loss(y, y + diff_y))
     print(F.mse_loss(w, w + diff_w))

{lycoris_lora-3.0.0.dev3 → lycoris_lora-3.0.0.dev5}/lycoris/kohya.py

@@ -1,16 +1,10 @@
-# network module for kohya
-# reference:
-# https://github.com/microsoft/LoRA/blob/main/loralib/layers.py
-# https://github.com/cloneofsimo/lora/blob/master/lora_diffusion/lora.py
-# https://github.com/kohya-ss/sd-scripts/blob/main/networks/lora.py
-
 import os
 import re
 from typing import List
 
 import torch
 
-from library.train_util import *
+from .utils import precalculate_safetensors_hashes
 from .wrapper import LycorisNetwork
 from .modules.locon import LoConModule
 from .modules.loha import LohaModule
@@ -150,10 +144,6 @@ def create_network(
             unbalanced_factorization=unbalanced_factorization,
         )
 
-    if algo == "dylora":
-        # dylora didn't support scale weight norm yet
-        delattr(type(network), "apply_max_norm_regularization")
-
     return network
 
 

{lycoris_lora-3.0.0.dev3 → lycoris_lora-3.0.0.dev5}/lycoris/modules/__init__.py

@@ -1,13 +1,19 @@
 import torch
 import torch.nn as nn
 
-from .full import FullModule
-from .norms import NormModule
 from .locon import LoConModule
 from .loha import LohaModule
-from .lokr import LokrModule, factorization
+from .lokr import LokrModule
+from .full import FullModule
+from .norms import NormModule
+from .diag_oft import DiagOFTModule
+from .boft import ButterflyOFTModule
+from .glora import GLoRAModule
+from .dylora import DyLoraModule
 from .ia3 import IA3Module
 
+from ..functional.general import factorization
+
 
 @torch.no_grad()
 def make_module(lyco_type, params, lora_name, orig_module):
@@ -51,12 +57,14 @@ def make_module(lyco_type, params, lora_name, orig_module):
             module.dora_scale.copy_(dora_scale)
     elif lyco_type == "kron":
         w1, w1a, w1b, w2, w2a, w2b, _, t2, alpha, dora_scale = params
+        full_matrix=False
         if w1a is not None:
             lora_dim = w1a.size(1)
         elif w2a is not None:
             lora_dim = w2a.size(1)
         else:
-            lora_dim = 10000000000000
+            full_matrix = True
+            lora_dim = 1
 
         if w1 is None:
             out_dim = w1a.size(0)
@@ -91,6 +99,7 @@ def make_module(lyco_type, params, lora_name, orig_module):
             decompose_both=w1 is None and w2 is None,
             factor=factor,
             weight_decompose=dora_scale is not None,
+            full_matrix=full_matrix,
         )
         if w1 is not None:
             module.lokr_w1.copy_(w1)

{lycoris_lora-3.0.0.dev3 → lycoris_lora-3.0.0.dev5}/lycoris/modules/base.py

@@ -256,7 +256,7 @@ class LycorisBaseModule(ModuleCustomSD):
 
     @torch.no_grad()
     def apply_max_norm(self, max_norm, device=None):
-        raise NotImplementedError
+        return None, None
 
     def bypass_forward_diff(self, x, scale=1):
         raise NotImplementedError

lycoris_lora-3.0.0.dev5/lycoris/modules/dylora.py

@@ -0,0 +1,198 @@
+import math
+import random
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .base import LycorisBaseModule
+from ..utils import product
+from ..utils.bnb import LinearNF4
+
+
+class DyLoraModule(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,
+        block_size=4,
+        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.")
+        assert lora_dim % block_size == 0, "lora_dim must be a multiple of block_size"
+        self.block_count = lora_dim // block_size
+        self.block_size = block_size
+
+        shape = (
+            self.shape[0],
+            product(self.shape[1:]),
+        )
+
+        self.lora_dim = lora_dim
+        self.up_list = nn.ParameterList(
+            [torch.empty(shape[0], self.block_size) for i in range(self.block_count)]
+        )
+        self.down_list = nn.ParameterList(
+            [torch.empty(self.block_size, shape[1]) for i in range(self.block_count)]
+        )
+
+        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))  # 定数として扱える
+
+        # Need more experiences on init method
+        for v in self.down_list:
+            torch.nn.init.kaiming_uniform_(v, a=math.sqrt(5))
+        for v in self.up_list:
+            torch.nn.init.zeros_(v)
+
+    def custom_state_dict(self):
+        destination = {}
+        destination["alpha"] = self.alpha
+        destination["lora_up.weight"] = nn.Parameter(
+            torch.concat(list(self.up_list), dim=1)
+        )
+        destination["lora_down.weight"] = nn.Parameter(
+            torch.concat(list(self.down_list)).reshape(
+                self.lora_dim, -1, *self.shape[2:]
+            )
+        )
+        return destination
+
+    def get_weight(self, rank):
+        b = math.ceil(rank / self.block_size)
+        down = torch.concat(
+            list(i.data for i in self.down_list[:b]) + list(self.down_list[b : (b + 1)])
+        )
+        up = torch.concat(
+            list(i.data for i in self.up_list[:b]) + list(self.up_list[b : (b + 1)]),
+            dim=1,
+        )
+        return down, up, self.alpha / (b + 1)
+
+    def get_random_rank_weight(self):
+        b = random.randint(0, self.block_count - 1)
+        return self.get_weight(b * self.block_size)
+
+    def get_diff_weight(self, multiplier=1, shape=None, device=None, rank=None):
+        if rank is None:
+            down, up, scale = self.get_random_rank_weight()
+        else:
+            down, up, scale = self.get_random_rank_weight()
+        w = up @ (down * (scale * multiplier))
+        if device is not None:
+            w = w.to(device)
+        if shape is not None:
+            w = w.view(shape)
+        else:
+            w = w.view(self.shape)
+        return w, None
+
+    def get_merged_weight(self, multiplier=1, shape=None, device=None, rank=None):
+        diff, _ = self.get_diff_weight(multiplier, shape, device, rank)
+        return diff + self.org_weight, None
+
+    def bypass_forward_diff(self, x, scale=1, rank=None):
+        if rank is None:
+            down, up, gamma = self.get_random_rank_weight()
+        else:
+            down, up, gamma = self.get_random_rank_weight()
+        down = down.view(self.lora_dim, -1, *self.shape[2:])
+        up = up.view(-1, self.lora_dim, *(1 for _ in self.shape[2:]))
+        scale = scale * gamma
+        return self.op(self.op(x, down, **self.kw_dict), up)
+
+    def bypass_forward(self, x, scale=1, rank=None):
+        return self.org_forward(x) + self.bypass_forward_diff(x, scale, rank)
+
+    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 = 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.dev3 → lycoris_lora-3.0.0.dev5}/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.dev5/lycoris/modules/ia3.py

@@ -0,0 +1,173 @@
+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.")
+
+        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.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.dev5}/lycoris/modules/locon.py

@@ -154,12 +154,12 @@ class LoConModule(LycorisBaseModule):
             )
 
     def make_weight(self, device=None):
-        wa = self.lora_up.weight.to(device)
-        wb = self.lora_down.weight.to(device)
+        wa = self.lora_up.weight
+        wb = self.lora_down.weight
         if self.tucker:
-            t = self.lora_mid.weight.to(device)
-            wa = wa.reshape(wa.size(0), -1).transpose(0, 1)
-            wb = wb.reshape(wb.size(0), -1)
+            t = self.lora_mid.weight
+            wa = wa.view(wa.size(0), -1).transpose(0, 1)
+            wb = wb.view(wb.size(0), -1)
             weight = rebuild_tucker(t, wa, wb)
         else:
             weight = wa.view(wa.size(0), -1) @ wb.view(wb.size(0), -1)

{lycoris_lora-3.0.0.dev3 → lycoris_lora-3.0.0.dev5}/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.dev5}/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.dev5}/lycoris/utils/__init__.py

@@ -1,4 +1,6 @@
 import re
+import hashlib
+from io import BytesIO
 from typing import Dict, Tuple, Union
 
 import numpy as np
@@ -6,10 +8,37 @@ import numpy as np
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-
 import torch.linalg as linalg
 
+import safetensors.torch
+
 from tqdm import tqdm
+from .general import *
+
+
+def load_bytes_in_safetensors(tensors):
+    bytes = safetensors.torch.save(tensors)
+    b = BytesIO(bytes)
+
+    b.seek(0)
+    header = b.read(8)
+    n = int.from_bytes(header, "little")
+
+    offset = n + 8
+    b.seek(offset)
+
+    return b.read()
+
+
+def precalculate_safetensors_hashes(state_dict):
+    # calculate each tensor one by one to reduce memory usage
+    hash_sha256 = hashlib.sha256()
+    for tensor in state_dict.values():
+        single_tensor_sd = {"tensor": tensor}
+        bytes_for_tensor = load_bytes_in_safetensors(single_tensor_sd)
+        hash_sha256.update(bytes_for_tensor)
+
+    return f"0x{hash_sha256.hexdigest()}"
 
 
 def str_bool(val):

lycoris_lora-3.0.0.dev5/lycoris/utils/general.py

@@ -0,0 +1,5 @@
+def product(xs: list[int | float]):
+    res = 1
+    for x in xs:
+        res *= x
+    return res

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

@@ -468,11 +468,12 @@ class LycorisNetwork(torch.nn.Module):
     def apply_max_norm_regularization(self, max_norm_value, device):
         key_scaled = 0
         norms = []
-        for model in self.loras:
-            if hasattr(model, "apply_max_norm"):
-                scaled, norm = model.apply_max_norm(max_norm_value, device)
-                norms.append(norm)
-                key_scaled += scaled
+        for module in self.loras:
+            scaled, norm = module.apply_max_norm(max_norm_value, device)
+            if scaled is None:
+                continue
+            norms.append(norm)
+            key_scaled += scaled
 
         if key_scaled == 0:
             return key_scaled, 0, 0

{lycoris_lora-3.0.0.dev3 → lycoris_lora-3.0.0.dev5/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.dev5
 Summary: Lora beYond Conventional methods, Other Rank adaptation Implementations for Stable diffusion
 Home-page: https://github.com/KohakuBlueleaf/LyCORIS
 Author: Shih-Ying Yeh(KohakuBlueLeaf), Yu-Guan Hsieh, Zhidong Gao
@@ -10,3 +10,4 @@ License-File: LICENSE.md
 Requires-Dist: torch
 Requires-Dist: einops
 Requires-Dist: toml
+Requires-Dist: tqdm

{lycoris_lora-3.0.0.dev3 → lycoris_lora-3.0.0.dev5}/lycoris_lora.egg-info/SOURCES.txt

@@ -24,6 +24,7 @@ lycoris/modules/lokr.py
 lycoris/modules/norms.py
 lycoris/utils/__init__.py
 lycoris/utils/bnb.py
+lycoris/utils/general.py
 lycoris/utils/logger.py
 lycoris/utils/preset.py
 lycoris/utils/xformers_utils.py

{lycoris_lora-3.0.0.dev3 → lycoris_lora-3.0.0.dev5}/lycoris_lora.egg-info/requires.txt

@@ -1,3 +1,4 @@
 torch
 einops
 toml
+tqdm

{lycoris_lora-3.0.0.dev3 → lycoris_lora-3.0.0.dev5}/setup.py

@@ -4,16 +4,12 @@ from setuptools import setup, find_packages
 setup(
     name="lycoris_lora",
     packages=find_packages(),
-    version="3.0.0.dev3",
+    version="3.0.0.dev5",
     url="https://github.com/KohakuBlueleaf/LyCORIS",
     description="Lora beYond Conventional methods, Other Rank adaptation Implementations for Stable diffusion",
     author="Shih-Ying Yeh(KohakuBlueLeaf), Yu-Guan Hsieh, Zhidong Gao",
     author_email="apolloyeh0123@gmail.com",
     zip_safe=False,
-    install_requires=[
-        "torch",
-        "einops",
-        "toml",
-    ],
+    install_requires=["torch", "einops", "toml", "tqdm"],
     python_requires=">=3.10",
 )

lycoris_lora-3.0.0.dev3/lycoris/modules/dylora.py

@@ -1,132 +0,0 @@
-import math
-import random
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-from .base import ModuleCustomSD
-
-
-class DyLoraModule(ModuleCustomSD):
-    """
-    Hadamard product Implementaion for Dynamic Low Rank adaptation
-    """
-
-    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,
-        block_size=1,
-        **kwargs,
-    ):
-        """if alpha == 0 or None, alpha is rank (no scaling)."""
-        super().__init__()
-        self.lora_name = lora_name
-        self.lora_dim = lora_dim
-        assert lora_dim % block_size == 0, "lora_dim must be a multiple of block_size"
-        self.block_count = lora_dim // block_size
-        self.block_size = block_size
-
-        self.shape = org_module.weight.shape
-        if org_module.__class__.__name__ == "Conv2d":
-            in_dim = org_module.in_channels
-            k_size = org_module.kernel_size
-            out_dim = org_module.out_channels
-            shape = (out_dim, in_dim * k_size[0] * k_size[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
-            shape = (out_dim, in_dim)
-            self.op = F.linear
-            self.extra_args = {}
-
-        self.lora_dim = lora_dim
-        self.up_list = nn.ParameterList(
-            [torch.empty(shape[0], 1) for i in range(lora_dim)]
-        )
-
-        self.down_list = nn.ParameterList(
-            [torch.empty(1, shape[1]) for i in range(lora_dim)]
-        )
-
-        self.index = 0
-
-        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))  # 定数として扱える
-
-        # Need more experiences on init method
-
-        for v in self.down_list:
-            torch.nn.init.kaiming_uniform_(v, a=math.sqrt(5))
-        for v in self.up_list:
-            torch.nn.init.zeros_(v)
-
-        self.multiplier = multiplier
-        self.org_module = [org_module]  # remove in applying
-        self.org_forward = self.org_module[0].forward
-        self.grad_ckpt = False
-        self.state_dict()
-
-    def custom_state_dict(self):
-        destination = {}
-        destination["alpha"] = self.alpha
-        destination["lora_up.weight"] = nn.Parameter(
-            torch.concat(list(self.up_list), dim=1)
-        )
-        destination["lora_down.weight"] = nn.Parameter(
-            torch.concat(list(self.down_list))
-        )
-        return destination
-
-    def apply_to(self):
-        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):
-        up = torch.concat(list(self.up_list), dim=1)
-        down = torch.concat(list(self.down_list))
-        diff_weight = up @ down * self.scale * multiplier / up.size(1)
-        self.org_module[0].weight.data.add_(diff_weight)
-
-    @torch.enable_grad()
-    def forward(self, x):
-        b = random.randint(0, self.block_count - 1)
-
-        down = torch.concat(
-            list(i.data for i in self.down_list[: b * self.block_size])
-            + list(self.down_list[b * self.block_size : (b + 1) * self.block_size])
-        )
-        up = torch.concat(
-            list(i.data for i in self.up_list[: b * self.block_size])
-            + list(self.up_list[b * self.block_size : (b + 1) * self.block_size]),
-            dim=1,
-        )
-
-        bias = None if self.org_module[0].bias is None else self.org_module[0].bias.data
-        return self.op(
-            x,
-            self.org_module[0].weight
-            + (up @ down).view(self.shape) * self.alpha / (b + 1),
-            bias,
-            **self.extra_args,
-        )

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