use np.int32 in scheduling

Author: hlky

src/diffusers/schedulers/scheduling_consistency_models.py

@@ -207,7 +207,7 @@ def set_timesteps(
                     f" {self.config.num_train_timesteps}."
                 )
 
-            timesteps = np.array(timesteps, dtype=np.int64)
+            timesteps = np.array(timesteps, dtype=np.int32)
             self.custom_timesteps = True
         else:
             if num_inference_steps > self.config.num_train_timesteps:
@@ -220,7 +220,7 @@ def set_timesteps(
             self.num_inference_steps = num_inference_steps
 
             step_ratio = self.config.num_train_timesteps // self.num_inference_steps
-            timesteps = (np.arange(0, num_inference_steps) * step_ratio).round()[::-1].copy().astype(np.int64)
+            timesteps = (np.arange(0, num_inference_steps) * step_ratio).round()[::-1].copy().astype(np.int32)
             self.custom_timesteps = False
 
         # Map timesteps to Karras sigmas directly for multistep sampling

src/diffusers/schedulers/scheduling_ddim.py

@@ -231,7 +231,7 @@ def __init__(
 
         # setable values
         self.num_inference_steps = None
-        self.timesteps = torch.from_numpy(np.arange(0, num_train_timesteps)[::-1].copy().astype(np.int64))
+        self.timesteps = torch.from_numpy(np.arange(0, num_train_timesteps)[::-1].copy().astype(np.int32))
 
     def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None) -> torch.Tensor:
         """
@@ -318,19 +318,19 @@ def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.devic
                 np.linspace(0, self.config.num_train_timesteps - 1, num_inference_steps)
                 .round()[::-1]
                 .copy()
-                .astype(np.int64)
+                .astype(np.int32)
             )
         elif self.config.timestep_spacing == "leading":
             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().astype(np.int64)
+            timesteps = (np.arange(0, num_inference_steps) * step_ratio).round()[::-1].copy().astype(np.int32)
             timesteps += self.config.steps_offset
         elif self.config.timestep_spacing == "trailing":
             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.round(np.arange(self.config.num_train_timesteps, 0, -step_ratio)).astype(np.int64)
+            timesteps = np.round(np.arange(self.config.num_train_timesteps, 0, -step_ratio)).astype(np.int32)
             timesteps -= 1
         else:
             raise ValueError(

src/diffusers/schedulers/scheduling_ddim_cogvideox.py

@@ -228,7 +228,7 @@ def __init__(
 
         # setable values
         self.num_inference_steps = None
-        self.timesteps = torch.from_numpy(np.arange(0, num_train_timesteps)[::-1].copy().astype(np.int64))
+        self.timesteps = torch.from_numpy(np.arange(0, num_train_timesteps)[::-1].copy().astype(np.int32))
 
     def _get_variance(self, timestep, prev_timestep):
         alpha_prod_t = self.alphas_cumprod[timestep]
@@ -281,19 +281,19 @@ def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.devic
                 np.linspace(0, self.config.num_train_timesteps - 1, num_inference_steps)
                 .round()[::-1]
                 .copy()
-                .astype(np.int64)
+                .astype(np.int32)
             )
         elif self.config.timestep_spacing == "leading":
             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().astype(np.int64)
+            timesteps = (np.arange(0, num_inference_steps) * step_ratio).round()[::-1].copy().astype(np.int32)
             timesteps += self.config.steps_offset
         elif self.config.timestep_spacing == "trailing":
             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.round(np.arange(self.config.num_train_timesteps, 0, -step_ratio)).astype(np.int64)
+            timesteps = np.round(np.arange(self.config.num_train_timesteps, 0, -step_ratio)).astype(np.int32)
             timesteps -= 1
         else:
             raise ValueError(

src/diffusers/schedulers/scheduling_ddim_inverse.py

@@ -228,7 +228,7 @@ def __init__(
 
         # setable values
         self.num_inference_steps = None
-        self.timesteps = torch.from_numpy(np.arange(0, num_train_timesteps).copy().astype(np.int64))
+        self.timesteps = torch.from_numpy(np.arange(0, num_train_timesteps).copy().astype(np.int32))
 
     # Copied from diffusers.schedulers.scheduling_ddim.DDIMScheduler.scale_model_input
     def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None) -> torch.Tensor:
@@ -271,13 +271,13 @@ def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.devic
             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().copy().astype(np.int64)
+            timesteps = (np.arange(0, num_inference_steps) * step_ratio).round().copy().astype(np.int32)
             timesteps += self.config.steps_offset
         elif self.config.timestep_spacing == "trailing":
             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.round(np.arange(self.config.num_train_timesteps, 0, -step_ratio)[::-1]).astype(np.int64)
+            timesteps = np.round(np.arange(self.config.num_train_timesteps, 0, -step_ratio)[::-1]).astype(np.int32)
             timesteps -= 1
         else:
             raise ValueError(

src/diffusers/schedulers/scheduling_ddim_parallel.py

@@ -238,7 +238,7 @@ def __init__(
 
         # setable values
         self.num_inference_steps = None
-        self.timesteps = torch.from_numpy(np.arange(0, num_train_timesteps)[::-1].copy().astype(np.int64))
+        self.timesteps = torch.from_numpy(np.arange(0, num_train_timesteps)[::-1].copy().astype(np.int32))
 
     # Copied from diffusers.schedulers.scheduling_ddim.DDIMScheduler.scale_model_input
     def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None) -> torch.Tensor:
@@ -341,19 +341,19 @@ def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.devic
                 np.linspace(0, self.config.num_train_timesteps - 1, num_inference_steps)
                 .round()[::-1]
                 .copy()
-                .astype(np.int64)
+                .astype(np.int32)
             )
         elif self.config.timestep_spacing == "leading":
             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().astype(np.int64)
+            timesteps = (np.arange(0, num_inference_steps) * step_ratio).round()[::-1].copy().astype(np.int32)
             timesteps += self.config.steps_offset
         elif self.config.timestep_spacing == "trailing":
             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.round(np.arange(self.config.num_train_timesteps, 0, -step_ratio)).astype(np.int64)
+            timesteps = np.round(np.arange(self.config.num_train_timesteps, 0, -step_ratio)).astype(np.int32)
             timesteps -= 1
         else:
             raise ValueError(

src/diffusers/schedulers/scheduling_ddpm.py

@@ -283,7 +283,7 @@ def set_timesteps(
                     f" {self.config.num_train_timesteps}."
                 )
 
-            timesteps = np.array(timesteps, dtype=np.int64)
+            timesteps = np.array(timesteps, dtype=np.int32)
             self.custom_timesteps = True
         else:
             if num_inference_steps > self.config.num_train_timesteps:
@@ -302,19 +302,19 @@ def set_timesteps(
                     np.linspace(0, self.config.num_train_timesteps - 1, num_inference_steps)
                     .round()[::-1]
                     .copy()
-                    .astype(np.int64)
+                    .astype(np.int32)
                 )
             elif self.config.timestep_spacing == "leading":
                 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().astype(np.int64)
+                timesteps = (np.arange(0, num_inference_steps) * step_ratio).round()[::-1].copy().astype(np.int32)
                 timesteps += self.config.steps_offset
             elif self.config.timestep_spacing == "trailing":
                 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.round(np.arange(self.config.num_train_timesteps, 0, -step_ratio)).astype(np.int64)
+                timesteps = np.round(np.arange(self.config.num_train_timesteps, 0, -step_ratio)).astype(np.int32)
                 timesteps -= 1
             else:
                 raise ValueError(

src/diffusers/schedulers/scheduling_ddpm_parallel.py

@@ -293,7 +293,7 @@ def set_timesteps(
                     f" {self.config.num_train_timesteps}."
                 )
 
-            timesteps = np.array(timesteps, dtype=np.int64)
+            timesteps = np.array(timesteps, dtype=np.int32)
             self.custom_timesteps = True
         else:
             if num_inference_steps > self.config.num_train_timesteps:
@@ -312,19 +312,19 @@ def set_timesteps(
                     np.linspace(0, self.config.num_train_timesteps - 1, num_inference_steps)
                     .round()[::-1]
                     .copy()
-                    .astype(np.int64)
+                    .astype(np.int32)
                 )
             elif self.config.timestep_spacing == "leading":
                 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().astype(np.int64)
+                timesteps = (np.arange(0, num_inference_steps) * step_ratio).round()[::-1].copy().astype(np.int32)
                 timesteps += self.config.steps_offset
             elif self.config.timestep_spacing == "trailing":
                 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.round(np.arange(self.config.num_train_timesteps, 0, -step_ratio)).astype(np.int64)
+                timesteps = np.round(np.arange(self.config.num_train_timesteps, 0, -step_ratio)).astype(np.int32)
                 timesteps -= 1
             else:
                 raise ValueError(

src/diffusers/schedulers/scheduling_deis_multistep.py

@@ -248,19 +248,19 @@ def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.devic
                 np.linspace(0, self.config.num_train_timesteps - 1, num_inference_steps + 1)
                 .round()[::-1][:-1]
                 .copy()
-                .astype(np.int64)
+                .astype(np.int32)
             )
         elif self.config.timestep_spacing == "leading":
             step_ratio = self.config.num_train_timesteps // (num_inference_steps + 1)
             # 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 + 1) * step_ratio).round()[::-1][:-1].copy().astype(np.int64)
+            timesteps = (np.arange(0, num_inference_steps + 1) * step_ratio).round()[::-1][:-1].copy().astype(np.int32)
             timesteps += self.config.steps_offset
         elif self.config.timestep_spacing == "trailing":
             step_ratio = self.config.num_train_timesteps / 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(self.config.num_train_timesteps, 0, -step_ratio).round().copy().astype(np.int64)
+            timesteps = np.arange(self.config.num_train_timesteps, 0, -step_ratio).round().copy().astype(np.int32)
             timesteps -= 1
         else:
             raise ValueError(

src/diffusers/schedulers/scheduling_dpm_cogvideox.py

@@ -229,7 +229,7 @@ def __init__(
 
         # setable values
         self.num_inference_steps = None
-        self.timesteps = torch.from_numpy(np.arange(0, num_train_timesteps)[::-1].copy().astype(np.int64))
+        self.timesteps = torch.from_numpy(np.arange(0, num_train_timesteps)[::-1].copy().astype(np.int32))
 
     def _get_variance(self, timestep, prev_timestep):
         alpha_prod_t = self.alphas_cumprod[timestep]
@@ -282,19 +282,19 @@ def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.devic
                 np.linspace(0, self.config.num_train_timesteps - 1, num_inference_steps)
                 .round()[::-1]
                 .copy()
-                .astype(np.int64)
+                .astype(np.int32)
             )
         elif self.config.timestep_spacing == "leading":
             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().astype(np.int64)
+            timesteps = (np.arange(0, num_inference_steps) * step_ratio).round()[::-1].copy().astype(np.int32)
             timesteps += self.config.steps_offset
         elif self.config.timestep_spacing == "trailing":
             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.round(np.arange(self.config.num_train_timesteps, 0, -step_ratio)).astype(np.int64)
+            timesteps = np.round(np.arange(self.config.num_train_timesteps, 0, -step_ratio)).astype(np.int32)
             timesteps -= 1
         else:
             raise ValueError(

src/diffusers/schedulers/scheduling_dpmsolver_multistep.py

@@ -359,7 +359,7 @@ def set_timesteps(
             raise ValueError("Cannot set `timesteps` with `config.use_beta_sigmas = True`.")
 
         if timesteps is not None:
-            timesteps = np.array(timesteps).astype(np.int64)
+            timesteps = np.array(timesteps).astype(np.int32)
         else:
             # Clipping the minimum of all lambda(t) for numerical stability.
             # This is critical for cosine (squaredcos_cap_v2) noise schedule.
@@ -372,21 +372,21 @@ def set_timesteps(
                     np.linspace(0, last_timestep - 1, num_inference_steps + 1)
                     .round()[::-1][:-1]
                     .copy()
-                    .astype(np.int64)
+                    .astype(np.int32)
                 )
             elif self.config.timestep_spacing == "leading":
                 step_ratio = last_timestep // (num_inference_steps + 1)
                 # 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 + 1) * step_ratio).round()[::-1][:-1].copy().astype(np.int64)
+                    (np.arange(0, num_inference_steps + 1) * step_ratio).round()[::-1][:-1].copy().astype(np.int32)
                 )
                 timesteps += self.config.steps_offset
             elif self.config.timestep_spacing == "trailing":
                 step_ratio = self.config.num_train_timesteps / 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(last_timestep, 0, -step_ratio).round().copy().astype(np.int64)
+                timesteps = np.arange(last_timestep, 0, -step_ratio).round().copy().astype(np.int32)
                 timesteps -= 1
             else:
                 raise ValueError(

src/diffusers/schedulers/scheduling_dpmsolver_multistep_inverse.py

@@ -260,19 +260,19 @@ def set_timesteps(self, num_inference_steps: int = None, device: Union[str, torc
         # "linspace", "leading", "trailing" corresponds to annotation of Table 2. of https://arxiv.org/abs/2305.08891
         if self.config.timestep_spacing == "linspace":
             timesteps = (
-                np.linspace(0, self.noisiest_timestep, num_inference_steps + 1).round()[:-1].copy().astype(np.int64)
+                np.linspace(0, self.noisiest_timestep, num_inference_steps + 1).round()[:-1].copy().astype(np.int32)
             )
         elif self.config.timestep_spacing == "leading":
             step_ratio = (self.noisiest_timestep + 1) // (num_inference_steps + 1)
             # 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 + 1) * step_ratio).round()[:-1].copy().astype(np.int64)
+            timesteps = (np.arange(0, num_inference_steps + 1) * step_ratio).round()[:-1].copy().astype(np.int32)
             timesteps += self.config.steps_offset
         elif self.config.timestep_spacing == "trailing":
             step_ratio = self.config.num_train_timesteps / 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(self.noisiest_timestep + 1, 0, -step_ratio).round()[::-1].copy().astype(np.int64)
+            timesteps = np.arange(self.noisiest_timestep + 1, 0, -step_ratio).round()[::-1].copy().astype(np.int32)
             timesteps -= 1
         else:
             raise ValueError(
@@ -286,7 +286,7 @@ def set_timesteps(self, num_inference_steps: int = None, device: Union[str, torc
         if self.config.use_karras_sigmas:
             sigmas = self._convert_to_karras(in_sigmas=sigmas, num_inference_steps=num_inference_steps)
             timesteps = np.array([self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas]).round()
-            timesteps = timesteps.copy().astype(np.int64)
+            timesteps = timesteps.copy().astype(np.int32)
             sigmas = np.concatenate([sigmas, sigmas[-1:]]).astype(np.float32)
         elif self.config.use_exponential_sigmas:
             sigmas = self._convert_to_exponential(in_sigmas=sigmas, num_inference_steps=num_inference_steps)