Flux2klein Onboarding

[amepas]

This branch contains all of the code for onboarding the Flux2.Klein model into this repo. PR contains support for both the 4B and 9B model variants.

Images can be generated with:
python src/maxdiffusion/generate_flux2klein.py src/maxdiffusion/configs/base_flux2klein.yml

For the 9B model, point to the corresponding .yml file
python src/maxdiffusion/generate_flux2klein.py src/maxdiffusion/configs/base_flux2klein_9B.yml

Running this implementation on a v7-1 with FSDP=2 on Batch-size 8 (1024 by 1024) gives the following latency:

• Model Loading: 546.69s
• Dry-Run/Compilation: 73.94s
• Text Embedding (Qwen3): 0.32s
• Denoising: 2.96s
• VAE: 0.24s
• E2E Latency: 3.51s

Some call-outs on the implementation:

• src/maxdiffusion/generate_flux2klein.py is the main user-facing file for running either the 4B and 9B files. It invokes src/maxdiffusion/pipelines/flux/flux2klein_pipeline.py.
• src/maxdiffusion/models/flux/transformers/transformer_flux_flax.py is heavily modified to add support for the new variants of the Attention blocks used in the Flux2.Klein models over the older generation of models.
• We add a custom Jax implementation of the Qwen3 model: src/maxdiffusion/models/qwen3_flax.py
• We add two sets of tests: a smoke test to check if the implementation has changed in any way, and a e2e parity test that generates images from the Pytorch diffusers library and compares the generated outputs from our new implementation for the 4B and 9B models.

[github-actions]

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

[eltsai]

I still see a few print statement instead of max_logging, but I think it's fine for now, we can send out a small PR for that later. RN let's prioritize merging the model onboarding PR.

I was able to reproduce the same e2e generation time as reported.

[github-actions]

🤖 Hi @eltsai, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

[github-actions]

🤖 I'm sorry @eltsai, but I was unable to process your request. Please see the logs for more details.

[eltsai]

Seems like the linter test failed because of #443 , but #443 itself passed the linter tests, @syhuang22 is investigating it now

[entrpn]

left some comments + please change all print statements to max_logging instead.

[eltsai]

I see a import/copybara error, this might be caused by missing license file in src/maxdiffusion/tests/nnx_flux2klein_test.py