text encoder used from hf

Author: prishajain1

src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py

@@ -402,13 +402,7 @@ def _get_gemma_prompt_embeds(
            text_encoder_hidden_states = jnp.array(text_encoder_hidden_states.cpu().numpy())
            prompt_attention_mask = jnp.array(prompt_attention_mask.cpu().numpy())
       else:
-          # Mock hidden states
-          # Should be removed once we have actual text_encoder ready to port
-          hidden_dim = 1024
-          num_layers = 2
-          text_encoder_hidden_states = jnp.zeros(
-              (batch_size, max_sequence_length, hidden_dim, num_layers), dtype=dtype or jnp.float32
-          )
+          raise ValueError("`text_encoder` is required to encode prompts.")
 
       sequence_lengths = prompt_attention_mask.sum(axis=-1)
 
@@ -605,28 +599,6 @@ def _create_noised_state(
   @staticmethod
   def _pack_audio_latents(
       latents: jax.Array, patch_size: Optional[int] = None, patch_size_t: Optional[int] = None
-  ) -> jax.Array:
-      if patch_size is not None and patch_size_t is not None:
-          batch_size, num_channels, latent_length, latent_mel_bins = latents.shape
-          post_patch_latent_length = latent_length // patch_size_t
-          post_patch_mel_bins = latent_mel_bins // patch_size
-          latents = latents.reshape(
-              batch_size, -1, post_patch_latent_length, patch_size_t, post_patch_mel_bins, patch_size
-          )
-          latents = latents.transpose(0, 2, 4, 1, 3, 5).reshape(batch_size, post_patch_latent_length * post_patch_mel_bins, -1)
-      else:
-          latents = latents.transpose(0, 2, 1).reshape(batch_size, latents.shape[2], -1) 
-          # Wait, original was transpose(1,2).flatten(2,3) -> (Batch, Channels, Length) -> (Batch, Length, Channels)?
-          # Diffusers: latents = latents.transpose(1, 2).flatten(2, 3) 
-          # (B, C, L) -> (B, L, C). 
-          # If 4D: (B, C, L, M) -> (B, C, L, P_t, M, P) -> ...
-          pass
-      return latents
-      
-  # Redefining _pack_audio_latents properly for JAX
-  @staticmethod
-  def _pack_audio_latents_jax(
-      latents: jax.Array, patch_size: Optional[int] = None, patch_size_t: Optional[int] = None
   ) -> jax.Array:
       if patch_size is not None and patch_size_t is not None:
           batch_size, num_channels, latent_length, latent_mel_bins = latents.shape

src/maxdiffusion/tests/ltx2_pipeline_test.py

@@ -285,18 +285,39 @@ def test_check_inputs(self):
                  negative_prompt_embeds=jnp.zeros((1, 5, 64)), # Mismatch length
                  negative_prompt_attention_mask=jnp.ones((1, 5))
              )
-        rngs = nnx.Rngs(0)
+    
+    def test_audio_packing_unpacking(self):
+        # (Batch, Channels, Length, Mel)
+        batch_size = 1
+        channels = 128
+        length = 32
+        mel = 64
+        patch_size = 4
+        patch_size_t = 1 # Audio typically has patch_size_t=1 in LTX logic, let's test that
 
-        pipeline = LTX2Pipeline.load_transformer(
-            devices_array=jnp.array(jax.devices()),
-            mesh=self.mesh,
-            rngs=rngs,
-            config=config,
-            subfolder="transformer"
+        latents = jax.random.normal(jax.random.key(0), (batch_size, channels, length, mel))
+        
+        packed = LTX2Pipeline._pack_audio_latents(latents, patch_size=patch_size, patch_size_t=patch_size_t)
+        
+        # Verify packed shape
+        # original logic: (B, T', F', C, p_t, p) -> (B, T' * F', -1)
+        # T' = 32 // 1 = 32
+        # F' = 64 // 4 = 16
+        # shape should be (1, 32 * 16, 128 * 1 * 4) = (1, 512, 512)
+        expected_seq_len = (length // patch_size_t) * (mel // patch_size)
+        expected_dim = channels * patch_size * patch_size_t
+        self.assertEqual(packed.shape, (batch_size, expected_seq_len, expected_dim))
+        
+        unpacked = LTX2Pipeline._unpack_audio_latents(
+            packed, 
+            latent_length=length, 
+            num_mel_bins=mel, 
+            patch_size=patch_size, 
+            patch_size_t=patch_size_t
         )
 
-        mock_create.assert_called_once()
-        self.assertEqual(pipeline, mock_create.return_value)
+        self.assertEqual(unpacked.shape, latents.shape)
+        np.testing.assert_allclose(unpacked, latents, atol=1e-6)
 
 if __name__ == "__main__":
     unittest.main()