[inureyes]

Closes #2424

Summary

Ports FLUX.2 Klein (image-generation family flux2) from mlx-vlm models/flux2 (git 73d366ec) into src/imagegen/flux2/. It covers text-to-image and reference-image editing for the Klein 4B and 9B distilled and base checkpoints, the 9b-kv reference K/V cache, Diffusers bf16 trees, and MLX affine 4-bit exports. The family is registered in imagegen::FAMILIES and in mlxcel arch (image_generate runtime, image output), so both mlxcel generate-image and mlxcel-server --image-model (/v1/images/generations, /v1/images/edits) serve it.

• Pipeline. Ports of the empirical-mu scheduler, the Qwen-template prompt encoding (sequence-length buckets, layers 9/18/27), the double- and single-stream transformer with joint RoPE and modulation, the tiled VAE decode with cosine blending, and editing with reference latents at t = 10 + 10 * i. The 9b-kv path stores reference K/V on step 0 and uses a zero-timestep modulation blend for the reference tokens.
• Memory. The text encoder is evicted after each encode, and prompts are cached. The transformer and VAE load on first use and stay resident. A batch of n images is decoded one image at a time with single-image tiling, so n > 1 at 2048x2048 peaks at 0.55 GiB in the VAE instead of about 15 GiB per image.
• Batched generation. Follows mlx-vlm#2355: one noise tensor of batch n, and every image reports seed. Edits loop with seed + i.
• Edit size. An edit without a size uses upstream's _resolve_edit_dimensions (first reference scaled to at most 1 MP, sides floored to 16 within 256..=2048) on the CLI and /v1/images/edits, through a new ImageGenerationModel::default_edit_size hook whose default keeps the reference size for other families.
• --gen-kwargs. max_sequence_length (1 to 512, default 512) and tiled_vae.
• Bonsai seam (feat: Port Prism Bonsai (bonsai) ternary 2-bit FLUX.2 text-to-image #2428). The transformer builds its linears through a LinearSource trait, so a ternary source can reuse every block.
• Shared code, additive only.
  • latent::prepare_packed_latents and prepare_grid_ids.
  • Flux2Vae::bn_normalize and denormalize_packed_latents.
  • A default-accepting ImageGenerationModel::check_request hook. The image worker maps a failure there to 400 instead of 500, so FLUX.2's size range and option errors are client errors.

Precision

Component	dtype	f16 measurement (GB10, MLXCEL_CUDA_F16_NORMALIZE=1)
Transformer	bf16 on every backend	Does not hold: 4B generate decodes NaN latents (10.2 dB).
VAE	Flux2Vae::default_dtype (bf16 on CUDA, f16 on Apple Silicon)	Holds: VAE-only 54.76 dB tiled, 52.37 dB untiled.
Qwen3 text encoder	the Qwen3 loader's policy	Holds: with a bf16 transformer, generate 34.47 dB and edit 38.54 dB.

Edits run the transformer in f32 by promotion from the f32 reference latents, as upstream does.

Parity against the CUDA oracle

Oracle: mlx-vlm 0.7.8 at git 73d366ec, mlx 0.32.4, CUDA backend on this GB10 host, MLX_ENABLE_TF32=0. Fixtures are in tests/fixtures/flux2/ (about 2.6 MB), reproducible with generate_flux2_fixtures.py. Tests are in src/imagegen/flux2/parity_tests.rs (#[ignore], checkpoint-gated).

Check (4B bf16, 512x512, 4 steps)	Result
Initial noise, generate seed 9909 and edit seed 123	bit-identical
Transformer in f32 vs upstream in f32 (ModelConfig.precision = f32), identical inputs	bit-identical, 16384 of 16384 values
VAE decoding upstream's final latents, tiled / untiled	57.59 / 58.45 dB
Generate "A cinematic portrait of a glass bonsai tree on a moonlit desk", seed 9909	32.27 dB
Edit "Add black sunglasses", seed 123	40.54 dB
Final latents rel RMS, generate / edit	1.03e-1 / 4.11e-2

Deviation from the issue: the 1e-2 final-latent bound does not hold in bf16, and the test does not assert it. The first divergent tensor is the context embedder output, a single 7680-wide bf16 matmul on identical inputs and weights; the image embedder (128-wide) is bit-identical. The cause is reduction order in bf16 GEMMs, which differs between this build's MLX (pin 7ca40033) and mlx 0.32.4. Four Euler steps then amplify it. Upstream's own f32 run differs from its bf16 run by 7.1e-2 on the same case, and the f32 transformer check shows the port's graph is exact. The end-to-end test therefore asserts PSNR >= 30 dB and a 0.2 sanity bound on the latents, which a wrong port (about 1.4) cannot pass. Generate is about 2 dB over the gate, so a future MLX pin change can move it. TF32 is off in every run here.

Edit fixture and upstream's encoder. Upstream's _match_conv_layout leaves the VAE encoder conv_in weight (128, 3, 3, 3) in PyTorch order. The edit fixture was generated with that tensor transposed (--fix-conv-in, the default), which is what this port computes. edit_reference.png is itself an upstream FLUX.2-klein-4B generation (prompt "A studio photo portrait of a smiling man with short brown hair, plain grey background", seed 7, 512x512, 4 steps) written by the same script; it shows no real person, and the 4B checkpoint is Apache-2.0.

Fixture generator fix: mx.array.view reinterprets the physical buffer of a strided array, so the generator now reads bf16 bits through f32.

Real-checkpoint runs (CUDA, GB10)

Run	Result	What the PNG shows
CLI acceptance command, 4B bf16, 512x512, 4 steps, seed 9909	512x512 PNG, 9.0 s wall including load; 32.27 dB vs oracle	A leafy green bonsai in a glass tray on a wooden desk, full moon behind it on a dark starry background.
Server /v1/images/generations, same request	200, base64 PNG pixel-identical to the CLI output (PSNR inf)	Same image.
CLI and server /v1/images/edits, 4B, "Add black sunglasses", seed 123	pixel-identical to each other; 40.54 dB vs oracle	The reference portrait (smiling man, short brown hair, white t-shirt, grey backdrop) now wearing black wayfarer sunglasses.
Server n=2, 256x256	both images report seed 5, and they are distinct (11.5 dB apart)	Two different bonsai compositions with the moon.
4B 4-bit generate	35.47 dB vs upstream 4-bit	Bonsai in a clear glass bowl, moon upper left.
base-4B, resolves to 50 steps / guidance 4.0, run at 12 steps	32.80 dB vs upstream at 12 steps	A softer, glass-like bonsai on an ice-like base, moon behind it.
9B bf16 generate / 9B 4-bit generate	25.78 / 20.70 dB vs upstream	Glass bonsai with crystal leaves on a desk with books and a pen. Side by side with upstream it has the same composition; only fine detail (sparkle, a cable) differs.
9B edit / 9b-kv edit	44.39 / 42.94 dB vs upstream	Both show the same man wearing black sunglasses.
9b-kv generate, 1024x768 generate	succeed	Glass bonsai with the moon; a landscape bonsai with a candle and a large moon.
Tiled vs untiled decode of the same 512 latents	25.05 dB, equal to upstream's own 25.05 dB

The issue asks that the 9b-kv edit match the non-cached 9B edit within 35 dB. That cannot hold, because the 9b-kv checkpoint has its own weights (none of the transformer tensors compared are identical to 9B). Upstream's 9b-kv and 9B edits differ by 25.15 dB, and this build's differ by 25.09 dB. The meaningful check is 9b-kv against upstream 9b-kv: 42.94 dB.

Edge errors (CLI and server).

• Server status codes:
  • 4096x512 and 128x128 are 400 width must be in [256, 2048], got ....
  • 520x512 is 400 multiples of 16.
  • An edit without an image is 400.
  • max_sequence_length 0 or 600, and a non-boolean tiled_vae, are 400.
• CLI:
  • --task edit without --image errors.
  • A checkpoint directory without vae/ fails at load: FLUX.2 checkpoint is missing .../vae.

Tests

• New unit tests (they do not exist on main):
  • scheduler at 4 and 50 steps for seq 1024 and 5000, against oracle f32 bit patterns;
  • prompt buckets for lengths 1, 33 and 600;
  • the prepare_text_ids layout;
  • tiled vs whole decode on a local decoder;
  • kv extract and cached equivalence, and RoPE;
  • config, variant, quantization metadata and text-encoder config synthesis;
  • edit dimension rounding, --gen-kwargs parsing, and registration;
  • a server worker test that a check_request failure is 400 and never runs the model.
• Proof against the missing integration: with the FAMILIES row removed, four tests fail:
  • flux2_is_registered_in_both_tables
  • every_shipping_row_is_a_probe_family
  • load_reports_a_flux2_checkpoint_missing_a_component (family flux2 is not supported by this build)
  • no_registered_family_claims_image_generation_before_a_model_ports_it

Verification

All commands ran with MLX_CUDA_ARCHITECTURES=121 and MLX_ENABLE_TF32=0. Test binaries ran under gpu-lock with --test-threads=1.

• cargo clippy --profile test-fast --features cuda -p mlxcel --lib --bins --tests --no-deps -- -D warnings: clean.

• cargo fmt --check: clean.

• Lib test binary filters:

  Filter	Passed	Ignored
  imagegen::	167	8
  flux2	35	6
  server::image	13	0
  server::routes::images	18	0
  models::registry	16	0
  flux2::parity_tests --ignored (real checkpoints)	6	0

• Bin test binary: generate_image 15 passed.

Changes during review

• Edit size: resolve_edit_dimensions was ported but unused; it is now wired through ImageGenerationModel::default_edit_size on the CLI and /v1/images/edits (a 1000x667 reference gives a 992x656 output). For every family, an edit without a size is now size-checked after model load instead of before.
• Security: a batched request decoded the whole batch untiled (upstream behaviour), about 15 GiB per 2048x2048 image, so n = 10 could exhaust host memory. Flux2Pipeline::decode now decodes per image; a checkpoint-gated test asserts the batched decode is byte-identical to decoding each image alone. For n > 1 this differs from upstream's untiled batch decode by the tiled vs untiled gap (57.59 vs 58.45 dB VAE-only).
• The transformer no longer reads the timestep back from the device on every forward call to choose the sigma scale.
• MLXCEL_TEST_FLUX2_DIR (test-only checkpoint path) is allowlisted in scripts/ci/check_env_var_docs.py.
• Open MEDIUM findings recorded in the security comment, not fixed here: absolute checkpoint paths in 500 responses for load failures (shared behaviour across routes), no per-request compute budget beyond the shared caps, and text-encoder reload per new prompt (matches upstream eviction).

Not covered here

• Metal and ROCm runs are not recorded. The arch row says Metal partial.

[inureyes]

Implementation Review Summary

Intent

Port FLUX.2 Klein (flux2) text-to-image and editing from mlx-vlm 73d366ec, served by mlxcel generate-image and mlxcel-server --image-model.

Findings Addressed

• flux2::resolve_edit_dimensions was orphan code: an edit with no size used the reference's native size, so FLUX.2 rejected a 1000x667 or 4000x3000 photo where upstream scales it to 1 MP and rounds to 16. It is now wired through ImageGenerationModel::default_edit_size (default: reference size, so other families keep their behavior) on both the CLI and /v1/images/edits. Smoke run on 4B with a 1000x667 reference and no size: a 992x656 PNG showing the portrait with sunglasses added. 4B edit parity is unchanged (40.54 dB). Fixed in d24f442. (HIGH)
• The trait docs in src/imagegen/pipeline.rs said seed + i unconditionally; they now note the FLUX.2 batched-seed exception. Fixed in d24f442. (MEDIUM)
• Fixture provenance was only in the generator docstring. The PR body and the docs Parity section now say that edit_reference.png is an upstream FLUX.2-klein-4B generation (seed 7, no real person) and that the 4B checkpoint is Apache-2.0, so committing it is license-safe. The fixtures total 2.6 MB for four 512x512 parity cases, which is acceptable. (MEDIUM)

Remaining Items

• run::generate_to_files is now called only by its own tests; the docs run row still lists it. Harmless. (LOW)
• generate_flux2_fixtures.py hardcodes the 4B variant, so the committed script cannot reproduce the 9B, 9b-kv, 4-bit and base PSNRs in the body. They are reported numbers, not gates. (LOW)

Acceptance deviations

Both are sound and documented honestly.

• Final-latent rel RMS < 1e-2. Three checks isolate the cause to bf16 arithmetic between MLX pins: the f32 transformer is bit-exact on upstream's own embeddings, the VAE scores 57 to 58 dB on upstream's latents, and the initial noise is bit-identical. Upstream's own f32 and bf16 runs differ by 7.1e-2. The 0.2 bound is labeled as a sanity bound in the test, the docs and the PR body.
• 9b-kv within 35 dB of the 9B edit. Upstream fails this against itself (25.15 dB), because the checkpoints have different weights. The replacement, 9b-kv against upstream 9b-kv, scores 42.94 dB, and the port's 9B vs 9b-kv gap matches upstream's (25.09 dB).

Other checks:

• The upstream conv_in claim holds. The checkpoint stores encoder.conv_in.weight as [128, 3, 3, 3] in PyTorch order (decoder.conv_in is [512, 32, 3, 3]), and _match_conv_layout cannot tell the two layouts apart at that shape.
• 9B generate scores 25.78 dB, below the 4B gate. 9B edits run the same weights in f32 by promotion and score 44.39 dB, which shows the 9B graph is right.

Verification

• All stated requirements implemented
• No placeholder/mock code remaining
• Integrated into project code flow
• Project conventions followed
• Existing modules reused where applicable (dit_blocks RoPE and timestep embedding are not reused: they differ in layout and precision, and the LinearSource seam feat: Port Prism Bonsai (bonsai) ternary 2-bit FLUX.2 text-to-image #2428 needs is separate)
• No unintended structural changes
• Tests pass at d24f442: lib imagegen:: 167 passed, 7 ignored; server:: 3768 passed; models::registry 16 passed; bin generate_image 15 passed; clippy and fmt clean (CUDA, GB10, MLX_ENABLE_TF32=0)

Also checked against upstream, no change needed:

• the scheduler (f64 mu, f32 tables) and prepare_packed_latents draw order;
• CFG with the " " negative prompt;
• the edit token order on both the kv and non-kv paths, with t = 10 + 10i;
• the modulation blends;
• the edit defaults (unbucketed, untiled), matching upstream's edit model;
• tiling, line for line;
• variant resolution;
• both registry rows;
• the 400/500 split.

[inureyes]

Security and performance review

GB10, CUDA, MLX_ENABLE_TF32=0, flux.2-klein-4b-bf16.

Fixed in c75cffd

• HIGH: batched generation could exhaust host memory. decode_image_tiled skips tiling for batch != 1 (as upstream), so n > 1 decoded the whole batch at full resolution. Measured peak at 2048x2048: 30 GiB for n = 2, so n = 10 needs about 150 GiB. Each image is now decoded alone with single-image tiling: 0.55 GiB for the same request. New ignored test checks batched and per-image decodes are byte-identical. For n > 1 this departs from upstream's untiled batch decode by the tiled/untiled gap (57.59 vs 58.45 dB); documented.
• LOW: per-forward device read of the timestep to pick the x1000 sigma scale; the host value now decides it.

Parity numbers unchanged (generate 32.27 dB, edit 40.54 dB, f32 transformer bit-identical). Unit filters flux2, imagegen::, server::image, server::routes::images pass; clippy and fmt clean.

Checked, no change needed

• Edit attention runs in bf16 at head dim 128 (probed), so cuDNN's fused kernel applies; a 2048x2048 one-reference edit completes in 61 to 82 s. VAE mid-block attention (512-wide head) is not the memory peak: 12 GiB encode and 15 GiB decode at 2048x2048 with or without query chunking.
• Reference images use the shared media limits and are resized to the validated 256..=2048 size; max_sequence_length and tiled_vae are validated (400); prompts truncate at 512 tokens; worker panics are caught. Shard/index parsing is pre-existing shared code on an operator-chosen directory.

Remaining (below the auto-fix threshold)

• MEDIUM: lazy component loads run inside requests, so a load failure returns its absolute checkpoint path in the 500 body. Forwarding Internal text is existing, test-asserted route behavior shared with embeddings, rerank and audio; worth one decision across all of them.
• MEDIUM: no per-request cost budget beyond the shared caps (16 references at 2048x2048 is about 279k image tokens per transformer call; up to 1000 steps with CFG), and the timeout does not cancel MLX work. A token budget in check_request would bound it.
• MEDIUM: the text encoder (16 GB for 9B) reloads for every unseen prompt, as upstream's eviction does.
• LOW: edits with n > 1 re-encode the same references per image; 9b-kv step 0 evals cached K/V arrays one by one; a directory name containing kv selects 9b-kv (upstream heuristic).