fix(flux2,z-image): run Diffusers schedulers on the node's sigma schedule

[Pfannkuchensack]

Summary

#9724 fixed FLUX.1's Diffusers scheduler paths. FLUX.2 and Z-Image still had the same problems:

• FLUX.2 Euler/LCM dropped the terminal sigma, so denoising_end < 1 finished the image instead of stopping. txt2img also clipped the unshifted schedule, while img2img clipped the shifted one, so a pass ending at denoising_end = x and a pass continuing from denoising_start = x did not meet at the same sigma.
• FLUX.2 Heun ran its own static shift of 3 and ignored mu and denoising_end.
• Z-Image Heun ran a linear schedule (shift 1.0) from sigma 1.0, ignoring the shifted, clipped schedule; img2img started at sigma 1.0 on latents blended at a lower sigma.
• Heun progress on both stopped at (n−1)/n and moved in a staircase (order=2 is halved by calc_percentage).

Now:

• The FLUX.2 node shifts the linear schedule with time_shift_flux2(mu) once, before clipping, on every path. Its Diffusers schedulers are built without a shift of their own. Euler and LCM run those sigmas plus the requested terminal sigma via a new set_timesteps_from_sigmas() helper (shared with FLUX.1's Euler); Heun runs them via set_heun_timesteps_from_sigmas() from Fix FLUX scheduler integration #9724.
• Z-Image Heun runs the node's shifted, index-clipped sigmas through the same Heun helper, so img2img starts where the blend left the latents. LCM keeps building its own schedule, as on FLUX.1.
• Heun progress reports 1/n … n/n with order=1 on both; FLUX.2 previews are taken from the sample and sigma of the model evaluation, as on FLUX.1.
• Because Z-Image Heun now starts at the img2img blend, the Upscale widget offers Euler and Heun for Z-Image and FLUX.1 (FLUX.1 Heun already did since Fix FLUX scheduler integration #9724; the two lists are coupled). LCM stays excluded. The Upscale docs are updated.

Same seed on an RTX 4090, mean absolute pixel difference (0–255) against the Euler result of the same graph, before → after:

Case	Before	After
FLUX.2 Klein 4B, Euler txt2img (before vs after)	–	0.00 (unchanged)
FLUX.2 Klein 4B, Heun txt2img vs Euler	32.5 (grainy sky)	20.7
FLUX.2 two-stage chain (end 0.6 → img2img start 0.6, add_noise=false) vs single Euler run	21.1 (oversaturated)	7.7
Z-Image Turbo, Euler txt2img (before vs after)	–	0.00 (unchanged)
Z-Image Turbo, Heun txt2img vs Euler	28.9 (grainy, different composition)	10.8
Z-Image Turbo, Heun img2img strength 0.5 vs Euler img2img	46.4 (blurry dark blob)	2.6

Related Issues / Discussions

Closes #9792

Follows #9724, which added the Heun helper for FLUX.1.

QA Instructions

Automated:

• uv run --no-sync pytest tests/app/invocations/test_denoise_noise_inputs.py tests/app/invocations/test_flux2_denoise_wiring.py tests/app/invocations/test_flux2_denoise_img2img_normalization.py tests/app/invocations/test_flux2_working_memory.py tests/app/invocations/test_z_image_schedule_shift.py tests/app/invocations/test_z_image_denoise_working_memory.py tests/app/invocations/test_z_image_denoise_schedulers.py tests/app/invocations/test_upscale_index_clipped_starts.py tests/app/invocations/test_module_compatibility.py tests/backend/flux tests/backend/flux2 tests/backend/architectures/test_scheduler_set.py → 525 passed, 4 xfailed (the known LCM preblend cases).
• New/changed tests fail against main: terminal sigma, Heun schedule and progress, previews from the evaluation sample, shift-before-clip on txt2img, Z-Image Heun schedule/img2img start/progress, and a node built with a scheduler that shifts again (shift=3.0).
• uv tool run ruff@0.11.2 check / format --check on the changed files: clean.
• webv2: pnpm lint passed; vitest run src/features/upscale 97 passed; vitest run --config vitest.browser.config.mts src/app/UpscaleCapabilityGate.browser.test.tsx passed (fails when Z-Image's list is reverted to Euler only).

Manual (RTX 4090, FLUX.2 Klein 4B fp8 and Z-Image Turbo int8, 1024², 8 steps, seed 42): the graphs in the table above, run against main and this branch on the same data root. In the Upscale widget (Vite dev server against this branch), Z-Image's Scheduler field lists "Euler" and "Heun (2nd order)"; a 2× Z-Image upscale with Heun ran end to end and matched the Euler upscale visually.

Not run: the full Python suite and pnpm check:release; no GPU check on AMD/ROCm or macOS.

Review

Resolved findings:

• The Upscale capability browser test still expected Z-Image to show no scheduler field.
• Coverage for a FLUX.2 scheduler that shifts again on Euler/LCM was missing; it is now checked at node level.
• The Upscale docs still said FLUX.1 and Z-Image run Euler only.
• A Z-Image preblend test case only exercised the helper, not the node; it was removed in favour of node-level tests.
• Z-Image Heun at an explicit shift=0 (allowed by the field's ge=0) evaluated the model at sigma 0 and returned NaN latents. Heun now skips the schedule's zero-length intervals, so it takes the single 1 → 0 step Euler takes and its progress counts the steps it runs. The field keeps its range, so saved workflows still load.

Remaining limitations:

• FLUX.2 img2img and inpainting still run manual Euler whatever scheduler is selected (unchanged; a candidate follow-up now that every scheduler starts exactly at timesteps[0]).
• FLUX.2's manual Euler path reports timestep=int(t_curr), i.e. 0, in its progress state (pre-existing).

Compatibility / Rollout

• No persisted-state, schema or OpenAPI change.
• Same-seed results change for FLUX.2 Heun and Z-Image Heun, and for FLUX.2 txt2img with denoising_end < 1, which now returns latents at sigma 1 − denoising_end on the shifted schedule instead of a finished image. Generate, Canvas and Upscale always send denoising_end = 1; only workflow-editor graphs use it.
• Full-range FLUX.2 Euler/LCM txt2img is unchanged: the pre-shifted schedule matches the old dynamic-shifting one within one float32 ulp (6e-8), and the images were identical.
• Internal API: invokeai.backend.flux2.denoise.denoise() no longer takes mu, and FLUX2_TXT2IMG_SCHEDULER_KWARGS is gone from invokeai.backend.flux2.sampling_utils. Neither is part of invokeai.invocation_api and nothing in the repo uses them; out-of-tree code that forked _run_diffusion fails with a TypeError/ImportError instead of shifting twice.
• fix(flux): make LCM reproducible with a fixed seed on FLUX.1, FLUX.2 and Z-Image #9837 (seeded LCM step noise) touches the same scheduler.step(...) line in invokeai/backend/flux2/denoise.py and the removed test_flux2_lcm_scheduler_setup_passes_mu; whichever merges second keeps generator=generator on the step call and drops that test. There is no semantic conflict.

Checklist

• The PR has a short but descriptive title, suitable for a changelog
• Meaningful regression coverage added / updated where needed; obsolete tests/code removed
• Persisted-state and API changes include required migrations / compatibility validation
• Relevant performance/efficiency opportunities considered; material claims have evidence
• Material review findings resolved and relevant checks rerun
• Documentation added / updated (if applicable)
• Updated What's New copy (if doing a release after this PR)

---

^{Need help on this PR? Tag @codesmith-bot with what you need.}

[Pfannkuchensack]

[joshistoast]

P2 Handle zero-shift schedules before passing them to Heun

shift=0 is accepted by this invocation, and _get_sigmas(0, 8) returns [1, 0, 0, …, 0]. Passing that schedule to Heun now evaluates the model at zero sigma: Diffusers 0.41.0 divides by sigma_next during its second-order correction, producing NaN latents on the second evaluation even with a finite constant prediction. I reproduced this using the pinned schedule builder and helper; the base revision's scheduler setup returns finite latents for the same configuration. Handle the degenerate schedule before configuring Heun so an accepted shift override cannot corrupt the generated result.