Bug of some timestep indices cannot be found in index_for_timestep() of scheduling_flow_match_euler_discrete.py

[Liang-ZX]

Describe the bug

Description

In the index_for_timestep() method, there's a precision issue when comparing floating-point timesteps with integer schedule_timesteps. The comparison schedule_timesteps == timestep can fail to find matching indices due to floating-point truncation, causing some timestep indices to be missed.

Location: scheduling_flow_match_euler_discrete.py, line 303:

indices = (schedule_timesteps == timestep).nonzero()

Reproduction

import torch
import numpy as np
from diffusers.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler

num_train_timesteps = 1000
batch_size = 1

noise_scheduler = FlowMatchEulerDiscreteScheduler(
    num_train_timesteps=num_train_timesteps,
)

device = "cuda:0"
for i in range(100000):
    print(i)
    action_gt = torch.rand(64, 128).to(device)
    noise = torch.randn(
            action_gt.shape, dtype=action_gt.dtype, device=device
        )
    timesteps = torch.randint(
            1, num_train_timesteps,
            (batch_size,), device=device
        ).long()
    
    noisy_action = noise_scheduler.scale_noise(
            sample=action_gt,
            noise=noise,
            timestep=timesteps,
        )

Logs

  File "/anaconda3/envs/rdt/lib/python3.10/site-packages/diffusers/schedulers/scheduling_flow_match_euler_discrete.py", line 175, in scale_noise
    step_indices = [self.index_for_timestep(t, schedule_timesteps) for t in timestep]
  File "/anaconda3/envs/rdt/lib/python3.10/site-packages/diffusers/schedulers/scheduling_flow_match_euler_discrete.py", line 175, in <listcomp>
    step_indices = [self.index_for_timestep(t, schedule_timesteps) for t in timestep]
  File "/anaconda3/envs/rdt/lib/python3.10/site-packages/diffusers/schedulers/scheduling_flow_match_euler_discrete.py", line 292, in index_for_timestep
    return indices[pos].item()
IndexError: index 0 is out of bounds for dimension 0 with size 0

System Info

diffusers version: 0.32.1
Python version: 3.10
OS: Linux / Ubuntu 22

Who can help?

@yiyixuxu

[Liang-ZX]

#10587

[hlky]

@Liang-ZX Passing IntTensor/LongTensor to schedulers is deprecated. We can't round timesteps in index_for_timestep as this would affect other cases.

If you want a random timestep try something like this

import torch
from diffusers.schedulers.scheduling_flow_match_euler_discrete import (
    FlowMatchEulerDiscreteScheduler,
)

num_train_timesteps = 1000
batch_size = 1

noise_scheduler = FlowMatchEulerDiscreteScheduler(
    num_train_timesteps=num_train_timesteps,
)

for i in range(10000):
    print(i)
    action_gt = torch.rand(64, 128)
    noise = torch.randn(action_gt.shape, dtype=action_gt.dtype)
    indices = torch.randint(1, num_train_timesteps, (batch_size,)).long()
    timesteps = noise_scheduler.timesteps[indices]
    noisy_action = noise_scheduler.scale_noise(
        sample=action_gt,
        noise=noise,
        timestep=timesteps,
    )