pkgs/development/python-modules/k-diffusion/default.nix

   1 {
   2   lib,
   3   accelerate,
   4   buildPythonPackage,
   5   clean-fid,
   6   clip-anytorch,
   7   dctorch,
   8   einops,
   9   fetchFromGitHub,
  10   jsonmerge,
  11   kornia,
  12   pillow,
  13   pythonOlder,
  14   rotary-embedding-torch,
  15   safetensors,
  16   scikit-image,
  17   scipy,
  18   torch,
  19   torchdiffeq,
  20   torchsde,
  21   torchvision,
  22   tqdm,
  23   wandb,
  24 }:
  25
  26 buildPythonPackage rec {
  27   pname = "k-diffusion";
  28   version = "0.1.1.post1";
  29   format = "setuptools";
  30
  31   disabled = pythonOlder "3.7";
  32
  33   src = fetchFromGitHub {
  34     owner = "crowsonkb";
  35     repo = "k-diffusion";
  36     rev = "refs/tags/v${version}";
  37     hash = "sha256-x/UHzobQv5ov0luUHqC8OA5YbtF+aWL39/SQtzTm0RM=";
  38   };
  39
  40   propagatedBuildInputs = [
  41     accelerate
  42     clean-fid
  43     clip-anytorch
  44     dctorch
  45     einops
  46     jsonmerge
  47     kornia
  48     pillow
  49     rotary-embedding-torch
  50     scikit-image
  51     scipy
  52     safetensors
  53     torch
  54     torchdiffeq
  55     torchsde
  56     torchvision
  57     tqdm
  58     wandb
  59   ];
  60
  61   pythonImportsCheck = [ "k_diffusion" ];
  62
  63   # no tests
  64   doCheck = false;
  65
  66   meta = with lib; {
  67     description = "Karras et al. (2022) diffusion models for PyTorch";
  68     homepage = "https://github.com/crowsonkb/k-diffusion";
  69     license = licenses.mit;
  70     maintainers = teams.tts.members;
  71   };
  72 }