nix-stable-diffusion/flake.nix
2023-03-19 11:08:36 +10:00

356 lines
14 KiB
Nix

{
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.1.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
]
++ 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
]
++ nixlib.optional webui [
pip
addict
future
lmdb
pyyaml
scikitimage
tqdm
yapf
gdown
lpips
fastapi
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;
};
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_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
];
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) //
(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.1";
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" ];
pythonRemoveDeps = [ "opencv-python" "flaskwebgui" "pyreadline3" ];
};
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;
};
};
};
}