diff --git a/flake.nix b/flake.nix index 8dee943..f672acc 100644 --- a/flake.nix +++ b/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, ... }: - let - pkgs = import inputs.nixpkgs { - inherit system; - 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; - }; - }; - }; - }); + imports = [ + ./modules/dependency-sets + ./modules/aipython3 + ./projects/invokeai + ]; + }; } diff --git a/modules/aipython3/default.nix b/modules/aipython3/default.nix new file mode 100644 index 0000000..ccb82fc --- /dev/null +++ b/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 + ]; + }; + }; +} diff --git a/modules/aipython3/overlays.nix b/modules/aipython3/overlays.nix new file mode 100644 index 0000000..dec97d3 --- /dev/null +++ b/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; + }; +} diff --git a/modules/dependency-sets/default.nix b/modules/dependency-sets/default.nix new file mode 100644 index 0000000..a1c650d --- /dev/null +++ b/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 = {}; + }; + }; +} diff --git a/projects/invokeai/default.nix b/projects/invokeai/default.nix new file mode 100644 index 0000000..9a5ab81 --- /dev/null +++ b/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; + }; + }; +}