embeddings_flax.py changes

Author: prishajain1

src/maxdiffusion/models/embeddings_flax.py

@@ -486,19 +486,13 @@ def __init__(
       weights_dtype: jnp.dtype = jnp.float32,
       precision: jax.lax.Precision = None,
   ):
-    if out_features is None:
-      out_features = hidden_size
-
-    self.linear = nnx.Linear(
+    self.text_embedder = NNXPixArtAlphaTextProjection(
         rngs=rngs,
         in_features=in_features,
-        out_features=out_features,
-        use_bias=True,
-        dtype=jnp.float32,
-        param_dtype=weights_dtype,
+        hidden_size=embedding_dim,
+        dtype=dtype,
+        weights_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)
@@ -516,7 +510,7 @@ def __init__(self, rngs: nnx.Rngs, embedding_dim: int, weights_dtype: jnp.dtype)
     self.emb = EmbWrapper(rngs, embedding_dim, weights_dtype)
 
   def __call__(self, caption, timestep):
-    hidden_states = self.linear(caption)
+    hidden_states = self.text_embedder(caption)
 
     timesteps_proj = self.time_proj(timestep)
     timesteps_emb = self.emb.timestep_embedder(timesteps_proj)