pyrox.dev
1{
2 lib,
3 symlinkJoin,
4 buildPythonPackage,
5 fetchFromGitHub,
6 fetchpatch,
7
8 # nativeBuildInputs
9 cmake,
10 pkg-config,
11 ninja,
12
13 # buildInputs
14 ffmpeg_6-full,
15 pybind11,
16 sox,
17 torch,
18 llvmPackages,
19
20 cudaSupport ? torch.cudaSupport,
21 cudaPackages,
22 rocmSupport ? torch.rocmSupport,
23 rocmPackages,
24
25 gpuTargets ? [ ],
26}:
27
28let
29 # TODO: Reuse one defined in torch?
30 # Some of those dependencies are probably not required,
31 # but it breaks when the store path is different between torch and torchaudio
32 rocmtoolkit_joined = symlinkJoin {
33 name = "rocm-merged";
34
35 paths = with rocmPackages; [
36 rocm-core
37 clr
38 rccl
39 miopen
40 rocrand
41 rocblas
42 rocsparse
43 hipsparse
44 rocthrust
45 rocprim
46 hipcub
47 roctracer
48 rocfft
49 rocsolver
50 hipfft
51 hipsolver
52 hipblas-common
53 hipblas
54 rocminfo
55 rocm-comgr
56 rocm-device-libs
57 rocm-runtime
58 clr.icd
59 hipify
60 ];
61
62 # Fix `setuptools` not being found
63 postBuild = ''
64 rm -rf $out/nix-support
65 '';
66 };
67 # Only used for ROCm
68 gpuTargetString = lib.strings.concatStringsSep ";" (
69 if gpuTargets != [ ] then
70 # If gpuTargets is specified, it always takes priority.
71 gpuTargets
72 else if rocmSupport then
73 rocmPackages.clr.gpuTargets
74 else
75 throw "No GPU targets specified"
76 );
77in
78buildPythonPackage rec {
79 pname = "torchaudio";
80 version = "2.8.0";
81 pyproject = true;
82
83 stdenv = torch.stdenv;
84
85 src = fetchFromGitHub {
86 owner = "pytorch";
87 repo = "audio";
88 tag = "v${version}";
89 hash = "sha256-SPa6ZWA2AWawfL4Z4mb1nddGaAsGEl/0dwweBpex2Wo=";
90 };
91
92 patches = [
93 ./0001-setup.py-propagate-cmakeFlags.patch
94 ];
95
96 postPatch = lib.optionalString rocmSupport ''
97 # There is no .info/version-dev, only .info/version
98 substituteInPlace cmake/LoadHIP.cmake \
99 --replace-fail "/.info/version-dev" "/.info/version"
100 '';
101
102 env = {
103 TORCH_CUDA_ARCH_LIST = "${lib.concatStringsSep ";" torch.cudaCapabilities}";
104 };
105
106 # https://github.com/pytorch/audio/blob/v2.1.0/docs/source/build.linux.rst#optional-build-torchaudio-with-a-custom-built-ffmpeg
107 FFMPEG_ROOT = symlinkJoin {
108 name = "ffmpeg";
109 paths = [
110 ffmpeg_6-full.bin
111 ffmpeg_6-full.dev
112 ffmpeg_6-full.lib
113 ];
114 };
115
116 nativeBuildInputs = [
117 cmake
118 pkg-config
119 ninja
120 ]
121 ++ lib.optionals cudaSupport [ cudaPackages.cuda_nvcc ]
122 ++ lib.optionals rocmSupport (
123 with rocmPackages;
124 [
125 clr
126 rocblas
127 hipblas
128 ]
129 );
130
131 buildInputs = [
132 ffmpeg_6-full
133 pybind11
134 sox
135 torch.cxxdev
136 ]
137 ++ lib.optionals stdenv.cc.isClang [ llvmPackages.openmp ];
138
139 dependencies = [ torch ];
140
141 BUILD_SOX = 0;
142 BUILD_KALDI = 0;
143 BUILD_RNNT = 0;
144 BUILD_CTC_DECODER = 0;
145
146 preConfigure = lib.optionalString rocmSupport ''
147 export ROCM_PATH=${rocmtoolkit_joined}
148 export PYTORCH_ROCM_ARCH="${gpuTargetString}"
149 '';
150
151 dontUseCmakeConfigure = true;
152
153 doCheck = false; # requires sox backend
154
155 pythonImportsCheck = [ "torchaudio" ];
156
157 meta = {
158 description = "PyTorch audio library";
159 homepage = "https://pytorch.org/";
160 changelog = "https://github.com/pytorch/audio/releases/tag/v${version}";
161 license = lib.licenses.bsd2;
162 platforms =
163 lib.platforms.linux ++ lib.optionals (!cudaSupport && !rocmSupport) lib.platforms.darwin;
164 maintainers = with lib.maintainers; [
165 GaetanLepage
166 junjihashimoto
167 ];
168 };
169}