Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
6 changes: 4 additions & 2 deletions src/diffusers/schedulers/scheduling_flow_match_lcm.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
4 changes: 3 additions & 1 deletion src/diffusers/schedulers/scheduling_unipc_multistep.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
59 changes: 59 additions & 0 deletions tests/schedulers/test_scheduler_flow_match_euler_discrete.py
Original file line number Diff line number Diff line change
@@ -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())
9 changes: 9 additions & 0 deletions tests/schedulers/test_scheduler_unipc.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Loading