Fix FLUX scheduler integration

[JPPhoto]

Summary

Fix FLUX denoising so Euler uses Diffusers by default, respects InvokeAI's full and partial sigma schedules, and previews use a matching sample and sigma. Apply InvokeAI's custom schedule to Heun while preserving its first and second order evaluations. Keep LCM's scheduler-defined schedule.

Related Issues / Discussions

QA Instructions

• Focused Flux and invocation suites: 139 passed, 5 expected xfails.
• Ruff lint, formatting, and git diff --check passed.
• Tests use synthetic CPU tensors; no accelerator generation was run.

Review

No material review findings remain.

Compatibility / Rollout

LCM continues to generate its own schedule. For partial denoising, its exact start and terminal sigmas can differ from InvokeAI's clipped schedule.

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)

[chatgpt-codex-connector]

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-10-09T20:57:08.770708Z	163cac1	PR opened

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

[Pfannkuchensack]

Thanks for this. I reviewed it at 300fd9a (base 3289f163af) and found no blockers. The fix holds up both on CPU and on the GPU.

Verification

CPU:

• Ruff check/format and git diff --check are clean.
• tests/backend/flux + test_denoise_noise_inputs.py: 165 passed, 5 xfailed.
• Adjacent FLUX/FLUX.2 invocation tests: 67 passed.
• Running the real denoise() with Euler, Heun and the default scheduler against independent fp64 Euler/Heun references (full, partial start/end, CFG, soft-mask inpainting) gives max error ≤ 3.7e-7.

GPU (RTX 4090, FLUX.1 dev int8, 1024², 20 steps, seed 42, same DB, base vs PR):

Case	flux_denoise time, base → PR	Image, base vs PR
Euler txt2img (mean of 3 warm runs)	21.4 s → 20.6 s	bit-identical
Euler img2img (start 0.5)	6.3 s → 5.3 s	bit-identical
Euler denoising_end=0.6	18.5 s → 17.2 s	base ran on to σ=0; PR correctly stops at σ≈0.4
Heun txt2img	40.0 s → 39.8 s	different (now on the shifted schedule); quality comparable
Heun img2img (start 0.5)	9.3 s → 9.2 s	base is visibly washed out / blurry; PR is sharp
LCM (8 steps)	8.3 s → 8.3 s	differs, but LCM differs run-to-run anyway (#9791)

Euler and Heun were bit-reproducible within each variant. The saved Euler step is the old wasted evaluation at σ=0, so the speedup costs no quality.

Findings (all non-blocking)

1. Heun progress is still a staircase. denoise.py reports order=2, but step/total_steps are already user-facing, and calc_percentage() halves both for order == 2.

   • Observed over the socket at 20 steps: 0.0, 0.1, 0.1, 0.2, 0.2, … 0.9, 1.0. With an odd step count it hits 100 % one step early.
   • Base stopped at 90 %, so this is an improvement, but reporting order=1 would make it linear.
   • The new tests currently pin [2, 2, 1].

2. The LCM exclusion is untested. FlowMatchLCMScheduler.set_timesteps accepts sigmas, so the not is_lcm check carries weight. If it is removed, the whole suite still passes, because test_lcm_preview_uses_the_model_evaluation_sample_and_sigma derives its expectations from scheduler.timesteps. Asserting the literal LCM schedule (e.g. model timesteps [1.0, 0.001] for 2 steps) would catch that.

3. Other test gaps (mutation-checked):

   • A Heun preview taken one evaluation late passes, because previews are only checked with a constant-velocity model. Checking callback.latents in the sample-dependent Heun test would catch it.
   • Indexing CFG by step_index instead of user_step passes, because no test uses a non-uniform cfg_scale with Heun. Expected timestep_index per evaluation for 3 steps: [0, 0, 1, 1, 2].
   • The Heun case of test_flux_img2img_preblend_matches_scheduler_first_sigma calls set_heun_timesteps_from_sigmas directly, so it can't fail. Routing Heun back to num_inference_steps in denoise() leaves it green. The real coverage is test_heun_partial_start_and_end_use_custom_sigmas.

4. set_heun_timesteps_from_sigmas resets private Diffusers state by hand (_step_index, _begin_index, dt, prev_derivative, sample).

   • The reset is complete for 0.41.0, and the independent Heun-recurrence test would catch layout drift.
   • A smaller surface: call scheduler.set_timesteps(num_inference_steps=len(sigmas), device=device) first, then overwrite only timesteps/sigmas.
   • Nit: the comment says "Diffusers 0.40.0", but the pin is 0.41.0 (the Heun/Euler/LCM sources are byte-identical).

5. Two non-obvious comments were lost when the internal Euler loop was removed:

   • why the negative CFG pass must include the Kontext img_cond_seq tokens (otherwise CFG artifacts);
   • the ControlNet-residual-sum peak-memory TODO.

   Both are worth keeping on the surviving loop.

6. Nits:

   • The inspect.signature + hasattr(scheduler, "sigmas") branch looks generic but assumes Euler's "append one terminal" layout. An explicit isinstance(..., FlowMatchEulerDiscreteScheduler) that raises for unknown schedulers would be clearer, and would match the isinstance Heun dispatch rather than the hasattr(scheduler, "state_in_first_order") duck-typing.
   • _FakeHeunScheduler no longer models a reachable path.
   • The two Euler tests overlap heavily.
   • torch.equal on float schedules relies on the fixture values rounding to whole numbers.

Worth mentioning in the PR description

• FLUX Euler with denoising_end < 1 now returns partially denoised latents. Base silently denoised to 0, so saved workflows that relied on that will look noisier. This is correct, but visible. The webv2 UI always sends denoising_end: 1.
• Heun images change for the same seed, because base ignored the shifted, clipped schedule.

Out of scope (filed separately, not caused by this PR)

• #9790: chaining two FLUX Denoise nodes with add_noise=false crashes (noise=None → pack()). This is why I could only test denoising_end single-stage.
• #9791: LCM is not seed-reproducible (unseeded noise in scheduler.step()).
• #9792: FLUX.2 and Z-Image have the same schedule/denoising_end/Heun-progress problems this PR fixes for FLUX.1. set_heun_timesteps_from_sigmas could be reused there.