fix(flux2): stop the sampler from advancing reference image tokens

[Pfannkuchensack]

Summary

fix(flux2): stop the sampler from advancing reference image tokens

FLUX.2 reference image conditioning (kontext_conditioning) is context: the encoded reference latents are concatenated onto the image stream for every forward pass, but the sampler must never advance them.

invokeai/backend/flux2/denoise.py concatenated them onto img once, before the denoising loop, and never sliced the model prediction back to the generated tokens. Both sampler paths then integrated the whole sequence:

• scheduler path: scheduler.step(model_output=pred, sample=img)
• manual Euler path: img = img + (t_prev - t_curr) * pred

So the reference latents were denoised along with the generated ones. After the first step the reference no longer is the encoded image — it drifts along the model's velocity field for the rest of the schedule, and feeds that drift back into the generated image at every step. The tensor shapes stay consistent throughout, so nothing ever failed; the conditioning just silently degraded.

Both reference implementations do it differently, and one of them is our own code:

• diffusers Flux2KleinPipeline rebuilds torch.cat([latents, image_latents]) inside the loop, slices noise_pred = noise_pred[:, : latents.size(1)], and steps only latents. image_latents stay pristine.
• our FLUX.1 Kontext path (invokeai/backend/flux/denoise.py) builds a temporary img_input inside the loop and slices pred = pred[:, :original_seq_len].

The FLUX.2 port lost that pattern. This PR restores it:

• position IDs for the concatenated sequence are precomputed once (they are constant)
• each step builds a temporary img_input, used by both the positive and the negative (CFG) forward pass
• pred and neg_pred are sliced to original_seq_len before CFG and before the sampler update
• img therefore only ever holds generated tokens, which also removes the split/recombine dances around the inpaint merge, the step previews, and the final slice

Behavior with no reference image is unchanged. With a single reference image the position-ID geometry is unchanged — only the contamination of the conditioning is gone.

One drive-by in the same block: in the manual Euler path the inpaint merge for the preview was written into a dead local (preview_gen) and thrown away. It now applies to preview_img, matching the scheduler path.

Related Issues / Discussions

Community report of FLUX.2 Klein reference-image edits coming out spatially displaced. This PR is not claimed to fix that report — see QA Instructions. It fixes a defect found while investigating it.

QA Instructions

Unit tests — tests/backend/flux2/test_denoise_ref_image_conditioning.py (new):

pytest tests/backend/flux2/ --no-cov

A fake transformer records its inputs. The core assertion is that the reference slice of hidden_states is bit-identical to the encoded latents at every step. Parametrized over both sampler paths, plus a CFG case (which also covers the negative-prediction slice) and a control case without reference images.

Verified the tests actually pin the bug: with denoise.py reverted to main, 4 of the 6 fail (all reference cases, both paths); the 2 no-reference control cases stay green.

End-to-end — FLUX.2 Klein 9B FP8 + FLUX.2 VAE (diffusers) + Qwen3 text_encoder, single reference image (1200×800), target 1200×800, 9 steps Euler, prompt asking for a haze/color change and explicitly no reframing. Four seeds (424242–424245), run with and without the fix, compared against the reference image by phase correlation (high-pass filtered, brute-forced over scale 0.90–1.20):

seed	with fix	without fix	delta
424242	0.2976	0.2686	+0.0290
424243	0.4795	0.4808	−0.0013
424244	0.3622	0.3843	−0.0220
424245	0.2940	0.3098	−0.0158
mean	0.3583 (SD 0.0867)	0.3609 (SD 0.0932)	−0.0025 (SD 0.0227)

No measurable difference in image quality, in either direction — the delta is smaller than its own spread and flips sign across seeds. Geometry is identical in all eight images (dx = 0, dy = 0/+1). This is a correctness fix, not a quality fix: the reference conditioning is supposed to stay constant across the schedule, and it now does.

What this PR does not establish: the reported ~20% horizontal displacement did not reproduce in any configuration tried — matching reference/target size, an oversized 1.88 MP reference, and an aspect-ratio mismatch (3:2 reference into a 16:9 target) all came out with < 1% horizontal offset. So the fix is justified on correctness grounds, not as a repro-and-fix of that report.

Merge Plan

Nothing special. Backend only, no schema or API changes, no frontend impact.

Out of scope, noted while investigating and worth a separate PR: with two or more reference images, invokeai/backend/flux2/ref_image_extension.py gives each additional reference a spatial h_offset/w_offset on top of the per-reference T offset. diffusers places every reference at H/W 0 and separates them by T alone; BFL instead concatenates the images into one picture and uses one latent. Our hybrid matches neither training convention. Untouched here because it changes output for every multi-reference user and needs a visual A/B, not a unit test. The single-reference path is unaffected (idx == 0 means both offsets are 0).

Checklist

• The PR has a short but descriptive title, suitable for a changelog
• Tests added / updated (if applicable)
• ❗Changes to a redux slice have a corresponding migration
• Documentation added / updated (if applicable)
• Updated What's New copy (if doing a release after this PR)

[the-space-fish]

Ahhh, I guess this probably explains why reference likeness in Klein was better preserved with fewer steps. Great find!