test

Author: Perseus14

src/maxdiffusion/models/lora_nnx.py

@@ -289,8 +289,15 @@ def merge_lora_for_scanned(
       continue
 
     nnx_path_str = ".".join(map(str, path))
+    nnx_path_str = nnx_path_str[len("blocks.") :]
     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: {module}")
+      max_logging.log(f"Module Type: {type(module)}")
+
     if not lora_key_template:
       continue