Upload Bearings checkpoint

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:4bc76bc5e4711ff5857f31d0cc4e89083110d232a9e58aee42f14935e8b22422
-size 44552888
+oid sha256:83cffe91cf2c20b3fa59763dc7b3e04d6d991302137d00a3662eba282cf7a0d5
+size 44552864

modeling_bearings.py

@@ -17,7 +17,6 @@ from typing import Optional
 import torch
 import torch.nn.functional as F
 from torch import nn
-from torchaudio.transforms import MelScale, Spectrogram
 from transformers import PreTrainedModel
 from transformers.utils import ModelOutput
 
@@ -30,20 +29,54 @@ from .configuration_bearings import BearingsConfig
 
 class FeatureExtractor(nn.Module):
     """Input (B, 4, N) ACN/SN3D FOA -> (B, 7, M, T): log-mel of W, Y, Z, X
-    and the per-band energy-normalised active intensity vector."""
+    and the per-band energy-normalised active intensity vector.
+
+    Plain torch.stft + matmul, equivalent to torchaudio's Spectrogram
+    (center, reflect pad, onesided, power=None) followed by MelScale (htk,
+    no norm).  The Hann window and the mel filterbank are persistent buffers
+    shipped in model.safetensors, so nothing is computed at construction
+    time: transformers >= 5 builds models on the meta device, where
+    torchaudio's filterbank constructor fails.
+    """
 
     def __init__(self, sample_rate, n_fft, win_length, hop_length, f_min, f_max,
                  n_mels, power=2.0):
         super().__init__()
+        self.sample_rate, self.n_fft, self.win_length = sample_rate, n_fft, win_length
+        self.hop_length, self.f_min, self.f_max, self.n_mels = hop_length, f_min, f_max, n_mels
         self.power = power
         self.eps = 1e-6
-        self.spectrogram = Spectrogram(
-            n_fft=n_fft, win_length=win_length, hop_length=hop_length,
-            window_fn=torch.hann_window, power=None, center=True,
-            pad_mode="reflect", onesided=True,
-        )
-        self.mel_scale = MelScale(n_mels, sample_rate, f_min, f_max,
-                                  n_fft // 2 + 1, None, "htk")
+        self.register_buffer("window", torch.zeros(win_length), persistent=True)
+        self.register_buffer("fb", torch.zeros(n_fft // 2 + 1, n_mels), persistent=True)
+        if self.fb.device.type != "meta":
+            self.reset_buffers()
+
+    @torch.no_grad()
+    def reset_buffers(self):
+        """Recompute window and filterbank from the hyper-parameters (same
+        values as the shipped buffers; torchaudio htk mel scale)."""
+        dev = self.fb.device
+        self.window.copy_(torch.hann_window(self.win_length, periodic=True, device=dev))
+        n_freqs = self.n_fft // 2 + 1
+        all_freqs = torch.linspace(0, self.sample_rate // 2, n_freqs, device=dev)
+        hz_to_mel = lambda f: 2595.0 * math.log10(1.0 + f / 700.0)
+        m_pts = torch.linspace(hz_to_mel(self.f_min), hz_to_mel(self.f_max), self.n_mels + 2, device=dev)
+        f_pts = 700.0 * (10.0 ** (m_pts / 2595.0) - 1.0)
+        f_diff = f_pts[1:] - f_pts[:-1]
+        slopes = f_pts.unsqueeze(0) - all_freqs.unsqueeze(1)             # (n_freqs, n_mels + 2)
+        down = (-1.0 * slopes[:, :-2]) / f_diff[:-1]
+        up = slopes[:, 2:] / f_diff[1:]
+        self.fb.copy_(torch.clamp(torch.min(down, up), min=0.0))
+
+    def _stft(self, audio):
+        B, C, N = audio.shape
+        spec = torch.stft(audio.reshape(B * C, N), self.n_fft, self.hop_length, self.win_length,
+                          self.window, center=True, pad_mode="reflect", normalized=False,
+                          onesided=True, return_complex=True)
+        return spec.reshape(B, C, spec.shape[-2], spec.shape[-1])       # (B, C, F, T)
+
+    def _mel(self, x):                                                   # (B, C, F, T) real
+        return torch.matmul(x.transpose(-1, -2), self.fb).transpose(-1, -2)
 
     def _intensity_energy(self, spec):
         W, YZX = spec[:, [0]], spec[:, 1:]
@@ -52,10 +85,10 @@ class FeatureExtractor(nn.Module):
         return I, E
 
     def forward(self, audio):
-        spec = self.spectrogram(audio)                                   # complex (B, 4, F, T)
-        logmel = torch.log(self.mel_scale(torch.abs(spec) ** self.power) + self.eps)
+        spec = self._stft(audio)                                         # complex (B, 4, F, T)
+        logmel = torch.log(self._mel(torch.abs(spec) ** self.power) + self.eps)
         I, E = self._intensity_energy(spec)
-        aiv = self.mel_scale(I) / (self.mel_scale(E) + self.eps)         # (B, 3, M, T)
+        aiv = self._mel(I) / (self._mel(E) + self.eps)                   # (B, 3, M, T)
         return torch.cat([logmel, aiv], dim=1)                           # (B, 7, M, T)