From 83e0624c36e856d23075f079c37cf17e07612ed2 Mon Sep 17 00:00:00 2001 From: promptsmith1990 Date: Sun, 23 Aug 2026 11:41:02 -0700 Subject: [PATCH] Fix nan sigmas when num_inference_steps=1 with shift_terminal FlowMatchEulerDiscreteScheduler.stretch_shift_to_terminal computes one_minus_z[-1] / scale_factor, where scale_factor is itself derived from one_minus_z[-1]. With num_inference_steps=1 the single sigma in the schedule is both the first and last point, and before any terminal stretching it is exactly 1.0 (pure noise), so one_minus_z[-1] is 0 and the division degenerates to 0/0, producing nan sigmas and timesteps. Any pipeline configured with shift_terminal (LTX, LTX2, ...) that is run with a single denoising step hits this. The same unguarded `if self.config.shift_terminal:` call-site pattern was copy-pasted into two other schedulers that share the identical stretch_shift_to_terminal implementation, UniPCMultistepScheduler's use_flow_sigmas branch and FlowMatchLCMScheduler, and both reproduce the same nan crash under the same conditions. Fixes this by skipping the stretch when there is only one point in the schedule to stretch, in all three call sites. Multi-step schedules are unaffected and still terminate at shift_terminal as before. Fixes #14411 Test Plan: pytest tests/schedulers/test_scheduler_flow_match_euler_discrete.py -v pytest tests/schedulers/test_scheduler_unipc.py -k flow_sigmas_single_step -v New test_scheduler_flow_match_euler_discrete.py covers: multi-step still reaches shift_terminal (regression guard), single-step no longer produces nan sigmas/timesteps, and a single scheduler.step() call runs to completion without nan output. A matching regression test was added to test_scheduler_unipc.py for the use_flow_sigmas branch. Also ran utils/check_copies.py and ruff check/format on the changed files (clean) to make sure the stretch_shift_to_terminal "# Copied from" block itself was left untouched. This PR was prepared with AI assistance (Claude Code): the assistant found the bug while investigating issue #14411, traced the same defect into the two other schedulers, wrote the fix and tests, and ran the verification above. I reviewed the diff and the verification output before submitting. Self-review (against .ai/references/review-rules.md): no blocking issues found. One scope note left for the actual review: the new test_scheduler_flow_match_euler_discrete.py uses a plain unittest.TestCase rather than SchedulerCommonTest, since FlowMatchEulerDiscreteScheduler currently has no dedicated test file at all and adopting the full common-test harness (dummy sample generation, forward-pass parity checks, etc.) for the first time is a separate, larger effort than this bug fix. --- .../scheduling_flow_match_euler_discrete.py | 7 ++- .../schedulers/scheduling_flow_match_lcm.py | 6 +- .../schedulers/scheduling_unipc_multistep.py | 4 +- ...est_scheduler_flow_match_euler_discrete.py | 59 +++++++++++++++++++ tests/schedulers/test_scheduler_unipc.py | 9 +++ 5 files changed, 80 insertions(+), 5 deletions(-) create mode 100644 tests/schedulers/test_scheduler_flow_match_euler_discrete.py diff --git a/src/diffusers/schedulers/scheduling_flow_match_euler_discrete.py b/src/diffusers/schedulers/scheduling_flow_match_euler_discrete.py index 0e4f5c6a1f97..e1f40807bd20 100644 --- a/src/diffusers/schedulers/scheduling_flow_match_euler_discrete.py +++ b/src/diffusers/schedulers/scheduling_flow_match_euler_discrete.py @@ -350,8 +350,11 @@ def set_timesteps( else: sigmas = self.shift * sigmas / (1 + (self.shift - 1) * sigmas) - # 3. If required, stretch the sigmas schedule to terminate at the configured `shift_terminal` value - if self.config.shift_terminal: + # 3. If required, stretch the sigmas schedule to terminate at the configured `shift_terminal` value. + # With a single step, the lone sigma is both the first and last point of the schedule, so + # `one_minus_z[-1]` is 0 and the stretch factor is a 0/0 division -> nan. Skip stretching in that + # case; there is nothing to stretch a single-point schedule against. + if self.config.shift_terminal and len(sigmas) > 1: sigmas = self.stretch_shift_to_terminal(sigmas) # 4. If required, convert sigmas to one of karras, exponential, or beta sigma schedules diff --git a/src/diffusers/schedulers/scheduling_flow_match_lcm.py b/src/diffusers/schedulers/scheduling_flow_match_lcm.py index 97d4ebbc8e42..8de9ad2199c6 100644 --- a/src/diffusers/schedulers/scheduling_flow_match_lcm.py +++ b/src/diffusers/schedulers/scheduling_flow_match_lcm.py @@ -359,8 +359,10 @@ def set_timesteps( else: sigmas = self.shift * sigmas / (1 + (self.shift - 1) * sigmas) # type: ignore - # 3. If required, stretch the sigmas schedule to terminate at the configured `shift_terminal` value - if self.config.shift_terminal: + # 3. If required, stretch the sigmas schedule to terminate at the configured `shift_terminal` value. + # With a single step, the lone sigma is both the first and last point of the schedule, so + # `stretch_shift_to_terminal` divides 0 by 0 and produces nan. Skip stretching in that case. + if self.config.shift_terminal and len(sigmas) > 1: sigmas = self.stretch_shift_to_terminal(sigmas) # type: ignore # 4. If required, convert sigmas to one of karras, exponential, or beta sigma schedules diff --git a/src/diffusers/schedulers/scheduling_unipc_multistep.py b/src/diffusers/schedulers/scheduling_unipc_multistep.py index 5c2cbcc13ff1..b309b192be7f 100644 --- a/src/diffusers/schedulers/scheduling_unipc_multistep.py +++ b/src/diffusers/schedulers/scheduling_unipc_multistep.py @@ -432,7 +432,9 @@ def set_timesteps( sigmas = self.time_shift(mu, 1.0, sigmas) else: sigmas = self.config.flow_shift * sigmas / (1 + (self.config.flow_shift - 1) * sigmas) - if self.config.shift_terminal: + # With a single step, the lone sigma is both the first and last point of the schedule, so + # `stretch_shift_to_terminal` divides 0 by 0 and produces nan. Skip stretching in that case. + if self.config.shift_terminal and len(sigmas) > 1: sigmas = self.stretch_shift_to_terminal(sigmas) eps = 1e-6 if np.fabs(sigmas[0] - 1) < eps: diff --git a/tests/schedulers/test_scheduler_flow_match_euler_discrete.py b/tests/schedulers/test_scheduler_flow_match_euler_discrete.py new file mode 100644 index 000000000000..6a5719aeb87d --- /dev/null +++ b/tests/schedulers/test_scheduler_flow_match_euler_discrete.py @@ -0,0 +1,59 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import unittest + +import torch + +from diffusers import FlowMatchEulerDiscreteScheduler + + +class FlowMatchEulerDiscreteSchedulerTest(unittest.TestCase): + scheduler_class = FlowMatchEulerDiscreteScheduler + + def get_default_config(self, **kwargs): + config = { + "num_train_timesteps": 1000, + "shift": 3.0, + } + config.update(**kwargs) + return config + + def test_set_timesteps_multi_step_reaches_shift_terminal(self): + # Sanity check that the multi-step schedule still stretches to shift_terminal; + # guards against the single-step fix below accidentally disabling stretching generally. + scheduler = self.scheduler_class(**self.get_default_config(shift_terminal=0.1)) + scheduler.set_timesteps(num_inference_steps=10) + self.assertFalse(torch.isnan(scheduler.sigmas).any()) + self.assertAlmostEqual(scheduler.sigmas[-2].item(), 0.1, places=5) + + def test_set_timesteps_single_step_with_shift_terminal_is_finite(self): + # With num_inference_steps=1 the lone sigma is both the first and last point of the + # schedule, so stretch_shift_to_terminal's `one_minus_z[-1] / scale_factor` degenerates to + # a 0/0 division and produced nan timesteps/sigmas. See gh-14411. + scheduler = self.scheduler_class(**self.get_default_config(shift_terminal=0.1)) + scheduler.set_timesteps(num_inference_steps=1) + self.assertFalse(torch.isnan(scheduler.sigmas).any()) + self.assertFalse(torch.isnan(scheduler.timesteps).any()) + + def test_step_single_step_with_shift_terminal_runs(self): + scheduler = self.scheduler_class(**self.get_default_config(shift_terminal=0.1)) + scheduler.set_timesteps(num_inference_steps=1) + + sample = torch.randn(1, 4, 4, 4) + model_output = torch.randn_like(sample) + + prev_sample = scheduler.step(model_output, scheduler.timesteps[0], sample).prev_sample + self.assertEqual(prev_sample.shape, sample.shape) + self.assertFalse(torch.isnan(prev_sample).any()) diff --git a/tests/schedulers/test_scheduler_unipc.py b/tests/schedulers/test_scheduler_unipc.py index ac7e1d3f88b4..5e53467df34f 100644 --- a/tests/schedulers/test_scheduler_unipc.py +++ b/tests/schedulers/test_scheduler_unipc.py @@ -403,6 +403,15 @@ def test_exponential_sigmas(self): def test_flow_and_karras_sigmas(self): self.check_over_configs(use_flow_sigmas=True, use_karras_sigmas=True) + def test_flow_sigmas_single_step_with_shift_terminal_is_finite(self): + # With num_inference_steps=1 the lone sigma is both the first and last point of the + # schedule, so stretch_shift_to_terminal's `one_minus_z[-1] / scale_factor` degenerates to + # a 0/0 division and produced nan timesteps/sigmas. See gh-14411. + scheduler = UniPCMultistepScheduler(use_flow_sigmas=True, flow_shift=3.0, shift_terminal=0.1) + scheduler.set_timesteps(num_inference_steps=1) + self.assertFalse(torch.isnan(scheduler.sigmas).any()) + self.assertFalse(torch.isnan(scheduler.timesteps).any()) + def test_flow_and_karras_sigmas_values(self): num_train_timesteps = 1000 num_inference_steps = 5