Add inference support for `mps` device (Apple Silicon)

[pcuenca]

There is a lot of community interest in running diffusers on Apple Silicon. A first step could be to introduce support for the mps device, which currently requires PyTorch nightly builds. Another step down the line could be to convert to Core ML and optimize to make sure that the use of the ANE (Neural Engine) is maximized.

This PR deals with the first approach.

Describe the solution you'd like
The following should work:

pipe = StableDiffusionPipeline.from_pretrained(model_id, use_auth_token=True)
pipe = pipe.to("mps")

# Rest of inference code remains the same

pipe.to would determine whether mps is available in the computer, and would raise an error if it's not.

One consideration is that perhaps not all the components have to be moved to the mps device. Perhaps it's more efficient or practical to keep the text encoder in CPU, for instance. If so, to() would move the required components and the pipelines would need to be adapted to move the tensors transparently.

Describe alternatives you've considered
Conversion to Core ML, as mentioned above.

Additional context
I have tested the unet module in mps vs cpu, and these are some preliminary results in my computer (M1 Max with 64 GB of unified RAM), when iterating through the 51 default steps of Stable Diffusion with the default scheduler and no classifier-free guidance:

Device: cpu, torch: 1.12.1,  time: 92.09s
Device: cpu, torch: 1.13.0.dev20220830,  time: 92.12s
Device: mps:0, torch: 1.13.0.dev20220830,  time: 17.20s

Execution in mps complains about one operation having to be performed in CPU:

The operator 'aten::masked_select' 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.)

We need to investigate the reason and whether an alternative would yield better performance.

[apolinario]

If it is of any help with this issue, here are some community-implemented or contributed Silicon compatibility with Stable Diffusion:

Using diffusers
PR to this repo: #278
Guide by wandb: https://wandb.ai/morgan/stable-diffusion/reports/Running-Stable-Diffusion-on-an-Apple-M1-Mac-with-Hugging-Face-Diffusers--VmlldzoyNTU2ODc2

Using the CompVis Stable Diffusion repo:

• https://github.com/einanao/stable-diffusion/tree/apple-silicon
• https://github.com/magnusviri/stable-diffusion/tree/apple-mps-support
• https://github.com/lstein/stable-diffusion

[mja]

Here's everything I've learned from the above linked resources as a Jupyter notebook that can be run in VSCode on MacOS: StableDiffusion-MPS.ipynb.

When running the pipeline I've confirmed using Activity Monitor that the Python process is using the GPU.

[mlajtos]

Regarding the alternative, i.e. converting the model to CoreML:

@madebyollin made gist that can convert the U-net to a CoreML model, which can be injected into the pipeline, or used as a standalone model in a native app. However, it is not a complete model as this library (diffusers) is doing an extra work around the model.

[pcuenca]

From what I've seen so far this task entails the following:

• Make the pipeline compatible when using classifier-free guidance. This is the contiguous() patch mentioned by mja in his notebook and addressed in Run x.contiguous to support Apple MPS #278. Interestingly, the unet works fine when not using CFG (i.e., input tensors are not concatenated).
• Reproducibility. We need a workaround for seeds / generators not working properly, see MPS: torch.manual_seed not working on metal (mps) for torch.randn pytorch/pytorch#84288
• "Warmup". The first time you run a batch through the forward pass the results are different to those from subsequent passes (for the same input). This is most puzzling, and occurs in both the unet and vae models. This breaks determinism tests, as pointed out by @kashif. I'd like to do two things here:
  • Create a temporary workaround so we can use mps in the codebase soon.
  • Identify the root cause and open an issue in PyTorch (if it's a PyTorch issue, which we don't know yet).
• Disable training tests for the mps device. There are additional problems with gradients when using mps, so we'll work on training later.
• Attempt to replace the ops still not ported (aten::masked_select).

Regarding @mja's notebook, he mentions that it's very slow, but it's much faster than the CPU in my tests. These are my preliminary results on M1 Max using StableDiffusionPipeline with 50 inference steps for a single image:

[mps] -  30.81s
[cpu] - 213.46s

Is this similar to what you are seeing, @mja?

[mja]

[mps] -  30.81s
[cpu] - 213.46s

Is this similar to what you are seeing, @mja?

On an M1 Mac Mini I'm getting 188s on MPS and 343s on CPU for 50 steps at 512x512.

Edit: Doing a bit more testing, I think it's the RAM that might be the limit on base M1 devices that is causing poor performance. I see the Python process grow to around 14GB and lots of red in the Activity Monitor Memory Pressure graph.