pyrox.dev
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
26buildPythonPackage 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 tag = "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 teams = [ teams.tts ];
71 };
72}