From 9efe640760a79708d3677fda5572211a7996e5b0 Mon Sep 17 00:00:00 2001 From: Chase Xu <80196056+Chase-Xuu@users.noreply.github.com> Date: Mon, 9 Mar 2026 21:02:11 -0500 Subject: [PATCH] fix: deduplicate timesteps in DPMSolverMultistep for squaredcos_cap_v2 + karras/lu sigmas MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When using `beta_schedule='squaredcos_cap_v2'` with `use_karras_sigmas=True` or `use_lu_lambdas=True`, the sigma-to-timestep mapping produces float values that are extremely close together near the end of the cosine schedule (e.g., 998.05 to 999.0). Since timesteps are cast to int64, these all collapse to the same integer value (e.g., 998), creating duplicate timesteps. The duplicate timesteps cause `index_for_timestep` to start at the wrong index (it picks the second occurrence), and after N steps the step_index overshoots the sigmas array, producing an IndexError: `index 21 is out of bounds for dimension 0 with size 21` This fix: 1. Always rounds timesteps for karras/lu sigma schedules (matching the behavior of DPMSolverSinglestepScheduler which already rounds unconditionally) 2. Deduplicates timesteps and their corresponding sigmas after rounding to prevent step index drift The number of effective inference steps may be reduced when deduplication occurs (e.g., 20 requested → 12 unique), which is correct behavior since the redundant timesteps would have produced identical denoising steps. Fixes #12771 Signed-off-by: Chase Xu Signed-off-by: Chase Xu <80196056+Chase-Xuu@users.noreply.github.com> --- .../schedulers/scheduling_dpmsolver_multistep.py | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py b/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py index 9c15df4569ca..051275ad946c 100644 --- a/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py +++ b/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py @@ -442,15 +442,13 @@ 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() + 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() + 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) @@ -467,6 +465,16 @@ def set_timesteps( else: sigmas = np.interp(timesteps, np.arange(0, len(sigmas)), sigmas) + # When using karras or lu sigmas with certain beta schedules (e.g. squaredcos_cap_v2), + # the sigma-to-timestep mapping can produce duplicate integer timesteps. Deduplicate + # them to prevent the step index from drifting out of bounds during multistep updates. + timesteps_int = np.round(timesteps).astype(np.int64) + _, unique_indices = np.unique(timesteps_int, return_index=True) + if len(unique_indices) < len(timesteps_int): + unique_indices = np.sort(unique_indices) + timesteps = timesteps[unique_indices] + sigmas = sigmas[unique_indices] + if self.config.final_sigmas_type == "sigma_min": sigma_last = ((1 - self.alphas_cumprod[0]) / self.alphas_cumprod[0]) ** 0.5 elif self.config.final_sigmas_type == "zero":