Skip to content

Set begin_index for flow-matching pipelines to avoid slow index_for_timestep#13206

Open
kirillsst wants to merge 1 commit intohuggingface:mainfrom
kirillsst:fix/set-begin-index-flow-matching-pipelines
Open

Set begin_index for flow-matching pipelines to avoid slow index_for_timestep#13206
kirillsst wants to merge 1 commit intohuggingface:mainfrom
kirillsst:fix/set-begin-index-flow-matching-pipelines

Conversation

@kirillsst
Copy link
Contributor

Summary

  • Add scheduler.set_begin_index(0) before the denoising loop in all flow-matching pipelines that were missing it
  • For img2img pipelines with strength-based schedule truncation (wan_video2video, ltx_condition), set the begin index to the appropriate start offset inside get_timesteps()
  • Without begin_index, _init_step_index() falls back to index_for_timestep() which uses .nonzero() causing an expensive DtoH sync on every first scheduler.step() call
  • This was already done in Flux and a few other pipelines but was missing across most flow-matching pipelines (56 files)

Details

The FlowMatchEulerDiscreteScheduler._init_step_index method has two paths:

if self.begin_index is None:
    self._step_index = self.index_for_timestep(timestep)  # slow: .nonzero()
else:
    self._step_index = self._begin_index  # fast: direct assignment

This PR ensures all flow-matching pipelines take the fast path, matching the pattern already established in pipeline_flux.py.

Fixes #9417

Test plan

  • ruff check passes on all 56 modified files
  • ruff format --check passes on all 56 modified files
  • Existing pipeline tests should pass (no behavioral change, only performance)

…r_timestep`

The `_init_step_index` method falls back to `index_for_timestep()` which
uses `.nonzero()` causing an expensive DtoH sync when `begin_index` is
not set. This was already fixed in Flux and a few other pipelines but
was missing across most flow-matching pipelines.

Add `scheduler.set_begin_index(0)` before the denoising loop in all
text-to-image/video pipelines that use FlowMatch schedulers. For
img2img pipelines with strength-based schedule truncation
(wan_video2video, ltx_condition), set the begin index to the
appropriate start offset inside `get_timesteps()`.

Fixes huggingface#9417
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Suggestion for speeding up index_for_timestep by removing sequential nonzero() calls in samplers

1 participant