diff --git a/pkgs/development/python-modules/k-diffusion/default.nix b/pkgs/development/python-modules/k-diffusion/default.nix new file mode 100644 index 00000000000..0d1410a5ced --- /dev/null +++ b/pkgs/development/python-modules/k-diffusion/default.nix @@ -0,0 +1,69 @@ +{ lib +, buildPythonPackage +, fetchFromGitHub + +# dependencies +, accelerate +, clean-fid +, clip-anytorch +, einops +, jsonmerge +, kornia +, pillow +, resize-right +, scikit-image +, scipy +, torch +, torchdiffeq +, torchsde +, torchvision +, tqdm +, wandb + +}: + +buildPythonPackage rec { + pname = "k-diffusion"; + version = "0.0.14"; + format = "setuptools"; + + src = fetchFromGitHub { + owner = "crowsonkb"; + repo = "k-diffusion"; + rev = "v${version}"; + hash = "sha256-KKVgk+1hidDBVaRnXjoqwuSRydI10OPHK3YModAizZU="; + }; + + propagatedBuildInputs = [ + accelerate + clean-fid + clip-anytorch + einops + jsonmerge + kornia + pillow + resize-right + scikit-image + scipy + torch + torchdiffeq + torchsde + torchvision + tqdm + wandb + ]; + + pythonImportsCheck = [ + "k_diffusion" + ]; + + # no tests + doCheck = false; + + meta = with lib; { + description = "Karras et al. (2022) diffusion models for PyTorch"; + homepage = "https://github.com/crowsonkb/k-diffusion"; + license = licenses.mit; + maintainers = teams.tts.members; + }; +} diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix index ec38b8a1d62..e4de921690d 100644 --- a/pkgs/top-level/python-packages.nix +++ b/pkgs/top-level/python-packages.nix @@ -5324,6 +5324,8 @@ self: super: with self; { jxmlease = callPackage ../development/python-modules/jxmlease { }; + k-diffusion = callPackage ../development/python-modules/k-diffusion { }; + k5test = callPackage ../development/python-modules/k5test { inherit (pkgs) krb5 findutils; };