Merge pull request #3 from rprospero/rocm

Add rocm support

flake.nix

@@ -15,7 +15,8 @@
                 config.allowUnfree = true;
             };
         in { 
-            devShells.default = import ./impl.nix pkgs;
+            devShells.default = import ./impl.nix {inherit pkgs;};
+            devShells.rocm = import ./impl.nix {inherit pkgs; isCUDA=false;};
         }
     );
 }

impl.nix

@@ -1,8 +1,28 @@
-pkgs:
+{ pkgs, isCUDA ? true, ... }:
 
+let
+  hardware_deps = with pkgs;
+    if isCUDA then [
+      cudatoolkit
+      linuxPackages.nvidia_x11
+      xorg.libXi
+      xorg.libXmu
+      freeglut
+      xorg.libXext
+      xorg.libX11
+      xorg.libXv
+      xorg.libXrandr
+      zlib
+    ] else [
+      rocm-runtime
+      pciutils
+    ];
+
+in
 pkgs.mkShell rec {
     name = "stable-diffusion-webui";
-    buildInputs = with pkgs; [
+    buildInputs = with pkgs;
+      hardware_deps ++ [
         git # The program instantly crashes if git is not present, even if everything is already downloaded
         python310
         stdenv.cc.cc.lib
@@ -11,15 +31,10 @@ pkgs.mkShell rec {
         binutils
         gitRepo gnupg autoconf curl
         procps gnumake util-linux m4 gperf unzip
-        cudatoolkit linuxPackages.nvidia_x11
         libGLU libGL
-        xorg.libXi xorg.libXmu freeglut
-        xorg.libXext xorg.libX11 xorg.libXv xorg.libXrandr zlib
         glib
     ];
-    shellHook = ''
-        export LD_LIBRARY_PATH="${pkgs.lib.makeLibraryPath buildInputs}"
-        export CUDA_PATH=${pkgs.cudatoolkit}
-        export EXTRA_LDFLAGS="-L${pkgs.linuxPackages.nvidia_x11}/lib"
-    '';
+    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";
 }

shell.nix

@@ -1 +1,4 @@
-let pkgs = import <nixpkgs> {}; in import ./impl.nix pkgs
+{ isCUDA ? true }:
+
+let pkgs = import <nixpkgs> { };
+in import ./impl.nix { inherit pkgs isCUDA; }