From 71584763ff18b8cd8e4801a68552147434576b1a Mon Sep 17 00:00:00 2001 From: olehtika Date: Tue, 15 Sep 2026 13:04:22 +0000 Subject: [PATCH 1/5] Stop double-projecting the FLUX timestep and guidance embeddings FluxTransformer2DModel runs timestep and guidance through its own timestep_embedding() before calling CombinedTimestepGuidanceTextProjEmbeddings, which since the Flux.2-Klein onboarding projects them a second time. That shapes the result (batch, frequency_embedding_size, embedding_dim) and FLUX-dev fails at transformer init with: ValueError: Incompatible shapes for broadcasting: shapes=[(8, 256, 3072), (8, 3072)] Take the already-projected values as-is, as this module did before. The NNX variant keeps its own projection, which is what Flux.2-Klein uses. --- src/maxdiffusion/models/embeddings_flax.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/src/maxdiffusion/models/embeddings_flax.py b/src/maxdiffusion/models/embeddings_flax.py index 17ac2b07a..9a2a9772e 100644 --- a/src/maxdiffusion/models/embeddings_flax.py +++ b/src/maxdiffusion/models/embeddings_flax.py @@ -546,14 +546,19 @@ class CombinedTimestepGuidanceTextProjEmbeddings(nn.Module): @nn.compact def __call__(self, timestep, guidance, pooled_projection=None): - timesteps_proj = FlaxTimesteps(dim=self.frequency_embedding_size, flip_sin_to_cos=True, freq_shift=0)(timestep) + # timestep and guidance arrive already projected: FluxTransformer2DModel runs + # both through its own timestep_embedding(), which carries a time_factor the + # sinusoidal helper here does not. Projecting again would shape them + # (batch, frequency_embedding_size, embedding_dim) and break the sum with the + # pooled projection below. The NNX variant owns its projection instead. + timesteps_proj = timestep dtype = pooled_projection.dtype if pooled_projection is not None else jnp.float32 timestep_emb = FlaxTimestepEmbedding( time_embed_dim=self.embedding_dim, dtype=self.dtype, weights_dtype=self.weights_dtype )(timesteps_proj.astype(dtype)) if self.guidance_embeds and guidance is not None: - guidance_proj = FlaxTimesteps(dim=self.frequency_embedding_size, flip_sin_to_cos=True, freq_shift=0)(guidance) + guidance_proj = guidance guidance_emb = FlaxTimestepEmbedding( time_embed_dim=self.embedding_dim, dtype=self.dtype, weights_dtype=self.weights_dtype )(guidance_proj.astype(dtype)) From d6a25c735db65246ce970170c317e31e51ebe08f Mon Sep 17 00:00:00 2001 From: olehtika Date: Tue, 15 Sep 2026 13:16:23 +0000 Subject: [PATCH 2/5] Load the FLUX final-layer modulation into its renamed submodule The Flux.2-Klein onboarding gave the Dense inside AdaLayerNormContinuous the name "linear", so the params tree now carries norm_out/linear. The original-format FLUX loader still rewrote adaLN_modulation_1 to norm_out.Dense_0, so loading pretrained transformer weights (train_new_flux=False) built a params tree that no longer matched the sharding tree: ValueError: device_put device specification must be a tree prefix of the corresponding value (- ('linear',) + ('Dense_0',)) The diffusers-format paths in this file were already updated to norm_out/linear. --- src/maxdiffusion/models/flux/util.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/maxdiffusion/models/flux/util.py b/src/maxdiffusion/models/flux/util.py index 77e34099c..f651368ec 100644 --- a/src/maxdiffusion/models/flux/util.py +++ b/src/maxdiffusion/models/flux/util.py @@ -260,7 +260,7 @@ def load_flow_model(name: str, eval_shapes: dict, device: str, hf_download: bool renamed_pt_key = renamed_pt_key.replace("out_layer", "linear_2") elif "final_layer" in renamed_pt_key: renamed_pt_key = renamed_pt_key.replace("final_layer.linear", "proj_out") - renamed_pt_key = renamed_pt_key.replace("final_layer.adaLN_modulation_1", "norm_out.Dense_0") + renamed_pt_key = renamed_pt_key.replace("final_layer.adaLN_modulation_1", "norm_out.linear") pt_tuple_key = tuple(renamed_pt_key.split(".")) flax_key, flax_tensor = rename_key_and_reshape_tensor(pt_tuple_key, tensor, eval_shapes) From 3c818f87947d5baf395099400dddc025a9882716 Mon Sep 17 00:00:00 2001 From: olehtika Date: Tue, 15 Sep 2026 13:16:24 +0000 Subject: [PATCH 3/5] Keep the legacy FLUX adaptive norm on shift_scale order AdaLayerNormContinuous gained a scale_shift_order option whose default, "scale_shift", matches Flux.2-Klein. FLUX's norm_out passed no order and so silently switched convention, swapping shift and scale for weights loaded from the original FLUX checkpoint, which emits shift first. Flux.2-Klein's Flax transformer already passes the option explicitly, so state it here too. --- .../models/flux/transformers/transformer_flux_flax.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/maxdiffusion/models/flux/transformers/transformer_flux_flax.py b/src/maxdiffusion/models/flux/transformers/transformer_flux_flax.py index 2d536ca1c..026aaded6 100644 --- a/src/maxdiffusion/models/flux/transformers/transformer_flux_flax.py +++ b/src/maxdiffusion/models/flux/transformers/transformer_flux_flax.py @@ -545,6 +545,9 @@ def setup(self): self.inner_dim, elementwise_affine=False, eps=self.eps, + # FLUX loads adaLN_modulation_1 from the original checkpoint, which emits + # shift before scale. The module default follows Flux.2-Klein instead. + scale_shift_order="shift_scale", dtype=self.dtype, weights_dtype=self.weights_dtype, precision=self.precision, From 0c4f9e061f9c607be3dc4374ab339420b862a7a9 Mon Sep 17 00:00:00 2001 From: olehtika Date: Tue, 15 Sep 2026 13:54:14 +0000 Subject: [PATCH 4/5] Add CPU tests for the FLUX conditioning conventions The double-projection regression reproduces at module level in seconds with random weights, so it is worth guarding. The other two cases pin what the original-format FLUX loader relies on: the modulation Dense is reachable as norm_out/linear, and shift_scale order applies the first half of the modulation as shift. --- .../tests/flux_conditioning_test.py | 95 +++++++++++++++++++ 1 file changed, 95 insertions(+) create mode 100644 src/maxdiffusion/tests/flux_conditioning_test.py diff --git a/src/maxdiffusion/tests/flux_conditioning_test.py b/src/maxdiffusion/tests/flux_conditioning_test.py new file mode 100644 index 000000000..813f17840 --- /dev/null +++ b/src/maxdiffusion/tests/flux_conditioning_test.py @@ -0,0 +1,95 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for the FLUX conditioning path (CPU backend). + +FLUX shares CombinedTimestepGuidanceTextProjEmbeddings and +AdaLayerNormContinuous with Flux.2-Klein, which reads them with different +conventions. These cover the conventions FLUX needs. +""" + +import os +import unittest + +os.environ.setdefault("JAX_PLATFORMS", "cpu") + +import flax +import jax +import jax.numpy as jnp +import numpy as np + +from maxdiffusion.models.embeddings_flax import CombinedTimestepGuidanceTextProjEmbeddings +from maxdiffusion.models.normalization_flax import AdaLayerNormContinuous + + +def _layer_norm(x, eps=1e-5): + mean = np.mean(x, axis=-1, keepdims=True) + variance = np.var(x, axis=-1, keepdims=True) + return (x - mean) / np.sqrt(variance + eps) + + +class FluxConditioningTest(unittest.TestCase): + + def test_guidance_embeddings_accept_projected_timesteps(self): + """FluxTransformer2DModel projects timestep and guidance before this module.""" + batch, frequency_embedding_size, embedding_dim, pooled_projection_dim = 2, 256, 64, 32 + module = CombinedTimestepGuidanceTextProjEmbeddings( + embedding_dim=embedding_dim, pooled_projection_dim=pooled_projection_dim + ) + timestep = jnp.zeros((batch, frequency_embedding_size), jnp.float32) + guidance = jnp.zeros((batch, frequency_embedding_size), jnp.float32) + pooled_projection = jnp.zeros((batch, pooled_projection_dim), jnp.float32) + + variables = module.init(jax.random.PRNGKey(0), timestep, guidance, pooled_projection) + conditioning = module.apply(variables, timestep, guidance, pooled_projection) + + self.assertEqual(conditioning.shape, (batch, embedding_dim)) + + def test_norm_out_dense_is_named_linear(self): + """load_flow_model writes final_layer.adaLN_modulation_1 to norm_out/linear.""" + module = AdaLayerNormContinuous(embedding_dim=4, elementwise_affine=False) + x = jnp.zeros((2, 3, 4), jnp.float32) + conditioning_embedding = jnp.zeros((2, 8), jnp.float32) + + variables = module.init(jax.random.PRNGKey(0), x, conditioning_embedding) + + self.assertIn("linear", variables["params"]) + + def test_shift_scale_order_reads_shift_first(self): + """The original FLUX checkpoint emits shift ahead of scale.""" + embedding_dim, conditioning_dim, batch, sequence = 4, 8, 2, 3 + x = jax.random.normal(jax.random.PRNGKey(1), (batch, sequence, embedding_dim)) + conditioning_embedding = jax.random.normal(jax.random.PRNGKey(2), (batch, conditioning_dim)) + + def modulate(scale_shift_order): + module = AdaLayerNormContinuous( + embedding_dim=embedding_dim, elementwise_affine=False, scale_shift_order=scale_shift_order + ) + variables = flax.linen.meta.unbox(flax.core.unfreeze(module.init(jax.random.PRNGKey(0), x, conditioning_embedding))) + # A zero kernel leaves the bias as the whole modulation, so the halves are + # known: 2.0 in the first, 0.0 in the second. + variables["params"]["linear"]["kernel"] = jnp.zeros_like(variables["params"]["linear"]["kernel"]) + variables["params"]["linear"]["bias"] = jnp.concatenate([jnp.full((embedding_dim,), 2.0), jnp.zeros((embedding_dim,))]) + return module.apply(variables, x, conditioning_embedding) + + normalized = _layer_norm(np.asarray(x)) + + # shift first: (1 + 0) * norm + 2 + np.testing.assert_allclose(np.asarray(modulate("shift_scale")), normalized + 2.0, atol=1e-4) + # scale first: (1 + 2) * norm + 0 + np.testing.assert_allclose(np.asarray(modulate("scale_shift")), 3.0 * normalized, atol=1e-4) + + +if __name__ == "__main__": + unittest.main() From df1788c64eb4ccb3aae68e5ccbbe95ef94ef5555 Mon Sep 17 00:00:00 2001 From: olehtika Date: Tue, 15 Sep 2026 16:41:45 +0000 Subject: [PATCH 5/5] Let FlaxTimestepEmbedding handle the conditioning dtype nn.Dense promotes its inputs to its own dtype, so casting the already-projected timestep and guidance to the pooled projection's dtype first only rounds them early: a no-op when activations are bfloat16, and a needless loss of precision when they are float32. CombinedTimestepTextProjEmbeddings already passes its projected timestep through untouched. --- src/maxdiffusion/models/embeddings_flax.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/src/maxdiffusion/models/embeddings_flax.py b/src/maxdiffusion/models/embeddings_flax.py index 9a2a9772e..e51308213 100644 --- a/src/maxdiffusion/models/embeddings_flax.py +++ b/src/maxdiffusion/models/embeddings_flax.py @@ -552,16 +552,15 @@ def __call__(self, timestep, guidance, pooled_projection=None): # (batch, frequency_embedding_size, embedding_dim) and break the sum with the # pooled projection below. The NNX variant owns its projection instead. timesteps_proj = timestep - dtype = pooled_projection.dtype if pooled_projection is not None else jnp.float32 timestep_emb = FlaxTimestepEmbedding( time_embed_dim=self.embedding_dim, dtype=self.dtype, weights_dtype=self.weights_dtype - )(timesteps_proj.astype(dtype)) + )(timesteps_proj) if self.guidance_embeds and guidance is not None: guidance_proj = guidance guidance_emb = FlaxTimestepEmbedding( time_embed_dim=self.embedding_dim, dtype=self.dtype, weights_dtype=self.weights_dtype - )(guidance_proj.astype(dtype)) + )(guidance_proj) time_guidance_emb = timestep_emb + guidance_emb else: time_guidance_emb = timestep_emb