translateion add . after startswith

Author: tsvikas

src/maxdiffusion/loaders/ltx2_lora_nnx_loader.py

@@ -60,14 +60,14 @@ def translate_fn(nnx_path_str):
     if hasattr(pipeline, "transformer"):
       max_logging.log(f"Merging LoRA into transformer with rank={rank}")
       # 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")}
+      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.")
 
     if hasattr(pipeline, "connectors"):
       max_logging.log(f"Merging LoRA into connectors with rank={rank}")
-      connector_state_dict = {k: v for k, v in h_state_dict.items() if k.startswith("text_embedding_projection")}
+      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("connectors not found.")