lycoris-lora 2.2.0.dev11__tar.gz → 2.2.0.post1__tar.gz

{lycoris_lora-2.2.0.dev11/lycoris_lora.egg-info → lycoris_lora-2.2.0.post1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: lycoris_lora
-Version: 2.2.0.dev11
+Version: 2.2.0.post1
 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-2.2.0.dev11 → lycoris_lora-2.2.0.post1}/README.md

@@ -153,6 +153,10 @@ forward_with_lyco = your_model(x)
 
 You can check my [HakuPhi](https://github.com/KohakuBlueleaf/HakuPhi) project to see how I utilize LyCORIS to finetune the Phi-1.5 models.
 
+#### Bitsandbytes support
+
+See [bnb_example.py](bnb_example.py) for example. Basically as same as standalone wrapper.
+
 #### Graphical interfaces and Colabs (via kohya trainer)
 
 You can also train LyCORIS with the following graphical interfaces
@@ -236,30 +240,30 @@ See [docs/Conversion-scripts.md](docs/Conversion-scripts.md) for more informatio
 
 For full log, please see [Change.md](Change.md)
 
-## 2024/02/18 update to 2.1.0
+## 2024/03/15 update to 2.2.0 - QLyCORIS and DoRA
 
 #### New Algo
 
-* [BOFT (Butterfly OFT)](https://arxiv.org/abs/2311.06243)
+* DoRA
+  * Ref: [DoRA: Weight-Decomposed Low-Rank Adaptation](https://github.com/KohakuBlueleaf/LyCORIS)
+* Weight decompose for LoHa and LoKr. (A.K.A DoHa/DoKr)
+  * DoRA/DoHa/DoKr will require smaller Learning rate!
 
-#### Improvements
+#### New Features
 
-* Faster, better extract script
-* support kohya-ss/sd-scripts image gen
-* support regex name in kohya-ss/sd-scripts
-* support resume on:
-  * full
-  * loha
-  * oft
-  * boft
-* Add logger into LyCORIS
+* Support "bypass" (a.k.a. adapter) mode for LoHa/LoKr/OFT/BOFT
+  * LoHa will require 2xFLOPs since we rebuild full diff weight and then do one more forward.
+  * LoKr, OFT, BOFT should be more efficient than LoHa in bypass mode.
+* Support [bnb 8bit/4bit Linear layer](https://github.com/TimDettmers/bitsandbytes) (a.k.a. QLyCORIS) with LoHa/LoKr/OFT/BOFT.
+  * This will force module to enable bypass mode.
 
 #### Fixes, slight changes
 
-* Update HCP convert for the case where only UNet or TE is trained.
-* Change arg names for conversion scripts.
-* Fix wrong TE prefix in merge scripts.
-* Fix warnings and confusing logging.
+* Refine some details about code quality. Based on the report from GitRoll. (Thx you gitroll!)
+* Remove redundant calculation in BOFT
+* rank_dropout has been removed from OFT/BOFT temporarily untill we ensure how to apply it.
+* Fix bugs in lokr when `lokr_w1_a` not exist.
+* Fix bugs in conversion scritps.
 
 ## Todo list
 

{lycoris_lora-2.2.0.dev11 → lycoris_lora-2.2.0.post1}/lycoris/kohya/__init__.py

@@ -350,10 +350,10 @@ class LycorisNetworkKohya(LycorisNetwork):
                     **kwargs,
                 )
             lora = None
-            if module.__class__.__name__ == "Linear" and lora_dim > 0:
+            if isinstance(module, torch.nn.Linear) and lora_dim > 0:
                 dim = dim or lora_dim
                 alpha = alpha or self.alpha
-            elif module.__class__.__name__ == "Conv2d":
+            elif isinstance(module, torch.nn.Conv2d):
                 k_size, *_ = module.kernel_size
                 if k_size == 1 and lora_dim > 0:
                     dim = dim or lora_dim

{lycoris_lora-2.2.0.dev11 → lycoris_lora-2.2.0.post1}/lycoris/modules/boft.py

@@ -8,7 +8,8 @@ import torch.nn.functional as F
 from einops import rearrange
 
 from .base import ModuleCustomSD
-from lycoris.logging import logger
+from ..logging import logger
+from ..utils.bnb import LinearNF4, QuantLinears, log_bypass
 
 
 @cache
@@ -67,6 +68,7 @@ class ButterflyOFTModule(ModuleCustomSD):
         rank_dropout_scale=False,
         constrain=0,
         rescaled=False,
+        bypass_mode=False,
         **kwargs,
     ):
         """if alpha == 0 or None, alpha is rank (no scaling)."""
@@ -89,7 +91,13 @@ class ButterflyOFTModule(ModuleCustomSD):
         else:
             raise NotImplementedError
 
-        out_dim = org_module.weight.shape[0]
+        if isinstance(org_module, QuantLinears):
+            if not bypass_mode:
+                log_bypass()
+            bypass_mode = True
+        self.bypass_mode = bypass_mode
+
+        out_dim = self.dim
         b, m_exp = butterfly_factor(out_dim, lora_dim)
         self.block_size = b
         self.block_num = m_exp
@@ -146,14 +154,6 @@ class ButterflyOFTModule(ModuleCustomSD):
         return r
 
     def make_weight(self, scale=1, device=None):
-        if self.rank_dropout and self.training:
-            drop = torch.rand(self.dim, device=device) < self.rank_dropout
-            drop = drop.to(self.oft_blocks.dtype)
-            if self.rank_dropout_scale:
-                drop /= drop.mean()
-        else:
-            drop = 1
-
         m = self.boft_m
         b = self.boft_b
         r_b = b // 2
@@ -173,10 +173,7 @@ class ButterflyOFTModule(ModuleCustomSD):
             inp = rearrange(inp, "d b ... -> (d b) ...")
             inp = rearrange(inp, "(c k g) ... -> (c g k) ...", g=2, k=2**i * r_b)
 
-        weight = inp
-        if self.rescaled:
-            weight = self.rescale * weight
-        return weight * drop
+        return inp
 
     @torch.no_grad()
     def apply_max_norm(self, max_norm, device=None):
@@ -191,23 +188,67 @@ class ButterflyOFTModule(ModuleCustomSD):
 
         return scaled, orig_norm * ratio
 
+    def bypass_forward(self, x, scale=1):
+        m = self.boft_m
+        b = self.boft_b
+        r_b = b // 2
+        r = self.get_r()
+        inp = self.org_forward(x)
+
+        for i in range(m):
+            bi = r[i]  # b_num, b_size, b_size
+            if i == 0:
+                # Apply multiplier/scale and rescale into first weight
+                bi = bi * scale + (1 - scale) * self.I
+                if self.rescaled:
+                    bi = bi * self.rescale
+            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
+
     def forward(self, x):
         if self.module_dropout and self.training:
             if torch.rand(1) < self.module_dropout:
                 return self.org_forward(x)
         scale = self.multiplier
 
-        w = self.make_weight(scale, x.device)
-        kw_dict = self.kw_dict | {"weight": w, "bias": self.org_module[0].bias}
-        return self.op(x, **kw_dict)
+        if self.bypass_mode:
+            return self.bypass_forward(x, scale)
+        else:
+            w = self.make_weight(scale, x.device)
+            kw_dict = self.kw_dict | {"weight": w, "bias": self.org_module[0].bias}
+            return self.op(x, **kw_dict)
 
 
 if __name__ == "__main__":
-    dim = 320
-    test_layer = nn.Linear(dim, dim)
-    test_layer.weight = nn.Parameter(torch.eye(dim))
-    test_boft = ButterflyOFTModule("test", test_layer, lora_dim=16)
-    test_boft.oft_blocks = nn.Parameter(torch.randn_like(test_boft.oft_blocks))
-
-    print(test_boft.make_weight(1, "cpu").shape)
-    print(torch.count_nonzero(test_boft.make_weight(1, "cpu")), dim * dim)
+    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)
+    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(
+        "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(test_output.shape)
+
+    base = nn.Conv2d(128, 128, 3, 1, 1)
+    lokr = ButterflyOFTModule(
+        "test", base, 1, 4, 1, weight_decompose=True, use_tucker=True
+    )
+    print(lokr)
+    test_input = torch.randn(1, 128, 16, 16)
+    test_output = lokr(test_input)
+    print(test_output.shape)

{lycoris_lora-2.2.0.dev11 → lycoris_lora-2.2.0.post1}/lycoris/modules/diag_oft.py

@@ -9,6 +9,7 @@ from einops import rearrange
 from .base import ModuleCustomSD
 from .lokr import factorization
 from ..logging import logger
+from ..utils.bnb import LinearNF4, QuantLinears, log_bypass
 
 
 @cache
@@ -40,6 +41,7 @@ class DiagOFTModule(ModuleCustomSD):
         rank_dropout_scale=False,
         constrain=0,
         rescaled=False,
+        bypass_mode=False,
         **kwargs,
     ):
         """if alpha == 0 or None, alpha is rank (no scaling)."""
@@ -62,7 +64,13 @@ class DiagOFTModule(ModuleCustomSD):
         else:
             raise NotImplementedError
 
-        out_dim = org_module.weight.shape[0]
+        if isinstance(org_module, QuantLinears):
+            if not bypass_mode:
+                log_bypass()
+            bypass_mode = True
+        self.bypass_mode = bypass_mode
+
+        out_dim = self.dim
         self.block_size, self.block_num = factorization(out_dim, lora_dim)
         # block_num > block_size
         self.rescaled = rescaled
@@ -120,14 +128,6 @@ class DiagOFTModule(ModuleCustomSD):
         return r
 
     def make_weight(self, scale=1, device=None):
-        if self.rank_dropout and self.training:
-            drop = torch.rand(self.dim, device=device) < self.rank_dropout
-            drop = drop.to(self.oft_blocks.dtype)
-            if self.rank_dropout_scale:
-                drop /= drop.mean()
-        else:
-            drop = 1
-
         r = self.get_r()
         org_weight = self.org_module[0].weight.to(device, dtype=r.dtype)
         org_weight = rearrange(
@@ -160,12 +160,59 @@ class DiagOFTModule(ModuleCustomSD):
 
         return scaled, orig_norm * ratio
 
+    def bypass_forward(self, x, scale=1):
+        r = self.get_r()
+        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,
+        )
+        oft_out = torch.einsum(
+            "k n m, ... k n -> ... k m", r * scale + (1 - scale) * self.I, org_out
+        )
+        out = rearrange(oft_out, "... k m -> ... (k m)")
+        if self.op == F.conv2d:
+            out = out.transpose(1, -1)
+        return out
+
     def forward(self, x):
         if self.module_dropout and self.training:
             if torch.rand(1) < self.module_dropout:
                 return self.org_forward(x)
         scale = self.multiplier
 
-        w = self.make_weight(scale, x.device)
-        kw_dict = self.kw_dict | {"weight": w, "bias": self.org_module[0].bias}
-        return self.op(x, **kw_dict)
+        if self.bypass_mode:
+            return self.bypass_forward(x, scale)
+        else:
+            w = self.make_weight(scale, x.device)
+            kw_dict = self.kw_dict | {"weight": w, "bias": self.org_module[0].bias}
+            return self.op(x, **kw_dict)
+
+
+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)
+    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)
+    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)
+    test_input = torch.randn(1, 128, 16, 16)
+    test_output = lokr(test_input)
+    print(test_output.shape)

{lycoris_lora-2.2.0.dev11 → lycoris_lora-2.2.0.post1}/lycoris/modules/locon.py

@@ -133,7 +133,6 @@ class LoConModule(ModuleCustomSD):
         if use_scalar:
             self.scalar = nn.Parameter(torch.tensor(0.0))
         else:
-            self.scalar: torch.FloatTensor
             self.register_buffer("scalar", torch.tensor(1.0), persistent=False)
         # same as microsoft's
         torch.nn.init.kaiming_uniform_(self.lora_down.weight, a=math.sqrt(5))
@@ -225,11 +224,11 @@ class LoConModule(ModuleCustomSD):
                 return self.org_forward(x)
         scale = self.scale * self.multiplier
 
-        dtype_device = next(self.parameters())
+        dtype = next(self.parameters()).dtype
         if not self.bypass_mode:
             weight = (
-                self.org_module[0].weight.data.to(dtype_device)
-                + self.make_weight(x.device).to(dtype_device) * scale
+                self.org_module[0].weight.data.to(device=x.device, dtype=dtype)
+                + self.make_weight(x.device).to(device=x.device, dtype=dtype) * scale
             )
             if self.wd:
                 weight = self.apply_weight_decompose(weight)

{lycoris_lora-2.2.0.dev11 → lycoris_lora-2.2.0.post1}/lycoris/modules/loha.py

@@ -281,7 +281,7 @@ class LohaModule(ModuleCustomSD):
 
     @torch.no_grad()
     def apply_max_norm(self, max_norm, device=None):
-        orig_norm = (self.get_weight() * self.scalar).norm()
+        orig_norm = (self.get_weight(self.shape) * self.scalar).norm()
         norm = torch.clamp(orig_norm, max_norm / 2)
         desired = torch.clamp(norm, max=max_norm)
         ratio = desired.cpu() / norm.cpu()

{lycoris_lora-2.2.0.dev11 → lycoris_lora-2.2.0.post1}/lycoris/modules/lokr.py

@@ -352,7 +352,7 @@ class LokrModule(ModuleCustomSD):
 
     @torch.no_grad()
     def apply_max_norm(self, max_norm, device=None):
-        orig_norm = self.get_weight().norm()
+        orig_norm = self.get_weight(self.shape).norm()
         norm = torch.clamp(orig_norm, max_norm / 2)
         desired = torch.clamp(norm, max=max_norm)
         ratio = desired.cpu() / norm.cpu()
@@ -448,9 +448,7 @@ class LokrModule(ModuleCustomSD):
 
 if __name__ == "__main__":
     base = nn.Linear(128, 128).cuda()
-    lokr = LokrModule(
-        "test", base, 1, 4, 1, weight_decompose=True, factor=8
-    ).cuda()
+    lokr = LokrModule("test", base, 1, 4, 1, weight_decompose=True, factor=8).cuda()
     print(lokr)
     test_input = torch.randn(1, 77, 128).cuda()
     test_output = lokr(test_input)

{lycoris_lora-2.2.0.dev11 → lycoris_lora-2.2.0.post1}/lycoris/wrapper.py

@@ -292,10 +292,10 @@ class LycorisNetwork(torch.nn.Module):
                     **kwargs,
                 )
             lora = None
-            if module.__class__.__name__ == "Linear" and lora_dim > 0:
+            if isinstance(module, torch.nn.Linear) and lora_dim > 0:
                 dim = dim or lora_dim
                 alpha = alpha or self.alpha
-            elif module.__class__.__name__ == "Conv2d":
+            elif isinstance(module, torch.nn.Conv2d):
                 k_size, *_ = module.kernel_size
                 if k_size == 1 and lora_dim > 0:
                     dim = dim or lora_dim

{lycoris_lora-2.2.0.dev11 → lycoris_lora-2.2.0.post1/lycoris_lora.egg-info}/PKG-INFO

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

@@ -4,7 +4,7 @@ from setuptools import setup, find_packages
 setup(
     name="lycoris_lora",
     packages=find_packages(),
-    version="2.2.0.dev11",
+    version="2.2.0.post1",
     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",