fix(flux2): img2img posterization - noise scaling and sigma schedule

[Pfannkuchensack]

Summary

fix — FLUX.2 img2img and inpainting posterize photographic detail into flat patches, progressively worse the higher the denoise strength. Two bugs in the same code path, the second masked by the first.

1. The img2img noise was divided by bn_std

The FLUX.2 transformer operates on BN-normalized latents, while flux2_vae_encode emits raw ones. Flux2DenoiseInvocation computed the rectified-flow preblend in raw space and BN-normalized the result afterwards:

x = t_0 * noise + (1.0 - t_0) * init_latents   # raw space
...
x = self._bn_normalize(x, bn_mean, bn_std)     # (x - bn_mean) / bn_std

Expanded, that is t_0 * (noise - bn_mean) / bn_std + (1 - t_0) * normalize(init_latents) — the noise term gets divided by bn_std too. Measured on the BFL FLUX.2 VAE, bn.running_var is 2.91…3.47 uniformly across all 128 packed channels, so bn_std ≈ 1.77 and 1 / bn_std ≈ 0.566: the start latents carried only ~57 % of the noise their timestep claims. The model is handed timestep t_0 but a sample noised to roughly 0.57 · t_0, so it over-denoises, eats real image detail and flattens it into posterized patches. The absolute error scales with t_0, which is exactly the reported "the greater the level of denoise, the more posterization occurs".

The noise handed to RectifiedFlowInpaintExtension was normalized the same way, so inpainting and outpainting were affected as well.

Fix: _prepare_normalized_start_latents() builds the start latents from the normalized operands instead of normalizing the blended tensor. Only init_latents are normalized; the noise stays at N(0, 1), which is precisely the distribution the transformer expects in normalized space. pack_flux2 is a pure rearrange, so blending before or after packing is equivalent.

2. img2img ran an unshifted sigma schedule

get_schedule_flux2() returns an unshifted linear schedule, because the FlowMatchEulerDiscreteScheduler used for txt2img applies the exponential shift from mu itself. img2img and inpainting bypass that scheduler and step the schedule manually, so they ran the model on a different sigma trajectory entirely. With 9 steps:

txt2img:  1.0  0.984  0.964  0.938  0.904  0.858  0.790  0.683  0.485 → 0
img2img:  1.0  0.889  0.778  0.667  0.556  0.444  0.333  0.222  0.111 → 0

txt2img never evaluates the model below sigma 0.485 — the final step jumps from there straight to 0. img2img called it at 0.44, 0.33, 0.22 and 0.11, outside the range a distilled model like FLUX.2 Klein is trained for, and it leaves a grainy residue there.

This was invisible while bug 1 was present: with only ~57 % of the nominal noise there was little left to remove, so the residue never showed. Fixing bug 1 alone therefore trades posterization for graininess — the two fixes belong together.

Fix: apply the same shift in the manual path, before clipping, so denoising_start/denoising_end select a fraction of the schedule the model is actually run on. The new time_shift_flux2() mirrors diffusers' _time_shift_exponential(mu, 1.0, t), which is how the scheduler invokes it; a test pins the two against each other across step counts.

Not affected

• txt2img — no init latents, so nothing was normalized; and the shift is applied only on the manual Euler path, which is exactly when the scheduler is not used, so there is no double shift.
• Reference images — Flux2RefImageExtension normalizes plain VAE latents with no noise mixed in, which is correct. This matches the issue thread, where the reporter notes reference images look fine and only canvas img2img posterizes.
• VAE formats that expose no BN stats — the raw-space preblend is preserved unchanged.

Related Issues / Discussions

Closes #8964

QA Instructions

Automated

pytest tests/app/invocations/test_flux2_denoise_img2img_normalization.py tests/backend/flux2/test_flux2_schedule_shift.py

15 tests. The normalization ones cover: the noise term keeps unit scale at t_0 ∈ {0.2, 0.5, 0.85, 1.0}; the result provably differs from normalizing the raw mixture by exactly t_0 * (noise - normalize(noise)); init latents are normalized; the add_noise=False path; the no-BN-stats fallback. The schedule ones pin time_shift_flux2() against a real FlowMatchEulerDiscreteScheduler at 4/9/20/30 steps, plus endpoints, monotonicity, and the schedule floor.

Measured end-to-end

FLUX.2 Klein 9B fp8, seed 424242, 9 steps, euler, 1024×1024, one txt2img base image re-run through img2img at four strengths on main and on this branch.

Endpoint test. At strength 1.0, denoising_start = 0, so x = 1·noise + 0·init — the init image drops out mathematically and img2img must reproduce txt2img exactly at the same seed:

	L1 vs txt2img	identical pixels
main	33.66	0.0 %
this branch	0.0000	100.0 %

Detail and tonality. HF energy is Laplacian variance over the frame; "sky levels" counts distinct 8-bit grey levels in a fixed 160×280 smooth-sky patch — the direct posterization measure.

run	HF energy	sky levels
txt2img base (reference)	653	50
strength 0.30 — main / branch	821 / 887	60 / 69
strength 0.50 — main / branch	785 / 858	61 / 69
strength 0.75 — main / branch	566 / 788	43 / 66
strength 1.00 — main / branch	550 / 653	8 / 50

On main, detail energy falls below the source image from strength 0.75 up, and the distinct tone count in the sky collapses monotonically (61 → 43 → 8). At strength 1.0 the sky is an 8-level flat field. On this branch the counts stay near or above the reference at every strength.

Canvas path, at the strengths users actually pick

The runs above drive the denoise node directly. The canvas adds two things: it composites the visible raster layers, and with Scale Before Processing it resizes the composite down before the VAE encode and back up after the decode. It also maps the UI strength non-linearly — for FLUX/FLUX.2 with Optimized Denoising (on by default):

denoising_start = 1 - strength ** 0.2

So the reporter's "denoise around 5+" is denoising_start ≈ 0.13, not 0.5 — a far more aggressive setting than a linear reading suggests, and squarely where this bug is worst.

Re-run through a graph mirroring addImageToImage() (1536×1536 raster layer, Scale Before Processing to 1024, resized back), at the UI strengths a user would actually pick:

	HF energy	sky levels	sky σ
source raster layer	161	50	6.78
UI 0.30 (ds 0.214) — main / branch	92 / 137	40 / 62	5.94 / 9.06
UI 0.50 (ds 0.129) — main / branch	90 / 132	34 / 58	5.02 / 8.37
UI 0.75 (ds 0.056) — main / branch	107 / 135	27 / 52	3.55 / 7.35

On main, every canvas result loses roughly 40 % of the source's detail energy and the distinct tone count in the sky falls monotonically to 27 — below the source's 50 at every strength. On this branch the tone count stays at or above the source and detail energy tracks it. Visually: on main the sky loses all cloud structure and the knit of the sweater flattens into grey; on the branch both survive.

Scheduler coverage

All three FLUX.2 schedulers were run on main and on this branch, txt2img and img2img (strength 0.5), same seed:

	result
txt2img, euler — main vs branch	bit-identical
txt2img, heun — main vs branch	bit-identical
txt2img, lcm — main vs branch	differs, but see below
img2img, euler / heun / lcm — within either branch	bit-identical to each other

txt2img is untouched for euler and heun, as intended. LCM differs — but two runs of the same code differ by the same margin (L1 43.5 vs 48.9 for main vs branch), so FlowMatchLCMScheduler is simply not reproducible from a seed here. That is pre-existing and unrelated to this PR; it may deserve an issue of its own.

The img2img rows are bit-identical across schedulers because the scheduler is only constructed when the manual Euler path is not taken — so for img2img and inpainting the scheduler dropdown has no effect at all and Euler is always used. Also pre-existing, and worth noting here because it is what makes the exponential/mu shift the correct one to apply on that path: it is Euler's own shift, and Euler is what actually runs.

Manual

1. Load FLUX.2 Klein, generate a photorealistic image with txt2img — already good, must stay bit-identical.
2. Put it on the canvas as a raster layer, denoise strength ~0.5, Invoke. Flat banded regions in skin, sky and smooth gradients should be gone, fine detail preserved.
3. Repeat at ~0.75. On main the artifact worsens with strength; it should not here.
4. Run an inpaint and an outpaint over a photographic image — both fixes apply to the masked-region merge as well.
5. Sanity check that txt2img and reference-image generations are visually unchanged.

Merge Plan

Two commits, best kept separate: the first is the correctness fix, the second removes the graininess it exposes. Shipping only the first would trade one artifact for another. Backend only, no schema or frontend impact.

Note this changes img2img/inpaint sigma schedules for all FLUX.2 models, [dev] included. That is the intended direction — txt2img already shifts for both — but it is a visible behaviour change. Verified end-to-end on Klein 9B; [dev] has not been run.

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 — n/a, backend only
• Documentation added / updated (if applicable) — n/a
• Updated What's New copy (if doing a release after this PR)

---

🤖 Generated with Claude Code

[Pfannkuchensack]