adaln.py for nnx, classes NNXPixArtAlphaCombinedTimestepSizeEmbeddings and NNXTimesteps added

Author: prishajain1

src/maxdiffusion/models/embeddings_flax.py

@@ -501,3 +501,38 @@ def __call__(self, timestep, guidance, pooled_projection):
     conditioning = time_guidance_emb + pooled_projections
 
     return conditioning
+
+
+class NNXTimesteps(nnx.Module):
+  def __init__(self, num_channels: int, flip_sin_to_cos: bool, downscale_freq_shift: float, scale: int = 1):
+    self.num_channels = num_channels
+    self.flip_sin_to_cos = flip_sin_to_cos
+    self.downscale_freq_shift = downscale_freq_shift
+    self.scale = scale
+
+  def __call__(self, timesteps: jax.Array) -> jax.Array:
+    return get_sinusoidal_embeddings(
+        timesteps=timesteps,
+        embedding_dim=self.num_channels,
+        freq_shift=self.downscale_freq_shift,
+        flip_sin_to_cos=self.flip_sin_to_cos,
+        scale=self.scale
+    )
+
+
+class NNXPixArtAlphaCombinedTimestepSizeEmbeddings(nnx.Module):
+  def __init__(self, rngs: nnx.Rngs, embedding_dim: int, size_emb_dim: int, dtype: jnp.dtype = jnp.float32, weights_dtype: jnp.dtype = jnp.float32):
+    self.outdim = size_emb_dim
+    self.time_proj = NNXTimesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0)
+    self.timestep_embedder = NNXTimestepEmbedding(
+        rngs=rngs,
+        in_channels=256,
+        time_embed_dim=embedding_dim,
+        dtype=dtype,
+        weights_dtype=weights_dtype
+    )
+
+  def __call__(self, timestep: jax.Array, hidden_dtype: jnp.dtype = jnp.float32) -> jax.Array:
+    timesteps_proj = self.time_proj(timestep)
+    timesteps_emb = self.timestep_embedder(timesteps_proj.astype(hidden_dtype))
+    return timesteps_emb

src/maxdiffusion/models/ltx_2/adaln.py

@@ -0,0 +1,35 @@
+
+import jax
+import jax.numpy as jnp
+from flax import nnx
+from typing import Optional, Tuple
+from maxdiffusion.models.embeddings_flax import NNXPixArtAlphaCombinedTimestepSizeEmbeddings
+
+class AdaLayerNormSingle(nnx.Module):
+    """
+    Norm layer adaptive layer norm single (adaLN-single).
+    As proposed in PixArt-Alpha (see: https://arxiv.org/abs/2310.00426; Section 2.3).
+    """
+    def __init__(self, rngs: nnx.Rngs, embedding_dim: int, embedding_coefficient: int = 6, dtype: jnp.dtype = jnp.float32, weights_dtype: jnp.dtype = jnp.float32):
+        self.emb = NNXPixArtAlphaCombinedTimestepSizeEmbeddings(
+            rngs=rngs,
+            embedding_dim=embedding_dim,
+            size_emb_dim=embedding_dim // 3,
+            dtype=dtype,
+            weights_dtype=weights_dtype
+        )
+        self.silu = nnx.silu
+        self.linear = nnx.Linear(
+            rngs=rngs,
+            in_features=embedding_dim,
+            out_features=embedding_coefficient * embedding_dim,
+            use_bias=True,
+            dtype=dtype,
+            param_dtype=weights_dtype,
+            kernel_init=nnx.initializers.zeros,
+            bias_init=nnx.initializers.zeros
+        )
+
+    def __call__(self, timestep: jax.Array, hidden_dtype: Optional[jnp.dtype] = None) -> Tuple[jax.Array, jax.Array]:
+        embedded_timestep = self.emb(timestep, hidden_dtype=hidden_dtype)
+        return self.linear(self.silu(embedded_timestep)), embedded_timestep