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

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