From a75694fdac311ac48ba212400a520becfba10c71 Mon Sep 17 00:00:00 2001 From: Taksh Date: Fri, 23 Oct 2026 09:24:38 +0530 Subject: [PATCH] use the PLMS schedule when PNDM has fewer than four steps the order-4 PRK tile cannot be built for 1, 2, or 3 steps, and set_timesteps crashed. --- src/diffusers/schedulers/scheduling_pndm.py | 3 ++- tests/schedulers/test_scheduler_pndm.py | 7 +++++++ 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/src/diffusers/schedulers/scheduling_pndm.py b/src/diffusers/schedulers/scheduling_pndm.py index f87dff14375d..a2256a843fb3 100644 --- a/src/diffusers/schedulers/scheduling_pndm.py +++ b/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 diff --git a/tests/schedulers/test_scheduler_pndm.py b/tests/schedulers/test_scheduler_pndm.py index 13c690468222..516bddf93986 100644 --- a/tests/schedulers/test_scheduler_pndm.py +++ b/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)