fix(schedulers): prevent NaN in ancestral schedulers with squaredcos_cap_v2#14221
Open
sergioperezcheco wants to merge 1 commit into
Open
Conversation
github-actions
Bot
added
fixes-issue
size/M
sergioperezcheco
force-pushed
the
fix/ancestral-scheduler-nan-squaredcos
branch
from
July 18, 2026 04:19
f8ea78b to
cff1528
Compare
…cap_v2 EulerAncestralDiscreteScheduler.step() and KDPM2AncestralDiscreteScheduler.set_timesteps() compute sigma_down = sqrt(sigma_to**2 - sigma_up**2). sigma_up is mathematically <= sigma_to, but under beta_schedule='squaredcos_cap_v2' the sigma dynamic range is so large that in float32 sigma_up can round fractionally larger than sigma_to, making the radicand negative and producing NaN that propagates into the sample (EulerAncestral) or into self.timesteps via log(0) in sigmas_interpol (KDPM2Ancestral). For EulerAncestral, clamp the radicand to >= 0 before the square root; the sigma_up == sigma_to limit gives sigma_down == 0, which is the correct value. For KDPM2Ancestral, clamp the sigma_up radicand to >= 0, compute sigmas_down via the algebraically identical and numerically stable sqrt(sigmas_next**4 / sigmas**2), and express sigmas_interpol as the geometric mean sqrt(sigmas * sigmas_down) instead of the log-space lerp whose -inf end term produces NaN wherever sigmas_down == 0. The linear and scaled_linear paths are unaffected and produce identical output. Adds regression tests for both schedulers covering the squaredcos_cap_v2 + low-step-count regime. Fixes huggingface#14213
sergioperezcheco
force-pushed
the
fix/ancestral-scheduler-nan-squaredcos
branch
from
July 18, 2026 17:43
cff1528 to
46511b4
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
EulerAncestralDiscreteScheduler and KDPM2AncestralDiscreteScheduler both produce all-NaN output when configured with beta_schedule="squaredcos_cap_v2" at low step counts (the standard ancestral fast-sampling regime). The root cause is that sigma_down is computed as sqrt(sigma_to2 - sigma_up2): sigma_up is mathematically <= sigma_to, but the huge sigma dynamic range from squaredcos_cap_v2 means float32 rounding can make sigma_up fractionally larger than sigma_to, driving the radicand negative. In EulerAncestral the NaN reaches the sample via step(); in KDPM2Ancestral it reaches self.timesteps during set_timesteps() because the same cancellation feeds log(0) into sigmas_interpol.
For EulerAncestral I clamp the radicand to >= 0 before the square root (the sigma_up == sigma_to limit correctly gives sigma_down == 0). For KDPM2Ancestral I clamp the sigma_up radicand, compute sigmas_down via the algebraically identical and numerically stable sqrt(sigmas_next4 / sigmas2), and express sigmas_interpol as the geometric mean sqrt(sigmas * sigmas_down) instead of the log-space lerp whose -inf end term produces NaN wherever sigmas_down == 0. The linear and scaled_linear paths are unaffected and produce identical output. I verified the fix against the reproduction in the issue and added regression tests for both schedulers covering the squaredcos_cap_v2 + 4-step regime.
Fixes #14213