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