Refactor timesteps calculation in ddim.py

ytl0623 · web-flow · commit 616f5f2349c8 · 2025-11-05T18:20:08.000+08:00
Signed-off-by: ytl0623 &lt;david89062388@gmail.com&gt;
diff --git a/monai/networks/schedulers/ddim.py b/monai/networks/schedulers/ddim.py
@@ -125,9 +125,7 @@ def set_timesteps(self, num_inference_steps: int, device: str | torch.device | N
                 f" the max train timestep."
             )
 
-        # creates integer timesteps by multiplying by ratio
-        # casting to int to avoid issues when num_inference_step is power of 3
-        timesteps = np.linspace(num_train_timesteps - 1 - steps_offset, 0, num_inference_steps).round().astype(np.int64)
+        timesteps = np.linspace((self.num_train_timesteps - 1) - self.steps_offset, 0, num_inference_steps).round().astype(np.int64)
         self.timesteps = torch.from_numpy(timesteps).to(device)
         self.timesteps += self.steps_offset
 

Original file line number	Diff line number	Diff line change
`@@ -125,9 +125,7 @@ def set_timesteps(self, num_inference_steps: int, device: str \| torch.device \| N`
`125`	`125`	`f" the max train timestep."`
`126`	`126`	`)`
`127`	`127`
`128`		`- # creates integer timesteps by multiplying by ratio`
`129`		`- # casting to int to avoid issues when num_inference_step is power of 3`
`130`		`- timesteps = np.linspace(num_train_timesteps - 1 - steps_offset, 0, num_inference_steps).round().astype(np.int64)`
	`128`	`+ timesteps = np.linspace((self.num_train_timesteps - 1) - self.steps_offset, 0, num_inference_steps).round().astype(np.int64)`
`131`	`129`	`self.timesteps = torch.from_numpy(timesteps).to(device)`
`132`	`130`	`self.timesteps += self.steps_offset`
`133`	`131`