Adding KV-Cache Support to Flux2.Klein Image Editing - #465
Draft
amepas wants to merge 11 commits into
Draft
Conversation
…weight loading, causal splash attention, and end-to-end optimizations - Implemented Flax NNX Transformer architecture (NNXFlux2KleinTransformer2DModel) with support for both 4B (5 double / 20 single layers) and 9B (8 double / 24 single layers) configurations. - Integrated Flax Qwen3 text encoder with 3-layer intermediate hidden states extraction (layers 9, 18, 27), custom causal splash attention, and proper sharding constraints. - Implemented FlaxAutoencoderKL VAE decoder with fused batch normalization unscaling and channel re-layout. - Added fused end-to-end denoising loop scan with Flow Match Euler scheduler. - Added concurrent AOT XLA compilation across Qwen3, Flux transformer, and VAE. - Implemented fast host-memory streaming weight converter for safetensors shards directly into NNX State PyTree. - Optimized splash attention block sizes and Ulysses context parallelism sharding. - Added comprehensive unit tests (nnx_flux2klein_test.py) and end-to-end smoke test suite (generate_flux2klein_smoke_test.py).
… & 9B) - Implement NNXAutoencoderKLFlux2 (pure Flax NNX VAE encoder and decoder) with zero-copy weight loading and verified float32 parity - Implement multi-image conditioning helpers (prepare_multi_image_ids, prepare_image_latents, patchify/unpatchify) - Implement FlaxFlux2KleinImageEditPipeline for end-to-end multi-image editing inference - Implement generate_flux2klein_image_edit.py CLI runner and default YAML configs for 4B and 9B - Add comprehensive unit tests and parity test suite across all 7 implementation phases
…o single pipeline and CLI - Refactored FlaxFlux2KleinPipeline to conditionally handle multi-image editing when images are provided, eliminating duplicate pipeline definitions. - Unified generate_flux2klein.py and base configs (base_flux2klein.yml, base_flux2klein_9B.yml) to accept image_paths for multi-image conditioning. - Added test_flux2klein_4b_image_edit_smoke to generate_flux2klein_smoke_test.py. - Added test_flux2klein_image_edit_e2e_parity.py for end-to-end parity verification against PyTorch Diffusers reference.
… obsolete dev tests - Deleted separate YAML configs (base_flux2klein_image_edit.yml, base_flux2klein_image_edit_9B.yml). - Deleted separate CLI runner (generate_flux2klein_image_edit.py) and separate pipeline (flux2klein_image_edit_pipeline.py). - Removed obsolete intermediate parity/scratch scripts in favor of test_flux2klein_image_edit_e2e_parity.py and generate_flux2klein_smoke_test.py. - Cleaned up export in maxdiffusion.pipelines.flux.__init__.
- Renamed test_flux2klein_image_edit_e2e_parity.py to edit_flux2klein_e2e_test.py. - Added test_flux2klein_4b_image_edit_smoke and test_flux2klein_9b_image_edit_smoke in generate_flux2klein_smoke_test.py. - Added golden reference images ref_flux2klein_4b_image_edit.png and ref_flux2klein_9b_image_edit.png conditioned on ref_flux2klein_4b.png with prompt 'change the lighting to evening'. - Verified SSIM >= 0.95 assertion on both 4B and 9B image edit smoke tests on TPU VM.
…e loading - Replace img.max() > 1.0 heuristic with dtype-aware normalization in FlaxFlux2KleinPipeline - Handle integer uint8 and floating point ranges explicitly - Add try-except error handling for reference image loading in generate_flux2klein - Remove duplicate unpatchify_latents in models/flux/util.py - Apply pyink formatting and resolve ruff linter issues
…s batch dimension
…ial scaling on TPU
…tion, and E2E parity test
…ing on KV cache branch
![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.
Comment on lines
+88
to
+89
| 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.
Suggested change
| 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.
Suggested change
| 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.
Suggested change
| ref_temb_silu = nnx.silu(ref_temb) | |
| ref_temb_silu = jax.nn.silu(ref_temb) |
Comment on lines
+155
to
+159
| 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.
Suggested change
| 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) |
Comment on lines
+452
to
+454
| 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.
Suggested change
| 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) |
Comment on lines
+536
to
+538
| 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.
Suggested change
| 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) |
…and clean attention block sizing - Added test_flux2klein_9b_kv_image_edit_smoke in generate_flux2klein_smoke_test.py - Added flux2klein_kv_pipeline_e2e_test.py achieving 0.9037 SSIM and 23.30 dB PSNR cross-framework parity - Added test reference images in tests/images/flux2klein/ - Refactored _select_flash_block_sizes to mathematically bound block sizes without breaking cross-attention or large KV-cache sequences - Added hierarchical jax.named_scope annotations to Flux transformer blocks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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.