partial support for xformers without triton

flake.nix

@@ -19,7 +19,7 @@
     let
       nixlib = inputs.nixlib.outputs.lib;
       system = "x86_64-linux";
-      requirementsFor = { pkgs, webui ? false }: with pkgs; with pkgs.python3.pkgs; [
+      requirementsFor = { pkgs, webui ? false, nvidia ? false }: with pkgs; with pkgs.python3.pkgs; [
         python3
         torch
         torchvision
@@ -47,6 +47,7 @@
         pillow
         safetensors
       ]
+      ++ nixlib.optional (nvidia) [ xformers ]
       ++ nixlib.optional (!webui) [
         npyscreen
         huggingface-hub
@@ -90,6 +91,7 @@
         blip
         psutil
         openclip
+        blendmodes
       ];
       overlay_default = nixpkgs: pythonPackages:
         {
@@ -150,6 +152,10 @@
             "safetensors"
             "picklescan"
             "openclip"
+            "blendmodes"
+            "xformers"
+            "pyre-extensions"
+            # "triton" nixpkgs is missing required llvm parts - mlir
           ];
         in
         {
@@ -233,7 +239,7 @@
                 };
               })
           ];
-        };
+        } // { inherit nvidia; };
     in
     {
       packages.${system} =
@@ -246,7 +252,7 @@
             version = "2.3.1";
             src = invokeai-repo;
             format = "pyproject";
-            propagatedBuildInputs = requirementsFor { pkgs = nixpkgs; };
+            propagatedBuildInputs = requirementsFor { pkgs = nixpkgs; nvidia = nixpkgs.nvidia; };
             nativeBuildInputs = [ nixpkgs.pkgs.pythonRelaxDepsHook ];
             pythonRelaxDeps = [ "torch" "pytorch-lightning" "flask-socketio" "flask" "dnspython" ];
             pythonRemoveDeps = [ "opencv-python" "flaskwebgui" "pyreadline3" ];

packages/accelerate/default.nix

@@ -6,11 +6,11 @@
 
 buildPythonPackage rec {
   pname = "accelerate";
-  version = "0.13.1";
+  version = "0.17.0";
 
   src = fetchPypi {
     inherit pname version;
-    sha256 = "1dk82s80rq8xp3v4hr9a27vgj9k3gy9yssp7ww7i3c0vc07gx2cv";
+    sha256 = "sha256-FvImuCGswe/obw4NXY8lZT3/VSNzXs4JkP7+p+t2B2o=";
   };
 
   propagatedBuildInputs = [ numpy packaging psutil pyyaml torch ];

packages/blendmodes/default.nix

@@ -0,0 +1,26 @@
+# WARNING: This file was automatically generated. You should avoid editing it.
+# If you run pynixify again, the file will be either overwritten or
+# deleted, and you will lose the changes you made to it.
+
+{ aenum, buildPythonPackage, deprecation, fetchPypi, lib, numpy, pillow }:
+
+buildPythonPackage rec {
+  pname = "blendmodes";
+  version = "2022";
+
+  src = fetchPypi {
+    inherit pname version;
+    sha256 = "00zx3025rlpgdsy78znm9xqml3kfd41sq34gjss391qfx70z2s4k";
+  };
+
+  propagatedBuildInputs = [ pillow aenum deprecation numpy ];
+
+  # TODO FIXME
+  doCheck = false;
+
+  meta = with lib; {
+    description =
+      "Use this module to apply a number of blending modes to a background and foreground image";
+    homepage = "https://github.com/FHPythonUtils/BlendModes";
+  };
+}

packages/pyre-extensions/default.nix

@@ -0,0 +1,25 @@
+# WARNING: This file was automatically generated. You should avoid editing it.
+# If you run pynixify again, the file will be either overwritten or
+# deleted, and you will lose the changes you made to it.
+
+{ buildPythonPackage, fetchPypi, lib, typing-extensions, typing-inspect }:
+
+buildPythonPackage rec {
+  pname = "pyre-extensions";
+  version = "0.0.30";
+
+  src = fetchPypi {
+    inherit pname version;
+    sha256 = "sha256-unkjxIbgia+zehBiOo9K6C1zz/QkJtcRxIrwcOW8MbI=";
+  };
+
+  propagatedBuildInputs = [ typing-inspect typing-extensions ];
+
+  # TODO FIXME
+  doCheck = false;
+
+  meta = with lib; {
+    description = "Type system extensions for use with the pyre type checker";
+    homepage = "https://pyre-check.org";
+  };
+}

packages/triton/default.nix

@@ -0,0 +1,39 @@
+# WARNING: This file was automatically generated. You should avoid editing it.
+# If you run pynixify again, the file will be either overwritten or
+# deleted, and you will lose the changes you made to it.
+
+{ buildPythonPackage, fetchFromGitHub, lib, cmake, fetchPypi, llvm }:
+
+buildPythonPackage rec {
+  pname = "triton";
+  version = "2.0.0";
+
+  src = fetchFromGitHub {
+    owner = "openai";
+    repo = "triton";
+    rev = "v${version}";
+    sha256 = "sha256-9GZzugab+Pdt74Dj6zjlEzjj4BcJ69rzMJmqcVMxsKU=";
+  };
+  #sourceRoot = ["source/python"];
+
+  LLVM_ROOT_DIR="${llvm}";
+  nativeBuildInputs = [ cmake llvm ]; #llvm clang git ];
+  #propagatedBuildInputs = [ torch filelock ncurses zlib pybind11 cmake-py ];
+
+  /* postConfigure = '' */
+  /*   cd .. */
+  /* ''; */
+  /*  */
+
+  # TODO FIXME
+  doCheck = false;
+
+  meta = with lib; {
+    description = "A language and compiler for custom Deep Learning operations";
+    homepage = "https://github.com/openai/triton/";
+    # there are no mlir support for LLVM in nixpkgs (yet). Don't think it's worth to pull in unfinished PRs and force users recompile half of llvm. After all, triton support is nice to have, not mandatory.
+    # https://github.com/NixOS/nixpkgs/pull/163878
+    # https://github.com/openai/triton/issues/1154#issuecomment-1419455783
+    broken = true;
+  };
+}

packages/xformers/default.nix

@@ -0,0 +1,23 @@
+# WARNING: This file was automatically generated. You should avoid editing it.
+# If you run pynixify again, the file will be either overwritten or
+# deleted, and you will lose the changes you made to it.
+
+{ buildPythonPackage, fetchPypi, lib, torch, numpy, pyre-extensions, pythonRelaxDepsHook, which }:
+
+buildPythonPackage rec {
+  pname = "xformers";
+  version = "0.0.16";
+
+  src = fetchPypi {
+    inherit pname version;
+    sha256 = "sha256-ksfwVWpo9EhkkmkbP1ZxQO4ZK1Y9kEGFtmabH4u4rlM=";
+  };
+  nativeBuildInputs = [ pythonRelaxDepsHook which ];
+  pythonRelaxDeps = [ "pyre-extensions" ];
+  propagatedBuildInputs = [ torch numpy pyre-extensions /*triton*/ ];
+
+  # TODO FIXME
+  doCheck = false;
+
+  meta = with lib; { };
+}