Bug in DPMSolverSinglestepScheduler: scheduler.add_noise does not work for some settings

[AlexandreBinninger]

Describe the bug

Hello!

When using DPMSolverSinglestepScheduler and the function set_timesteps, the resulting timesteps might have two times the same value. A bug occurs when we use the function add_noise. An error is thrown if the value of t appears multiple times in self.timesteps
Setting scheduler.use_karras_sigmas = True does not fix the issue.

The problem probably comes from the line

timesteps = (
 np.linspace(0, self.config.num_train_timesteps - 1 - clipped_idx, num_inference_steps + 1)
 .round()[::-1][:-1]
 .copy()
 .astype(np.int64)
 )

I get that the bug might not appear with more carefully chosen parameters, but when loading with default parameters (like SDXL 1.0), it seems to exist. So, just wanted to let you know!

Thank you!

Reproduction

import diffusers
import torch

scheduler = diffusers.DPMSolverSinglestepScheduler.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", 
 subfolder="scheduler")
scheduler.set_timesteps(scheduler.config.num_train_timesteps)

# returns 2
print(torch.sum(scheduler.timesteps == 500).item())

# bug:
t = torch.tensor([500])
latents = torch.rand((1, 4, 128, 128))
noise = torch.randn_like(latents)
scheduler.add_noise(latents, noise, t)

Logs

2
Traceback (most recent call last):
 File "XX/test_single_dpm.py", line 10, in <module>
 scheduler.add_noise(latents, noise, t)
 File "/home/XXX/anaconda3/envs/condaenv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_dpmsolver_singlestep.py", line 899, in add_noise
 step_indices = [(schedule_timesteps == t).nonzero().item() for t in timesteps]
 File "/home/XX/anaconda3/envs/condaenv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_dpmsolver_singlestep.py", line 899, in <listcomp>
 step_indices = [(schedule_timesteps == t).nonzero().item() for t in timesteps]
RuntimeError: a Tensor with 2 elements cannot be converted to Scalar

System Info

• diffusers version: 0.24.0
• Platform: Linux-6.2.0-39-generic-x86_64-with-glibc2.35
• Python version: 3.10.13
• PyTorch version (GPU?): 2.1.1+cu121 (True)
• Huggingface_hub version: 0.19.4
• Transformers version: 4.36.0
• Accelerate version: 0.25.0
• xFormers version: not installed
• Using GPU in script?: RTX4090
• Using distributed or parallel set-up in script?: No

Who can help?

@yiyixuxu @PATR

[sayakpaul]

Cc: @yiyixuxu

[patrickvonplaten]

@yiyixuxu can you take a look here?

[yiyixuxu]

@patrickvonplaten
i think it is related to
#6085
#5746

[patrickvonplaten]

Let's try to solve this soon, this bug seems to be affecting lots of use cases

[yiyixuxu]

actually this should already be solved with this PR #6085

[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.