pyrox.dev
1{
2 lib,
3 stdenv,
4 buildPythonPackage,
5 braceexpand,
6 ftfy,
7 huggingface-hub,
8 pandas,
9 protobuf,
10 pytestCheckHook,
11 regex,
12 sentencepiece,
13 timm,
14 torch,
15 torchvision,
16 tqdm,
17 transformers,
18 setuptools,
19 webdataset,
20 wheel,
21 fetchFromGitHub,
22}:
23buildPythonPackage rec {
24 pname = "open-clip-torch";
25 version = "2.24.0";
26 pyproject = true;
27
28 src = fetchFromGitHub {
29 owner = "mlfoundations";
30 repo = "open_clip";
31 rev = "refs/tags/v${version}";
32 hash = "sha256-ugbXnXiOY9FrNvr8ZxnAgZO/SLCVoXbRgupi8cUwflU=";
33 };
34
35 nativeBuildInputs = [
36 setuptools
37 wheel
38 ];
39
40 propagatedBuildInputs = [
41 ftfy
42 huggingface-hub
43 protobuf
44 regex
45 sentencepiece
46 timm
47 torch
48 torchvision
49 tqdm
50 ];
51
52 nativeCheckInputs = [
53 pytestCheckHook
54 braceexpand
55 pandas
56 transformers
57 webdataset
58 ];
59
60 pythonImportsCheck = [ "open_clip" ];
61
62 disabledTestPaths = lib.optionals (stdenv.isAarch64 || stdenv.isDarwin) [ "tests/test_wds.py" ];
63
64 disabledTests =
65 [
66 # requires network
67 "test_download_pretrained_from_hfh"
68 "test_inference_simple"
69 "test_inference_with_data"
70 "test_pretrained_text_encoder"
71 "test_training_mt5"
72 # fails due to type errors
73 "test_num_shards"
74 ]
75 ++ lib.optionals (stdenv.isAarch64 && stdenv.isLinux) [
76 "test_training"
77 "test_training_coca"
78 "test_training_unfreezing_vit"
79 "test_training_clip_with_jit"
80 ];
81
82 meta = with lib; {
83 description = "An open source implementation of CLIP";
84 homepage = "https://github.com/mlfoundations/open_clip";
85 license = licenses.asl20;
86 maintainers = with maintainers; [ iynaix ];
87 mainProgram = "open-clip";
88 };
89}