From 7a6ed103921d19283f2c9e4c74a615b1d7768410 Mon Sep 17 00:00:00 2001 From: Chase Xu <80196056+Chase-Xuu@users.noreply.github.com> Date: Mon, 9 Mar 2026 20:55:42 -0500 Subject: [PATCH] fix: deduplicate timesteps in DPMSolverMultistepScheduler to prevent IndexError When using `beta_schedule='squaredcos_cap_v2'` with `use_karras_sigmas=True`, the sigma-to-timestep mapping produces many nearly identical float values that collapse to the same integer timestep after rounding (e.g., 998.90, 998.80 both become 998). These duplicate timesteps cause the step_index counter to drift past the sigmas array, resulting in an IndexError at `self.sigmas[self.step_index + 1]`. Fix: - Round timesteps unconditionally (remove the squaredcos_cap_v2 exception) - Deduplicate timesteps after rounding, keeping only the first occurrence - Filter corresponding sigmas to match the deduplicated timesteps - Update num_inference_steps to reflect the actual number of unique steps This mirrors the existing deduplication logic in DPMSolverMultistepInverseScheduler. Fixes #12771 --- .../schedulers/scheduling_dpmsolver_multistep.py | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py b/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py index 9c15df4569ca..8f5ae1b04620 100644 --- a/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py +++ b/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py @@ -442,15 +442,11 @@ def set_timesteps( sigmas = np.flip(sigmas).copy() sigmas = self._convert_to_karras(in_sigmas=sigmas, num_inference_steps=num_inference_steps) timesteps = np.array([self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas]) - if self.config.beta_schedule != "squaredcos_cap_v2": - timesteps = timesteps.round() elif self.config.use_lu_lambdas: lambdas = np.flip(log_sigmas.copy()) lambdas = self._convert_to_lu(in_lambdas=lambdas, num_inference_steps=num_inference_steps) sigmas = np.exp(lambdas) timesteps = np.array([self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas]) - if self.config.beta_schedule != "squaredcos_cap_v2": - timesteps = timesteps.round() elif self.config.use_exponential_sigmas: sigmas = np.flip(sigmas).copy() sigmas = self._convert_to_exponential(in_sigmas=sigmas, num_inference_steps=num_inference_steps) @@ -476,6 +472,17 @@ def set_timesteps( f"`final_sigmas_type` must be one of 'zero', or 'sigma_min', but got {self.config.final_sigmas_type}" ) + # When sigma-to-timestep mapping produces duplicates after rounding to integer + # (e.g. cosine schedule with karras/lu sigmas where many sigmas map to nearly the same timestep), + # deduplicate while preserving order to prevent index out-of-bounds errors in multistep solvers. + # This mirrors the same fix in DPMSolverMultistepInverseScheduler. + timesteps = np.round(timesteps).astype(np.int64) + _, unique_indices = np.unique(timesteps, return_index=True) + unique_indices = np.sort(unique_indices) + if len(unique_indices) < len(timesteps): + timesteps = timesteps[unique_indices] + sigmas = sigmas[unique_indices] + sigmas = np.concatenate([sigmas, [sigma_last]]).astype(np.float32) self.sigmas = torch.from_numpy(sigmas)