Fix FlowMatchEulerDiscreteScheduler.index_for_timestep float precision lookup

[ATOM00blue]

What does this PR do?

FlowMatchEulerDiscreteScheduler.index_for_timestep raises IndexError for timesteps that are conceptually integers (e.g. 254) because of float32 rounding in self.timesteps.

Fixes #11749

Bug

import torch
from diffusers import FlowMatchEulerDiscreteScheduler

scheduler = FlowMatchEulerDiscreteScheduler()
print(scheduler.timesteps[746].item())   # 254.00001525878906

# direct lookup
scheduler.index_for_timestep(torch.tensor(254.0))
# IndexError: index 0 is out of bounds for dimension 0 with size 0

# public training API hits the same path (begin_index is None during training)
sample, noise = torch.randn(2, 4, 8, 8), torch.randn(2, 4, 8, 8)
scheduler.scale_noise(sample, torch.tensor([254, 500]), noise)
# IndexError: index 0 is out of bounds for dimension 0 with size 0

Root cause

self.timesteps is computed as sigmas * num_train_timesteps and stored in float32, so values that should be integers carry a tiny rounding error (254 -> 254.00001). index_for_timestep looks them up with an exact equality:

indices = (schedule_timesteps == timestep).nonzero()
...
return indices[pos].item()

When the caller passes the integer value (very common in the training path of scale_noise/add_noise, where begin_index is None), the exact match is empty and indices[pos] raises IndexError. The normal inference loop is unaffected because it passes the exact float elements of scheduler.timesteps.

Fix

Fall back to a torch.isclose match only when the exact lookup returns nothing. The float tolerance is far smaller than the spacing between consecutive timesteps (~1), so:

• conceptually integer timesteps now resolve to their index,
• exact elements still resolve to their own index (no drift to a neighbor),
• a value that is genuinely not in the schedule still raises (real mistakes are not masked).

Testing

Added tests/schedulers/test_scheduler_flow_match_euler.py covering: integer-timestep lookup, the scale_noise training path, exact-value lookups, and rejection of an unknown timestep. The two integer-timestep tests fail on main and pass with this change.

$ python -m pytest tests/schedulers/test_scheduler_flow_match_euler.py -q
4 passed

ruff check/ruff format are clean and utils/check_copies.py reports no issues.

Before submitting

• Did you read the contributor guideline?
• Did you write any new necessary tests?

Who can review?

@yiyixuxu @hlky

Developed with AI coding assistance; reviewed and submitted by the author.

[Copilot]

Pull request overview

This PR fixes a float-precision edge case in FlowMatchEulerDiscreteScheduler.index_for_timestep where timesteps that are conceptually integer values (but stored as float32 with small rounding error) could fail an exact == lookup and trigger an IndexError. This affects training-path calls like scale_noise/add_noise when integer timesteps are passed.

Changes:

• Add a fallback torch.isclose-based lookup in index_for_timestep when exact equality finds no matches.
• Add a dedicated scheduler test file covering integer-timestep lookup, exact-value lookup, unknown-timestep rejection, and the scale_noise training path.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

File	Description
src/diffusers/schedulers/scheduling_flow_match_euler_discrete.py	Adds an isclose fallback when exact timestep lookup fails to avoid IndexError from float32 rounding.
tests/schedulers/test_scheduler_flow_match_euler.py	Adds regression tests for the float-precision timestep indexing bug and related training-path behavior.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.