use the PLMS schedule when PNDM has fewer than four steps

src/diffusers/schedulers/scheduling_pndm.py

@@ -205,7 +205,8 @@ def set_timesteps(self, num_inference_steps: int, device: str | torch.device = N
                 f"{self.config.timestep_spacing} is not supported. Please make sure to choose one of 'linspace', 'leading' or 'trailing'."
             )
 
-        if self.config.skip_prk_steps:
+        # The PRK tile is order 4. With 1, 2, or 3 steps that tile cannot be built, so use PLMS.
+        if self.config.skip_prk_steps or len(self._timesteps) < self.pndm_order:
             # for some models like stable diffusion the prk steps can/should be skipped to
             # produce better results. When using PNDM with `self.config.skip_prk_steps` the implementation
             # is based on crowsonkb's PLMS sampler implementation: https://github.com/CompVis/latent-diffusion/pull/51

tests/schedulers/test_scheduler_pndm.py

@@ -234,6 +234,13 @@ def test_full_loop_with_set_alpha_to_one(self):
         assert abs(result_sum.item() - 230.0399) < 1e-2
         assert abs(result_mean.item() - 0.2995) < 1e-3
 
+    def test_set_timesteps_fewer_than_prk_order(self):
+        # pndm_order is 4. 1, 2, or 3 steps used to crash in the PRK tile.
+        scheduler = self.scheduler_classes[0](**self.get_scheduler_config())
+        for steps in (1, 2, 3, 4):
+            scheduler.set_timesteps(steps)
+            self.assertGreater(len(scheduler.timesteps), 0)
+
     def test_full_loop_with_no_set_alpha_to_one(self):
         # We specify different beta, so that the first alpha is 0.99
         sample = self.full_loop(set_alpha_to_one=False, beta_start=0.01)