chnage + testing for distilled loras

Author: prishajain1

src/maxdiffusion/configs/ltx2_video.yml

@@ -107,9 +107,9 @@ audio_format: "s16"
 # LoRA parameters
 enable_lora: False
 lora_config: {
-  lora_model_name_or_path: ["Lightricks/LTX-2-19b-IC-LoRA-Depth-Control", "Lightricks/LTX-2-19b-LoRA-Camera-Control-Dolly-Out"],
-  weight_name: ["ltx-2-19b-ic-lora-depth-control.safetensors","ltx-2-19b-lora-camera-control-dolly-out.safetensors"],
-  adapter_name: ["Depth-control","Camera-control-dolly-out"],
-  rank: [64, 32]
+  lora_model_name_or_path: ["Lightricks/LTX-2"],
+  weight_name: ["ltx-2-19b-distilled-lora-384.safetensors"],
+  adapter_name: ["distilled-lora-384"],
+  rank: [384]
 }
 

src/maxdiffusion/loaders/lora_conversion_utils.py

@@ -757,6 +757,22 @@ def translate_ltx2_nnx_path_to_diffusers_lora(nnx_path_str, scan_layers=False):
   }
 
   # --- 3. Translation Logic ---
+  global_map = {
+      "proj_in": "diffusion_model.patchify_proj",
+      "audio_proj_in": "diffusion_model.audio_patchify_proj",
+      "proj_out": "diffusion_model.proj_out",
+      "audio_proj_out": "diffusion_model.audio_proj_out",
+      "time_embed.linear": "diffusion_model.adaln_single.linear",
+      "audio_time_embed.linear": "diffusion_model.audio_adaln_single.linear",
+      "av_cross_attn_video_a2v_gate.linear": "diffusion_model.av_ca_a2v_gate_adaln_single.linear",
+      "av_cross_attn_audio_v2a_gate.linear": "diffusion_model.av_ca_v2a_gate_adaln_single.linear",
+      "av_cross_attn_audio_scale_shift.linear": "diffusion_model.av_ca_audio_scale_shift_adaln_single.linear",
+      "av_cross_attn_video_scale_shift.linear": "diffusion_model.av_ca_video_scale_shift_adaln_single.linear",
+  }
+  
+  if nnx_path_str in global_map:
+    return global_map[nnx_path_str]
+
   if scan_layers:
     if nnx_path_str.startswith("transformer_blocks."):
       inner_suffix = nnx_path_str[len("transformer_blocks.") :]