diff --git a/default.nix b/default.nix index 2a9c8e8..02e6e56 100644 --- a/default.nix +++ b/default.nix @@ -95,4 +95,10 @@ in }; getnative = callPythonPackage ./tools/getnative { }; + + pythonPackagesExtensions = prev.pythonPackagesExtensions ++ [ + (python-final: python-prev: { + styler00dollar-vsgan-trt = callPythonPackage ./tools/styler00dollar-vsgan-trt { }; + }) + ]; } diff --git a/tools/styler00dollar-vsgan-trt/default.nix b/tools/styler00dollar-vsgan-trt/default.nix new file mode 100644 index 0000000..b9cc966 --- /dev/null +++ b/tools/styler00dollar-vsgan-trt/default.nix @@ -0,0 +1,63 @@ +{ lib +, fetchFromGitHub +, buildPythonPackage +, python3 + # Needed to explicitly get the VapourSynth *Python* package. +, python3Packages +, numpy +, torch +}: + +let + pythonPackageName = "styler00dollar"; +in +buildPythonPackage rec { + pname = "styler00dollar-vsgan-trt"; + version = "unstable-2022-09-19"; + + src = fetchFromGitHub { + owner = "styler00dollar"; + repo = "VSGAN-tensorrt-docker"; + rev = "cf9c45cdacba9f19b25e26d0e4767bcfe5d48db0"; + sha256 = "sha256-L2HwE5r9OMiuupolJJbvPRIzWIRXc9uDSPC2WTZwWl8="; + }; + + format = "other"; + + propagatedBuildInputs = [ + numpy + torch + python3Packages.vapoursynth + ]; + + postPatch = '' + sed --in-place 's/src\./${pythonPackageName}./' convert_*_to_onnx.py + ''; + + installPhase = '' + runHook preInstall + + mkdir -p $out/${python3.sitePackages} + # "src" is the actual name of the subdirectory from upstream. + cp -r src $out/${python3.sitePackages}/${pythonPackageName} + + mkdir -p $out/examples + cp convert_*_to_onnx.py $out/examples + + runHook postInstall + ''; + + meta = with lib; { + description = "VSGAN utility modules from styler00dollar"; + longDescription = '' + Various utility modules by styler00dollar for working with TensorRT. + + This can be used to convert ESRGAN models to to ONXX models, which is + required if you want to use ESRGAN models with TensorRT. + ''; + homepage = "https://github.com/styler00dollar/VSGAN-tensorrt-docker"; + license = licenses.bsd3; + maintainers = with maintainers; [ aidalgol ]; + platforms = platforms.all; + }; +}