Update inference.py

inference.py

@@ -21,43 +21,6 @@ else:
     XLA_AVAILABLE = False
 
 
-def retrieve_timesteps(
-        scheduler,
-        num_inference_steps: Optional[int] = None,
-        device: Optional[Union[str, torch.device]] = None,
-        timesteps: Optional[List[int]] = None,
-        sigmas: Optional[List[float]] = None,
-        **kwargs,
-):
-    if timesteps is not None and sigmas is not None:
-        raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
-    if timesteps is not None:
-        accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
-        if not accepts_timesteps:
-            raise ValueError(
-                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
-                f" timestep schedules. Please check whether you are using the correct scheduler."
-            )
-        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
-        timesteps = scheduler.timesteps
-        num_inference_steps = len(timesteps)
-    elif sigmas is not None:
-        accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
-        if not accept_sigmas:
-            raise ValueError(
-                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
-                f" sigmas schedules. Please check whether you are using the correct scheduler."
-            )
-        scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
-        timesteps = scheduler.timesteps
-        num_inference_steps = len(timesteps)
-    else:
-        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
-        timesteps = scheduler.timesteps
-
-    return timesteps, num_inference_steps
-
-
 @torch.no_grad()
 def run(
         self,
@@ -68,6 +31,7 @@ def run(
         width: Optional[int] = None,
         num_inference_steps: int = 28,
         sigmas: Optional[List[float]] = None,
+        timesteps: Optional[List[float]] = None,
         scales: List[float] = None,
         guidance_scale: float = 7.0,
         negative_prompt: Optional[Union[str, List[str]]] = None,
@@ -196,13 +160,6 @@ def run(
         scheduler_kwargs["mu"] = mu
     elif mu is not None:
         scheduler_kwargs["mu"] = mu
-    timesteps, num_inference_steps = retrieve_timesteps(
-        self.scheduler,
-        num_inference_steps,
-        device,
-        sigmas=sigmas,
-        **scheduler_kwargs,
-    )
     num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
     self._num_timesteps = len(timesteps)
 
@@ -269,8 +226,8 @@ def run(
 
             # compute the previous noisy sample x_t -> x_t-1
             latents_dtype = latents.dtype
-            sigma = self.scheduler.sigmas[i]
-            sigma_next = self.scheduler.sigmas[i + 1]
+            sigma = sigmas[i]
+            sigma_next = sigmas[i + 1]
             x0_pred = (latents - sigma * noise_pred)
             try:
                 x0_pred = torch.nn.functional.interpolate(x0_pred, size=scales[i + 1])