distilled lora connector key fix

Author: prishajain1

src/maxdiffusion/loaders/lora_conversion_utils.py

@@ -786,6 +786,9 @@ def translate_ltx2_nnx_path_to_diffusers_lora(nnx_path_str, scan_layers=False):
       "caption_projection.linear_2": "diffusion_model.caption_projection.linear_2",
       "audio_caption_projection.linear_1": "diffusion_model.audio_caption_projection.linear_1",
       "audio_caption_projection.linear_2": "diffusion_model.audio_caption_projection.linear_2",
+      
+      # Connectors
+      "feature_extractor.linear": "text_embedding_projection.aggregate_embed",
   }
 
   if nnx_path_str in global_map:

src/maxdiffusion/loaders/ltx2_lora_nnx_loader.py

@@ -50,12 +50,22 @@ def load_lora_weights(
     def translate_fn(nnx_path_str):
       return lora_conversion_utils.translate_ltx2_nnx_path_to_diffusers_lora(nnx_path_str, scan_layers=scan_layers)
 
+    h_state_dict = None
     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)
-      # We assume keys match the translation function output.
       merge_fn(pipeline.transformer, h_state_dict, rank, scale, translate_fn, dtype=dtype)
     else:
       max_logging.log("transformer not found or no weight name provided for LoRA.")
 
+    if hasattr(pipeline, "connectors"):
+      max_logging.log(f"Merging LoRA into connectors with rank={rank}")
+      if h_state_dict is None and transformer_weight_name:
+        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)
+      else:
+        max_logging.log("Could not load LoRA state dict for connectors.")
+
     return pipeline