lycoris-lora 3.0.0.dev9__py3-none-any.whl

lycoris/functional/locon.py

@@ -0,0 +1,102 @@
+import math
+
+import torch
+import torch.nn as nn
+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
+
+    Args:
+        org_weight (torch.Tensor): the weight tensor
+        rank (int): low rank
+
+    Returns:
+        torch.Tensor: down, up[, mid]
+    """
+    out_dim, in_dim, *k = org_weight.shape
+    if k and tucker:
+        down = torch.empty(rank, in_dim, *(1 for _ in k))
+        up = torch.empty(out_dim, rank, *(1 for _ in k))
+        mid = torch.empty(rank, rank, *k)
+        nn.init.kaiming_uniform_(down, a=math.sqrt(5))
+        nn.init.constant_(up, 0)
+        nn.init.kaiming_uniform_(mid, a=math.sqrt(5))
+        return down, up, mid
+    else:
+        down = torch.empty(rank, in_dim)
+        up = torch.empty(out_dim, rank)
+        nn.init.kaiming_uniform_(down, a=math.sqrt(5))
+        nn.init.constant_(up, 0)
+        return down, up
+
+
+def lora_diff_weight(d, u, m=None, gamma=1.0):
+    """### lora_diff_weight
+
+    Get ΔW = BA, where BA is low rank decomposition
+
+    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): middle weight of tucker decomposition
+        gamma (float, optional): scale factor, normally alpha/rank here
+
+    Returns:
+        torch.Tensor: ΔW
+    """
+    R, I, *k = d.shape
+    O, R, *_ = u.shape
+    u = u * gamma
+
+    if m is None:
+        result = u.reshape(-1, u.size(1)) @ d.reshape(d.size(0), -1)
+    else:
+        R, R, *k = m.shape
+        u = u.reshape(u.size(0), -1).transpose(0, 1)
+        d = d.reshape(d.size(0), -1)
+        result = rebuild_tucker(m, u, d)
+    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
+
+    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): 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
+    """
+    if m is not None:
+        down = FUNC_LIST[d.dim()](x, d)
+        mid = FUNC_LIST[d.dim()](down, m, **extra_args)
+        up = FUNC_LIST[d.dim()](mid, u)
+    else:
+        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/functional/loha.py

@@ -0,0 +1,187 @@
+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
+    def forward(ctx, w1d, w1u, w2d, w2u, scale=torch.tensor(1)):
+        ctx.save_for_backward(w1d, w1u, w2d, w2u, scale)
+        diff_weight = ((w1u @ w1d) * (w2u @ w2d)) * scale
+        return diff_weight
+
+    @staticmethod
+    def backward(ctx, grad_out):
+        (w1d, w1u, w2d, w2u, scale) = ctx.saved_tensors
+        grad_out = grad_out * scale
+        temp = grad_out * (w2u @ w2d)
+        grad_w1u = temp @ w1d.T
+        grad_w1d = w1u.T @ temp
+
+        temp = grad_out * (w1u @ w1d)
+        grad_w2u = temp @ w2d.T
+        grad_w2d = w2u.T @ temp
+
+        del temp
+        return grad_w1d, grad_w1u, grad_w2d, grad_w2u, None
+
+
+class HadaWeightTucker(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, t1, w1d, w1u, t2, w2d, w2u, scale=torch.tensor(1)):
+        ctx.save_for_backward(t1, w1d, w1u, t2, w2d, w2u, scale)
+
+        rebuild1 = torch.einsum("i j ..., j r, i p -> p r ...", t1, w1d, w1u)
+        rebuild2 = torch.einsum("i j ..., j r, i p -> p r ...", t2, w2d, w2u)
+
+        return rebuild1 * rebuild2 * scale
+
+    @staticmethod
+    def backward(ctx, grad_out):
+        (t1, w1d, w1u, t2, w2d, w2u, scale) = ctx.saved_tensors
+        grad_out = grad_out * scale
+
+        temp = torch.einsum("i j ..., j r -> i r ...", t2, w2d)
+        rebuild = torch.einsum("i j ..., i r -> r j ...", temp, w2u)
+
+        grad_w = rebuild * grad_out
+        del rebuild
+
+        grad_w1u = torch.einsum("r j ..., i j ... -> r i", temp, grad_w)
+        grad_temp = torch.einsum("i j ..., i r -> r j ...", grad_w, w1u.T)
+        del grad_w, temp
+
+        grad_w1d = torch.einsum("i r ..., i j ... -> r j", t1, grad_temp)
+        grad_t1 = torch.einsum("i j ..., j r -> i r ...", grad_temp, w1d.T)
+        del grad_temp
+
+        temp = torch.einsum("i j ..., j r -> i r ...", t1, w1d)
+        rebuild = torch.einsum("i j ..., i r -> r j ...", temp, w1u)
+
+        grad_w = rebuild * grad_out
+        del rebuild
+
+        grad_w2u = torch.einsum("r j ..., i j ... -> r i", temp, grad_w)
+        grad_temp = torch.einsum("i j ..., i r -> r j ...", grad_w, w2u.T)
+        del grad_w, temp
+
+        grad_w2d = torch.einsum("i r ..., i j ... -> r j", t2, grad_temp)
+        grad_t2 = torch.einsum("i j ..., j r -> i r ...", grad_temp, w2d.T)
+        del grad_temp
+        return grad_t1, grad_w1d, grad_w1u, grad_t2, grad_w2d, grad_w2u, None
+
+
+def make_weight(w1d, w1u, w2d, w2u, scale):
+    return HadaWeight.apply(w1d, w1u, w2d, w2u, scale)
+
+
+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
+
+    Get ΔW = BA, where BA is low rank decomposition
+
+    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:
+        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:
+        R, I, *k = w1d.shape
+        O, R, *_ = w1u.shape
+        w1d = w1d.reshape(w1d.size(0), -1)
+        w1u = w1u.reshape(-1, w1u.size(1))
+        w2d = w2d.reshape(w2d.size(0), -1)
+        w2u = w1u.reshape(-1, w2u.size(1))
+        result = make_weight(w1d, w1u, w2d, w2u, gamma)
+
+    result = result.reshape(O, I, *k)
+    return result
+
+
+def loha_bypass_forward_diff(
+    x, w1d, w1u, w2d, w2u, t1=None, t2=None, gamma=1.0, extra_args={}
+):
+    """### loha_bypass_forward_diff
+
+    Args:
+        x (torch.Tensor): input tensor
+        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
+    """
+    diff_w = loha_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/functional/lokr.py

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

lycoris/hcp/__init__.py

@@ -0,0 +1,4 @@
+from .base import LycorisPluginContainer, LycorisPluginBlock
+from .loha import LohaBlock
+from .lokr import LokrBlock
+from .diag_oft import DiagOFTBlock

lycoris/hcp/base.py

@@ -0,0 +1,153 @@
+from typing import Tuple, Callable, Any, Optional
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from diffusers.models.lora import LoRACompatibleLinear, LoRACompatibleConv
+
+try:
+    import hcpdiff
+    from hcpdiff.models.plugin import PatchPluginContainer, PatchPluginBlock
+    from hcpdiff.utils.utils import isinstance_list
+except ImportError:
+    print(
+        "HCP-Diffusion is not installed, hcp feature will be disabled. "
+        "Use `pip install hcpdiff` to install hcpdiff."
+    )
+
+from ..utils.logger import logger
+
+
+class LycorisPluginContainer(PatchPluginContainer):
+    def __init__(self, host_name: str, host: nn.Module, parent_block: nn.Module):
+        super(LycorisPluginContainer, self).__init__(host_name, host, parent_block)
+
+        self.op = None
+        self.module_type = ""
+        self.extra_args = {}
+        if isinstance(host, nn.Conv2d):
+            self.op = F.conv2d
+            self.module_type = "conv"
+            self.extra_args = {
+                "stride": host.stride,
+                "padding": host.padding,
+                "dilation": host.dilation,
+                "groups": host.groups,
+            }
+        elif isinstance(host, nn.Linear):
+            self.op = F.linear
+            self.module_type = "linear"
+            self.extra_args = {}
+        elif isinstance(host, nn.LayerNorm):
+            self.op = F.layer_norm
+            self.module_type = "norm"
+            self.extra_args = {
+                "normalized_shape": host.normalized_shape,
+                "eps": host.eps,
+            }
+        elif isinstance(host, nn.GroupNorm):
+            self.op = F.group_norm
+            self.module_type = "norm"
+            self.extra_args = {"num_groups": host.num_groups, "eps": host.eps}
+        else:
+            logger.warning(
+                f"Unsupported host type: '{type(host)}' for LyCORIS plugin. "
+                "Will ignore this block."
+            )
+
+    def forward(self, *args, **kwargs):
+        if self.op is None:
+            return self._host(*args, **kwargs)
+
+        if (
+            isinstance(self._host, (LoRACompatibleLinear, LoRACompatibleConv))
+            and len(args) == 2
+        ):
+            scale = args[1]
+            args = [args[0]]
+        else:
+            scale = 1
+
+        org_weight = self._host.weight
+        org_bias = getattr(self._host, "bias", None)
+        new_weight = org_weight
+        new_bias = org_bias
+        total_diff_output = 0
+
+        for name in self.plugin_names:
+            lyco_plugin_block = getattr(self, name)
+            (
+                diff_weight,
+                diff_bias,
+                diff_output,
+                update_weight,
+                update_bias,
+            ) = lyco_plugin_block(
+                org_weight, org_bias, new_weight, new_bias, *args, **kwargs
+            )
+            # logger.debug(f'diff weight norm: {torch.norm(diff_weight).detach().cpu().item()}')
+            if diff_weight is not None:
+                if update_weight:
+                    new_weight = new_weight + diff_weight * scale
+                else:
+                    new_weight = diff_weight
+            if diff_bias is not None:
+                if update_bias:
+                    new_bias = (
+                        diff_bias * scale
+                        if new_bias is None
+                        else new_bias + diff_bias * scale
+                    )
+                else:
+                    new_bias = diff_bias
+            if diff_output is not None:
+                total_diff_output = total_diff_output + diff_output * scale
+
+        weight_dict = {
+            "weight": new_weight,
+            "bias": new_bias,
+        }
+        output = self.op(*args, **weight_dict, **self.extra_args)
+        return output + total_diff_output
+
+
+class LycorisPluginBlock(PatchPluginBlock):
+    container_cls = LycorisPluginContainer
+
+    def __init__(
+        self,
+        *args,
+        dim=4,
+        alpha=1.0,
+        dropout=0.0,
+        rank_dropout=0.0,
+        module_dropout=0.0,
+        use_tucker=False,
+        **kwargs,
+    ):
+        super(LycorisPluginBlock, self).__init__(*args, **kwargs)
+        self.module_type = self.container().module_type
+        self.dim = dim
+        self.alpha = alpha
+        self.dropout = dropout
+        self.rank_dropout = rank_dropout
+        self.module_dropout = module_dropout
+        self.tucker = False
+
+        if self.module_type == "conv":
+            self.out_dim = self.host().out_channels
+            self.in_dim = self.host().in_channels
+            self.k_size = self.host().kernel_size
+            self.tucker = use_tucker
+            self.shape = [self.out_dim, self.in_dim, *self.k_size]
+        elif self.module_type == "linear":
+            self.out_dim = self.host().out_features
+            self.in_dim = self.host().in_features
+            self.shape = [self.out_dim, self.in_dim]
+        elif self.module_type == "norm":
+            self.dim = self.host().weight.shape[0]
+        else:
+            raise NotImplementedError
+
+    def forward(self, org_weight, org_bias, new_weight, new_bias, *args, **kwargs):
+        raise NotImplementedError