Add Transformers loading for the existing AECC 2026 encoder

CODE_LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 audio-embeddings contributors
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 audio-embeddings contributors
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -1,35 +1,149 @@
-# BEST-RQ-2 (xares-llm encoder)
+---
+library_name: transformers
+license: mit
+tags:
+  - audio
+  - feature-extraction
+  - custom_code
+  - self-supervised-learning
+  - audio-embeddings
+  - best-rq-2
+  - audioset
+---
+# BEST-RQ-2
 
-This folder contains the BEST-RQ-2 audio encoder integration for `xares-llm`.
+BEST-RQ-2 is a self-supervised audio encoder trained on AudioSet for
+a configured budget of **200,000 optimizer steps**. It produces **768-dimensional** clip and frame
+embeddings from mono **16 kHz** waveforms, and supports downstream fine-tuning.
 
-Benchmark repository: [xiaomi-research/xares-llm](https://github.com/xiaomi-research/xares-llm)
+This repository contains the trained encoder, its preprocessing configuration,
+and the custom Transformers implementation. No installation of the research
+repository is needed.
 
-## Setup
+## Model and training
 
-1. Make sure the environment to run `xares-llm` is set up properly (virtual environment initialized and the `xares-llm` package downloaded/installed).
+| Property | Value |
+| --- | --- |
+| Architecture | 12-layer Transformer, 768 dimensions, 12 attention heads |
+| Input frontend | 128-bin mel spectrogram with a linear patch projection |
+| Patch shape | 16 mel bins by 16 time frames |
+| Transformer | Sinusoidal positional embeddings, LayerNorm, GELU MLP |
+| Training data | AudioSet |
+| Training objective | Masked prediction of frozen codebook targets |
+| Masking ratio | 40–60% |
+| Checkpoint step metadata | Not recorded in this older safetensors file |
+| Exported weight dtype | float32 |
+| Extraction policy | `overlap50_two_phase` |
 
-2. Add the `BEST-RQ-2` folder to the `xares-llm` (current) directory so it is available at `./BEST-RQ-2`.
+## Load the model
 
-3. Before running a `xares-llm` evaluation, you must install the required packages for BEST-RQ-2:
+Install the runtime dependencies in your Python environment:
 
 ```bash
-uv pip install -r BEST-RQ-2/audio-embeddings/pyproject.toml
+pip install "torch>=2.9.1" "torchaudio>=2.9.1" "timm>=0.9" "einops>=0.7" "transformers>=4.57,<6"
 ```
 
-## Run an evaluation
+Use matching PyTorch and torchaudio versions. GPU installations may require the
+appropriate PyTorch build for your CUDA version.
 
-Single task (e.g. to test that everything works):
+```python
+import torch
+from transformers import AutoFeatureExtractor, AutoModel
 
-```bash
-uv run -m xares_llm.run BEST-RQ-2.BEST-RQ-2_encoder.BestRQ2Encoder <task> <task> # e.g. replace <task> with esc-50, score should be around 0.48
+model_name = "ltuncay/BEST-RQ-2"
+processor = AutoFeatureExtractor.from_pretrained(model_name, trust_remote_code=True)
+model = AutoModel.from_pretrained(model_name, trust_remote_code=True).eval()
+audio = torch.zeros(processor.sampling_rate)  # Replace with real mono audio.
+inputs = processor(audio, sampling_rate=processor.sampling_rate, return_tensors="pt")
+with torch.inference_mode():
+    outputs = model(**inputs)
+clip_embeddings = outputs.pooler_output
+frame_embeddings = outputs.last_hidden_state
+output_dim = model.config.encoder_kwargs["embed_dim"]
 ```
 
-All tasks:
+Resample to `processor.sampling_rate` and downmix stereo before preprocessing.
+The extractor performs padding only; spectrogram/convolution features are computed
+inside the model. Pass the returned sample attention mask for variable durations.
+Outputs include a frame attention mask and timestamps in milliseconds (-1 for padding).
+`pooler_output` uses the saved HEAR extraction preset, including phase-balanced pooling.
+Frame features average frequency patches; they are not the raw frequency-time ViT grid.
 
-```bash
-uv run -m xares_llm.run BEST-RQ-2.BEST-RQ-2_encoder.BestRQ2Encoder all
+## Fine-tuning
+
+For fine-tuning call `model.train()`, attach a task head and optimize its parameters
+alongside the model. Save with `model.save_pretrained(path)` and
+`processor.save_pretrained(path)`. This encoder export excludes pretraining predictors,
+quantizers, teachers and optimizer state. Continue self-supervised research training
+with the original Lightning code and checkpoints.
+
+## Reproducibility and provenance
+
+`revision` is optional; pin both loaders to the same full commit hash for reproducibility.
+Custom Python code is included in this repository and requires `trust_remote_code=True`.
+The weights use safetensors. See `export_manifest.json` for source and validation details.
+Source: the existing **AECC 2026 submission** in this repository, pinned at
+[`a03eeb5c4433f4bf7a7e6e8b4724af862789959c`](https://huggingface.co/ltuncay/BEST-RQ-2/tree/a03eeb5c4433f4bf7a7e6e8b4724af862789959c).
+The export uses its `BEST-RQ-2.safetensors` and matching `config.yaml`.
+The saved configuration specifies a 200,000-step training budget, but this older
+safetensors file has no saved-step metadata. The actual checkpoint step therefore
+remains unspecified in the export manifest. Source file hashes are recorded there.
+
+## AECC 2026 compatibility
+
+The [original submission](https://huggingface.co/ltuncay/BEST-RQ-2/tree/aecc-2026-submission)
+is preserved at revision `aecc-2026-submission`, including its model card, checkpoint,
+technical report and xares-llm integration. The legacy files also remain available
+on `main` for existing consumers.
+
+This Transformers export uses the **same pretrained encoder weights**. Its public
+outputs follow the shared BEST-RQ-2 family extraction policy: frequency-averaged
+frame embeddings and phase-balanced clip embeddings with overlapping windows.
+The original AECC wrapper returned spectrogram patch tokens with sequential audio
+chunking, so the two public interfaces are not interchangeable for reproducing
+benchmark results. Use the archived revision for the original evaluation setup.
+
+A comparison against the archived frontend and encoder on a 10-second waveform
+found exact spectrogram and patch-projection agreement, with maximum encoder
+absolute difference `1.67e-6` (within floating-point tolerance). The validation
+manifest records this comparison.
+
+The exported model matched the pre-export encoder exactly on the exporter's
+variable-duration test batch. These are integration checks, not a downstream
+benchmark evaluation of this release.
+
+## Intended uses and limitations
+
+Use these embeddings as features for audio research, classification, retrieval,
+or downstream fine-tuning. This is an encoder: it does not generate transcripts,
+audio, or class labels without a downstream model.
+
+AudioSet training does not establish performance on every language, acoustic
+domain, demographic group, or downstream task. Evaluate the model on your target
+data. No new downstream benchmark scores are claimed for this export.
+
+The feature extractor pads waveforms and checks the sample rate; it does not
+resample or downmix. Unequal-length clips are processed individually after padding
+is removed. The output embeddings follow the saved windowing policy rather than
+exposing the raw spectrogram token grid.
+
+Browse the [BEST-RQ-2 family collection](https://huggingface.co/collections/ltuncay/best-rq-2-family-6aad761396dfe31da48ff91a).
+
+## Source and citation
+
+Research code: [audio-embeddings](https://github.com/LudovicTuncay/audio-embeddings).
+BEST-RQ-2 belongs to the BEST-RQ-2 family. Cite the family paper when using it:
+
+```bibtex
+@inproceedings{tuncay2026best,
+  title={BEST-RQ-2: Contextualize-Then-Predict, a Two-Step Approach for Self-Supervised Audio Representations},
+  author={Tuncay, Ludovic K and Labb{\'e}, Etienne and Pellegrini, Thomas},
+  booktitle={Interspeech 2026},
+  year={2026}
+}
 ```
 
-## Help
+## License
 
-If you encounter any problems, contact: [ludovic.tuncay@irit.fr](mailto:ludovic.tuncay@irit.fr)
+The model weights and bundled implementation are released under the MIT license.
+See `LICENSE` and `CODE_LICENSE`.

adapters.py

@@ -0,0 +1,473 @@
+# MIT License
+# 
+# Copyright (c) 2026 audio-embeddings contributors
+# 
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+# 
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+# 
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+from __future__ import annotations
+
+from dataclasses import asdict, dataclass
+from typing import Any, Mapping
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .extraction import fuse_context_windows
+from .extraction import get_preset
+from .extraction import merge_phases
+from .patch_embed import PatchEmbed
+from .spectrogram import Spectrogram
+from .vit import ViT
+from .vit import vit_config_with_patch_geometry
+from .waveform_feature_encoder import WaveformFeatureEncoder
+
+SPECTROGRAM_TARGETS = {
+    "src.models.audio_jepa_module.AudioJEPAModule": "student",
+    "src.models.rqa_jepa_module.RQAJEPAModule": "student",
+    "src.models.best_rq_module.BestRQModule": "encoder",
+    "src.models.best_rq2_module.BestRQ2Module": "encoder",
+    "src.models.best_rq22_module.BestRQ22Module": "encoder",
+    "src.models.best_rq23_module.BestRQ23Module": "encoder",
+}
+WAVEFORM_TARGETS = {
+    "src.models.best_rq3_module.BestRQ3Module": "encoder",
+}
+
+
+@dataclass(frozen=True)
+class AdapterSpec:
+    adapter_key: str
+    model_target: str
+    encoder_prefix: str
+    sample_rate: int
+    embedding_dim: int
+    max_context_tokens: int
+    temporal_grid_tokens: int
+    token_hop_samples: int
+    receptive_field_samples: int
+    supported_phase_offsets_samples: tuple[int, ...]
+    channel_policy: str = "mono_mean"
+
+    def to_dict(self) -> dict[str, Any]:
+        return asdict(self)
+
+
+@dataclass(frozen=True)
+class EmbeddingOutput:
+    timestamp_embeddings: torch.Tensor
+    timestamps_ms: torch.Tensor
+    scene_embedding: torch.Tensor
+
+
+def _mapping(value: Any, path: str) -> Mapping[str, Any]:
+    if not isinstance(value, Mapping):
+        raise ValueError(f"Expected mapping at {path}, got {type(value).__name__}")
+    return value
+
+
+def _positive_int(value: Any, path: str) -> int:
+    try:
+        normalized = int(value)
+    except (TypeError, ValueError) as error:
+        raise ValueError(f"Expected integer at {path}, got {value!r}") from error
+    if normalized <= 0:
+        raise ValueError(f"Expected positive integer at {path}, got {normalized}")
+    return normalized
+
+
+def resolve_adapter_spec(config: Mapping[str, Any]) -> AdapterSpec:
+    model = _mapping(config.get("model"), "model")
+    target = str(model.get("_target_", ""))
+    net = _mapping(model.get("net"), "model.net")
+    encoder = _mapping(net.get("encoder"), "model.net.encoder")
+
+    if target in SPECTROGRAM_TARGETS:
+        adapter_key = "spectrogram_patch"
+        encoder_prefix = SPECTROGRAM_TARGETS[target]
+        frontend = _mapping(
+            net.get("spectrogram"),
+            "model.net.spectrogram",
+        )
+        sample_rate_path = "model.net.spectrogram.sample_rate"
+        patch = _mapping(net.get("patch_embed"), "model.net.patch_embed")
+        patch_size = tuple(patch.get("patch_size", ()))
+        image_size = tuple(patch.get("img_size", ()))
+        if len(patch_size) != 2 or len(image_size) != 2:
+            raise ValueError("Spectrogram HEAR adapters require 2-D patch/image sizes")
+        patch_height, patch_width = map(int, patch_size)
+        frequency_tokens = int(image_size[0]) // patch_height
+        n_fft = _positive_int(
+            frontend.get("n_fft", 4096), "model.net.spectrogram.n_fft"
+        )
+        if frontend.get("win_length") is not None:
+            win_length = _positive_int(
+                frontend.get("win_length"), "model.net.spectrogram.win_length"
+            )
+        elif frontend.get("win_length_ms") is not None:
+            win_length = int(
+                int(frontend["sample_rate"]) * float(frontend["win_length_ms"]) / 1000
+            )
+        else:
+            win_length = n_fft
+        if frontend.get("hop_length") is not None:
+            frontend_hop = _positive_int(
+                frontend.get("hop_length"), "model.net.spectrogram.hop_length"
+            )
+        elif frontend.get("hop_length_ms") is not None:
+            frontend_hop = int(
+                int(frontend["sample_rate"]) * float(frontend["hop_length_ms"]) / 1000
+            )
+        else:
+            frontend_hop = win_length // 2
+        token_hop = frontend_hop * patch_width
+        receptive_field = n_fft + (patch_width - 1) * frontend_hop
+    elif target in WAVEFORM_TARGETS:
+        adapter_key = "waveform_conv"
+        encoder_prefix = WAVEFORM_TARGETS[target]
+        frontend = _mapping(net.get("sampling"), "model.net.sampling")
+        sample_rate_path = "model.net.sampling.sample_rate"
+        feature_config = dict(
+            _mapping(net.get("feature_encoder"), "model.net.feature_encoder")
+        )
+        feature_encoder = WaveformFeatureEncoder(**feature_config)
+        receptive_field = 1
+        token_hop = 1
+        for _, kernel, stride in feature_encoder.conv_layers_spec:
+            receptive_field += (kernel - 1) * token_hop
+            token_hop *= stride
+        frequency_tokens = 1
+    else:
+        supported = sorted((*SPECTROGRAM_TARGETS, *WAVEFORM_TARGETS))
+        raise ValueError(
+            f"No HEAR adapter is registered for model target {target!r}. "
+            f"Register one for the new model. Current targets: {supported}"
+        )
+
+    sample_rate = _positive_int(frontend.get("sample_rate"), sample_rate_path)
+    data = config.get("data")
+    if isinstance(data, Mapping) and data.get("target_sample_rate") is not None:
+        data_sample_rate = _positive_int(
+            data.get("target_sample_rate"),
+            "data.target_sample_rate",
+        )
+        if data_sample_rate != sample_rate:
+            raise ValueError(
+                "Model/data sampling-rate mismatch: "
+                f"{sample_rate_path}={sample_rate}, "
+                f"data.target_sample_rate={data_sample_rate}"
+            )
+
+    max_context_tokens = _positive_int(
+        encoder.get("num_patches"),
+        "model.net.encoder.num_patches",
+    )
+    if frequency_tokens <= 0 or max_context_tokens < frequency_tokens:
+        raise ValueError(
+            "Encoder context cannot hold one complete frequency-token column"
+        )
+    return AdapterSpec(
+        adapter_key=adapter_key,
+        model_target=target,
+        encoder_prefix=encoder_prefix,
+        sample_rate=sample_rate,
+        embedding_dim=_positive_int(
+            encoder.get("embed_dim"),
+            "model.net.encoder.embed_dim",
+        ),
+        max_context_tokens=max_context_tokens,
+        temporal_grid_tokens=max_context_tokens // frequency_tokens,
+        token_hop_samples=token_hop,
+        receptive_field_samples=receptive_field,
+        supported_phase_offsets_samples=(
+            (0, token_hop // 2) if token_hop % 2 == 0 else (0,)
+        ),
+    )
+
+
+class HearEncoderAdapter(nn.Module):
+    spec: AdapterSpec
+
+    @property
+    def sample_rate(self) -> int:
+        return self.spec.sample_rate
+
+    @property
+    def embedding_dim(self) -> int:
+        return self.spec.embedding_dim
+
+    def extract(self, waveform: torch.Tensor, *, preset_name: str) -> EmbeddingOutput:
+        raise NotImplementedError
+
+
+def _single_waveform(waveform: torch.Tensor) -> torch.Tensor:
+    if waveform.ndim == 1:
+        waveform = waveform.unsqueeze(0)
+    if waveform.ndim != 2 or waveform.shape[0] != 1:
+        raise ValueError(
+            "Adapter extraction expects one mono waveform [samples] or [1, samples], "
+            f"got {tuple(waveform.shape)}"
+        )
+    if waveform.shape[-1] == 0:
+        raise ValueError("Cannot embed an empty waveform")
+    return waveform.unsqueeze(0)
+
+
+class SpectrogramPatchAdapter(HearEncoderAdapter):
+    def __init__(self, config: Mapping[str, Any], spec: AdapterSpec) -> None:
+        super().__init__()
+        self.spec = spec
+        model = _mapping(config.get("model"), "model")
+        net = _mapping(model.get("net"), "model.net")
+        spectrogram_config = dict(
+            _mapping(net.get("spectrogram"), "model.net.spectrogram")
+        )
+        patch_config = dict(_mapping(net.get("patch_embed"), "model.net.patch_embed"))
+        encoder_config = dict(_mapping(net.get("encoder"), "model.net.encoder"))
+        self.spectrogram = Spectrogram(**spectrogram_config)
+        self.patch_embed = PatchEmbed(**patch_config)
+        self.encoder = ViT(
+            **vit_config_with_patch_geometry(
+                encoder_config,
+                img_size=self.patch_embed.img_size,
+                patch_size=self.patch_embed.patch_size,
+            )
+        )
+        self.adjustment_mode = str(model.get("spectrogram_adjustment_mode", "pad"))
+        if self.adjustment_mode not in {"pad", "truncate"}:
+            raise ValueError(
+                f"Unknown spectrogram_adjustment_mode {self.adjustment_mode!r}"
+            )
+
+    def _phase_tokens(
+        self,
+        spectrogram: torch.Tensor,
+        *,
+        frame_offset: int,
+    ) -> tuple[torch.Tensor, int]:
+        patch_height, patch_width = self.patch_embed.patch_size
+        phase = spectrogram[..., frame_offset:]
+        original_frames = phase.shape[-1]
+        if original_frames < patch_width:
+            phase = F.pad(phase, (0, patch_width - original_frames))
+        else:
+            remainder = original_frames % patch_width
+            if remainder:
+                if self.adjustment_mode == "pad":
+                    phase = F.pad(phase, (0, patch_width - remainder))
+                else:
+                    phase = phase[..., : original_frames - remainder]
+        tokens = self.patch_embed(phase)
+        frequency = phase.shape[-2] // patch_height
+        time = phase.shape[-1] // patch_width
+        return tokens.reshape(frequency, time, -1), original_frames
+
+    def extract(self, waveform: torch.Tensor, *, preset_name: str) -> EmbeddingOutput:
+        preset = get_preset(preset_name)
+        waveform = _single_waveform(waveform)
+        duration_samples = waveform.shape[-1]
+        spectrogram = self.spectrogram(waveform)
+        patch_width = self.patch_embed.patch_size[1]
+        if preset.num_phases == 2 and patch_width % 2:
+            raise ValueError(
+                "Two-phase extraction requires an exact half-hop temporal patch offset"
+            )
+        frame_offsets = [0]
+        if preset.num_phases == 2:
+            frame_offsets.append(patch_width // 2)
+
+        hop_samples = int(self.spectrogram.mel_spec.hop_length)
+        phases: list[tuple[torch.Tensor, torch.Tensor]] = []
+        for frame_offset in frame_offsets:
+            token_grid, _ = self._phase_tokens(
+                spectrogram,
+                frame_offset=frame_offset,
+            )
+            frequency = token_grid.shape[0]
+
+            def encode_window(
+                window: torch.Tensor,
+                position_ids: torch.Tensor,
+            ) -> torch.Tensor:
+                width = window.shape[1] // frequency
+                return self.encoder(
+                    window,
+                    pos_ids=position_ids,
+                    grid_size=(frequency, width),
+                )
+
+            embeddings = fuse_context_windows(
+                token_grid,
+                max_context_tokens=self.spec.max_context_tokens,
+                overlap=preset.overlap,
+                encode_window=encode_window,
+            )
+            centers_in_frames = (
+                torch.arange(embeddings.shape[0], device=embeddings.device)
+                * patch_width
+                + frame_offset
+                + (patch_width - 1) / 2.0
+            )
+            centers_in_samples = centers_in_frames * hop_samples
+            centers_in_samples = torch.clamp(
+                centers_in_samples,
+                max=max(0, duration_samples - 1),
+            )
+            timestamps_ms = centers_in_samples * (1000.0 / self.sample_rate)
+            phases.append((embeddings, timestamps_ms))
+
+        merged, timestamps, scene = merge_phases(phases)
+        return EmbeddingOutput(merged, timestamps, scene)
+
+
+class WaveformConvAdapter(HearEncoderAdapter):
+    def __init__(self, config: Mapping[str, Any], spec: AdapterSpec) -> None:
+        super().__init__()
+        self.spec = spec
+        model = _mapping(config.get("model"), "model")
+        net = _mapping(model.get("net"), "model.net")
+        feature_config = dict(
+            _mapping(net.get("feature_encoder"), "model.net.feature_encoder")
+        )
+        encoder_config = dict(_mapping(net.get("encoder"), "model.net.encoder"))
+        self.feature_encoder = WaveformFeatureEncoder(**feature_config)
+        feature_dim = self.feature_encoder.embedding_dim
+        self.encoder_input_proj: nn.Module
+        if feature_dim == spec.embedding_dim:
+            self.encoder_input_proj = nn.Identity()
+        else:
+            self.encoder_input_proj = nn.Linear(feature_dim, spec.embedding_dim)
+        self.encoder = ViT(**encoder_config)
+        receptive_field = 1
+        stride = 1
+        for _, kernel, layer_stride in self.feature_encoder.conv_layers_spec:
+            receptive_field += (kernel - 1) * stride
+            stride *= layer_stride
+        self.receptive_field_samples = receptive_field
+        self.token_hop_samples = stride
+        if self.receptive_field_samples != spec.receptive_field_samples:
+            raise ValueError("Waveform adapter receptive-field metadata mismatch")
+        if self.token_hop_samples != spec.token_hop_samples:
+            raise ValueError("Waveform adapter hop metadata mismatch")
+
+    def extract(self, waveform: torch.Tensor, *, preset_name: str) -> EmbeddingOutput:
+        preset = get_preset(preset_name)
+        waveform = _single_waveform(waveform)
+        duration_samples = waveform.shape[-1]
+        offsets = [0]
+        if preset.num_phases == 2:
+            if self.token_hop_samples % 2:
+                raise ValueError(
+                    "Two-phase extraction requires an exact half-hop sample offset"
+                )
+            offsets.append(self.token_hop_samples // 2)
+
+        phases: list[tuple[torch.Tensor, torch.Tensor]] = []
+        for offset in offsets:
+            phase = waveform[..., offset:]
+            if phase.shape[-1] < self.receptive_field_samples:
+                phase = F.pad(
+                    phase,
+                    (0, self.receptive_field_samples - phase.shape[-1]),
+                )
+            local_features = self.feature_encoder(phase)
+            tokens = self.encoder_input_proj(local_features).squeeze(0).unsqueeze(0)
+
+            def encode_window(
+                window: torch.Tensor,
+                position_ids: torch.Tensor,
+            ) -> torch.Tensor:
+                return self.encoder(
+                    window,
+                    pos_ids=position_ids,
+                    grid_size=(1, window.shape[1]),
+                )
+
+            embeddings = fuse_context_windows(
+                tokens,
+                max_context_tokens=self.spec.max_context_tokens,
+                overlap=preset.overlap,
+                encode_window=encode_window,
+            )
+            centers = (
+                torch.arange(embeddings.shape[0], device=embeddings.device)
+                * self.token_hop_samples
+                + offset
+                + (self.receptive_field_samples - 1) / 2.0
+            )
+            centers = torch.clamp(centers, max=max(0, duration_samples - 1))
+            phases.append((embeddings, centers * (1000.0 / self.sample_rate)))
+
+        merged, timestamps, scene = merge_phases(phases)
+        return EmbeddingOutput(merged, timestamps, scene)
+
+
+def _normalized_source_state(
+    state_dict: Mapping[str, torch.Tensor],
+) -> dict[str, torch.Tensor]:
+    normalized: dict[str, torch.Tensor] = {}
+    for key, value in state_dict.items():
+        name = str(key)
+        for prefix in ("module.", "model."):
+            if name.startswith(prefix):
+                name = name.removeprefix(prefix)
+        normalized[name] = value
+    return normalized
+
+
+def _load_inference_weights(
+    adapter: HearEncoderAdapter,
+    source_state: Mapping[str, torch.Tensor],
+) -> None:
+    source = _normalized_source_state(source_state)
+    canonical: dict[str, torch.Tensor] = {}
+    missing: list[str] = []
+    for expected_key in adapter.state_dict():
+        source_key = expected_key
+        if expected_key not in source and expected_key.startswith("encoder."):
+            source_key = (
+                f"{adapter.spec.encoder_prefix}.{expected_key.removeprefix('encoder.')}"
+            )
+        if source_key not in source:
+            missing.append(source_key)
+        else:
+            canonical[expected_key] = source[source_key]
+    if missing:
+        raise ValueError(
+            "Checkpoint is missing inference weights: " + ", ".join(missing[:12])
+        )
+    adapter.load_state_dict(canonical, strict=True)
+
+
+def build_encoder_adapter(
+    config: Mapping[str, Any],
+    state_dict: Mapping[str, torch.Tensor],
+) -> HearEncoderAdapter:
+    spec = resolve_adapter_spec(config)
+    if spec.adapter_key == "spectrogram_patch":
+        adapter: HearEncoderAdapter = SpectrogramPatchAdapter(config, spec)
+    elif spec.adapter_key == "waveform_conv":
+        adapter = WaveformConvAdapter(config, spec)
+    else:
+        raise AssertionError(f"Unsupported registered adapter {spec.adapter_key}")
+    _load_inference_weights(adapter, state_dict)
+    adapter.eval()
+    for parameter in adapter.parameters():
+        parameter.requires_grad = False
+    return adapter

config.json

@@ -0,0 +1,52 @@
+{
+  "architectures": [
+    "AudioEmbeddingModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_audio.AudioEmbeddingConfig",
+    "AutoModel": "modeling_audio.AudioEmbeddingModel"
+  },
+  "dtype": "float32",
+  "encoder_kwargs": {
+    "attn_drop_rate": 0.0,
+    "depth": 12,
+    "drop_path_rate": 0.1,
+    "drop_rate": 0.0,
+    "embed_dim": 768,
+    "mlp_ratio": 4.0,
+    "num_heads": 12,
+    "num_patches": 128,
+    "pos_embed_type": "sincos",
+    "qkv_bias": true
+  },
+  "extraction_preset": "overlap50_two_phase",
+  "feature_encoder_kwargs": {},
+  "hidden_size": 768,
+  "model_target": "src.models.best_rq2_module.BestRQ2Module",
+  "model_type": "audio_embeddings",
+  "patch_embed_kwargs": {
+    "embed_dim": 768,
+    "img_size": [
+      128,
+      256
+    ],
+    "in_chans": 1,
+    "patch_size": [
+      16,
+      16
+    ]
+  },
+  "sampling_rate": 16000,
+  "spectrogram_adjustment_mode": "truncate",
+  "spectrogram_kwargs": {
+    "f_max": 8000,
+    "f_min": 0,
+    "hop_length_ms": 39.0625,
+    "n_fft": 2048,
+    "n_mels": 128,
+    "power": 2.0,
+    "sample_rate": 16000,
+    "win_length_ms": 128
+  },
+  "transformers_version": "5.17.0"
+}

configuration_audio.py

@@ -0,0 +1,114 @@
+# MIT License
+# 
+# Copyright (c) 2026 audio-embeddings contributors
+# 
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+# 
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+# 
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+"""Serializable architecture settings; no Hydra execution at inference time."""
+
+from __future__ import annotations
+
+from copy import deepcopy
+from typing import Any
+
+from transformers import PretrainedConfig
+
+from .adapters import SPECTROGRAM_TARGETS, WAVEFORM_TARGETS
+from .extraction import get_preset
+
+
+class AudioEmbeddingConfig(PretrainedConfig):
+    model_type = "audio_embeddings"
+
+    def __init__(
+        self,
+        model_target: str = "src.models.best_rq2_module.BestRQ2Module",
+        encoder_kwargs: dict[str, Any] | None = None,
+        spectrogram_kwargs: dict[str, Any] | None = None,
+        patch_embed_kwargs: dict[str, Any] | None = None,
+        feature_encoder_kwargs: dict[str, Any] | None = None,
+        sampling_rate: int = 16000,
+        spectrogram_adjustment_mode: str = "pad",
+        extraction_preset: str = "overlap50_two_phase",
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        if model_target not in {*SPECTROGRAM_TARGETS, *WAVEFORM_TARGETS}:
+            raise ValueError(f"Unsupported model_target: {model_target!r}")
+        if not isinstance(sampling_rate, int) or sampling_rate <= 0:
+            raise ValueError("sampling_rate must be a positive integer")
+        if spectrogram_adjustment_mode not in {"pad", "truncate"}:
+            raise ValueError("spectrogram_adjustment_mode must be 'pad' or 'truncate'")
+        get_preset(extraction_preset)
+        self.model_target = model_target
+        self.encoder_kwargs = deepcopy(
+            encoder_kwargs
+            or {
+                "embed_dim": 768,
+                "num_patches": 128,
+            }
+        )
+        self.spectrogram_kwargs = deepcopy(spectrogram_kwargs or {})
+        self.patch_embed_kwargs = deepcopy(patch_embed_kwargs or {})
+        self.feature_encoder_kwargs = deepcopy(feature_encoder_kwargs or {})
+        # Remote configuration must contain data, never Python expressions.
+        if isinstance(self.feature_encoder_kwargs.get("conv_layers_spec"), str):
+            raise ValueError(
+                "conv_layers_spec must be a JSON list, not a Python expression"
+            )
+        self.sampling_rate = sampling_rate
+        self.spectrogram_adjustment_mode = spectrogram_adjustment_mode
+        self.extraction_preset = extraction_preset
+        self.hidden_size = int(self.encoder_kwargs["embed_dim"])
+        if self.hidden_size <= 0:
+            raise ValueError("encoder_kwargs.embed_dim must be positive")
+
+    def to_adapter_config(self) -> dict[str, Any]:
+        net: dict[str, Any] = {"encoder": deepcopy(self.encoder_kwargs)}
+        if self.model_target in SPECTROGRAM_TARGETS:
+            net.update(
+                spectrogram=deepcopy(self.spectrogram_kwargs),
+                patch_embed=deepcopy(self.patch_embed_kwargs),
+            )
+        else:
+            net.update(
+                feature_encoder=deepcopy(self.feature_encoder_kwargs),
+                sampling={"sample_rate": self.sampling_rate},
+            )
+        return {
+            "data": {"target_sample_rate": self.sampling_rate},
+            "model": {
+                "_target_": self.model_target,
+                "spectrogram_adjustment_mode": self.spectrogram_adjustment_mode,
+                "net": net,
+            },
+        }
+
+
+AudioEmbeddingConfig.register_for_auto_class()
+
+from .adapters import __name__ as _bundled_adapters  # noqa: F401
+from .extraction import __name__ as _bundled_extraction  # noqa: F401
+from .patch_embed import __name__ as _bundled_patch_embed  # noqa: F401
+from .spectrogram import __name__ as _bundled_spectrogram  # noqa: F401
+from .vit import __name__ as _bundled_vit  # noqa: F401
+from .rope import __name__ as _bundled_rope  # noqa: F401
+from .transformer import __name__ as _bundled_transformer  # noqa: F401
+from .normalization import __name__ as _bundled_normalization  # noqa: F401
+from .waveform_feature_encoder import __name__ as _bundled_waveform_feature_encoder  # noqa: F401

export_manifest.json

@@ -0,0 +1,60 @@
+{
+  "schema_version": 1,
+  "model_name": "BEST-RQ-2",
+  "mock": false,
+  "model_target": "src.models.best_rq2_module.BestRQ2Module",
+  "source": {
+    "huggingface_repository": "ltuncay/BEST-RQ-2",
+    "huggingface_revision": "a03eeb5c4433f4bf7a7e6e8b4724af862789959c",
+    "checkpoint_name": "BEST-RQ-2.safetensors",
+    "training_config_sha256": "9673abdbd7fc0df52803d3ca7703687d6ef271208f11dd2e4f547d31af6296db",
+    "configured_max_steps": 200000,
+    "original_distribution": "AECC 2026 submission",
+    "checkpoint_sha256": "7111465e6c868e3d0b55c5fe9a23dc5069ac80beba8444c5ee9db4691d796899",
+    "global_step": null
+  },
+  "runtime_versions": {
+    "torch": "2.9.1",
+    "torchaudio": "2.9.1",
+    "transformers": "5.17.0",
+    "timm": "1.0.22",
+    "einops": "0.8.1"
+  },
+  "validation": {
+    "max_absolute_error": 0.0,
+    "batch_size": 2,
+    "aecc_component_parity": {
+      "source_revision": "a03eeb5c4433f4bf7a7e6e8b4724af862789959c",
+      "input_samples": 160000,
+      "encoder_output_shape": [
+        1,
+        128,
+        768
+      ],
+      "spectrogram_max_absolute_error": 0.0,
+      "patch_max_absolute_error": 0.0,
+      "encoder_max_absolute_error": 1.6689300537109375e-06
+    }
+  },
+  "files": {
+    "CODE_LICENSE": "8fe9e8b749cd4abedabcb3100df445db899e72192394d14cc2dbf24a40811af6",
+    "README.md": "7537e8c7103f6c735f84dc0bfc781fac2eb8b2f9f831b35eb3f00f7cec15b83f",
+    "adapters.py": "ca0f26826763c0e48b7508da732242bb83adfeb4814e389ed3fde93ad648c728",
+    "config.json": "d37b85c37f17aa954f715b3b3d956c9ed18215ec6ea9565bcd4a403e9b51a134",
+    "configuration_audio.py": "59f3a0b8db0df5af85e677ac33a4431595e9b2eff6d34c1e6a1778dd75a4568d",
+    "extraction.py": "8949389abc52278e24634518f72eae1270bb1a649ad016b71e0d6e8113acaa7c",
+    "feature_extraction_audio.py": "5ed9004727be8b61bd6c0efbeb703b634c20c1b790d9fd59d6a9910438e89b39",
+    "model.safetensors": "f79b5c53ae7749420b8533ec459dcfd5ee0e82858d1df56896a5f74faa845236",
+    "modeling_audio.py": "ba394cd7f80929274090bc1b8986298247f61494db5ed99c0c5469f4e30bb874",
+    "normalization.py": "918b0eb9d547e77ba1adc154deb33d29861532e2a0cd7663497bfc9df6155cef",
+    "patch_embed.py": "825cf2ca0a9d8d10e12f40dc69d2293aabf04b627b926d42314a8b5193e231f7",
+    "preprocessor_config.json": "e37df5cada8190b709335e4c7d45e714617cfcdd532d0e270562225b173b26af",
+    "requirements.txt": "08c9ad278bc06a31d7c7ed3a2255dca48bd63faf29ae7c27635c284d082bb818",
+    "rope.py": "d74bdc7a7dd4ba98414f93c4f1937b95614b57394984d265af5999327efc2d82",
+    "spectrogram.py": "b1e6de8e8ec8220ad9369a765f80a1504b13a6a54aae2ab5bc18cd8dd1e0a894",
+    "transformer.py": "03ea1cb7509669f91003f8cd94edd3660b0c6788e14335ca995c60834a2017d4",
+    "vit.py": "cb6b05b5f617def68e05c40e6bf4d8a24a7b43e5acc4481649e3541bd172e727",
+    "waveform_feature_encoder.py": "69c811624650441e7f46d4fc4e8177c9bac2f65970aef5c6fbd826522f109e22",
+    "LICENSE": "8fe9e8b749cd4abedabcb3100df445db899e72192394d14cc2dbf24a40811af6"
+  }
+}

extraction.py

@@ -0,0 +1,157 @@
+# MIT License
+# 
+# Copyright (c) 2026 audio-embeddings contributors
+# 
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+# 
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+# 
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Callable
+
+import torch
+
+
+@dataclass(frozen=True)
+class ExtractionPreset:
+    name: str
+    overlap: float
+    num_phases: int
+
+
+PRESETS = {
+    preset.name: preset
+    for preset in (
+        ExtractionPreset("native_one_phase", 0.0, 1),
+        ExtractionPreset("overlap25_one_phase", 0.25, 1),
+        ExtractionPreset("overlap50_one_phase", 0.5, 1),
+        ExtractionPreset("native_two_phase", 0.0, 2),
+        ExtractionPreset("overlap25_two_phase", 0.25, 2),
+        ExtractionPreset("overlap50_two_phase", 0.5, 2),
+    )
+}
+
+
+def get_preset(name: str) -> ExtractionPreset:
+    try:
+        return PRESETS[name]
+    except KeyError as error:
+        raise ValueError(
+            f"Unknown extraction preset {name!r}. Expected one of {sorted(PRESETS)}"
+        ) from error
+
+
+def _window_starts(length: int, window: int, overlap: float) -> list[int]:
+    if length <= window:
+        return [0]
+    if not 0.0 <= overlap < 1.0:
+        raise ValueError(f"overlap must be in [0, 1), got {overlap}")
+    stride = max(1, round(window * (1.0 - overlap)))
+    if overlap == 0.0:
+        return list(range(0, length, stride))
+    starts = list(range(0, length - window + 1, stride))
+    final_start = length - window
+    if starts[-1] != final_start:
+        starts.append(final_start)
+    return starts
+
+
+def _positive_triangular_weights(
+    length: int,
+    *,
+    device: torch.device,
+    dtype: torch.dtype,
+) -> torch.Tensor:
+    positions = torch.arange(length, device=device, dtype=dtype)
+    return 1.0 - torch.abs((2.0 * positions) - (length - 1)) / (length + 1)
+
+
+def fuse_context_windows(
+    tokens: torch.Tensor,
+    *,
+    max_context_tokens: int,
+    overlap: float,
+    encode_window: Callable[[torch.Tensor, torch.Tensor], torch.Tensor],
+) -> torch.Tensor:
+    """Contextualize a [frequency, time, dim] grid and fuse repeated tokens."""
+
+    if tokens.ndim != 3:
+        raise ValueError(
+            f"Expected token grid [frequency, time, dim], got {tuple(tokens.shape)}"
+        )
+    frequency, time, dimension = tokens.shape
+    max_time = max_context_tokens // frequency
+    if max_time <= 0:
+        raise ValueError(
+            f"Encoder context {max_context_tokens} cannot hold {frequency} frequency tokens"
+        )
+    if time == 0:
+        raise ValueError("Token grid has no time steps")
+
+    accumulator = torch.zeros_like(tokens)
+    denominator = torch.zeros(time, device=tokens.device, dtype=tokens.dtype)
+    for start in _window_starts(time, max_time, overlap):
+        end = min(time, start + max_time)
+        width = end - start
+        window = tokens[:, start:end, :]
+        flattened = window.reshape(1, frequency * width, dimension)
+        position_ids = torch.arange(
+            frequency * width,
+            device=tokens.device,
+        )
+        encoded = encode_window(flattened, position_ids)
+        if encoded.shape != flattened.shape:
+            raise ValueError(
+                "Encoder changed the token-grid shape: "
+                f"expected {tuple(flattened.shape)}, got {tuple(encoded.shape)}"
+            )
+        encoded_grid = encoded.reshape(frequency, width, dimension)
+        weights = _positive_triangular_weights(
+            width,
+            device=tokens.device,
+            dtype=tokens.dtype,
+        )
+        accumulator[:, start:end, :] += encoded_grid * weights.view(1, -1, 1)
+        denominator[start:end] += weights
+
+    if torch.any(denominator <= 0):
+        raise RuntimeError("At least one token received zero overlap weight")
+    fused = accumulator / denominator.view(1, -1, 1)
+    return fused.mean(dim=0)
+
+
+def merge_phases(
+    phases: list[tuple[torch.Tensor, torch.Tensor]],
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    if not phases:
+        raise ValueError("At least one embedding phase is required")
+    embedding_dim = phases[0][0].shape[-1]
+    for embeddings, timestamps in phases:
+        if embeddings.ndim != 2 or embeddings.shape[-1] != embedding_dim:
+            raise ValueError("All phase embeddings must have shape [time, dimension]")
+        if timestamps.ndim != 1 or timestamps.shape[0] != embeddings.shape[0]:
+            raise ValueError("Each phase needs one timestamp per embedding")
+
+    merged_embeddings = torch.cat([phase[0] for phase in phases], dim=0)
+    merged_timestamps = torch.cat([phase[1] for phase in phases], dim=0)
+    order = torch.argsort(merged_timestamps, stable=True)
+    scene = torch.stack([embeddings.mean(dim=0) for embeddings, _ in phases]).mean(
+        dim=0
+    )
+    return merged_embeddings[order], merged_timestamps[order], scene

feature_extraction_audio.py

@@ -0,0 +1,114 @@
+# MIT License
+# 
+# Copyright (c) 2026 audio-embeddings contributors
+# 
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+# 
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+# 
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+"""Prepare mono waveforms without changing the learned frontend's numerics."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import numpy as np
+import torch
+from transformers.feature_extraction_sequence_utils import SequenceFeatureExtractor
+from transformers.feature_extraction_utils import BatchFeature
+
+
+class AudioEmbeddingFeatureExtractor(SequenceFeatureExtractor):
+    model_input_names = ["input_values", "attention_mask"]
+
+    def __init__(
+        self,
+        sampling_rate: int = 16000,
+        padding_value: float = 0.0,
+        return_attention_mask: bool = True,
+        **kwargs: Any,
+    ) -> None:
+        if not isinstance(sampling_rate, int) or sampling_rate <= 0:
+            raise ValueError("sampling_rate must be a positive integer")
+        kwargs.pop("feature_size", None)
+        super().__init__(
+            feature_size=1,
+            sampling_rate=sampling_rate,
+            padding_value=padding_value,
+            return_attention_mask=return_attention_mask,
+            **kwargs,
+        )
+
+    def __call__(
+        self,
+        raw_speech: Any,
+        *,
+        sampling_rate: int | None = None,
+        padding: bool | str = True,
+        max_length: int | None = None,
+        truncation: bool = False,
+        pad_to_multiple_of: int | None = None,
+        return_attention_mask: bool | None = None,
+        return_tensors: str | None = "pt",
+    ) -> BatchFeature:
+        if sampling_rate != self.sampling_rate:
+            raise ValueError(
+                f"Pass sampling_rate={self.sampling_rate}; got {sampling_rate}. "
+                "Resample audio to the model rate before calling the feature extractor."
+            )
+        if isinstance(raw_speech, torch.Tensor):
+            raw_speech = raw_speech.detach().cpu().float().numpy()
+        if isinstance(raw_speech, np.ndarray):
+            if raw_speech.ndim not in {1, 2}:
+                raise ValueError(
+                    "Expected mono audio [samples] or a batch [batch, samples]"
+                )
+            batch = [raw_speech] if raw_speech.ndim == 1 else list(raw_speech)
+        elif isinstance(raw_speech, (list, tuple)) and len(raw_speech):
+            batch = [raw_speech] if np.isscalar(raw_speech[0]) else list(raw_speech)
+        else:
+            raise ValueError("Provide a nonempty waveform or batch of mono waveforms")
+        waveforms = []
+        for waveform in batch:
+            if isinstance(waveform, torch.Tensor):
+                waveform = waveform.detach().cpu().float().numpy()
+            array = np.asarray(waveform, dtype=np.float32)
+            if array.ndim != 1 or array.size == 0 or not np.isfinite(array).all():
+                raise ValueError(
+                    "Each waveform must be a nonempty, finite, mono 1-D array"
+                )
+            waveforms.append(array)
+        if not waveforms:
+            raise ValueError("The audio batch cannot be empty")
+        if max_length is not None and max_length <= 0:
+            raise ValueError("max_length must be positive")
+        return self.pad(
+            BatchFeature({"input_values": waveforms}),
+            padding=padding,
+            max_length=max_length,
+            truncation=truncation,
+            pad_to_multiple_of=pad_to_multiple_of,
+            return_attention_mask=(
+                self.return_attention_mask
+                if return_attention_mask is None
+                else return_attention_mask
+            ),
+            return_tensors=return_tensors,
+        )
+
+
+AudioEmbeddingFeatureExtractor.register_for_auto_class("AutoFeatureExtractor")

model.safetensors

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

modeling_audio.py

@@ -0,0 +1,182 @@
+# MIT License
+# 
+# Copyright (c) 2026 audio-embeddings contributors
+# 
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+# 
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+# 
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+"""Trainable audio embeddings with the same extraction policy as HEAR."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import torch
+from torch.nn.utils.rnn import pad_sequence
+from transformers import PreTrainedModel
+from transformers.utils import ModelOutput
+
+from .adapters import SpectrogramPatchAdapter, WaveformConvAdapter
+from .adapters import resolve_adapter_spec
+from .configuration_audio import AudioEmbeddingConfig
+
+
+@dataclass
+class AudioEmbeddingOutput(ModelOutput):
+    last_hidden_state: torch.Tensor | None = None
+    pooler_output: torch.Tensor | None = None
+    attention_mask: torch.Tensor | None = None
+    timestamps_ms: torch.Tensor | None = None
+
+
+class AudioEmbeddingModel(PreTrainedModel):
+    config_class = AudioEmbeddingConfig
+    base_model_prefix = "adapter"
+    main_input_name = "input_values"
+    # RoPE modules are shared by every attention block. Save all buffer keys,
+    # so loading does not need special tied-buffer handling.
+    _supports_assign_param_buffer = False
+
+    def __init__(self, config: AudioEmbeddingConfig) -> None:
+        super().__init__(config)
+        adapter_config = config.to_adapter_config()
+        # Transformers 5 loads under a default meta device, but torchaudio's
+        # filter-bank constructors need real values. Build on CPU; HF subsequently
+        # loads the checkpoint tensors onto the requested device/dtype.
+        with torch.device("cpu"):
+            spec = resolve_adapter_spec(adapter_config)
+            adapter_type = (
+                SpectrogramPatchAdapter
+                if spec.adapter_key == "spectrogram_patch"
+                else WaveformConvAdapter
+            )
+            self.adapter = adapter_type(adapter_config, spec)
+        self.post_init()
+
+    def _init_weights(self, module: torch.nn.Module) -> None:
+        """Preserve initialization performed by the research components themselves."""
+
+    def save_pretrained(
+        self, save_directory: str | Path, *args: Any, **kwargs: Any
+    ) -> None:
+        if kwargs.get("state_dict") is None:
+            kwargs["state_dict"] = {
+                key: value.detach().clone().contiguous()
+                for key, value in self.state_dict().items()
+            }
+        return super().save_pretrained(save_directory, *args, **kwargs)
+
+    def forward(
+        self,
+        input_values: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        return_dict: bool | None = None,
+    ) -> AudioEmbeddingOutput | tuple[torch.Tensor, ...]:
+        if input_values.ndim != 2 or min(input_values.shape) <= 0:
+            raise ValueError(
+                "input_values must have shape [batch, samples] with nonempty axes"
+            )
+        if (
+            not input_values.is_floating_point()
+            or not torch.isfinite(input_values).all()
+        ):
+            raise ValueError(
+                "input_values must contain finite floating-point waveforms"
+            )
+        if attention_mask is None:
+            lengths = [input_values.shape[1]] * input_values.shape[0]
+        else:
+            if attention_mask.shape != input_values.shape:
+                raise ValueError(
+                    "attention_mask must have the same shape as input_values"
+                )
+            if not torch.all((attention_mask == 0) | (attention_mask == 1)):
+                raise ValueError("attention_mask must contain only zeros and ones")
+            lengths_tensor = attention_mask.long().sum(dim=1)
+            expected = (
+                torch.arange(input_values.shape[1], device=attention_mask.device)[None]
+                < lengths_tensor[:, None]
+            )
+            if not torch.equal(attention_mask.bool(), expected) or torch.any(
+                lengths_tensor == 0
+            ):
+                raise ValueError(
+                    "attention_mask must describe nonempty, right-padded waveforms"
+                )
+            lengths = lengths_tensor.tolist()
+        # Clear non-buffer RoPE caches between calls: inference-mode caches cannot
+        # be reused for autograd, and .to(device/dtype) does not move these caches.
+        rope = self.adapter.encoder.rope
+        if rope is not None:
+            for name in ("cached_cos_sin", "cached_cos_sin_h", "cached_cos_sin_w"):
+                if hasattr(rope, name):
+                    setattr(rope, name, None)
+        outputs = []
+        for waveform, length in zip(input_values, lengths):
+            if isinstance(self.adapter, SpectrogramPatchAdapter):
+                minimum = self.adapter.spectrogram.mel_spec.n_fft // 2 + 1
+                if length < minimum:
+                    raise ValueError(
+                        f"Audio requires at least {minimum} samples for this spectrogram; got {length}"
+                    )
+            outputs.append(
+                self.adapter.extract(
+                    waveform[:length], preset_name=self.config.extraction_preset
+                )
+            )
+        hidden = pad_sequence(
+            [item.timestamp_embeddings for item in outputs], batch_first=True
+        )
+        frame_lengths = torch.tensor(
+            [item.timestamp_embeddings.shape[0] for item in outputs],
+            device=hidden.device,
+        )
+        frame_mask = (
+            torch.arange(hidden.shape[1], device=hidden.device)[None]
+            < frame_lengths[:, None]
+        )
+        result = AudioEmbeddingOutput(
+            last_hidden_state=hidden,
+            pooler_output=torch.stack([item.scene_embedding for item in outputs]),
+            attention_mask=frame_mask.long(),
+            timestamps_ms=pad_sequence(
+                [item.timestamps_ms for item in outputs],
+                batch_first=True,
+                padding_value=-1.0,
+            ),
+        )
+        return (
+            result
+            if (self.config.return_dict if return_dict is None else return_dict)
+            else result.to_tuple()
+        )
+
+
+AudioEmbeddingModel.register_for_auto_class("AutoModel")
+
+from .adapters import __name__ as _bundled_adapters  # noqa: F401
+from .extraction import __name__ as _bundled_extraction  # noqa: F401
+from .patch_embed import __name__ as _bundled_patch_embed  # noqa: F401
+from .spectrogram import __name__ as _bundled_spectrogram  # noqa: F401
+from .vit import __name__ as _bundled_vit  # noqa: F401
+from .rope import __name__ as _bundled_rope  # noqa: F401
+from .transformer import __name__ as _bundled_transformer  # noqa: F401
+from .normalization import __name__ as _bundled_normalization  # noqa: F401
+from .waveform_feature_encoder import __name__ as _bundled_waveform_feature_encoder  # noqa: F401

normalization.py

@@ -0,0 +1,42 @@
+# MIT License
+# 
+# Copyright (c) 2026 audio-embeddings contributors
+# 
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+# 
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+# 
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class MixedPrecisionRMSNorm(nn.RMSNorm):
+    """RMSNorm that keeps its master weight while matching activation dtype.
+
+    PyTorch's fused RMSNorm requires the input and weight to have the same dtype.
+    Mixed-precision training keeps parameters in FP32, so use a differentiable
+    low-precision view of the weight for the operation itself.
+    """
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        weight = self.weight
+        if weight is not None and weight.dtype != input.dtype:
+            weight = weight.to(dtype=input.dtype)
+        return F.rms_norm(input, self.normalized_shape, weight, self.eps)

patch_embed.py

@@ -0,0 +1,246 @@
+# MIT License
+# 
+# Copyright (c) 2026 audio-embeddings contributors
+# 
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+# 
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+# 
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+from collections.abc import Sequence
+from math import prod
+
+import torch
+import torch.nn as nn
+from timm.layers import PatchEmbed as TimmPatchEmbed
+
+
+def _is_power_of_two(value: int) -> bool:
+    return value > 0 and value & (value - 1) == 0
+
+
+def _build_hmlp_kernel_schedule(
+    patch_size: tuple[int, int],
+) -> tuple[tuple[int, int], ...]:
+    """Build non-overlapping aggregation steps for one target patch.
+
+    The paper starts from 4x4 subpatches and doubles both axes until reaching
+    16x16. For rectangular patches, an axis stops growing once it reaches its
+    target while the other axis keeps doubling.
+    """
+    patch_height, patch_width = patch_size
+    if not _is_power_of_two(patch_height) or not _is_power_of_two(patch_width):
+        raise ValueError(
+            "hMLP patch dimensions must be powers of two so each axis can be "
+            f"doubled exactly, got {patch_size}"
+        )
+
+    first_kernel = (min(4, patch_height), min(4, patch_width))
+    schedule = [first_kernel]
+    current_height, current_width = first_kernel
+
+    while (current_height, current_width) != patch_size:
+        kernel_height = 2 if current_height < patch_height else 1
+        kernel_width = 2 if current_width < patch_width else 1
+        schedule.append((kernel_height, kernel_width))
+        current_height *= kernel_height
+        current_width *= kernel_width
+
+    return tuple(schedule)
+
+
+def _resolve_hmlp_kernel_schedule(
+    patch_size: tuple[int, int],
+    kernel_schedule: Sequence[Sequence[int]] | None,
+) -> tuple[tuple[int, int], ...]:
+    if kernel_schedule is None:
+        return _build_hmlp_kernel_schedule(patch_size)
+
+    schedule: list[tuple[int, int]] = []
+    for stage_index, stage in enumerate(kernel_schedule):
+        try:
+            stage_values = tuple(stage)
+        except TypeError as error:
+            raise ValueError(
+                "Each hMLP stage must contain [kernel_height, kernel_width], "
+                f"stage {stage_index} has {stage}"
+            ) from error
+        if len(stage_values) != 2:
+            raise ValueError(
+                "Each hMLP stage must contain [kernel_height, kernel_width], "
+                f"stage {stage_index} has {stage}"
+            )
+        try:
+            kernel_size = tuple(int(value) for value in stage_values)
+        except (TypeError, ValueError) as error:
+            raise ValueError(
+                f"hMLP stage kernels must be integers, got {stage_values}"
+            ) from error
+        if any(
+            isinstance(original, bool) or normalized != original
+            for normalized, original in zip(kernel_size, stage_values)
+        ):
+            raise ValueError(f"hMLP stage kernels must be integers, got {stage_values}")
+        if any(value <= 0 for value in kernel_size):
+            raise ValueError(f"hMLP stage kernels must be positive, got {kernel_size}")
+        schedule.append(kernel_size)
+
+    if not schedule:
+        raise ValueError("hMLP kernel_schedule must contain at least one stage")
+
+    aggregated_patch_size = tuple(
+        prod(kernel_size[axis] for kernel_size in schedule) for axis in range(2)
+    )
+    if aggregated_patch_size != patch_size:
+        raise ValueError(
+            "hMLP kernel_schedule stages must multiply to patch_size; "
+            f"got {aggregated_patch_size} from {tuple(schedule)}, expected {patch_size}"
+        )
+
+    return tuple(schedule)
+
+
+class HierarchicalMLPPatchEmbed(nn.Module):
+    """hMLP patch stem with independent, hierarchical patch aggregation."""
+
+    def __init__(
+        self,
+        img_size: tuple[int, int] = (128, 256),
+        patch_size: tuple[int, int] = (16, 16),
+        in_chans: int = 1,
+        embed_dim: int = 768,
+        bias: bool = True,
+        kernel_schedule: Sequence[Sequence[int]] | None = None,
+    ) -> None:
+        super().__init__()
+        self.img_size = tuple(img_size)
+        self.patch_size = tuple(patch_size)
+        self.in_chans = in_chans
+        self.embed_dim = embed_dim
+        self.bias = bias
+        self.kernel_schedule = _resolve_hmlp_kernel_schedule(
+            self.patch_size,
+            kernel_schedule,
+        )
+        self.num_patches = (self.img_size[0] // self.patch_size[0]) * (
+            self.img_size[1] // self.patch_size[1]
+        )
+
+        hidden_dim = max(1, embed_dim // 4)
+        layers: list[nn.Module] = []
+        input_dim = in_chans
+        for stage_index, kernel_size in enumerate(self.kernel_schedule):
+            is_last = stage_index == len(self.kernel_schedule) - 1
+            output_dim = embed_dim if is_last else hidden_dim
+            layers.extend(
+                [
+                    nn.Conv2d(
+                        input_dim,
+                        output_dim,
+                        kernel_size=kernel_size,
+                        stride=kernel_size,
+                        bias=bias,
+                    ),
+                    nn.SyncBatchNorm(output_dim),
+                ]
+            )
+            if not is_last:
+                layers.append(nn.GELU())
+            input_dim = output_dim
+
+        self.proj = nn.Sequential(*layers)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if x.ndim != 4:
+            raise ValueError(f"Expected input with shape [B, C, H, W], got {x.shape}")
+        if x.shape[1] != self.in_chans:
+            raise ValueError(
+                f"Expected {self.in_chans} input channels, got {x.shape[1]}"
+            )
+        if x.shape[2] % self.patch_size[0] != 0 or x.shape[3] % self.patch_size[1] != 0:
+            raise ValueError(
+                "Input spatial dimensions must be divisible by the hMLP patch "
+                f"size {self.patch_size}, got {tuple(x.shape[2:])}"
+            )
+
+        return self.proj(x).flatten(2).transpose(1, 2)
+
+
+class PatchEmbed(nn.Module):
+    """
+    2D Image to Patch Embedding.
+
+    Args:
+        img_size (tuple[int, int]): Input image size (H, W).
+        patch_size (tuple[int, int]): Patch size (H, W).
+        in_chans (int): Number of input channels.
+        embed_dim (int): Embedding dimension.
+    """
+
+    def __init__(
+        self,
+        img_size: tuple[int, int] = (128, 256),
+        patch_size: tuple[int, int] = (16, 16),
+        in_chans: int = 1,
+        embed_dim: int = 768,
+        bias: bool = True,
+        stem_type: str = "linear",
+        hmlp_kernel_schedule: Sequence[Sequence[int]] | None = None,
+    ):
+        super().__init__()
+        self.img_size = tuple(img_size)
+        self.patch_size = tuple(patch_size)
+        self.in_chans = in_chans
+        self.embed_dim = embed_dim
+        self.bias = bias
+        self.stem_type = stem_type.strip().lower().replace("-", "_")
+
+        if self.stem_type == "linear":
+            self.patch_embed = TimmPatchEmbed(
+                img_size=img_size,
+                patch_size=patch_size,
+                in_chans=in_chans,
+                embed_dim=embed_dim,
+                flatten=True,
+                bias=bias,
+                strict_img_size=False,
+            )
+        elif self.stem_type == "hmlp":
+            self.patch_embed = HierarchicalMLPPatchEmbed(
+                img_size=img_size,
+                patch_size=patch_size,
+                in_chans=in_chans,
+                embed_dim=embed_dim,
+                bias=bias,
+                kernel_schedule=hmlp_kernel_schedule,
+            )
+        else:
+            raise ValueError(
+                f"Unknown stem_type={stem_type!r}; expected 'linear' or 'hmlp'"
+            )
+        self.num_patches = self.patch_embed.num_patches
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass.
+
+        Args:
+            x (torch.Tensor): Input tensor [B, C, H, W].
+
+        Returns:
+            torch.Tensor: Patch embeddings [B, N, D].
+        """
+        return self.patch_embed(x)

preprocessor_config.json

@@ -0,0 +1,11 @@
+{
+  "auto_map": {
+    "AutoFeatureExtractor": "feature_extraction_audio.AudioEmbeddingFeatureExtractor"
+  },
+  "feature_extractor_type": "AudioEmbeddingFeatureExtractor",
+  "feature_size": 1,
+  "padding_side": "right",
+  "padding_value": 0.0,
+  "return_attention_mask": true,
+  "sampling_rate": 16000
+}

requirements.txt

@@ -0,0 +1,5 @@
+torch>=2.9.1
+torchaudio>=2.9.1
+timm>=0.9
+einops>=0.7
+transformers>=4.57,<6

rope.py

@@ -0,0 +1,307 @@
+# MIT License
+# 
+# Copyright (c) 2026 audio-embeddings contributors
+# 
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+# 
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+# 
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .normalization import MixedPrecisionRMSNorm
+
+
+def _build_qk_norm(norm_type: str, head_dim: int) -> nn.Module:
+    normalized = norm_type.strip().lower().replace("_", "")
+    if normalized == "layernorm":
+        return nn.LayerNorm(head_dim)
+    if normalized == "rmsnorm":
+        return MixedPrecisionRMSNorm(head_dim)
+    raise ValueError(
+        f"Unknown qk_norm_type={norm_type!r}; expected 'layernorm' or 'rmsnorm'"
+    )
+
+
+class RotaryEmbedding2D(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        max_res: Tuple[int, int] = (128, 256),
+        temperature: float = 10000.0,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.max_h, self.max_w = max_res
+        self.temperature = temperature
+
+        # Check if dim is divisible by 4 (since we split into 2 for H/W, and each needs 2 for complex)
+        assert dim % 4 == 0, "Embedding dimension must be divisible by 4 for 2D RoPE"
+
+        dim_h = dim // 2
+        dim_w = dim // 2
+
+        # Generate frequencies for H and W
+        # inv_freq_h: [dim_h // 2]
+        inv_freq_h = 1.0 / (temperature ** (torch.arange(0, dim_h, 2).float() / dim_h))
+        inv_freq_w = 1.0 / (temperature ** (torch.arange(0, dim_w, 2).float() / dim_w))
+
+        self.register_buffer("inv_freq_h", inv_freq_h)
+        self.register_buffer("inv_freq_w", inv_freq_w)
+
+        # Cache
+        self.cached_cos_sin_h = None
+        self.cached_cos_sin_w = None
+
+    def _update_cache(self, h: int, w: int, device: torch.device, dtype: torch.dtype):
+        # Generate grid
+        # We need to support arbitrary positions, but usually we just precompute for max_res
+        # or compute on the fly for the given indices.
+        # Let's compute for max_res and index into it.
+
+        if self.cached_cos_sin_h is None or self.cached_cos_sin_h[0].shape[0] < h:
+            t_h = torch.arange(h, device=device, dtype=dtype)
+            freqs_h = torch.einsum("i,j->ij", t_h, self.inv_freq_h)  # [H, dim_h/2]
+            emb_h = torch.cat((freqs_h, freqs_h), dim=-1)  # [H, dim_h]
+            self.cached_cos_sin_h = (emb_h.cos(), emb_h.sin())
+
+        if self.cached_cos_sin_w is None or self.cached_cos_sin_w[0].shape[0] < w:
+            t_w = torch.arange(w, device=device, dtype=dtype)
+            freqs_w = torch.einsum("i,j->ij", t_w, self.inv_freq_w)  # [W, dim_w/2]
+            emb_w = torch.cat((freqs_w, freqs_w), dim=-1)  # [W, dim_w]
+            self.cached_cos_sin_w = (emb_w.cos(), emb_w.sin())
+
+    def forward(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        pos_ids: torch.Tensor,
+        grid_size: Tuple[int, int],
+    ):
+        # q, k: [B, num_heads, N, head_dim]
+        # pos_ids: [B, N] or [N] (indices of patches)
+        # grid_size: (H, W) - original grid size to decode pos_ids
+
+        B, num_heads, N, D = q.shape
+        H_grid, W_grid = grid_size
+
+        # Decode pos_ids to (h, w)
+        # pos_ids are indices in flattened grid [0, H*W-1]
+        # h = pos_ids // W_grid
+        # w = pos_ids % W_grid
+
+        h_idx = pos_ids.div(W_grid, rounding_mode="floor")  # [B, N]
+        w_idx = pos_ids % W_grid  # [B, N]
+
+        # Ensure cache is large enough
+        self._update_cache(H_grid, W_grid, q.device, q.dtype)
+
+        # Fetch cos/sin for H and W
+        # cos_h: [B, N, dim_h]
+        # We need to gather from cached [max_h, dim_h] using h_idx
+
+        # Handle shared pos_ids (if [N])
+        if h_idx.ndim == 1:
+            h_idx = h_idx.unsqueeze(0).expand(B, -1)
+            w_idx = w_idx.unsqueeze(0).expand(B, -1)
+
+        cos_h = F.embedding(h_idx, self.cached_cos_sin_h[0])  # [B, N, dim_h]
+        sin_h = F.embedding(h_idx, self.cached_cos_sin_h[1])
+        cos_w = F.embedding(w_idx, self.cached_cos_sin_w[0])  # [B, N, dim_w]
+        sin_w = F.embedding(w_idx, self.cached_cos_sin_w[1])
+
+        # Split q, k into halves
+        # q: [B, num_heads, N, D] -> [B, N, num_heads, D] for easier manipulation?
+        # Usually RoPE is applied on [B, num_heads, N, D] or [N, B, num_heads, D]
+        # Let's keep [B, num_heads, N, D]
+
+        dim_half = D // 2
+        q_h, q_w = q.split(dim_half, dim=-1)
+        k_h, k_w = k.split(dim_half, dim=-1)
+
+        # Apply RoPE
+        # We need to reshape cos/sin to broadcast over num_heads
+        # cos_h: [B, N, dim_h] -> [B, 1, N, dim_h]
+        cos_h = cos_h.unsqueeze(1)
+        sin_h = sin_h.unsqueeze(1)
+        cos_w = cos_w.unsqueeze(1)
+        sin_w = sin_w.unsqueeze(1)
+
+        q_h_rot = self._apply_rotary(q_h, cos_h, sin_h)
+        k_h_rot = self._apply_rotary(k_h, cos_h, sin_h)
+
+        q_w_rot = self._apply_rotary(q_w, cos_w, sin_w)
+        k_w_rot = self._apply_rotary(k_w, cos_w, sin_w)
+
+        q_rot = torch.cat((q_h_rot, q_w_rot), dim=-1)
+        k_rot = torch.cat((k_h_rot, k_w_rot), dim=-1)
+
+        return q_rot, k_rot
+
+    def _apply_rotary(
+        self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
+    ) -> torch.Tensor:
+        # x: [B, num_heads, N, dim_half]
+        # cos, sin: [B, 1, N, dim_half]
+        # Standard RoPE rotation:
+        # x = [x1, x2]
+        # out = [x1*cos - x2*sin, x1*sin + x2*cos]
+        # This assumes pairs are adjacent.
+        # My inv_freq generation: cat(freqs, freqs).
+        # This corresponds to x = [x_first_half, x_second_half] pairing?
+        # Usually RoPE pairs even/odd or first/second half.
+        # "The standard implementation ... pairs feature i with i + d/2"
+        # My emb generation: cat(freqs, freqs) -> [f0, f1, ..., f0, f1, ...] ? No.
+        # freqs is [0, 2, ...]
+        # cat(freqs, freqs) -> [f0, f2, ..., f0, f2, ...]
+        # So it expects x to be split into two halves and rotated.
+        # rotate_half(x) = [-x2, x1]
+
+        return (x * cos) + (self._rotate_half(x) * sin)
+
+    def _rotate_half(self, x: torch.Tensor) -> torch.Tensor:
+        x1, x2 = x.chunk(2, dim=-1)
+        return torch.cat((-x2, x1), dim=-1)
+
+
+class RotaryEmbedding1D(nn.Module):
+    """Apply RoPE over the temporal axis using the full attention head."""
+
+    def __init__(
+        self,
+        dim: int,
+        max_seq_len: int = 256,
+        temperature: float = 10000.0,
+    ) -> None:
+        super().__init__()
+        assert dim % 2 == 0, "Embedding dimension must be divisible by 2 for 1D RoPE"
+
+        self.dim = dim
+        self.max_seq_len = max_seq_len
+        self.temperature = temperature
+        inv_freq = 1.0 / (temperature ** (torch.arange(0, dim, 2).float() / dim))
+        self.register_buffer("inv_freq", inv_freq)
+        self.cached_cos_sin: tuple[torch.Tensor, torch.Tensor] | None = None
+
+    def _update_cache(
+        self,
+        seq_len: int,
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> None:
+        if self.cached_cos_sin is None or self.cached_cos_sin[0].shape[0] < seq_len:
+            positions = torch.arange(seq_len, device=device, dtype=dtype)
+            freqs = torch.einsum("i,j->ij", positions, self.inv_freq)
+            embedding = torch.cat((freqs, freqs), dim=-1)
+            self.cached_cos_sin = (embedding.cos(), embedding.sin())
+
+    def forward(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        pos_ids: torch.Tensor,
+        grid_size: Tuple[int, int],
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Rotate using each patch's time coordinate from the flattened grid."""
+        batch_size = q.shape[0]
+        _, width = grid_size
+        time_indices = pos_ids % width
+
+        if time_indices.ndim == 1:
+            time_indices = time_indices.unsqueeze(0).expand(batch_size, -1)
+
+        self._update_cache(width, q.device, q.dtype)
+        cos = F.embedding(time_indices, self.cached_cos_sin[0]).unsqueeze(1)
+        sin = F.embedding(time_indices, self.cached_cos_sin[1]).unsqueeze(1)
+
+        return self._apply_rotary(q, cos, sin), self._apply_rotary(k, cos, sin)
+
+    def _apply_rotary(
+        self,
+        x: torch.Tensor,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+    ) -> torch.Tensor:
+        return (x * cos) + (self._rotate_half(x) * sin)
+
+    def _rotate_half(self, x: torch.Tensor) -> torch.Tensor:
+        x1, x2 = x.chunk(2, dim=-1)
+        return torch.cat((-x2, x1), dim=-1)
+
+
+class RoPEAttention(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int = 8,
+        qkv_bias: bool = False,
+        proj_bias: bool = True,
+        attn_drop: float = 0.0,
+        proj_drop: float = 0.0,
+        rope: Optional[RotaryEmbedding1D | RotaryEmbedding2D] = None,
+        qk_norm: bool = False,
+        qk_norm_type: str = "layernorm",
+    ):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = head_dim**-0.5
+        self.rope = rope
+        self.qk_norm = qk_norm
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.q_norm = _build_qk_norm(qk_norm_type, head_dim) if qk_norm else None
+        self.k_norm = _build_qk_norm(qk_norm_type, head_dim) if qk_norm else None
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim, bias=proj_bias)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        pos_ids: torch.Tensor = None,
+        grid_size: Tuple[int, int] = None,
+    ) -> torch.Tensor:
+        B, N, C = x.shape
+        qkv = (
+            self.qkv(x)
+            .reshape(B, N, 3, self.num_heads, C // self.num_heads)
+            .permute(2, 0, 3, 1, 4)
+        )
+        q, k, v = qkv[0], qkv[1], qkv[2]  # [B, num_heads, N, head_dim]
+
+        if self.rope is not None and pos_ids is not None and grid_size is not None:
+            q, k = self.rope(q, k, pos_ids, grid_size)
+
+        if self.q_norm is not None and self.k_norm is not None:
+            q = self.q_norm(q)
+            k = self.k_norm(k)
+
+        x = F.scaled_dot_product_attention(
+            q,
+            k,
+            v,
+            dropout_p=self.attn_drop.p if self.training else 0.0,
+            scale=self.scale,
+        )
+        x = x.transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x

spectrogram.py

@@ -0,0 +1,220 @@
+# MIT License
+# 
+# Copyright (c) 2026 audio-embeddings contributors
+# 
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+# 
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+# 
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torchaudio
+
+
+class Spectrogram(nn.Module):
+    """
+    Mel-frequency audio representation with optional temporal derivatives.
+
+    Args:
+        sample_rate (int): Sample rate of the audio.
+        n_fft (int): Size of FFT.
+        win_length (Optional[int]): Window length. Defaults to n_fft.
+        win_length_ms (Optional[float]): Window length in milliseconds. Overrides win_length if provided.
+        hop_length (Optional[int]): Hop length. Defaults to win_length // 2.
+        hop_length_ms (Optional[float]): Hop length in milliseconds. Overrides hop_length if provided.
+        n_mels (int): Number of mel filterbanks.
+        f_min (float): Minimum frequency.
+        f_max (Optional[float]): Maximum frequency.
+        power (float): Power of the magnitude.
+        representation (str): ``log_mel`` or phase-aware ``complex_mel``.
+        complex_log_dynamic_range_db (float): Retained complex-mel log range.
+        add_delta (bool): Append one temporal derivative per base channel.
+        add_delta_delta (bool): Append a second derivative per base channel.
+        delta_win_length (int): Odd regression window used for derivatives.
+    """
+
+    def __init__(
+        self,
+        sample_rate: int = 32000,
+        n_fft: int = 4096,
+        win_length: Optional[int] = None,
+        win_length_ms: Optional[float] = None,
+        hop_length: Optional[int] = None,
+        hop_length_ms: Optional[float] = None,
+        n_mels: int = 128,
+        f_min: float = 0.0,
+        f_max: Optional[float] = None,
+        power: float = 2.0,
+        representation: str = "log_mel",
+        complex_log_dynamic_range_db: float = 80.0,
+        add_delta: bool = False,
+        add_delta_delta: bool = False,
+        delta_win_length: int = 5,
+    ):
+        super().__init__()
+
+        if win_length is None:
+            if win_length_ms is None:
+                win_length = n_fft
+            else:
+                win_length = int(sample_rate * win_length_ms / 1000)
+
+        if hop_length is None:
+            if hop_length_ms is None:
+                hop_length = win_length // 2
+            else:
+                hop_length = int(sample_rate * hop_length_ms / 1000)
+
+        representation = representation.strip().lower().replace("-", "_")
+        if representation not in {"log_mel", "complex_mel"}:
+            raise ValueError(
+                "representation must be 'log_mel' or 'complex_mel', "
+                f"got {representation!r}"
+            )
+        if complex_log_dynamic_range_db <= 0.0:
+            raise ValueError(
+                "complex_log_dynamic_range_db must be positive, "
+                f"got {complex_log_dynamic_range_db}"
+            )
+        if delta_win_length < 3 or delta_win_length % 2 == 0:
+            raise ValueError(
+                f"delta_win_length must be an odd integer >= 3, got {delta_win_length}"
+            )
+
+        self.representation = representation
+        self.complex_log_dynamic_range_db = float(complex_log_dynamic_range_db)
+        self.add_delta = bool(add_delta)
+        self.add_delta_delta = bool(add_delta_delta)
+        self.delta_win_length = int(delta_win_length)
+        self.base_output_channels = 1 if representation == "log_mel" else 2
+        derivative_orders = int(self.add_delta) + int(self.add_delta_delta)
+        self.output_channels = self.base_output_channels * (1 + derivative_orders)
+
+        if representation == "log_mel":
+            self.mel_spec = torchaudio.transforms.MelSpectrogram(
+                sample_rate=sample_rate,
+                n_fft=n_fft,
+                win_length=win_length,
+                hop_length=hop_length,
+                n_mels=n_mels,
+                f_min=f_min,
+                f_max=f_max,
+                power=power,
+                normalized=True,
+            )
+            self.amplitude_to_db = torchaudio.transforms.AmplitudeToDB()
+        else:
+            # Keep the historical ``mel_spec`` attribute for sample-rate and
+            # hop-length discovery in callbacks and HEAR adapters.
+            self.mel_spec = torchaudio.transforms.Spectrogram(
+                n_fft=n_fft,
+                win_length=win_length,
+                hop_length=hop_length,
+                power=None,
+                normalized=True,
+            )
+            self.mel_spec.sample_rate = sample_rate
+            self.amplitude_to_db = None
+            mel_fb = torchaudio.functional.melscale_fbanks(
+                n_freqs=n_fft // 2 + 1,
+                f_min=f_min,
+                f_max=float(sample_rate / 2 if f_max is None else f_max),
+                n_mels=n_mels,
+                sample_rate=sample_rate,
+                norm=None,
+                mel_scale="htk",
+            )
+            self.register_buffer("mel_fb", mel_fb)
+
+    def _complex_mel(self, x: torch.Tensor) -> torch.Tensor:
+        if x.shape[1] != 1:
+            raise ValueError(
+                "complex_mel expects mono waveform input [B, 1, T], "
+                f"got {tuple(x.shape)}"
+            )
+
+        complex_spec = self.mel_spec(x[:, 0])
+        mel_fb = self.mel_fb.to(dtype=complex_spec.dtype)
+        complex_mel = torch.matmul(
+            complex_spec.transpose(-1, -2),
+            mel_fb,
+        ).transpose(-1, -2)
+
+        magnitude = complex_mel.abs()
+        eps = torch.finfo(magnitude.dtype).eps
+        reference = magnitude.amax(dim=(-2, -1), keepdim=True)
+        relative_magnitude = magnitude / reference.clamp_min(eps)
+        floor_ratio = 10.0 ** (-self.complex_log_dynamic_range_db / 20.0)
+        log_magnitude_db = 20.0 * torch.log10(relative_magnitude.clamp_min(floor_ratio))
+        log_radius = (
+            log_magnitude_db + self.complex_log_dynamic_range_db
+        ) / self.complex_log_dynamic_range_db
+        log_radius = log_radius.clamp_(0.0, 1.0)
+        log_radius = torch.where(
+            reference > eps, log_radius, torch.zeros_like(log_radius)
+        )
+
+        unit_phase = complex_mel / magnitude.clamp_min(eps)
+        compressed = unit_phase * log_radius
+        return torch.stack([compressed.real, compressed.imag], dim=1)
+
+    def _append_deltas(self, spec: torch.Tensor) -> torch.Tensor:
+        if not self.add_delta and not self.add_delta_delta:
+            return spec
+
+        delta = torchaudio.functional.compute_deltas(
+            spec,
+            win_length=self.delta_win_length,
+            mode="replicate",
+        )
+        channels = [spec]
+        if self.add_delta:
+            channels.append(delta)
+        if self.add_delta_delta:
+            channels.append(
+                torchaudio.functional.compute_deltas(
+                    delta,
+                    win_length=self.delta_win_length,
+                    mode="replicate",
+                )
+            )
+        return torch.cat(channels, dim=1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass.
+
+        Args:
+            x (torch.Tensor): Input waveform [B, C, T] or [B, T].
+
+        Returns:
+            torch.Tensor: Mel-frequency features [B, C, F, T].
+        """
+        if x.ndim == 2:
+            x = x.unsqueeze(1)
+        if x.ndim != 3:
+            raise ValueError(f"Expected waveform [B, C, T], got {tuple(x.shape)}")
+
+        if self.representation == "complex_mel":
+            spec = self._complex_mel(x)
+        else:
+            spec = self.mel_spec(x)
+            spec = self.amplitude_to_db(spec)
+
+        return self._append_deltas(spec)

transformer.py

@@ -0,0 +1,274 @@
+# MIT License
+# 
+# Copyright (c) 2026 audio-embeddings contributors
+# 
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+# 
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+# 
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import Callable
+
+import torch
+import torch.nn as nn
+from timm.layers import DropPath
+from timm.layers import Mlp
+
+from .normalization import MixedPrecisionRMSNorm
+from .rope import RoPEAttention
+from .rope import RotaryEmbedding1D
+from .rope import RotaryEmbedding2D
+
+
+class FullAttentionResidual(nn.Module):
+    """Softmax attention over all preceding outputs along model depth."""
+
+    def __init__(self, dim: int) -> None:
+        super().__init__()
+        self.dim = dim
+        self.norm = MixedPrecisionRMSNorm(dim)
+        self.query = nn.Parameter(torch.zeros(dim))
+
+    def forward(self, values: Sequence[torch.Tensor]) -> torch.Tensor:
+        if not values:
+            raise ValueError("FullAttentionResidual requires at least one value")
+
+        expected_shape = values[0].shape
+        if len(expected_shape) != 3 or expected_shape[-1] != self.dim:
+            raise ValueError(
+                "FullAttentionResidual values must have shape [B, N, D] with "
+                f"D={self.dim}, got {expected_shape}"
+            )
+        for index, value in enumerate(values[1:], start=1):
+            if value.shape != expected_shape:
+                raise ValueError(
+                    "FullAttentionResidual values must have identical shapes; "
+                    f"value 0 has {expected_shape}, value {index} has {value.shape}"
+                )
+
+        stacked_values = torch.stack(tuple(values), dim=0)
+        keys = self.norm(stacked_values)
+        logits = torch.einsum("d,l b n d->l b n", self.query, keys)
+        weights = logits.softmax(dim=0)
+        return torch.einsum("l b n,l b n d->b n d", weights, stacked_values)
+
+
+def build_norm_layer(
+    *,
+    dim: int,
+    norm_type: str = "layernorm",
+    norm_layer: Callable[[int], nn.Module] | None = None,
+    norm_eps: float | None = None,
+) -> nn.Module:
+    """Build a token-channel normalization layer.
+
+    `norm_layer` is kept for backward compatibility with direct Python
+    construction. Configs should prefer `norm_type` so choices are explicit in
+    Hydra overrides and experiment files.
+    """
+    if norm_layer is not None:
+        if norm_eps is not None:
+            raise ValueError("Set norm_eps or a custom norm_layer, not both")
+        return norm_layer(dim)
+
+    if norm_eps is not None and norm_eps <= 0:
+        raise ValueError(f"norm_eps must be positive, got {norm_eps}")
+    kwargs = {} if norm_eps is None else {"eps": norm_eps}
+    normalized = norm_type.strip().lower().replace("_", "")
+    if normalized == "layernorm":
+        return nn.LayerNorm(dim, **kwargs)
+    if normalized == "rmsnorm":
+        return MixedPrecisionRMSNorm(dim, **kwargs)
+
+    raise ValueError(
+        f"Unknown norm_type={norm_type!r}; expected 'layernorm' or 'rmsnorm'"
+    )
+
+
+def build_gelu_mlp(
+    *,
+    dim: int,
+    mlp_ratio: float,
+    act_layer: type[nn.Module],
+    drop: float,
+    bias: bool = True,
+) -> nn.Module:
+    """Build the baseline ViT feed-forward layer used before ablations."""
+    return Mlp(
+        in_features=dim,
+        hidden_features=int(dim * mlp_ratio),
+        act_layer=act_layer,
+        bias=bias,
+        drop=drop,
+    )
+
+
+class SwiGLUMlp(nn.Module):
+    """SwiGLU feed-forward layer for transformer ablations.
+
+    The hidden dimension is controlled by `mlp_ratio` in the same place as the
+    baseline MLP. GLU variants normally use a smaller ratio such as 8/3 because
+    they have two input projections before the output projection.
+    """
+
+    def __init__(
+        self,
+        *,
+        dim: int,
+        hidden_features: int,
+        drop: float,
+        bias: bool = True,
+    ) -> None:
+        super().__init__()
+        self.gate = nn.Linear(dim, hidden_features, bias=bias)
+        self.value = nn.Linear(dim, hidden_features, bias=bias)
+        self.proj = nn.Linear(hidden_features, dim, bias=bias)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = nn.functional.silu(self.gate(x)) * self.value(x)
+        x = self.drop(x)
+        x = self.proj(x)
+        return self.drop(x)
+
+
+def build_mlp(
+    *,
+    dim: int,
+    mlp_ratio: float,
+    mlp_type: str = "gelu_mlp",
+    act_layer: type[nn.Module],
+    drop: float,
+    bias: bool = True,
+) -> nn.Module:
+    """Build the transformer feed-forward layer selected by config."""
+    hidden_features = int(dim * mlp_ratio)
+    normalized = mlp_type.strip().lower().replace("-", "_")
+
+    if normalized == "gelu_mlp":
+        return build_gelu_mlp(
+            dim=dim,
+            mlp_ratio=mlp_ratio,
+            act_layer=act_layer,
+            drop=drop,
+            bias=bias,
+        )
+    if normalized == "swiglu":
+        return SwiGLUMlp(
+            dim=dim,
+            hidden_features=hidden_features,
+            drop=drop,
+            bias=bias,
+        )
+
+    raise ValueError(f"Unknown mlp_type={mlp_type!r}; expected 'gelu_mlp' or 'swiglu'")
+
+
+class RoPEBlock(nn.Module):
+    """Pre-norm transformer block with RoPE-aware attention."""
+
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int,
+        mlp_ratio: float = 4.0,
+        qkv_bias: bool = False,
+        proj_bias: bool = True,
+        mlp_bias: bool = True,
+        qk_norm: bool = False,
+        qk_norm_type: str = "layernorm",
+        mlp_type: str = "gelu_mlp",
+        proj_drop: float = 0.0,
+        attn_drop: float = 0.0,
+        drop_path: float = 0.0,
+        act_layer: type[nn.Module] = nn.GELU,
+        norm_type: str = "layernorm",
+        norm_layer: Callable[[int], nn.Module] | None = None,
+        rope: RotaryEmbedding1D | RotaryEmbedding2D | None = None,
+        residual_type: str = "standard",
+        norm_eps: float | None = None,
+    ) -> None:
+        super().__init__()
+        self.residual_type = residual_type.strip().lower().replace("-", "_")
+        if self.residual_type not in {"standard", "full_attnres"}:
+            raise ValueError(
+                f"Unknown residual_type={residual_type!r}; expected 'standard' "
+                "or 'full_attnres'"
+            )
+        self.norm1 = build_norm_layer(
+            dim=dim,
+            norm_type=norm_type,
+            norm_layer=norm_layer,
+            norm_eps=norm_eps,
+        )
+        self.attn = RoPEAttention(
+            dim,
+            num_heads=num_heads,
+            qkv_bias=qkv_bias,
+            proj_bias=proj_bias,
+            attn_drop=attn_drop,
+            proj_drop=proj_drop,
+            rope=rope,
+            qk_norm=qk_norm,
+            qk_norm_type=qk_norm_type,
+        )
+        self.norm2 = build_norm_layer(
+            dim=dim,
+            norm_type=norm_type,
+            norm_layer=norm_layer,
+            norm_eps=norm_eps,
+        )
+        self.mlp = build_mlp(
+            dim=dim,
+            mlp_ratio=mlp_ratio,
+            mlp_type=mlp_type,
+            act_layer=act_layer,
+            drop=proj_drop,
+            bias=mlp_bias,
+        )
+        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+        if self.residual_type == "full_attnres":
+            self.attention_residual = FullAttentionResidual(dim)
+            self.mlp_residual = FullAttentionResidual(dim)
+        else:
+            self.attention_residual = None
+            self.mlp_residual = None
+
+    def attention_output(
+        self,
+        x: torch.Tensor,
+        pos_ids: torch.Tensor | None = None,
+        grid_size: tuple[int, int] | None = None,
+    ) -> torch.Tensor:
+        return self.drop_path(
+            self.attn(self.norm1(x), pos_ids=pos_ids, grid_size=grid_size)
+        )
+
+    def mlp_output(self, x: torch.Tensor) -> torch.Tensor:
+        return self.drop_path(self.mlp(self.norm2(x)))
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        pos_ids: torch.Tensor | None = None,
+        grid_size: tuple[int, int] | None = None,
+    ) -> torch.Tensor:
+        x = x + self.attention_output(x, pos_ids=pos_ids, grid_size=grid_size)
+        x = x + self.mlp_output(x)
+        return x

vit.py

@@ -0,0 +1,298 @@
+# MIT License
+# 
+# Copyright (c) 2026 audio-embeddings contributors
+# 
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+# 
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+# 
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+from collections.abc import Mapping
+from typing import Any, Optional, Tuple
+
+import torch
+import torch.nn as nn
+from timm.layers import build_sincos2d_pos_embed
+
+from .rope import RotaryEmbedding1D
+from .rope import RotaryEmbedding2D
+from .transformer import build_norm_layer
+from .transformer import FullAttentionResidual
+from .transformer import RoPEBlock
+
+
+def vit_config_with_patch_geometry(
+    config: Mapping[str, Any],
+    *,
+    img_size: tuple[int, int],
+    patch_size: tuple[int, int],
+) -> dict[str, Any]:
+    """Return a ViT config aligned with the patch embedding's actual geometry."""
+    resolved = dict(config)
+    resolved["img_size"] = tuple(img_size)
+    resolved["patch_size"] = tuple(patch_size)
+    resolved.setdefault("rope_mode", "auto")
+    return resolved
+
+
+class ViT(nn.Module):
+    """
+    Vision Transformer with support for RoPE and 2D positional embeddings.
+
+    Args:
+        embed_dim (int): Embedding dimension.
+        depth (int): Number of transformer blocks.
+        num_heads (int): Number of attention heads.
+        mlp_ratio (float): Ratio of MLP hidden dim to embedding dim.
+        qkv_bias (bool): Enable bias for QKV projections.
+        drop_rate (float): Dropout rate.
+        attn_drop_rate (float): Attention dropout rate.
+        drop_path_rate (float): Stochastic depth rate.
+        norm_layer (nn.Module): Normalization layer.
+        norm_eps (float | None): Explicit normalization epsilon; None preserves
+            the normalization implementation's default.
+        act_layer (nn.Module): Activation layer.
+        num_patches (int): Total number of patches (used for learnable/sincos pos embed).
+        img_size (tuple[int, int]): Input image size (H, W).
+        patch_size (tuple[int, int]): Patch size (H, W).
+        pos_embed_type (str): Type of positional embedding ("rope", "sincos", "learnable").
+        rope_mode (str | None): RoPE axes ("2d" or temporal "1d"). ``None``
+            or ``"auto"`` selects 1D when the patch height equals the image height,
+            otherwise 2D.
+    """
+
+    def __init__(
+        self,
+        embed_dim: int = 768,
+        depth: int = 12,
+        num_heads: int = 12,
+        mlp_ratio: float = 4.0,
+        mlp_type: str = "gelu_mlp",
+        qkv_bias: bool = True,
+        proj_bias: bool = True,
+        mlp_bias: bool = True,
+        qk_norm: bool = False,
+        qk_norm_type: str = "layernorm",
+        drop_rate: float = 0.0,
+        attn_drop_rate: float = 0.0,
+        drop_path_rate: float = 0.0,
+        norm_layer: nn.Module | None = None,
+        norm_type: str = "layernorm",
+        act_layer: nn.Module = nn.GELU,
+        num_patches: int = 128,
+        img_size: tuple[int, int] = (128, 256),
+        patch_size: tuple[int, int] = (16, 16),
+        pos_embed_type: str = "rope",
+        rope_mode: str | None = "2d",
+        residual_type: str = "standard",
+        norm_eps: float | None = None,
+    ):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.num_patches = num_patches
+        self.grid_size = (img_size[0] // patch_size[0], img_size[1] // patch_size[1])
+        self.pos_embed_type = pos_embed_type
+        requested_rope_mode = (
+            "auto" if rope_mode is None else rope_mode.strip().lower().replace("-", "_")
+        )
+        if requested_rope_mode == "auto":
+            self.rope_mode = "1d" if patch_size[0] == img_size[0] else "2d"
+        else:
+            self.rope_mode = requested_rope_mode
+        self.norm_type = norm_type
+        self.mlp_type = mlp_type
+        self.residual_type = residual_type.strip().lower().replace("-", "_")
+        if self.residual_type not in {"standard", "full_attnres"}:
+            raise ValueError(
+                f"Unknown residual_type={residual_type!r}; expected 'standard' "
+                "or 'full_attnres'"
+            )
+
+        # Positional Embeddings
+        if pos_embed_type == "rope":
+            head_dim = embed_dim // num_heads
+            if self.rope_mode == "2d":
+                self.rope = RotaryEmbedding2D(dim=head_dim, max_res=self.grid_size)
+            elif self.rope_mode == "1d":
+                self.rope = RotaryEmbedding1D(
+                    dim=head_dim,
+                    max_seq_len=self.grid_size[1],
+                )
+            else:
+                raise ValueError(
+                    f"Unknown rope_mode: {rope_mode!r}; expected 'auto', '1d', or '2d'"
+                )
+            self.pos_embed = None
+        elif pos_embed_type == "sincos":
+            self.rope = None
+            # build_sincos2d_pos_embed(feat_shape, dim, ...)
+            # We assume grid_size matches num_patches
+            pos_embed = build_sincos2d_pos_embed(self.grid_size, embed_dim)
+            self.register_buffer("pos_embed", pos_embed.unsqueeze(0))  # [1, N, D]
+        elif pos_embed_type == "learnable":
+            self.rope = None
+            self.pos_embed = nn.Parameter(torch.zeros(1, num_patches, embed_dim))
+            nn.init.trunc_normal_(self.pos_embed, std=0.02)
+        else:
+            raise ValueError(f"Unknown pos_embed_type: {pos_embed_type}")
+
+        # Stochastic Depth
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]
+
+        self.blocks = nn.ModuleList(
+            [
+                RoPEBlock(
+                    dim=embed_dim,
+                    num_heads=num_heads,
+                    mlp_ratio=mlp_ratio,
+                    mlp_type=mlp_type,
+                    qkv_bias=qkv_bias,
+                    proj_bias=proj_bias,
+                    mlp_bias=mlp_bias,
+                    qk_norm=qk_norm,
+                    qk_norm_type=qk_norm_type,
+                    proj_drop=drop_rate,
+                    attn_drop=attn_drop_rate,
+                    drop_path=dpr[i],
+                    norm_type=norm_type,
+                    norm_layer=norm_layer,
+                    norm_eps=norm_eps,
+                    act_layer=act_layer,
+                    rope=self.rope,
+                    residual_type=self.residual_type,
+                )
+                for i in range(depth)
+            ]
+        )
+
+        self.norm = build_norm_layer(
+            dim=embed_dim,
+            norm_type=norm_type,
+            norm_layer=norm_layer,
+            norm_eps=norm_eps,
+        )
+        self.output_residual = (
+            FullAttentionResidual(embed_dim)
+            if self.residual_type == "full_attnres"
+            else None
+        )
+
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m: nn.Module) -> None:
+        if isinstance(m, nn.Linear):
+            nn.init.trunc_normal_(m.weight, std=0.02)
+            if m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+        elif isinstance(m, nn.RMSNorm):
+            nn.init.constant_(m.weight, 1.0)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        pos_ids: Optional[torch.Tensor] = None,
+        add_pos_embed: bool = True,
+        grid_size: Optional[Tuple[int, int]] = None,
+    ) -> torch.Tensor:
+        """
+        Forward pass.
+
+        Args:
+            x (torch.Tensor): Input tensor [B, N, D].
+            pos_ids (Optional[torch.Tensor]): Positional indices [B, N] or [N].
+            add_pos_embed (bool): Whether to add positional embeddings (for non-RoPE).
+            grid_size (Optional[Tuple[int, int]]): Grid size for RoPE/PosEmbed.
+
+        Returns:
+            torch.Tensor: Output tensor [B, N, D].
+        """
+        # Determine grid size
+        if grid_size is None:
+            if pos_ids is None:
+                # Infer from x assuming full sequence
+                B, N, D = x.shape
+                H_grid = self.grid_size[0]
+                W_grid = N // H_grid
+                current_grid_size = (H_grid, W_grid)
+            else:
+                # Cannot infer, use default (might be wrong if variable length)
+                current_grid_size = self.grid_size
+        else:
+            current_grid_size = grid_size
+
+        if self.pos_embed_type != "rope" and add_pos_embed:
+            if pos_ids is not None:
+                # Select positional embeddings
+                if pos_ids.ndim == 1:
+                    # Shared pos_ids across batch
+                    pos_embed = self.pos_embed[:, pos_ids, :]  # [1, N_subset, D]
+                else:
+                    # Different pos_ids per sample
+                    pos_embed = self.pos_embed.expand(x.shape[0], -1, -1)
+                    pos_embed = torch.gather(
+                        pos_embed,
+                        1,
+                        pos_ids.unsqueeze(-1).expand(-1, -1, self.embed_dim),
+                    )
+                x = x + pos_embed
+            else:
+                # Assume full sequence
+                if x.shape[1] == self.num_patches:
+                    x = x + self.pos_embed
+                elif (
+                    self.pos_embed is not None and x.shape[1] <= self.pos_embed.shape[1]
+                ):
+                    x = x + self.pos_embed[:, : x.shape[1], :]
+
+        # For RoPE, we need pos_ids. If not provided, generate them.
+        if self.pos_embed_type == "rope" and pos_ids is None:
+            device = x.device
+            # We need to generate pos_ids for the current grid
+            # If we inferred current_grid_size, we should use it.
+            # pos_ids should be 0..N-1
+            B, N, D = x.shape
+            pos_ids = torch.arange(N, device=device)
+
+        if self.residual_type == "full_attnres":
+            values = [x]
+            for block in self.blocks:
+                if block.attention_residual is None or block.mlp_residual is None:
+                    raise RuntimeError(
+                        "Full AttnRes block is missing depth aggregators"
+                    )
+                attention_input = block.attention_residual(values)
+                values.append(
+                    block.attention_output(
+                        attention_input,
+                        pos_ids=pos_ids,
+                        grid_size=current_grid_size,
+                    )
+                )
+                mlp_input = block.mlp_residual(values)
+                values.append(block.mlp_output(mlp_input))
+
+            if self.output_residual is None:
+                raise RuntimeError("Full AttnRes ViT is missing its output aggregator")
+            x = self.output_residual(values)
+        else:
+            for block in self.blocks:
+                x = block(x, pos_ids=pos_ids, grid_size=current_grid_size)
+
+        x = self.norm(x)
+        return x

waveform_feature_encoder.py

@@ -0,0 +1,171 @@
+# MIT License
+# 
+# Copyright (c) 2026 audio-embeddings contributors
+# 
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+# 
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+# 
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+from __future__ import annotations
+
+from typing import Sequence
+
+import torch
+from einops import rearrange
+from einops.layers.torch import Rearrange
+from torch import nn
+
+
+def _parse_conv_layers_spec(
+    conv_layers_spec: str | Sequence[Sequence[int]] | Sequence[tuple[int, int, int]],
+) -> list[tuple[int, int, int]]:
+    if isinstance(conv_layers_spec, str):
+        # Config-driven expression style used by wavjepa, e.g.
+        # "[(512, 10, 5)] + [(512, 3, 2)] * 4 + [(512, 2, 2)]"
+        parsed = eval(conv_layers_spec, {"__builtins__": {}}, {})  # noqa: S307
+    else:
+        parsed = conv_layers_spec
+
+    out: list[tuple[int, int, int]] = []
+    for layer in parsed:
+        if len(layer) != 3:
+            raise ValueError(f"Invalid conv layer spec {layer}, expected (dim, k, s)")
+        dim, kernel, stride = layer
+        out.append((int(dim), int(kernel), int(stride)))
+    if len(out) == 0:
+        raise ValueError("conv_layers_spec must contain at least one layer")
+    return out
+
+
+class WaveformFeatureEncoder(nn.Module):
+    """
+    Convolutional waveform feature encoder that outputs a token sequence.
+
+    Input shape: [B, C, T]
+    Output shape: [B, N, F]
+    """
+
+    def __init__(
+        self,
+        conv_layers_spec: str
+        | Sequence[Sequence[int]]
+        | Sequence[
+            tuple[int, int, int]
+        ] = "[(512, 10, 5)] + [(512, 3, 2)] * 4 + [(512, 2, 2)]",
+        in_channels: int = 1,
+        dropout: float = 0.0,
+        mode: str = "default",
+        conv_bias: bool = False,
+        depthwise: bool = False,
+    ) -> None:
+        super().__init__()
+        if mode not in {"default", "layer_norm"}:
+            raise ValueError(
+                f"Unknown mode='{mode}', expected 'default' or 'layer_norm'"
+            )
+        self.conv_layers_spec = _parse_conv_layers_spec(conv_layers_spec)
+        self.in_channels = in_channels
+        self.depthwise = depthwise
+
+        layers: list[nn.Module] = []
+        in_dim = in_channels
+        for idx, (out_dim, kernel, stride) in enumerate(self.conv_layers_spec):
+            layers.append(
+                self._make_block(
+                    in_dim=in_dim,
+                    out_dim=out_dim,
+                    kernel=kernel,
+                    stride=stride,
+                    dropout=dropout,
+                    mode=mode,
+                    conv_bias=conv_bias,
+                    depthwise=depthwise,
+                    is_first=idx == 0,
+                )
+            )
+            in_dim = out_dim
+
+        self.cnn = nn.Sequential(*layers)
+        self.embedding_dim = self.conv_layers_spec[-1][0]
+
+    @staticmethod
+    def _make_block(
+        in_dim: int,
+        out_dim: int,
+        kernel: int,
+        stride: int,
+        dropout: float,
+        mode: str,
+        conv_bias: bool,
+        depthwise: bool,
+        is_first: bool,
+    ) -> nn.Module:
+        if depthwise:
+            if out_dim % in_dim != 0:
+                raise ValueError(
+                    "Depthwise mode requires out_dim to be a multiple of in_dim, "
+                    f"got out_dim={out_dim}, in_dim={in_dim}"
+                )
+            conv = nn.Conv1d(
+                in_dim,
+                out_dim,
+                kernel_size=kernel,
+                stride=stride,
+                bias=conv_bias,
+                groups=in_dim,
+            )
+        else:
+            conv = nn.Conv1d(
+                in_dim,
+                out_dim,
+                kernel_size=kernel,
+                stride=stride,
+                bias=conv_bias,
+            )
+        nn.init.kaiming_normal_(conv.weight)
+
+        if mode == "layer_norm":
+            return nn.Sequential(
+                conv,
+                nn.Dropout(p=dropout),
+                Rearrange("... c t -> ... t c"),
+                nn.LayerNorm(out_dim, elementwise_affine=True),
+                Rearrange("... t c -> ... c t"),
+                nn.GELU(),
+            )
+
+        if mode == "default" and is_first:
+            return nn.Sequential(
+                conv,
+                nn.Dropout(p=dropout),
+                nn.GroupNorm(out_dim, out_dim, affine=True),
+                nn.GELU(),
+            )
+
+        return nn.Sequential(conv, nn.Dropout(p=dropout), nn.GELU())
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.cnn(x)
+        return rearrange(x, "b f n -> b n f")
+
+    def total_patches(self, time_samples: int) -> int:
+        n = int(time_samples)
+        for _, kernel, stride in self.conv_layers_spec:
+            if n < kernel:
+                return 0
+            n = (n - kernel) // stride + 1
+        return int(n)