Add NNXCombinedTimestepGuidanceTextProjEmbeddings class

Author: amepas

src/maxdiffusion/models/embeddings_flax.py

@@ -535,6 +535,80 @@ def __call__(self, timestep, guidance, pooled_projection=None):
     return conditioning
 
 
+class NNXCombinedTimestepGuidanceTextProjEmbeddings(nnx.Module):
+
+  def __init__(
+      self,
+      rngs: nnx.Rngs,
+      embedding_dim: int,
+      pooled_projection_dim: int = 0,
+      guidance_embeds: bool = True,
+      frequency_embedding_size: int = 256,
+      dtype: jnp.dtype = jnp.float32,
+      weights_dtype: jnp.dtype = jnp.float32,
+      precision: Optional[jax.lax.Precision] = None,
+  ):
+    self.embedding_dim = embedding_dim
+    self.pooled_projection_dim = pooled_projection_dim
+    self.guidance_embeds = guidance_embeds
+    self.frequency_embedding_size = frequency_embedding_size
+    self.dtype = dtype
+
+    self.time_proj = NNXTimesteps(num_channels=frequency_embedding_size, flip_sin_to_cos=True, downscale_freq_shift=0)
+    self.timestep_embedder = NNXTimestepEmbedding(
+        rngs=rngs,
+        in_channels=frequency_embedding_size,
+        time_embed_dim=embedding_dim,
+        dtype=dtype,
+        weights_dtype=weights_dtype,
+    )
+
+    if guidance_embeds:
+      self.guidance_proj = NNXTimesteps(num_channels=frequency_embedding_size, flip_sin_to_cos=True, downscale_freq_shift=0)
+      self.guidance_embedder = NNXTimestepEmbedding(
+          rngs=rngs,
+          in_channels=frequency_embedding_size,
+          time_embed_dim=embedding_dim,
+          dtype=dtype,
+          weights_dtype=weights_dtype,
+      )
+
+    if pooled_projection_dim > 0:
+      self.pooled_embedder = NNXPixArtAlphaTextProjection(
+          rngs=rngs,
+          in_features=pooled_projection_dim,
+          out_features=embedding_dim,
+          act_fn="silu",
+          dtype=dtype,
+          weights_dtype=weights_dtype,
+      )
+
+  def __call__(
+      self,
+      timestep: jax.Array,
+      guidance: Optional[jax.Array] = None,
+      pooled_projection: Optional[jax.Array] = None,
+  ) -> jax.Array:
+    timesteps_proj = self.time_proj(timestep)
+    dtype = pooled_projection.dtype if pooled_projection is not None else jnp.float32
+    timestep_emb = self.timestep_embedder(timesteps_proj.astype(dtype))
+
+    if self.guidance_embeds and guidance is not None:
+      guidance_proj = self.guidance_proj(guidance)
+      guidance_emb = self.guidance_embedder(guidance_proj.astype(dtype))
+      time_guidance_emb = timestep_emb + guidance_emb
+    else:
+      time_guidance_emb = timestep_emb
+
+    if pooled_projection is not None and self.pooled_projection_dim > 0:
+      pooled_projections = self.pooled_embedder(pooled_projection)
+      conditioning = time_guidance_emb + pooled_projections
+    else:
+      conditioning = time_guidance_emb
+
+    return conditioning
+
+
 class NNXTimesteps(nnx.Module):
 
   def __init__(self, num_channels: int, flip_sin_to_cos: bool, downscale_freq_shift: float, scale: int = 1):