Scheduler crashes when num_training_steps == num_inference_steps

[lucala]

Describe the bug

scheduling_dpmsolver_multistep.set_timesteps() computes a timesteps array, which the scheduler accesses during forward steps, the array is rounded to the nearest int which will give duplicate entries in certain cases (e.g. when num_train_timesteps is equal to num_inference_steps). This causes a bug down the line when the assumption of unique entries no longer holds when obtaining a specific value using step_index = (self.timesteps == timestep).nonzero(). I suspect this problem might also appear in other schedulers.

Reproduction

>>> import numpy as np
>>> num_train_timesteps = 20
>>> num_inference_steps = 20
>>> np.linspace(0, num_train_timesteps - 1, num_inference_steps + 1) # copied from set_timesteps()
array([ 0.  ,  0.95,  1.9 ,  2.85,  3.8 ,  4.75,  5.7 ,  6.65,  7.6 ,
        8.55,  9.5 , 10.45, 11.4 , 12.35, 13.3 , 14.25, 15.2 , 16.15,
       17.1 , 18.05, 19.  ])
>>> np.linspace(0, num_train_timesteps - 1, num_inference_steps + 1).round()[::-1][:-1]
array([19., 18., 17., 16., 15., 14., 13., 12., 11., 10., 10.,  9.,  8.,
        7.,  6.,  5.,  4.,  3.,  2.,  1.])

Logs

No response

System Info

• diffusers version: 0.14.0
• Platform: Linux-4.19.0-23-amd64-x86_64-with-glibc2.28
• Python version: 3.9.16
• PyTorch version (GPU?): 2.0.0+cu117 (True)
• Huggingface_hub version: 0.13.3
• Transformers version: 4.27.3
• Accelerate version: 0.18.0
• xFormers version: 0.0.18
• Using GPU in script?: yes
• Using distributed or parallel set-up in script?: no

[nipunjindal]

There are multiple ways to solve this. Maybe @patrickvonplaten can weigh in on a preferred way?
I'd propose something like the following will make sure we have unique (int) timesteps

        timesteps = (
            np.linspace(0, self.num_train_timesteps - 1, num_inference_steps + 1)
            .round()[::-1][:-1]
            .copy()
            .astype(np.int64)
        )
        _, unique_indices = np.unique(timesteps, return_index=True)
        timesteps = timesteps[np.sort(unique_indices)]

[patrickvonplaten]

Think this should be solved by #2969 no?

[nipunjindal]

@patrickvonplaten Awesome. Yes, that should solve it, happy it's the same fix.
Regarding the other schedulers that may need this fix, it would be a good idea to create a ticket maybe?

[nipunjindal]

Just saw I think Kashif added it today.

[lucala]

Thanks guys, I will close the issue as it is indeed a duplicate of #2969