pyrox.dev
1{
2 lib,
3 buildPythonPackage,
4 fetchFromGitHub,
5
6 # build-system
7 pdm-backend,
8
9 # dependencies
10 ftfy,
11 huggingface-hub,
12 protobuf,
13 regex,
14 safetensors,
15 sentencepiece,
16 timm,
17 torch,
18 torchvision,
19 tqdm,
20
21 # checks
22 pytestCheckHook,
23 braceexpand,
24 pandas,
25 transformers,
26 webdataset,
27
28 stdenv,
29}:
30buildPythonPackage rec {
31 pname = "open-clip-torch";
32 version = "2.32.0";
33 pyproject = true;
34
35 src = fetchFromGitHub {
36 owner = "mlfoundations";
37 repo = "open_clip";
38 tag = "v${version}";
39 hash = "sha256-HXzorEAVPieCHfW3xzXqNTTIzJSbIuaZhcfcp0htdCk=";
40 };
41
42 build-system = [ pdm-backend ];
43
44 dependencies = [
45 ftfy
46 huggingface-hub
47 protobuf
48 regex
49 safetensors
50 sentencepiece
51 timm
52 torch
53 torchvision
54 tqdm
55 ];
56
57 nativeCheckInputs = [
58 pytestCheckHook
59 braceexpand
60 pandas
61 transformers
62 webdataset
63 ];
64
65 pythonImportsCheck = [ "open_clip" ];
66
67 # -> On Darwin:
68 # AttributeError: Can't pickle local object 'build_params.<locals>.<lambda>'
69 # -> On Linux:
70 # KeyError: Caught KeyError in DataLoader worker process 0
71 disabledTestPaths = [ "tests/test_wds.py" ];
72
73 disabledTests =
74 [
75 # requires network
76 "test_download_pretrained_from_hfh"
77 "test_inference_simple"
78 "test_inference_with_data"
79 "test_pretrained_text_encoder"
80 "test_training_mt5"
81 # fails due to type errors
82 "test_num_shards"
83 ]
84 ++ lib.optionals (stdenv.hostPlatform.isAarch64 && stdenv.hostPlatform.isLinux) [
85 "test_training"
86 "test_training_coca"
87 "test_training_unfreezing_vit"
88 "test_training_clip_with_jit"
89 ];
90
91 meta = {
92 description = "Open source implementation of CLIP";
93 homepage = "https://github.com/mlfoundations/open_clip";
94 changelog = "https://github.com/mlfoundations/open_clip/releases/tag/${src.tag}";
95 license = lib.licenses.asl20;
96 maintainers = with lib.maintainers; [ iynaix ];
97 mainProgram = "open-clip";
98 };
99}