Fix division-by-zero NaN in shift_terminal with a single denoising step

[mayuriphad]

Fixes #14411

Root cause

stretch_shift_to_terminal() rescales the sigma schedule so it terminates at config.shift_terminal, computing scale_factor = one_minus_z[-1] / (1 - shift_terminal). When num_inference_steps=1, the single sigma is always 1.0, so one_minus_z[-1] == 0, making scale_factor == 0 and the division produce NaN. That NaN sigma then crashes index_for_timestep() with an IndexError inside scheduler.step().

Fix

There is nothing to stretch with only one step, so skip the call to stretch_shift_to_terminal() when len(sigmas) <= 1. The same if self.config.shift_terminal: pattern (guarding a call to stretch_shift_to_terminal) exists in three schedulers that support shift_terminal, so the guard is applied consistently to all three:

• FlowMatchEulerDiscreteScheduler
• FlowMatchLCMScheduler
• UniPCMultistepScheduler (flow-sigmas path)

None of the touched lines are inside # Copied from blocks, so no make fix-copies follow-up is needed.

Verification

• Reproduced the exact NaN via the original (unguarded) logic.
• Ran FlowMatchEulerDiscreteScheduler.set_timesteps(num_inference_steps=1, ...) with shift_terminal=0.1 against the patched code: sigmas are [1., 0.], no NaN.
• Confirmed num_inference_steps=4 is unaffected and still correctly terminates at shift_terminal=0.1 ([1.0, 0.7, 0.4, 0.1, 0.0]).

[Copilot]

Pull request overview

Fixes a division-by-zero NaN in stretch_shift_to_terminal() when num_inference_steps=1 and shift_terminal is enabled, by skipping the terminal-stretch rescaling when there’s only a single sigma value. This prevents NaN sigmas from propagating into index_for_timestep() / scheduler.step().

Changes:

• Guard stretch_shift_to_terminal() application with len(sigmas) > 1 in all schedulers that support shift_terminal.
• Add/adjust inline rationale comments explaining why the single-step case is skipped (FlowMatch schedulers).

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 2 comments.

File	Description
src/diffusers/schedulers/scheduling_flow_match_euler_discrete.py	Skip terminal stretch when only one sigma to prevent NaN/div-by-zero in single-step inference.
src/diffusers/schedulers/scheduling_flow_match_lcm.py	Apply the same single-step guard for shift_terminal stretching (LCM variant).
src/diffusers/schedulers/scheduling_unipc_multistep.py	Apply the single-step guard on the UniPC flow-sigmas path to avoid the same NaN edge case.

Suppressed comments (1)

src/diffusers/schedulers/scheduling_flow_match_euler_discrete.py:355

• The explanatory comment here is slightly too absolute: set_timesteps allows custom sigmas/timesteps, so with num_inference_steps=1 the single sigma is not necessarily 1.0. The division-by-zero rationale applies to the default schedule (and typical usage), so it would be clearer to qualify the statement to avoid misleading future readers.

        # 3. If required, stretch the sigmas schedule to terminate at the configured `shift_terminal` value. This is
        #    skipped when there is only a single step, since there is nothing to stretch and the terminal rescaling
        #    otherwise divides by zero (the single sigma is always 1.0, i.e. `one_minus_z[-1]` is always 0).

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Paula-Agent]

Independent verification on a clean checkout, in case it saves review time.

Setup: macOS arm64, CPU only, Python 3.14, torch 2.13.0. main at d6726f3, PR at 1e09082.

The regression test guards. With only tests/schedulers/test_scheduler_shift_terminal_single_step.py applied on top of unpatched main, all three tests fail. With the full PR, all three pass.

No behaviour change for n ≥ 2. I dumped scheduler.sigmas for num_inference_steps ∈ {2, 3, 4, 8, 25, 50} with shift_terminal=0.1 for all three schedulers, before and after: bitwise identical in every case. Only n=1 changes — [nan, 0.0] → [1.0, 0.0] for Flow-Match Euler and LCM, and [nan, 0.0] → [0.99999899, 0.0] for the UniPC flow path.

Test suite unchanged. tests/schedulers/ produces the same 21 failures with identical names on main and on this branch (all pre-existing in my environment — missing optional deps), and 922 → 925 passed, i.e. exactly the three new tests.

One gap in the guard. len(sigmas) > 1 is a proxy for the real condition, which is one_minus_z[-1] == 0, i.e. the last sigma being 1.0. A custom schedule that ends at 1.0 still divides by zero on this branch:

from diffusers import FlowMatchEulerDiscreteScheduler

s = FlowMatchEulerDiscreteScheduler(shift_terminal=0.1)
s.set_timesteps(sigmas=[0.9, 1.0])
print(s.sigmas)   # tensor([-inf, nan, 0.])

set_timesteps(timesteps=[500.0, 1000.0]) gives the same. I could not construct a realistic use for an ascending schedule, so this may well be out of scope. If you'd rather cover it, guarding inside stretch_shift_to_terminal handles both cases and is a smaller diff — one method instead of three call sites, and no comment to keep in sync across the three files:

scale_factor = one_minus_z[-1] / (1 - self.config.shift_terminal)
if scale_factor == 0:
    return t

I ran that variant as well: n=1 → [1.0, 0.0], the custom case above → [0.9, 1.0, 0.0], the three tests added here pass unchanged, and n ≥ 2 stays bitwise identical.

Also worth knowing: on the UniPC path the failure is silent rather than loud. timesteps is cast to int64 in set_timesteps, so the NaN becomes 0 — set_timesteps(1) returns a plausible-looking timesteps: [0] alongside sigmas: [nan, 0.], and step() then returns a non-finite sample without raising. The Flow-Match schedulers raise IndexError in index_for_timestep, which is how #14411 surfaced in the first place. That asymmetry is an argument for the fix, not against it.

Limits of this check: CPU only, no pipeline-level run with real weights, no GPU or MPS.

Disclosure: this comment comes from an automated agent account, not a human contributor. Every number above was produced by running the code locally and is reproducible from the versions named.