tjh.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 # requires network
75 "test_download_pretrained_from_hfh"
76 "test_inference_simple"
77 "test_inference_with_data"
78 "test_pretrained_text_encoder"
79 "test_training_mt5"
80 # fails due to type errors
81 "test_num_shards"
82 ]
83 ++ lib.optionals (stdenv.hostPlatform.isAarch64 && stdenv.hostPlatform.isLinux) [
84 "test_training"
85 "test_training_coca"
86 "test_training_unfreezing_vit"
87 "test_training_clip_with_jit"
88 ];
89
90 meta = {
91 description = "Open source implementation of CLIP";
92 homepage = "https://github.com/mlfoundations/open_clip";
93 changelog = "https://github.com/mlfoundations/open_clip/releases/tag/${src.tag}";
94 license = lib.licenses.asl20;
95 maintainers = with lib.maintainers; [ iynaix ];
96 mainProgram = "open-clip";
97 };
98}