change for use_prompt_embedding param

Author: prishajain1

src/maxdiffusion/models/ltx2/transformer_ltx2.py

@@ -737,10 +737,25 @@ def __init__(
 
     # 2. Prompt embeddings
     if self.use_prompt_embeddings:
+      self.caption_projection = NNXPixArtAlphaTextProjection(
+          rngs=rngs,
+          in_features=self.caption_channels,
+          hidden_size=inner_dim,
+          dtype=self.dtype,
+          weights_dtype=self.weights_dtype,
+      )
+      self.audio_caption_projection = NNXPixArtAlphaTextProjection(
+          rngs=rngs,
+          in_features=self.caption_channels,
+          hidden_size=audio_inner_dim,
+          dtype=self.dtype,
+          weights_dtype=self.weights_dtype,
+      )
+    else:
       self.caption_projection = None
       self.audio_caption_projection = None
-      self.cross_attn_mod = True  # Force True for LTX-2.3 prompt modulation
 
+    if self.cross_attn_mod:
       self.prompt_adaln = LTX2AdaLayerNormSingle(
           rngs=rngs,
           embedding_dim=inner_dim,
@@ -757,9 +772,6 @@ def __init__(
           dtype=self.dtype,
           weights_dtype=self.weights_dtype,
       )
-    else:
-      self.caption_projection = None
-      self.audio_caption_projection = None
     # 3. Timestep Modulation Params and Embedding
     num_mod_params = 9 if self.cross_attn_mod else 6
     self.time_embed = LTX2AdaLayerNormSingle(

src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py

@@ -157,12 +157,13 @@ def create_model(rngs: nnx.Rngs, ltx2_config: dict):
   ltx2_config["remat_policy"] = config.remat_policy
   ltx2_config["names_which_can_be_saved"] = config.names_which_can_be_saved
   ltx2_config["names_which_can_be_offloaded"] = config.names_which_can_be_offloaded
+  ltx2_config["use_prompt_embeddings"] = True
 
   if getattr(config, "model_name", "") == "ltx2.3":
     ltx2_config["gated_attn"] = True
     ltx2_config["cross_attn_mod"] = True
     ltx2_config["perturbed_attn"] = True
-    ltx2_config["use_prompt_embeddings"] = True
+    ltx2_config["use_prompt_embeddings"] = False
 
   # 2. eval_shape
   p_model_factory = partial(create_model, ltx2_config=ltx2_config)