Skip to content

Fix NaN corruption on final step in DPMSolverMultistep and UniPCMultistep with final_sigmas_type="sigma_min" - #15072

Open
deep1blue wants to merge 1 commit into
huggingface:mainfrom
deep1blue:fix-dpmsolver-unipc-sigma-min-nan
Open

deep1blue wants to merge 1 commit into
huggingface:mainfrom
deep1blue:fix-dpmsolver-unipc-sigma-min-nan

Conversation

@deep1blue

Copy link
Copy Markdown

What does this PR do?

Fixes #14890.

Root Cause

When final_sigmas_type="sigma_min" is enabled alongside sigma schedules whose final step already reaches sigma_min (use_karras_sigmas=True, use_exponential_sigmas=True, use_beta_sigmas=True, or use_lu_lambdas=True), set_timesteps() appends sigma_last = sigmas[-1] to the sigma array, making the last two entries identical (sigmas[-2] == sigmas[-1]).

On the final denoising step (step_index == num_inference_steps - 1), this produces a zero-width log-SNR step (h = lambda_t - lambda_s == 0.0). While dpm_solver_first_order_update safely multiplies by expm1(-h) = 0.0 (leaving the sample unchanged), the 2nd-order and 3rd-order multistep updates in DPMSolverMultistepScheduler (and the multistep predictor/corrector in UniPCMultistepScheduler) divide h = 0.0 by r0 = h_0 / h = inf (or compute hh / h_0 where h_0 = 0.0), producing 0.0 / 0.0 = NaN and corrupting the entire output tensor on the very last step.

Solution

  • In DPMSolverMultistepScheduler.step(): include (self.step_index == self.num_inference_steps - 1 and self.sigmas[self.step_index + 1] == self.sigmas[self.step_index]) in lower_order_final so zero-width terminal steps use dpm_solver_first_order_update.
  • In UniPCMultistepScheduler.step(): bypass the multistep corrector when the previous step had zero width (self.sigmas[self.step_index] == self.sigmas[self.step_index - 1]) and fall back to first-order multistep_uni_p_bh_update(..., order=1) when the final step has zero width (self.sigmas[self.step_index + 1] == self.sigmas[self.step_index]).
  • Added unit tests test_zero_width_final_step in both tests/schedulers/test_scheduler_dpm_multi.py and tests/schedulers/test_scheduler_unipc.py.

…step with final_sigmas_type=sigma_min

When final_sigmas_type='sigma_min' is combined with schedules whose last step already reaches sigma_min (e.g., use_karras_sigmas, use_exponential_sigmas, use_beta_sigmas, or use_lu_lambdas), the appended terminal sigma duplicates the penultimate sigma, producing h = lambda_t - lambda_s = 0.0 on the final step and 0.0 / 0.0 = NaN in higher-order updates. Fall back to the first-order update on the final step when sigmas[step_index + 1] == sigmas[step_index].

Fixes huggingface#14890

Signed-off-by: Parth Vadera <parthvadera@gmail.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
@deep1blue

Copy link
Copy Markdown
Author

Gentle, polite ping @stevhliu @sayakpaul @yiyixuxu @DN6 @hlky — whenever you have a spare moment, would you mind taking a quick look at this small, self-contained scheduler fix for #14890? When final_sigmas_type="sigma_min" duplicates the terminal sigma (sigmas[-2] == sigmas[-1], h == 0), falling back to the exact first-order update on the zero-width terminal step avoids the 0.0 / 0.0 = NaN division in DPMSolverMultistepScheduler and UniPCMultistepScheduler, verified with unit tests across both schedulers. Thank you so much for your time!

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

1 participant