lycoris-lora 2.3.0.dev11__tar.gz → 3.0.0.dev1__tar.gz

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

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

@@ -0,0 +1,13 @@
+from .general import (
+    rebuild_tucker,
+    factorization,
+    power2factorization,
+)
+from .locon import (
+    lora_diff_weight,
+    lora_bypass_forward_diff,
+)
+from .loha import (
+    loha_diff_weight,
+    loha_bypass_forward_diff,
+)

lycoris_lora-3.0.0.dev1/lycoris/functional/general.py

@@ -0,0 +1,74 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def rebuild_tucker(t, wa, wb):
+    rebuild2 = torch.einsum("i j ..., i p, j r -> p r ...", t, wa, wb)
+    return rebuild2
+
+
+def factorization(dimension: int, factor: int = -1) -> tuple[int, int]:
+    """
+    return a tuple of two value of input dimension decomposed by the number closest to factor
+    second value is higher or equal than first value.
+
+    In LoRA with Kroneckor Product, first value is a value for weight scale.
+    second value is a value for weight.
+
+    Because of non-commutative property, A⊗B ≠ B⊗A. Meaning of two matrices is slightly different.
+
+    examples)
+    factor
+        -1               2                4               8               16               ...
+    127 -> 1, 127   127 -> 1, 127    127 -> 1, 127   127 -> 1, 127   127 -> 1, 127
+    128 -> 8, 16    128 -> 2, 64     128 -> 4, 32    128 -> 8, 16    128 -> 8, 16
+    250 -> 10, 25   250 -> 2, 125    250 -> 2, 125   250 -> 5, 50    250 -> 10, 25
+    360 -> 8, 45    360 -> 2, 180    360 -> 4, 90    360 -> 8, 45    360 -> 12, 30
+    512 -> 16, 32   512 -> 2, 256    512 -> 4, 128   512 -> 8, 64    512 -> 16, 32
+    1024 -> 32, 32  1024 -> 2, 512   1024 -> 4, 256  1024 -> 8, 128  1024 -> 16, 64
+    """
+
+    if factor > 0 and (dimension % factor) == 0:
+        m = factor
+        n = dimension // factor
+        if m > n:
+            n, m = m, n
+        return m, n
+    if factor < 0:
+        factor = dimension
+    m, n = 1, dimension
+    length = m + n
+    while m < n:
+        new_m = m + 1
+        while dimension % new_m != 0:
+            new_m += 1
+        new_n = dimension // new_m
+        if new_m + new_n > length or new_m > factor:
+            break
+        else:
+            m, n = new_m, new_n
+    if m > n:
+        n, m = m, n
+    return m, n
+
+
+def power2factorization(dimension: int, factor: int = -1) -> tuple[int, int]:
+    """
+    m = 2k
+    n = 2**p
+    m*n = dim
+    """
+
+    # Find the first solution and check if it is even doable
+    m = n = 0
+    while m <= factor:
+        m += 2
+        while dimension % m != 0 and m < dimension:
+            m += 2
+        if m > factor:
+            break
+        if sum(int(i) for i in f"{dimension//m:b}") == 1:
+            n = dimension // m
+
+    return dimension // n, n

{lycoris_lora-2.3.0.dev11 → lycoris_lora-3.0.0.dev1}/lycoris/functional/locon.py

@@ -97,7 +97,7 @@ def lora_bypass_forward_diff(x, d, u, m=None, gamma=1.0, extra_args={}):
 if __name__ == "__main__":
     w = torch.randn(32, 32, 3, 3, 3)
     d, u, m = lora_weight_gen(w, 4)
-    u = u+0.01
+    u = u + 0.01
     extra_args = {"padding": 1}
 
     x = torch.randn(1, 32, 8, 8, 8)

{lycoris_lora-2.3.0.dev11 → lycoris_lora-3.0.0.dev1}/lycoris/functional/loha.py

@@ -15,12 +15,12 @@ class HadaWeight(torch.autograd.Function):
         (w1d, w1u, w2d, w2u, scale) = ctx.saved_tensors
         grad_out = grad_out * scale
         temp = grad_out * (w2u @ w2d)
-        grad_w1d = temp @ w1d.T
-        grad_w1u = w1u.T @ temp
+        grad_w1u = temp @ w1d.T
+        grad_w1d = w1u.T @ temp
 
         temp = grad_out * (w1u @ w1d)
-        grad_w2d = temp @ w2d.T
-        grad_w2u = w2u.T @ temp
+        grad_w2u = temp @ w2d.T
+        grad_w2d = w2u.T @ temp
 
         del temp
         return grad_w1d, grad_w1u, grad_w2d, grad_w2u, None
@@ -47,11 +47,11 @@ class HadaWeightTucker(torch.autograd.Function):
         grad_w = rebuild * grad_out
         del rebuild
 
-        grad_w1d = torch.einsum("r j ..., i j ... -> r i", temp, grad_w)
+        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_w1u = torch.einsum("i r ..., i j ... -> r j", t1, grad_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
 
@@ -61,11 +61,11 @@ class HadaWeightTucker(torch.autograd.Function):
         grad_w = rebuild * grad_out
         del rebuild
 
-        grad_w2d = torch.einsum("r j ..., i j ... -> r i", temp, grad_w)
+        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_w2u = torch.einsum("i r ..., i j ... -> r j", t2, grad_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

lycoris_lora-2.3.0.dev11/lycoris/kohya/__init__.py → lycoris_lora-3.0.0.dev1/lycoris/kohya.py

@@ -10,24 +10,24 @@ from typing import List
 
 import torch
 
-from .utils import *
-from ..wrapper import LycorisNetwork
-from ..modules.locon import LoConModule
-from ..modules.loha import LohaModule
-from ..modules.ia3 import IA3Module
-from ..modules.lokr import LokrModule
-from ..modules.dylora import DyLoraModule
-from ..modules.glora import GLoRAModule
-from ..modules.norms import NormModule
-from ..modules.full import FullModule
-from ..modules.diag_oft import DiagOFTModule
-from ..modules.boft import ButterflyOFTModule
-from ..modules import make_module
-
-from ..config import PRESET
-from ..utils.preset import read_preset
-from ..utils import get_module, str_bool
-from ..logging import logger
+from library.train_util import *
+from .wrapper import LycorisNetwork
+from .modules.locon import LoConModule
+from .modules.loha import LohaModule
+from .modules.ia3 import IA3Module
+from .modules.lokr import LokrModule
+from .modules.dylora import DyLoraModule
+from .modules.glora import GLoRAModule
+from .modules.norms import NormModule
+from .modules.full import FullModule
+from .modules.diag_oft import DiagOFTModule
+from .modules.boft import ButterflyOFTModule
+from .modules import make_module
+
+from .config import PRESET
+from .utils.preset import read_preset
+from .utils import get_module, str_bool
+from .logging import logger
 
 
 network_module_dict = {

{lycoris_lora-2.3.0.dev11 → lycoris_lora-3.0.0.dev1}/lycoris/modules/base.py

@@ -61,6 +61,7 @@ class ModuleCustomSD(nn.Module):
 
 
 class LycorisBaseModule(ModuleCustomSD):
+    dtype_tensor: torch.Tensor
     support_module = {}
 
     def __init__(
@@ -150,6 +151,8 @@ class LycorisBaseModule(ModuleCustomSD):
         else:
             self.module_type = "unknown"
 
+        self.register_buffer("dtype_tensor", torch.tensor(0.0), persistent=False)
+
         self.is_bnb = False
         if isinstance(org_module, QuantLinears):
             if not bypass_mode:
@@ -162,9 +165,30 @@ class LycorisBaseModule(ModuleCustomSD):
         self.rank_dropout_scale = rank_dropout_scale
         self.module_dropout = module_dropout
 
+        ## Dropout things
+        # Since LoKr/LoHa/OFT/BOFT are hard to follow the rank_dropout definition from kohya
+        # We redefine the dropout procedure here.
+        # g(x) = WX + drop(Brank_drop(AX)) for LoCon(lora), bypass
+        # g(x) = WX + drop(ΔWX) for any algo except LoCon(lora), bypass
+        # g(x) = (W + Brank_drop(A))X for LoCon(lora), rebuid
+        # g(x) = (W + rank_drop(ΔW))X for any algo except LoCon(lora), rebuild
+        self.drop = nn.Identity() if dropout == 0.0 else nn.Dropout(dropout)
+        self.rank_drop = (
+            nn.Identity() if rank_dropout == 0.0 else nn.Dropout(rank_dropout)
+        )
+
         self.multiplier = multiplier
+        self.org_forward = org_module.forward
         self.org_module = [org_module]
 
+    @property
+    def dtype(self):
+        return self.dtype_tensor.dtype
+
+    @property
+    def device(self):
+        return self.dtype_tensor.device
+
     @property
     def org_weight(self):
         return self.org_module[0].weight

{lycoris_lora-2.3.0.dev11 → lycoris_lora-3.0.0.dev1}/lycoris/modules/boft.py

@@ -8,6 +8,7 @@ import torch.nn.functional as F
 from einops import rearrange
 
 from .base import LycorisBaseModule
+from ..functional import power2factorization
 from ..logging import logger
 from ..utils.bnb import LinearNF4
 
@@ -22,30 +23,15 @@ def log_butterfly_factorize(dim, factor, result):
 
 
 def butterfly_factor(dimension: int, factor: int = -1) -> tuple[int, int]:
-    """
-    m = 2k
-    n = 2**p
-    m*n = dim
-    """
-
-    # Find the first solution and check if it is even doable
-    m = n = 0
-    while m <= factor:
-        m += 2
-        while dimension % m != 0 and m < dimension:
-            m += 2
-        if m > factor:
-            break
-        if sum(int(i) for i in f"{dimension//m:b}") == 1:
-            n = dimension // m
+    m, n = power2factorization(dimension, factor)
 
     if n == 0:
         raise ValueError(
             f"It is impossible to decompose {dimension} with factor {factor} under BOFT constrains."
         )
 
-    log_butterfly_factorize(dimension, factor, (dimension // n, n))
-    return dimension // n, n
+    log_butterfly_factorize(dimension, factor, (m, n))
+    return m, n
 
 
 class ButterflyOFTModule(LycorisBaseModule):
@@ -108,7 +94,7 @@ class ButterflyOFTModule(LycorisBaseModule):
 
     @property
     def I(self):
-        return torch.eye(self.block_size, device=next(self.parameters()).device)
+        return torch.eye(self.block_size, device=self.device)
 
     def get_r(self):
         I = self.I
@@ -215,29 +201,32 @@ class ButterflyOFTModule(LycorisBaseModule):
 
 if __name__ == "__main__":
     base = nn.Linear(128, 128).cuda()
-    lokr = ButterflyOFTModule("test", base, 1, 4, 1, weight_decompose=True).cuda()
-    print(lokr)
-    test_input = torch.randn(1, 128).cuda()
-    test_output = lokr(test_input)
+    boft = ButterflyOFTModule(
+        "test", base, 1, 4, 1, weight_decompose=True, factor=8
+    ).cuda()
+    print(boft)
+    test_input = torch.randn(1, 77, 128).cuda()
+    test_output = boft(test_input)
+    torch.sum(test_output).backward()
     print(test_output.shape)
-    print(F.mse_loss(test_output, lokr.bypass_forward(test_input)))
 
     base_4bit = LinearNF4(128, 128)
     base_4bit.load_state_dict(base.state_dict())
     base_4bit.cuda()
-    qlocon = ButterflyOFTModule(
+    qboft = ButterflyOFTModule(
         "test", base_4bit, 1, 4, 1, weight_decompose=False
     ).cuda()
-    print(qlocon)
-    test_input = torch.randn(1, 128).cuda()
-    test_output = qlocon(test_input)
+    print(qboft)
+    test_output = qboft(test_input)
+    torch.sum(test_output).backward()
     print(test_output.shape)
 
     base = nn.Conv2d(128, 128, 3, 1, 1)
-    lokr = ButterflyOFTModule(
+    boft = ButterflyOFTModule(
         "test", base, 1, 4, 1, weight_decompose=True, use_tucker=True
     )
-    print(lokr)
+    print(boft)
     test_input = torch.randn(1, 128, 16, 16)
-    test_output = lokr(test_input)
+    test_output = boft(test_input)
+    torch.sum(test_output).backward()
     print(test_output.shape)

{lycoris_lora-2.3.0.dev11 → lycoris_lora-3.0.0.dev1}/lycoris/modules/diag_oft.py

@@ -7,7 +7,7 @@ import torch.nn.functional as F
 from einops import rearrange
 
 from .base import LycorisBaseModule
-from .lokr import factorization
+from ..functional import factorization
 from ..logging import logger
 from ..utils.bnb import LinearNF4
 
@@ -83,7 +83,7 @@ class DiagOFTModule(LycorisBaseModule):
 
     @property
     def I(self):
-        return torch.eye(self.block_size, device=next(self.parameters()).device)
+        return torch.eye(self.block_size, device=self.device)
 
     def get_r(self):
         I = self.I
@@ -100,22 +100,19 @@ class DiagOFTModule(LycorisBaseModule):
 
     def make_weight(self, scale=1, device=None, diff=False):
         r = self.get_r()
+        _, *shape = self.org_weight.shape
         org_weight = self.org_weight.to(device, dtype=r.dtype)
-        org_weight = rearrange(
-            org_weight,
-            "(k n) ... -> k n ...",
-            k=self.block_num,
-            n=self.block_size,
-        )
+        org_weight = org_weight.view(self.block_num, self.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 * scale - scale * self.I + (self.I if diff else 0),
+            self.rank_drop(r * scale) - scale * self.I + (0 if diff else self.I),
             org_weight,
-        )
-        weight = rearrange(weight, "k m ... -> (k m) ...")
+        ).view(-1, *shape)
         if self.rescaled:
             weight = self.rescale * weight
+            if diff:
+                weight = weight + (self.rescale - 1) * org_weight
         return weight
 
     def get_diff_weight(self, multiplier=1, shape=None, device=None):
@@ -148,18 +145,21 @@ class DiagOFTModule(LycorisBaseModule):
         org_out = self.org_forward(x)
         if self.op == F.conv2d:
             org_out = org_out.transpose(1, -1)
-        org_out = rearrange(
-            org_out,
-            "... (k n) ->... k n",
-            k=self.block_num,
-            n=self.block_size,
-        )
+        *shape, _ = org_out.shape
+        org_out = org_out.view(*shape, self.block_num, self.block_size)
+        mask = neg_mask = 1
+        if self.dropout != 0 and self.training:
+            mask = torch.ones_like(org_out)
+            mask = self.drop(mask)
+            neg_mask = torch.max(mask) - mask
         oft_out = torch.einsum(
             "k n m, ... k n -> ... k m",
-            r * scale - scale * self.I + (self.I if diff else 0),
+            r * scale * mask + (1 - scale) * self.I * neg_mask,
             org_out,
         )
-        out = rearrange(oft_out, "... k m -> ... (k m)")
+        if diff:
+            out = out - org_out
+        out = oft_out.view(*shape, -1)
         if self.op == F.conv2d:
             out = out.transpose(1, -1)
         return out
@@ -186,24 +186,26 @@ class DiagOFTModule(LycorisBaseModule):
 
 if __name__ == "__main__":
     base = nn.Linear(128, 128).cuda()
-    lokr = DiagOFTModule("test", base, 1, 4, 1, weight_decompose=True).cuda()
-    print(lokr)
-    test_input = torch.randn(1, 128).cuda()
-    test_output = lokr(test_input)
+    doft = DiagOFTModule("test", base, 1, 4, 1, weight_decompose=True, factor=8).cuda()
+    print(doft)
+    test_input = torch.randn(1, 77, 128).cuda()
+    test_output = doft(test_input)
+    torch.sum(test_output).backward()
     print(test_output.shape)
 
     base_4bit = LinearNF4(128, 128)
     base_4bit.load_state_dict(base.state_dict())
     base_4bit.cuda()
-    qlocon = DiagOFTModule("test", base_4bit, 1, 4, 1, weight_decompose=False).cuda()
-    print(qlocon)
-    test_input = torch.randn(1, 128).cuda()
-    test_output = qlocon(test_input)
+    qdoft = DiagOFTModule("test", base_4bit, 1, 4, 1, weight_decompose=False).cuda()
+    print(qdoft)
+    test_output = qdoft(test_input)
+    torch.sum(test_output).backward()
     print(test_output.shape)
 
     base = nn.Conv2d(128, 128, 3, 1, 1)
-    lokr = DiagOFTModule("test", base, 1, 4, 1, weight_decompose=True, use_tucker=True)
-    print(lokr)
+    doft = DiagOFTModule("test", base, 1, 4, 1, weight_decompose=True, use_tucker=True)
+    print(doft)
     test_input = torch.randn(1, 128, 16, 16)
-    test_output = lokr(test_input)
+    test_output = doft(test_input)
+    torch.sum(test_output).backward()
     print(test_output.shape)

{lycoris_lora-2.3.0.dev11 → lycoris_lora-3.0.0.dev1}/lycoris/modules/full.py

@@ -41,6 +41,11 @@ class FullModule(LycorisBaseModule):
         if self.module_type not in self.support_module:
             raise ValueError(f"{self.module_type} is not supported in Full algo.")
 
+        if self.is_bnb:
+            raise ValueError(
+                "Quant Linear is not supported and meaningless in Full algo."
+            )
+
         if self.bypass_mode:
             raise ValueError("bypass mode is not supported in Full algo.")
 
@@ -50,6 +55,14 @@ class FullModule(LycorisBaseModule):
         else:
             self.bias = None
 
+    @property
+    def org_weight(self):
+        return self._org_weight[0]
+
+    @org_weight.setter
+    def org_weight(self, value):
+        self.org_module[0].weight.data.copy_(value)
+
     def apply_to(self, **kwargs):
         self.org_forward = self.org_module[0].forward
         self.org_module[0].forward = self.forward
@@ -85,32 +98,40 @@ class FullModule(LycorisBaseModule):
             diff_bias = state_dict["diff_b"]
             self.bias.data.add_(diff_bias)
 
-    def make_weight(self, scale=1, device=None, original=False):
-        if original:
-            weight = self._org_weight[0].to(device, dtype=self.weight.dtype)
-            if self.org_bias is not None:
-                bias = self.org_bias[0].to(device, dtype=self.bias.dtype)
-            else:
-                bias = None
-            return weight, bias
+    def make_weight(self, scale=1, device=None):
         drop = (
             torch.rand(self.dim, device=device) > self.rank_dropout
             if self.rank_dropout and self.training
             else 1
         )
         if drop != 1 or scale != 1:
-            org_weight = self.org_module[0].weight.to(device, dtype=self.weight.dtype)
-            diff = self.weight.to(device) - org_weight
-            weight = diff * drop * scale + org_weight
-            if self.bias is not None:
-                org_bias = self.org_module[0].bias.to(device, dtype=self.bias.dtype)
-                diff_b = self.bias.to(device) - org_bias
-                bias = diff_b * drop * scale + org_bias
+            diff_w, diff_b = self.get_diff_weight(scale, device=device)
+            weight = self.org_module[0].weight + diff_w * drop
+            bias = self.org_module[0].bias + diff_b * drop
         else:
             weight = self.weight
             bias = self.bias
         return weight, bias
 
+    def get_diff_weight(self, multiplier=1, shape=None, device=None):
+        org_weight = self.org_module[0].weight.to(device, dtype=self.weight.dtype)
+        diff = self.weight.to(device) - org_weight
+        diff_b = None
+        if shape:
+            diff = diff.view(shape)
+        if self.bias is not None:
+            org_bias = self.org_module[0].bias.to(device, dtype=self.bias.dtype)
+            diff_b = self.bias.to(device) - org_bias
+        if device is not None:
+            diff = diff.to(device)
+            if self.bias is not None:
+                diff_b = diff_b.to(device)
+        if multiplier != 1:
+            diff = diff * multiplier
+            if diff_b is not None:
+                diff_b = diff_b * multiplier
+        return diff * multiplier, diff_b
+
     def get_merged_weight(self, multiplier=1, shape=None, device=None):
         weight, bias = self.make_weight(multiplier, device)
         if shape is not None:
@@ -128,7 +149,27 @@ class FullModule(LycorisBaseModule):
             original = True
         else:
             original = False
+        if original:
+            return self.org_forward(x)
         scale = self.multiplier
-        weight, bias = self.make_weight(scale, x.device, original=original)
+        weight, bias = self.make_weight(scale, x.device)
         kw_dict = self.kw_dict | {"weight": weight, "bias": bias}
         return self.op(x, **kw_dict)
+
+
+if __name__ == "__main__":
+    base = nn.Linear(128, 128).cuda()
+    full = FullModule("test", base, 1, 4, 1, weight_decompose=True, factor=8).cuda()
+    print(full)
+    test_input = torch.randn(1, 77, 128).cuda()
+    test_output = full(test_input)
+    torch.sum(test_output).backward()
+    print(test_output.shape)
+
+    base = nn.Conv2d(128, 128, 3, 1, 1)
+    full = FullModule("test", base, 1, 4, 1, weight_decompose=True, use_tucker=True)
+    print(full)
+    test_input = torch.randn(1, 128, 16, 16)
+    test_output = full(test_input)
+    torch.sum(test_output).backward()
+    print(test_output.shape)