pyrox.dev
1{
2 lib,
3 stdenv,
4 python,
5 buildPythonPackage,
6 fetchurl,
7 pythonOlder,
8 pythonAtLeast,
9
10 # buildInputs
11 cudaPackages,
12
13 # nativeBuildInputs
14 addDriverRunpath,
15 autoPatchelfHook,
16
17 # dependencies
18 pillow,
19 torch-bin,
20}:
21
22let
23 pyVerNoDot = builtins.replaceStrings [ "." ] [ "" ] python.pythonVersion;
24 srcs = import ./binary-hashes.nix version;
25 unsupported = throw "Unsupported system";
26 version = "0.23.0";
27in
28buildPythonPackage {
29 inherit version;
30
31 pname = "torchvision";
32
33 format = "wheel";
34
35 src = fetchurl srcs."${stdenv.system}-${pyVerNoDot}" or unsupported;
36
37 disabled = (pythonOlder "3.9") || (pythonAtLeast "3.14");
38
39 # Note that we don't rely on config.cudaSupport here, because the Linux wheels all come built with CUDA support.
40 buildInputs =
41 with cudaPackages;
42 lib.optionals stdenv.hostPlatform.isLinux [
43 # $out/${sitePackages}/torchvision/_C.so wants libcudart.so.11.0 but torchvision.libs only ships
44 # libcudart.$hash.so.11.0
45 cuda_cudart
46 ];
47
48 nativeBuildInputs = lib.optionals stdenv.hostPlatform.isLinux [
49 addDriverRunpath
50 autoPatchelfHook
51 ];
52
53 dependencies = [
54 pillow
55 torch-bin
56 ];
57
58 # The wheel-binary is not stripped to avoid the error of `ImportError: libtorch_cuda_cpp.so: ELF load command address/offset not properly aligned.`.
59 dontStrip = true;
60
61 pythonImportsCheck = [ "torchvision" ];
62
63 preInstall = lib.optionalString stdenv.hostPlatform.isLinux ''
64 addAutoPatchelfSearchPath "${torch-bin}/${python.sitePackages}/torch"
65 '';
66
67 meta = {
68 description = "PyTorch vision library";
69 homepage = "https://pytorch.org/";
70 changelog = "https://github.com/pytorch/vision/releases/tag/v${version}";
71 # Includes CUDA and Intel MKL, but redistributions of the binary are not limited.
72 # https://docs.nvidia.com/cuda/eula/index.html
73 # https://www.intel.com/content/www/us/en/developer/articles/license/onemkl-license-faq.html
74 license = lib.licenses.bsd3;
75 sourceProvenance = with lib.sourceTypes; [ binaryNativeCode ];
76 platforms = [
77 "aarch64-darwin"
78 "x86_64-linux"
79 "aarch64-linux"
80 ];
81 maintainers = with lib.maintainers; [
82 GaetanLepage
83 junjihashimoto
84 ];
85 };
86}