StableDiffusionPipeline producing unexpected output with MPS device using diffusers==0.4.0

[jncasey]

Describe the bug

I tried testing the potential speed updates of diffusers 0.4.0 on my M1 mac using an existing StableDiffusionPipeline-based script, and I found that a large image that would take ~3 min to generate in diffusers 0.3.0 was estimated to take more than 10x as long.

Since my existing script had a lot going on (e.g. large resolutions, attention slicing), I tried to diagnose the problem with a minimal script (see below), running in two identical environments, with the only difference being the diffusers version.

In diffusers 0.3.0, it takes ~35 seconds to generate a reasonable result like this:

In diffusers 0.4.0, it takes ~50 seconds (which is slower than 0.3.0, but better that the 10x performance hit I was getting before), but each attempt (with varying seeds) triggered the NSFW filter. Disabling the filter, the results appear to be just noise:

I'm not sure if the initial 10x performance hit I initially observed in my original script would be fixed by fixing this bug, but this certainly seems to be at least part of it.

Reproduction

import torch
from diffusers import StableDiffusionPipeline

pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4")
pipe.to("mps")

result = pipe("dogs playing poker", generator=torch.manual_seed(1))

result.images[0].save("test.png")

Logs

Under 0.4.0 there's also this warning:

/opt/homebrew/Caskroom/miniforge/base/envs/sd/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py:222: UserWarning: The operator 'aten::repeat_interleave.self_int' is not currently supported on the MPS backend and will fall back to run on the CPU. This may have performance implications. (Triggered internally at /Users/runner/work/pytorch/pytorch/pytorch/aten/src/ATen/mps/MPSFallback.mm:11.)
  text_embeddings = text_embeddings.repeat_interleave(num_images_per_prompt, dim=0)

System Info

• diffusers version: 0.4.0
• Platform: macOS-12.6-arm64-arm-64bit
• Python version: 3.10.6
• PyTorch version (GPU?): 1.13.0.dev20220911 (False)
• Huggingface_hub version: 0.10.0
• Transformers version: 4.21.3
• Using GPU in script?: MPS
• Using distributed or parallel set-up in script?: no

[pcuenca]

Hi @jncasey, thanks for reporting this!

From the details you describe it looks like one of the last-minute optimizations we did for the release is not compatible with mps, so it's falling back to the CPU and that's why it's slower. It's probably the reason why it fails to generate images, I'll test and work on a fix today.

I do apologize for the trouble, I tested on my Mac a few hours before the release, but not after all the changes were merged together. Thanks again for letting us know!

[pcuenca]

(For reference, I believe this is what's affecting mps: #741 (comment))

[pcuenca]

Hi @jncasey! #766 will replace the slow operation (once it's approved and released), but I had to upgrade my PyTorch nightly version to the latest one (1.14.0.dev20221007) to fix the broken images issue. Would it be possible for you to do the same, following the instructions here?

Thank you!

[jncasey]

Okay, so updating to the latest daily pytorch has solved the noise issue – the image generated by my test script with 0.4.0 is now reasonable and very similar (thought not identical) to the one generated by the old pytorch and 0.3.0 using the same seed. The speed under 0.4.0 remains roughly the same.

Interestingly, with this pytorch, my original script now seems to be slower under 0.3.0 than even it was under 0.4.0
0.3.0, pytorch 1.13.0.dev20220911: ~3 min
0.4.0, pytorch 1.13.0.dev20220911: ~30 min
0.3.0, pytorch 1.14.0.dev20221007: ~40 min
0.4.0, pytorch 1.14.0.dev20221007: ~24 min

I have to run to work now, so I can't dig any deeper, but I'll keep poking around this weekend.

[tcapelle]

I was expecting faster performance with the new optimization, but the exported CoreML version of the Unet (after tracing) is roughly the same on my machine.

More info here http://wandb.me/mac_diffusion

[pcuenca]

Okay, so updating to the latest daily pytorch has solved the noise issue – the image generated by my test script with 0.4.0 is now reasonable and very similar (thought not identical) to the one generated by the old pytorch and 0.3.0 using the same seed. The speed under 0.4.0 remains roughly the same.

Interestingly, with this pytorch, my original script now seems to be slower under 0.3.0 than even it was under 0.4.0 0.3.0, pytorch 1.13.0.dev20220911: ~3 min 0.4.0, pytorch 1.13.0.dev20220911: ~30 min 0.3.0, pytorch 1.14.0.dev20221007: ~40 min 0.4.0, pytorch 1.14.0.dev20221007: ~24 min

I have to run to work now, so I can't dig any deeper, but I'll keep poking around this weekend.

HI @jncasey, the change in #766 has not been released yet, it may help with performance as the operation should not be moved to the CPU. However, it looks like the version of PyTorch has had a big impact. I haven't observed that in my system, so I'm not sure what could be going on.

[pcuenca]

Hi @tcapelle! Your analyses on this matter are always very interesting, I have yet to try your most recent one on my system. I'm slightly surprised about your findings of mps-PyTorch vs Tensorflow, even considering that PyTorch is a pre-release. I suspect that memory might be a factor, but we should find out for sure as my computer has 64 GB.

Regarding the problems mentioned in this issue, I'd recommend you use the branch in #766, which is still not part of v0.4.1. I'm not sure what you mean by your Core ML metrics with respect to diffusers, as we are using mps for now. I'll review your findings tomorrow and try to reproduce them :)

[jncasey]

However, it looks like the version of PyTorch has had a big impact. I haven't observed that in my system, so I'm not sure what could be going on.

Thanks, @pcuenca. That certainly seems to be the case. I'll keep poking around this weekend and report back if I can figure anything out. Thankfully, I've got access to another machine with a 3090, so I can work around the mps issues.

[jncasey]

Okay, here are a few more data points. I haven't yet tried #766 yet, but I've done some methodical (if not totally scientific) tests using the simple script in the first post, comparing diffusers 0.3.0 and 0.4.1, and using my original torch 1.13.0.dev20220911, the latest 1.14.0.dev20221007, and 1.13.0.dev20221006 (to compare 1.13 and 1.14 nightlies from similar dates). These were all run in the same environment, and using pip to change the versions up and down. The times were repeatable no matter which order the versions were upgraded or downgraded.

diffusers 0.3.0, torch 1.13.0.dev2022911 averaged 0:33 over five runs, and produced the image in the first post.
diffusers 0.4.1, torch 1.13.0.dev2022911 averaged 0:47 over five runs, and produced the noise in the first post.

There was no difference in time between the two more recent torch nightlies when using the same version of diffusers.
diffusers 0.3.0, and either of the recent torch nightlies averaged 3:58 over five runs.
diffusers 0.4.1, and either of the recent torch nightlies averaged 0:45 over five runs.
All four of those runs produced identical images:

Which is slightly different from the image generated with diffusers 0.3.0, torch 1.13.0.dev2022911:

I don't know what any of this means, but it's clear that something very different was happening with diffusers 0.3.0 and torch 1.13.0.dev2022911 that resulted in different math but much faster execution.

[patrickvonplaten]

Thanks for the very in-detail benchmarking here @jncasey !
Seems like things are still a bit unstable both from the diffusers and native PyTorch side. I guess we have to opt always for the newest PyTorch nightly.

[jncasey]

Thanks, @pcuenca! FYI it looks like the other Stable Diffusion pipelines (e.g. Img2Img, inpaint) are still using repeat_interleave, so they may want your fix, too.

The change hasn't returned performance to what I'm getting with diffusers 0.3.0 and the torch nightly from September 11, but that seems to be on the pytorch side based on this other report.