/
/
/
1# mypy: ignore-errors
2# ruff: noqa
3"""Vendored S-KEY model for musical key detection.
4
5Source: https://github.com/deezer/skey (MIT License)
6Paper: S-KEY (ICASSP 2025), STONE (ISMIR 2024)
7ConvNeXt blocks originally from Meta FAIR (https://github.com/facebookresearch/ConvNeXt)
8"""
9
10from __future__ import annotations
11
12import math
13from pathlib import Path
14from typing import Any
15
16import torch
17import torchaudio
18from einops import rearrange
19from nnAudio.features.vqt import VQT as _VQT
20from torch import nn
21
22KEY_MAP: dict[int, str] = {
23 0: "A Major",
24 1: "Bb Major",
25 2: "B Major",
26 3: "C Major",
27 4: "C# Major",
28 5: "D Major",
29 6: "D# Major",
30 7: "E Major",
31 8: "F Major",
32 9: "F# Major",
33 10: "G Major",
34 11: "G# Major",
35 12: "B minor",
36 13: "C minor",
37 14: "C# minor",
38 15: "D minor",
39 16: "D# minor",
40 17: "E minor",
41 18: "F minor",
42 19: "F# minor",
43 20: "G minor",
44 21: "G# minor",
45 22: "A minor",
46 23: "Bb minor",
47}
48
49_CHECKPOINT_PATH = Path(__file__).parent / "skey.pt"
50
51
52# ---------------------------------------------------------------------------
53# ConvNeXt components
54# ---------------------------------------------------------------------------
55
56
57class DropPath(nn.Module):
58 """Drop paths (Stochastic Depth) per sample."""
59
60 def __init__(self, drop_prob: float = 0.0, scale_by_keep: bool = True) -> None:
61 super().__init__()
62 self.drop_prob = drop_prob
63 self.scale_by_keep = scale_by_keep
64
65 def forward(self, x: torch.Tensor) -> torch.Tensor:
66 """Apply stochastic depth."""
67 if self.drop_prob == 0.0 or not self.training:
68 return x
69 keep_prob = 1 - self.drop_prob
70 shape = (x.shape[0],) + (1,) * (x.ndim - 1)
71 random_tensor = x.new_empty(shape).bernoulli_(keep_prob)
72 if keep_prob > 0.0 and self.scale_by_keep:
73 random_tensor.div_(keep_prob)
74 return x * random_tensor
75
76
77class ConvNeXtBlock(nn.Module):
78 """ConvNeXt Block: DwConv -> Permute -> LayerNorm -> Linear -> GELU -> Linear -> Permute."""
79
80 def __init__(
81 self,
82 in_channels: int,
83 out_channels: int,
84 kernel_size: int = 7,
85 padding: int = 3,
86 drop_path: float = 0.1,
87 layer_scale_init_value: float = 1e-1,
88 ) -> None:
89 super().__init__()
90 self.dwconv = nn.Conv2d(
91 in_channels,
92 out_channels,
93 kernel_size=kernel_size,
94 padding=padding,
95 groups=in_channels,
96 padding_mode="replicate",
97 )
98 self.norm = nn.functional.layer_norm
99 self.pwconv1 = nn.Linear(out_channels, 4 * out_channels)
100 self.act = nn.GELU()
101 self.pwconv2 = nn.Linear(4 * out_channels, in_channels)
102 self.gamma = (
103 nn.Parameter(layer_scale_init_value * torch.ones(out_channels), requires_grad=True)
104 if layer_scale_init_value > 0
105 else None
106 )
107 self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
108
109 def forward(self, x: torch.Tensor) -> torch.Tensor:
110 """Forward pass."""
111 input_x = x
112 x = self.dwconv(x)
113 x = self.norm(x, x.shape[1:])
114 x = x.permute(0, 2, 3, 1)
115 x = self.pwconv1(x)
116 x = self.act(x)
117 x = self.pwconv2(x)
118 if self.gamma is not None:
119 x = self.gamma * x
120 x = x.permute(0, 3, 1, 2)
121 return input_x + self.drop_path(x)
122
123
124class TimeDownsamplingBlock(nn.Module):
125 """Time Downsampling Block: LayerNorm -> 1x2 strided Conv -> GELU."""
126
127 def __init__(self, in_channels: int, out_channels: int, bias: bool = True) -> None:
128 super().__init__()
129 self.norm = nn.functional.layer_norm
130 self.conv = nn.Conv2d(
131 in_channels, out_channels, kernel_size=(1, 2), stride=(1, 2), bias=bias
132 )
133 self.act = nn.GELU()
134
135 def forward(self, x: torch.Tensor) -> torch.Tensor:
136 """Forward pass."""
137 x = self.norm(x, x.shape[1:])
138 x = self.conv(x)
139 return self.act(x)
140
141
142# ---------------------------------------------------------------------------
143# Octave pooling and ChromaNet
144# ---------------------------------------------------------------------------
145
146
147class OctavePool(nn.Module):
148 """Average log-frequency axis across octaves, producing a chromagram."""
149
150 def __init__(self, bins_per_octave: int) -> None:
151 super().__init__()
152 self.bins_per_octave = bins_per_octave
153
154 def forward(self, x: torch.Tensor) -> torch.Tensor:
155 """Forward pass."""
156 x = rearrange(x, "B C (j k) W -> B C k j W", k=self.bins_per_octave)
157 return x.mean(dim=3)
158
159
160class ChromaNet(nn.Module):
161 """ChromaNet neural network from STONE (ISMIR 2024)."""
162
163 def __init__(
164 self,
165 n_bins: int,
166 n_harmonics: int,
167 out_channels: list[int],
168 kernels: list[int],
169 temperature: float,
170 ) -> None:
171 super().__init__()
172 assert len(kernels) == len(out_channels)
173 self.n_harmonics = n_harmonics
174 self.n_bins = n_bins
175 in_channel = self.n_harmonics
176 self.out_channels = out_channels
177 self.kernels = kernels
178 self.temperature = temperature
179 self.drop_path = 0.1
180 time_downsampling_blocks = []
181 convnext_blocks = []
182 for i, out_channel in enumerate(self.out_channels):
183 time_downsampling_blocks.append(TimeDownsamplingBlock(in_channel, out_channel))
184 kernel = self.kernels[i]
185 convnext_blocks.append(
186 ConvNeXtBlock(
187 out_channel,
188 out_channel,
189 kernel_size=kernel,
190 padding=kernel // 2,
191 drop_path=self.drop_path,
192 )
193 )
194 in_channel = out_channel
195 self.convnext_blocks = nn.ModuleList(convnext_blocks)
196 self.time_downsampling_blocks = nn.ModuleList(time_downsampling_blocks)
197 self.octave_pool = OctavePool(12)
198 self.global_average_pool = nn.AdaptiveAvgPool2d((12, 1))
199 self.classifier = nn.Conv2d(out_channel, 2, kernel_size=(1, 1))
200 self.flatten = nn.Flatten()
201 self.batch_norm = nn.BatchNorm2d(2, affine=False)
202 self.softmax = nn.Softmax(dim=-1)
203
204 def forward(self, x: torch.Tensor) -> torch.Tensor:
205 """Forward pass returning shape (batch_size, 24) key probabilities."""
206 for convnext_block, time_downsampling_block in zip(
207 self.convnext_blocks, self.time_downsampling_blocks, strict=False
208 ):
209 x = time_downsampling_block(x)
210 x = convnext_block(x)
211 x = self.octave_pool(x)
212 x = self.global_average_pool(x)
213 x = self.classifier(x)
214 x = self.batch_norm(x)
215 x = self.flatten(x)
216 return self.softmax(x / self.temperature)
217
218
219# ---------------------------------------------------------------------------
220# Harmonic VQT and CQT cropping
221# ---------------------------------------------------------------------------
222
223
224class VQT(_VQT): # type: ignore[misc]
225 """Harmonic VQT: a collection of VQTs with different harmonic shifts."""
226
227 def __init__(
228 self,
229 *,
230 harmonics: list[float],
231 fmin: float,
232 n_bins: int,
233 bins_per_octave: int = 12,
234 **kwargs: Any,
235 ) -> None:
236 self.harmonics = harmonics
237 self.bin_shifts: list[int] = []
238 self.n_bins_per_slice = n_bins
239 self.fmin = fmin
240 for harmonic in harmonics:
241 shift = round(bins_per_octave * math.log2(harmonic))
242 self.bin_shifts.append(shift)
243 low_octave_shift = min([0, *self.bin_shifts]) / bins_per_octave
244 fmin = fmin * (2**low_octave_shift)
245 n_bins = n_bins + max([0, *self.bin_shifts]) - min([0, *self.bin_shifts])
246 super().__init__(fmin=fmin, n_bins=n_bins, bins_per_octave=bins_per_octave, **kwargs)
247
248 def forward(
249 self,
250 x: torch.Tensor,
251 output_format: str = "Magnitude",
252 normalization_type: str = "librosa",
253 ) -> torch.Tensor:
254 """Compute harmonic VQT and return log-scaled output."""
255 vqt = super().forward(x, output_format, normalization_type)
256 hvqt = []
257 for shift in self.bin_shifts:
258 bin_start = shift - min([0, *self.bin_shifts])
259 bin_stop = bin_start + self.n_bins_per_slice
260 hvqt.append(vqt[:, bin_start:bin_stop, ...])
261 hvqt_tensor = torch.stack(hvqt, dim=1)
262 return ( # type: ignore[no-any-return]
263 (1.0 / 80.0) * torchaudio.transforms.AmplitudeToDB(top_db=80)(hvqt_tensor)
264 ) + 1.0
265
266
267class CropCQT(nn.Module):
268 """Crop Constant-Q Transform spectrograms to a fixed height."""
269
270 def __init__(self, height: int) -> None:
271 super().__init__()
272 self.height = height
273
274 def forward(self, spectrograms: torch.Tensor, transpose: torch.Tensor) -> torch.Tensor:
275 """Crop spectrograms based on transpose values."""
276 return torch.stack(
277 [
278 s[:, int(start_idx) : int(start_idx) + self.height, :]
279 for s, start_idx in zip(spectrograms, transpose, strict=False)
280 ]
281 )
282
283
284# ---------------------------------------------------------------------------
285# Checkpoint loading
286# ---------------------------------------------------------------------------
287
288
289def load_skey_components(
290 device: str = "cpu",
291) -> tuple[VQT, ChromaNet, CropCQT]:
292 """Load VQT, ChromaNet, and CropCQT from the bundled S-KEY checkpoint.
293
294 :param device: Device to load the models onto.
295 """
296 ckpt = torch.load(_CHECKPOINT_PATH, map_location=device, weights_only=False)
297
298 hcqt = VQT(harmonics=[1], fmin=27.5, n_bins=99, verbose=False).to(device)
299 chromanet = ChromaNet(
300 n_bins=84,
301 n_harmonics=1,
302 out_channels=[2, 3, 40, 40, 30, 10, 3],
303 kernels=[7, 7, 7, 7, 7, 5, 5],
304 temperature=1,
305 ).to(device)
306
307 hcqt.load_state_dict(
308 {k.replace("hcqt.", ""): v for k, v in ckpt["stone"].items() if "hcqt" in k}
309 )
310 chromanet.load_state_dict(
311 {k.replace("chromanet.", ""): v for k, v in ckpt["stone"].items() if "chromanet" in k}
312 )
313
314 hcqt.eval()
315 chromanet.eval()
316
317 return hcqt, chromanet, CropCQT(84)
318