Fix

Author: Perseus14

src/maxdiffusion/models/lora_nnx.py

@@ -207,17 +207,27 @@ def merge_lora(model: nnx.Module, state_dict: dict, scale: float):
     """
     lora_params = {}
     for k, v in state_dict.items():
+      # Try matching diffusers rename format: "some.thing_lora.down.weight"
       m = re.match(r"^(.*?)_lora\.(down|up)\.weight$", k)
-      if not m:
-        m = re.match(r"^(.*?)\.lora\.(down|up)\.weight$", k)
-
       if m:
-        module_path_str, weight_type = m.group(1), m.group(2)
-        if module_path_str not in lora_params:
-          lora_params[module_path_str] = {}
-        lora_params[module_path_str][weight_type] = jnp.array(v)
+          module_path_str, weight_type = m.group(1), m.group(2)
       else:
-        max_logging.log(f"Could not parse LoRA key: {k}")
+          # Try matching diffusers format: "some.thing.lora.down.weight"
+          m = re.match(r"^(.*?)\.lora\.(down|up)\.weight$", k)
+          if m:
+              module_path_str, weight_type = m.group(1), m.group(2)
+          else:
+              # Try matching kohya/lightning format: "some.thing.lora_down.weight"
+              m = re.match(r"^(.*?)\.(lora_down|lora_up)\.weight$", k)
+              if m:
+                  module_path_str, weight_type = m.group(1), m.group(2).replace("lora_", "")
+              else:
+                  max_logging.debug(f"Could not parse LoRA key: {k}")
+                  continue
+
+      if module_path_str not in lora_params:
+          lora_params[module_path_str] = {}
+      lora_params[module_path_str][weight_type] = jnp.array(v)
 
     assigned_count = 0
     for path, module in nnx.iter_graph(model):
@@ -257,5 +267,5 @@ def merge_lora(model: nnx.Module, state_dict: dict, scale: float):
                           f"with kernel_size {module.kernel_size} > 1 is not supported."
                       )
           else:
-              max_logging.log(f"LoRA weights for {matched_key} incomplete.")
+              max_logging.warning(f"LoRA weights for {matched_key} incomplete.")
     max_logging.log(f"Merged weights into {assigned_count} layers in {type(model).__name__}.")