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,
    )

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

[ChanHyeok-Choi]

I got the same issue