Fix scheduler inference steps error with power of 3 (huggingface#466)

* initial attempt at solving

* fix pndm power of 3 inference_step

* add power of 3 test

* fix index in pndm test, remove ddim test

* add comments, change to round()

Author: Nathan Lambert

schedulers/scheduling_ddim.py

@@ -145,9 +145,10 @@ def set_timesteps(self, num_inference_steps: int, offset: int = 0):
                 optional value to shift timestep values up by. A value of 1 is used in stable diffusion for inference.
         """
         self.num_inference_steps = num_inference_steps
-        self.timesteps = np.arange(
-            0, self.config.num_train_timesteps, self.config.num_train_timesteps // self.num_inference_steps
-        )[::-1].copy()
+        step_ratio = self.config.num_train_timesteps // self.num_inference_steps
+        # creates integer timesteps by multiplying by ratio
+        # casting to int to avoid issues when num_inference_step is power of 3
+        self.timesteps = (np.arange(0, num_inference_steps) * step_ratio).round()[::-1].copy()
         self.timesteps += offset
         self.set_format(tensor_format=self.tensor_format)
 

schedulers/scheduling_pndm.py

@@ -143,9 +143,10 @@ def set_timesteps(self, num_inference_steps: int, offset: int = 0) -> torch.Floa
                 optional value to shift timestep values up by. A value of 1 is used in stable diffusion for inference.
         """
         self.num_inference_steps = num_inference_steps
-        self._timesteps = list(
-            range(0, self.config.num_train_timesteps, self.config.num_train_timesteps // num_inference_steps)
-        )
+        step_ratio = self.config.num_train_timesteps // self.num_inference_steps
+        # creates integer timesteps by multiplying by ratio
+        # casting to int to avoid issues when num_inference_step is power of 3
+        self._timesteps = (np.arange(0, num_inference_steps) * step_ratio).round().tolist()
         self._offset = offset
         self._timesteps = np.array([t + self._offset for t in self._timesteps])