Adding Image Editing to Flux2.Klein Models

[amepas]

Summary:

Adds Image Editing Support to the Flux2.Klein 4B and 9B models. Supports both single-image and multi-reference image editing (conditioning image generation on up to 4 reference images alongside textual prompts).

Additions

• Updated VAE: Added support for VAE Encoder to encode input images.
• Multi-Reference Conditioning: Accepts up to 4 input images via CLI.
• Zero T2I Regression: Verified zero accuracy loss or performance regressions on existing pipelines.
• Kernel Stability: Explicitly set use_base_exp2=True to eliminate mathematical precision loss when running the custom 2D Ring Ulysses attention kernel (ulysses_ring_custom_fixed_m).

Usage

Code is configured to re-use all possible pipeline components and existing entry points. Calling image editing only requires a flag pointing to the image file locations.

python generate_flux2klein.py \
  prompt="add a bunch of animals swimming in the water in front of the castle and get rid of all birds" \
  image_paths="['/path/to/img1.png', '/path/to/img2.png']"

Runtime:

Editing single image of size 1024 by 1024 with the Flux2.Klein 9B using the Optimal TPU-v7 T2I settings from PR #458.

Pipeline Stage	TPU v7-2 (Optimal T2I settings)	TPU v6-4 (TP=4)
VAE Encoding	24.28 ms	26.47 ms
Qwen3 Encoding	8.60 ms	7.35 ms
Denoising Loop	640.19 ms	1204.24 ms
VAE Decoding	31.11 ms	34.02 ms
Image Saving	107.76 ms	113.21 ms
Total Inference (Excl. Saving)	707.75 ms	1276.39 ms
Total E2E Pipeline	815.51 ms	1389.50 ms

Correctness:

E2E Parity Test against Diffusers pytorch baseline shows SSIM of 0.80 (both running bfloat16). Test included!

Visual Verification

Example 1:

Prompt: "add a bunch of animals swimming in the water in front of the castle and get rid of all birds"

Original:

Edited:

Example 2:

Prompt: "change the painting so she is facing forward instead of looking over her shoulder"

Original:

Edited:

Future Work

KV Cache not supported in this PR.

[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 support for the FLUX.2-Klein model in MaxDiffusion, adding the top-level NNXFlux2KleinTransformer2DModel and NNXAutoencoderKLFlux2 VAE models, along with image editing capabilities, concurrent AOT compilation, and end-to-end parity tests. The review feedback highlights several critical runtime issues that must be addressed, including missing imports (snapshot_download, Image, and time) and the use of non-existent Flax NNX API functions (nnx.to_flat_state and nnx.from_flat_state), which should be replaced with .flat() and nnx.State.from_flat() respectively.