lycoris-lora 3.2.0.dev1__tar.gz → 3.2.0.dev2__tar.gz

{lycoris_lora-3.2.0.dev1/lycoris_lora.egg-info → lycoris_lora-3.2.0.dev2}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.2
 Name: lycoris_lora
-Version: 3.2.0.dev1
+Version: 3.2.0.dev2
 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
@@ -23,3 +23,10 @@ Requires-Dist: torch
 Requires-Dist: einops
 Requires-Dist: toml
 Requires-Dist: tqdm
+Dynamic: author
+Dynamic: author-email
+Dynamic: classifier
+Dynamic: home-page
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary

{lycoris_lora-3.2.0.dev1 → lycoris_lora-3.2.0.dev2}/lycoris/config.py

@@ -6,12 +6,12 @@ PRESET = {
             "ResnetBlock2D",
             "Downsample2D",
             "Upsample2D",
-            "HunYuanDiTBlock", #HunYuanDiT
-            "DoubleStreamBlock", #Flux
-            "SingleStreamBlock", #Flux
-            "SingleDiTBlock", #SD3.5
-            "MMDoubleStreamBlock", #HunYuanVideo
-            "MMSingleStreamBlock", #HunYuanVideo
+            "HunYuanDiTBlock",  # HunYuanDiT
+            "DoubleStreamBlock",  # Flux
+            "SingleStreamBlock",  # Flux
+            "SingleDiTBlock",  # SD3.5
+            "MMDoubleStreamBlock",  # HunYuanVideo
+            "MMSingleStreamBlock",  # HunYuanVideo
         ],
         "unet_target_name": [
             "conv_in",
@@ -37,8 +37,8 @@ PRESET = {
             "DoubleStreamBlock",
             "SingleStreamBlock",
             "SingleDiTBlock",
-            "MMDoubleStreamBlock", #HunYuanVideo
-            "MMSingleStreamBlock", #HunYuanVideo
+            "MMDoubleStreamBlock",  # HunYuanVideo
+            "MMSingleStreamBlock",  # HunYuanVideo
         ],
         "unet_target_name": [
             "time_embedding.linear_1",
@@ -61,8 +61,8 @@ PRESET = {
             "DoubleStreamBlock",
             "SingleStreamBlock",
             "SingleDiTBlock",
-            "MMDoubleStreamBlock", #HunYuanVideo
-            "MMSingleStreamBlock", #HunYuanVideo
+            "MMDoubleStreamBlock",  # HunYuanVideo
+            "MMSingleStreamBlock",  # HunYuanVideo
         ],
         "unet_target_name": [],
         "text_encoder_target_module": [
@@ -100,8 +100,8 @@ PRESET = {
             "DoubleStreamBlock",
             "SingleStreamBlock",
             "SingleDiTBlock",
-            "MMDoubleStreamBlock", #HunYuanVideo
-            "MMSingleStreamBlock", #HunYuanVideo
+            "MMDoubleStreamBlock",  # HunYuanVideo
+            "MMSingleStreamBlock",  # HunYuanVideo
         ],
         "unet_target_name": [
             "conv_in",
@@ -120,8 +120,8 @@ PRESET = {
             "DoubleStreamBlock",
             "SingleStreamBlock",
             "SingleDiTBlock",
-            "MMDoubleStreamBlock", #HunYuanVideo
-            "MMSingleStreamBlock", #HunYuanVideo
+            "MMDoubleStreamBlock",  # HunYuanVideo
+            "MMSingleStreamBlock",  # HunYuanVideo
         ],
         "unet_target_name": [],
         "text_encoder_target_module": [],

{lycoris_lora-3.2.0.dev1 → lycoris_lora-3.2.0.dev2}/lycoris/kohya.py

@@ -63,10 +63,22 @@ def create_network(
     rs_lora = str_bool(kwargs.get("rs_lora", False))
     unbalanced_factorization = str_bool(kwargs.get("unbalanced_factorization", False))
     train_t5xxl = str_bool(kwargs.get("train_t5xxl", False))
-    #lora_plus
-    loraplus_lr_ratio = float(kwargs.get("loraplus_lr_ratio", None)) if kwargs.get("loraplus_lr_ratio", None) is not None else None
-    loraplus_unet_lr_ratio = float(kwargs.get("loraplus_unet_lr_ratio", None)) if kwargs.get("loraplus_unet_lr_ratio", None) is not None else None
-    loraplus_text_encoder_lr_ratio = float(kwargs.get("loraplus_text_encoder_lr_ratio", None)) if kwargs.get("loraplus_text_encoder_lr_ratio", None) is not None else None
+    # lora_plus
+    loraplus_lr_ratio = (
+        float(kwargs.get("loraplus_lr_ratio", None))
+        if kwargs.get("loraplus_lr_ratio", None) is not None
+        else None
+    )
+    loraplus_unet_lr_ratio = (
+        float(kwargs.get("loraplus_unet_lr_ratio", None))
+        if kwargs.get("loraplus_unet_lr_ratio", None) is not None
+        else None
+    )
+    loraplus_text_encoder_lr_ratio = (
+        float(kwargs.get("loraplus_text_encoder_lr_ratio", None))
+        if kwargs.get("loraplus_text_encoder_lr_ratio", None) is not None
+        else None
+    )
 
     if unbalanced_factorization:
         logger.info("Unbalanced factorization for LoKr is enabled")
@@ -121,8 +133,14 @@ def create_network(
         unbalanced_factorization=unbalanced_factorization,
         train_t5xxl=train_t5xxl,
     )
-    if loraplus_lr_ratio is not None or loraplus_unet_lr_ratio is not None or loraplus_text_encoder_lr_ratio is not None:
-        network.set_loraplus_lr_ratio(loraplus_lr_ratio, loraplus_unet_lr_ratio, loraplus_text_encoder_lr_ratio)
+    if (
+        loraplus_lr_ratio is not None
+        or loraplus_unet_lr_ratio is not None
+        or loraplus_text_encoder_lr_ratio is not None
+    ):
+        network.set_loraplus_lr_ratio(
+            loraplus_lr_ratio, loraplus_unet_lr_ratio, loraplus_text_encoder_lr_ratio
+        )
 
     return network
 
@@ -232,8 +250,8 @@ class LycorisNetworkKohya(LycorisNetwork):
         "DoubleStreamBlock",
         "SingleStreamBlock",
         "SingleDiTBlock",
-        "MMDoubleStreamBlock", #HunYuanVideo
-        "MMSingleStreamBlock", #HunYuanVideo
+        "MMDoubleStreamBlock",  # HunYuanVideo
+        "MMSingleStreamBlock",  # HunYuanVideo
     ]
     UNET_TARGET_REPLACE_NAME = [
         "conv_in",
@@ -301,7 +319,7 @@ class LycorisNetworkKohya(LycorisNetwork):
         self.multiplier = multiplier
         self.lora_dim = lora_dim
         self.train_t5xxl = train_t5xxl
-        
+
         # 初始化LoRA+相关属性
         self.loraplus_lr_ratio = None
         self.loraplus_unet_lr_ratio = None
@@ -621,15 +639,23 @@ class LycorisNetworkKohya(LycorisNetwork):
 
         return key_scaled, sum(norms) / len(norms), max(norms)
 
-    def set_loraplus_lr_ratio(self, loraplus_lr_ratio, loraplus_unet_lr_ratio, loraplus_text_encoder_lr_ratio):
+    def set_loraplus_lr_ratio(
+        self, loraplus_lr_ratio, loraplus_unet_lr_ratio, loraplus_text_encoder_lr_ratio
+    ):
         self.loraplus_lr_ratio = loraplus_lr_ratio
         self.loraplus_unet_lr_ratio = loraplus_unet_lr_ratio
         self.loraplus_text_encoder_lr_ratio = loraplus_text_encoder_lr_ratio
 
-        logger.info(f"LoRA+ UNet LR Ratio: {self.loraplus_unet_lr_ratio or self.loraplus_lr_ratio}")
-        logger.info(f"LoRA+ Text Encoder LR Ratio: {self.loraplus_text_encoder_lr_ratio or self.loraplus_lr_ratio}")
+        logger.info(
+            f"LoRA+ UNet LR Ratio: {self.loraplus_unet_lr_ratio or self.loraplus_lr_ratio}"
+        )
+        logger.info(
+            f"LoRA+ Text Encoder LR Ratio: {self.loraplus_text_encoder_lr_ratio or self.loraplus_lr_ratio}"
+        )
 
-    def prepare_optimizer_params(self, text_encoder_lr=None, unet_lr: float = 1e-4, learning_rate=None):
+    def prepare_optimizer_params(
+        self, text_encoder_lr=None, unet_lr: float = 1e-4, learning_rate=None
+    ):
         self.requires_grad_(True)
 
         all_params = []
@@ -658,7 +684,10 @@ class LycorisNetworkKohya(LycorisNetwork):
                     else:
                         param_data["lr"] = lr
 
-                if param_data.get("lr", None) == 0 or param_data.get("lr", None) is None:
+                if (
+                    param_data.get("lr", None) == 0
+                    or param_data.get("lr", None) is None
+                ):
                     logger.info("NO LR skipping!")
                     continue
 
@@ -674,7 +703,9 @@ class LycorisNetworkKohya(LycorisNetwork):
                 self.loraplus_text_encoder_lr_ratio or self.loraplus_lr_ratio,
             )
             all_params.extend(params)
-            lr_descriptions.extend(["textencoder" + (" " + d if d else "") for d in descriptions])
+            lr_descriptions.extend(
+                ["textencoder" + (" " + d if d else "") for d in descriptions]
+            )
 
         if self.unet_loras:
             params, descriptions = assemble_params(
@@ -683,7 +714,9 @@ class LycorisNetworkKohya(LycorisNetwork):
                 self.loraplus_unet_lr_ratio or self.loraplus_lr_ratio,
             )
             all_params.extend(params)
-            lr_descriptions.extend(["unet" + (" " + d if d else "") for d in descriptions])
+            lr_descriptions.extend(
+                ["unet" + (" " + d if d else "") for d in descriptions]
+            )
 
         return all_params, lr_descriptions
 
@@ -691,7 +724,7 @@ class LycorisNetworkKohya(LycorisNetwork):
         # not supported
         pass
 
-    def prepare_grad_etc(self, unet):
+    def prepare_grad_etc(self, *args):
         self.requires_grad_(True)
 
     def on_epoch_start(self, unet):

{lycoris_lora-3.2.0.dev1 → lycoris_lora-3.2.0.dev2}/lycoris/modules/base.py

@@ -290,6 +290,38 @@ class LycorisBaseModule(ModuleCustomSD):
                 self.org_module[0].bias = nn.Parameter(bias)
         self.to(self_device, self_dtype)
 
+    def onfly_merge(self, multiplier=1.0):
+        if self.not_supported:
+            return
+        self_device = next(self.parameters()).device
+        self_dtype = next(self.parameters()).dtype
+        self.to(self.org_weight)
+        self.cached_org_weight = self.org_weight.data.cpu()
+        self.cached_org_bias = None
+        weight, bias = self.get_merged_weight(
+            multiplier, self.org_weight.shape, self.org_weight.device
+        )
+        self.org_weight = weight
+        if bias is not None:
+            bias = bias.to(self.org_weight)
+            if self.org_module[0].bias is not None:
+                self.org_module[0].bias.data.copy_(bias)
+                self.cached_org_bias = self.org_module[0].bias.data.cpu()
+            else:
+                self.org_module[0].bias = nn.Parameter(bias)
+        if self.org_module[0].bias is not None:
+            self.org_module[0].bias = self.org_module[0].bias.to(self.org_weight)
+        self.to(self_device, self_dtype)
+
+    def onfly_restore(self):
+        if self.not_supported:
+            return
+        self.org_weight = self.cached_org_weight.to(self.org_weight)
+        if self.cached_org_bias is not None:
+            self.org_module[0].bias.data.copy_(self.cached_org_bias.to(self.org_weight))
+        del self.cached_org_weight
+        del self.cached_org_bias
+
     def get_diff_weight(self, multiplier=1.0, shape=None, device=None):
         raise NotImplementedError
 

{lycoris_lora-3.2.0.dev1 → lycoris_lora-3.2.0.dev2}/lycoris/wrapper.py

@@ -568,6 +568,14 @@ class LycorisNetwork(torch.nn.Module):
         for lora in self.loras:
             lora.merge_to(weight)
 
+    def onfly_merge(self, weight=1.0):
+        for lora in self.loras:
+            lora.onfly_merge(weight)
+
+    def onfly_restore(self):
+        for lora in self.loras:
+            lora.onfly_restore()
+
     def apply_max_norm_regularization(self, max_norm_value, device):
         key_scaled = 0
         norms = []

{lycoris_lora-3.2.0.dev1 → lycoris_lora-3.2.0.dev2/lycoris_lora.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.2
 Name: lycoris_lora
-Version: 3.2.0.dev1
+Version: 3.2.0.dev2
 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
@@ -23,3 +23,10 @@ Requires-Dist: torch
 Requires-Dist: einops
 Requires-Dist: toml
 Requires-Dist: tqdm
+Dynamic: author
+Dynamic: author-email
+Dynamic: classifier
+Dynamic: home-page
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary

{lycoris_lora-3.2.0.dev1 → lycoris_lora-3.2.0.dev2}/lycoris_lora.egg-info/SOURCES.txt

@@ -38,7 +38,9 @@ lycoris_lora.egg-info/not-zip-safe
 lycoris_lora.egg-info/requires.txt
 lycoris_lora.egg-info/top_level.txt
 test/__init__.py
+test/compile.py
 test/functional.py
 test/kohya.py
 test/module.py
+test/restore.py
 test/wrapper.py

{lycoris_lora-3.2.0.dev1 → lycoris_lora-3.2.0.dev2}/setup.py

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

lycoris_lora-3.2.0.dev2/test/compile.py

@@ -0,0 +1,69 @@
+import sys
+
+sys.setrecursionlimit(10000000)
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from torchao.quantization.quant_api import int8_weight_only, quantize_
+
+from lycoris import create_lycoris
+
+
+class TestNet(nn.Module):
+    def __init__(self):
+        super(TestNet, self).__init__()
+        self.fc1 = nn.Linear(8, 1024)
+        self.fc2 = nn.Linear(1024, 8)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = F.relu(x)
+        x = self.fc2(x)
+        return x
+
+
+def test_lycoris_compile():
+    net = TestNet().cuda()
+
+    test_x = torch.randn(1, 8).cuda()
+    out = net(test_x)
+    print("\nOriginal\n", out)
+
+    quantize_(net, int8_weight_only())
+    out = net(test_x)
+    print("\nQuantized\n", out)
+    compiled_net = torch.compile(net)
+    out = compiled_net(test_x)
+    print("\nQuantized compiled\n", out)
+
+    compiled_net._orig_mod.cpu()
+    compiled_net._orig_mod.cuda()
+
+    compiled_net.cpu()
+    compiled_net.cuda()
+
+    lycoris = create_lycoris(
+        net, algo="lokr", factor=2, bypass_mode=True, full_matrix=True
+    )
+    lycoris.apply_to()
+    lycoris.cuda()
+    out = net(test_x)
+    print("\nBypass Lokr Applied\n", out)
+
+    compiled_net = torch.compile(net)
+    compiled_out = compiled_net(test_x)
+    print("\nLyCORIS+torchao quantize_ compiled\n", compiled_out)
+    torch.sum(compiled_out).backward()
+
+    diff = torch.sum(torch.abs(out - compiled_out))
+    print("\nDiff: ", diff)
+    net.cpu()
+    lycoris.cpu()
+
+    assert torch.allclose(out, compiled_out, atol=1e-6)
+
+
+if __name__ == "__main__":
+    test_lycoris_compile()

lycoris_lora-3.2.0.dev2/test/restore.py

@@ -0,0 +1,56 @@
+import sys
+
+sys.setrecursionlimit(10000000)
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from lycoris import create_lycoris
+
+
+class TestNet(nn.Module):
+    def __init__(self):
+        super(TestNet, self).__init__()
+        self.fc1 = nn.Linear(8, 1024)
+        self.fc2 = nn.Linear(1024, 8)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = F.relu(x)
+        x = self.fc2(x)
+        return x
+
+
+def test_lycoris_restore():
+    net = TestNet().cuda()
+    compiled_net = torch.compile(net)
+
+    test_x = torch.randn(1, 8).cuda()
+    out = compiled_net(test_x)
+    print("\nOriginal\n", out)
+
+    lycoris = create_lycoris(
+        net, algo="lokr", factor=2, bypass_mode=True, full_matrix=True
+    )
+    lycoris.apply_to()
+    lycoris.cuda()
+    for parameter in lycoris.parameters():
+        parameter.data = torch.randn_like(parameter) * 0.1
+    net = torch.compile(net)
+    out = net(test_x)
+    print("\nBypass Lokr Applied\n", out)
+
+    lycoris.restore()
+    # lycoris.set_multiplier(0.0)
+    out = net(test_x)
+    print("\nRestored\n", out)
+
+    lycoris.apply_to()
+    # lycoris.set_multiplier(1.0)
+    out = net(test_x)
+    print("\nBypass Lokr Applied\n", out)
+
+
+if __name__ == "__main__":
+    test_lycoris_restore()