MPS crash when using LMSDiscreteScheduler

[FahimF]

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

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

[patrickvonplaten]

Also cc @anton-l

[FahimF]

@pcuenca I've pulled the latest code from the main branch, verified at least one of the changes from the pull request which was merged is there in my locally installed package and then ran the above code (from the bug report) and it still crashes. So there might be something which was missed?

Here's the new crash:

Traceback (most recent call last):
 File "/Users/fahim/Code/Python/sd-gui/test2.py", line 35, 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 257, in add_noise
 timesteps = 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.

Just as a note, do note that the line number for the crash has now changed because of the new code that was added.

Edit: It just occurred to me that I should probably look at my sample code and see where the actual error is coming from now. In a rush at the moment and so will leave this here, but will do a more detailed check tomorrow and get back to you.

[pcuenca]

You are right, it still fails. Sorry, I'm not sure how that happened.

[pcuenca]

Should be fixed now.

[FahimF]

Thanks @pcuenca 🙂 I can confirm that it works correctly now.