timesteps in set_timesteps of DDPM and DDIM schedulers doesn't reach end point T (e.g. T= 1000)

[AbrightWay]

I am reading the source code implementation of DDPM and DDIM schedulers, and I found this line defining timesteps for the schedulers timesteps = (np.arange(0, num_inference_steps) * step_ratio).round()[::-1].astype(np.int64)

np.arange(0, 100) doesn't return an array from 0 to 1000, it actually returns an array from 0 to 999. Therefore, if we have, say, step_ratio = 10, and num_inference_steps=100 for a model trained on DDPM of num_train_timesteps=1000, then the final self.timesteps would be a 1D array from 990 to 0, NOT 1000 to 0.

I think the right one is: timesteps = (np.arange(0, num_inference_steps + 1e-3) * step_ratio).round()[::-1].astype(np.int64)
Here's the links for the schedulers:

1. DDPM:

   MONAI/monai/networks/schedulers/ddpm.py

   Line 128 in 69f3dd2

   timesteps = (np.arange(0, num_inference_steps) * step_ratio).round()[::-1].astype(np.int64)

2. DDIM:

   MONAI/monai/networks/schedulers/ddim.py

   Line 130 in 69f3dd2

   timesteps = (np.arange(0, num_inference_steps) * step_ratio).round()[::-1].copy().astype(np.int64)

[ytl0623]

Hi @AbrightWay , I’ve submitted a fix for this issue in PR #8623.
Please take a look when you have a moment.

Thanks in advance!!

[AbrightWay]

Hi @ytl0623 , I have reviewed the source code of diffusers, and there's actually 3 modes for setting the timesteps. The timesteps I mentioned in the error is "trailing" and what I'm trying to achieve (or what you have pulled the request is "linspace").

You can take a look at diffusers implementation
For DDIM (lines 316- 338): https://github.com/huggingface/diffusers/blob/5a47442f9221eba3abe3e27c7af8304241ac58ca/src/diffusers/schedulers/scheduling_ddim.py#L316
For DDPM (lines 306 - 321): https://github.com/huggingface/diffusers/blob/5a47442f9221eba3abe3e27c7af8304241ac58ca/src/diffusers/schedulers/scheduling_ddpm.py#L306

Therefore, a complete fix should be adding an additional argument of how to choose timestep spacing for the users as implemented in diffusers.

Best,
Trinh D. Pham