Repository navigation
Several FlowMatch pipelines miss the set_begin_index(0) DtoH-sync fix from #11696 #14573
Description
Summary
#11696 ("Avoid DtoH sync from access of nonzero() item in scheduler") added
self.scheduler.set_begin_index(0) before the denoising loop in pipeline_flux.py. The same fix was
later propagated to pipeline_wan.py and pipeline_qwenimage.py. Several other pipelines that use
FlowMatchEulerDiscreteScheduler never received it.
Opening this as one systematic issue rather than per-pipeline PRs, per the contributing guide.
Current state
I checked every FlowMatchEulerDiscreteScheduler pipeline on main for set_begin_index:
| Pipeline | set_begin_index |
Status |
|---|---|---|
flux/pipeline_flux.py |
✅ | fixed by #11696 |
wan/pipeline_wan.py |
✅ | propagated |
qwenimage/pipeline_qwenimage.py |
✅ | propagated |
stable_diffusion_3/pipeline_stable_diffusion_3.py |
❌ | already covered by open PR #14176 |
hunyuan_video/pipeline_hunyuan_video.py |
❌ | unclaimed |
mochi/pipeline_mochi.py |
❌ | unclaimed |
lumina2/pipeline_lumina2.py |
❌ | unclaimed |
aura_flow/pipeline_aura_flow.py |
❌ | unclaimed |
chroma/pipeline_chroma.py |
❌ | unclaimed |
So five pipelines are unclaimed, and SD3 is already in flight — I'd leave that one to #14176 and
scope this to the other five.
Why it matters
Without a pre-set begin_index, the first step() call routes through
FlowMatchEulerDiscreteScheduler._init_step_index() → index_for_timestep(), which does
(schedule_timesteps == timestep).nonzero() followed by .item(). That .item() forces a
device-to-host sync on the first denoising step.
It's a latency cost rather than a correctness bug, and it's most noticeable exactly where these
pipelines are used — long video generations and torch.compile'd runs, where an avoidable sync
stalls the pipeline. pipeline_wan.py even carries a comment pointing back at #11696, so the intent
to propagate seems already established.
Proposed scope
Add the same single line before the denoising loop in each of the five unclaimed pipelines, matching
#11696's placement and style exactly. No behavioral change beyond removing the sync.
Questions before I open anything
- Do you want this as one PR across the five, or one PR per pipeline? The guide says fix
patterns rather than one-offs, so I'd default to a single PR — but five files touched by a new
contributor may not be what you want to review, so I'd rather ask. - Are any of these five deliberately excluded for a reason I can't see from the source?
- Should SD3 stay with [DtoH sync] set scheduler begin index in SD3 T2I #14176? I've assumed yes.
Happy to take it from here once someone acknowledges the scope.
Disclosure per the contributing guide: I used AI assistance to survey the codebase for instances of
this pattern. The table above I verified by reading each file on main directly. If this moves to a
PR I'll run the self-review skill and post the notes with it, along with the test commands and
their output.
Activity
dineshyadav03 commentedon Aug 24, 2026
Agree with scoping this as one systemic fix rather than five near-identical PRs. Given the root cause is in
FlowMatchEulerDiscreteScheduler.index_for_timestep()'snonzero().item()call, patching call sites is treating a scheduler bug as a pipeline bug — it just reappears in the next FlowMatch pipeline someone adds, exactly like it already did here across flux/wan/qwenimage vs. the five unclaimed ones. Better fix is inset_timesteps()(or__init__): default_begin_indexto 0 there instead of requiring callers to know to callset_begin_index(0)afterward. That closes the gap by construction instead of by per-pipeline audit, and it's a much smaller diff than five pipeline edits.Worth double-checking the Flux variant pipelines too —
pipeline_flux_img2img.py,pipeline_flux_inpaint.py,pipeline_flux_controlnet.py— since #11696 only touched the basepipeline_flux.py; the fix may not have been carried into the siblings even though they share the same scheduler and loop structure.Set scheduler begin index in remaining FlowMatch pipelines to avoid D…