Apple MPS error in unet_2d_condition.py

[FahimF]

Describe the bug

When you use an LMSDiscreteScheduler on an Apple Silicon machine, you'll get the following error:
Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64. Please use float32 instead

The offending line is 134 in unet_2d_condition.py.

The current code is:
timesteps = timesteps[None].to(sample.device)

Changing that to the following stops the crash:
timesteps = timesteps[None].long().to(sample.device)

However, I believe you'd really want to do a check to see if the current device is MPS and only do the format conversion if you are on MPS?

Reproduction

When you use an LMSDiscreteScheduler on an Apple Silicon machine you should see the crash.

Logs

No response

System Info

The current main branch from the repo since that appears to be different from the current release version (0.2.4?)

[anton-l]

Hi @FahimF! There's ongoing work to support mps for schedulers, feel free to check in on the progress at #355
cc @pcuenca

[pcuenca]

This will fixed by 12f6670 when the branch is merged. Thanks again for your help @FahimF!

[FahimF]

@pcuenca Thank you for the fixes! Much obliged! Will wait for the merge eagerly 😄

[FahimF]

Unfortunately, with the 0.3.0 release installed, this issue crops up on line 95 in scheduling_utils.py. Sorry 😢

Update: tagging @pcuenca since the ticket is closed and not sure if anybody gets notified.

[pcuenca]

Hi @FahimF! Works for me. Would you mind sharing a code snippet so I can try to reproduce? Also, some information about your setup could be useful. Thanks a lot!

[FahimF]

@pcuenca Thank you for taking a look. Let me try to remove the fix I put in the code and come up with something simple to demonstrate the issue. I know at least one of those was while using the StableDiffusionImg2ImgPipeline but I don't know if that was requisite ... I was in a rush and so simply fixed the code and moved on. Should have kept better records. Sorry.

Here's what I have in my console from that particular crash:

Traceback (most recent call last):
  File "/Users/fahim/miniforge3/envs/ml/lib/python3.8/tkinter/__init__.py", line 1892, in __call__
    return self.func(*args)
  File "gui.py", line 191, in generate_images
    result = pipe(prompt=g_prompt, init_image=image, strength=strength,
  File "/Users/fahim/miniforge3/envs/ml/lib/python3.8/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File "/Users/fahim/miniforge3/envs/ml/lib/python3.8/site-packages/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py", line 205, in __call__
    init_latents = self.scheduler.add_noise(init_latents, noise, timesteps).to(self.device)
  File "/Users/fahim/miniforge3/envs/ml/lib/python3.8/site-packages/diffusers/schedulers/scheduling_lms_discrete.py", line 189, in add_noise
    sigmas = self.match_shape(self.sigmas[ts], noise)
  File "/Users/fahim/miniforge3/envs/ml/lib/python3.8/site-packages/diffusers/schedulers/scheduling_utils.py", line 95, in match_shape
    values = values.to(broadcast_array.device)
TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64. Please use float32 instead.

I see from the error above that I was also using LMSDiscreteScheduler. Perhaps that would be enough to replicate the issue at your end? If not, please let me know and I'll come up with a simple bit of code since my current code has all sorts of other stuff like a tkinter GUI 😄

As far as set up goes, latest pytorch nightly (installed today) and diffusers 0.3.0 (installed today) and on a 2021 MBP. If you need any additional info (I don't know what would help and what won't) please let me know and I'll provide.

Update:
As far as the code goes, I believe the most relevant bits would be the two following lines:

sched = LMSDiscreteScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear")
pipe = StableDiffusionImg2ImgPipeline.from_pretrained("stable-diffusion-v1-4", scheduler=sched).to(device)

[pcuenca]

Thanks a lot, @FahimF, I'll test the image to image pipeline and report back.

[FahimF]

Just updating that this issue might have been due to something in the PyTorch nightlies. I could not generate valid images via img2img either and then updated PyTorch nightly (and did a clean install) and that issue went away. So I tested for this one and this one's gone too ...