From a2e79fcaf7e08dbf9fc5056732c96edaf9fbc2c5 Mon Sep 17 00:00:00 2001
From: Jari Vetoniemi <jari.vetoniemi@cloudef.pw>
Date: Tue, 21 Nov 2023 05:35:25 +0800
Subject: [PATCH] feat: add CPU and Apple Silicon support

On Apple silicon, the CPU path is still hardware accelerated.
---
 .gitignore |  1 +
 README.md  | 40 +++++++++++++++++++++++++++++++---------
 flake.nix  | 13 ++++++-------
 impl.nix   | 14 +++++++-------
 shell.nix  |  4 ++--
 5 files changed, 47 insertions(+), 25 deletions(-)
 create mode 100644 .gitignore

diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..8352b6b
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1 @@
+stable-diffusion-webui
diff --git a/README.md b/README.md
index 8e7ea62..14b6516 100644
--- a/README.md
+++ b/README.md
@@ -1,17 +1,39 @@
 # [AUTOMATIC1111/stable-diffusion-webui](https://github.com/AUTOMATIC1111/stable-diffusion-webui) for CUDA and ROCm on NixOS
 
 This is literally just a `shell.nix`/`flake.nix` for stable-diffusion-webui that also enables CUDA/ROCm on NixOS.
-This supports both NVIDIA GPUs (using CUDA) and AMD GPUs (using ROCm).
+This supports NVIDIA GPUs (using CUDA), AMD GPUs (using ROCm), and CPU compute (including Apple silicon).
 
 ## Usage
+
+### Setup
+
 ```bash
-git clone https://github.com/AUTOMATIC1111/stable-diffusion-webui
 git clone https://github.com/virchau13/automatic1111-webui-nix
-cp automatic1111-webui-nix/*.nix stable-diffusion-webui/
-cd stable-diffusion-webui
-git add *.nix
-nix develop # or `nix-shell` if you're not using flakes
-# just use `./webui.sh` to run it, it'll install all the rest automatically
+cd automatic1111-webui-nix
+git clone https://github.com/AUTOMATIC1111/stable-diffusion-webui
+```
+
+### Flakes
+
+```bash
+nix develop # CPU or Apple silicon
+nix develop .#cuda # CUDA
+nix develop .#rocm # ROCm
+```
+
+### Non-flakes
+
+```bash
+nix-shell # CPU
+nix-shell --argstr variant CUDA # CUDA
+nix-shell --argstr variant ROCM # ROCm
+```
+
+### Launching webui
+
+```bash
+cd automatic1111-webui-nix
+./webui.sh
 # follow the tutorials at the original project for setting up Stable Diffusion / GFPGAN / whatever
 ```
 
@@ -22,8 +44,7 @@ You might want to switch to high performance mode on battery-powered devices.
 This is just a Nix shell for bootstrapping the web UI, not an actual pure flake; the `./webui.sh` will still install
 a bunch of Python packages (into a venv, so not polluting your system) when you run it.
 
-
-## Troubleshooting 
+## Troubleshooting
 
 ### xformers
 Run `./webui.sh --xformers` to set up xformers. If this command fails, try git pulling from the source repository, then deleting the venv (`rm -r venv`) and reinstalling everything from scratch again (sometimes the venv will keep old packages around forever even if they should be updated.)
@@ -31,4 +52,5 @@ Run `./webui.sh --xformers` to set up xformers. If this command fails, try git p
 ## Credits
 - AUTOMATIC1111 for obvious reasons.
 - rprospero for [ROCm support](https://github.com/virchau13/automatic1111-webui-nix/pull/3).
+- Cloudef for [CPU compute / Apple Silicon support](https://github.com/virchau13/automatic1111-webui-nix/pull/10)
 - polypoyo for [the original draft of this](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/4736).
diff --git a/flake.nix b/flake.nix
index cfb118a..ee34e10 100644
--- a/flake.nix
+++ b/flake.nix
@@ -6,17 +6,16 @@
         flake-utils.url = github:numtide/flake-utils;
     };
 
-    outputs = { self, nixpkgs, flake-utils }: let
-        isLinux = system: builtins.match ".*linux.*" system != null;
-        linuxSystems = builtins.filter isLinux flake-utils.lib.defaultSystems;
-    in flake-utils.lib.eachSystem linuxSystems (system: let
+    outputs = { self, nixpkgs, flake-utils }:
+    flake-utils.lib.eachSystem flake-utils.lib.defaultSystems (system: let
             pkgs = import nixpkgs {
                 inherit system;
                 config.allowUnfree = true;
             };
-        in { 
-            devShells.default = import ./impl.nix {inherit pkgs;};
-            devShells.rocm = import ./impl.nix {inherit pkgs; isCUDA=false;};
+        in {
+            devShells.default = import ./impl.nix { inherit pkgs; variant = "CPU"; };
+            devShells.cuda = import ./impl.nix { inherit pkgs; variant = "CUDA"; };
+            devShells.rocm = import ./impl.nix { inherit pkgs; variant = "ROCM"; };
         }
     );
 }
diff --git a/impl.nix b/impl.nix
index 9279fb7..3df2c5a 100644
--- a/impl.nix
+++ b/impl.nix
@@ -1,8 +1,8 @@
-{ pkgs, isCUDA ? true, ... }:
+{ pkgs, variant ? "CUDA", ... }:
 
 let
   hardware_deps = with pkgs;
-    if isCUDA then [
+    if variant == "CUDA" then [
       cudatoolkit
       linuxPackages.nvidia_x11
       xorg.libXi
@@ -13,13 +13,13 @@ let
       xorg.libXv
       xorg.libXrandr
       zlib
-      
+
       # for xformers
       gcc
-    ] else [
+    ] else if variant == "ROCM" then [
       rocm-runtime
       pciutils
-    ];
+    ] else [];
 
 in
 pkgs.mkShell rec {
@@ -38,6 +38,6 @@ pkgs.mkShell rec {
         glib
       ];
     LD_LIBRARY_PATH = pkgs.lib.makeLibraryPath buildInputs;
-    CUDA_PATH = pkgs.lib.optionalString isCUDA pkgs.cudatoolkit;
-    EXTRA_LDFLAGS = pkgs.lib.optionalString isCUDA "-L${pkgs.linuxPackages.nvidia_x11}/lib";
+    CUDA_PATH = pkgs.lib.optionalString (variant == "CUDA") pkgs.cudatoolkit;
+    EXTRA_LDFLAGS = pkgs.lib.optionalString (variant == "CUDA") "-L${pkgs.linuxPackages.nvidia_x11}/lib";
 }
diff --git a/shell.nix b/shell.nix
index 52257bb..7afbe7f 100644
--- a/shell.nix
+++ b/shell.nix
@@ -1,4 +1,4 @@
-{ isCUDA ? true }:
+{ variant ? "CUDA" }:
 
 let pkgs = import <nixpkgs> { };
-in import ./impl.nix { inherit pkgs isCUDA; }
+in import ./impl.nix { inherit pkgs variant; }