MPS crash when using LMSDiscreteScheduler #940
Description
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
