add support for Apple hardware using MPS acceleration

Author: lstein

.gitignore

@@ -180,4 +180,6 @@ outputs
 # created from generated embeddings.
 logs
 testtube
-checkpoints
+checkpoints
+# If it's a Mac
+.DS_Store

README-Mac-MPS.md

@@ -0,0 +1,228 @@
+# Apple Silicon Mac Users
+
+Several people have gotten Stable Diffusion to work on Apple Silicon
+Macs using Anaconda. I've gathered up most of their instructions and
+put them in this fork (and readme). I haven't tested anything besides
+Anaconda, and I've read about issues with things like miniforge, so if
+you have an issue that isn't dealt with in this fork then head on over
+to the [Apple
+Silicon](https://github.com/CompVis/stable-diffusion/issues/25) issue
+on GitHub (that page is so long that GitHub hides most of it by
+default, so you need to find the hidden part and expand it to view the
+whole thing). This fork would not have been possible without the work
+done by the people on that issue.
+
+You have to have macOS 12.3 Monterey or later. Anything earlier than that won't work.
+
+BTW, I haven't tested any of this on Intel Macs.
+
+How to:
+
+```
+git clone https://github.com/lstein/stable-diffusion.git
+cd stable-diffusion
+
+mkdir -p models/ldm/stable-diffusion-v1/
+ln -s /path/to/ckpt/sd-v1-1.ckpt models/ldm/stable-diffusion-v1/model.ckpt
+
+conda env create -f environment-mac.yaml
+conda activate ldm
+```
+
+These instructions are identical to the main repo except I added
+environment-mac.yaml because Mac doesn't have cudatoolkit.
+
+After you follow all the instructions and run txt2img.py you might get several errors. Here's the errors I've seen and found solutions for.
+
+### Doesn't work anymore?
+
+We are using PyTorch nightly, which includes support for MPS. I don't
+know exactly how Anaconda does updates, but I woke up one morning and
+Stable Diffusion crashed and I couldn't think of anything I did that
+would've changed anything the night before, when it worked. A day and
+a half later I finally got it working again. I don't know what changed
+overnight. PyTorch-nightly changes overnight but I'm pretty sure I
+didn't manually update it. Either way, things are probably going to be
+bumpy on Apple Silicon until PyTorch releases a firm version that we
+can lock to.
+
+To manually update to the latest version of PyTorch nightly (which could fix issues), run this command.
+
+	conda install pytorch torchvision torchaudio -c pytorch-nightly
+
+## Debugging?
+
+Tired of waiting for your renders to finish before you can see if it
+works? Reduce the steps! The picture wont look like anything but if it
+finishes, hey, it works! This could also help you figure out if you've
+got a memory problem, because I'm betting 1 step doesn't use much
+memory.
+
+	python ./scripts/txt2img.py --prompt "ocean" --ddim_steps 1
+
+### "No module named cv2" (or some other module)
+
+Did you remember to `conda activate ldm`? If your terminal prompt
+begins with "(ldm)" then you activated it. If it begins with "(base)"
+or something else you haven't.
+
+If it says you're missing taming you need to rebuild your virtual
+environment.
+
+	conda env remove -n ldm
+	conda env create -f environment-mac.yaml
+
+If you have activated the ldm virtual environment and tried rebuilding
+it, maybe the problem could be that I have something installed that
+you don't and you'll just need to manually install it. Make sure you
+activate the virtual environment so it installs there instead of
+globally.
+
+	conda activate ldm
+	pip install *name*
+
+You might also need to install Rust (I mention this again below).
+
+### "The operator [name] is not current implemented for the MPS device." (sic)
+
+Example error.
+
+```
+...
+NotImplementedError: The operator 'aten::index.Tensor' is not current implemented for the MPS device. If you want this op to be added in priority during the prototype phase of this feature, please comment on [https://github.com/pytorch/pytorch/issues/77764](https://github.com/pytorch/pytorch/issues/77764). As a temporary fix, you can set the environment variable `PYTORCH_ENABLE_MPS_FALLBACK=1` to use the CPU as a fallback for this op. WARNING: this will be slower than running natively on MPS.
+```
+
+Just do what it says:
+
+	export PYTORCH_ENABLE_MPS_FALLBACK=1
+
+### "Could not build wheels for tokenizers"
+
+I have not seen this error because I had Rust installed on my computer before I started playing with Stable Diffusion. The fix is to install Rust.
+
+	curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh
+
+### How come `--seed` doesn't work?
+
+> Completely reproducible results are not guaranteed across PyTorch
+releases, individual commits, or different platforms. Furthermore,
+results may not be reproducible between CPU and GPU executions, even
+when using identical seeds.
+
+[PyTorch docs](https://pytorch.org/docs/stable/notes/randomness.html)
+
+There is an [open issue](https://github.com/pytorch/pytorch/issues/78035) (as of August 2022) in pytorch regarding gradient inconsistency. I am guessing that's what is causing this.
+
+### libiomp5.dylib error?
+
+	OMP: Error #15: Initializing libiomp5.dylib, but found libomp.dylib already initialized.
+
+There are several things you can do. First, you could use something
+besides Anaconda like miniforge. I read a lot of things online telling
+people to use something else, but I am stuck with Anaconda for other
+reasons.
+
+Or you can try this.
+
+	export KMP_DUPLICATE_LIB_OK=True
+
+Or this (which takes forever on my computer and didn't work anyway).
+
+	conda install nomkl
+
+This error happens with Anaconda on Macs, and
+[nomkl](https://stackoverflow.com/questions/66224879/what-is-the-nomkl-python-package-used-for)
+is supposed to fix the issue (it isn't a module but a fix of some
+sort). [There's more
+suggestions](https://stackoverflow.com/questions/53014306/error-15-initializing-libiomp5-dylib-but-found-libiomp5-dylib-already-initial),
+like uninstalling tensorflow and reinstalling. I haven't tried them.
+
+### Not enough memory.
+
+This seems to be a common problem and is probably the underlying
+problem for a lot of symptoms (listed below). The fix is to lower your
+image size or to add `model.half()` right after the model is loaded. I
+should probably test it out. I've read that the reason this fixes
+problems is because it converts the model from 32-bit to 16-bit and
+that leaves more RAM for other things. I have no idea how that would
+affect the quality of the images though.
+
+See [this issue](https://github.com/CompVis/stable-diffusion/issues/71).
+
+### "Error: product of dimension sizes > 2**31'"
+
+This error happens with img2img, which I haven't played with too much
+yet. But I know it's because your image is too big or the resolution
+isn't a multiple of 32x32. Because the stable-diffusion model was
+trained on images that were 512 x 512, it's always best to use that
+output size (which is the default). However, if you're using that size
+and you get the above error, try 256 x 256 or 512 x 256 or something
+as the source image.
+
+BTW, 2**31-1 = [2,147,483,647](https://en.wikipedia.org/wiki/2,147,483,647#In_computing), which is also 32-bit signed [LONG_MAX](https://en.wikipedia.org/wiki/C_data_types) in C.
+
+### I just got Rickrolled! Do I have a virus?
+
+You don't have a virus. It's part of the project. Here's
+[Rick](https://github.com/lstein/stable-diffusion/blob/main/assets/rick.jpeg)
+and here's [the
+code](https://github.com/lstein/stable-diffusion/blob/69ae4b35e0a0f6ee1af8bb9a5d0016ccb27e36dc/scripts/txt2img.py#L79)
+that swaps him in. It's a NSFW filter, which IMO, doesn't work very
+good (and we call this "computer vision", sheesh).
+
+Actually, this could be happening because there's not enough RAM. You could try the `model.half()` suggestion or specify smaller output images.
+
+### My images come out black
+
+I haven't solved this issue. I just throw away my black
+images. There's a [similar
+issue](https://github.com/CompVis/stable-diffusion/issues/69) on CUDA
+GPU's where the images come out green. Maybe it's the same issue?
+Someone in that issue says to use "--precision full", but this fork
+actually disables that flag. I don't know why, someone else provided
+that code and I don't know what it does. Maybe the `model.half()`
+suggestion above would fix this issue too. I should probably test it.
+
+### "view size is not compatible with input tensor's size and stride"
+
+```
+  File "/opt/anaconda3/envs/ldm/lib/python3.10/site-packages/torch/nn/functional.py", line 2511, in layer_norm
+    return torch.layer_norm(input, normalized_shape, weight, bias, eps, torch.backends.cudnn.enabled)
+RuntimeError: view size is not compatible with input tensor's size and stride (at least one dimension spans across two contiguous subspaces). Use .reshape(...) instead.
+```
+
+Update to the latest version of lstein/stable-diffusion. We were
+patching pytorch but we found a file in stable-diffusion that we could
+change instead. This is a 32-bit vs 16-bit problem.
+
+### The processor must support the Intel bla bla bla
+
+What? Intel? On an Apple Silicon?
+
+	Intel MKL FATAL ERROR: This system does not meet the minimum requirements for use of the Intel(R) Math Kernel Library.
+	The processor must support the Intel(R) Supplemental Streaming SIMD Extensions 3 (Intel(R) SSSE3) instructions.██████████████| 50/50 [02:25<00:00,  2.53s/it]
+	The processor must support the Intel(R) Streaming SIMD Extensions 4.2 (Intel(R) SSE4.2) instructions.
+	The processor must support the Intel(R) Advanced Vector Extensions (Intel(R) AVX) instructions.
+
+This fixed it for me:
+
+	conda clean --yes --all
+
+### Still slow?
+
+I changed the defaults of n_samples and n_iter to 1 so that it uses
+less RAM and makes less images so it will be faster the first time you
+use it. I don't actually know what n_samples does internally, but I
+know it consumes a lot more RAM. The n_iter flag just loops around the
+image creation code, so it shouldn't consume more RAM (it should be
+faster if you're going to do multiple images because the libraries and
+model will already be loaded--use a prompt file to get this speed
+boost).
+
+These flags are the default sample and iter settings in this fork/branch:
+
+~~~~
+python scripts/txt2img.py --prompt "ocean" --n_samples=1 --n_iter=1
+~~~
+
+

README.md

@@ -387,7 +387,7 @@ Credit goes to @rinongal and the repository located at
 https://github.com/rinongal/textual_inversion Please see the
 repository and associated paper for details and limitations.
 
-# Latest
+# Latest Changes
 
 - v1.13 (in process)
 
@@ -403,9 +403,9 @@ For older changelogs, please visit **[CHANGELOGS](CHANGELOG.md)**.
 
 # Installation
 
-There are separate installation walkthroughs for [Linux/Mac](#linuxmac) and [Windows](#windows).
+There are separate installation walkthroughs for [Linux](#linux), [Windows](#windows) and [Macintosh](#Macintosh)
 
-## Linux/Mac
+## Linux
 
 1. You will need to install the following prerequisites if they are not already available. Use your
    operating system's preferred installer
@@ -580,7 +580,15 @@ python scripts\dream.py -l
 python scripts\dream.py
 ```
 
-10. Subsequently, to relaunch the script, first activate the Anaconda command window (step 3), enter the stable-diffusion directory (step 5, "cd \path\to\stable-diffusion"), run "conda activate ldm" (step 6b), and then launch the dream script (step 9).
+10. Subsequently, to relaunch the script, first activate the Anaconda
+command window (step 3), enter the stable-diffusion directory (step 5,
+"cd \path\to\stable-diffusion"), run "conda activate ldm" (step 6b),
+and then launch the dream script (step 9).
+
+**Note:** Tildebyte has written an alternative ["Easy peasy Windows
+install"](https://github.com/lstein/stable-diffusion/wiki/Easy-peasy-Windows-install)
+which uses the Windows Powershell and pew. If you are having trouble
+with Anaconda on Windows, give this a try (or try it first!)
 
 ### Updating to newer versions of the script
 
@@ -595,11 +603,16 @@ git pull
 
 This will bring your local copy into sync with the remote one.
 
-## Simplified API for text to image generation
+## Macintosh
+
+See (README-Mac-MPS)[README-Mac-MPS.md] for instructions.
+
+# Simplified API for text to image generation
 
 For programmers who wish to incorporate stable-diffusion into other
-products, this repository includes a simplified API for text to image generation, which
-lets you create images from a prompt in just three lines of code:
+products, this repository includes a simplified API for text to image
+generation, which lets you create images from a prompt in just three
+lines of code:
 
 ```
 from ldm.simplet2i import T2I
@@ -608,9 +621,10 @@ outputs = model.txt2img("a unicorn in manhattan")
 ```
 
 Outputs is a list of lists in the format [[filename1,seed1],[filename2,seed2]...]
-Please see ldm/simplet2i.py for more information.
+Please see ldm/simplet2i.py for more information. A set of example scripts is
+coming RSN.
 
-## Workaround for machines with limited internet connectivity
+# Workaround for machines with limited internet connectivity
 
 My development machine is a GPU node in a high-performance compute
 cluster which has no connection to the internet. During model

environment-mac.yaml

@@ -0,0 +1,32 @@
+name: ldm
+channels:
+  - apple
+  - conda-forge
+  - pytorch-nightly
+  - defaults
+dependencies:
+  - python=3.10.4
+  - pip=22.1.2
+  - pytorch
+  - torchvision
+  - numpy=1.23.1
+  - pip:
+    - albumentations==0.4.6
+    - opencv-python==4.6.0.66
+    - pudb==2019.2
+    - imageio==2.9.0
+    - imageio-ffmpeg==0.4.2
+    - pytorch-lightning==1.4.2
+    - omegaconf==2.1.1
+    - test-tube>=0.7.5
+    - streamlit==1.12.0
+    - pillow==9.2.0
+    - einops==0.3.0
+    - torch-fidelity==0.3.0
+    - transformers==4.19.2
+    - torchmetrics==0.6.0
+    - kornia==0.6.0
+    - -e git+https://github.com/openai/CLIP.git@main#egg=clip
+    - -e git+https://github.com/CompVis/taming-transformers.git@master#egg=taming-transformers
+    - -e git+https://github.com/lstein/k-diffusion.git@master#egg=k-diffusion
+    - -e .

ldm/dream/devices.py

@@ -0,0 +1,11 @@
+import torch
+
+def choose_torch_device() -> str:
+    '''Convenience routine for guessing which GPU device to run model on'''
+    if torch.cuda.is_available():
+        return 'cuda'
+    if torch.backends.mps.is_available():
+        return 'mps'
+    return 'cpu'
+
+    

ldm/models/diffusion/ddim.py

@@ -4,6 +4,7 @@
 import numpy as np
 from tqdm import tqdm
 from functools import partial
+from ldm.dream.devices import choose_torch_device
 
 from ldm.modules.diffusionmodules.util import (
     make_ddim_sampling_parameters,
@@ -14,17 +15,17 @@
 
 
 class DDIMSampler(object):
-    def __init__(self, model, schedule='linear', device='cuda', **kwargs):
+    def __init__(self, model, schedule='linear', device=None, **kwargs):
         super().__init__()
         self.model = model
         self.ddpm_num_timesteps = model.num_timesteps
         self.schedule = schedule
-        self.device = device
+        self.device   = device or choose_torch_device()
 
     def register_buffer(self, name, attr):
         if type(attr) == torch.Tensor:
             if attr.device != torch.device(self.device):
-                attr = attr.to(torch.device(self.device))
+                attr = attr.to(dtype=torch.float32, device=self.device)
         setattr(self, name, attr)
 
     def make_schedule(