Skip to content

fix: deduplicate timesteps in DPMSolverMultistep for squaredcos_cap_v2 + karras/lu sigmas#13240

Open
Chase-Xuu wants to merge 1 commit intohuggingface:mainfrom
Chase-Xuu:fix/dpmsolver-squaredcos-karras-timestep-dedup
Open

fix: deduplicate timesteps in DPMSolverMultistep for squaredcos_cap_v2 + karras/lu sigmas#13240
Chase-Xuu wants to merge 1 commit intohuggingface:mainfrom
Chase-Xuu:fix/dpmsolver-squaredcos-karras-timestep-dedup

Conversation

@Chase-Xuu
Copy link

Summary

Fixes #12771DPMSolverMultistepScheduler crashes with IndexError: index 21 is out of bounds for dimension 0 with size 21 when using beta_schedule='squaredcos_cap_v2' with use_karras_sigmas=True.

Root Cause

When using karras (or lu) sigmas with the cosine (squaredcos_cap_v2) beta schedule, _sigma_to_t() returns float timestep values that are extremely close together at the high end of the schedule:

[998.999, 998.903, 998.802, 998.696, 998.583, 998.464, 998.338, 998.203, 998.059, ...]

These were intentionally not rounded for the cosine schedule (unlike linear and scaled_linear), but since they're cast to int64 downstream, they all collapse to the same value (998), creating 9 duplicate timesteps.

The duplicate timesteps cause index_for_timestep() to start at index 1 (it picks the second occurrence of duplicates). After 20 step iterations, step_index reaches 21, and the second-order update's self.sigmas[self.step_index + 1] goes out of bounds on a size-21 array.

Fix

  1. Always round timesteps for karras/lu sigma schedules — this matches the behavior of DPMSolverSinglestepScheduler, which already rounds unconditionally (line 386)
  2. Deduplicate timesteps and their corresponding sigmas after rounding to prevent step index drift

When deduplication occurs, the number of effective inference steps is reduced (e.g., 20 requested → 12 unique). This is correct behavior since the redundant timesteps would have produced identical denoising steps anyway.

Test Results

Configuration Before After
squaredcos_cap_v2 + karras (20 steps) ❌ IndexError at step 19 ✅ 12 unique steps, completes
linear + karras (20 steps) ✅ Unaffected (no duplicates)
squaredcos_cap_v2 without karras ✅ Unaffected
squaredcos_cap_v2 + lu_lambdas ❌ Same index drift ✅ 17 unique steps, completes

Reproduction

from diffusers import DPMSolverMultistepScheduler
import torch

scheduler = DPMSolverMultistepScheduler(
    beta_schedule='squaredcos_cap_v2',
    use_karras_sigmas=True,
)
scheduler.set_timesteps(20)

latents = torch.randn(1, 4, 64, 64)
noise_pred = torch.randn_like(latents)
for i, t in enumerate(scheduler.timesteps):
    result = scheduler.step(noise_pred, t, latents)
    latents = result.prev_sample
# Before: IndexError at step 19
# After: Completes successfully

…2 + karras/lu sigmas

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 huggingface#12771

Signed-off-by: Chase Xu <chase_xu@outlook.com>
Signed-off-by: Chase Xu <80196056+Chase-Xuu@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

DPMSolverMultistepScheduler out-of-bounds error when using beta_schedule="squaredcos_cap_v2" + use_karras_sigmas=True

1 participant