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.