test

Author: Perseus14

src/maxdiffusion/loaders/lora_conversion_utils.py

@@ -651,7 +651,7 @@ def translate_wan_nnx_path_to_diffusers_lora(nnx_path_str, scan_layers=False):
     return "diffusion_model.patch_embedding"
   if nnx_path_str == "proj_out":
     return "diffusion_model.head.head"
-  if nnx_path_str == "adaln_scale_shift_table":
+  if nnx_path_str == "scale_shift_table":
     return "diffusion_model.head.modulation"
   if nnx_path_str == "condition_embedder.time_proj":
     return "diffusion_model.time_projection.1"

src/maxdiffusion/models/lora_nnx.py

@@ -295,12 +295,6 @@ def merge_lora_for_scanned(
     nnx_path_str = ".".join(map(str, path))
     lora_key_template = translate_fn(nnx_path_str) if translate_fn else None
 
-    if 'adaln_scale_shift_table' in nnx_path_str:
-      max_logging.log(f"adaln_scale_shift_table: {nnx_path_str}")
-      max_logging.log(f"lora_key_template: {lora_key_template}")
-      max_logging.log(f"Module ndim: {module.ndim}")
-      max_logging.log(f"Module Type: {type(module)}")
-
     if not lora_key_template:
       continue
 
@@ -315,7 +309,8 @@ def merge_lora_for_scanned(
     elif isinstance(module, nnx.Conv):
       is_scanned = module.kernel.ndim == 5
     elif isinstance(module, nnx.Param):
-      is_scanned = module.ndim > 1  # Heuristic for scanned Param
+      # Use template format to disambiguate: if template has {}, then it is scanned.
+      is_scanned = ("{}" in lora_key_template)
 
     if not is_scanned:
       lora_key = lora_key_template