{ description = "Nix Flake for runnig Stable Diffusion on NixOS"; inputs = { nixlib.url = "github:nix-community/nixpkgs.lib"; nixpkgs = { url = "github:NixOS/nixpkgs"; #?rev=33919d25f0c873b0c73e2f8d0859fab3bd0d1b26"; }; stable-diffusion-repo = { url = "github:Stability-AI/stablediffusion?rev=47b6b607fdd31875c9279cd2f4f16b92e4ea958e"; flake = false; }; invokeai-repo = { url = "github:invoke-ai/InvokeAI?ref=v2.3.5.post2"; flake = false; }; webui-repo = { #url = "github:AUTOMATIC1111/stable-diffusion-webui"; url = "github:gbtb/stable-diffusion-webui"; #instead of patching in flake, better use a fork. CRLF bullshit makes patching this repo almost impossible flake = false; }; }; outputs = { self, nixpkgs, nixlib, stable-diffusion-repo, invokeai-repo, webui-repo }@inputs: let nixlib = inputs.nixlib.outputs.lib; system = "x86_64-linux"; requirementsFor = { pkgs, webui ? false, nvidia ? false }: with pkgs; with pkgs.python3.pkgs; [ python3 torch torchvision numpy albumentations opencv4 pudb imageio imageio-ffmpeg pytorch-lightning omegaconf test-tube streamlit protobuf einops taming-transformers-rom1504 torch-fidelity torchmetrics transformers kornia k-diffusion diffusers # following packages not needed for vanilla SD but used by both UIs realesrgan pillow safetensors fastapi ] ++ nixlib.optional (nvidia) [ xformers ] #probably won't fully work ++ nixlib.optional (!webui) [ npyscreen huggingface-hub dnspython datasets click pypatchmatch torchsde compel send2trash flask flask-socketio flask-cors gfpgan eventlet clipseg getpass-asterisk picklescan peft packaging python-multipart fastapi-socketio fastapi-events ] ++ nixlib.optional webui [ pip addict future lmdb pyyaml scikitimage tqdm yapf gdown lpips lark analytics-python ffmpy markdown-it-py shap gradio fonts font-roboto piexif codeformer blip psutil openclip blendmodes ]; overlay_default = nixpkgs: pythonPackages: { pytorch-lightning = pythonPackages.pytorch-lightning.overrideAttrs (old: { nativeBuildInputs = old.nativeBuildInputs ++ [ nixpkgs.python3Packages.pythonRelaxDepsHook ]; pythonRelaxDeps = [ "protobuf" ]; }); scikit-image = pythonPackages.scikitimage; moto = pythonPackages.moto.overrideAttrs (old: { doCheck = false; dontUsePytestCheck = true; }); }; overlay_webui = nixpkgs: pythonPackages: { transformers = pythonPackages.transformers.overrideAttrs (old: { src = nixpkgs.fetchFromGitHub { owner = "huggingface"; repo = "transformers"; rev = "refs/tags/v4.19.2"; hash = "sha256-9r/1vW7Rhv9+Swxdzu5PTnlQlT8ofJeZamHf5X4ql8w="; }; }); }; overlay_invoke = nixpkgs: pythonPackages: let ifNotMinVersion = pkg: ver: overlay: if ( nixlib.versionOlder pkg.version ver ) then pkg.overrideAttrs overlay else pkg; in { huggingface-hub = ifNotMinVersion pythonPackages.huggingface-hub "0.13.2" ( old: rec { version = "0.14.1"; src = nixpkgs.fetchFromGitHub { owner = "huggingface"; repo = "huggingface_hub"; rev = "refs/tags/v${version}"; hash = "sha256-+BtXi+O+Ef4p4b+8FJCrZFsxX22ZYOPXylexFtsldnA="; }; propagatedBuildInputs = old.propagatedBuildInputs ++ [ pythonPackages.fsspec ]; } ); transformers = ifNotMinVersion pythonPackages.transformers "4.26" ( old: rec { version = "4.28.1"; src = nixpkgs.fetchFromGitHub { inherit (old.src) owner repo; rev = "refs/tags/v${version}"; hash = "sha256-FmiuWfoFZjZf1/GbE6PmSkeshWWh+6nDj2u2PMSeDk0="; }; } ); }; overlay_pynixify = self: let rm = d: d.overrideAttrs (old: { nativeBuildInputs = old.nativeBuildInputs ++ [ self.pythonRelaxDepsHook ]; pythonRemoveDeps = [ "opencv-python-headless" "opencv-python" "tb-nightly" "clip" ]; }); callPackage = self.callPackage; rmCallPackage = path: args: rm (callPackage path args); mapCallPackage = pnames: builtins.listToAttrs (builtins.map (pname: { name = pname; value = (callPackage (./packages + "/${pname}") { }); }) pnames); simplePackages = [ "filterpy" "kornia" "lpips" "ffmpy" "shap" "fonts" "font-roboto" "analytics-python" "markdown-it-py" "gradio" "hatch-requirements-txt" "timm" "blip" "fairscale" "torch-fidelity" "resize-right" "torchdiffeq" "accelerate" "clip-anytorch" "jsonmerge" "clean-fid" "getpass-asterisk" "pypatchmatch" "trampoline" "torchsde" "compel" "diffusers" "safetensors" "picklescan" "openclip" "blendmodes" "xformers" "pyre-extensions" # "triton" TODO: nixpkgs is missing required llvm parts - mlir. https://github.com/NixOS/nixpkgs/pull/163878 "peft" "fastapi-events" "fastapi-socketio" ]; in { pydeprecate = callPackage ./packages/pydeprecate { }; taming-transformers-rom1504 = callPackage ./packages/taming-transformers-rom1504 { }; albumentations = rmCallPackage ./packages/albumentations { opencv-python-headless = self.opencv4; }; qudida = rmCallPackage ./packages/qudida { opencv-python-headless = self.opencv4; }; gfpgan = rmCallPackage ./packages/gfpgan { opencv-python = self.opencv4; }; basicsr = rmCallPackage ./packages/basicsr { opencv-python = self.opencv4; }; facexlib = rmCallPackage ./packages/facexlib { opencv-python = self.opencv4; }; codeformer = callPackage ./packages/codeformer { opencv-python = self.opencv4; }; realesrgan = rmCallPackage ./packages/realesrgan { opencv-python = self.opencv4; }; clipseg = rmCallPackage ./packages/clipseg { opencv-python = self.opencv4; }; k-diffusion = callPackage ./packages/k-diffusion { clean-fid = self.clean-fid; }; } // mapCallPackage simplePackages; overlay_amd = nixpkgs: pythonPackages: rec { #IMPORTANT: you can browse available wheels on the server, but only if you add trailing "/" - e.g. https://download.pytorch.org/whl/rocm5.2/ torch-bin = pythonPackages.torch-bin.overrideAttrs (old: { src = nixpkgs.fetchurl { name = "torch-1.13.1+rocm5.2-cp310-cp310-linux_x86_64.whl"; url = "https://download.pytorch.org/whl/rocm5.2/torch-1.13.1%2Brocm5.2-cp310-cp310-linux_x86_64.whl"; hash = "sha256-82hdCKwNjJUcw2f5vUsskkxdRRdmnEdoB3SKvNlmE28="; }; }); torchvision-bin = pythonPackages.torchvision-bin.overrideAttrs (old: { src = nixpkgs.fetchurl { name = "torchvision-0.14.1+rocm5.2-cp310-cp310-linux_x86_64.whl"; url = "https://download.pytorch.org/whl/rocm5.2/torchvision-0.14.1%2Brocm5.2-cp310-cp310-linux_x86_64.whl"; hash = "sha256-oBYG/K7bgkxu0UvmyS2U1ud2LkFQ/CarcxpEJ9xzMYQ="; }; }); torch = torch-bin; torchvision = torchvision-bin; }; overlay_nvidia = nixpkgs: pythonPackages: { torch = pythonPackages.torch-bin; torchvision = pythonPackages.torchvision-bin; }; in let nixpkgs_ = { amd ? false, nvidia ? false, webui ? false }: import inputs.nixpkgs { inherit system; config.allowUnfree = nvidia; #CUDA is unfree. overlays = [ (final: prev: let optional = nixlib.optionalAttrs; sl = (prev.streamlit.override ({ protobuf3 = prev.protobuf; })); makePythonHook = args: final.makeSetupHook ({ passthru.provides.setupHook = true; } // args); pythonRelaxDepsHook = prev.callPackage ({ wheel }: #upstream hook doesn't work properly with non-standard wheel names #which means that some packages from pip silently fail to be overriden #https://github.com/NixOS/nixpkgs/issues/198342 makePythonHook { name = "python-relax-deps-hook"; propagatedBuildInputs = [ wheel ]; substitutions = { pythonInterpreter = nixlib.getExe prev.python3Packages.python; }; } ./python-relax-deps-hook.sh) { wheel = prev.python3.pkgs.wheel; }; in { inherit pythonRelaxDepsHook; streamlit = sl.overrideAttrs (old: { nativeBuildInputs = old.nativeBuildInputs ++ [ pythonRelaxDepsHook ]; pythonRemoveDeps = [ "protobuf" ]; }); python3 = prev.python3.override { packageOverrides = python-self: python-super: (overlay_default prev python-super) // optional amd (overlay_amd prev python-super) // optional nvidia (overlay_nvidia prev python-super) // optional webui (overlay_webui prev python-super) // optional (!webui) (overlay_invoke prev python-super) // (overlay_pynixify python-self); }; }) ]; } // { inherit nvidia; }; in { packages.${system} = let nixpkgsAmd = (nixpkgs_ { amd = true; }); nixpkgsNvidia = (nixpkgs_ { nvidia = true; }); invokeaiF = nixpkgs: nixpkgs.python3.pkgs.buildPythonApplication { pname = "invokeai"; version = "2.3.5"; src = invokeai-repo; format = "pyproject"; meta.mainProgram = "invokeai"; propagatedBuildInputs = requirementsFor { pkgs = nixpkgs; nvidia = nixpkgs.nvidia; }; nativeBuildInputs = [ nixpkgs.pkgs.pythonRelaxDepsHook ]; pythonRelaxDeps = [ "torch" "pytorch-lightning" "flask-socketio" "flask" "dnspython" "fastapi" ]; pythonRemoveDeps = [ "opencv-python" "flaskwebgui" "pyreadline3" ]; postPatch = '' # Add subprocess to the imports substituteInPlace ./ldm/invoke/config/invokeai_configure.py --replace \ 'import shutil' \ ' import shutil import subprocess ' # shutil.copytree will inherit the permissions of files in the /nix/store # which are read only, so we subprocess.call cp instead and tell it not to # preserve the mode substituteInPlace ./ldm/invoke/config/invokeai_configure.py --replace \ "shutil.copytree(configs_src, configs_dest, dirs_exist_ok=True)" \ "subprocess.call(f'${nixpkgs.pkgs.coreutils}/bin/cp -r --no-preserve=mode {configs_src}/* {configs_dest}', shell=True)" ''; }; webuiF = nixpkgs: let submodel = pkg: nixpkgs.pkgs.python3.pkgs.${pkg} + "/lib/python3.10/site-packages"; taming-transformers = submodel "taming-transformers-rom1504"; k_diffusion = submodel "k-diffusion"; codeformer = (submodel "codeformer") + "/codeformer"; blip = (submodel "blip") + "/blip"; in nixpkgs.python3.pkgs.buildPythonApplication { pname = "stable-diffusion-webui"; version = "2023-03-12"; src = webui-repo; format = "other"; propagatedBuildInputs = requirementsFor { pkgs = nixpkgs; webui = true; nvidia = nixpkgs.nvidia; }; nativeBuildInputs = [ nixpkgs.pkgs.makeWrapper ]; meta.mainProgram = "flake-launch"; buildPhase = '' runHook preBuild cp -r . $out chmod -R +w $out cd $out #firstly, we need to make launch.py runnable by adding python shebang cat <<-EOF > exec_launch.py.unwrapped $(echo "#!/usr/bin/python") $(cat launch.py) EOF chmod +x exec_launch.py.unwrapped #creating wrapper around launch.py with PYTHONPATH correctly set makeWrapper "$(pwd)/exec_launch.py.unwrapped" exec_launch.py \ --set-default PYTHONPATH $PYTHONPATH mkdir $out/bin pushd $out/bin ln -s ../exec_launch.py launch.py buck='$' #escaping $ inside shell inside shell is tricky #next is an additional shell wrapper, which sets sensible default args for CLI #additional arguments will be passed further cat <<-EOF > flake-launch #!/usr/bin/env bash pushd $out #For some reason, fastapi only works when current workdir is set inside the repo trap "popd" EXIT "$out/bin/launch.py" --skip-install "$buck{@}" EOF # below lie remnants of my attempt to make webui use similar paths as InvokeAI for models download # additions of such options in upstream is a welcome sign, however they're mostly ignored and therefore useless # TODO: check in 6 months, maybe it'll work # For now, your best bet is to use ZFS dataset with dedup enabled or make symlinks after the fact #--codeformer-models-path "\$mp/codeformer" \ #--gfpgan-models-path "\$mp/gfpgan" --esrgan-models-path "\$mp/esrgan" \ #--bsrgan-models-path "\$mp/bsrgan" --realesrgan-models-path "\$mp/realesrgan" \ #--clip-models-path "\$mp/clip" chmod +x flake-launch popd runHook postBuild ''; installPhase = '' runHook preInstall rm -rf repositories/ mkdir repositories pushd repositories ln -s ${inputs.stable-diffusion-repo}/ stable-diffusion-stability-ai ln -s ${taming-transformers}/ taming-transformers ln -s ${k_diffusion}/ k-diffusion ln -s ${codeformer}/ CodeFormer ln -s ${blip}/ BLIP popd runHook postInstall ''; }; in { invokeai = { amd = invokeaiF nixpkgsAmd; default = invokeaiF nixpkgsNvidia; }; webui = { amd = webuiF nixpkgsAmd; default = webuiF nixpkgsNvidia; }; }; }; }