changes to transformer and test

Author: prishajain1

src/maxdiffusion/models/ltx2/transformer_ltx2.py

@@ -704,11 +704,11 @@ def init_block(rngs):
                 dim=inner_dim,
                 num_attention_heads=self.num_attention_heads,
                 attention_head_dim=self.attention_head_dim,
-                cross_attention_dim=self.cross_attention_dim,
+                cross_attention_dim=inner_dim,
                 audio_dim=audio_inner_dim,
                 audio_num_attention_heads=self.audio_num_attention_heads,
                 audio_attention_head_dim=self.audio_attention_head_dim,
-                audio_cross_attention_dim=self.audio_cross_attention_dim,
+                audio_cross_attention_dim=audio_inner_dim,
                 activation_fn=self.activation_fn,
                 qk_norm=self.qk_norm,
                 attention_bias=self.attention_bias,
@@ -735,11 +735,11 @@ def init_block(rngs):
                     dim=inner_dim,
                     num_attention_heads=self.num_attention_heads,
                     attention_head_dim=self.attention_head_dim,
-                    cross_attention_dim=self.cross_attention_dim,
+                    cross_attention_dim=inner_dim,
                     audio_dim=audio_inner_dim,
                     audio_num_attention_heads=self.audio_num_attention_heads,
                     audio_attention_head_dim=self.audio_attention_head_dim,
-                    audio_cross_attention_dim=self.audio_cross_attention_dim,
+                    audio_cross_attention_dim=audio_inner_dim,
                     activation_fn=self.activation_fn,
                     qk_norm=self.qk_norm,
                     attention_bias=self.attention_bias,

src/maxdiffusion/tests/ltx_2_transformer_test.py

@@ -91,7 +91,7 @@ def test_transformer_block_shapes(self):
 
             # Create dummy inputs
             hidden_states = jnp.zeros((self.batch_size, self.seq_len, dim))
-            audio_hidden_states = jnp.zeros((self.batch_size, 10, audio_dim)) # 10 audio frames
+            audio_hidden_states = jnp.zeros((self.batch_size, 128, audio_dim)) # 128 audio frames for TPFA
             encoder_hidden_states = jnp.zeros((self.batch_size, 5, cross_dim))
             audio_encoder_hidden_states = jnp.zeros((self.batch_size, 5, cross_dim)) # reusing cross_dim for audio context 
 
@@ -221,7 +221,7 @@ def test_transformer_3d_model_instantiation_and_forward(self):
 
         # Let's pass (B, L, C).
         hidden_states = jnp.zeros((self.batch_size, self.seq_len, self.in_channels))
-        audio_hidden_states = jnp.zeros((self.batch_size, 10, self.audio_in_channels))
+        audio_hidden_states = jnp.zeros((self.batch_size, 128, self.audio_in_channels))
 
         timestep = jnp.array([1.0]) # (B,)
 
@@ -241,7 +241,9 @@ def test_transformer_3d_model_instantiation_and_forward(self):
                 num_frames=self.num_frames,
                 height=self.height,
                 width=self.width,
-                audio_num_frames=10,
+                width=self.width,
+                audio_num_frames=128,
+                fps=24.0,
                 fps=24.0,
                 return_dict=True,
                 encoder_attention_mask=encoder_attention_mask,
@@ -255,7 +257,8 @@ def test_transformer_3d_model_instantiation_and_forward(self):
         print(f"Model Output Audio Shape: {audio_sample.shape}")
 
         self.assertEqual(sample.shape, (self.batch_size, self.seq_len, self.out_channels))
-        self.assertEqual(audio_sample.shape, (self.batch_size, 10, self.audio_in_channels))
+        self.assertEqual(sample.shape, (self.batch_size, self.seq_len, self.out_channels))
+        self.assertEqual(audio_sample.shape, (self.batch_size, 128, self.audio_in_channels))
 
     def test_scan_remat_parity(self):
         """
@@ -300,7 +303,7 @@ def test_scan_remat_parity(self):
 
         # Inputs
         hidden_states = jnp.ones((self.batch_size, self.seq_len, self.in_channels)) * 0.5
-        audio_hidden_states = jnp.ones((self.batch_size, 10, self.audio_in_channels)) * 0.5
+        audio_hidden_states = jnp.ones((self.batch_size, 128, self.audio_in_channels)) * 0.5
         timestep = jnp.array([1.0])
         encoder_hidden_states = jnp.ones((self.batch_size, 5, 32)) * 0.1
         audio_encoder_hidden_states = jnp.ones((self.batch_size, 5, 32)) * 0.1
@@ -314,7 +317,10 @@ def test_scan_remat_parity(self):
             num_frames=self.num_frames,
             height=self.height,
             width=self.width,
-            audio_num_frames=10,
+            height=self.height,
+            width=self.width,
+            audio_num_frames=128,
+            fps=24.0,
             fps=24.0,
             return_dict=True,
         )