Adding KV-Cache Support to Flux2.Klein Image Editing - #465
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
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.
| try: | ||
| fb_dir = snapshot_download(repo_id=repo_id, local_files_only=True) |
There was a problem hiding this comment.
snapshot_download is called here but is not imported in this scope (it is only imported inside main). This will cause a NameError when encode_prompt is executed. Please import snapshot_download.
try:
from huggingface_hub import snapshot_download
fb_dir = snapshot_download(repo_id=repo_id, local_files_only=True)| def __call__(self, x: jax.Array) -> jax.Array: | ||
| x = self.linear_in(x) | ||
| x1, x2 = jnp.split(x, 2, axis=-1) | ||
| hidden = nnx.silu(x1) * x2 |
There was a problem hiding this comment.
nnx.silu is not a valid attribute of flax.nnx. Please use jax.nn.silu instead.
| hidden = nnx.silu(x1) * x2 | |
| hidden = jax.nn.silu(x1) * x2 |
| temb = temb.astype(hidden_states.dtype) | ||
|
|
||
| temb_silu = jax.nn.silu(temb) | ||
| temb_silu = nnx.silu(temb) |
There was a problem hiding this comment.
nnx.silu is not a valid attribute of flax.nnx. Please use jax.nn.silu instead.
| temb_silu = nnx.silu(temb) | |
| temb_silu = jax.nn.silu(temb) |
| num_img_tokens = hidden_states.shape[1] - num_ref_tokens | ||
| ref_timestep = jnp.full_like(timestep_scaled, ref_fixed_timestep * 1000.0) | ||
| ref_temb = self.time_text_embed(ref_timestep, guidance_scaled, pooled_projections).astype(hidden_states.dtype) | ||
| ref_temb_silu = nnx.silu(ref_temb) |
There was a problem hiding this comment.
nnx.silu is not a valid attribute of flax.nnx. Please use jax.nn.silu instead.
| ref_temb_silu = nnx.silu(ref_temb) | |
| ref_temb_silu = jax.nn.silu(ref_temb) |
| h = self.norm1(x) | ||
| h = nnx.silu(h) | ||
| h = self.conv1(h) | ||
| h = self.norm2(h) | ||
| h = nnx.silu(h) |
There was a problem hiding this comment.
nnx.silu is not a valid attribute of flax.nnx. Please use jax.nn.silu instead.
| h = self.norm1(x) | |
| h = nnx.silu(h) | |
| h = self.conv1(h) | |
| h = self.norm2(h) | |
| h = nnx.silu(h) | |
| h = self.norm1(x) | |
| h = jax.nn.silu(h) | |
| h = self.conv1(h) | |
| h = self.norm2(h) | |
| h = jax.nn.silu(h) |
| x = self.conv_norm_out(x) | ||
| x = nnx.silu(x) | ||
| x = self.conv_out(x) |
There was a problem hiding this comment.
nnx.silu is not a valid attribute of flax.nnx. Please use jax.nn.silu instead.
| x = self.conv_norm_out(x) | |
| x = nnx.silu(x) | |
| x = self.conv_out(x) | |
| x = self.conv_norm_out(x) | |
| x = jax.nn.silu(x) | |
| x = self.conv_out(x) |
| x = self.conv_norm_out(x) | ||
| x = nnx.silu(x) | ||
| x = self.conv_out(x) |
There was a problem hiding this comment.
nnx.silu is not a valid attribute of flax.nnx. Please use jax.nn.silu instead.
| x = self.conv_norm_out(x) | |
| x = nnx.silu(x) | |
| x = self.conv_out(x) | |
| x = self.conv_norm_out(x) | |
| x = jax.nn.silu(x) | |
| x = self.conv_out(x) |
…X (4B & 9B) with E2E parity and smoke tests - Add multi-image editing pipeline (FlaxFlux2KleinPipeline) supporting arbitrary reference image conditioning - Port VAE encoder and decoder to Flax NNX (AutoencoderKLFlux2NNX) - Implement efficient image preprocessing, tiling, aspect-ratio-aware padding, and packing utilities - Support fast sharded safetensors loading and TP sharding for Klein 4B & 9B - Use standard HF_HOME resolution without machine-specific hardcoded paths - Ensure clean compatibility across Transformers 4.x and 5.x via lazy module loading and dynamic FlaxPreTrainedModel lookup - Add end-to-end multi-image editing parity test and preprocessing unit test suite
…age editing with E2E parity and smoke tests - Implement prefix extraction phase (step 0) and cached denoising scan loop in FlaxFlux2KleinPipeline - Support KV cache slicing and concatenation across Flax NNX Double and Single transformer blocks - Safely bound Splash/Flash attention block sizes for asymmetric cross-attention sequences - Add CLI and config support for use_kv with dynamic FLUX.2-klein-9b-kv repository resolution - Add cross-framework E2E parity test achieving 0.9037 SSIM / 23.30 dB PSNR - Add 9B KV-cache image editing smoke test with verified golden reference image
709b764 to
6238524
Compare
1 participant
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.
Usage
Just adding the use_kv=True flag to the 9B model call will switch to the KV-cache model support.
Speed-Ups
Block-sizes can be tuned for each setting. Using default settings still gives substantial improvements
Visual Verification
Prompt: change the painting so she is facing forward instead of looking over her shoulder
Original:
Edited:
Correctness
E2E parity test against reference diffusers implementation shows high visual similarity (SSIM>0.9). Smoke test also included to prevent any implementation regressions in future PRs.