NNXCombinedTimestepTextProjEmbeddings addition

Author: prishajain1

src/maxdiffusion/models/embeddings_flax.py

@@ -473,6 +473,71 @@ def __call__(self, timestep, pooled_projection):
     conditioning = timestep_emb + pooled_projections
     return conditioning
 
+class NNXCombinedTimestepTextProjEmbeddings(nnx.Module):
+  def __init__(
+      self,
+      rngs: nnx.Rngs,
+      in_features: int,
+      hidden_size: int,
+      embedding_dim: int,
+      out_features: int = None,
+      act_fn: str = "gelu_tanh",
+      dtype: jnp.dtype = jnp.float32,
+      weights_dtype: jnp.dtype = jnp.float32,
+      precision: jax.lax.Precision = None,
+  ):
+    if out_features is None:
+      out_features = hidden_size
+
+    self.linear_1 = nnx.Linear(
+        rngs=rngs,
+        in_features=in_features,
+        out_features=hidden_size,
+        use_bias=True,
+        dtype=jnp.float32,
+        param_dtype=weights_dtype,
+        precision=precision,
+        kernel_init=nnx.with_partitioning(nnx.initializers.xavier_uniform(), ("embed", "mlp")),
+        bias_init=nnx.with_partitioning(nnx.initializers.zeros, ("mlp",)),
+    )
+    self.act_1 = get_activation(act_fn)
+
+    self.linear_2 = nnx.Linear(
+        rngs=rngs,
+        in_features=hidden_size,
+        out_features=out_features,
+        use_bias=True,
+        dtype=jnp.float32,
+        param_dtype=weights_dtype,
+        precision=precision,
+        kernel_init=nnx.with_partitioning(nnx.initializers.xavier_uniform(), ("mlp", "embed")),
+        bias_init=nnx.with_partitioning(nnx.initializers.zeros, ("embed",)),
+    )
+
+    self.time_proj = NNXTimesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0)
+
+    class EmbWrapper(nnx.Module):
+      def __init__(self, rngs: nnx.Rngs, embedding_dim: int, weights_dtype: jnp.dtype):
+        self.timestep_embedder = NNXTimestepEmbedding(
+            rngs=rngs,
+            in_channels=256,
+            time_embed_dim=embedding_dim,
+            dtype=jnp.float32,
+            weights_dtype=weights_dtype,
+        )
+
+    self.emb = EmbWrapper(rngs, embedding_dim, weights_dtype)
+
+  def __call__(self, caption, timestep):
+    hidden_states = self.linear_1(caption)
+    hidden_states = self.act_1(hidden_states)
+    hidden_states = self.linear_2(hidden_states)
+
+    timesteps_proj = self.time_proj(timestep)
+    timesteps_emb = self.emb.timestep_embedder(timesteps_proj)
+
+    return hidden_states + timesteps_emb
+
 
 class CombinedTimestepGuidanceTextProjEmbeddings(nn.Module):
   embedding_dim: int

src/maxdiffusion/models/ltx2/transformer_ltx2.py

@@ -21,7 +21,7 @@
 
 from maxdiffusion.models.ltx2.attention_ltx2 import LTX2Attention, LTX2RotaryPosEmbed
 from maxdiffusion.models.attention_flax import NNXSimpleFeedForward
-from maxdiffusion.models.embeddings_flax import NNXPixArtAlphaCombinedTimestepSizeEmbeddings, NNXPixArtAlphaTextProjection
+from maxdiffusion.models.embeddings_flax import NNXPixArtAlphaCombinedTimestepSizeEmbeddings, NNXPixArtAlphaTextProjection, NNXCombinedTimestepTextProjEmbeddings
 from maxdiffusion.models.gradient_checkpoint import GradientCheckpointType
 from maxdiffusion.configuration_utils import ConfigMixin, register_to_config
 from maxdiffusion.common_types import BlockSizes
@@ -692,20 +692,38 @@ def __init__(
     )
 
     # 2. Prompt embeddings
-    self.caption_projection = NNXPixArtAlphaTextProjection(
-        rngs=rngs,
-        in_features=self.caption_channels,
-        hidden_size=inner_dim,
-        dtype=self.dtype,
-        weights_dtype=self.weights_dtype,
-    )
-    self.audio_caption_projection = NNXPixArtAlphaTextProjection(
-        rngs=rngs,
-        in_features=self.caption_channels,
-        hidden_size=audio_inner_dim,
-        dtype=self.dtype,
-        weights_dtype=self.weights_dtype,
-    )
+    if self.cross_attn_mod:
+      self.caption_projection = NNXCombinedTimestepTextProjEmbeddings(
+          rngs=rngs,
+          in_features=self.caption_channels,
+          hidden_size=inner_dim,
+          embedding_dim=inner_dim,
+          dtype=self.dtype,
+          weights_dtype=self.weights_dtype,
+      )
+      self.audio_caption_projection = NNXCombinedTimestepTextProjEmbeddings(
+          rngs=rngs,
+          in_features=self.caption_channels,
+          hidden_size=audio_inner_dim,
+          embedding_dim=audio_inner_dim,
+          dtype=self.dtype,
+          weights_dtype=self.weights_dtype,
+      )
+    else:
+      self.caption_projection = NNXPixArtAlphaTextProjection(
+          rngs=rngs,
+          in_features=self.caption_channels,
+          hidden_size=inner_dim,
+          dtype=self.dtype,
+          weights_dtype=self.weights_dtype,
+      )
+      self.audio_caption_projection = NNXPixArtAlphaTextProjection(
+          rngs=rngs,
+          in_features=self.caption_channels,
+          hidden_size=audio_inner_dim,
+          dtype=self.dtype,
+          weights_dtype=self.weights_dtype,
+      )
     # 3. Timestep Modulation Params and Embedding
     num_mod_params = 9 if self.cross_attn_mod else 6
     self.time_embed = LTX2AdaLayerNormSingle(