Skip to content

Adding KV-Cache Support to Flux2.Klein Image Editing #465

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
amepas wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Adding KV-Cache Support to Flux2.Klein Image Editing #465

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
e86be47
feat(flux2klein): Onboard multi-image editing support in pure Flax NN…
amepas Aug 26, 2026
6121acd
feat(flux2klein): Add KV-cache acceleration for FLUX.2-Klein multi-im…
amepas Aug 26, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 20 additions & 0 deletions README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
[![Unit Tests](https://github.com/AI-Hypercomputer/maxdiffusion/actions/workflows/UnitTests.yml/badge.svg)](https://github.com/AI-Hypercomputer/maxdiffusion/actions/workflows/UnitTests.yml)

# What's new?
- **`2026/08/28`**: Flux2.Klein text to image and image editing (w/ KV Cache) is now supported.
- **`2026/07/14`**: Automatic attention tile-size (`block_q`/`block_kv`) search for Wan is now supported.
- **`2026/06/26`**: 2D ring (USP) attention with a custom splash kernel is now supported for Wan (`ulysses_ring_custom`), splitting context parallelism into an intra-chip Ulysses axis and a cross-chip ring axis.
- **`2026/04/16`**: Support for Tokamax Ring Attention kernel is now added.
Expand Down Expand Up @@ -49,6 +50,7 @@ MaxDiffusion supports
* Stable Diffusion 2.1 (training and inference)
* Stable Diffusion XL (training and inference).
* Flux Dev and Schnell (Training and inference).
* Flux.2-Klein 4B & 9B (text-to-image and multi-image editing with KV-Cache).
* Stable Diffusion Lightning (inference).
* Hyper-SD XL LoRA loading (inference).
* Load Multiple LoRA (SDXL inference).
Expand Down Expand Up @@ -759,6 +761,8 @@ The optimal attention tile sizes (`block_q` / `block_kv`) depend on the sequence

Flux.2-Klein provides ultra-fast 4-step image generation using Qwen3 text embeddings and FLUX.2 transformer blocks.

#### Text-to-Image Generation:

Flux.2-Klein 4B:

```bash
Expand All @@ -770,6 +774,22 @@ The optimal attention tile sizes (`block_q` / `block_kv`) depend on the sequence
```bash
python src/maxdiffusion/generate_flux2klein.py src/maxdiffusion/configs/base_flux2klein_9B.yml run_name=flux2klein_9b prompt="A detailed vector illustration of a robotic hummingbird"
```

#### Multi-Reference Image Editing:

Flux.2-Klein supports multi-reference image editing conditioned on up to 4 reference images via the `image_paths` CLI flag.

Flux.2-Klein 9B Image Editing:

```bash
python src/maxdiffusion/generate_flux2klein.py src/maxdiffusion/configs/base_flux2klein_9B.yml run_name=flux2klein_9b_image_edit prompt="change the lighting to evening" image_paths="['src/maxdiffusion/tests/images/flux2klein/ref_flux2klein_9b.png']"
```

The 9B model also supports KV-Cache for faster inference, and can be toggled with the `use_kv=True` CLI flag:

```bash
python src/maxdiffusion/generate_flux2klein.py src/maxdiffusion/configs/base_flux2klein_9B.yml run_name=flux2klein_9b_kv_edit prompt="change the lighting to evening" image_paths="['src/maxdiffusion/tests/images/flux2klein/ref_flux2klein_9b.png']" use_kv=True
```
## Fused Attention for GPU:
Fused Attention for GPU is supported via TransformerEngine. Installation instructions:

Expand Down
6 changes: 5 additions & 1 deletion src/maxdiffusion/configs/base_flux2klein.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,8 @@ metrics_file: "" # for testing, local file that stores scalar metrics. If empty,
write_metrics: True

timing_metrics_file: "" # for testing, local file that stores function timing metrics such as state creation, compilation. If empty, no metrics are written.
write_timing_metrics: True
write_timing_metrics: False
timing: False

gcs_metrics: False
# If true save config to GCS in {base_output_directory}/{run_name}/
Expand All @@ -40,6 +41,9 @@ max_sequence_length: 512
time_shift: True
base_shift: 0.5
max_shift: 1.15
image_paths: []
use_base2_exp: True
use_kv: False


unet_checkpoint: ''
Expand Down
6 changes: 5 additions & 1 deletion src/maxdiffusion/configs/base_flux2klein_9B.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,8 @@ metrics_file: "" # for testing, local file that stores scalar metrics. If empty,
write_metrics: True

timing_metrics_file: "" # for testing, local file that stores function timing metrics such as state creation, compilation. If empty, no metrics are written.
write_timing_metrics: True
write_timing_metrics: False
timing: False

gcs_metrics: False
# If true save config to GCS in {base_output_directory}/{run_name}/
Expand All @@ -40,6 +41,9 @@ max_sequence_length: 512
time_shift: True
base_shift: 0.5
max_shift: 1.15
image_paths: []
use_base2_exp: True
Comment thread
amepas marked this conversation as resolved.
use_kv: False


unet_checkpoint: ''
Expand Down
Loading
Loading