From 1fbe16bffecbcec8c2243a1487596cfd20d9348d Mon Sep 17 00:00:00 2001 From: kyo-zzz <1147385157@qq.com> Date: Mon, 24 Aug 2026 05:39:27 +0800 Subject: [PATCH] fix: raise clear error in UnCLIPScheduler.set_timesteps for fewer than 2 steps The karlo-style step ratio divides by num_inference_steps - 1, so num_inference_steps=1 crashed with a bare ZeroDivisionError and 0 silently produced an empty schedule. Validate the input up front (same approach as MiniMaxH3Scheduler) and document the constraint. n=2 still spans the full training schedule: [999, 0]. --- src/diffusers/schedulers/scheduling_unclip.py | 6 +++++- tests/schedulers/test_scheduler_unclip.py | 16 ++++++++++++++++ 2 files changed, 21 insertions(+), 1 deletion(-) diff --git a/src/diffusers/schedulers/scheduling_unclip.py b/src/diffusers/schedulers/scheduling_unclip.py index 747453eb2d40..9d689fec9715 100644 --- a/src/diffusers/schedulers/scheduling_unclip.py +++ b/src/diffusers/schedulers/scheduling_unclip.py @@ -183,10 +183,14 @@ def set_timesteps(self, num_inference_steps: int, device: str | torch.device | N Args: num_inference_steps (`int`): - The number of diffusion steps used when generating samples with a pre-trained model. + The number of diffusion steps used when generating samples with a pre-trained model. Must be at + least 2: the step ratio interpolates between the two ends of the training schedule + (`num_train_timesteps - 1` divided by `num_inference_steps - 1`). device (`str` or `torch.device`, *optional*): The device to which the timesteps are moved. If `None`, the timesteps are not moved. """ + if num_inference_steps < 2: + raise ValueError(f"`set_timesteps` requires `num_inference_steps` >= 2, but got {num_inference_steps}.") self.num_inference_steps = num_inference_steps step_ratio = (self.config.num_train_timesteps - 1) / (self.num_inference_steps - 1) timesteps = (np.arange(0, num_inference_steps) * step_ratio).round()[::-1].copy().astype(np.int64) diff --git a/tests/schedulers/test_scheduler_unclip.py b/tests/schedulers/test_scheduler_unclip.py index 9e66a328f42e..8b788a830fbc 100644 --- a/tests/schedulers/test_scheduler_unclip.py +++ b/tests/schedulers/test_scheduler_unclip.py @@ -27,6 +27,22 @@ def test_timesteps(self): for timesteps in [1, 5, 100, 1000]: self.check_over_configs(num_train_timesteps=timesteps) + def test_set_timesteps_at_least_two_steps(self): + scheduler_class = self.scheduler_classes[0] + scheduler_config = self.get_scheduler_config() + scheduler = scheduler_class(**scheduler_config) + + # 0 and 1 steps have no valid karlo-style schedule (the step ratio divides by + # num_inference_steps - 1) and must fail with a clear error instead of a ZeroDivisionError + for num_inference_steps in [0, 1]: + with self.assertRaises(ValueError) as context: + scheduler.set_timesteps(num_inference_steps) + self.assertIn("num_inference_steps", str(context.exception)) + + # 2 steps is the lowest valid input and spans the full training schedule + scheduler.set_timesteps(2) + self.assertEqual(scheduler.timesteps.tolist(), [999, 0]) + def test_variance_type(self): for variance in ["fixed_small_log", "learned_range"]: self.check_over_configs(variance_type=variance)