fix(schedulers): prevent NaN in ancestral schedulers with squaredcos_cap_v2

[sergioperezcheco]

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

[MohammadSadeghSalehi]

You already use the closed form in scheduling_k_dpm_2_ancestral_discrete.py:

sigmas_down = sigmas_next**2 / sigmas.clamp(min=1e-30)  # sqrt(sigmas_next**4/sigmas**2)

The same identity applies to scheduling_euler_ancestral_discrete.py, where the diff still clamps instead:

sigma_up^2   = sigma_to^2 (sigma_from^2 - sigma_to^2) / sigma_from^2
sigma_down^2 = sigma_to^2 - sigma_up^2
             = sigma_to^2 [1 - (sigma_from^2 - sigma_to^2) / sigma_from^2]
             = sigma_to^4 / sigma_from^2

so sigma_down = sigma_to**2 / sigma_from exactly, and the subtraction of two nearly equal squares never happens.

The distinction matters beyond tidiness, because the clamp does not only suppress the NaN. It returns zero where the true value is small but non zero. In float32:

sigma_from	sigma_to	clamped subtractive	closed form	float64 reference
2.0e4	1.0e-3	0	5.0e-11	5.0e-11
2.0e4	1.0	0	5.0e-05	5.0e-05
1.0e3	1.0e-2	0	1.0e-07	1.0e-07
14.6	13.9	13.2336	13.2336	13.2336

squaredcos_cap_v2 gives sigma[0] around 2.0e4, so the first rows are the regime this PR is about. In those cases sigma_down becomes exactly 0, which sets dt = -sigma and sends the step to the origin in one move rather than taking the small step the schedule asks for. That is a quiet wrong trajectory rather than a visible NaN, and it is harder to notice than the bug being fixed.

The closed form also makes sigma_up cheaper, since sigma_up^2 = sigma_to^2 - sigma_down^2 follows directly and no longer needs the difference of squares either.

Checked on CPython 3.12, torch float32, macOS arm64. Happy to send this as a patch against your branch if that is useful.