fix: raise a clear error in UnCLIPScheduler.set_timesteps for fewer than 2 steps

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)

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)