lycoris-lora 3.1.1.post1__tar.gz → 3.2.0.dev2__tar.gz

{lycoris_lora-3.1.1.post1/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.1.1.post1
+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.1.1.post1 → lycoris_lora-3.2.0.dev2}/lycoris/config.py

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

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

@@ -63,6 +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
+    )
 
     if unbalanced_factorization:
         logger.info("Unbalanced factorization for LoKr is enabled")
@@ -117,6 +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
+        )
 
     return network
 
@@ -157,22 +181,23 @@ def create_network_from_weights(
         if lora_name in unet_loras:
             unet_loras[lora_name] = modules
 
-    if isinstance(text_encoder, list):
-        text_encoders = text_encoder
-        use_index = True
-    else:
-        text_encoders = [text_encoder]
-        use_index = False
-
-    for idx, te in enumerate(text_encoders):
-        if use_index:
-            prefix = f"{LycorisNetworkKohya.LORA_PREFIX_TEXT_ENCODER}{idx+1}"
+    if text_encoder:
+        if isinstance(text_encoder, list):
+            text_encoders = text_encoder
+            use_index = True
         else:
-            prefix = LycorisNetworkKohya.LORA_PREFIX_TEXT_ENCODER
-        for name, modules in te.named_modules():
-            lora_name = f"{prefix}_{name}".replace(".", "_")
-            if lora_name in te_loras:
-                te_loras[lora_name] = modules
+            text_encoders = [text_encoder]
+            use_index = False
+
+        for idx, te in enumerate(text_encoders):
+            if use_index:
+                prefix = f"{LycorisNetworkKohya.LORA_PREFIX_TEXT_ENCODER}{idx+1}"
+            else:
+                prefix = LycorisNetworkKohya.LORA_PREFIX_TEXT_ENCODER
+            for name, modules in te.named_modules():
+                lora_name = f"{prefix}_{name}".replace(".", "_")
+                if lora_name in te_loras:
+                    te_loras[lora_name] = modules
 
     original_level = logger.level
     logger.setLevel(logging.ERROR)
@@ -192,14 +217,16 @@ def create_network_from_weights(
     logger.info(f"{len(network.unet_loras)} Modules Loaded")
 
     logger.info("Loading TE Modules from state dict...")
-    for lora_name, orig_modules in te_loras.items():
-        if orig_modules is None:
-            continue
-        lyco_type, params = get_module(weights_sd, lora_name)
-        module = make_module(lyco_type, params, lora_name, orig_modules)
-        if module is not None:
-            network.text_encoder_loras.append(module)
-    logger.info(f"{len(network.text_encoder_loras)} Modules Loaded")
+
+    if text_encoder:
+        for lora_name, orig_modules in te_loras.items():
+            if orig_modules is None:
+                continue
+            lyco_type, params = get_module(weights_sd, lora_name)
+            module = make_module(lyco_type, params, lora_name, orig_modules)
+            if module is not None:
+                network.text_encoder_loras.append(module)
+        logger.info(f"{len(network.text_encoder_loras)} Modules Loaded")
 
     for lora in network.unet_loras + network.text_encoder_loras:
         lora.multiplier = multiplier
@@ -223,6 +250,8 @@ class LycorisNetworkKohya(LycorisNetwork):
         "DoubleStreamBlock",
         "SingleStreamBlock",
         "SingleDiTBlock",
+        "MMDoubleStreamBlock",  # HunYuanVideo
+        "MMSingleStreamBlock",  # HunYuanVideo
     ]
     UNET_TARGET_REPLACE_NAME = [
         "conv_in",
@@ -291,6 +320,11 @@ class LycorisNetworkKohya(LycorisNetwork):
         self.lora_dim = lora_dim
         self.train_t5xxl = train_t5xxl
 
+        # 初始化LoRA+相关属性
+        self.loraplus_lr_ratio = None
+        self.loraplus_unet_lr_ratio = None
+        self.loraplus_text_encoder_lr_ratio = None
+
         if not self.ENABLE_CONV:
             conv_lora_dim = 0
 
@@ -460,27 +494,28 @@ class LycorisNetworkKohya(LycorisNetwork):
             ]
             LycorisNetworkKohya.UNET_TARGET_REPLACE_NAME = []
 
-        if isinstance(text_encoder, list):
-            text_encoders = text_encoder
-            use_index = True
-        else:
-            text_encoders = [text_encoder]
-            use_index = False
-
         self.text_encoder_loras = []
-        for i, te in enumerate(text_encoders):
-            self.text_encoder_loras.extend(
-                create_modules(
-                    LycorisNetworkKohya.LORA_PREFIX_TEXT_ENCODER
-                    + (f"{i+1}" if use_index else ""),
-                    te,
-                    LycorisNetworkKohya.TEXT_ENCODER_TARGET_REPLACE_MODULE,
-                    LycorisNetworkKohya.TEXT_ENCODER_TARGET_REPLACE_NAME,
+        if text_encoder:
+            if isinstance(text_encoder, list):
+                text_encoders = text_encoder
+                use_index = True
+            else:
+                text_encoders = [text_encoder]
+                use_index = False
+
+            for i, te in enumerate(text_encoders):
+                self.text_encoder_loras.extend(
+                    create_modules(
+                        LycorisNetworkKohya.LORA_PREFIX_TEXT_ENCODER
+                        + (f"{i+1}" if use_index else ""),
+                        te,
+                        LycorisNetworkKohya.TEXT_ENCODER_TARGET_REPLACE_MODULE,
+                        LycorisNetworkKohya.TEXT_ENCODER_TARGET_REPLACE_NAME,
+                    )
                 )
+            logger.info(
+                f"create LyCORIS for Text Encoder: {len(self.text_encoder_loras)} modules."
             )
-        logger.info(
-            f"create LyCORIS for Text Encoder: {len(self.text_encoder_loras)} modules."
-        )
 
         self.unet_loras = create_modules(
             LycorisNetworkKohya.LORA_PREFIX_UNET,
@@ -604,29 +639,102 @@ class LycorisNetworkKohya(LycorisNetwork):
 
         return key_scaled, sum(norms) / len(norms), max(norms)
 
-    def prepare_optimizer_params(self, text_encoder_lr, unet_lr, learning_rate):
-        def enumerate_params(loras):
-            params = []
-            for lora in loras:
-                params.extend(lora.parameters())
-            return params
+    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}"
+        )
 
+    def prepare_optimizer_params(
+        self, text_encoder_lr=None, unet_lr: float = 1e-4, learning_rate=None
+    ):
         self.requires_grad_(True)
+
         all_params = []
+        lr_descriptions = []
+
+        def assemble_params(loras, lr, ratio):
+            param_groups = {"lora": {}, "plus": {}}
+            for lora in loras:
+                for name, param in lora.named_parameters():
+                    if ratio is not None and "lora_up" in name:
+                        param_groups["plus"][f"{lora.lora_name}.{name}"] = param
+                    else:
+                        param_groups["lora"][f"{lora.lora_name}.{name}"] = param
+
+            params = []
+            descriptions = []
+            for key in param_groups.keys():
+                param_data = {"params": param_groups[key].values()}
+
+                if len(param_data["params"]) == 0:
+                    continue
+
+                if lr is not None:
+                    if key == "plus":
+                        param_data["lr"] = lr * ratio
+                    else:
+                        param_data["lr"] = lr
+
+                if (
+                    param_data.get("lr", None) == 0
+                    or param_data.get("lr", None) is None
+                ):
+                    logger.info("NO LR skipping!")
+                    continue
+
+                params.append(param_data)
+                descriptions.append("plus" if key == "plus" else "")
+
+            return params, descriptions
 
         if self.text_encoder_loras:
-            param_data = {"params": enumerate_params(self.text_encoder_loras)}
-            if text_encoder_lr is not None:
-                param_data["lr"] = text_encoder_lr
-            all_params.append(param_data)
+            params, descriptions = assemble_params(
+                self.text_encoder_loras,
+                text_encoder_lr if text_encoder_lr is not None else default_lr,
+                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]
+            )
 
         if self.unet_loras:
-            param_data = {"params": enumerate_params(self.unet_loras)}
-            if unet_lr is not None:
-                param_data["lr"] = unet_lr
-            all_params.append(param_data)
+            params, descriptions = assemble_params(
+                self.unet_loras,
+                unet_lr if unet_lr is not None else default_lr,
+                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]
+            )
+
+        return all_params, lr_descriptions
+
+    def enable_gradient_checkpointing(self):
+        # not supported
+        pass
+
+    def prepare_grad_etc(self, *args):
+        self.requires_grad_(True)
+
+    def on_epoch_start(self, unet):
+        self.train()
+
+    def on_step_start(self):
+        pass
 
-        return all_params
+    def get_trainable_params(self):
+        return self.parameters()
 
     def save_weights(self, file, dtype, metadata):
         if metadata is not None and len(metadata) == 0:

{lycoris_lora-3.1.1.post1 → 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.1.1.post1 → 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.1.1.post1 → 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.1.1.post1
+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.1.1.post1 → 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.1.1.post1 → 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.1.1.post1",
+    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()