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

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

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

@@ -4,10 +4,17 @@ from .general import (
     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/lycoris/functional/boft.py

@@ -0,0 +1,113 @@
+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 boft_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:
+        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):
+    """### 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)
+
+    return inp * rescale - org
+
+
+def boft_bypass_forward_diff(
+    org_out, need_transpose, oft_blocks, rescale=None, constraint=None
+):
+    """### 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_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)
+
+    inp = inp * rescale.transpose(0, -1) - 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/functional/diag_oft.py

@@ -0,0 +1,119 @@
+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 > 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 daig_oft_weight_gen(org_weight, max_block_size, rescale=False):
+    """### daig_oft_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 daig_oft_diff_weight(org_weight, oft_blocks, rescale=None, constraint=None):
+    """### daig_oft_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 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
+):
+    """### daig_oft_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)
+    *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
+
+
+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.dev11 → lycoris_lora-3.0.0.dev13}/lycoris/functional/lokr.py

@@ -17,7 +17,7 @@ def make_kron(w1, w2, scale):
     return rebuild * scale
 
 
-def weight_gen(
+def lokr_weight_gen(
     org_weight,
     rank,
     tucker=True,
@@ -245,7 +245,7 @@ def lokr_bypass_forward_diff(
 
 if __name__ == "__main__":
     w = torch.randn(32, 32, 3, 3, 3)
-    w1, w1a, w1b, w2, w2a, w2b, t = weight_gen(w, 2, tucker=False)
+    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)