Change timestep device to cpu for xla (#11501)

* Change timestep device to cpu for xla

* Add all pipelines

* ruff format

* Apply style fixes

---------

Co-authored-by: YiYi Xu <[email protected]>
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>

Author: 3 people

src/diffusers/pipelines/allegro/pipeline_allegro.py

@@ -887,7 +887,13 @@ def __call__(
             prompt_embeds = prompt_embeds.unsqueeze(1)  # b l d -> b 1 l d
 
         # 4. Prepare timesteps
-        timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps)
+        if XLA_AVAILABLE:
+            timestep_device = "cpu"
+        else:
+            timestep_device = device
+        timesteps, num_inference_steps = retrieve_timesteps(
+            self.scheduler, num_inference_steps, timestep_device, timesteps
+        )
         self.scheduler.set_timesteps(num_inference_steps, device=device)
 
         # 5. Prepare latents.

src/diffusers/pipelines/animatediff/pipeline_animatediff_video2video.py

@@ -897,16 +897,20 @@ def __call__(
         dtype = self.dtype
 
         # 3. Prepare timesteps
+        if XLA_AVAILABLE:
+            timestep_device = "cpu"
+        else:
+            timestep_device = device
         if not enforce_inference_steps:
             timesteps, num_inference_steps = retrieve_timesteps(
-                self.scheduler, num_inference_steps, device, timesteps, sigmas
+                self.scheduler, num_inference_steps, timestep_device, timesteps, sigmas
             )
             timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, timesteps, strength, device)
             latent_timestep = timesteps[:1].repeat(batch_size * num_videos_per_prompt)
         else:
             denoising_inference_steps = int(num_inference_steps / strength)
             timesteps, denoising_inference_steps = retrieve_timesteps(
-                self.scheduler, denoising_inference_steps, device, timesteps, sigmas
+                self.scheduler, denoising_inference_steps, timestep_device, timesteps, sigmas
             )
             timesteps = timesteps[-num_inference_steps:]
             latent_timestep = timesteps[:1].repeat(batch_size * num_videos_per_prompt)

src/diffusers/pipelines/animatediff/pipeline_animatediff_video2video_controlnet.py

@@ -1100,16 +1100,20 @@ def __call__(
         dtype = self.dtype
 
         # 3. Prepare timesteps
+        if XLA_AVAILABLE:
+            timestep_device = "cpu"
+        else:
+            timestep_device = device
         if not enforce_inference_steps:
             timesteps, num_inference_steps = retrieve_timesteps(
-                self.scheduler, num_inference_steps, device, timesteps, sigmas
+                self.scheduler, num_inference_steps, timestep_device, timesteps, sigmas
             )
             timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, timesteps, strength, device)
             latent_timestep = timesteps[:1].repeat(batch_size * num_videos_per_prompt)
         else:
             denoising_inference_steps = int(num_inference_steps / strength)
             timesteps, denoising_inference_steps = retrieve_timesteps(
-                self.scheduler, denoising_inference_steps, device, timesteps, sigmas
+                self.scheduler, denoising_inference_steps, timestep_device, timesteps, sigmas
             )
             timesteps = timesteps[-num_inference_steps:]
             latent_timestep = timesteps[:1].repeat(batch_size * num_videos_per_prompt)

src/diffusers/pipelines/aura_flow/pipeline_aura_flow.py

@@ -586,7 +586,13 @@ def __call__(
         # 4. Prepare timesteps
 
         # sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps)
-        timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, sigmas=sigmas)
+        if XLA_AVAILABLE:
+            timestep_device = "cpu"
+        else:
+            timestep_device = device
+        timesteps, num_inference_steps = retrieve_timesteps(
+            self.scheduler, num_inference_steps, timestep_device, sigmas=sigmas
+        )
 
         # 5. Prepare latents.
         latent_channels = self.transformer.config.in_channels

src/diffusers/pipelines/cogvideo/pipeline_cogvideox.py

@@ -664,7 +664,13 @@ def __call__(
             prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
 
         # 4. Prepare timesteps
-        timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps)
+        if XLA_AVAILABLE:
+            timestep_device = "cpu"
+        else:
+            timestep_device = device
+        timesteps, num_inference_steps = retrieve_timesteps(
+            self.scheduler, num_inference_steps, timestep_device, timesteps
+        )
         self._num_timesteps = len(timesteps)
 
         # 5. Prepare latents

src/diffusers/pipelines/cogvideo/pipeline_cogvideox_fun_control.py

@@ -717,7 +717,13 @@ def __call__(
             prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
 
         # 4. Prepare timesteps
-        timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps)
+        if XLA_AVAILABLE:
+            timestep_device = "cpu"
+        else:
+            timestep_device = device
+        timesteps, num_inference_steps = retrieve_timesteps(
+            self.scheduler, num_inference_steps, timestep_device, timesteps
+        )
         self._num_timesteps = len(timesteps)
 
         # 5. Prepare latents

src/diffusers/pipelines/cogvideo/pipeline_cogvideox_image2video.py

@@ -762,7 +762,13 @@ def __call__(
             prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
 
         # 4. Prepare timesteps
-        timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps)
+        if XLA_AVAILABLE:
+            timestep_device = "cpu"
+        else:
+            timestep_device = device
+        timesteps, num_inference_steps = retrieve_timesteps(
+            self.scheduler, num_inference_steps, timestep_device, timesteps
+        )
         self._num_timesteps = len(timesteps)
 
         # 5. Prepare latents

src/diffusers/pipelines/cogvideo/pipeline_cogvideox_video2video.py

@@ -737,7 +737,13 @@ def __call__(
             prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
 
         # 4. Prepare timesteps
-        timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps)
+        if XLA_AVAILABLE:
+            timestep_device = "cpu"
+        else:
+            timestep_device = device
+        timesteps, num_inference_steps = retrieve_timesteps(
+            self.scheduler, num_inference_steps, timestep_device, timesteps
+        )
         timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, timesteps, strength, device)
         latent_timestep = timesteps[:1].repeat(batch_size * num_videos_per_prompt)
         self._num_timesteps = len(timesteps)

src/diffusers/pipelines/cogview3/pipeline_cogview3plus.py

@@ -566,7 +566,13 @@ def __call__(
             prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
 
         # 4. Prepare timesteps
-        timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps)
+        if XLA_AVAILABLE:
+            timestep_device = "cpu"
+        else:
+            timestep_device = device
+        timesteps, num_inference_steps = retrieve_timesteps(
+            self.scheduler, num_inference_steps, timestep_device, timesteps
+        )
         self._num_timesteps = len(timesteps)
 
         # 5. Prepare latents.

src/diffusers/pipelines/cogview4/pipeline_cogview4.py

@@ -599,8 +599,12 @@ def __call__(
             self.scheduler.config.get("base_shift", 0.25),
             self.scheduler.config.get("max_shift", 0.75),
         )
+        if XLA_AVAILABLE:
+            timestep_device = "cpu"
+        else:
+            timestep_device = device
         timesteps, num_inference_steps = retrieve_timesteps(
-            self.scheduler, num_inference_steps, device, timesteps, sigmas, mu=mu
+            self.scheduler, num_inference_steps, timestep_device, timesteps, sigmas, mu=mu
         )
         self._num_timesteps = len(timesteps)