cleanup and modularize

flake.nix

@@ -3,7 +3,7 @@
 
   inputs = {
     nixpkgs = {
-      url = "github:NixOS/nixpkgs/nixos-unstable"; # ?rev=fd54651f5ffb4a36e8463e0c327a78442b26cbe7";
+      url = "github:NixOS/nixpkgs/nixos-unstable";
     };
     invokeai-src = {
       url = "github:invoke-ai/InvokeAI/v2.2.5";
@@ -15,195 +15,14 @@
     };
   };
   outputs = { flake-parts, invokeai-src, ... }@inputs:
-    flake-parts.lib.mkFlake { inherit inputs; } ({ inputs, lib, ... }: {
+    flake-parts.lib.mkFlake { inherit inputs; } {
       systems = [
         "x86_64-linux"
       ];
-      perSystem = { system, ... }:
+      imports = [
-    let
+        ./modules/dependency-sets
-      pkgs = import inputs.nixpkgs {
+        ./modules/aipython3
-        inherit system;
+        ./projects/invokeai
-        config.allowUnfree = true;
+      ];
-      };
+  };
-      python3 = pkgs.python3.override {
-        packageOverrides = lib.composeManyExtensions [
-          overlay_default
-          overlay_pynixify
-          overlay_amd
-        ];
-      };
-      requirements = with python3.pkgs; rec {
-        base = [
-          numpy
-          albumentations
-          opencv4
-          pudb
-          imageio
-          imageio-ffmpeg
-          pytorch-lightning
-          protobuf
-          omegaconf
-          test-tube
-          ((pkgs.streamlit.overrideAttrs (old: {
-            nativeBuildInputs = old.nativeBuildInputs ++ [ pythonRelaxDepsHook ];
-            pythonRelaxDeps = [ "protobuf" ];
-          })).override { protobuf3 = protobuf; })
-          einops
-          taming-transformers-rom1504
-          torch-fidelity
-          torchmetrics
-          transformers
-          kornia
-          k-diffusion
-          picklescan
-          diffusers
-          pypatchmatch
-
-          # following packages not needed for vanilla SD but used by both UIs
-          realesrgan
-          pillow
-        ];
-        invokeai = base ++ [
-          send2trash
-          flask
-          flask-socketio
-          flask-cors
-          dependency-injector
-          gfpgan
-          eventlet
-          clipseg
-          getpass-asterisk
-        ];
-      };
-      overlay_default = final: prev:
-        {
-          pytorch-lightning = prev.pytorch-lightning.overrideAttrs (old: {
-            nativeBuildInputs = old.nativeBuildInputs ++ [ final.pythonRelaxDepsHook ];
-            pythonRelaxDeps = [ "protobuf" ];
-          });
-          wandb = prev.wandb.overrideAttrs (old: {
-            nativeBuildInputs = old.nativeBuildInputs ++ [ final.pythonRelaxDepsHook ];
-            pythonRelaxDeps = [ "protobuf" ];
-          });
-          scikit-image = final.scikitimage;
-        };
-      overlay_pynixify = final: prev:
-        let
-          rm = d: d.overrideAttrs (old: {
-            nativeBuildInputs = old.nativeBuildInputs ++ [ final.pythonRelaxDepsHook ];
-            pythonRemoveDeps = [ "opencv-python-headless" "opencv-python" "tb-nightly" "clip" ];
-          });
-          callPackage = final.callPackage;
-          rmCallPackage = path: args: rm (callPackage path args);
-        in
-        rec {
-          pydeprecate = callPackage ./packages/pydeprecate { };
-          taming-transformers-rom1504 =
-            callPackage ./packages/taming-transformers-rom1504 { };
-          albumentations = rmCallPackage ./packages/albumentations { opencv-python-headless = final.opencv4; };
-          qudida = rmCallPackage ./packages/qudida { opencv-python-headless = final.opencv4; };
-          gfpgan = rmCallPackage ./packages/gfpgan { opencv-python = final.opencv4; };
-          basicsr = rmCallPackage ./packages/basicsr { opencv-python = final.opencv4; };
-          facexlib = rmCallPackage ./packages/facexlib { opencv-python = final.opencv4; };
-          realesrgan = rmCallPackage ./packages/realesrgan { opencv-python = final.opencv4; };
-          codeformer = callPackage ./packages/codeformer { opencv-python = final.opencv4; };
-          clipseg = rmCallPackage ./packages/clipseg { opencv-python = final.opencv4; };
-          filterpy = callPackage ./packages/filterpy { };
-          kornia = callPackage ./packages/kornia { };
-          lpips = callPackage ./packages/lpips { };
-          ffmpy = callPackage ./packages/ffmpy { };
-          shap = callPackage ./packages/shap { };
-          picklescan = callPackage ./packages/picklescan { };
-          diffusers = callPackage ./packages/diffusers { };
-          pypatchmatch = callPackage ./packages/pypatchmatch { };
-          fonts = callPackage ./packages/fonts { };
-          font-roboto = callPackage ./packages/font-roboto { };
-          analytics-python = callPackage ./packages/analytics-python { };
-          markdown-it-py = callPackage ./packages/markdown-it-py { };
-          gradio = callPackage ./packages/gradio { };
-          hatch-requirements-txt = callPackage ./packages/hatch-requirements-txt { };
-          timm = callPackage ./packages/timm { };
-          blip = callPackage ./packages/blip { };
-          fairscale = callPackage ./packages/fairscale { };
-          torch-fidelity = callPackage ./packages/torch-fidelity { };
-          resize-right = callPackage ./packages/resize-right { };
-          torchdiffeq = callPackage ./packages/torchdiffeq { };
-          k-diffusion = callPackage ./packages/k-diffusion { clean-fid = final.clean-fid; };
-          accelerate = callPackage ./packages/accelerate { };
-          clip-anytorch = callPackage ./packages/clip-anytorch { };
-          jsonmerge = callPackage ./packages/jsonmerge { };
-          clean-fid = callPackage ./packages/clean-fid { };
-          getpass-asterisk = callPackage ./packages/getpass-asterisk { };
-        };
-      overlay_amd = final: prev:
-        rec {
-          # TODO: figure out how to patch torch-bin trying to access /opt/amdgpu
-          # there might be an environment variable for it, can use a wrapper for that
-          # otherwise just grep the world for /opt/amdgpu or something and substituteInPlace the path
-          # you can run this thing without the fix by creating /opt and running nix build nixpkgs#libdrm --inputs-from . --out-link /opt/amdgpu
-          torch-bin = prev.torch-bin.overrideAttrs (old: {
-            src = pkgs.fetchurl {
-              name = "torch-1.13.1+rocm5.1.1-cp310-cp310-linux_x86_64.whl";
-              url = "https://download.pytorch.org/whl/rocm5.1.1/torch-1.13.1%2Brocm5.1.1-cp310-cp310-linux_x86_64.whl";
-              hash = "sha256-qUwAL3L9ODy9hjne8jZQRoG4BxvXXLT7cAy9RbM837A=";
-            };
-          });
-          torchvision-bin = prev.torchvision-bin.overrideAttrs (old: {
-            src = pkgs.fetchurl {
-              name = "torchvision-0.14.1+rocm5.1.1-cp310-cp310-linux_x86_64.whl";
-              url = "https://download.pytorch.org/whl/rocm5.1.1/torchvision-0.14.1%2Brocm5.1.1-cp310-cp310-linux_x86_64.whl";
-              hash = "sha256-8CM1QZ9cZfexa+HWhG4SfA/PTGB2475dxoOtGZ3Wa2E=";
-            };
-          });
-          torch = torch-bin;
-          torchvision = torchvision-bin;
-          #torch = pythonPackages.torch.override {
-          #  rocmSupport = true;
-          #  magma = nixpkgs.magma-hip;
-          #};
-          huggingface-hub = prev.huggingface-hub.overrideAttrs(_: {
-            src = pkgs.fetchFromGitHub {
-              owner = "huggingface";
-              repo = "huggingface_hub";
-              rev = "refs/tags/v0.11.0";
-              hash = "sha256-d+X4hGt4K6xmRFw8mevKpZ6RDv+U1PJ8WbmdKGDbVNs=";
-              };
-          });
-          #overriding because of https://github.com/NixOS/nixpkgs/issues/196653
-          opencv4 = prev.opencv4.override { openblas = pkgs.blas; };
-        };
-      overlay_nvidia = final: prev:
-        {
-          torch = final.torch-bin;
-          torchvision = final.torchvision-bin;
-          huggingface-hub = prev.huggingface-hub.overrideAttrs(_: {
-            src = pkgs.fetchFromGitHub {
-              owner = "huggingface";
-              repo = "huggingface_hub";
-              rev = "refs/tags/v0.11.0";
-              hash = "sha256-d+X4hGt4K6xmRFw8mevKpZ6RDv+U1PJ8WbmdKGDbVNs=";
-              };
-          });
-          opencv4 = prev.opencv4.override { openblas = pkgs.blas; };
-        };
-    in
-    {
-      packages = {
-        invokeai-amd = python3.pkgs.buildPythonPackage {
-          pname = "InvokeAI";
-          version = toString (builtins.match ".+VERSION = '([^']+)'.+" (builtins.readFile "${inputs.invokeai-src}/setup.py"));
-          src = invokeai-src;
-          propagatedBuildInputs = requirements.invokeai;
-          nativeBuildInputs = [ python3.pkgs.pythonRelaxDepsHook ];
-          pythonRemoveDeps = [ "clip" "pyreadline3" "flaskwebgui" ];
-          pythonRelaxDeps = [ "protobuf" ];
-          postFixup = ''
-            chmod +x $out/bin/*
-            wrapPythonPrograms
-          '';
-          doCheck = false;
-        };
-      };
-    };
-  });
 }

modules/aipython3/default.nix

@@ -0,0 +1,28 @@
+{ lib, ... }:
+
+{
+  perSystem = { pkgs, ... }: {
+    dependencySets = let
+      overlays = import ./overlays.nix pkgs;
+
+      mkPythonPackages = overlayList: let
+        python3' = pkgs.python3.override {
+          packageOverrides = lib.composeManyExtensions overlayList;
+        };
+      in python3'.pkgs;
+
+    in {
+      aipython3-amd = mkPythonPackages [
+        overlays.fixPackages
+        overlays.extraDeps
+        overlays.torchRocm
+      ];
+
+      aipython3-nvidia = mkPythonPackages [
+        overlays.fixPackages
+        overlays.extraDeps
+        overlays.torchCuda
+      ];
+    };
+  };
+}

modules/aipython3/overlays.nix

@@ -0,0 +1,90 @@
+pkgs: {
+  fixPackages = final: prev: {
+    pytorch-lightning = prev.pytorch-lightning.overrideAttrs (old: {
+      nativeBuildInputs = old.nativeBuildInputs ++ [ final.pythonRelaxDepsHook ];
+      pythonRelaxDeps = [ "protobuf" ];
+    });
+    wandb = prev.wandb.overrideAttrs (old: {
+      nativeBuildInputs = old.nativeBuildInputs ++ [ final.pythonRelaxDepsHook ];
+      pythonRelaxDeps = [ "protobuf" ];
+    });
+    scikit-image = final.scikitimage;
+    #overriding because of https://github.com/NixOS/nixpkgs/issues/196653
+    opencv4 = prev.opencv4.override { openblas = pkgs.blas; };
+  };
+
+  extraDeps = final: prev: let
+    rm = d: d.overrideAttrs (old: {
+      nativeBuildInputs = old.nativeBuildInputs ++ [ final.pythonRelaxDepsHook ];
+      pythonRemoveDeps = [ "opencv-python-headless" "opencv-python" "tb-nightly" "clip" ];
+    });
+    callPackage = final.callPackage;
+    rmCallPackage = path: args: rm (callPackage path args);
+  in {
+    pydeprecate = callPackage ../../packages/pydeprecate { };
+    taming-transformers-rom1504 =
+      callPackage ../../packages/taming-transformers-rom1504 { };
+    albumentations = rmCallPackage ../../packages/albumentations { opencv-python-headless = final.opencv4; };
+    qudida = rmCallPackage ../../packages/qudida { opencv-python-headless = final.opencv4; };
+    gfpgan = rmCallPackage ../../packages/gfpgan { opencv-python = final.opencv4; };
+    basicsr = rmCallPackage ../../packages/basicsr { opencv-python = final.opencv4; };
+    facexlib = rmCallPackage ../../packages/facexlib { opencv-python = final.opencv4; };
+    realesrgan = rmCallPackage ../../packages/realesrgan { opencv-python = final.opencv4; };
+    codeformer = callPackage ../../packages/codeformer { opencv-python = final.opencv4; };
+    clipseg = rmCallPackage ../../packages/clipseg { opencv-python = final.opencv4; };
+    filterpy = callPackage ../../packages/filterpy { };
+    kornia = callPackage ../../packages/kornia { };
+    lpips = callPackage ../../packages/lpips { };
+    ffmpy = callPackage ../../packages/ffmpy { };
+    shap = callPackage ../../packages/shap { };
+    picklescan = callPackage ../../packages/picklescan { };
+    diffusers = callPackage ../../packages/diffusers { };
+    pypatchmatch = callPackage ../../packages/pypatchmatch { };
+    fonts = callPackage ../../packages/fonts { };
+    font-roboto = callPackage ../../packages/font-roboto { };
+    analytics-python = callPackage ../../packages/analytics-python { };
+    markdown-it-py = callPackage ../../packages/markdown-it-py { };
+    gradio = callPackage ../../packages/gradio { };
+    hatch-requirements-txt = callPackage ../../packages/hatch-requirements-txt { };
+    timm = callPackage ../../packages/timm { };
+    blip = callPackage ../../packages/blip { };
+    fairscale = callPackage ../../packages/fairscale { };
+    torch-fidelity = callPackage ../../packages/torch-fidelity { };
+    resize-right = callPackage ../../packages/resize-right { };
+    torchdiffeq = callPackage ../../packages/torchdiffeq { };
+    k-diffusion = callPackage ../../packages/k-diffusion { clean-fid = final.clean-fid; };
+    accelerate = callPackage ../../packages/accelerate { };
+    clip-anytorch = callPackage ../../packages/clip-anytorch { };
+    jsonmerge = callPackage ../../packages/jsonmerge { };
+    clean-fid = callPackage ../../packages/clean-fid { };
+    getpass-asterisk = callPackage ../../packages/getpass-asterisk { };
+  };
+
+  torchRocm = final: prev: rec {
+    # TODO: figure out how to patch torch-bin trying to access /opt/amdgpu
+    # there might be an environment variable for it, can use a wrapper for that
+    # otherwise just grep the world for /opt/amdgpu or something and substituteInPlace the path
+    # you can run this thing without the fix by creating /opt and running nix build nixpkgs#libdrm --inputs-from . --out-link /opt/amdgpu
+    torch-bin = prev.torch-bin.overrideAttrs (old: {
+      src = pkgs.fetchurl {
+        name = "torch-1.13.1+rocm5.1.1-cp310-cp310-linux_x86_64.whl";
+        url = "https://download.pytorch.org/whl/rocm5.1.1/torch-1.13.1%2Brocm5.1.1-cp310-cp310-linux_x86_64.whl";
+        hash = "sha256-qUwAL3L9ODy9hjne8jZQRoG4BxvXXLT7cAy9RbM837A=";
+      };
+    });
+    torchvision-bin = prev.torchvision-bin.overrideAttrs (old: {
+      src = pkgs.fetchurl {
+        name = "torchvision-0.14.1+rocm5.1.1-cp310-cp310-linux_x86_64.whl";
+        url = "https://download.pytorch.org/whl/rocm5.1.1/torchvision-0.14.1%2Brocm5.1.1-cp310-cp310-linux_x86_64.whl";
+        hash = "sha256-8CM1QZ9cZfexa+HWhG4SfA/PTGB2475dxoOtGZ3Wa2E=";
+      };
+    });
+    torch = torch-bin;
+    torchvision = torchvision-bin;
+  };
+
+  torchCuda = final: prev: {
+    torch = final.torch-bin;
+    torchvision = final.torchvision-bin;
+  };
+}

modules/dependency-sets/default.nix

@@ -0,0 +1,15 @@
+{ lib, ... }:
+
+let
+  inherit (lib) mkOption types;
+in
+
+{
+  perSystem.options = {
+    dependencySets = mkOption {
+      description = "Specially instantiated dependency sets for customized builds";
+      type = with types; lazyAttrsOf unspecified;
+      default = {};
+    };
+  };
+}

projects/invokeai/default.nix

@@ -0,0 +1,70 @@
+{ inputs, lib, ... }:
+
+{
+  perSystem = { config, pkgs, ... }: let
+    inherit (config.dependencySets) aipython3-amd aipython3-nvidia;
+
+    getVersion = lib.flip lib.pipe [
+      (src: builtins.readFile "${src}/setup.py")
+      (builtins.match ".+VERSION = '([^']+)'.+")
+      builtins.head
+    ];
+
+    mkInvokeAIVariant = aipython3: aipython3.buildPythonPackage rec {
+      pname = "InvokeAI";
+      version = getVersion src;
+      src = inputs.invokeai-src;
+      propagatedBuildInputs = with aipython3; [
+        numpy
+        albumentations
+        opencv4
+        pudb
+        imageio
+        imageio-ffmpeg
+        pytorch-lightning
+        protobuf
+        omegaconf
+        test-tube
+        ((pkgs.streamlit.overrideAttrs (old: {
+          nativeBuildInputs = old.nativeBuildInputs ++ [ pythonRelaxDepsHook ];
+          pythonRelaxDeps = [ "protobuf" ];
+        })).override { protobuf3 = protobuf; })
+        einops
+        taming-transformers-rom1504
+        torch-fidelity
+        torchmetrics
+        transformers
+        kornia
+        k-diffusion
+        picklescan
+        diffusers
+        pypatchmatch
+        realesrgan
+        pillow
+        send2trash
+        flask
+        flask-socketio
+        flask-cors
+        dependency-injector
+        gfpgan
+        eventlet
+        clipseg
+        getpass-asterisk
+      ];
+      nativeBuildInputs = [ aipython3.pythonRelaxDepsHook ];
+      pythonRemoveDeps = [ "clip" "pyreadline3" "flaskwebgui" ];
+      pythonRelaxDeps = [ "protobuf" ];
+      postFixup = ''
+        chmod +x $out/bin/*
+        wrapPythonPrograms
+      '';
+      doCheck = false;
+    };
+
+  in {
+    packages = {
+      invokeai-amd = mkInvokeAIVariant aipython3-amd;
+      invokeai-nvidia = mkInvokeAIVariant aipython3-nvidia;
+    };
+  };
+}