transformer weight

Author: prishajain1

src/maxdiffusion/models/ltx2/ltx2_utils.py

@@ -243,6 +243,8 @@ def load_transformer_weights(
         if not pt_key.startswith("model.diffusion_model."):
           continue
         pt_key = pt_key.replace("model.diffusion_model.", "")
+        if pt_key.startswith("audio_embeddings_connector") or pt_key.startswith("video_embeddings_connector"):
+          continue
 
       renamed_pt_key = rename_key(pt_key)
       renamed_pt_key = rename_for_ltx2_transformer(renamed_pt_key)

src/maxdiffusion/models/ltx2/transformer_ltx2.py

@@ -707,18 +707,19 @@ def __init__(
         weights_dtype=self.weights_dtype,
     )
     # 3. Timestep Modulation Params and Embedding
+    num_mod_params = 9 if self.cross_attn_mod else 6
     self.time_embed = LTX2AdaLayerNormSingle(
         rngs=rngs,
         embedding_dim=inner_dim,
-        num_mod_params=6,
+        num_mod_params=num_mod_params,
         use_additional_conditions=False,
         dtype=self.dtype,
         weights_dtype=self.weights_dtype,
     )
     self.audio_time_embed = LTX2AdaLayerNormSingle(
         rngs=rngs,
         embedding_dim=audio_inner_dim,
-        num_mod_params=6,
+        num_mod_params=num_mod_params,
         use_additional_conditions=False,
         dtype=self.dtype,
         weights_dtype=self.weights_dtype,