huggingface / huggingface/diffusers
Suggestion for speeding up `index_for_timestep` by removing sequential `nonzero()` calls in samplers
- Lingua principale
- Python
- Stelle
- 34.5k
- Fork
- 7.3k
- Merge medio
- 3g 3h
- PR unite (30g)
- 91
Descrizione
**Is your feature request related to a problem? Please describe.**
First off, thanks for the great codebase and providing so many resources! I just wanted to provide some insight into an improvement I made for myself, in case you'd like to include it for all samplers. I'm using the `FlowMatchEulerDiscreteScheduler` and after profiling, I've noticed that it's unexpectedly slowing down my training speeds. I'll describe the issue and proposed solution here rather than making a PR, since this would touch a lot of code and perhaps someone on the diffusers team would like to implement it.
**Describe the solution you'd like.**
This line in particular is very slow because it is a for loop `step_indices = [self.index_for_timestep(t, schedule_timesteps) for t in timestep]` and the `self.index_for_timestep()` is calling a nonzero() function which is slow.
https://github.com/huggingface/diffusers/blob/b9e2f886cd6e9182f1bf1bf7421c6363956f94c5/src/diffusers/schedulers/scheduling_flow_match_euler_discrete.py#L149
**Describe alternatives you've considered.**
I've changed the code as follows:
```python
# huggingface code
def index_for_timestep(self, timestep, schedule_timesteps=None):
if schedule_timesteps is None:
schedule_timesteps = self.timesteps
indices = (schedule_timesteps == timestep).nonzero()
# The sigma index that is taken for the **very** first `step`
# is always the second index (or the last index if there is only 1)
# This way we can ensure we don't accidentally skip a sigma in
# case we start in the middle of the denoising schedule (e.g. for image-to-image)
pos = 1 if len(indices) > 1 else 0
return indices[pos].item()
```
changed to =>
```python
# my code
def index_for_timestep(self, timestep, schedule_timesteps=None):
if schedule_timesteps is None:
schedule_timesteps = self.timesteps
num_steps = len(schedule_timesteps)
start = schedule_timesteps[0].item()
end = schedule_timesteps[-1].item()
indices = torch.round(((timestep - start) / (end - start)) * (num_steps - 1)).long()
return indices
```
and
```python
# huggingface code
# self.begin_index is None when scheduler is used for training, or pipeline does not implement set_begin_index
if self.begin_index is None:
step_indices = [self.index_for_timestep(t, schedule_timesteps) for t in timestep]
```
changed to =>
```python
# my code
# self.begin_index is None when scheduler is used for training, or pipeline does not implement set_begin_index
if self.begin_index is None:
step_indices = self.index_for_timestep(timestep, schedule_timesteps)
```
**Additional context.**
Just wanted to bring this modification to your attention since it could be a training speedup for folks. 🙂 Especially when someone has a large batch size > 1 and this for loop it occurring with nonzero search operations. Some other small changes might be necessary to ensure compatibility of the function changes, but I suspect it could help everyone. Thanks for the consideration!
Guida per i contributori
Apri la guida per i contributori
Direzione di ricerca
Start with src/diffusers/schedulers/scheduling_flow_match_euler_discrete.py at index_for_timestep and its step_indices call, then inspect equivalent usages across samplers. Compare the proposed batched lookup with existing timestep and sigma semantics, especially begin-index behavior, and benchmark large batches. Done means compatible scheduler behavior with a measured improvement and appropriate validation across affected samplers.
Scritto dal modello di indicizzazione a partire dal testo della issue.
Valutazione
- Stack tecnologico
- python, pytorch
- Ambito
- machine-learning, performance
- Tipo di issue
- Funzionalità
- Difficoltà
- 4/5
- Tempo stimato
- 3-5 giorni
- Stato di attività
- Ferma
- Chiarezza
- Abbastanza chiara
- Idoneità per principianti
- 35/100