Refactor LTX2 text encoders: replace Video/AV classes with unified EmbeddingsProcessor; move tests to tests/ltx2/

Signed-off-by: James Huang <syhuang1201@gmail.com>

Author: syhuang22

src/maxdiffusion/models/ltx2/text_encoders/text_encoders_ltx2.py

@@ -27,85 +27,29 @@
 DType = common_types.DType
 
 
-class LTX2VideoGemmaTextEncoder(nnx.Module):
+class LTX2EmbeddingsProcessor(nnx.Module):
   """
-  Encoder for Video-only tasks.
-  Pipeline: Gemma Hidden States -> Feature Extractor -> Video Connector -> Output
-  """
-
-  def __init__(
-      self,
-      # Feature Extractor Config
-      gemma_dim: int = 3840,  # Gemma-3-12b
-      gemma_layers: int = 49,  # Gemma-3 has 48 layers + 1 embedding layer output = 49 hidden states
-      projection_dim: int = 3840,  # LTX-2 conditioning dim
-      # Connector Config
-      connector_heads: int = 32,
-      connector_head_dim: int = 128,
-      connector_layers: int = 2,
-      num_thinking_tokens: int = 128,
-      dtype: DType = jnp.float32,
-      attention_kernel: str = "flash",
-      mesh: jax.sharding.Mesh = None,
-      rngs: nnx.Rngs = None,
-  ):
-    input_dim = gemma_dim * gemma_layers
-
-    self.feature_extractor = LTX2GemmaFeatureExtractor(
-        input_dim=input_dim,
-        output_dim=projection_dim,
-        dtype=dtype,
-        rngs=rngs,
-    )
-
-    self.embeddings_connector = Embeddings1DConnector(
-        input_dim=projection_dim,
-        heads=connector_heads,
-        head_dim=connector_head_dim,
-        layers=connector_layers,
-        num_learnable_registers=num_thinking_tokens,
-        rope_type="interleaved",
-        attention_kernel=attention_kernel,
-        mesh=mesh,
-        rngs=rngs,
-    )
-
-  def __call__(
-      self,
-      hidden_states: Union[Tuple[Array, ...], List[Array]],
-      attention_mask: Array,
-  ) -> Array:
-    """
-    Args:
-        hidden_states: From Gemma output.hidden_states (Tuple of [B, T, D])
-        attention_mask: [B, T]
-    """
-    # 1. Feature Extraction (Stack -> Norm -> Project)
-    features = self.feature_extractor(hidden_states, attention_mask)
+  Wraps feature extractor + video connector + audio connector.
+  Mirrors diffusers LTX2TextConnectors.
 
-    # 2. Connection (Refine + Thinking Tokens)
-    video_embeds = self.embeddings_connector(features, attention_mask)
-
-    return video_embeds
-
-
-class LTX2AudioVideoGemmaTextEncoder(nnx.Module):
-  """
-  Encoder for Audio-Video tasks.
   Pipeline: Gemma Hidden States -> Feature Extractor -> [Video Connector, Audio Connector]
   """
 
   def __init__(
       self,
-      # Feature Extractor Config (Shared)
+      # Feature Extractor Config
       gemma_dim: int = 3840,  # Gemma-3-12b
       gemma_layers: int = 49,  # Gemma-3 has 48 layers + 1 embedding layer output = 49 hidden states
-      projection_dim: int = 3840,
-      # Connector Config
+      projection_dim: int = 3840,  # LTX-2 conditioning dim
+      # Video Connector Config
       connector_heads: int = 30,
       connector_head_dim: int = 128,
       connector_layers: int = 2,
       num_thinking_tokens: int = 128,
+      # Audio Connector Config (defaults to same as video if not specified)
+      audio_connector_heads: int = 30,
+      audio_connector_head_dim: int = 128,
+      audio_connector_layers: int = 2,
       dtype: DType = jnp.float32,
       attention_kernel: str = "flash",
       mesh: jax.sharding.Mesh = None,
@@ -120,8 +64,8 @@ def __init__(
         rngs=rngs,
     )
 
-    # Two independent connectors
-    self.video_embeddings_connector = Embeddings1DConnector(
+    # Video connector
+    self.video_connector = Embeddings1DConnector(
         input_dim=projection_dim,
         heads=connector_heads,
         head_dim=connector_head_dim,
@@ -133,11 +77,12 @@ def __init__(
         rngs=rngs,
     )
 
-    self.audio_embeddings_connector = Embeddings1DConnector(
+    # Audio connector
+    self.audio_connector = Embeddings1DConnector(
         input_dim=projection_dim,
-        heads=connector_heads,
-        head_dim=connector_head_dim,
-        layers=connector_layers,
+        heads=audio_connector_heads,
+        head_dim=audio_connector_head_dim,
+        layers=audio_connector_layers,
         num_learnable_registers=num_thinking_tokens,
         rope_type="interleaved",
         attention_kernel=attention_kernel,
@@ -151,14 +96,20 @@ def __call__(
       attention_mask: Array,
   ) -> Tuple[Array, Array]:
     """
+    Args:
+        hidden_states: From Gemma output.hidden_states (Tuple of [B, T, D])
+        attention_mask: [B, T]
+
     Returns:
         (video_embeds, audio_embeds)
     """
-    # 1. Shared Feature Extraction
+    # 1. Feature Extraction (Stack -> Norm -> Project)
     features = self.feature_extractor(hidden_states, attention_mask)
 
-    # 2. Parallel Connection
-    video_embeds = self.video_embeddings_connector(features, attention_mask)
-    audio_embeds = self.audio_embeddings_connector(features, attention_mask)
+    # 2. Video Connector
+    video_embeds = self.video_connector(features, attention_mask)
+
+    # 3. Audio Connector
+    audio_embeds = self.audio_connector(features, attention_mask)
 
     return video_embeds, audio_embeds

…/tests/test_embeddings_connector_ltx2.py …s/ltx2/test_embeddings_connector_ltx2.pysrc/maxdiffusion/tests/test_embeddings_connector_ltx2.py renamed to src/maxdiffusion/tests/ltx2/test_embeddings_connector_ltx2.py src/maxdiffusion/tests/test_embeddings_connector_ltx2.py renamed to src/maxdiffusion/tests/ltx2/test_embeddings_connector_ltx2.py

@@ -18,7 +18,7 @@
 import jax.numpy as jnp
 import numpy as np
 from flax import nnx
-from ..models.ltx2.text_encoders.embeddings_connector_ltx2 import Embeddings1DConnector
+from maxdiffusion.models.ltx2.text_encoders.embeddings_connector_ltx2 import Embeddings1DConnector
 
 
 class Embeddings1DConnectorTest(unittest.TestCase):

…ion/tests/test_feature_extractor_ltx2.py …ests/ltx2/test_feature_extractor_ltx2.pysrc/maxdiffusion/tests/test_feature_extractor_ltx2.py renamed to src/maxdiffusion/tests/ltx2/test_feature_extractor_ltx2.py src/maxdiffusion/tests/test_feature_extractor_ltx2.py renamed to src/maxdiffusion/tests/ltx2/test_feature_extractor_ltx2.py

@@ -20,7 +20,7 @@
 import jax.numpy as jnp
 from flax import nnx
 
-from ..models.ltx2.text_encoders.feature_extractor_ltx2 import LTX2GemmaFeatureExtractor, _norm_and_concat_padded_batch
+from maxdiffusion.models.ltx2.text_encoders.feature_extractor_ltx2 import LTX2GemmaFeatureExtractor, _norm_and_concat_padded_batch
 
 
 # ==========================================

…ffusion/tests/test_text_encoders_ltx2.py …on/tests/ltx2/test_text_encoders_ltx2.pysrc/maxdiffusion/tests/test_text_encoders_ltx2.py renamed to src/maxdiffusion/tests/ltx2/test_text_encoders_ltx2.py src/maxdiffusion/tests/test_text_encoders_ltx2.py renamed to src/maxdiffusion/tests/ltx2/test_text_encoders_ltx2.py

@@ -18,7 +18,7 @@
 import jax.numpy as jnp
 import numpy as np
 from flax import nnx
-from ..models.ltx2.text_encoders.text_encoders_ltx2 import LTX2VideoGemmaTextEncoder, LTX2AudioVideoGemmaTextEncoder
+from maxdiffusion.models.ltx2.text_encoders.text_encoders_ltx2 import LTX2EmbeddingsProcessor
 
 
 class LTX2TextEncodersTest(unittest.TestCase):
@@ -36,53 +36,35 @@ def setUp(self):
 
     self.attention_mask = jnp.ones((self.B, self.T))
 
-  def test_video_encoder_forward(self):
-    encoder = LTX2VideoGemmaTextEncoder(
+  def test_embeddings_processor_forward(self):
+    processor = LTX2EmbeddingsProcessor(
         gemma_dim=self.gemma_dim,
         gemma_layers=self.gemma_layers,
         projection_dim=self.proj_dim,
         connector_heads=4,
         connector_head_dim=16,
         connector_layers=1,
         num_thinking_tokens=8,
+        audio_connector_heads=4,
+        audio_connector_head_dim=16,
+        audio_connector_layers=1,
         attention_kernel="dot_product",
         mesh=None,
         rngs=self.rng,
     )
 
-    output = encoder(tuple(self.hidden_states), self.attention_mask)
-
-    # Expected shape: [B, T, proj_dim]
-    self.assertEqual(output.shape, (self.B, self.T, self.proj_dim))
-    print("\n[PASS] Video Encoder Forward Pass Verified.")
-
-  def test_av_encoder_forward(self):
-    encoder = LTX2AudioVideoGemmaTextEncoder(
-        gemma_dim=self.gemma_dim,
-        gemma_layers=self.gemma_layers,
-        projection_dim=self.proj_dim,
-        connector_heads=4,
-        connector_head_dim=16,
-        connector_layers=1,
-        num_thinking_tokens=8,
-        attention_kernel="dot_product",
-        mesh=None,
-        rngs=self.rng,
-    )
-
-    video_out, audio_out = encoder(tuple(self.hidden_states), self.attention_mask)
+    video_out, audio_out = processor(tuple(self.hidden_states), self.attention_mask)
 
     # Expected shapes: Both [B, T, proj_dim]
     self.assertEqual(video_out.shape, (self.B, self.T, self.proj_dim))
     self.assertEqual(audio_out.shape, (self.B, self.T, self.proj_dim))
 
     # Ensure they are different (different random init for connectors)
-    # Note: In reality they are initialized differently, so outputs should differ
     self.assertFalse(
         jnp.allclose(video_out, audio_out), "Video and Audio outputs should differ due to different connector weights"
     )
 
-    print("\n[PASS] Audio-Video Encoder Forward Pass Verified.")
+    print("\n[PASS] Embeddings Processor Forward Pass Verified.")
 
 
 if __name__ == "__main__":