Repository: huggingface/diffusers

Issue:
MPS crash when using LMSDiscreteScheduler
### Describe the bug

If you run the following code using `LMSDiscreteScheduler` a crash occurs under Apple Silicon/MPS:

```
import torch
from diffusers import AutoencoderKL, UNet2DConditionModel, LMSDiscreteScheduler
from PIL import Image
from torchvision import transforms as tfms

# Set device
torch_device = "cuda" if torch.cuda.is_available() else "mps" if torch.has_mps else "cpu"

# Load the autoencoder model which will be used to decode the latents into image space.
vae = AutoencoderKL.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="vae").to(torch_device)

# The UNet model for generating the latents.
unet = UNet2DConditionModel.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="unet").to(torch_device)

# The noise scheduler
scheduler = LMSDiscreteScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000)

def pil_to_latent(input_im):
    # Single image -> single latent in a batch (so size 1, 4, 64, 64)
    with torch.no_grad():
        latent = vae.encode(tfms.ToTensor()(input_im).unsqueeze(0).to(torch_device)*2-1) # Note scaling
    return 0.18215 * latent.latent_dist.sample()

# Load the image with PIL
input_image = Image.open('macaw.jpg').resize((512, 512))

# Encode to the latent space
encoded = pil_to_latent(input_image)

# Setting the number of sampling steps:
scheduler.set_timesteps(15)

noise = torch.randn_like(encoded)
sampling_step = 10
encoded_and_noised = scheduler.add_noise(encoded, noise, timesteps=torch.tensor([scheduler.timesteps[sampling_step]]))
```

The crash is as follows:
```
Traceback (most recent call last):
  File "/Users/fahim/Code/Python/fastai/diffusion-nbs/t.py", line 36, in <module>
    encoded_and_noised = scheduler.add_noise(encoded, noise, timesteps=torch.tensor([scheduler.timesteps[sampling_step]]))
  File "/Users/fahim/miniconda3/envs/ml/lib/python3.9/site-packages/diffusers/schedulers/scheduling_lms_discrete.py", line 255, in add_noise
    self.timesteps = self.timesteps.to(original_samples.device)
TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64. Please use float32 instead.
```

### Reproduction

Code is provided above

### Logs

```shell
Traceback (most recent call last):
  File "/Users/fahim/Code/Python/fastai/diffusion-nbs/t.py", line 36, in <module>
    encoded_and_noised = scheduler.add_noise(encoded, noise, timesteps=torch.tensor([scheduler.timesteps[sampling_step]]))
  File "/Users/fahim/miniconda3/envs/ml/lib/python3.9/site-packages/diffusers/schedulers/scheduling_lms_discrete.py", line 255, in add_noise
    self.timesteps = self.timesteps.to(original_samples.device)
TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64. Please use float32 instead.
```


### System Info

Python 3.9.13
diffusers                         0.6.0
torch                             1.14.0.dev20221021
macOS 12.6
Apple M1 Max 32GB


Discussion:

Resolving patch:
diff --git a/src/diffusers/schedulers/scheduling_lms_discrete.py b/src/diffusers/schedulers/scheduling_lms_discrete.py
--- a/src/diffusers/schedulers/scheduling_lms_discrete.py
+++ b/src/diffusers/schedulers/scheduling_lms_discrete.py
@@ -252,7 +252,8 @@ def add_noise(
     ) -> torch.FloatTensor:
         # Make sure sigmas and timesteps have the same device and dtype as original_samples
         self.sigmas = self.sigmas.to(device=original_samples.device, dtype=original_samples.dtype)
-        self.timesteps = self.timesteps.to(original_samples.device)
+        dtype = torch.float32 if original_samples.device.type == "mps" else timesteps.dtype
+        self.timesteps = self.timesteps.to(original_samples.device, dtype=dtype)
         timesteps = timesteps.to(original_samples.device)
 
         schedule_timesteps = self.timesteps


Test patch:
diff --git a/tests/test_scheduler.py b/tests/test_scheduler.py
--- a/tests/test_scheduler.py
+++ b/tests/test_scheduler.py
@@ -27,6 +27,7 @@
     PNDMScheduler,
     ScoreSdeVeScheduler,
 )
+from diffusers.utils import torch_device
 
 
 torch.backends.cuda.matmul.allow_tf32 = False
@@ -258,6 +259,23 @@ def test_scheduler_public_api(self):
             scaled_sample = scheduler.scale_model_input(sample, 0.0)
             self.assertEqual(sample.shape, scaled_sample.shape)
 
+    def test_add_noise_device(self):
+        for scheduler_class in self.scheduler_classes:
+            if scheduler_class == IPNDMScheduler:
+                # Skip until #990 is addressed
+                continue
+            scheduler_config = self.get_scheduler_config()
+            scheduler = scheduler_class(**scheduler_config)
+
+            sample = self.dummy_sample.to(torch_device)
+            scaled_sample = scheduler.scale_model_input(sample, 0.0)
+            self.assertEqual(sample.shape, scaled_sample.shape)
+
+            noise = torch.randn_like(scaled_sample).to(torch_device)
+            t = torch.tensor([10]).to(torch_device)
+            noised = scheduler.add_noise(scaled_sample, noise, t)
+            self.assertEqual(noised.shape, scaled_sample.shape)
+
 
 class DDPMSchedulerTest(SchedulerCommonTest):
     scheduler_classes = (DDPMScheduler,)