Several FlowMatch pipelines miss the set_begin_index(0) DtoH-sync fix from #11696

[shoemoney]

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

1. 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.
2. Are any of these five deliberately excluded for a reason I can't see from the source?
3. 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.

[dineshyadav03]

Agree with scoping this as one systemic fix rather than five near-identical PRs. Given the root cause is in FlowMatchEulerDiscreteScheduler.index_for_timestep()'s nonzero().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 in set_timesteps() (or __init__): default _begin_index to 0 there instead of requiring callers to know to call set_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 base pipeline_flux.py; the fix may not have been carried into the siblings even though they share the same scheduler and loop structure.