feat(flux2klein): Faster Loading + Compliation

[amepas]

Summary
Significantly improving the model loading and compilation times for the Flux Klein models. This enables faster start-up and lets us iterate faster on model improvements. We also include bug fixes in Flux2.Klein implementation (handling of batch-sizes, making sure all code goes through NNX modules rather than linen).

Optimizations

1. Loading weights directly into bfloat16: Previous benchmarking of implementation required us to be able to run
   the model in fp32, but this is not needed anymore.
2. Parallelizing compilation of Qwen3 text encoder, denoising steps, vae
3. Use only 1 step of the denoiser instead of 4
4. Fixes to VAE: replaced jax.image.resize(hidden_states, shape=(batch, height * 2, width * 2, channels), method="nearest") with 6D broadcast + reshape 2x upsampling . This saved a lot of time in compilation
5. Latent tensor sharding: Ensuring the sharding is maintained across the layers so that there is no re-compilation or re-sharding that needs to be triggered when we do a warm-up pass.

None of these changes affect the model implementation, so the outputs are still identical.

Results
On a v6-4, we see the following speedups (BS=8, fsdp) for the 9B model:

• Warm Model Loading: 154s -> 12s
• XLA compilation/warmup-pass: 61s -> 11s
• Inference: unchanged

On a v7-8:
-Warm model loading: 540s -> 140s (for cold loads, use a Hyperdisk mount)
-XLA compilation/warmup-pass: 82s -> 49s
-Inference unchanged

So, on the v6-4: We bring the loading+compilation time down to <30s. On a v7-8, <4 min. The main bottleneck is the hardware constraint of how fast it can load the model.

Correctness
Visual results are high-quality, SSIM changes from 1.00 to 0.9 (still passing all parity checks against PyTorch). Slight changes to SSIM are expected given that using NNX will lead to different graphs than linen, and thus cause different handling of floats. Reference images updated to maintain 1.00 sim on the repo.

[github-actions]

e2e testgrid: https://8bcf50593faf4ea38060e236169827e5-dot-us-central1.composer.googleusercontent.com/dags/maxdiffusion_tpu_e2e/grid

[gemini-code-assist]

Code Review

This pull request introduces several performance optimizations and robustness improvements to the FLUX.2-Klein pipeline, including concurrent Ahead-of-Time (AOT) compilation of XLA graphs, fused VAE decoding, direct target-dtype weight loading, and nearest-neighbor upsampling using JAX broadcasting. It also adds support for multiple inference repetitions, profiling, and fallback configurations when local files are missing. The review feedback highlights two critical issues: first, removing the return statement in partition_prompts causes the function to return None when prompt truncation is triggered; second, both the try and except blocks for snapshot downloading use local_files_only=True, which will fail completely if the model is not already cached locally instead of falling back to an online download.

[entrpn]

left some comments.