nix: add shell.nix file

2024-08-30 20:32:17 +00:00 · 2022-10-20 04:10:37 +09:00 · 2022-10-20 04:10:37 +09:00 · ce1bfbc32d
commit ce1bfbc32d
parent 7a923beb4c
1 changed files with 162 additions and 0 deletions
--- a/shell.nix
+++ b/shell.nix
@ -0,0 +1,162 @@
+{ pkgs ? import <nixpkgs> {}
+  , lib ? pkgs.lib
+  , stdenv ? pkgs.stdenv
+  , fetchurl ? pkgs.fetchurl
+  , runCommand ? pkgs.runCommand
+  , makeWrapper ? pkgs.makeWrapper
+  , mkShell ? pkgs.mkShell
+  , buildFHSUserEnv ? pkgs.buildFHSUserEnv
+  , frameworks ? pkgs.darwin.apple_sdk.frameworks
+}:
+
+# Setup InvokeAI environment using nix
+# Simple usage:
+# nix-shell
+# python3 scripts/preload_models.py
+# python3 scripts/invoke.py -h
+
+let
+  conda-shell = { url, sha256, installPath, packages, shellHook }:
+  let
+    src = fetchurl { inherit url sha256; };
+    libPath = lib.makeLibraryPath ([] ++ lib.optionals (stdenv.isLinux) [ pkgs.zlib ]);
+    condaArch = if stdenv.system == "aarch64-darwin" then "osx-arm64" else "";
+    installer =
+    if stdenv.isDarwin then
+      runCommand "conda-install" {
+        nativeBuildInputs = [ makeWrapper ];
+      } ''
+        mkdir -p $out/bin
+        cp ${src} $out/bin/miniconda-installer.sh
+        chmod +x $out/bin/miniconda-installer.sh
+        makeWrapper                       \
+          $out/bin/miniconda-installer.sh \
+          $out/bin/conda-install          \
+          --add-flags "-p ${installPath}" \
+          --add-flags "-b"
+      ''
+    else if stdenv.isLinux then
+      runCommand "conda-install" {
+        nativeBuildInputs = [ makeWrapper ];
+        buildInputs = [ pkgs.zlib ];
+      }
+      # on line 10, we have 'unset LD_LIBRARY_PATH'
+      # we have to comment it out however in a way that the number of bytes in the
+      # file does not change. So we replace the 'u' in the line with a '#'
+      # The reason is that the binary payload is encoded as number
+      # of bytes from the top of the installer script
+      # and unsetting the library path prevents the zlib library from being discovered
+      ''
+        mkdir -p $out/bin
+        sed 's/unset LD_LIBRARY_PATH/#nset LD_LIBRARY_PATH/' ${src} > $out/bin/miniconda-installer.sh
+        chmod +x $out/bin/miniconda-installer.sh
+        makeWrapper                       \
+          $out/bin/miniconda-installer.sh \
+          $out/bin/conda-install          \
+          --add-flags "-p ${installPath}" \
+          --add-flags "-b"                \
+          --prefix "LD_LIBRARY_PATH" : "${libPath}"
+      ''
+    else {};
+
+    hook = ''
+      export CONDA_SUBDIR=${condaArch}
+    '' +  shellHook;
+
+    fhs = buildFHSUserEnv {
+      name = "conda-shell";
+      targetPkgs = pkgs: [ stdenv.cc pkgs.git installer ] ++ packages;
+      profile = hook;
+      runScript = "bash";
+    };
+
+    shell = mkShell {
+      shellHook = if stdenv.isDarwin then hook else "conda-shell; exit";
+      packages = if stdenv.isDarwin then [ pkgs.git installer ] ++ packages else [ fhs ];
+    };
+  in shell;
+
+  packages = with pkgs; [
+    cmake
+    protobuf
+    libiconv
+    rustc
+    cargo
+    rustPlatform.bindgenHook
+  ];
+
+  env = {
+    aarch64-darwin = {
+      envFile = "environment-mac.yml";
+      condaPath = (builtins.toString ./.) + "/.conda";
+      ptrSize = "8";
+    };
+    x86_64-linux =  {
+      envFile = "environment.yml";
+      condaPath = (builtins.toString ./.) + "/.conda";
+      ptrSize = "8";
+    };
+  };
+
+  envFile = env.${stdenv.system}.envFile;
+  installPath = env.${stdenv.system}.condaPath;
+  ptrSize = env.${stdenv.system}.ptrSize;
+  shellHook = ''
+    conda-install
+
+    # tmpdir is too small in nix
+    export TMPDIR="${installPath}/tmp"
+
+    # Add conda to PATH
+    export PATH="${installPath}/bin:$PATH"
+
+    # Allows `conda activate` to work properly
+    source ${installPath}/etc/profile.d/conda.sh
+
+    # Paths for gcc if compiling some C sources with pip
+    export NIX_CFLAGS_COMPILE="-I${installPath}/include -I$TMPDIR/include"
+    export NIX_CFLAGS_LINK="-L${installPath}/lib $BINDGEN_EXTRA_CLANG_ARGS"
+
+    export PIP_EXISTS_ACTION=w
+
+    # rust-onig fails (think it writes config.h to wrong location)
+    mkdir -p "$TMPDIR/include"
+    cat <<'EOF' > "$TMPDIR/include/config.h"
+    #define HAVE_PROTOTYPES 1
+    #define STDC_HEADERS 1
+    #define HAVE_STRING_H 1
+    #define HAVE_STDARG_H 1
+    #define HAVE_STDLIB_H 1
+    #define HAVE_LIMITS_H 1
+    #define HAVE_INTTYPES_H 1
+    #define SIZEOF_INT 4
+    #define SIZEOF_SHORT 2
+    #define SIZEOF_LONG ${ptrSize}
+    #define SIZEOF_VOIDP ${ptrSize}
+    #define SIZEOF_LONG_LONG 8
+    EOF
+
+    conda env create -f "${envFile}" || conda env update --prune -f "${envFile}"
+    conda activate invokeai
+  '';
+
+  version = "4.12.0";
+  conda = {
+    aarch64-darwin = {
+      shell = conda-shell {
+        inherit shellHook installPath;
+        url = "https://repo.anaconda.com/miniconda/Miniconda3-py39_${version}-MacOSX-arm64.sh";
+        sha256 = "4bd112168cc33f8a4a60d3ef7e72b52a85972d588cd065be803eb21d73b625ef";
+        packages = [ frameworks.Security ] ++ packages;
+      };
+    };
+    x86_64-linux =  {
+      shell = conda-shell {
+        inherit shellHook installPath;
+        url = "https://repo.continuum.io/miniconda/Miniconda3-py39_${version}-Linux-x86_64.sh";
+        sha256 = "78f39f9bae971ec1ae7969f0516017f2413f17796670f7040725dd83fcff5689";
+        packages = with pkgs; [ libGL glib ] ++ packages;
+      };
+    };
+  };
+in conda.${stdenv.system}.shell