lycoris-lora 3.0.0.dev12__tar.gz → 3.0.0.dev14__tar.gz

{lycoris_lora-3.0.0.dev12/lycoris_lora.egg-info → lycoris_lora-3.0.0.dev14}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: lycoris_lora
-Version: 3.0.0.dev12
+Version: 3.0.0.dev14
 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.dev12 → lycoris_lora-3.0.0.dev14}/lycoris/config.py

@@ -108,4 +108,15 @@ PRESET = {
         "text_encoder_target_module": [],
         "text_encoder_target_name": [],
     },
+    "ia3": {
+        "enable_conv": False,
+        "unet_target_module": [],
+        "unet_target_name": ["to_k", "to_v", "ff.net.2"],
+        "text_encoder_target_module": [],
+        "text_encoder_target_name": ["k_proj", "v_proj", "mlp.fc2"],
+        "name_algo_map": {
+            "mlp.fc2": {"train_on_input": True},
+            "ff.net.2": {"train_on_input": True},
+        },
+    },
 }

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

@@ -0,0 +1 @@
+from .general import rebuild_tucker, factorization, power2factorization, FUNC_LIST

lycoris_lora-3.0.0.dev14/lycoris/functional/boft.py

@@ -0,0 +1,99 @@
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from einops import rearrange
+
+from .general import power2factorization, FUNC_LIST
+from .diag_oft import get_r
+
+
+def weight_gen(org_weight, max_block_size, boft_m=-1, rescale=False):
+    """### boft_weight_gen
+
+    Args:
+        org_weight (torch.Tensor): the weight tensor
+        max_block_size (int): max block size
+        rescale (bool, optional): whether to rescale the weight. Defaults to False.
+
+    Returns:
+        torch.Tensor: oft_blocks[, rescale_weight]
+    """
+    out_dim, *rest = org_weight.shape
+    block_size, block_num = power2factorization(out_dim, max_block_size)
+    max_boft_m = sum(int(i) for i in f"{block_num-1:b}") + 1
+    if boft_m == -1:
+        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 is not None:
+        return oft_blocks, torch.ones(out_dim, *[1] * len(rest))
+    else:
+        return oft_blocks, None
+
+
+def diff_weight(org_weight, oft_blocks, rescale=None, constraint=None):
+    """### boft_diff_weight
+
+    Args:
+        TODO
+
+    Returns:
+        torch.Tensor: ΔW
+    """
+    m, num, b, _ = oft_blocks.shape
+    r_b = b // 2
+    I = torch.eye(b, device=oft_blocks.device)
+    r = get_r(oft_blocks, I, constraint)
+    inp = org = org_weight.to(dtype=r.dtype)
+
+    for i in range(m):
+        bi = r[i]  # b_num, b_size, b_size
+        inp = rearrange(inp, "(c g k) ... -> (c k g) ...", g=2, k=2**i * r_b)
+        inp = rearrange(inp, "(d b) ... -> d b ...", b=b)
+        inp = torch.einsum("b i j, b j ... -> b i ...", bi, inp)
+        inp = rearrange(inp, "d b ... -> (d b) ...")
+        inp = rearrange(inp, "(c k g) ... -> (c g k) ...", g=2, k=2**i * r_b)
+
+    if rescale is not None:
+        inp = inp * rescale
+
+    return inp - org
+
+
+def bypass_forward_diff(
+    org_out, oft_blocks, rescale=None, constraint=None, need_transpose=False
+):
+    """### boft_bypass_forward_diff
+
+    Args:
+        TODO
+
+    Returns:
+        torch.Tensor: output tensor
+    """
+    m, num, b, _ = oft_blocks.shape
+    r_b = b // 2
+    I = torch.eye(b, device=oft_blocks.device)
+    r = get_r(oft_blocks, I, constraint)
+    inp = org = org_out.to(dtype=r.dtype)
+    if need_transpose:
+        inp = org = inp.transpose(1, -1)
+
+    for i in range(m):
+        bi = r[i]  # b_num, b_size, b_size
+        inp = rearrange(inp, "... (c g k) ->... (c k g)", g=2, k=2**i * r_b)
+        inp = rearrange(inp, "... (d b) -> ... d b", b=b)
+        inp = torch.einsum("b i j, ... b j -> ... b i", bi, inp)
+        inp = rearrange(inp, "... d b -> ... (d b)")
+        inp = rearrange(inp, "... (c k g) -> ... (c g k)", g=2, k=2**i * r_b)
+
+    if rescale is not None:
+        inp = inp * rescale.transpose(0, -1)
+
+    inp = inp - org
+    if need_transpose:
+        inp = inp.transpose(1, -1)
+    return inp

lycoris_lora-3.0.0.dev14/lycoris/functional/diag_oft.py

@@ -0,0 +1,103 @@
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .general import factorization, FUNC_LIST
+
+
+def get_r(oft_blocks, I=None, constraint=0):
+    if I is None:
+        I = torch.eye(oft_blocks.shape[-1], device=oft_blocks.device)
+    if I.ndim < oft_blocks.ndim:
+        for _ in range(oft_blocks.ndim - I.ndim):
+            I = I.unsqueeze(0)
+    # for Q = -Q^T
+    q = oft_blocks - oft_blocks.transpose(-1, -2)
+    normed_q = q
+    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
+    # use float() to prevent unsupported type
+    r = (I + normed_q) @ (I - normed_q).float().inverse()
+    return r
+
+
+def weight_gen(org_weight, max_block_size, rescale=False):
+    """### weight_gen
+
+    Args:
+        org_weight (torch.Tensor): the weight tensor
+        max_block_size (int): max block size
+        rescale (bool, optional): whether to rescale the weight. Defaults to False.
+
+    Returns:
+        torch.Tensor: oft_blocks[, rescale_weight]
+    """
+    out_dim, *rest = org_weight.shape
+    block_size, block_num = factorization(out_dim, max_block_size)
+    oft_blocks = torch.zeros(block_num, block_size, block_size)
+    if rescale:
+        return oft_blocks, torch.ones(out_dim, *[1] * len(rest))
+    else:
+        return oft_blocks, None
+
+
+def diff_weight(org_weight, oft_blocks, rescale=None, constraint=None):
+    """### diff_weight
+
+    Args:
+        TODO
+
+    Returns:
+        torch.Tensor: ΔW
+    """
+    I = torch.eye(oft_blocks.shape[1], device=oft_blocks.device)
+    r = get_r(oft_blocks, I, constraint)
+
+    block_num, block_size, _ = oft_blocks.shape
+    _, *shape = org_weight.shape
+    org_weight = org_weight.to(dtype=r.dtype)
+    org_weight = org_weight.view(block_num, block_size, *shape)
+    # Init R=0, so add I on it to ensure the output of step0 is original model output
+    weight = torch.einsum(
+        "k n m, k n ... -> k m ...",
+        r - I,
+        org_weight,
+    ).view(-1, *shape)
+    if rescale is not None:
+        weight = rescale * weight
+        weight = weight + (rescale - 1) * org_weight
+    return weight
+
+
+def bypass_forward_diff(
+    org_out, oft_blocks, rescale=None, constraint=None, need_transpose=False
+):
+    """### bypass_forward_diff
+
+    Args:
+        TODO
+
+    Returns:
+        torch.Tensor: output tensor
+    """
+    block_num, block_size, _ = oft_blocks.shape
+    I = torch.eye(block_size, device=oft_blocks.device)
+    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)
+    )
+    out = oft_out.view(*shape, -1)
+    if rescale is not None:
+        out = rescale.transpose(-1, 0) * out
+        out = out + (rescale - 1).transpose(-1, 0) * org_out
+    if need_transpose:
+        out = out.transpose(1, -1)
+    return out

{lycoris_lora-3.0.0.dev12 → lycoris_lora-3.0.0.dev14}/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.dev12 → lycoris_lora-3.0.0.dev14}/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.dev12 → lycoris_lora-3.0.0.dev14}/lycoris/functional/lokr.py

@@ -118,8 +118,8 @@ def 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 = 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))