Fix NaN in KDPM2DiscreteScheduler.sigmas_interpol (all-NaN output on MPS)

[4ktLuffy]

What this fixes

Fixes #14368 (the KDPM2DiscreteScheduler half).

sigmas_interpol is the geometric mean of each sigma and its predecessor, computed through log space:

sigmas_interpol = sigmas.log().lerp(sigmas.roll(1).log(), 0.5).exp()

The sigma schedule always ends in a zero sigma, so log(0) = -inf enters that lerp. torch resolves the -inf differently on CPU and MPS (pytorch/pytorch#111374, open since 2023), which is why the same line behaves differently per device:

sigmas_interpol, 4 steps, default config
cpu:  [nan, 10.632364, 10.632364, 2.336412, 2.336412, 0.188194, 0.188194, 0.0, 0.0, 0.0]
mps:  [0.0, 10.632364, 10.632364, 2.336412, 2.336412, 0.188194, 0.188194, nan, nan, nan]

• On CPU the NaN lands on sigmas_interpol[0], a wrap-around entry from roll(1) that step() never reads. Harmless in effect, but sigmas_interpol is a public attribute and it was non-finite in 24 of 24 configurations I tested.
• On MPS the NaN lands on the trailing live entries instead, so dt = sigma_interpol - sigma_hat is NaN on the final step and the whole latent is destroyed.

End to end, on a tiny StableDiffusionPipeline with this scheduler:

	NaN pixels	mean
before, cpu	0 / 49152	0.5080
before, mps	49152 / 49152	nan
after, cpu	0 / 49152	0.5080
after, mps	0 / 49152	0.5080

The change

Compute the geometric mean directly, which never touches log(0):

sigmas_interpol = (sigmas * sigmas.roll(1)).sqrt()

exp(½·log a + ½·log b) == sqrt(a·b) exactly for positive sigmas, so this is the same quantity by a different route.

On accuracy — the new form is not just equivalent but slightly better. Against a float64 ground truth on a real 25-step schedule:

	max abs err	max rel err
old, exp(lerp(log, log))	3.02e-06	1.36e-07
new, sqrt(a*b)	1.12e-06	5.02e-08

More accurate on 8 of 25 entries, equal on 17, worse on none. Existing CPU/CUDA results therefore shift by a small amount (largest observed output delta 2.4e-04 after 25 steps) — toward the exact value, not away from it. Happy to keep the old expression under a device check instead if you would rather have bit-identical CUDA output; this seemed cleaner.

On overflow — the largest sigma across all configs I tried is 2.03e+04, whose square is 4.1e+08, against a float32 max of 3.4e+38.

Verification

• sigmas_interpol finite on both devices across 48 configurations (2 schedulers × 8 configs × 3 step counts): karras / exponential / beta sigmas, all three timestep_spacing modes, squaredcos_cap_v2 and scaled_linear.
• After the fix, MPS output matches CPU exactly — max delta 0.0 across the 15 configurations that run on both.
• Full tests/schedulers/ suite: 977 passed, 6 failed — identical list before and after this change. All 6 are pre-existing local MPS failures unrelated to this line (test_full_loop_with_v_prediction here is the separate CPU-scalar storage_offset bug; it passes on a torch 2.14 nightly, with and without this patch).

Test

test_set_timesteps_sigmas_interpol_no_nan asserts sigmas_interpol is finite at 4/10/25 steps. It fails on main and passes with this change — and because CPU also carried the NaN at index 0, it fails on CPU too, so CI will catch a regression without needing Apple hardware.

Scope

The ancestral variant has the same log/lerp pattern, but I have deliberately left it alone: #14221 is already rewriting that line (as (sigmas.clamp(min=0) * sigmas_down) ** 0.5) while fixing #14213. This PR touches only the non-ancestral scheduler to avoid conflicting with it.