[horizonbymuneeb]

Summary

ERNIE-Image offers Heun alongside Euler in its scheduler set, but every partial denoise on Heun failed before the first step:

ValueError: FlowMatchHeunDiscreteScheduler does not accept an explicit sigma schedule, so it cannot honor denoising_start/denoising_end. Use the euler or lcm scheduler for partial denoising.

denoise() handed Euler the requested window through set_timesteps(sigmas=...) and, for a scheduler whose set_timesteps takes only a step count, refused anything that did not run 1.0 → 0.0. So img2img, inpaint and Canvas passes with strength < 1 failed for a scheduler the model loader advertises.

After: the Heun branch rebuilds the scheduler's expanded two-evaluation layout from the requested window, applying the same flow shift set_timesteps applies, which is the layout the FLUX driver already installs through set_heun_timesteps_from_sigmas. A pass now opens at the denoising_start sigma and stops at the denoising_end sigma.

Full-window Heun output changes

A full-window Heun pass is not bit-identical to before. Previously set_timesteps(N) built diffusers' own Heun schedule; now txt2img also goes through the get_schedule window, so the sigmas differ. At shift=4, 8 steps, the tail goes from … 0.41, 0.016 → 0 to … 0.57, 0.36 → 0: the last interval becomes a first-order jump from 0.36. Measured mean absolute pixel difference against main is 12.5 at 8 steps and 2.8 at 20 steps.

That aligns Heun with the Euler/FLUX schedule the rest of the model already uses, and it is what makes the partial window well-defined, but existing Heun seeds no longer reproduce. Verified on an RTX 4090 with ERNIE-Image Turbo (1024², seed 7, shift=4): Euler is bit-identical, Heun changes as described, and the partial-window img2img Heun case that raised #10003 now returns a clean result.

Related Issues / Discussions

Closes #10003

QA Instructions

Reproduced on a stub transformer ([1, 128, 2, 2] latents, steps=8), which is the same shape the existing tests/backend/ernie_image/test_ernie_denoise.py harness uses. Before, euler partial → ok, heun partial → the ValueError above, heun full → ok. After the change all three run.

tests/backend/ernie_image/test_ernie_denoise.py:

• test_heun_rejects_a_partial_denoise_range asserted the old refusal; it is replaced by test_heun_honours_a_partial_denoise_range (first sigma 0.5, terminal 0.0, layout length 2N-1 timesteps, progress does not overrun).
• test_heun_keeps_the_schedulers_own_shift_on_a_partial_range pins both window ends under shift=3: start 0.5 → 0.75 and terminal 0.25 → 0.5. A denoising_end=1.0 case alone cannot catch a missing shift, because the shift maps 0 to 0.
• The Euler cases are unchanged and still pin their behaviour. The full-window Heun case still runs, but its sigmas change as described above; the tests pin the first and last sigma only, so they do not cover that difference.

The test file was run against the real denoise() path with the two heavy imports the module pulls in stubbed (the pipeline-state dataclass and the inpaint extension, neither of which the scheduler path touches), since uv sync does not build in this sparse checkout: 15 passed. Mutation check: making the new shift helper an identity turns both test_heun_keeps_the_schedulers_own_shift_on_a_partial_range cases red (0.5 != 0.75), so the shift coverage is not vacuous. ruff check and ruff format --check are clean on both changed files.

Review

No material findings: the change reuses the scheduler layout helper the FLUX driver already uses and keeps the ERNIE-specific shift in one private helper next to its only caller. The remaining risk is a schedule built from a window that is not a plain linspace subset: get_schedule is the only producer and it slices a linspace, so the shifted window matches what set_timesteps would have produced for those end sigmas.

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)

[horizonbymuneeb]

Gentle check-in on this one — it has been sitting for about a day and a half without review. The Heun partial-denoise change reuses the schedule helper the FLUX driver already uses, and the updated test_ernie_denoise.py passes (15 tests, with a mutation check confirming the shift coverage is not vacuous).

The branch is based on main; if it is now behind, I'm happy to merge main in and push. Let me know if anything here needs adjusting.

[horizonbymuneeb]

Checking in on this one — it has been a couple of days since the last note and main has moved a fair way ahead (about 90 commits), so the branch is showing as behind.

The fix itself is unaffected by that drift: the Heun partial-denoise path in ernie_image/denoise.py and the two new tests still describe current main, and I am happy to rebase and re-run the package tests whenever you want to look at it. If the change is not something you want to carry, say so and I will close it rather than leave it in the queue.

[Pfannkuchensack]

Tested on an RTX 4090 with ERNIE-Image Turbo (1024², seed 7, shift 4), calling denoise() from main and from this branch:

Case	main	this PR
Euler, 8 steps	ok	ok, bit-identical
Heun, 8 / 20 steps	ok	ok, different image
img2img Euler, strength 0.5	ok	ok
img2img Heun, strength 0.5	ValueError (#10003)	ok, clean result

pytest tests/backend/ernie_image: 15 passed. The partial-window fix works.

Two points:

1. Full-window Heun output changes. The description says the full-window case is unchanged, but it isn't. Before, set_timesteps(N) built diffusers' own Heun schedule. Now txt2img also goes through the get_schedule window, so the sigmas differ (shift 4, 8 steps):

   • main: … 0.41, 0.016 → 0
   • PR: … 0.57, 0.36 → 0

   The final interval is now a first-order jump from 0.36 to 0. Mean absolute pixel difference vs main: 12.5 at 8 steps, 2.8 at 20 steps. Visually the PR image is not worse; at 8 steps it is arguably cleaner. Aligning Heun with the Euler/FLUX schedule seems fine to me, but please state it in the description, since existing Heun seeds no longer reproduce. The tests only pin the first and last sigma, so they don't cover this.

2. ruff format --check fails on tests/backend/ernie_image/test_ernie_denoise.py: the signature of test_heun_keeps_the_schedulers_own_shift_on_a_partial_range fits on one line.

[horizonbymuneeb]

Thanks — that is a genuinely useful run, and you are right on both points.

1. ruff format. Fixed in 24905826: the signature fits on one line. ruff check and ruff format --check are clean on both changed files now. Sorry for the noise — I did not have ruff available when the branch was written and checked the 120-column limit by hand.

2. Full-window Heun output change. You are right and the description was wrong. I have corrected it rather than argued with it: a new "Full-window Heun output changes" section states the sigma tail change (… 0.41, 0.016 → 0 to … 0.57, 0.36 → 0 at shift=4, 8 steps), gives your measured pixel differences (12.5 at 8 steps, 2.8 at 20), and says plainly that existing Heun seeds no longer reproduce. I also corrected the QA section, which claimed the full-window case was unchanged — the tests pin the first and last sigma only, so they never covered that.

On the substance: aligning Heun with the Euler/FLUX schedule is what makes a partial window well-defined, since the window is a slice of the same linspace the Euler path already uses. Building it from diffusers' own set_timesteps would mean a second schedule that disagrees with the window boundaries. I think the trade is right, but it is a behaviour change for anyone with saved Heun seeds, so it belongs in the description rather than in a footnote.

Your table is also the first real-hardware evidence for the case that motivated this: img2img Heun at strength 0.5 went from ValueError to a clean result.