Repository: huggingface/diffusers

Issue:
Apple MPS and KDPM2AncestralDiscreteScheduler & KDPM2DiscreteScheduler
### Describe the bug

The above two schedulers have issues when running under Apple MPS. It looks as if `float64` to `float32` conversion is not handled correctly for MPS for either of these. See attached screenshots for errors. 
![CleanShot 2022-12-04 at 08 32 41@2x](https://user-images.githubusercontent.com/181110/205475110-a09e93f2-bd89-46ff-b43f-6ebd094942b8.jpg)
![CleanShot 2022-12-04 at 08 36 25@2x](https://user-images.githubusercontent.com/181110/205475115-823ce109-9e0e-47a6-b6fe-70c778218325.jpg)


### Reproduction

Simply use the KDPM2AncestralDiscreteScheduler or the KDPM2DiscreteScheduler with Stable Diffusion v2 base model. This is how I tested.

### Logs

```shell
Logs shown in screenshots attached above.
```


### System Info

- `diffusers` version: 0.10.0.dev0
- Platform: macOS-13.1-arm64-arm-64bit
- Python version: 3.9.13
- PyTorch version (GPU?): 1.14.0.dev20221124 (False)
- Huggingface_hub version: 0.10.1
- Transformers version: 4.24.0
- Using GPU in script?: Using MPS
- Using distributed or parallel set-up in script?: No


Discussion:
Thanks for opening cc @pcuenca here maybe
@pcuenca should we maybe open this one up for the community if you're busy? :-) 
This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/diffusers/blob/main/CONTRIBUTING.md) are likely to be ignored.
still exists in 0.12.1, stalebot.

Resolving patch:
diff --git a/src/diffusers/schedulers/scheduling_k_dpm_2_ancestral_discrete.py b/src/diffusers/schedulers/scheduling_k_dpm_2_ancestral_discrete.py
--- a/src/diffusers/schedulers/scheduling_k_dpm_2_ancestral_discrete.py
+++ b/src/diffusers/schedulers/scheduling_k_dpm_2_ancestral_discrete.py
@@ -161,16 +161,16 @@ def set_timesteps(
         # standard deviation of the initial noise distribution
         self.init_noise_sigma = self.sigmas.max()
 
-        timesteps = torch.from_numpy(timesteps).to(device)
-        timesteps_interpol = self.sigma_to_t(sigmas_interpol).to(device)
-        interleaved_timesteps = torch.stack((timesteps_interpol[:-2, None], timesteps[1:, None]), dim=-1).flatten()
-        timesteps = torch.cat([timesteps[:1], interleaved_timesteps])
-
         if str(device).startswith("mps"):
             # mps does not support float64
-            self.timesteps = timesteps.to(device, dtype=torch.float32)
+            timesteps = torch.from_numpy(timesteps).to(device, dtype=torch.float32)
         else:
-            self.timesteps = timesteps
+            timesteps = torch.from_numpy(timesteps).to(device)
+
+        timesteps_interpol = self.sigma_to_t(sigmas_interpol).to(device)
+        interleaved_timesteps = torch.stack((timesteps_interpol[:-2, None], timesteps[1:, None]), dim=-1).flatten()
+
+        self.timesteps = torch.cat([timesteps[:1], interleaved_timesteps])
 
         self.sample = None
 
diff --git a/src/diffusers/schedulers/scheduling_k_dpm_2_discrete.py b/src/diffusers/schedulers/scheduling_k_dpm_2_discrete.py
--- a/src/diffusers/schedulers/scheduling_k_dpm_2_discrete.py
+++ b/src/diffusers/schedulers/scheduling_k_dpm_2_discrete.py
@@ -149,18 +149,17 @@ def set_timesteps(
         # standard deviation of the initial noise distribution
         self.init_noise_sigma = self.sigmas.max()
 
-        timesteps = torch.from_numpy(timesteps).to(device)
+        if str(device).startswith("mps"):
+            # mps does not support float64
+            timesteps = torch.from_numpy(timesteps).to(device, dtype=torch.float32)
+        else:
+            timesteps = torch.from_numpy(timesteps).to(device)
 
         # interpolate timesteps
         timesteps_interpol = self.sigma_to_t(sigmas_interpol).to(device)
         interleaved_timesteps = torch.stack((timesteps_interpol[1:-1, None], timesteps[1:, None]), dim=-1).flatten()
-        timesteps = torch.cat([timesteps[:1], interleaved_timesteps])
 
-        if str(device).startswith("mps"):
-            # mps does not support float64
-            self.timesteps = timesteps.to(torch.float32)
-        else:
-            self.timesteps = timesteps
+        self.timesteps = torch.cat([timesteps[:1], interleaved_timesteps])
 
         self.sample = None
 


Test patch: