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