Adding Image Editing to Flux2.Klein Models - #461
Draft
amepas wants to merge 9 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
![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 FLUX.2-Klein model in MaxDiffusion, adding the top-level NNXFlux2KleinTransformer2DModel and NNXAutoencoderKLFlux2 VAE models, along with image editing capabilities, concurrent AOT compilation, and end-to-end parity tests. The review feedback highlights several critical runtime issues that must be addressed, including missing imports (snapshot_download, Image, and time) and the use of non-existent Flax NNX API functions (nnx.to_flat_state and nnx.from_flat_state), which should be replaced with .flat() and nnx.State.from_flat() respectively.
…nd AOT compilation
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:
Adds Image Editing Support to the Flux2.Klein 4B and 9B models. Supports both single-image and multi-reference image editing (conditioning image generation on up to 4 reference images alongside textual prompts).
Key Additions
use_base_exp2=Trueto eliminate mathematical precision loss when running the custom 2D Ring Ulysses attention kernel (ulysses_ring_custom_fixed_m).Usage
Code is configured to re-use all possible pipeline components and existing entry points. Calling image editing only requires a flag pointing to the image file locations.
Runtime:
Editing single image of size 1024 by 1024 with the Flux2.Klein 9B using the Optimal TPU-v7 T2I settings from PR #458.
TP=4)24.28 ms26.47 ms8.60 ms7.35 ms640.19 ms1204.24 ms31.11 ms34.02 ms107.76 ms113.21 ms707.75 ms1276.39 ms815.51 ms1389.50 msCorrectness:
E2E Parity Test against Diffusers pytorch baseline shows SSIM of 0.80 (both running bfloat16). Test included!
Visual Verification
Example 1:
Prompt: "add a bunch of animals swimming in the water in front of the castle and get rid of all birds"
Original:
Edited:
Example 2:
Prompt: "change the painting so she is facing forward instead of looking over her shoulder"
Original:
Edited:
Future Work
KV Cache not supported in this PR.