fix

Author: prishajain1

src/maxdiffusion/models/ltx2/ltx2_utils.py

@@ -305,6 +305,8 @@ def rename_for_ltx2_connector(key):
   key = key.replace("video_connector", "video_embeddings_connector")
   key = key.replace("audio_connector", "audio_embeddings_connector")
   key = key.replace("text_proj_in", "feature_extractor.linear")
+  key = key.replace("audio_feature_extractor.linear", "audio_text_proj_in")
+  key = key.replace("video_feature_extractor.linear", "video_text_proj_in")
 
   if "transformer_blocks" in key:
     key = key.replace("transformer_blocks", "stacked_blocks")

src/maxdiffusion/models/ltx2/text_encoders/text_encoders_ltx2.py

@@ -63,13 +63,15 @@ def __init__(
       proj_bias: bool = False,
       video_gated_attn: bool = False,
       audio_gated_attn: bool = False,
+      audio_hidden_dim: Optional[int] = None,
+      video_hidden_dim: Optional[int] = None,
       **kwargs,
   ):
     gemma_dim = 3840 if video_caption_channels is not None else caption_channels
     input_dim = gemma_dim * text_proj_in_factor
 
-    v_dim = video_caption_channels if video_caption_channels is not None else caption_channels
-    a_dim = audio_caption_channels if audio_caption_channels is not None else caption_channels
+    v_dim = video_hidden_dim if video_hidden_dim is not None else (video_caption_channels if video_caption_channels is not None else caption_channels)
+    a_dim = audio_hidden_dim if audio_hidden_dim is not None else (audio_caption_channels if audio_caption_channels is not None else caption_channels)
 
     self.per_modality_projections = per_modality_projections