self.scheduler.set_timesteps now uses device arg for schedulers that accept it

Author: NouamaneTazi

examples/community/clip_guided_stable_diffusion.py

@@ -259,7 +259,7 @@ def __call__(
         if accepts_offset:
             extra_set_kwargs["offset"] = 1
 
-        self.scheduler.set_timesteps(num_inference_steps, **extra_set_kwargs)
+        self.scheduler.set_timesteps(num_inference_steps, device=self.device, **extra_set_kwargs)
 
         # if we use LMSDiscreteScheduler, let's make sure latents are multiplied by sigmas
         if isinstance(self.scheduler, LMSDiscreteScheduler):

src/diffusers/pipelines/ddim/pipeline_ddim.py

@@ -80,7 +80,7 @@ def __call__(
         image = image.to(self.device)
 
         # set step values
-        self.scheduler.set_timesteps(num_inference_steps)
+        self.scheduler.set_timesteps(num_inference_steps, device=self.device)
 
         for t in self.progress_bar(self.scheduler.timesteps):
             # 1. predict noise model_output

src/diffusers/pipelines/ddpm/pipeline_ddpm.py

@@ -73,7 +73,7 @@ def __call__(
         image = image.to(self.device)
 
         # set step values
-        self.scheduler.set_timesteps(1000)
+        self.scheduler.set_timesteps(1000, device=self.device)
 
         for t in self.progress_bar(self.scheduler.timesteps):
             # 1. predict noise model_output

src/diffusers/pipelines/latent_diffusion/pipeline_latent_diffusion.py

@@ -118,7 +118,7 @@ def __call__(
         )
         latents = latents.to(self.device)
 
-        self.scheduler.set_timesteps(num_inference_steps)
+        self.scheduler.set_timesteps(num_inference_steps, device=self.device)
 
         # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
         accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())

src/diffusers/pipelines/latent_diffusion_uncond/pipeline_latent_diffusion_uncond.py

@@ -64,7 +64,7 @@ def __call__(
         )
         latents = latents.to(self.device)
 
-        self.scheduler.set_timesteps(num_inference_steps)
+        self.scheduler.set_timesteps(num_inference_steps, device=self.device)
 
         # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
         accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())

src/diffusers/pipelines/pndm/pipeline_pndm.py

@@ -14,6 +14,7 @@
 # limitations under the License.
 
 
+from os import device_encoding
 from typing import Optional, Tuple, Union
 
 import torch
@@ -80,7 +81,7 @@ def __call__(
         )
         image = image.to(self.device)
 
-        self.scheduler.set_timesteps(num_inference_steps)
+        self.scheduler.set_timesteps(num_inference_steps, device=self.device)
         for t in self.progress_bar(self.scheduler.timesteps):
             model_output = self.unet(image, t).sample
 

src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py

@@ -233,9 +233,7 @@ def __call__(
             latents = latents.to(latents_device)
 
         # set timesteps
-        self.scheduler.set_timesteps(num_inference_steps)
-        if isinstance(self.scheduler.timesteps, torch.Tensor):
-            self.scheduler.timesteps = self.scheduler.timesteps.to(self.device)
+        self.scheduler.set_timesteps(num_inference_steps, device=self.device)
 
         # if we use LMSDiscreteScheduler, let's make sure latents are multiplied by sigmas
         if isinstance(self.scheduler, LMSDiscreteScheduler):

src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py

@@ -189,7 +189,7 @@ def __call__(
             raise ValueError(f"The value of strength should in [0.0, 1.0] but is {strength}")
 
         # set timesteps
-        self.scheduler.set_timesteps(num_inference_steps)
+        self.scheduler.set_timesteps(num_inference_steps, device=self.device)
 
         if isinstance(init_image, PIL.Image.Image):
             init_image = preprocess(init_image)

src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py

@@ -209,7 +209,7 @@ def __call__(
             raise ValueError(f"The value of strength should in [0.0, 1.0] but is {strength}")
 
         # set timesteps
-        self.scheduler.set_timesteps(num_inference_steps)
+        self.scheduler.set_timesteps(num_inference_steps, device=self.device)
 
         # preprocess image
         if not isinstance(init_image, torch.FloatTensor):

src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_onnx.py

@@ -111,7 +111,7 @@ def __call__(
             raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {latents_shape}")
 
         # set timesteps
-        self.scheduler.set_timesteps(num_inference_steps)
+        self.scheduler.set_timesteps(num_inference_steps, device=self.device)
 
         # if we use LMSDiscreteScheduler, let's make sure latents are multiplied by sigmas
         if isinstance(self.scheduler, LMSDiscreteScheduler):