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