Lora

diffsynth/core/__init__.py

@@ -6,3 +6,4 @@
 from .vram import *
 from .device import *
 from .offload_training import *
+from .lora_train import *

diffsynth/core/lora_train/__init__.py

@@ -0,0 +1 @@
+from .patch import LoRALinear, inject_lora_into_model

diffsynth/core/lora_train/patch.py

@@ -0,0 +1,60 @@
+import math, re, torch
+
+
+class LoRALinear(torch.nn.Module):
+    def __init__(self, base_layer: torch.nn.Linear, lora_rank: int, lora_alpha: float = None, dtype: torch.dtype = None, lora_bias: bool = False):
+        super().__init__()
+        if lora_alpha is None:
+            lora_alpha = lora_rank
+        self.base_layer = base_layer
+        self.lora_rank = lora_rank
+        self.lora_alpha = lora_alpha
+        self.scaling = lora_alpha / lora_rank
+        weight = base_layer.weight
+        dtype = dtype or (weight.dtype if weight.dtype in (torch.float32, torch.float16, torch.bfloat16) else torch.float32)
+        self.lora_A = torch.nn.Linear(base_layer.in_features, lora_rank, bias=False)
+        self.lora_B = torch.nn.Linear(lora_rank, base_layer.out_features, bias=lora_bias)
+        torch.nn.init.kaiming_uniform_(self.lora_A.weight, a=math.sqrt(5))
+        torch.nn.init.zeros_(self.lora_B.weight)
+        if lora_bias:
+            torch.nn.init.zeros_(self.lora_B.bias)
+        self.lora_A.to(device=weight.device, dtype=dtype)
+        self.lora_B.to(device=weight.device, dtype=dtype)
+
+    def forward(self, x, *args, **kwargs):
+        out = self.base_layer(x, *args, **kwargs)
+        out_dtype = out.dtype
+        out = out + self.lora_B(self.lora_A(x.to(self.lora_A.weight.dtype))) * self.scaling
+        return out.to(out_dtype)
+
+
+def match_lora_target_module(name, target_modules):
+    if isinstance(target_modules, str):
+        return re.fullmatch(target_modules, name) is not None
+    return any(name == target or name.endswith("." + target) for target in target_modules)
+
+
+def inject_lora_into_model(model: torch.nn.Module, target_modules, lora_rank: int, lora_alpha: float = None, dtype: torch.dtype = None):
+    replaced_names, skipped_names = [], []
+    for name, module in list(model.named_modules()):
+        if name == "" or not match_lora_target_module(name, target_modules):
+            continue
+        if not isinstance(module, torch.nn.Linear):
+            skipped_names.append(f"{name} ({type(module).__name__})")
+            continue
+        parent_name, _, child_name = name.rpartition(".")
+        parent = model.get_submodule(parent_name) if parent_name != "" else model
+        setattr(parent, child_name, LoRALinear(module, lora_rank, lora_alpha=lora_alpha, dtype=dtype))
+        replaced_names.append(name)
+    if len(skipped_names) > 0:
+        print(f"These matched modules are not `torch.nn.Linear`, so LoRA is not patched on them: {skipped_names}.")
+    if len(replaced_names) == 0:
+        raise ValueError(f"No `torch.nn.Linear` module matches the LoRA target modules: {target_modules}.")
+    for param in model.parameters():
+        param.requires_grad = False
+    for module in model.modules():
+        if isinstance(module, LoRALinear):
+            module.lora_A.requires_grad_(True)
+            module.lora_B.requires_grad_(True)
+    print(f"LoRA is patched on {len(replaced_names)} modules.")
+    return model

diffsynth/diffusion/training_module.py

@@ -1,5 +1,5 @@
 import torch, json, os, inspect
-from ..core import ModelConfig, load_state_dict, QuantizeConfig
+from ..core import ModelConfig, load_state_dict, QuantizeConfig, inject_lora_into_model
 from ..utils.controlnet import ControlNetInput
 from .base_pipeline import PipelineUnit
 
@@ -91,27 +91,14 @@ def trainable_param_names(self):
 
 
     def add_lora_to_model(self, model, target_modules, lora_rank, lora_alpha=None, upcast_dtype=None):
-        from peft import LoraConfig, inject_adapter_in_model
-        if lora_alpha is None:
-            lora_alpha = lora_rank
-        if isinstance(target_modules, list) and len(target_modules) == 1:
-            target_modules = target_modules[0]
-        lora_config = LoraConfig(r=lora_rank, lora_alpha=lora_alpha, target_modules=target_modules)
-        model = inject_adapter_in_model(lora_config, model)
-        if upcast_dtype is not None:
-            for param in model.parameters():
-                if param.requires_grad:
-                    param.data = param.to(upcast_dtype)
-        return model
+        return inject_lora_into_model(model, target_modules, lora_rank, lora_alpha=lora_alpha, dtype=upcast_dtype)
 
 
     def mapping_lora_state_dict(self, state_dict):
         new_state_dict = {}
         for key, value in state_dict.items():
-            if "lora_A.weight" in key or "lora_B.weight" in key:
-                new_key = key.replace("lora_A.weight", "lora_A.default.weight").replace("lora_B.weight", "lora_B.default.weight")
-                new_state_dict[new_key] = value
-            elif "lora_A.default.weight" in key or "lora_B.default.weight" in key:
+            key = key.replace("lora_A.default.weight", "lora_A.weight").replace("lora_B.default.weight", "lora_B.weight")
+            if key.endswith("lora_A.weight") or key.endswith("lora_B.weight"):
                 new_state_dict[key] = value
         return new_state_dict
 

diffsynth/models/krea2_dit.py

@@ -85,10 +85,6 @@ class SingleMMDiTConfig:
     txtheads: int = 20
     txtkvheads: int = 20
 
-    def get(self, *args, **kwargs):
-        # For compatibility with low-version peft
-        return None
-
 
 class SimpleModulation(torch.nn.Module):
     def __init__(self, dim: int):

diffsynth/models/qwen_image_image2lora.py

@@ -66,8 +66,8 @@ def forward(self, x, residual=None):
         for lora_pattern, layer in zip(self.lora_patterns, self.layers):
             name = lora_pattern[0]
             lora_a, lora_b = layer(x, residual=residual)
-            lora[f"transformer_blocks.{self.block_id}.{name}.lora_A.default.weight"] = lora_a
-            lora[f"transformer_blocks.{self.block_id}.{name}.lora_B.default.weight"] = lora_b
+            lora[f"transformer_blocks.{self.block_id}.{name}.lora_A.weight"] = lora_a
+            lora[f"transformer_blocks.{self.block_id}.{name}.lora_B.weight"] = lora_b
         return lora
 
 

diffsynth/models/z_image_image2lora.py

@@ -23,8 +23,8 @@ def forward(self, x, residual=None):
         for lora_pattern, layer in zip(self.lora_patterns, self.layers):
             name = lora_pattern[0]
             lora_a, lora_b = layer(x, residual=residual)
-            lora[f"{self.prefix}.{self.block_id}.{name}.lora_A.default.weight"] = lora_a
-            lora[f"{self.prefix}.{self.block_id}.{name}.lora_B.default.weight"] = lora_b
+            lora[f"{self.prefix}.{self.block_id}.{name}.lora_A.weight"] = lora_a
+            lora[f"{self.prefix}.{self.block_id}.{name}.lora_B.weight"] = lora_b
         return lora
 
 
@@ -170,7 +170,7 @@ def forward(self, x, residual=None):
         lora = {}
         for name, module in self.module_dict.items():
             name = name.replace("___", ".")
-            name_a, name_b = f"{name}.lora_A.default.weight", f"{name}.lora_B.default.weight"
+            name_a, name_b = f"{name}.lora_A.weight", f"{name}.lora_B.weight"
             lora_a, lora_b = module(x)
             lora[name_a] = lora_a
             lora[name_b] = lora_b

diffsynth/utils/lora/krea2.py

@@ -108,10 +108,10 @@ def align_to_opensource_format(state_dict, alpha=None):
         rename_dict = Krea2LoRAConverter().rename_dict
         state_dict_ = {}
         for name in state_dict:
-            weight_name = name.replace(".lora_A.default.", ".").replace(".lora_B.default.", ".")
-            if name.endswith(".lora_A.default.weight"):
+            weight_name = name.replace(".lora_A.", ".").replace(".lora_B.", ".")
+            if name.endswith(".lora_A.weight"):
                 name_ = "transformer." + rename_dict[weight_name].replace(".weight", ".lora_A.weight")
-            elif name.endswith(".lora_B.default.weight"):
+            elif name.endswith(".lora_B.weight"):
                 name_ = "transformer." + rename_dict[weight_name].replace(".weight", ".lora_B.weight")
             state_dict_[name_] = state_dict[name]
         return state_dict_
@@ -124,8 +124,8 @@ def align_to_diffsynth_format(state_dict):
         for name in state_dict:
             weight_name = name.replace(".lora_A.", ".").replace(".lora_B.", ".").replace("transformer.", "")
             if name.endswith(".lora_A.weight"):
-                name_ = rename_dict[weight_name].replace(".weight", ".lora_A.default.weight")
+                name_ = rename_dict[weight_name].replace(".weight", ".lora_A.weight")
             elif name.endswith(".lora_B.weight"):
-                name_ = rename_dict[weight_name].replace(".weight", ".lora_B.default.weight")
+                name_ = rename_dict[weight_name].replace(".weight", ".lora_B.weight")
             state_dict_[name_] = state_dict[name]
         return state_dict_