From 785bcfaa2ecb774c3c5739908460a1563713baba Mon Sep 17 00:00:00 2001 From: Mohammad Sadegh Salehi Date: Sun, 2 Aug 2026 21:05:13 +0100 Subject: [PATCH 1/2] Derive DDIM prev_timestep from the inference schedule DDIMScheduler.step used a uniform stride, num_train_timesteps // num_inference_steps, to form prev_timestep. That matches leading and trailing spacing but not linspace: for 1000 train steps and 10 inference steps the schedule steps by 111 while the formula steps by 100, so alpha_prod_t_prev is read from the wrong index. Look up the previous value as the next entry of self.timesteps, as DDPMScheduler.previous_timestep already does. Leading and trailing stay bit-identical: intermediate prev values match the old stride, and the final step still selects final_alpha_cumprod whether the formula yields -100 or the list yields -1. Default spacing remains leading, so only callers that set linspace change. Apply the same helper to DDIMParallel, CogVideoXDDIM and CogVideoXDPM via # Copied from. PNDM and RePaint keep the stride formula: PLMS/PRK half-steps and RePaint jump schedules are not consecutive schedule entries. Fixes #12633 Fixes #11347 --- src/diffusers/schedulers/scheduling_ddim.py | 40 +++++++++++++++- .../schedulers/scheduling_ddim_cogvideox.py | 41 +++++++++++++++- .../schedulers/scheduling_ddim_parallel.py | 47 +++++++++++++++++-- .../schedulers/scheduling_dpm_cogvideox.py | 41 +++++++++++++++- tests/schedulers/test_scheduler_ddim.py | 38 +++++++++++++++ 5 files changed, 196 insertions(+), 11 deletions(-) diff --git a/src/diffusers/schedulers/scheduling_ddim.py b/src/diffusers/schedulers/scheduling_ddim.py index d29bb5b2593c..8d8fbd3c7146 100644 --- a/src/diffusers/schedulers/scheduling_ddim.py +++ b/src/diffusers/schedulers/scheduling_ddim.py @@ -381,6 +381,42 @@ def set_timesteps(self, num_inference_steps: int, device: str | torch.device = N self.timesteps = torch.from_numpy(timesteps).to(device) + def previous_timestep(self, timestep: int | torch.Tensor) -> int | torch.Tensor: + """ + Return the previous timestep from the schedule produced by [`set_timesteps`]. + + The schedule already encodes `timestep_spacing` (`leading`, `trailing`, or `linspace`). Deriving the previous + value from that list is a no-op for uniform strides and corrects `linspace`, where a fixed + `num_train_timesteps // num_inference_steps` step disagrees with the materialised list. + + Args: + timestep (`int` or `torch.Tensor`): + The current discrete timestep in the diffusion chain (an entry of `self.timesteps`). + + Returns: + `int` or `torch.Tensor`: + The previous timestep. The final schedule entry returns `-1`, which selects `final_alpha_cumprod`. + """ + if self.num_inference_steps is None: + raise ValueError( + "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler" + ) + + if isinstance(timestep, torch.Tensor) and timestep.ndim > 0: + flat = timestep.flatten() + prev = torch.stack([self.previous_timestep(t) for t in flat]).reshape(timestep.shape) + return prev + + index_candidates = (self.timesteps == timestep).nonzero(as_tuple=True)[0] + if len(index_candidates) == 0: + # Not an entry of the inference schedule (direct step() calls outside the loop). + # Keep the historical unit-stride fallback so off-schedule tooling still works. + return timestep - self.config.num_train_timesteps // self.num_inference_steps + index = index_candidates[0] + if index == self.timesteps.shape[0] - 1: + return torch.tensor(-1, device=self.timesteps.device, dtype=self.timesteps.dtype) + return self.timesteps[index + 1] + def step( self, model_output: torch.Tensor, @@ -441,8 +477,8 @@ def step( # - pred_sample_direction -> "direction pointing to x_t" # - pred_prev_sample -> "x_t-1" - # 1. get previous step value (=t-1) - prev_timestep = timestep - self.config.num_train_timesteps // self.num_inference_steps + # 1. get previous step value from the materialised schedule (= next entry of self.timesteps) + prev_timestep = self.previous_timestep(timestep) # 2. compute alphas, betas alpha_prod_t = self.alphas_cumprod[timestep] diff --git a/src/diffusers/schedulers/scheduling_ddim_cogvideox.py b/src/diffusers/schedulers/scheduling_ddim_cogvideox.py index 4cb9cedce29c..a6d22c898229 100644 --- a/src/diffusers/schedulers/scheduling_ddim_cogvideox.py +++ b/src/diffusers/schedulers/scheduling_ddim_cogvideox.py @@ -325,6 +325,43 @@ def set_timesteps( self.timesteps = torch.from_numpy(timesteps).to(device) + # Copied from diffusers.schedulers.scheduling_ddim.DDIMScheduler.previous_timestep + def previous_timestep(self, timestep: int | torch.Tensor) -> int | torch.Tensor: + """ + Return the previous timestep from the schedule produced by [`set_timesteps`]. + + The schedule already encodes `timestep_spacing` (`leading`, `trailing`, or `linspace`). Deriving the previous + value from that list is a no-op for uniform strides and corrects `linspace`, where a fixed + `num_train_timesteps // num_inference_steps` step disagrees with the materialised list. + + Args: + timestep (`int` or `torch.Tensor`): + The current discrete timestep in the diffusion chain (an entry of `self.timesteps`). + + Returns: + `int` or `torch.Tensor`: + The previous timestep. The final schedule entry returns `-1`, which selects `final_alpha_cumprod`. + """ + if self.num_inference_steps is None: + raise ValueError( + "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler" + ) + + if isinstance(timestep, torch.Tensor) and timestep.ndim > 0: + flat = timestep.flatten() + prev = torch.stack([self.previous_timestep(t) for t in flat]).reshape(timestep.shape) + return prev + + index_candidates = (self.timesteps == timestep).nonzero(as_tuple=True)[0] + if len(index_candidates) == 0: + # Not an entry of the inference schedule (direct step() calls outside the loop). + # Keep the historical unit-stride fallback so off-schedule tooling still works. + return timestep - self.config.num_train_timesteps // self.num_inference_steps + index = index_candidates[0] + if index == self.timesteps.shape[0] - 1: + return torch.tensor(-1, device=self.timesteps.device, dtype=self.timesteps.dtype) + return self.timesteps[index + 1] + def step( self, model_output: torch.Tensor, @@ -384,8 +421,8 @@ def step( # - pred_sample_direction -> "direction pointing to x_t" # - pred_prev_sample -> "x_t-1" - # 1. get previous step value (=t-1) - prev_timestep = timestep - self.config.num_train_timesteps // self.num_inference_steps + # 1. get previous step value from the materialised schedule (= next entry of self.timesteps) + prev_timestep = self.previous_timestep(timestep) # 2. compute alphas, betas alpha_prod_t = self.alphas_cumprod[timestep] diff --git a/src/diffusers/schedulers/scheduling_ddim_parallel.py b/src/diffusers/schedulers/scheduling_ddim_parallel.py index a5bbac60275c..014613a26196 100644 --- a/src/diffusers/schedulers/scheduling_ddim_parallel.py +++ b/src/diffusers/schedulers/scheduling_ddim_parallel.py @@ -268,7 +268,7 @@ def scale_model_input(self, sample: torch.Tensor, timestep: int | None = None) - def _get_variance(self, timestep: int, prev_timestep: int | None = None) -> torch.Tensor: if prev_timestep is None: - prev_timestep = timestep - self.config.num_train_timesteps // self.num_inference_steps + prev_timestep = self.previous_timestep(timestep) alpha_prod_t = self.alphas_cumprod[timestep] alpha_prod_t_prev = self.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else self.final_alpha_cumprod @@ -385,6 +385,43 @@ def set_timesteps(self, num_inference_steps: int, device: str | torch.device = N self.timesteps = torch.from_numpy(timesteps).to(device) + # Copied from diffusers.schedulers.scheduling_ddim.DDIMScheduler.previous_timestep + def previous_timestep(self, timestep: int | torch.Tensor) -> int | torch.Tensor: + """ + Return the previous timestep from the schedule produced by [`set_timesteps`]. + + The schedule already encodes `timestep_spacing` (`leading`, `trailing`, or `linspace`). Deriving the previous + value from that list is a no-op for uniform strides and corrects `linspace`, where a fixed + `num_train_timesteps // num_inference_steps` step disagrees with the materialised list. + + Args: + timestep (`int` or `torch.Tensor`): + The current discrete timestep in the diffusion chain (an entry of `self.timesteps`). + + Returns: + `int` or `torch.Tensor`: + The previous timestep. The final schedule entry returns `-1`, which selects `final_alpha_cumprod`. + """ + if self.num_inference_steps is None: + raise ValueError( + "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler" + ) + + if isinstance(timestep, torch.Tensor) and timestep.ndim > 0: + flat = timestep.flatten() + prev = torch.stack([self.previous_timestep(t) for t in flat]).reshape(timestep.shape) + return prev + + index_candidates = (self.timesteps == timestep).nonzero(as_tuple=True)[0] + if len(index_candidates) == 0: + # Not an entry of the inference schedule (direct step() calls outside the loop). + # Keep the historical unit-stride fallback so off-schedule tooling still works. + return timestep - self.config.num_train_timesteps // self.num_inference_steps + index = index_candidates[0] + if index == self.timesteps.shape[0] - 1: + return torch.tensor(-1, device=self.timesteps.device, dtype=self.timesteps.dtype) + return self.timesteps[index + 1] + def step( self, model_output: torch.Tensor, @@ -440,8 +477,8 @@ def step( # - pred_sample_direction -> "direction pointing to x_t" # - pred_prev_sample -> "x_t-1" - # 1. get previous step value (=t-1) - prev_timestep = timestep - self.config.num_train_timesteps // self.num_inference_steps + # 1. get previous step value from the materialised schedule (= next entry of self.timesteps) + prev_timestep = self.previous_timestep(timestep) # 2. compute alphas, betas alpha_prod_t = self.alphas_cumprod[timestep] @@ -565,9 +602,9 @@ def batch_step_no_noise( # - pred_sample_direction -> "direction pointing to x_t" # - pred_prev_sample -> "x_t-1" - # 1. get previous step value (=t-1) + # 1. get previous step value from the materialised schedule (= next entry of self.timesteps) t = timesteps - prev_t = t - self.config.num_train_timesteps // self.num_inference_steps + prev_t = self.previous_timestep(t) t = t.view(-1, *([1] * (model_output.ndim - 1))) prev_t = prev_t.view(-1, *([1] * (model_output.ndim - 1))) diff --git a/src/diffusers/schedulers/scheduling_dpm_cogvideox.py b/src/diffusers/schedulers/scheduling_dpm_cogvideox.py index 6dcbaa0b42b4..13dbdef574bd 100644 --- a/src/diffusers/schedulers/scheduling_dpm_cogvideox.py +++ b/src/diffusers/schedulers/scheduling_dpm_cogvideox.py @@ -328,6 +328,43 @@ def set_timesteps( self.timesteps = torch.from_numpy(timesteps).to(device) + # Copied from diffusers.schedulers.scheduling_ddim.DDIMScheduler.previous_timestep + def previous_timestep(self, timestep: int | torch.Tensor) -> int | torch.Tensor: + """ + Return the previous timestep from the schedule produced by [`set_timesteps`]. + + The schedule already encodes `timestep_spacing` (`leading`, `trailing`, or `linspace`). Deriving the previous + value from that list is a no-op for uniform strides and corrects `linspace`, where a fixed + `num_train_timesteps // num_inference_steps` step disagrees with the materialised list. + + Args: + timestep (`int` or `torch.Tensor`): + The current discrete timestep in the diffusion chain (an entry of `self.timesteps`). + + Returns: + `int` or `torch.Tensor`: + The previous timestep. The final schedule entry returns `-1`, which selects `final_alpha_cumprod`. + """ + if self.num_inference_steps is None: + raise ValueError( + "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler" + ) + + if isinstance(timestep, torch.Tensor) and timestep.ndim > 0: + flat = timestep.flatten() + prev = torch.stack([self.previous_timestep(t) for t in flat]).reshape(timestep.shape) + return prev + + index_candidates = (self.timesteps == timestep).nonzero(as_tuple=True)[0] + if len(index_candidates) == 0: + # Not an entry of the inference schedule (direct step() calls outside the loop). + # Keep the historical unit-stride fallback so off-schedule tooling still works. + return timestep - self.config.num_train_timesteps // self.num_inference_steps + index = index_candidates[0] + if index == self.timesteps.shape[0] - 1: + return torch.tensor(-1, device=self.timesteps.device, dtype=self.timesteps.dtype) + return self.timesteps[index + 1] + def get_variables( self, alpha_prod_t: torch.Tensor, @@ -463,8 +500,8 @@ def step( # - pred_sample_direction -> "direction pointing to x_t" # - pred_prev_sample -> "x_t-1" - # 1. get previous step value (=t-1) - prev_timestep = timestep - self.config.num_train_timesteps // self.num_inference_steps + # 1. get previous step value from the materialised schedule (= next entry of self.timesteps) + prev_timestep = self.previous_timestep(timestep) # 2. compute alphas, betas alpha_prod_t = self.alphas_cumprod[timestep] diff --git a/tests/schedulers/test_scheduler_ddim.py b/tests/schedulers/test_scheduler_ddim.py index 13b353a44b08..cdcfd6f42a24 100644 --- a/tests/schedulers/test_scheduler_ddim.py +++ b/tests/schedulers/test_scheduler_ddim.py @@ -73,6 +73,44 @@ def test_timestep_spacing(self): for timestep_spacing in ["trailing", "leading"]: self.check_over_configs(timestep_spacing=timestep_spacing) + def test_prev_timestep_matches_schedule(self): + # step must use the materialised schedule, not a uniform-stride formula. + # The stride formula matches leading and trailing intermediates, but not linspace. + for timestep_spacing in ["leading", "trailing", "linspace"]: + scheduler = DDIMScheduler( + num_train_timesteps=1000, + beta_start=0.0001, + beta_end=0.02, + beta_schedule="linear", + clip_sample=False, + timestep_spacing=timestep_spacing, + ) + scheduler.set_timesteps(10) + timesteps = scheduler.timesteps.tolist() + for i, t in enumerate(timesteps): + expected = -1 if i == len(timesteps) - 1 else timesteps[i + 1] + got = scheduler.previous_timestep(t) + got = int(got) if not isinstance(got, int) else got + if expected < 0: + self.assertLess(got, 0) + else: + self.assertEqual(got, expected) + + # leading / trailing: schedule prev equals the historical stride formula except + # at the final step, where both are negative and both select final_alpha_cumprod. + if timestep_spacing in ["leading", "trailing"]: + stride = scheduler.config.num_train_timesteps // scheduler.num_inference_steps + for i, t in enumerate(timesteps[:-1]): + self.assertEqual(timesteps[i + 1], t - stride) + hard_last = timesteps[-1] - stride + self.assertLess(hard_last, 0) + + # linspace: the historical stride formula disagrees on intermediate steps. + if timestep_spacing == "linspace": + stride = scheduler.config.num_train_timesteps // scheduler.num_inference_steps + mismatches = sum(1 for i, t in enumerate(timesteps[:-1]) if timesteps[i + 1] != t - stride) + self.assertGreater(mismatches, 0) + def test_rescale_betas_zero_snr(self): for rescale_betas_zero_snr in [True, False]: self.check_over_configs(rescale_betas_zero_snr=rescale_betas_zero_snr) From f98c147e6c47856c1051ec99712303512caeb844 Mon Sep 17 00:00:00 2001 From: Mohammad Sadegh Salehi Date: Sun, 2 Aug 2026 21:30:52 +0100 Subject: [PATCH 2/2] Wrap previous_timestep docstrings to doc-builder max line length doc-builder style rewrites these four copied docstrings to the project max length of 119. Keep the copies in sync with make fix-copies. --- src/diffusers/schedulers/scheduling_ddim.py | 4 ++-- src/diffusers/schedulers/scheduling_ddim_cogvideox.py | 4 ++-- src/diffusers/schedulers/scheduling_ddim_parallel.py | 4 ++-- src/diffusers/schedulers/scheduling_dpm_cogvideox.py | 4 ++-- 4 files changed, 8 insertions(+), 8 deletions(-) diff --git a/src/diffusers/schedulers/scheduling_ddim.py b/src/diffusers/schedulers/scheduling_ddim.py index 8d8fbd3c7146..7d21744b1352 100644 --- a/src/diffusers/schedulers/scheduling_ddim.py +++ b/src/diffusers/schedulers/scheduling_ddim.py @@ -386,8 +386,8 @@ def previous_timestep(self, timestep: int | torch.Tensor) -> int | torch.Tensor: Return the previous timestep from the schedule produced by [`set_timesteps`]. The schedule already encodes `timestep_spacing` (`leading`, `trailing`, or `linspace`). Deriving the previous - value from that list is a no-op for uniform strides and corrects `linspace`, where a fixed - `num_train_timesteps // num_inference_steps` step disagrees with the materialised list. + value from that list is a no-op for uniform strides and corrects `linspace`, where a fixed `num_train_timesteps + // num_inference_steps` step disagrees with the materialised list. Args: timestep (`int` or `torch.Tensor`): diff --git a/src/diffusers/schedulers/scheduling_ddim_cogvideox.py b/src/diffusers/schedulers/scheduling_ddim_cogvideox.py index a6d22c898229..83fdbd796bfa 100644 --- a/src/diffusers/schedulers/scheduling_ddim_cogvideox.py +++ b/src/diffusers/schedulers/scheduling_ddim_cogvideox.py @@ -331,8 +331,8 @@ def previous_timestep(self, timestep: int | torch.Tensor) -> int | torch.Tensor: Return the previous timestep from the schedule produced by [`set_timesteps`]. The schedule already encodes `timestep_spacing` (`leading`, `trailing`, or `linspace`). Deriving the previous - value from that list is a no-op for uniform strides and corrects `linspace`, where a fixed - `num_train_timesteps // num_inference_steps` step disagrees with the materialised list. + value from that list is a no-op for uniform strides and corrects `linspace`, where a fixed `num_train_timesteps + // num_inference_steps` step disagrees with the materialised list. Args: timestep (`int` or `torch.Tensor`): diff --git a/src/diffusers/schedulers/scheduling_ddim_parallel.py b/src/diffusers/schedulers/scheduling_ddim_parallel.py index 014613a26196..e8e368da516a 100644 --- a/src/diffusers/schedulers/scheduling_ddim_parallel.py +++ b/src/diffusers/schedulers/scheduling_ddim_parallel.py @@ -391,8 +391,8 @@ def previous_timestep(self, timestep: int | torch.Tensor) -> int | torch.Tensor: Return the previous timestep from the schedule produced by [`set_timesteps`]. The schedule already encodes `timestep_spacing` (`leading`, `trailing`, or `linspace`). Deriving the previous - value from that list is a no-op for uniform strides and corrects `linspace`, where a fixed - `num_train_timesteps // num_inference_steps` step disagrees with the materialised list. + value from that list is a no-op for uniform strides and corrects `linspace`, where a fixed `num_train_timesteps + // num_inference_steps` step disagrees with the materialised list. Args: timestep (`int` or `torch.Tensor`): diff --git a/src/diffusers/schedulers/scheduling_dpm_cogvideox.py b/src/diffusers/schedulers/scheduling_dpm_cogvideox.py index 13dbdef574bd..bb4009b017f2 100644 --- a/src/diffusers/schedulers/scheduling_dpm_cogvideox.py +++ b/src/diffusers/schedulers/scheduling_dpm_cogvideox.py @@ -334,8 +334,8 @@ def previous_timestep(self, timestep: int | torch.Tensor) -> int | torch.Tensor: Return the previous timestep from the schedule produced by [`set_timesteps`]. The schedule already encodes `timestep_spacing` (`leading`, `trailing`, or `linspace`). Deriving the previous - value from that list is a no-op for uniform strides and corrects `linspace`, where a fixed - `num_train_timesteps // num_inference_steps` step disagrees with the materialised list. + value from that list is a no-op for uniform strides and corrects `linspace`, where a fixed `num_train_timesteps + // num_inference_steps` step disagrees with the materialised list. Args: timestep (`int` or `torch.Tensor`):