logs parsing made better

Author: prishajain1

src/maxdiffusion/loaders/ltx2_lora_nnx_loader.py

@@ -54,7 +54,9 @@ def translate_fn(nnx_path_str):
     if hasattr(pipeline, "transformer") and transformer_weight_name:
       max_logging.log(f"Merging LoRA into transformer with rank={rank}")
       h_state_dict, _ = lora_loader.lora_state_dict(lora_model_path, weight_name=transformer_weight_name, **kwargs)
-      merge_fn(pipeline.transformer, h_state_dict, rank, scale, translate_fn, dtype=dtype)
+      # Filter state dict for transformer keys to avoid confusing warnings
+      transformer_state_dict = {k: v for k, v in h_state_dict.items() if k.startswith("diffusion_model")}
+      merge_fn(pipeline.transformer, transformer_state_dict, rank, scale, translate_fn, dtype=dtype)
     else:
       max_logging.log("transformer not found or no weight name provided for LoRA.")
 
@@ -64,7 +66,9 @@ def translate_fn(nnx_path_str):
         h_state_dict, _ = lora_loader.lora_state_dict(lora_model_path, weight_name=transformer_weight_name, **kwargs)
 
       if h_state_dict is not None:
-        merge_fn(pipeline.connectors, h_state_dict, rank, scale, translate_fn, dtype=dtype)
+        # Filter state dict for connector keys to avoid confusing warnings
+        connector_state_dict = {k: v for k, v in h_state_dict.items() if k.startswith("text_embedding_projection")}
+        merge_fn(pipeline.connectors, connector_state_dict, rank, scale, translate_fn, dtype=dtype)
       else:
         max_logging.log("Could not load LoRA state dict for connectors.")