Skip to content

Set scheduler begin index in remaining FlowMatch pipelines to avoid DtoH sync - #14576

Open
shoemoney wants to merge 1 commit into
huggingface:mainfrom
shoemoney:fix/begin-index-siblings
Open

Set scheduler begin index in remaining FlowMatch pipelines to avoid DtoH sync#14576
shoemoney wants to merge 1 commit into
huggingface:mainfrom
shoemoney:fix/begin-index-siblings

Conversation

@shoemoney

Copy link
Copy Markdown

Fixes#14573.

#11696 added self.scheduler.set_begin_index(0) before the denoising loop in FluxPipeline to avoid the device-to-host sync from nonzero().item() in FlowMatchEulerDiscreteScheduler.index_for_timestep(). That fix was later propagated to Wan, QwenImage, and others, but five FlowMatchEulerDiscreteScheduler pipelines never received it: HunyuanVideo, Mochi, Lumina2, AuraFlow, and Chroma (text-to-image).

This PR applies the identical one-line pattern (with the same comment pointing at #11696) to those five. SD3 is intentionally excluded since open PR #14176 already covers it.

Verified on CPU with the repo's dummy-component fixtures: before, each pipeline enters index_for_timestep (the .item() sync path) once on the first step; after, zero entries, begin_index == 0, outputs unchanged. Fast test suites for all five pipelines show no new failures versus main; make style / check_copies clean. Like #11696, no new tests, since the eliminated sync is only observable in a profiler trace.

…toH sync
Propagate the set_begin_index(0) fix from huggingface#11696 to HunyuanVideo, Mochi,
Lumina2, AuraFlow, and Chroma text-to-image/video pipelines. Without it,
the first scheduler.step() routes through _init_step_index() ->
index_for_timestep(), whose nonzero().item() forces a device-to-host
sync on the first denoising step. Fixeshuggingface#14573.
Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Labels

fixes-issuepipelinessize/SPR with diff < 50 LOC

Projects

None yet

Development

Successfully merging this pull request may close these issues.

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

1 participant

@shoemoney