NNX simple feed forward wrapper

Author: prishajain1

src/maxdiffusion/models/embeddings_flax.py

@@ -385,6 +385,25 @@ def __call__(self, caption):
     return hidden_states
 
 
+class NNXSimpleLinearWrapper(nnx.Module):
+
+  def __init__(self, rngs: nnx.Rngs, in_features: int, out_features: int, weights_dtype: jnp.dtype):
+    super().__init__()
+    self.linear = nnx.Linear(
+        rngs=rngs,
+        in_features=in_features,
+        out_features=out_features,
+        use_bias=True,
+        dtype=jnp.float32,
+        param_dtype=weights_dtype,
+        kernel_init=nnx.with_partitioning(nnx.initializers.xavier_uniform(), ("mlp", "embed")),
+        bias_init=nnx.with_partitioning(nnx.initializers.zeros, ("embed",)),
+    )
+
+  def __call__(self, x):
+    return self.linear(x)
+
+
 class PixArtAlphaTextProjection(nn.Module):
   """
   Projects caption embeddings. Also handles dropout for classifier-free guidance.

src/maxdiffusion/models/ltx2/transformer_ltx2.py

@@ -21,12 +21,13 @@
 
 from maxdiffusion.models.ltx2.attention_ltx2 import LTX2Attention, LTX2RotaryPosEmbed
 from maxdiffusion.models.attention_flax import NNXSimpleFeedForward
-from maxdiffusion.models.embeddings_flax import NNXPixArtAlphaCombinedTimestepSizeEmbeddings, NNXPixArtAlphaTextProjection, NNXCombinedTimestepTextProjEmbeddings
+from maxdiffusion.models.embeddings_flax import NNXPixArtAlphaCombinedTimestepSizeEmbeddings, NNXPixArtAlphaTextProjection, NNXCombinedTimestepTextProjEmbeddings, NNXSimpleLinearWrapper
 from maxdiffusion.models.gradient_checkpoint import GradientCheckpointType
 from maxdiffusion.configuration_utils import ConfigMixin, register_to_config
 from maxdiffusion.common_types import BlockSizes
 
 
+
 class LTX2AdaLayerNormSingle(nnx.Module):
 
   def __init__(
@@ -695,20 +696,16 @@ def __init__(
 
     # 2. Prompt embeddings
     if self.cross_attn_mod:
-      self.caption_projection = NNXCombinedTimestepTextProjEmbeddings(
+      self.caption_projection = NNXSimpleLinearWrapper(
           rngs=rngs,
           in_features=self.caption_channels,
-          hidden_size=self.cross_attention_dim,
-          embedding_dim=inner_dim,
-          dtype=self.dtype,
+          out_features=self.cross_attention_dim,
           weights_dtype=self.weights_dtype,
       )
-      self.audio_caption_projection = NNXCombinedTimestepTextProjEmbeddings(
+      self.audio_caption_projection = NNXSimpleLinearWrapper(
           rngs=rngs,
           in_features=self.audio_caption_channels,
-          hidden_size=self.audio_cross_attention_dim,
-          embedding_dim=audio_inner_dim,
-          dtype=self.dtype,
+          out_features=self.audio_cross_attention_dim,
           weights_dtype=self.weights_dtype,
       )
     else: