Fix three FLUX regressions from the Flux.2-Klein onboarding

[olehtika]

FLUX-dev training is broken on main. It fails at transformer init with:

ValueError: Incompatible shapes for broadcasting: shapes=[(8, 256, 3072), (8, 3072)]
  .../models/embeddings_flax.py, conditioning = time_guidance_emb + pooled_projections

Investigating it turned up three separate problems, all introduced by 64450df ("Onboard Flux.2-klein model family (4B/9B) into MaxDiffusion (JAX+TPU)"). Flux.2-Klein reads two shared modules with its own conventions, and FLUX was left reading them with the old ones. The third problem produces no error at all, just wrong numbers.

1. The timestep and guidance projections run twice

FluxTransformer2DModel.__call__ projects both values through its own timestep_embedding(..., 256) before calling CombinedTimestepGuidanceTextProjEmbeddings, which now projects them again. That shapes the embedding (batch, frequency_embedding_size, embedding_dim) and breaks the sum with the pooled projection, which is the crash above. Note that the transformer's own helper applies a time_factor the sinusoidal helper in the module does not, so the caller's projection is the one to keep.

This module now takes the already-projected values, as it did before. NNXCombinedTimestepGuidanceTextProjEmbeddings keeps its internal projection, which is what Flux.2-Klein uses, and the non-guidance CombinedTimestepTextProjEmbeddings was already a pass-through.

2. The original-format loader writes to a submodule name that no longer exists

The Dense inside AdaLayerNormContinuous gained name="linear", so the params tree carries norm_out/linear. load_flow_model still rewrote final_layer.adaLN_modulation_1 to norm_out.Dense_0, so loading pretrained transformer weights (train_new_flux=False) built a tree that no longer matched the sharding tree:

ValueError: device_put device specification must be a tree prefix of the corresponding value
  pytree structure error: different pytree metadata at key path
    device_put device.params['norm_out']
    - ('linear',)
    + ('Dense_0',)

validate_flax_state_dict did notice, logging key: ('norm_out', 'linear', 'bias') not found..., but it only logs, so the run continued to the opaque device_put failure. The diffusers-format paths in util.py were already updated to norm_out/linear; this is the one that was missed.

3. FLUX silently switched to Klein's shift/scale convention

AdaLayerNormContinuous gained a scale_shift_order option defaulting to "scale_shift". Flux.2-Klein's Flax transformer passes it explicitly; FLUX's norm_out passed nothing and so silently changed convention, reading the modulation's first half as scale. Weights loaded from the original FLUX checkpoint emit shift first, so shift and scale were swapped.

Nothing fails, which is what makes this one worth attention. Measured on FLUX-dev with pretrained weights, loss lands about 8x too high:

	loss, steps 0-5
with shift_scale	20864, 12288, 12544, 16064, 9856, 8512
with the inherited scale_shift	174080, 101888, 108032, 153600, 103936, 100864

FLUX now states the order at the call site, so every user of the Flax module is explicit. The module default is untouched, so Flux.2-Klein is unaffected.

Validation

Hardware was 8x MI355X, so these are the ROCm numbers rather than TPU ones, and the comparison baseline is 68e0696 (six months of main behind this change, and the last commit where FLUX training was known to work here).

FLUX-dev, synthetic data, 512 resolution, batch 14/device, 20 steps, train_new_flux=False:

	step time	loss, steps 0-5
68e0696 baseline	1.555-1.571 s	18432, 10816, 11712, 17536, 12160, 11072
this PR	1.563-1.855 s	20864, 12288, 12544, 16064, 9856, 8512

The losses are not expected to match exactly, since 200+ commits separate the two, but they sit in the same regime, where the unfixed order does not.

train_new_flux=True also trains, which is worth noting because it is the path the README's FLUX benchmark command uses: it skips the pretrained load, so problem 2 never fires there and only problem 1 blocks it.

WAN is unaffected. None of the changed files' symbols are referenced under models/wan, pipelines/wan, or the WAN trainers.

Tests

src/maxdiffusion/tests/flux_conditioning_test.py runs on CPU in about 7 seconds with random weights, no hub downloads. The double-projection case fails on main and passes here; the other two pin the conventions the original-format loader depends on.

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[gemini-code-assist]

Code Review

This pull request updates the FLUX conditioning path to align with original checkpoint conventions. Specifically, it avoids re-projecting already projected timesteps and guidance embeddings, configures the normalization module to use the 'shift_scale' order, updates weight loading key mappings, and adds corresponding unit tests. Feedback suggests removing redundant dtype casting on the already projected timestep and guidance inputs to avoid unnecessary overhead.

[gemini-code-assist]

Since timestep and guidance are now passed already projected (and thus already in the correct precision/dtype of the model), casting them to dtype (which defaults to float32 if pooled_projection is None) is redundant and can cause unnecessary precision casting overhead. We can simplify this by removing the dtype definition and passing timesteps_proj and guidance_proj directly to FlaxTimestepEmbedding, matching the behavior of CombinedTimestepTextProjEmbeddings.

    timesteps_proj = timestep
    timestep_emb = FlaxTimestepEmbedding(
        time_embed_dim=self.embedding_dim, dtype=self.dtype, weights_dtype=self.weights_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)

[olehtika]

Taken in e8170b1, thanks. One correction to the reasoning, and some measurements.

The inputs are not necessarily already in the model's dtype. The caller's timestep_embedding() returns t.dtype, which comes from the scheduler in the training step and from t_vec = jnp.full(bs, 0, dtype=self.dtype) at init, so it varies. The cast is unnecessary for a different reason: nn.Dense promotes its inputs to its own dtype, so FlaxTimestepEmbedding already handles this.

Measured on FlaxTimestepEmbedding with identical params, feeding a float32 projection:

activations / pooled	output dtype	bit-identical without the cast
bfloat16 / bfloat16	unchanged	yes
float32 / bfloat16	unchanged	no
bfloat16 / float32	unchanged	yes

So it is a no-op in the usual bfloat16 setup, and in the float32-activations case the pre-cast was rounding to bfloat16 ahead of a float32 matmul, which is the one configuration where this change alters anything, for the better.

End to end on FLUX-dev with pretrained weights, 20 steps, 8x MI355X: the first six steps are identical to the run before this change, and 7 of the 19 logged steps differ by at most 1%, starting at step 6. Given the module-level result above, those cannot come from this edit; the two runs were on different machines, and bfloat16 drift accumulates. For contrast, a real semantic change in this path shows up immediately and large: the scale_shift_order fix in this PR moves step 0 by 8x.

Left the NNX variant's casts alone, since that is Flux.2-Klein's path and outside this PR.

[entrpn]

@amepas can you please take a look.