Derive DDIM prev_timestep from the inference schedule

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]

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]

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)))

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]

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)