Setting up the CompVis/stable-diffusion Environment on Apple Silicon

If you just want to generate images on a Mac, I think you can do it by bringing what is running on Colab in Google Colab で はじめる Stable Diffusion v1.4｜npaka｜note to your local machine and using the CPU instead of CUDA.

main.py

from diffusers import StableDiffusionPipeline

pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", use_auth_token=YOUR_TOKEN)
pipe.to("cpu")

prompt = """
a photograph of an astronaut riding a horse
"""

images = pipe(prompt)["sample"]
for image in images:
    image.save(f"{prompt}.png")

This time, I built it from the repository because I wanted to modify the source code while understanding what CompVis/stable-diffusion is doing, such as switching to different models or checking PyTorch's MPS support.

The steps are almost the same as in M1 MacBook ProでStable Diffusionを動かすまでのメモ, but I am using a Python 3.10 environment installed under asdf rather than conda.

Checking out the repository

git clone --depth 1 [e-mail]:CompVis/stable-diffusion.git

https://github.com/CompVis/stable-diffusion/pull/47

Following magnusviri's MPS support branch. It appeared to be merged, but it seems that was a mistake.

Once this update is included in CompVis's main branch, these steps will not be necessary.

git remote add magnusviri [e-mail]:magnusviri/stable-diffusion.git
git fetch magnusviri
git rebase -i magnusviri/apple-silicon-mps-support

Installing libraries

pip install opencv-python omegaconf invisible-watermark einops clip kornia
pip install -e scripts/taming-transformers/

Adjusting PyTorch

PyTorch requires a nightly version with MPS support.

pip3 install --pre torch torchvision pytorch_lightning --extra-index-url https://download.pytorch.org/whl/nightly/cpu

$ python3 -c 'import torch; print(torch.__version__) '
1.13.0.dev20220824

$ python -c"import torch; print(torch.backends.mps.is_available())"
True

Additionally, it seems a monkey patch for site-packages/torch/nn/functional.py is required.

The input inside layer_norm() becomes input.contiguous().

def layer_norm(
    input: Tensor,
    normalized_shape: List[int],
    weight: Optional[Tensor] = None,
    bias: Optional[Tensor] = None,
    eps: float = 1e-5,
) -> Tensor:
    r"""Applies Layer Normalization for last certain number of dimensions.

    See :class:`~torch.nn.LayerNorm` for details.
    """
    if has_torch_function_variadic(input, weight, bias):
        return handle_torch_function(
            layer_norm, (input.contiguous(), weight, bias), input, normalized_shape, weight=weight, bias=bias, eps=eps
        )
    return torch.layer_norm(input.contiguous(), normalized_shape, weight, bias, eps, torch.backends.cudnn.enabled)

Downloading the model

git-lfs is a tool for handling large files in Git.

brew install git-lfs
# Enter huggingface login credentials at the prompt
git lfs clone https://huggingface.co/CompVis/stable-diffusion-v-1-4-original  models/ldm/stable-diffusion-v1

Executing txt2img.py

I am setting PYTHONPATH so that it can find ./ldm (note that it will not work with the ldm package available on PyPI).

PYTHONPATH=. python scripts/txt2img.py --ckpt models/ldm/stable-diffusion-v1/sd-v1-4.ckpt --prompt "a photograph of an astronaut riding a horse"  

With this, the environment setup is complete. However, I am wondering why running txt2img.py with MPS + DDI sampling is taking significantly longer than when I ran the initial StableDiffusionPipeline example on a plain CPU.