Apple MPS Compatibility

[FahimF]

Updated the "Stable Diffusion Deep Dive" and stable_diffusion notebooks to be Apple MPS compatible as much as possible. Added performance enhancements for MPS and added comments where necessary when the code diverged from the original code.

Should be fully MPS compatible and run on an MPS device except for one main issue. I have opened a ticket with Hugging Face to see if the issue can be resolved from their side ...

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[pcuenca]

Hi @FahimF!

First a note: the issue you mention will be resolved when this PR is merged: huggingface/diffusers#942. Thanks for your report!

Regarding the actual support of the mps device, I think it's a bit tangential to the main purpose of the notebooks, which is to illustrate how diffusion models work, and might be a distraction. Perhaps it would be more useful to create a copy of the notebooks instead, so students willing to try them out on their Macs can refer to them. How do you feel about this topic @johnowhitaker and @jph00?

[FahimF]

Hey @pcuenca 🙂 If the above issue and the other one where you can't pass an array of prompts to the pipeline on MPS (I believe that one's something which needs to be fixed from the PyTorch end) are fixed, then the only thing that the notebooks would require is to add a device based on GPU and use that device name instead of hardcoding for CUDA.

But you are right, as it stands now, it might be too much of a distraction and something which requires too many explanations. (And also a possible source of confusion). My original idea was to just have a separate branch on my forked repo with the changes and put up a link to that on the forums. But was wondering if perhaps somebody might go, "How come you didn't submit a PR?" 🙂

But I think that might be the least trouble overall till those issues are resolved. At that point, it shouldn't matter much anyway. Will wait in case there's another opinion, but if not, will close the PR and proceed with adding a link on the forums?

[jph00]

Thanks for the interesting PR! I agree with your conclusion. Hopefully some or all of these ideas will make their way into Diffusers given time - if supporting Mac was as simple as adding a device name at the top or there was an automatic device selection thingie, that would be awesome. Then we'd be able to support Mac inference without adding lots of code to the notebook!

[drscotthawley]

(Dropping an update rather than starting a brand-new Issue)

Those of interested in Mac support will do well to keep an eye on relevant active PyTorch Issues and Pull requests:

• autocast: [MPS] Add support for autocast in MPS  pytorch/pytorch#99272
• ConvNet3D: torch.nn.Conv3D on MPS backend pytorch/pytorch#77818
• BF16 support: torch.embedding: Trying to convert BFloat16 to the MPS backend but it does not have support for that dtype. pytorch/pytorch#104191