Bug in DPMSolverSinglestepScheduler: scheduler.add_noise does not work for some settings #6232
Description
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 ScalarSystem Info
diffusersversion: 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