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