allow tensors in several schedulers step() call (huggingface#8905)

Author: catwell

src/diffusers/schedulers/scheduling_deis_multistep.py

@@ -674,7 +674,7 @@ def _init_step_index(self, timestep):
     def step(
         self,
         model_output: torch.Tensor,
-        timestep: int,
+        timestep: Union[int, torch.Tensor],
         sample: torch.Tensor,
         return_dict: bool = True,
     ) -> Union[SchedulerOutput, Tuple]:
@@ -685,7 +685,7 @@ def step(
         Args:
             model_output (`torch.Tensor`):
                 The direct output from learned diffusion model.
-            timestep (`float`):
+            timestep (`int`):
                 The current discrete timestep in the diffusion chain.
             sample (`torch.Tensor`):
                 A current instance of a sample created by the diffusion process.

src/diffusers/schedulers/scheduling_dpmsolver_multistep.py

@@ -920,7 +920,7 @@ def _init_step_index(self, timestep):
     def step(
         self,
         model_output: torch.Tensor,
-        timestep: int,
+        timestep: Union[int, torch.Tensor],
         sample: torch.Tensor,
         generator=None,
         variance_noise: Optional[torch.Tensor] = None,

src/diffusers/schedulers/scheduling_dpmsolver_multistep_inverse.py

@@ -787,7 +787,7 @@ def _init_step_index(self, timestep):
     def step(
         self,
         model_output: torch.Tensor,
-        timestep: int,
+        timestep: Union[int, torch.Tensor],
         sample: torch.Tensor,
         generator=None,
         variance_noise: Optional[torch.Tensor] = None,

src/diffusers/schedulers/scheduling_dpmsolver_singlestep.py

@@ -927,7 +927,7 @@ def _init_step_index(self, timestep):
     def step(
         self,
         model_output: torch.Tensor,
-        timestep: int,
+        timestep: Union[int, torch.Tensor],
         sample: torch.Tensor,
         generator=None,
         return_dict: bool = True,

src/diffusers/schedulers/scheduling_edm_dpmsolver_multistep.py

@@ -594,7 +594,7 @@ def _init_step_index(self, timestep):
     def step(
         self,
         model_output: torch.Tensor,
-        timestep: int,
+        timestep: Union[int, torch.Tensor],
         sample: torch.Tensor,
         generator=None,
         return_dict: bool = True,

src/diffusers/schedulers/scheduling_ipndm.py

@@ -138,7 +138,7 @@ def _init_step_index(self, timestep):
     def step(
         self,
         model_output: torch.Tensor,
-        timestep: int,
+        timestep: Union[int, torch.Tensor],
         sample: torch.Tensor,
         return_dict: bool = True,
     ) -> Union[SchedulerOutput, Tuple]:

src/diffusers/schedulers/scheduling_unipc_multistep.py

@@ -822,7 +822,7 @@ def _init_step_index(self, timestep):
     def step(
         self,
         model_output: torch.Tensor,
-        timestep: int,
+        timestep: Union[int, torch.Tensor],
         sample: torch.Tensor,
         return_dict: bool = True,
     ) -> Union[SchedulerOutput, Tuple]: