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

Draft

amepas wants to merge 11 commits into

mainfrom

onboarding-imageedit-kv-flux2klein

Commits on Aug 14, 2026

feat(flux2klein): Onboard FLUX.2-Klein 4B and 9B with Flax NNX, fast weight loading, causal splash attention, and end-to-end optimizations
amepas
committed
3cb4455
Onboard FLUX.2-Klein multi-image editing support in pure Flax NNX (4B & 9B)
amepas
committed
284ff36

Commits on Aug 17, 2026

refactor(flux2klein): Unify text-to-image and multi-image editing into single pipeline and CLI
amepas
committed
1c4c1a9
cleanup(flux2klein): Remove redundant standalone image-edit files and obsolete dev tests
amepas
committed
c20a7dc
test(flux2klein): Add 4B/9B image edit smoke tests and rename E2E test
amepas
committed
622fc63
fix(flux2klein): Improve image normalization robustness and safe image loading
amepas
committed
e5303fe
fix(flux2klein): Broadcast reference image latents and RoPE IDs across batch dimension
amepas
committed
cb3912b

Commits on Aug 18, 2026

Commits on Aug 19, 2026

feat(flux2klein): preserve VAE reference image encoding latency tracking on KV cache branch
amepas
committed
1ada2c4

Commits on Aug 20, 2026

Add FLUX.2-Klein KV-cache image editing smoke test, e2e parity test, and clean attention block sizing
amepas
committed
709b764