Adding KV-Cache Support to Flux2.Klein Image Editing

[amepas]

Summary

Speeding up Image Editing on Flux2.Klein models by onboarding Flux2.Klein-9B-KV model variant. This is a separate model (same architecture, different weights) that was trained to support caching KV for reference images during image editing.

Builds on PR 461 so some overlapping code changes are included in this PR and its listed as two commits until 461 is merged.

Usage

Just adding the use_kv=True flag to the 9B model call will switch to the KV-cache model support.

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']" \
  use_kv=True

Speed-Ups

Block-sizes can be tuned further for each setting. Using default settings still gives substantial improvements (using CP=2 on a TPU v7-2)

Reference Images ($1024 \times 1024$)	Latency Speedup
1 Image	1.28x
2 Images	1.32x
4 Images	1.74x

Visual Verification

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

Original:

Edited:

Correctness

KV-Cache implementation does not cause any regression or change in behavior on other flux2.klein pipelines.
Below numbers generated on v6-4.

• 4B T2I Smoke Test: 0.982 SSIM from baseline
• 9B T2I Smoke Test: 0.881 SSIM from baseline
• 4B Image Edit Smoke Test: 0.996 from baseline
• 9B Image Edit Smoke Test: 0.997 from baseline
• 9B KV-Cache Image Edit Smoke Test: 0.999 from baseline
• Diffusers KV-Cache Image Edit E2E Parity Test: 0.889 from PyTorch Diffusers

Note: KV-Cache requires a separate model checkpoint/training setup and is only supported on the 9B model.

[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 Flux2Klein model, including new configuration files, NNX-based model implementations for the Transformer and VAE, and a dedicated pipeline. It also adds end-to-end parity and smoke tests. The review identified several critical issues: missing imports in flux2klein_pipeline.py and generate_flux2klein.py that will cause runtime errors, and the incorrect use of nnx.silu instead of jax.nn.silu across multiple model files.

[mbohlool]

looks great with a couple of nit comments

[mbohlool]

Nit: find_model_path() contains hardcoded machine paths (/mnt/hyperdisk_weights/..., /mnt/data/...). We can use get_model_snapshot_dir() (as in edit_flux2klein_e2e_test.py) for clean HF_HOME resolution.

[mbohlool]

Nit: import flax, import flax.linen as nn, and nn_partitioning appear to be unused

The merge-base changed after approval.