use NNXSimpleFeedForward, add mesh support

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

Author: syhuang22

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

@@ -19,28 +19,12 @@
 import jax.numpy as jnp
 from flax import nnx
 from maxdiffusion import common_types
-from ..attention_ltx2 import LTX2Attention
+from maxdiffusion.models.ltx2.attention_ltx2 import LTX2Attention
+from maxdiffusion.models.attention_flax import NNXSimpleFeedForward
 
 Array = common_types.Array
 DType = common_types.DType
 
-
-class FeedForward(nnx.Module):
-
-  def __init__(self, dim: int, dim_out: Optional[int] = None, mult: int = 4, dropout: float = 0.0, rngs: nnx.Rngs = None):
-    inner_dim = int(dim * mult)
-    dim_out = dim_out if dim_out is not None else dim
-
-    self.proj1 = nnx.Linear(dim, inner_dim, rngs=rngs)
-    self.proj2 = nnx.Linear(inner_dim, dim_out, rngs=rngs)
-
-  def __call__(self, x: Array) -> Array:
-    x = self.proj1(x)
-    x = jax.nn.gelu(x)
-    x = self.proj2(x)
-    return x
-
-
 class _BasicTransformerBlock1D(nnx.Module):
 
   def __init__(
@@ -50,6 +34,7 @@ def __init__(
       dim_head: int,
       rope_type: str = "interleaved",
       attention_kernel: str = "flash",
+      mesh: jax.sharding.Mesh = None,
       rngs: nnx.Rngs = None,
   ):
     self.attn1 = LTX2Attention(
@@ -60,9 +45,10 @@ def __init__(
         bias=True,  # LTX-2 default
         out_bias=True,
         attention_kernel=attention_kernel,
+        mesh=mesh,
         rngs=rngs,
     )
-    self.ff = FeedForward(dim, dim_out=dim, rngs=rngs)
+    self.ff = NNXSimpleFeedForward(rngs=rngs, dim=dim, dim_out=dim)
     self.norm1 = nnx.RMSNorm(dim, epsilon=1e-6, dtype=jnp.float32, param_dtype=jnp.float32, use_scale=True, rngs=rngs)
     self.norm2 = nnx.RMSNorm(dim, epsilon=1e-6, dtype=jnp.float32, param_dtype=jnp.float32, use_scale=True, rngs=rngs)
 
@@ -101,6 +87,7 @@ def __init__(
       num_learnable_registers: int = 128,
       rope_type: str = "interleaved",
       attention_kernel: str = "flash",
+      mesh: jax.sharding.Mesh = None,
       rngs: nnx.Rngs = None,
   ):
     self.dim = input_dim
@@ -115,7 +102,13 @@ def __init__(
     @nnx.vmap(in_axes=0, out_axes=0, axis_size=layers)
     def create_block(rngs):
       return _BasicTransformerBlock1D(
-          dim=input_dim, heads=heads, dim_head=head_dim, rope_type=rope_type, attention_kernel=attention_kernel, rngs=rngs
+          dim=input_dim,
+          heads=heads,
+          dim_head=head_dim,
+          rope_type=rope_type,
+          attention_kernel=attention_kernel,
+          mesh=mesh,
+          rngs=rngs,
       )
 
     # Call the vmapped constructor

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

@@ -46,6 +46,7 @@ def __init__(
       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
@@ -65,6 +66,7 @@ def __init__(
         num_learnable_registers=num_thinking_tokens,
         rope_type="interleaved",
         attention_kernel=attention_kernel,
+        mesh=mesh,
         rngs=rngs,
     )
 
@@ -106,6 +108,7 @@ def __init__(
       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
@@ -126,6 +129,7 @@ def __init__(
         num_learnable_registers=num_thinking_tokens,
         rope_type="interleaved",
         attention_kernel=attention_kernel,
+        mesh=mesh,
         rngs=rngs,
     )
 
@@ -137,6 +141,7 @@ def __init__(
         num_learnable_registers=num_thinking_tokens,
         rope_type="interleaved",
         attention_kernel=attention_kernel,
+        mesh=mesh,
         rngs=rngs,
     )
 

src/maxdiffusion/tests/test_embeddings_connector_ltx2.py

@@ -44,6 +44,7 @@ def test_thinking_tokens_replacement(self):
         head_dim=self.head_dim,
         layers=1,
         num_learnable_registers=self.num_learnable_registers,
+        mesh=None,
         rngs=self.rng,
     )
 
@@ -96,6 +97,7 @@ def test_forward_shape_and_run(self):
         layers=2,
         num_learnable_registers=self.num_learnable_registers,
         attention_kernel="dot_product",  # Use dot_product for testing on CPU
+        mesh=None,
         rngs=self.rng,
     )
 

src/maxdiffusion/tests/test_text_encoders_ltx2.py

@@ -47,6 +47,7 @@ def test_video_encoder_forward(self):
         connector_layers=1,
         num_thinking_tokens=8,
         attention_kernel="dot_product",
+        mesh=None,
         rngs=self.rng,
     )
 
@@ -66,6 +67,7 @@ def test_av_encoder_forward(self):
         connector_layers=1,
         num_thinking_tokens=8,
         attention_kernel="dot_product",
+        mesh=None,
         rngs=self.rng,
     )