Complete NNX module refactor and unit test suite for FLUX.2-klein

Author: amepas

src/maxdiffusion/models/embeddings_flax.py

@@ -12,12 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import math
-from typing import Optional, Any
+from typing import List, Optional, Tuple, Union, Any
 import flax.linen as nn
 from flax import nnx
-import jax.numpy as jnp
-from typing import List, Union
 import jax
+import jax.numpy as jnp
 from .modeling_flax_utils import get_activation
 from ..models.attention_flax import NNXSimpleFeedForward
 from ..models.normalization_flax import FP32LayerNorm
@@ -469,6 +468,48 @@ def __call__(self, ids):
       return out_freqs
 
 
+class NNXFluxPosEmbed(nnx.Module):
+
+  def __init__(
+      self,
+      theta: float = 10000.0,
+      axes_dim: Tuple[int, ...] = (16, 56, 56),
+      dtype: jnp.dtype = jnp.float32,
+      return_tuple: bool = True,
+  ):
+    self.theta = theta
+    self.axes_dim = axes_dim
+    self.dtype = dtype
+    self.return_tuple = return_tuple
+
+  def __call__(self, ids: jax.Array):
+    n_axes = len(self.axes_dim)
+    pos = ids.astype(self.dtype)
+    freqs_dtype = self.dtype
+
+    if self.return_tuple:
+      cos_out = []
+      sin_out = []
+      for i in range(n_axes):
+        dim = self.axes_dim[i]
+        p = pos[..., i]
+        freqs = 1.0 / (self.theta ** (jnp.arange(0, dim, 2, dtype=freqs_dtype) / dim))
+        freqs = jnp.outer(p, freqs)
+        freqs = jnp.repeat(freqs, 2, axis=-1)
+        cos_out.append(jnp.cos(freqs))
+        sin_out.append(jnp.sin(freqs))
+      freqs_cos = jnp.concatenate(cos_out, axis=-1)
+      freqs_sin = jnp.concatenate(sin_out, axis=-1)
+      return freqs_cos, freqs_sin
+    else:
+      out_freqs = []
+      for i in range(n_axes):
+        out = get_1d_rotary_pos_embed(self.axes_dim[i], pos[..., i], theta=self.theta, freqs_dtype=freqs_dtype)
+        out_freqs.append(out)
+      out_freqs = jnp.concatenate(out_freqs, axis=1)
+      return out_freqs
+
+
 class CombinedTimestepTextProjEmbeddings(nn.Module):
   embedding_dim: int
   pooled_projection_dim: int
@@ -554,7 +595,7 @@ def __init__(
     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.time_proj = NNXFlaxTimesteps(dim=frequency_embedding_size, flip_sin_to_cos=True, freq_shift=0.0)
     self.timestep_embedder = NNXTimestepEmbedding(
         rngs=rngs,
         in_channels=frequency_embedding_size,
@@ -564,7 +605,7 @@ def __init__(
     )
 
     if guidance_embeds:
-      self.guidance_proj = NNXTimesteps(num_channels=frequency_embedding_size, flip_sin_to_cos=True, downscale_freq_shift=0)
+      self.guidance_proj = NNXFlaxTimesteps(dim=frequency_embedding_size, flip_sin_to_cos=True, freq_shift=0.0)
       self.guidance_embedder = NNXTimestepEmbedding(
           rngs=rngs,
           in_channels=frequency_embedding_size,
@@ -577,6 +618,7 @@ def __init__(
       self.pooled_embedder = NNXPixArtAlphaTextProjection(
           rngs=rngs,
           in_features=pooled_projection_dim,
+          hidden_size=embedding_dim,
           out_features=embedding_dim,
           act_fn="silu",
           dtype=dtype,

src/maxdiffusion/models/flux/transformers/transformer_flux_flax.py

@@ -23,11 +23,19 @@
 from einops import repeat, rearrange
 from ....configuration_utils import ConfigMixin, flax_register_to_config
 from ...modeling_flax_utils import FlaxModelMixin
-from ...normalization_flax import AdaLayerNormZeroSingle, AdaLayerNormContinuous, AdaLayerNormZero
+from ...normalization_flax import (
+    AdaLayerNormZeroSingle,
+    AdaLayerNormContinuous,
+    AdaLayerNormZero,
+    NNXAdaLayerNormZeroSingle,
+    NNXAdaLayerNormContinuous,
+    NNXAdaLayerNormZero,
+)
 from ...attention_flax import FlaxFluxAttention as FluxAttention, FlaxFluxAttention, apply_rope
 from flax import nnx
 from ...embeddings_flax import (
     FluxPosEmbed,
+    NNXFluxPosEmbed,
     CombinedTimestepGuidanceTextProjEmbeddings,
     CombinedTimestepGuidanceTextProjEmbeddings as CombinedTimestepGuidanceTextEmbeddings,
     CombinedTimestepTextProjEmbeddings,
@@ -714,6 +722,7 @@ def init_weights(self, rngs, max_sequence_length, eval_only=True):
       )["params"]
 
 
+
 class FlaxSwiGluFeedForward(nn.Module):
   dim: int
   hidden_dim: int
@@ -1334,15 +1343,15 @@ def __init__(
     )
     self.query_norm = nnx.RMSNorm(
         num_features=dim_head,
-        eps=1e-6,
+        epsilon=1e-6,
         scale_init=nnx.with_partitioning(nnx.initializers.ones, ("heads",)),
         dtype=dtype,
         param_dtype=weights_dtype,
         rngs=rngs,
     )
     self.key_norm = nnx.RMSNorm(
         num_features=dim_head,
-        eps=1e-6,
+        epsilon=1e-6,
         scale_init=nnx.with_partitioning(nnx.initializers.ones, ("heads",)),
         dtype=dtype,
         param_dtype=weights_dtype,
@@ -1382,8 +1391,7 @@ def __call__(
     v = jnp.concatenate([v_txt, v_img], axis=1)
 
     if image_rotary_emb is not None:
-      cos, sin = image_rotary_emb
-      q, k = apply_rope(q, k, cos, sin)
+      q, k = apply_rope(q, k, image_rotary_emb)
 
     scale = self.dim_head**-0.5
     attn_weights = jnp.einsum("b q h d, b k h d -> b h q k", q, k, precision=None) * scale
@@ -1437,15 +1445,15 @@ def __init__(
     )
     self.norm_q = nnx.RMSNorm(
         num_features=attention_head_dim,
-        eps=1e-6,
+        epsilon=1e-6,
         scale_init=nnx.with_partitioning(nnx.initializers.ones, ("heads",)),
         dtype=dtype,
         param_dtype=weights_dtype,
         rngs=rngs,
     )
     self.norm_k = nnx.RMSNorm(
         num_features=attention_head_dim,
-        eps=1e-6,
+        epsilon=1e-6,
         scale_init=nnx.with_partitioning(nnx.initializers.ones, ("heads",)),
         dtype=dtype,
         param_dtype=weights_dtype,
@@ -1472,8 +1480,7 @@ def __call__(
     k = self.norm_k(k)
 
     if image_rotary_emb is not None:
-      cos, sin = image_rotary_emb
-      q, k = apply_rope(q, k, cos, sin)
+      q, k = apply_rope(q, k, image_rotary_emb)
 
     scale = self.dim_head**-0.5
     attn_weights = jnp.einsum("b q h d, b k h d -> b h q k", q, k, precision=None) * scale
@@ -1505,8 +1512,8 @@ def __init__(
     self.head_dim = attention_head_dim
     mlp_hidden_dim = int(dim * mlp_ratio)
 
-    self.img_norm1 = AdaLayerNormZero(dim, dtype=dtype, weights_dtype=weights_dtype)
-    self.txt_norm1 = AdaLayerNormZero(dim, dtype=dtype, weights_dtype=weights_dtype)
+    self.img_norm1 = NNXAdaLayerNormZero(dim, dtype=dtype, weights_dtype=weights_dtype)
+    self.txt_norm1 = NNXAdaLayerNormZero(dim, dtype=dtype, weights_dtype=weights_dtype)
 
     self.attn = NNXFluxDoubleAttention(
         rngs=rngs,
@@ -1605,7 +1612,7 @@ def __init__(
       weights_dtype: jnp.dtype = jnp.float32,
   ):
     self.dim = dim
-    self.norm = AdaLayerNormZeroSingle(dim, dtype=dtype, weights_dtype=weights_dtype)
+    self.norm = NNXAdaLayerNormZeroSingle(dim, dtype=dtype, weights_dtype=weights_dtype)
     self.attn = NNXFluxSingleAttention(
         rngs=rngs,
         dim=dim,
@@ -1660,7 +1667,7 @@ def __init__(
     self.inner_dim = num_attention_heads * attention_head_dim
     self.dtype = dtype
 
-    self.pos_embed = FluxPosEmbed(theta=theta, axes_dim=axes_dim, return_tuple=True)
+    self.pos_embed = NNXFluxPosEmbed(axes_dim=axes_dim, theta=theta, return_tuple=True)
     self.time_text_embed = NNXCombinedTimestepGuidanceTextProjEmbeddings(
         rngs=rngs,
         embedding_dim=self.inner_dim,
@@ -1710,7 +1717,7 @@ def __init__(
         rngs=rngs,
     )
 
-    self.double_blocks = [
+    self.double_blocks = nnx.List([
         NNXFluxDoubleTransformerBlock(
             rngs=rngs,
             dim=self.inner_dim,
@@ -1720,9 +1727,9 @@ def __init__(
             weights_dtype=weights_dtype,
         )
         for _ in range(num_layers)
-    ]
+    ])
 
-    self.single_blocks = [
+    self.single_blocks = nnx.List([
         NNXFluxSingleTransformerBlock(
             rngs=rngs,
             dim=self.inner_dim,
@@ -1732,11 +1739,11 @@ def __init__(
             weights_dtype=weights_dtype,
         )
         for _ in range(num_single_layers)
-    ]
+    ])
 
-    self.norm_out = AdaLayerNormContinuous(
-        self.inner_dim,
-        elementwise_affine=False,
+    self.norm_out = NNXAdaLayerNormContinuous(
+        rngs=rngs,
+        embedding_dim=self.inner_dim,
         eps=1e-6,
         dtype=dtype,
         weights_dtype=weights_dtype,

src/maxdiffusion/models/normalization_flax.py

@@ -173,3 +173,71 @@ def __init__(self, rngs: nnx.Rngs, dim: int, eps: float, elementwise_affine: boo
   def __call__(self, inputs: jax.Array) -> jax.Array:
     origin_dtype = inputs.dtype
     return self.layer_norm(inputs.astype(dtype=jnp.float32)).astype(dtype=origin_dtype)
+
+
+# =============================================================================
+# FLAX NNX ADALAYERNORM IMPLEMENTATIONS FOR FLUX.2-KLEIN
+# =============================================================================
+
+
+class NNXAdaLayerNormContinuous(nnx.Module):
+
+  def __init__(
+      self,
+      rngs: nnx.Rngs,
+      embedding_dim: int,
+      eps: float = 1e-6,
+      dtype: jnp.dtype = jnp.float32,
+      weights_dtype: jnp.dtype = jnp.float32,
+  ):
+    self.embedding_dim = embedding_dim
+    self.eps = eps
+    self.dtype = dtype
+    self.layer_norm = nnx.LayerNorm(
+        num_features=embedding_dim, epsilon=eps, use_bias=False, use_scale=False, dtype=dtype, rngs=rngs
+    )
+    self.linear = nnx.Linear(
+        in_features=embedding_dim, out_features=embedding_dim * 2, use_bias=True, dtype=dtype, param_dtype=weights_dtype, rngs=rngs
+    )
+
+  def __call__(self, x: jax.Array, conditioning_embedding: jax.Array) -> jax.Array:
+    emb = self.linear(jax.nn.silu(conditioning_embedding))
+    scale, shift = jnp.split(emb, 2, axis=-1)
+    x_norm = self.layer_norm(x)
+    return (1.0 + scale[:, None, :]) * x_norm + shift[:, None, :]
+
+
+class NNXAdaLayerNormZero(nnx.Module):
+
+  def __init__(self, embedding_dim: int, eps: float = 1e-6, dtype: jnp.dtype = jnp.float32, weights_dtype: jnp.dtype = jnp.float32):
+    self.embedding_dim = embedding_dim
+    self.eps = eps
+    self.dtype = dtype
+
+  def __call__(self, x: jax.Array, emb: jax.Array):
+    if emb.ndim == 2:
+      emb = emb[:, None, :]
+    shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = jnp.split(emb, 6, axis=-1)
+    mean = jnp.mean(x, axis=-1, keepdims=True)
+    variance = jnp.mean(jnp.square(x - mean), axis=-1, keepdims=True)
+    inv_std = jax.lax.rsqrt(variance + self.eps)
+    normed_x = (x - mean) * inv_std * (1.0 + scale_msa) + shift_msa
+    return normed_x, gate_msa, shift_mlp, scale_mlp, gate_mlp
+
+
+class NNXAdaLayerNormZeroSingle(nnx.Module):
+
+  def __init__(self, embedding_dim: int, eps: float = 1e-6, dtype: jnp.dtype = jnp.float32, weights_dtype: jnp.dtype = jnp.float32):
+    self.embedding_dim = embedding_dim
+    self.eps = eps
+    self.dtype = dtype
+
+  def __call__(self, x: jax.Array, emb: jax.Array):
+    if emb.ndim == 2:
+      emb = emb[:, None, :]
+    shift_msa, scale_msa, gate_msa = jnp.split(emb, 3, axis=-1)
+    mean = jnp.mean(x, axis=-1, keepdims=True)
+    variance = jnp.mean(jnp.square(x - mean), axis=-1, keepdims=True)
+    inv_std = jax.lax.rsqrt(variance + self.eps)
+    normed_x = (x - mean) * inv_std * (1.0 + scale_msa) + shift_msa
+    return normed_x, gate_msa

src/maxdiffusion/models/qwen3_flax.py

@@ -417,7 +417,7 @@ def __call__(self, x: jnp.ndarray) -> jnp.ndarray:
     x_float = x.astype(jnp.float32)
     variance = jnp.mean(jnp.square(x_float), axis=-1, keepdims=True)
     normed = x_float * jax.lax.rsqrt(variance + self.eps)
-    return (normed.astype(self.dtype)) * self.weight.value
+    return (normed.astype(self.dtype)) * self.weight[...]
 
 
 class NNXFlaxQwen3MLP(nnx.Module):
@@ -518,7 +518,9 @@ def __call__(
     k = self.k_norm(k)
 
     if cos_table is not None and sin_table is not None:
-      q, k = apply_qwen3_rotary_pos_emb(q, k, cos_table, sin_table)
+      cos_seq = cos_table[:seq_len, :]
+      sin_seq = sin_table[:seq_len, :]
+      q, k = apply_qwen3_rotary_pos_emb(q, k, cos_seq, sin_seq)
 
     if self.num_kv_heads != self.num_heads:
       num_repeats = self.num_heads // self.num_kv_heads
@@ -585,11 +587,15 @@ def __init__(self, rngs: nnx.Rngs, config: FlaxQwen3Config):
     self.embed_tokens = nnx.Embed(
         num_embeddings=config.vocab_size,
         features=config.hidden_size,
-        embedding_init=nnx.with_partitioning(nnx.initializers.normal(stddev=config.hidden_size**-0.5), ("vocab", "embed")),
+        embedding_init=nnx.with_partitioning(
+            nnx.initializers.normal(stddev=config.hidden_size**-0.5), ("vocab", "embed")
+        ),
         dtype=config.dtype,
         rngs=rngs,
     )
-    self.layers = [NNXFlaxQwen3DecoderLayer(rngs=rngs, config=config) for _ in range(config.num_hidden_layers)]
+    self.layers = nnx.List(
+        [NNXFlaxQwen3DecoderLayer(rngs=rngs, config=config) for _ in range(config.num_hidden_layers)]
+    )
     self.norm = NNXFlaxQwen3RMSNorm(rngs=rngs, dim=config.hidden_size, eps=config.rms_norm_eps, dtype=config.dtype)
 
   def __call__(

src/maxdiffusion/models/vae_flax.py

@@ -16,7 +16,7 @@
 
 import math
 from functools import partial
-from typing import Tuple
+from typing import Optional, Tuple
 
 import flax
 from flax import nnx
@@ -995,21 +995,24 @@ def __init__(
         dtype=dtype,
         weights_dtype=weights_dtype,
     )
-    self.post_quant_conv = nn.Conv(
-        latent_channels,
+    self.post_quant_conv = nnx.Conv(
+        in_features=latent_channels,
+        out_features=latent_channels,
         kernel_size=(1, 1),
         strides=(1, 1),
         padding="VALID",
         dtype=dtype,
         param_dtype=weights_dtype,
+        rngs=rngs,
     )
 
-  def decode(self, latents: jax.Array, deterministic: bool = True, return_dict: bool = True):
+  def decode(self, latents: jax.Array, decoder_params: Optional[dict] = None, deterministic: bool = True, return_dict: bool = True):
     if latents.shape[-1] != self.latent_channels:
       latents = jnp.transpose(latents, (0, 2, 3, 1))
 
     hidden_states = self.post_quant_conv(latents)
-    hidden_states = self.decoder(hidden_states, deterministic=deterministic)
+    if decoder_params is not None:
+      hidden_states = self.decoder.apply({"params": decoder_params}, hidden_states, deterministic=deterministic)
     hidden_states = jnp.transpose(hidden_states, (0, 3, 1, 2))
 
     if not return_dict: