lycoris-lora 3.0.0.dev13__tar.gz → 3.0.0.dev15__tar.gz

{lycoris_lora-3.0.0.dev13/lycoris_lora.egg-info → lycoris_lora-3.0.0.dev15}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: lycoris_lora
-Version: 3.0.0.dev13
+Version: 3.0.0.dev15
 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.dev15/lycoris/functional/__init__.py

@@ -0,0 +1,9 @@
+from .general import (
+    rebuild_tucker,
+    factorization,
+    power2factorization,
+    FUNC_LIST,
+    tucker_weight,
+    tucker_weight_from_conv,
+    apply_dora_scale,
+)

{lycoris_lora-3.0.0.dev13 → lycoris_lora-3.0.0.dev15}/lycoris/functional/boft.py

@@ -10,7 +10,7 @@ from .general import power2factorization, FUNC_LIST
 from .diag_oft import get_r
 
 
-def boft_weight_gen(org_weight, max_block_size, boft_m=-1, rescale=False):
+def weight_gen(org_weight, max_block_size, boft_m=-1, rescale=False):
     """### boft_weight_gen
 
     Args:
@@ -28,13 +28,13 @@ def boft_weight_gen(org_weight, max_block_size, boft_m=-1, rescale=False):
         boft_m = max_boft_m
     boft_m = min(boft_m, max_boft_m)
     oft_blocks = torch.zeros(boft_m, block_num, block_size, block_size)
-    if rescale:
+    if rescale is not None:
         return oft_blocks, torch.ones(out_dim, *[1] * len(rest))
     else:
         return oft_blocks, None
 
 
-def boft_diff_weight(org_weight, oft_blocks, rescale=None, constraint=None):
+def diff_weight(org_weight, oft_blocks, rescale=None, constraint=None):
     """### boft_diff_weight
 
     Args:
@@ -57,11 +57,14 @@ def boft_diff_weight(org_weight, oft_blocks, rescale=None, constraint=None):
         inp = rearrange(inp, "d b ... -> (d b) ...")
         inp = rearrange(inp, "(c k g) ... -> (c g k) ...", g=2, k=2**i * r_b)
 
-    return inp * rescale - org
+    if rescale is not None:
+        inp = inp * rescale
 
+    return inp - org
 
-def boft_bypass_forward_diff(
-    org_out, need_transpose, oft_blocks, rescale=None, constraint=None
+
+def bypass_forward_diff(
+    org_out, oft_blocks, rescale=None, constraint=None, need_transpose=False
 ):
     """### boft_bypass_forward_diff
 
@@ -75,7 +78,7 @@ def boft_bypass_forward_diff(
     r_b = b // 2
     I = torch.eye(b, device=oft_blocks.device)
     r = get_r(oft_blocks, I, constraint)
-    inp = org_out.to(dtype=r.dtype)
+    inp = org = org_out.to(dtype=r.dtype)
     if need_transpose:
         inp = org = inp.transpose(1, -1)
 
@@ -87,27 +90,10 @@ def boft_bypass_forward_diff(
         inp = rearrange(inp, "... d b -> ... (d b)")
         inp = rearrange(inp, "... (c k g) -> ... (c g k)", g=2, k=2**i * r_b)
 
-    inp = inp * rescale.transpose(0, -1) - org
+    if rescale is not None:
+        inp = inp * rescale.transpose(0, -1)
+
+    inp = inp - org
     if need_transpose:
         inp = inp.transpose(1, -1)
     return inp
-
-
-if __name__ == "__main__":
-    w = torch.randn(32, 32, 3, 3, 3)
-    blocks, rescale = boft_weight_gen(w, 4, -1, True)
-    blocks = blocks + torch.randn_like(blocks) * 0.01
-    extra_args = {"padding": 1}
-
-    x = torch.randn(1, 32, 8, 8, 8)
-    y = FUNC_LIST[x.dim()](x, w, **extra_args)
-    diff_w = boft_diff_weight(w, blocks, rescale, 0.1)
-    diff_y_rebuild = FUNC_LIST[x.dim()](x, diff_w, **extra_args)
-    diff_y = boft_bypass_forward_diff(y, w.dim() > 2, blocks, rescale, 0.1)
-
-    print(y.shape, diff_y.shape, diff_y_rebuild.shape)
-    print(w.shape, diff_w.shape)
-
-    print(F.mse_loss(y, y + diff_y))
-    print(F.mse_loss(w, w + diff_w))
-    print(F.mse_loss(diff_y, diff_y_rebuild))

{lycoris_lora-3.0.0.dev13 → lycoris_lora-3.0.0.dev15}/lycoris/functional/diag_oft.py

@@ -16,7 +16,7 @@ def get_r(oft_blocks, I=None, constraint=0):
     # for Q = -Q^T
     q = oft_blocks - oft_blocks.transpose(-1, -2)
     normed_q = q
-    if constraint > 0:
+    if constraint is not None and constraint > 0:
         q_norm = torch.norm(q) + 1e-8
         if q_norm > constraint:
             normed_q = q * constraint / q_norm
@@ -25,8 +25,8 @@ def get_r(oft_blocks, I=None, constraint=0):
     return r
 
 
-def daig_oft_weight_gen(org_weight, max_block_size, rescale=False):
-    """### daig_oft_weight_gen
+def weight_gen(org_weight, max_block_size, rescale=False):
+    """### weight_gen
 
     Args:
         org_weight (torch.Tensor): the weight tensor
@@ -45,8 +45,8 @@ def daig_oft_weight_gen(org_weight, max_block_size, rescale=False):
         return oft_blocks, None
 
 
-def daig_oft_diff_weight(org_weight, oft_blocks, rescale=None, constraint=None):
-    """### daig_oft_diff_weight
+def diff_weight(org_weight, oft_blocks, rescale=None, constraint=None):
+    """### diff_weight
 
     Args:
         TODO
@@ -67,16 +67,16 @@ def daig_oft_diff_weight(org_weight, oft_blocks, rescale=None, constraint=None):
         r - I,
         org_weight,
     ).view(-1, *shape)
-    if rescale is None:
+    if rescale is not None:
         weight = rescale * weight
         weight = weight + (rescale - 1) * org_weight
     return weight
 
 
-def daig_oft_bypass_forward_diff(
-    org_out, need_transpose, oft_blocks, rescale=None, constraint=None
+def bypass_forward_diff(
+    org_out, oft_blocks, rescale=None, constraint=None, need_transpose=False
 ):
-    """### daig_oft_bypass_forward_diff
+    """### bypass_forward_diff
 
     Args:
         TODO
@@ -89,6 +89,7 @@ def daig_oft_bypass_forward_diff(
     r = get_r(oft_blocks, I, constraint)
     if need_transpose:
         org_out = org_out.transpose(1, -1)
+    org_out = org_out.to(dtype=r.dtype)
     *shape, _ = org_out.shape
     oft_out = torch.einsum(
         "k n m, ... k n -> ... k m", r - I, org_out.view(*shape, block_num, block_size)
@@ -100,20 +101,3 @@ def daig_oft_bypass_forward_diff(
     if need_transpose:
         out = out.transpose(1, -1)
     return out
-
-
-if __name__ == "__main__":
-    w = torch.randn(32, 32, 3, 3, 3)
-    blocks, rescale = daig_oft_weight_gen(w, 4, True)
-    blocks = blocks + torch.randn_like(blocks) * 0.01
-    extra_args = {"padding": 1}
-
-    x = torch.randn(1, 32, 8, 8, 8)
-    y = FUNC_LIST[x.dim()](x, w, **extra_args)
-    diff_w = daig_oft_diff_weight(w, blocks, rescale, 0.1)
-    diff_y_rebuild = FUNC_LIST[x.dim()](x, diff_w, **extra_args)
-    diff_y = daig_oft_bypass_forward_diff(y, w.dim() > 2, blocks, rescale, 0.1)
-
-    print(F.mse_loss(y, y + diff_y))
-    print(F.mse_loss(w, w + diff_w))
-    print(F.mse_loss(diff_y, diff_y_rebuild))

{lycoris_lora-3.0.0.dev13 → lycoris_lora-3.0.0.dev15}/lycoris/functional/general.py

@@ -77,3 +77,30 @@ def power2factorization(dimension: int, factor: int = -1) -> tuple[int, int]:
             n = dimension // m
 
     return dimension // n, n
+
+
+def tucker_weight_from_conv(up, down, mid):
+    up = up.reshape(up.size(0), up.size(1))
+    down = down.reshape(down.size(0), down.size(1))
+    return torch.einsum("m n ..., i m, n j -> i j ...", mid, up, down)
+
+
+def tucker_weight(wa, wb, t):
+    temp = torch.einsum("i j ..., j r -> i r ...", t, wb)
+    return torch.einsum("i j ..., i r -> r j ...", temp, wa)
+
+
+def apply_dora_scale(org_weight, rebuild, dora_scale, scale):
+    dora_norm_dims = org_weight.dim() - 1
+    weight = org_weight + rebuild
+    weight = weight.to(dora_scale.dtype)
+    weight_norm = (
+        weight.transpose(0, 1)
+        .reshape(weight.shape[1], -1)
+        .norm(dim=1, keepdim=True)
+        .reshape(weight.shape[1], *[1] * dora_norm_dims)
+        .transpose(0, 1)
+    )
+    merged_scale1 = weight / weight_norm * dora_scale
+    diff_weight = merged_scale1 - org_weight
+    return org_weight + diff_weight * scale

{lycoris_lora-3.0.0.dev13 → lycoris_lora-3.0.0.dev15}/lycoris/functional/locon.py

@@ -7,8 +7,8 @@ import torch.nn.functional as F
 from .general import rebuild_tucker, FUNC_LIST
 
 
-def lora_weight_gen(org_weight, rank, tucker=True):
-    """### lora_weight_gen
+def weight_gen(org_weight, rank, tucker=True):
+    """### weight_gen
 
     Args:
         org_weight (torch.Tensor): the weight tensor
@@ -34,8 +34,8 @@ def lora_weight_gen(org_weight, rank, tucker=True):
         return down, up
 
 
-def lora_diff_weight(d, u, m=None, gamma=1.0):
-    """### lora_diff_weight
+def diff_weight(d, u, m=None, gamma=1.0):
+    """### diff_weight
 
     Get ΔW = BA, where BA is low rank decomposition
 
@@ -62,8 +62,8 @@ def lora_diff_weight(d, u, m=None, gamma=1.0):
     return result.reshape(O, I, *k)
 
 
-def lora_bypass_forward_diff(x, d, u, m=None, gamma=1.0, extra_args={}):
-    """### lora_bypass_forward_diff
+def bypass_forward_diff(x, d, u, m=None, gamma=1.0, extra_args={}):
+    """### bypass_forward_diff
 
     Args:
         x (torch.Tensor): input tensor
@@ -85,18 +85,3 @@ def lora_bypass_forward_diff(x, d, u, m=None, gamma=1.0, extra_args={}):
         down = FUNC_LIST[d.dim()](x, d, **extra_args)
         up = FUNC_LIST[d.dim()](down, u)
     return up * gamma
-
-
-if __name__ == "__main__":
-    w = torch.randn(32, 32, 3, 3, 3)
-    d, u, m = lora_weight_gen(w, 4)
-    u = u + 0.01
-    extra_args = {"padding": 1}
-
-    x = torch.randn(1, 32, 8, 8, 8)
-    y = FUNC_LIST[d.dim()](x, w, **extra_args)
-    diff_w = lora_diff_weight(d, u, m, 1)
-    diff_y = lora_bypass_forward_diff(x, d, u, m, 1, extra_args)
-
-    print(F.mse_loss(y, y + diff_y))
-    print(F.mse_loss(w, w + diff_w))

{lycoris_lora-3.0.0.dev13 → lycoris_lora-3.0.0.dev15}/lycoris/functional/loha.py

@@ -83,8 +83,8 @@ 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
+def weight_gen(org_weight, rank, tucker=True):
+    """### weight_gen
 
     Args:
         org_weight (torch.Tensor): the weight tensor
@@ -101,27 +101,23 @@ def loha_weight_gen(org_weight, rank, tucker=True):
         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)
+        w1u = torch.empty(out_dim, rank)
         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
+        w2u = torch.empty(out_dim, rank)
+        t1 = t2 = None
+    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, w1u, w2d, w2u, t1, t2
 
 
-def loha_diff_weight(w1d, w1u, w2d, w2u, t1=None, t2=None, gamma=1.0):
-    """### loha_diff_weight
+def diff_weight(w1d, w1u, w2d, w2u, t1=None, t2=None, gamma=1.0):
+    """### diff_weight
 
     Get ΔW = BA, where BA is low rank decomposition
 
@@ -152,10 +148,10 @@ def loha_diff_weight(w1d, w1u, w2d, w2u, t1=None, t2=None, gamma=1.0):
     return result
 
 
-def loha_bypass_forward_diff(
+def bypass_forward_diff(
     x, w1d, w1u, w2d, w2u, t1=None, t2=None, gamma=1.0, extra_args={}
 ):
-    """### loha_bypass_forward_diff
+    """### bypass_forward_diff
 
     Args:
         x (torch.Tensor): input tensor
@@ -168,20 +164,5 @@ def loha_bypass_forward_diff(
     Returns:
         torch.Tensor: output tensor
     """
-    diff_w = loha_diff_weight(w1d, w1u, w2d, w2u, t1, t2, gamma)
+    diff_w = diff_weight(w1d, w1u, w2d, w2u, t1, t2, gamma)
     return FUNC_LIST[w1d.dim() if t1 is None else t1.dim()](x, diff_w, **extra_args)
-
-
-if __name__ == "__main__":
-    w = torch.randn(128, 128, 3, 3, 3)
-    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(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.dev13 → lycoris_lora-3.0.0.dev15}/lycoris/functional/lokr.py

@@ -17,7 +17,7 @@ def make_kron(w1, w2, scale):
     return rebuild * scale
 
 
-def lokr_weight_gen(
+def weight_gen(
     org_weight,
     rank,
     tucker=True,
@@ -118,8 +118,8 @@ def lokr_weight_gen(
     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
+def diff_weight(w1, w1a, w1b, w2, w2a, w2b, t, gamma=1.0):
+    """### diff_weight
 
     Args:
         gamma (float, optional): scale factor, normally alpha/rank here
@@ -146,10 +146,8 @@ def lokr_diff_weight(w1, w1a, w1b, w2, w2a, w2b, t, gamma=1.0):
     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
+def bypass_forward_diff(h, w1, w1a, w1b, w2, w2a, w2b, t, gamma=1.0, extra_args={}):
+    """### bypass_forward_diff
 
     Args:
         gamma (float, optional): scale factor, normally alpha/rank here
@@ -163,7 +161,6 @@ def lokr_bypass_forward_diff(
     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
@@ -241,19 +238,3 @@ def lokr_bypass_forward_diff(
         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 = lokr_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.dev13 → lycoris_lora-3.0.0.dev15}/lycoris/modules/base.py

@@ -269,8 +269,9 @@ class LycorisBaseModule(ModuleCustomSD):
         weight, bias = self.get_merged_weight(
             multiplier, self.org_weight.shape, self.org_weight.device
         )
-        self.org_weight = weight
+        self.org_weight = weight.to(self.org_weight)
         if bias is not None:
+            bias = bias.to(self.org_weight)
             if self.org_module[0].bias is not None:
                 self.org_module[0].bias.data.copy_(bias)
             else:

{lycoris_lora-3.0.0.dev13 → lycoris_lora-3.0.0.dev15}/lycoris/modules/boft.py

@@ -197,7 +197,7 @@ class ButterflyOFTModule(LycorisBaseModule):
         r = self.get_r()
         inp = org = self.org_forward(x)
         if self.op in {F.conv2d, F.conv1d, F.conv3d}:
-            org_out = org_out.transpose(1, -1)
+            inp = inp.transpose(1, -1)
 
         for i in range(m):
             bi = r[i]  # b_num, b_size, b_size

{lycoris_lora-3.0.0.dev13 → lycoris_lora-3.0.0.dev15}/lycoris/modules/diag_oft.py

@@ -188,7 +188,7 @@ class DiagOFTModule(LycorisBaseModule):
         if diff:
             out = out - org_out
         out = oft_out.view(*shape, -1)
-        if self.rescale is not None:
+        if self.rescaled:
             out = self.rescale.transpose(-1, 0) * out
             out = out + (self.rescale.transpose(-1, 0) - 1) * org_out
         if self.op in {F.conv2d, F.conv1d, F.conv3d}:

{lycoris_lora-3.0.0.dev13 → lycoris_lora-3.0.0.dev15}/lycoris/modules/glora.py

@@ -5,7 +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 ..functional import tucker_weight_from_conv
 from ..utils.bnb import LinearNF4
 
 

{lycoris_lora-3.0.0.dev13 → lycoris_lora-3.0.0.dev15}/lycoris/modules/loha.py

@@ -4,7 +4,7 @@ import torch
 import torch.nn as nn
 
 from .base import LycorisBaseModule
-from ..functional.loha import loha_diff_weight
+from ..functional.loha import diff_weight as loha_diff_weight
 from ..utils.bnb import LinearNF4
 
 

{lycoris_lora-3.0.0.dev13 → lycoris_lora-3.0.0.dev15}/lycoris/utils/__init__.py

@@ -417,138 +417,10 @@ def convert_diffusers_name_to_compvis(key):
     return key
 
 
-def tucker_weight_from_conv(up, down, mid):
-    up = up.reshape(up.size(0), up.size(1))
-    down = down.reshape(down.size(0), down.size(1))
-    return torch.einsum("m n ..., i m, n j -> i j ...", mid, up, down)
-
-
-def tucker_weight(wa, wb, t):
-    temp = torch.einsum("i j ..., j r -> i r ...", t, wb)
-    return torch.einsum("i j ..., i r -> r j ...", temp, wa)
-
-
-def apply_dora_scale(org_weight, rebuild, dora_scale, scale):
-    dora_norm_dims = org_weight.dim() - 1
-    weight = org_weight + rebuild
-    weight = weight.to(dora_scale.dtype)
-    weight_norm = (
-        weight.transpose(0, 1)
-        .reshape(weight.shape[1], -1)
-        .norm(dim=1, keepdim=True)
-        .reshape(weight.shape[1], *[1] * dora_norm_dims)
-        .transpose(0, 1)
-    )
-    merged_scale1 = weight / weight_norm * dora_scale
-    diff_weight = merged_scale1 - org_weight
-    return org_weight + diff_weight * scale
-
-
-@torch.no_grad()
-def rebuild_weight(module_type, params, orig_weight, orig_bias, scale=1):
-    if orig_weight is None:
-        return None, None
-    merged = orig_weight
-    merged_bias = orig_bias
-    if module_type == "locon":
-        up, down, mid, alpha, dora_scale = params
-        if alpha is not None:
-            scale *= alpha / up.size(1)
-        if mid is not None:
-            rebuild = tucker_weight_from_conv(up, down, mid)
-        else:
-            rebuild = up.reshape(up.size(0), -1) @ down.reshape(down.size(0), -1)
-        rebuild = rebuild.reshape(orig_weight.shape)
-        if dora_scale is None:
-            merged = orig_weight + rebuild * scale
-        else:
-            merged = apply_dora_scale(orig_weight, rebuild, dora_scale, scale)
-        del up, down, mid, alpha, params, rebuild
-    elif module_type == "hada":
-        w1a, w1b, w2a, w2b, t1, t2, alpha, dora_scale = params
-        if alpha is not None:
-            scale *= alpha / w1b.size(0)
-        if t1 is not None:
-            rebuild1 = tucker_weight(w1a, w1b, t1)
-        else:
-            rebuild1 = w1a @ w1b
-        if t2 is not None:
-            rebuild2 = tucker_weight(w2a, w2b, t2)
-        else:
-            rebuild2 = w2a @ w2b
-        rebuild = (rebuild1 * rebuild2).reshape(orig_weight.shape)
-        if dora_scale is None:
-            merged = orig_weight + rebuild * scale
-        else:
-            merged = apply_dora_scale(orig_weight, rebuild, dora_scale, scale)
-        del w1a, w1b, w2a, w2b, t1, t2, alpha, params, rebuild, rebuild1, rebuild2
-    elif module_type == "ia3":
-        weight, on_input = params
-        if not on_input:
-            weight = weight.reshape(-1, 1)
-        merged = orig_weight + weight * orig_weight * scale
-        del weight, on_input, params
-    elif module_type == "kron":
-        w1, w1a, w1b, w2, w2a, w2b, t1, t2, alpha, dora_scale = params
-        if alpha is not None and (w1b is not None or w2b is not None):
-            scale *= alpha / (w1b.size(0) if w1b else w2b.size(0))
-        if w1a is not None and w1b is not None:
-            if t1 is not None:
-                w1 = tucker_weight(w1a, w1b, t1)
-            else:
-                w1 = w1a @ w1b
-        if w2a is not None and w2b is not None:
-            if t2 is not None:
-                w2 = tucker_weight(w2a, w2b, t2)
-            else:
-                w2 = w2a @ w2b
-        if len(w2.shape) == 4:
-            w1 = w1.unsqueeze(2).unsqueeze(2)
-        w2 = w2.contiguous()
-        rebuild = torch.kron(w1, w2).reshape(orig_weight.shape)
-        if dora_scale is None:
-            merged = orig_weight + rebuild * scale
-        else:
-            merged = apply_dora_scale(orig_weight, rebuild, dora_scale, scale)
-        del w1, w1a, w1b, w2, w2a, w2b, t1, t2, alpha, params, rebuild
-    elif module_type == "full":
-        rebuild = params[0].reshape(orig_weight.shape)
-        rebuild_b = (
-            params[1].reshape(orig_bias.shape) if orig_bias is not None else None
-        )
-        merged = orig_weight + rebuild * scale
-        if rebuild_b is not None:
-            merged_bias = orig_bias + rebuild_b * scale
-        del params, rebuild, rebuild_b
-    elif module_type == "norm":
-        rebuild = params[0].reshape(orig_weight.shape)
-        rebuild_b = params[1].reshape(orig_bias.shape)
-        merged = orig_weight + rebuild * scale
-        merged_bias = orig_bias + rebuild_b * scale
-    else:
-        return None, None
-
-    return merged, merged_bias
-
-
 @torch.no_grad()
 def merge(tes, unet, lyco_state_dict, scale: float = 1.0, device="cpu"):
-    UNET_TARGET_REPLACE_MODULE = [
-        "Linear",
-        "Conv2d",
-        "LayerNorm",
-        "GroupNorm",
-        "GroupNorm32",
-    ]
-    TEXT_ENCODER_TARGET_REPLACE_MODULE = [
-        "Embedding",
-        "Linear",
-        "Conv2d",
-        "LayerNorm",
-        "GroupNorm",
-        "GroupNorm32",
-        "Embedding",
-    ]
+    from ..modules import make_module, get_module
+
     LORA_PREFIX_UNET = "lora_unet"
     LORA_PREFIX_TEXT_ENCODER = "lora_te"
     merged = 0
@@ -557,29 +429,25 @@ def merge(tes, unet, lyco_state_dict, scale: float = 1.0, device="cpu"):
         prefix,
         root_module: torch.nn.Module,
         lyco_state_dict: Dict[str, torch.Tensor],
-        target_replace_modules,
     ):
         nonlocal merged
         for child_name, child_module in tqdm(
             list(root_module.named_modules()), desc=f"Merging {prefix}"
         ):
-            if child_module.__class__.__name__ in target_replace_modules:
-                lora_name = prefix + "." + child_name
-                lora_name = lora_name.replace(".", "_")
-
-                result, result_b = rebuild_weight(
-                    *get_module(lyco_state_dict, lora_name),
-                    getattr(child_module, "weight"),
-                    getattr(child_module, "bias", None),
-                    scale,
-                )
-                if result is not None:
-                    key_dict.pop(lora_name)
-                    merged += 1
-                    child_module.requires_grad_(False)
-                    child_module.weight.copy_(result)
-                if result_b is not None:
-                    child_module.bias.copy_(result_b)
+            lora_name = prefix + "." + child_name
+            lora_name = lora_name.replace(".", "_")
+
+            lyco_type, params = get_module(lyco_state_dict, lora_name)
+            if lyco_type is None:
+                continue
+            module = make_module(lyco_type, params, lora_name, child_module)
+            if module is None:
+                continue
+            module.to(device)
+            module.merge_to(scale)
+            key_dict.pop(convert_diffusers_name_to_compvis(lora_name), None)
+            key_dict.pop(lora_name, None)
+            merged += 1
 
     key_dict = {}
     for k, v in tqdm(list(lyco_state_dict.items()), desc="Converting Dtype and Device"):
@@ -592,12 +460,7 @@ def merge(tes, unet, lyco_state_dict, scale: float = 1.0, device="cpu"):
             k = f"{convert_key}.{weight_key}"
         else:
             key_dict[module] = key_dict.get(module, []) + [k]
-        if device == "cpu":
-            lyco_state_dict[k] = v.float().cpu()
-        else:
-            lyco_state_dict[k] = v.to(
-                device, dtype=tes[0].parameters().__next__().dtype
-            )
+        lyco_state_dict[k] = v.float().cpu()
 
     for idx, te in enumerate(tes):
         if len(tes) > 1:
@@ -606,15 +469,15 @@ def merge(tes, unet, lyco_state_dict, scale: float = 1.0, device="cpu"):
             prefix = LORA_PREFIX_TEXT_ENCODER
         merge_state_dict(
             prefix,
-            te.to(device),
+            te,
             lyco_state_dict,
-            TEXT_ENCODER_TARGET_REPLACE_MODULE,
         )
+        torch.cuda.empty_cache()
     merge_state_dict(
         LORA_PREFIX_UNET,
-        unet.to(device),
+        unet,
         lyco_state_dict,
-        UNET_TARGET_REPLACE_MODULE,
     )
+    torch.cuda.empty_cache()
     print(f"Unused state dict key: {key_dict}")
     print(f"{merged} Modules been merged")

{lycoris_lora-3.0.0.dev13 → lycoris_lora-3.0.0.dev15/lycoris_lora.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: lycoris_lora
-Version: 3.0.0.dev13
+Version: 3.0.0.dev15
 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.dev13 → lycoris_lora-3.0.0.dev15}/setup.py

@@ -4,7 +4,7 @@ from setuptools import setup, find_packages
 setup(
     name="lycoris_lora",
     packages=find_packages(),
-    version="3.0.0.dev13",
+    version="3.0.0.dev15",
     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.dev15/test/functional.py

@@ -0,0 +1,87 @@
+import unittest
+from itertools import product
+from parameterized import parameterized
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from lycoris.functional import locon, loha, lokr, diag_oft, boft
+
+
+EPS_DTYPE = {
+    torch.float32: 5e-6,
+    torch.float16: 5e-5,
+    torch.bfloat16: 5e-4,
+}
+
+
+modules = [locon, loha, lokr, diag_oft, boft]
+base_module_and_input_adn_weight = [
+    lambda dim: (F.linear, torch.randn(dim, dim), torch.randn(1, dim)),
+    lambda dim: (F.conv1d, torch.randn(dim, dim, 3), torch.randn(1, dim, 16)),
+    lambda dim: (F.conv2d, torch.randn(dim, dim, 3, 3), torch.randn(1, dim, 16, 16)),
+    lambda dim: (
+        F.conv3d,
+        torch.randn(dim, dim, 3, 3, 3),
+        torch.randn(1, dim, 16, 16, 16),
+    ),
+]
+device_and_dtype = [
+    (torch.device("cpu"), torch.float32),
+]
+
+if torch.cuda.is_available():
+    device_and_dtype.append((torch.device("cuda"), torch.float32))
+    device_and_dtype.append((torch.device("cuda"), torch.float16))
+    device_and_dtype.append((torch.device("cuda"), torch.bfloat16))
+
+if torch.backends.mps.is_available():
+    device_and_dtype.append((torch.device("mps"), torch.float32))
+
+
+patch_forward_param_list = list(
+    product(
+        modules,
+        base_module_and_input_adn_weight,
+        device_and_dtype,
+    )
+)
+
+
+class LycorisFunctionalTests(unittest.TestCase):
+    @parameterized.expand(patch_forward_param_list)
+    def test_lycoris_functional(self, module, base, device_dtype):
+        func, test_weight, test_input = base(16)
+        device, dtype = device_dtype
+        print(
+            f"{module.__name__: <27}",
+            f"{func.__name__: <7}",
+            f"device={str(device): <5}",
+            f"dtype={str(dtype): <15}",
+            sep="||",
+        )
+
+        w = test_weight.to(device, dtype)
+        x = test_input.to(device, dtype)
+        y = func(x, w)
+
+        params = list(module.weight_gen(w, 4))
+        for idx, param in enumerate(params):
+            if param is not None:
+                param = param.to(device, dtype)
+                params[idx] = param + torch.randn_like(param) * 0.01
+
+        if module in {boft, diag_oft}:
+            diff_w = module.diff_weight(w, *params)
+            diff_y = module.bypass_forward_diff(y, *params, need_transpose=w.ndim > 2)
+        else:
+            diff_w = module.diff_weight(*params)
+            diff_y = module.bypass_forward_diff(x, *params)
+
+        diff_y_from_diff_w = func(x, diff_w.to(x))
+        self.assertTrue(
+            F.mse_loss(diff_y, diff_y_from_diff_w).item() < EPS_DTYPE[dtype],
+            f"Error: {module.__name__} {base.__name__} {device} {dtype} ||"
+            f"{F.mse_loss(diff_y, diff_y_from_diff_w).item()}",
+        )

lycoris_lora-3.0.0.dev13/lycoris/functional/__init__.py

@@ -1,20 +0,0 @@
-from .general import (
-    rebuild_tucker,
-    factorization,
-    power2factorization,
-)
-from .locon import (
-    lora_weight_gen,
-    lora_diff_weight,
-    lora_bypass_forward_diff,
-)
-from .loha import (
-    loha_weight_gen,
-    loha_diff_weight,
-    loha_bypass_forward_diff,
-)
-from .lokr import (
-    lokr_weight_gen,
-    lokr_diff_weight,
-    lokr_bypass_forward_diff,
-)

lycoris_lora-3.0.0.dev13/test/functional.py

@@ -1,71 +0,0 @@
-import unittest
-from itertools import product
-from parameterized import parameterized
-
-import torch
-import torch.nn as nn
-
-from lycoris.modules import (
-    LycorisBaseModule,
-    LoConModule,
-    LohaModule,
-    LokrModule,
-    FullModule,
-    DiagOFTModule,
-    ButterflyOFTModule,
-    GLoRAModule,
-    DyLoraModule,
-    IA3Module,
-)
-
-
-modules: list[LycorisBaseModule] = [
-    LoConModule,
-    LohaModule,
-    LokrModule,
-    FullModule,
-    DiagOFTModule,
-    ButterflyOFTModule,
-    GLoRAModule,
-    DyLoraModule,
-    IA3Module,
-]
-base_module_and_input = [
-    lambda dim: (nn.Linear(dim, dim), torch.randn(1, dim)),
-    lambda dim: (nn.Conv1d(dim, dim, 3, 1, 1), torch.randn(1, dim, 16)),
-    lambda dim: (nn.Conv2d(dim, dim, (3, 3), 1, 1), torch.randn(1, dim, 16, 16)),
-    # lambda dim: (nn.Conv3d(dim, dim, (3, 3, 3), 1, 1), torch.randn(1, dim, 16, 16, 16)),
-]
-device_and_dtype = [
-    (torch.device("cpu"), torch.float32),
-]
-weight_decompose = [False, True]
-use_tucker = [False, True]
-use_scalar = [False, True]
-
-if torch.cuda.is_available():
-    device_and_dtype.append((torch.device("cuda"), torch.float32))
-    device_and_dtype.append((torch.device("cuda"), torch.float16))
-    device_and_dtype.append((torch.device("cuda"), torch.bfloat16))
-
-if torch.backends.mps.is_available():
-    device_and_dtype.append((torch.device("mps"), torch.float32))
-
-
-patch_forward_param_list = list(
-    product(
-        modules,
-        base_module_and_input,
-        device_and_dtype,
-        weight_decompose,
-        use_tucker,
-        use_scalar,
-    )
-)
-
-
-class LycorisModuleTests(unittest.TestCase):
-    @parameterized.expand(patch_forward_param_list)
-    def test_lycoris_functional(self, module, base, device_dtype, wd, tucker, scalar):
-        base, test_input = base(16)
-        device, dtype = device_dtype