transformer weight loading for scan false and timestepprojection issue

Author: prishajain1

src/maxdiffusion/models/ltx2/transformer_ltx2.py

@@ -698,15 +698,15 @@ def __init__(
       self.caption_projection = NNXCombinedTimestepTextProjEmbeddings(
           rngs=rngs,
           in_features=self.caption_channels,
-          hidden_size=inner_dim,
+          hidden_size=self.cross_attention_dim,
           embedding_dim=inner_dim,
           dtype=self.dtype,
           weights_dtype=self.weights_dtype,
       )
       self.audio_caption_projection = NNXCombinedTimestepTextProjEmbeddings(
           rngs=rngs,
           in_features=self.audio_caption_channels,
-          hidden_size=audio_inner_dim,
+          hidden_size=self.audio_cross_attention_dim,
           embedding_dim=audio_inner_dim,
           dtype=self.dtype,
           weights_dtype=self.weights_dtype,
@@ -1050,10 +1050,10 @@ def __call__(
       audio_cross_attn_v2a_gate = audio_cross_attn_v2a_gate.reshape(batch_size, -1, audio_cross_attn_v2a_gate.shape[-1])
 
       # 4. Prepare prompt embeddings
-      encoder_hidden_states = self.caption_projection(encoder_hidden_states)
+      encoder_hidden_states = self.caption_projection(encoder_hidden_states, timestep)
       encoder_hidden_states = encoder_hidden_states.reshape(batch_size, -1, hidden_states.shape[-1])
 
-      audio_encoder_hidden_states = self.audio_caption_projection(audio_encoder_hidden_states)
+      audio_encoder_hidden_states = self.audio_caption_projection(audio_encoder_hidden_states, audio_timestep if audio_timestep is not None else timestep)
       audio_encoder_hidden_states = audio_encoder_hidden_states.reshape(batch_size, -1, audio_hidden_states.shape[-1])
 
     # 5. Run transformer blocks