From a3efa433eac5feba842350c38a1db29244963fb5 Mon Sep 17 00:00:00 2001 From: Hamish Friedlander Date: Tue, 18 Oct 2022 23:06:46 +1300 Subject: [PATCH] Fix DDIM on Windows not using int64 for timesteps (#819) --- src/diffusers/schedulers/scheduling_ddim.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/diffusers/schedulers/scheduling_ddim.py b/src/diffusers/schedulers/scheduling_ddim.py index 0f1a4022..612c24f8 100644 --- a/src/diffusers/schedulers/scheduling_ddim.py +++ b/src/diffusers/schedulers/scheduling_ddim.py @@ -149,7 +149,7 @@ class DDIMScheduler(SchedulerMixin, ConfigMixin): # setable values self.num_inference_steps = None - self.timesteps = torch.from_numpy(np.arange(0, num_train_timesteps)[::-1].copy()) + self.timesteps = torch.from_numpy(np.arange(0, num_train_timesteps)[::-1].copy().astype(np.int64)) def scale_model_input(self, sample: torch.FloatTensor, timestep: Optional[int] = None) -> torch.FloatTensor: """ @@ -192,7 +192,7 @@ class DDIMScheduler(SchedulerMixin, ConfigMixin): 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 - timesteps = (np.arange(0, num_inference_steps) * step_ratio).round()[::-1].copy() + timesteps = (np.arange(0, num_inference_steps) * step_ratio).round()[::-1].copy().astype(np.int64) self.timesteps = torch.from_numpy(timesteps).to(device) self.timesteps += offset