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)