Upload Bearings checkpoint

README.md

@@ -1,5 +1,5 @@
 ---
-library_name: pytorch
+library_name: transformers
 tags:
 - audio
 - spatial-audio
@@ -22,7 +22,25 @@ The encoder needs its tokens fused with a content encoder such as [GRAM-T](https
 
 ## Usage
 
-Clone the repository and install `requirements.txt`, then:
+With transformers (no clone needed, `pip install transformers torchaudio`):
+
+```python
+import torch
+from transformers import AutoModel
+
+model = AutoModel.from_pretrained("labhamlet/bearings", trust_remote_code=True).eval()
+wav = torch.randn(1, 4, int(48000 * 5.3))     # (batch, [W, Y, Z, X], samples), any sample rate
+
+with torch.no_grad():
+    out = model(wav, sampling_rate=48000)
+out.last_hidden_state.shape   # (1, 600, 384): B x (T * F) x E, T = 75 time steps, F = 8 mel bands
+out.timestamps.shape          # (600,) seconds, one per token, window centres: 0.04, ..., 5.96
+out.n_valid_time              # tensor([67]): time steps that end before the zero padding
+```
+
+`strategy="frames"` returns (B, T, 8 * 384) per-frame vectors, `strategy="mean"` and `strategy="cls"` return one (B, 384) vector per clip. `AutoFeatureExtractor.from_pretrained("labhamlet/bearings", trust_remote_code=True)` does the resampling and padding separately if you want to cache waveforms.
+
+With the research code, clone the repository and install `requirements.txt`, then:
 
 ```python
 import torch
@@ -37,12 +55,10 @@ out["timestamps"].shape    # (600,) seconds, one per token, window centres: 0.04
 out["n_valid_time"]        # 67: time steps that end before the zero padding
 ```
 
-`encode_waveform` resamples to 32 kHz, applies per-clip W-RMS normalisation, zero-pads to a whole number of 2 s windows, encodes each window and concatenates along time.
+`encode_waveform` resamples to 32 kHz, applies per-clip W-RMS normalisation, zero-pads to a whole number of 2 s windows, encodes each window and concatenates along time. The AutoModel forward does exactly the same and the two paths give identical embeddings.
 
 **The token axis stacks the T x F grid time-major:** token `i` is time step `i // 8` and mel band `i % 8`. The 8 bands of a time step are adjacent and share a timestamp. Each time step covers 80 ms. `embeddings.view(B, n_time, n_freq, E)` restores the grid and `.reshape(B, n_time, 8 * 384)` gives the per-frame vectors the SELD probes consume.
 
 ## Evaluation
 
 Used as a frozen spatial stream next to a frozen content encoder in SELD probes on TAU-NIGENS 2021 and STARSS23 (DCASE 2023 metrics: ER20, F20, LE, LR). Scripts and task ids are in `seld-dcase2023-gram/` of the repository.
-
-

config.json

@@ -1,57 +1,29 @@
 {
-  "encoder_embedding_dim": 384,
+  "architectures": [
+    "BearingsModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_bearings.BearingsConfig",
+    "AutoFeatureExtractor": "feature_extraction_bearings.BearingsFeatureExtractor",
+    "AutoModel": "modeling_bearings.BearingsModel"
+  },
+  "dtype": "float32",
   "encoder_depth": 6,
-  "num_heads": 12,
-  "mlp_ratio": 4.0,
-  "decoder_depth": 4,
-  "decoder_num_heads": 8,
-  "decoder_embedding_dim": 256,
-  "q_loss_weight": 1.0,
-  "diffuseness_loss_weight": 1.0,
-  "leveldiff_loss_weight": 0.0,
-  "level_loss_weight": 0.0,
-  "lr": 0.0002,
-  "b1": 0.9,
-  "b2": 0.98,
-  "weight_decay": 0.01,
-  "warmup_steps": 10000,
+  "encoder_embedding_dim": 384,
+  "f_max": 14000.0,
+  "f_min": 50.0,
+  "fshape": 16,
+  "hop_length": 320,
   "in_channels": 7,
-  "sr": 32000,
+  "layer_norm_eps": 1e-05,
+  "mlp_ratio": 4.0,
+  "model_type": "bearings",
+  "n_fft": 1024,
+  "num_heads": 12,
   "num_mel_bins": 128,
-  "input_length": 500,
+  "sampling_rate": 32000,
   "target_length": 200,
-  "inter_channel_heads": 4,
-  "inter_channel_layers": 1,
-  "gramt_model_id": "labhamlet/gramt-mono",
-  "gramt_mask_context": false,
-  "n_grid": 256,
-  "vmf_kappa": 40.0,
-  "grid_chunk": 64,
-  "coh_tile_t": 5,
-  "coh_tile_f": 3,
-  "coh_long_t": 35,
-  "coh_long_f": 3,
-  "diffuseness_c": 1.0,
-  "mask_ratio": 0.8,
-  "mask_p_span": 0.4,
-  "mask_p_random": 0.2,
-  "mask_p_censor": 0.4,
-  "span_min_tokens": 2,
-  "span_max_tokens": null,
-  "rotation_mode": "so3",
-  "rotation_prob": 1.0,
-  "f_max": 14000,
-  "log_every_n_steps": 500,
-  "samples_per_clip": 2,
-  "native_sr": 48000,
-  "model_class": "src.model.SphereV5",
-  "fshape": 16,
+  "transformers_version": "4.57.6",
   "tshape": 8,
-  "fstride": 16,
-  "tstride": 8,
-  "global_step": 100000,
-  "source_checkpoint": "step=100000.ckpt",
-  "stripped_prefixes": [
-    "gram."
-  ]
-}
+  "win_length": 1024
+}

configuration_bearings.py

@@ -0,0 +1,52 @@
+"""Bearings configuration for the transformers AutoModel interface."""
+from transformers import PretrainedConfig
+
+
+class BearingsConfig(PretrainedConfig):
+    model_type = "bearings"
+
+    def __init__(
+        self,
+        encoder_embedding_dim: int = 384,
+        encoder_depth: int = 6,
+        num_heads: int = 12,
+        mlp_ratio: float = 4.0,
+        layer_norm_eps: float = 1e-5,
+        in_channels: int = 7,
+        num_mel_bins: int = 128,
+        target_length: int = 200,
+        fshape: int = 16,
+        tshape: int = 8,
+        sampling_rate: int = 32000,
+        n_fft: int = 1024,
+        win_length: int = 1024,
+        hop_length: int = 320,
+        f_min: float = 50.0,
+        f_max: float = 14000.0,
+        **kwargs,
+    ):
+        self.encoder_embedding_dim = encoder_embedding_dim
+        self.encoder_depth = encoder_depth
+        self.num_heads = num_heads
+        self.mlp_ratio = mlp_ratio
+        self.layer_norm_eps = layer_norm_eps
+        self.in_channels = in_channels
+        self.num_mel_bins = num_mel_bins
+        self.target_length = target_length
+        self.fshape = fshape
+        self.tshape = tshape
+        self.sampling_rate = sampling_rate
+        self.n_fft = n_fft
+        self.win_length = win_length
+        self.hop_length = hop_length
+        self.f_min = f_min
+        self.f_max = f_max
+        super().__init__(**kwargs)
+
+    @property
+    def n_freq_patches(self) -> int:
+        return self.num_mel_bins // self.fshape
+
+    @property
+    def n_time_patches(self) -> int:
+        return self.target_length // self.tshape

feature_extraction_bearings.py

@@ -0,0 +1,51 @@
+"""Waveform-level feature extractor for Bearings.
+
+Resamples first-order Ambisonic audio to the model rate and zero-pads it to a
+whole number of 2 s windows.  Log-mel / intensity-vector extraction and the
+per-clip W-RMS normalisation live inside the model so that the model can also
+be called on raw waveforms directly.
+"""
+import math
+from typing import List, Optional, Union
+
+import numpy as np
+import torch
+from transformers.feature_extraction_utils import BatchFeature, FeatureExtractionMixin
+
+
+class BearingsFeatureExtractor(FeatureExtractionMixin):
+    model_input_names = ["input_values"]
+
+    def __init__(self, sampling_rate: int = 32000, window_seconds: float = 2.0,
+                 **kwargs):
+        super().__init__(**kwargs)
+        self.sampling_rate = sampling_rate
+        self.window_seconds = window_seconds
+
+    def __call__(self, audio: Union[np.ndarray, torch.Tensor, List],
+                 sampling_rate: Optional[int] = None,
+                 return_tensors: str = "pt") -> BatchFeature:
+        """audio: (4, N) or (B, 4, N), channels W, Y, Z, X (ACN/SN3D)."""
+        if sampling_rate is None:
+            raise ValueError("pass sampling_rate= of the audio you are feeding")
+        wav = torch.as_tensor(np.asarray(audio) if not torch.is_tensor(audio) else audio)
+        wav = wav.float()
+        if wav.ndim == 2:
+            wav = wav.unsqueeze(0)
+        if wav.ndim != 3 or wav.shape[1] != 4:
+            raise ValueError(f"expected (B, 4, N) FOA audio, got {tuple(wav.shape)}")
+        if int(sampling_rate) != self.sampling_rate:
+            import torchaudio
+            wav = torchaudio.functional.resample(
+                wav, int(sampling_rate), self.sampling_rate, lowpass_filter_width=64)
+        n_valid = wav.shape[-1]
+        win = int(self.window_seconds * self.sampling_rate)
+        n_win = max(1, math.ceil(n_valid / win))
+        wav = torch.nn.functional.pad(wav, (0, n_win * win - n_valid))
+        out = {"input_values": wav,
+               "n_valid_samples": torch.full((wav.shape[0],), n_valid, dtype=torch.long)}
+        if return_tensors == "np":
+            out = {k: v.numpy() for k, v in out.items()}
+        elif return_tensors != "pt":
+            raise ValueError("return_tensors must be 'pt' or 'np'")
+        return BatchFeature(out)

model.safetensors

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

modeling_bearings.py

@@ -0,0 +1,262 @@
+"""Bearings: self-supervised sound-field encoder for first-order Ambisonics.
+
+Standalone transformers port of the SphereV5 encoder
+(https://github.com/root-goksenin/Embisonics, src/model/mae.py).  Only the
+inference path is here: log-mel + active-intensity front end, patch
+tokeniser, ViT encoder.  The pre-training decoder and targets are not.
+
+    model = AutoModel.from_pretrained("labhamlet/bearings", trust_remote_code=True)
+    out = model(wav, sampling_rate=48000)        # wav: (B, 4, N) FOA, W Y Z X
+    out.last_hidden_state                        # (B, T * F, E) time-major stacked tokens
+    out.timestamps                               # (T * F,) seconds
+"""
+import math
+from dataclasses import dataclass
+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
+
+from .configuration_bearings import BearingsConfig
+
+
+# =============================================================================
+# Front end: 4 log-mel + 3 normalised active-intensity channels
+# =============================================================================
+
+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."""
+
+    def __init__(self, sample_rate, n_fft, win_length, hop_length, f_min, f_max,
+                 n_mels, power=2.0):
+        super().__init__()
+        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")
+
+    def _intensity_energy(self, spec):
+        W, YZX = spec[:, [0]], spec[:, 1:]
+        I = 2.0 * torch.real(torch.conj(W) * YZX)                       # (B, 3, F, T)
+        E = torch.abs(W) ** 2 + (torch.abs(YZX) ** 2).sum(1, keepdim=True)
+        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)
+        I, E = self._intensity_energy(spec)
+        aiv = self.mel_scale(I) / (self.mel_scale(E) + self.eps)         # (B, 3, M, T)
+        return torch.cat([logmel, aiv], dim=1)                           # (B, 7, M, T)
+
+
+# =============================================================================
+# ViT pieces (parameter names match timm's Block)
+# =============================================================================
+
+class Attention(nn.Module):
+    def __init__(self, dim, num_heads):
+        super().__init__()
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.qkv = nn.Linear(dim, dim * 3, bias=True)
+        self.proj = nn.Linear(dim, dim)
+
+    def forward(self, x):
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv.unbind(0)
+        x = F.scaled_dot_product_attention(q, k, v)
+        return self.proj(x.transpose(1, 2).reshape(B, N, C))
+
+
+class Mlp(nn.Module):
+    def __init__(self, dim, hidden):
+        super().__init__()
+        self.fc1 = nn.Linear(dim, hidden)
+        self.act = nn.GELU()
+        self.fc2 = nn.Linear(hidden, dim)
+
+    def forward(self, x):
+        return self.fc2(self.act(self.fc1(x)))
+
+
+class Block(nn.Module):
+    def __init__(self, dim, num_heads, mlp_ratio, eps):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(dim, eps=eps)
+        self.attn = Attention(dim, num_heads)
+        self.norm2 = nn.LayerNorm(dim, eps=eps)
+        self.mlp = Mlp(dim, int(dim * mlp_ratio))
+
+    def forward(self, x):
+        x = x + self.attn(self.norm1(x))
+        return x + self.mlp(self.norm2(x))
+
+
+class PatchEmbed(nn.Module):
+    """Non-overlapping Conv2D tokeniser, (B, 7, F, T) -> (B, P, D), freq-major."""
+
+    def __init__(self, in_ch, dim, fshape, tshape):
+        super().__init__()
+        self.proj = nn.Conv2d(in_ch, dim, kernel_size=(fshape, tshape), stride=(fshape, tshape))
+
+    def forward(self, x):
+        x = self.proj(x)
+        B, D, Nf, Nt = x.shape
+        return x.permute(0, 2, 3, 1).reshape(B, Nf * Nt, D)
+
+
+# =============================================================================
+# Model
+# =============================================================================
+
+@dataclass
+class BearingsOutput(ModelOutput):
+    """last_hidden_state: (B, T * F, E) for strategy 'timestamps'; (B, T, F * E)
+    for 'frames'; (B, E) for 'mean' / 'cls'.  Token axis of 'timestamps' is
+    TIME-MAJOR: index = t * F + f (F = 8 mel bands per 80 ms time step)."""
+    last_hidden_state: torch.FloatTensor = None
+    timestamps: Optional[torch.FloatTensor] = None
+    n_time: Optional[int] = None
+    n_freq: Optional[int] = None
+    n_valid_time: Optional[torch.LongTensor] = None
+
+
+class BearingsPreTrainedModel(PreTrainedModel):
+    config_class = BearingsConfig
+    base_model_prefix = "bearings"
+    main_input_name = "input_values"
+    supports_gradient_checkpointing = False
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            nn.init.xavier_uniform_(module.weight)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+
+
+class BearingsModel(BearingsPreTrainedModel):
+    def __init__(self, config: BearingsConfig):
+        super().__init__(config)
+        c = config
+        D = c.encoder_embedding_dim
+        self.n_freq = c.n_freq_patches
+        self.n_time_per_window = c.n_time_patches
+        P = self.n_freq * self.n_time_per_window
+
+        self.melspec = FeatureExtractor(
+            c.sampling_rate, c.n_fft, c.win_length, c.hop_length, c.f_min, c.f_max,
+            c.num_mel_bins, power=2.0,
+        )
+        self.patch_embed = PatchEmbed(c.in_channels, D, c.fshape, c.tshape)
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, D))
+        # fixed 2-D sincos embedding, shipped with the weights
+        self.register_buffer("pos_embed", torch.zeros(1, P + 1, D), persistent=True)
+        self.encoder_blocks = nn.ModuleList([
+            Block(D, c.num_heads, c.mlp_ratio, c.layer_norm_eps) for _ in range(c.encoder_depth)
+        ])
+        self.encoder_norm = nn.LayerNorm(D, eps=c.layer_norm_eps)
+        self.post_init()
+
+    # ---- pieces -----------------------------------------------------------
+    @property
+    def window_samples(self) -> int:
+        return self.config.target_length * self.config.hop_length
+
+    def extract_features(self, wav: torch.Tensor) -> torch.Tensor:
+        """(B, 4, N) at config.sampling_rate -> (B, 7, T, M), W-RMS normalised."""
+        with torch.autocast(device_type=wav.device.type, enabled=False):
+            wav = wav.float()
+            rms = torch.sqrt(wav[:, :1].pow(2).mean(dim=-1, keepdim=True) + 1e-8)
+            return self.melspec(wav / rms).transpose(-1, -2)
+
+    def encode_window(self, x7: torch.Tensor) -> torch.Tensor:
+        """(B, 7, 200, 128) -> (B, 1 + 200, E); tokens after CLS are freq-major."""
+        x = self.patch_embed(x7.transpose(2, 3)) + self.pos_embed[:, 1:]
+        cls = self.cls_token.expand(x.shape[0], -1, -1) + self.pos_embed[:, :1]
+        x = torch.cat([cls, x], dim=1)
+        for blk in self.encoder_blocks:
+            x = blk(x)
+        return self.encoder_norm(x)
+
+    # ---- forward ----------------------------------------------------------
+    def forward(self, input_values: torch.Tensor, sampling_rate: Optional[int] = None,
+                n_valid_samples: Optional[torch.Tensor] = None,
+                strategy: str = "timestamps", chunk_batch: int = 64,
+                return_dict: bool = True):
+        """input_values: (B, 4, N) or (4, N) FOA waveform, channels W, Y, Z, X.
+        sampling_rate: rate of input_values; resampled to config.sampling_rate
+        when different (omit it if you used BearingsFeatureExtractor).
+        strategy: 'timestamps' (B, T*F, E) + timestamps | 'frames' (B, T, F*E)
+        | 'mean' (B, E) | 'cls' (B, E)."""
+        c = self.config
+        wav = input_values
+        if wav.ndim == 2:
+            wav = wav.unsqueeze(0)
+        if wav.ndim != 3 or wav.shape[1] != 4:
+            raise ValueError(f"expected (B, 4, N) FOA audio, got {tuple(wav.shape)}")
+        wav = wav.to(self.pos_embed.device, torch.float32)
+        if sampling_rate is not None and int(sampling_rate) != c.sampling_rate:
+            import torchaudio
+            wav = torchaudio.functional.resample(
+                wav, int(sampling_rate), c.sampling_rate, lowpass_filter_width=64)
+
+        B, _, N = wav.shape
+        if n_valid_samples is None:
+            n_valid_samples = torch.full((B,), N, dtype=torch.long, device=wav.device)
+        win = self.window_samples
+        n_win = max(1, math.ceil(N / win))
+        wav = F.pad(wav, (0, n_win * win - N))
+
+        x7 = self.extract_features(wav)[:, :, :n_win * c.target_length]    # (B, 7, T, M)
+        x7 = x7.reshape(B, 7, n_win, c.target_length, c.num_mel_bins)
+        x7 = x7.permute(0, 2, 1, 3, 4).reshape(B * n_win, 7, c.target_length, c.num_mel_bins)
+
+        zs = [self.encode_window(x7[i:i + chunk_batch]).float()
+              for i in range(0, x7.shape[0], chunk_batch)]
+        z = torch.cat(zs, 0)                                                 # (B*n, 1+P, E)
+        E = z.shape[-1]
+        Fb, Tw = self.n_freq, self.n_time_per_window
+        cls = z[:, 0].reshape(B, n_win, E)
+        tok = z[:, 1:].reshape(B, n_win, Fb, Tw, E).permute(0, 1, 3, 2, 4)  # (B, n, t, f, E)
+        T = n_win * Tw
+        tok = tok.reshape(B, T, Fb, E)
+
+        step = c.tshape * c.hop_length / c.sampling_rate                     # 80 ms
+        t_centre = (torch.arange(T, device=tok.device, dtype=torch.float32) + 0.5) * step
+        n_valid_time = torch.clamp(
+            torch.ceil(n_valid_samples.float() / (c.tshape * c.hop_length)).long(), max=T)
+
+        if strategy == "timestamps":
+            hidden = tok.reshape(B, T * Fb, E)
+            timestamps = t_centre.repeat_interleave(Fb)
+        elif strategy == "frames":
+            hidden = tok.reshape(B, T, Fb * E)
+            timestamps = t_centre
+        elif strategy in ("mean", "cls"):
+            # average over the valid time steps / windows only
+            if strategy == "mean":
+                mask = (torch.arange(T, device=tok.device)[None, :] < n_valid_time[:, None]).float()
+                hidden = (tok.mean(2) * mask[..., None]).sum(1) / mask.sum(1, keepdim=True)
+            else:
+                n_valid_win = torch.clamp(torch.ceil(n_valid_samples.float() / win).long(), max=n_win)
+                mask = (torch.arange(n_win, device=tok.device)[None, :] < n_valid_win[:, None]).float()
+                hidden = (cls * mask[..., None]).sum(1) / mask.sum(1, keepdim=True)
+            timestamps = None
+        else:
+            raise ValueError("strategy must be 'timestamps', 'frames', 'mean' or 'cls'")
+
+        if not return_dict:
+            return hidden, timestamps
+        return BearingsOutput(last_hidden_state=hidden, timestamps=timestamps,
+                              n_time=T, n_freq=Fb, n_valid_time=n_valid_time)

preprocessor_config.json

@@ -0,0 +1,8 @@
+{
+  "auto_map": {
+    "AutoFeatureExtractor": "feature_extraction_bearings.BearingsFeatureExtractor"
+  },
+  "feature_extractor_type": "BearingsFeatureExtractor",
+  "sampling_rate": 32000,
+  "window_seconds": 2.0
+}