[deep1blue]

What does this PR do?

Propagates two fixes from DPMSolverMultistepScheduler (scheduling_dpmsolver_multistep.py) to its sibling multistep schedulers (EDMDPMSolverMultistepScheduler, CosineDPMSolverMultistepScheduler, and DPMSolverMultistepInverseScheduler):

1. Prevent NaN division-by-zero on zero-width final step (self.sigmas[-1] == self.sigmas[-2]):

   • In DPMSolverMultistepScheduler.step() (scheduling_dpmsolver_multistep.py), lower_order_final is triggered when:

     (self.config.lower_order_final and self.step_index == len(self.timesteps) - 1)
     or (len(self.sigmas) > 1 and self.sigmas[-1] == self.sigmas[-2])

   • However, EDMDPMSolverMultistepScheduler, CosineDPMSolverMultistepScheduler, and DPMSolverMultistepInverseScheduler were missing the or (len(self.sigmas) > 1 and self.sigmas[-1] == self.sigmas[-2]) check.
   • When final_sigmas_type="sigma_min" (or algorithm_type in ["sde-dpmsolver++", "sde-dpmsolver"] in EDMDPMSolverMultistepScheduler), the last two entries of self.sigmas are identical (sigma_t == sigma_s0), so $h = \lambda_t - \lambda_{s_0} = 0$. When lower_order_final=False with solver_order >= 2, taking dpm_solver_second_order_update or dpm_solver_third_order_update divides by h = 0 (1.0 / h, (exp(-h) - 1.0) / h, r0 = h_0 / h) and produces NaNs across the entire output tensor.
   • In CosineDPMSolverMultistepScheduler.step(), even with first-order update on a zero-width final step (sigma_t == sigma_s), self.noise_sampler(sigma_s, sigma_t) (BrownianTreeNoiseSampler) divides by sqrt(|t1 - t0|) = 0, producing NaN noise (0.0 * NaN = NaN in IEEE-754). Guarding self.noise_sampler when self.sigmas[self.step_index + 1] == self.sigmas[self.step_index] avoids the NaN noise.
   • This directly resolves and unskips the long-standing @unittest.skip("TODO (patil-suraj): Fix this test") on TestEDMDPMSolverMultistepScheduler.test_solver_order_and_type.

2. Avoid unnecessary scalar .item() extraction in set_timesteps():

   • When final_sigmas_type == "sigma_min", EDMDPMSolverMultistepScheduler and CosineDPMSolverMultistepScheduler called sigma_last = sigmas[-1].item() followed by torch.tensor([sigma_last], ...), whereas DPMSolverMultistepScheduler slices sigmas[-1:] directly (torch.cat([sigmas, sigmas[-1:]])).

Before submitting

• Did you read the contributor guideline?
• Did you write any new necessary tests?
  • Unskipped TestEDMDPMSolverMultistepScheduler.test_solver_order_and_type and added test_zero_width_final_step in tests/schedulers/test_scheduler_edm_dpmsolver_multistep.py and tests/schedulers/test_scheduler_dpm_multi_inverse.py (66 passed).

Credits & Acknowledgments

• Inspired by @Nicholas022400701's root-cause analysis in Converted sigma schedules with final_sigmas_type="sigma_min" repeat the last sigma, and the third-order and heun updates of DPMSolverMultistepScheduler and UniPCMultistepScheduler return NaN on that zero-width final step #14890 (and UniPCMultistepScheduler returns nan with solver_type="bh1", and inf or nan with lower_order_final=False or predict_x0=False, on the final step to sigma = 0 #14887) identifying how repeated terminal sigma_min entries (h = 0) cause NaN division-by-zero in higher-order multistep updates when lower_order_final=False, which led to auditing and fixing the sibling schedulers (EDMDPMSolverMultistepScheduler, CosineDPMSolverMultistepScheduler, and DPMSolverMultistepInverseScheduler).

[github-actions]

Hi @deep1blue, thanks for the PR! It does not appear to link an issue it fixes. If this PR addresses an existing issue, please add a closing keyword (e.g. Fixes #1234) to the PR description so the issue is linked. See the contribution guide for more details. If this PR intentionally does not fix a tracked issue, a maintainer can add the no-issue-needed label to silence this reminder.

Please note that PRs without a linked issue are likely to be automatically closed 10 days after this notice.

Once the PR links an issue (or gets the no-issue-needed label), you can ignore this message — it stays here as a comment, but it no longer applies.