Add config.json, README, and CPU inference example

README.md

@@ -1,3 +1,81 @@
----
-license: mit
----
+---
+license: mit
+tags:
+  - audio-jepa
+  - jepa
+  - audio-representation-learning
+  - self-supervised
+paper:
+  title: "Audio-JEPA: Joint-Embedding Predictive Architecture for Audio Representation Learning"
+  authors:
+    - Ludovic Tuncay
+    - Etienne Labbé
+    - Emmanouil Benetos
+    - Thomas Pellegrini
+  venue: "ICME 2025"
+  arxiv: "2507.02915"
+  hal: "hal-05128180"
+---
+
+# Audio-JEPA: pretrained encoder
+
+Weights for the encoder trained in *Audio-JEPA: Joint-Embedding Predictive Architecture for Audio Representation Learning* (Tuncay, Labbé, Benetos, Pellegrini, ICME 2025). Source code: [`LudovicTuncay/Audio-JEPA`](https://github.com/LudovicTuncay/Audio-JEPA).
+
+## Files
+
+| File | Purpose |
+|------|---------|
+| `JEPA.ckpt` | PyTorch Lightning checkpoint (state_dict + trainer metadata). |
+| `config.json` | Standalone documentation of the architecture and audio pipeline. Machine-readable. |
+| `README.md` | This document. |
+
+## Model summary
+
+- Encoder: `VisionTransformer` (ViT-Base, 12 layers, 768-dim, 12 heads).
+- Input: log-mel spectrogram of shape `(target_time_bins=256, n_mels=128)` from a 10 s mono waveform at 32 kHz.
+- Patchification: `(16, 16)`, giving a grid `8 x 16 = 128` patches.
+- Each patch spans ~625 ms of audio × 16 mel bins.
+- Output: `(128, 768)` embeddings per 10 s clip.
+- Effective temporal resolution: **1.6 positions/s** (the 12.8 tokens/s counts 8 temporal × 16 frequency patches per second).
+
+## Minimal inference (CPU, no `flash-attn` install required)
+
+The training repo depends on `flash-attn`, which requires CUDA to build. For inference-only use, `flash_attn.modules.mha.MHA` can be substituted with a torch-native equivalent that matches the checkpoint's parameter names (`qkv`, `proj`). See `inference_example.py` in this repository for a ~150-line standalone script.
+
+```bash
+# 1. Clone the source code (needed for the ViT class)
+git clone --depth 1 https://github.com/LudovicTuncay/Audio-JEPA.git
+
+# 2. Install a lean set of deps
+pip install torch torchaudio numpy huggingface_hub
+
+# 3. Run the example (downloads JEPA.ckpt on first run)
+python inference_example.py --audio-jepa-src ./Audio-JEPA
+```
+
+The script prints the loading diagnostics (should be `0 missing, 0 unexpected`) and the embedding shape.
+
+## Domain fit: where the model excels vs where it doesn't
+
+Audio-JEPA is designed to learn a **generic audio representation**. The 16×16 patch shape is a compromise across speech, music, and environmental sounds. Consequences:
+
+- **Strong on**: music, environmental sounds, audio captioning, general audio tagging.
+- **Weaker on**: speech-only downstream tasks (see the paper's X-ARES tables). Speech-specific SSL models such as `wav2vec 2.0`, HuBERT and the Whisper encoder currently outperform Audio-JEPA on speech-centric benchmarks.
+- **A follow-up from the same author**, [`BEST-RQ-2`](https://huggingface.co/ltuncay/BEST-RQ-2) (Tuncay, Labbé, Pellegrini, Interspeech 2026, [arXiv 2606.30700](https://arxiv.org/abs/2606.30700)), combines the encoder-predictor decomposition of Audio-JEPA with BEST-RQ discrete targets. It yields substantially better cross-domain results while sharing the exact same encoder inference speed.
+
+## Citation
+
+```bibtex
+@inproceedings{tuncay2025audio,
+  title  = {Audio-JEPA: Joint-Embedding Predictive Architecture for Audio Representation Learning},
+  author = {Tuncay, Ludovic and Labb{\'e}, Etienne and Benetos, Emmanouil and Pellegrini, Thomas},
+  booktitle = {ICME 2025},
+  address = {Nantes, France},
+  year   = {2025},
+  url    = {https://hal.science/hal-05128180}
+}
+```
+
+## License
+
+MIT.

config.json

@@ -0,0 +1,57 @@
+{
+  "model_type": "audio-jepa",
+  "architecture": "VisionTransformer",
+  "citation": "Tuncay, Labbé, Benetos, Pellegrini. Audio-JEPA: Joint-Embedding Predictive Architecture for Audio Representation Learning. ICME 2025 (arXiv 2507.02915)",
+  "audio": {
+    "sample_rate": 32000,
+    "clip_length_s": 10,
+    "n_mels": 128,
+    "target_time_bins": 256,
+    "mel_spec_backend": "torchaudio.compliance.kaldi.fbank",
+    "f_min": 20,
+    "f_max": null,
+    "log_fbank": true,
+    "window": "hanning"
+  },
+  "patch_embed": {
+    "patch_size": [16, 16],
+    "in_chans": 1,
+    "num_patches": 128,
+    "grid_time": 8,
+    "grid_freq": 16
+  },
+  "encoder": {
+    "embed_dim": 768,
+    "depth": 12,
+    "num_heads": 12,
+    "mlp_ratio": 4.0,
+    "qkv_bias": true,
+    "use_flash_attn": true,
+    "norm_layer": "LayerNorm",
+    "cls_token": false,
+    "pos_embed": "2d_sincos"
+  },
+  "output": {
+    "shape_per_10s_clip": [128, 768],
+    "temporal_positions_per_second": 1.6,
+    "tokens_per_second": 12.8,
+    "note": "Each patch spans 16 mel bins x 16 time bins ≈ 625 ms. The 12.8 Hz rate counts 8 temporal x 16 frequency patches per second; effective temporal resolution is 1.6 positions/s."
+  },
+  "training": {
+    "framework": "PyTorch Lightning + Hydra",
+    "pretraining_corpus": "AudioSet unlabeled (~5.3k h after silence filtering)",
+    "target_module": "target_encoder (EMA of encoder)",
+    "predictor": {
+      "type": "VisionTransformerPredictor",
+      "depth": 6,
+      "num_heads": 12,
+      "mlp_ratio": 4.0
+    },
+    "loss": "norm_mse (with norm_pix_loss=True)",
+    "note": "Only the encoder is used downstream. Predictor + target encoder are discarded at inference."
+  },
+  "known_speech_gap": {
+    "description": "This model is designed for generic audio representation. The 16x16 patch shape reflects a compromise across speech, music, and environmental sounds. Speech-only downstream performance is lower than speech-specific SSL models (wav2vec2, HuBERT, Whisper). See the paper's X-ARES tables and the follow-up BEST-RQ-2 (arXiv 2606.30700) for a quantitative comparison.",
+    "suggested_alternative_for_speech": "For speech-heavy downstream tasks, consider BEST-RQ-2 (same author, arXiv 2606.30700, MIT weights at ltuncay/BEST-RQ-2) which shares the encoder architecture, or wait for the upcoming BEST-RQ-2.1 with explicit speech-oriented improvements."
+  }
+}

inference_example.py

@@ -0,0 +1,230 @@
+"""Minimal standalone inference for Audio-JEPA (ltuncay/Audio-JEPA).
+
+Loads the JEPA.ckpt weights, instantiates the ViT encoder with the correct
+input shape, and produces (T, D) embeddings from a mono waveform.
+
+Runs on CPU without installing flash-attn CUDA kernels: a small shim
+substitutes flash_attn.modules.mha.MHA with a torch-native equivalent that
+matches the checkpoint parameter names (qkv, proj).
+
+Requires the upstream code cloned locally:
+
+    git clone --depth 1 https://github.com/LudovicTuncay/Audio-JEPA.git /path/to/audio-jepa
+
+Usage:
+
+    python inference_example.py --audio-jepa-src /path/to/audio-jepa
+
+or, from a directory containing this file next to ``audio-jepa/``:
+
+    python inference_example.py
+"""
+
+from __future__ import annotations
+
+import argparse
+import sys
+import time
+import types
+from importlib.machinery import ModuleSpec
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torchaudio
+from huggingface_hub import hf_hub_download
+
+# Config values matching the shipped checkpoint (confirmed by the author).
+SAMPLE_RATE = 32_000
+CLIP_LENGTH_S = 10
+N_MELS = 128
+TARGET_TIME_BINS = 256
+PATCH_SIZE = (16, 16)
+EMBED_DIM = 768
+DEPTH = 12
+NUM_HEADS = 12
+MLP_RATIO = 4.0
+
+
+def install_flash_attn_shim() -> None:
+    """Replace ``flash_attn.modules.mha.MHA`` with a CPU-friendly torch class.
+
+    Matches the checkpoint's parameter naming (``qkv``, ``proj``) so
+    ``load_state_dict(..., strict=True)`` succeeds without any CUDA build.
+    """
+
+    class CpuMultiHeadAttention(nn.Module):
+        def __init__(
+            self,
+            embed_dim: int,
+            num_heads: int,
+            dropout: float = 0.0,
+            qkv_proj_bias: bool = True,
+            use_flash_attn: bool = False,
+            **_ignore,
+        ) -> None:
+            super().__init__()
+            assert embed_dim % num_heads == 0
+            self.embed_dim = embed_dim
+            self.num_heads = num_heads
+            self.head_dim = embed_dim // num_heads
+            self.dropout = dropout
+            self.qkv = nn.Linear(embed_dim, 3 * embed_dim, bias=qkv_proj_bias)
+            self.proj = nn.Linear(embed_dim, embed_dim, bias=True)
+
+        def forward(self, x: torch.Tensor) -> torch.Tensor:
+            b, n, d = x.shape
+            qkv = self.qkv(x).reshape(b, n, 3, self.num_heads, self.head_dim)
+            q, k, v = qkv.permute(2, 0, 3, 1, 4).unbind(0)
+            out = torch.nn.functional.scaled_dot_product_attention(
+                q, k, v, dropout_p=self.dropout if self.training else 0.0
+            )
+            return self.proj(out.transpose(1, 2).reshape(b, n, d))
+
+    def make(name: str) -> types.ModuleType:
+        m = types.ModuleType(name)
+        m.__spec__ = ModuleSpec(name, loader=None)
+        return m
+
+    flash_attn = make("flash_attn")
+    flash_attn.__version__ = "0.0.0-cpu-shim"
+    modules = make("flash_attn.modules")
+    mha = make("flash_attn.modules.mha")
+    mha.MHA = CpuMultiHeadAttention
+    sys.modules.update(
+        {
+            "flash_attn": flash_attn,
+            "flash_attn.modules": modules,
+            "flash_attn.modules.mha": mha,
+        }
+    )
+
+
+def stub_upstream_inits(root: Path) -> None:
+    """Pre-empt heavy ``__init__.py`` files in the upstream that pull hydra/wandb.
+
+    Only the leaf model file and the mel-spec transform are needed for inference.
+    """
+    for cached in list(sys.modules):
+        if cached == "src" or cached.startswith("src."):
+            del sys.modules[cached]
+
+    for name in (
+        "src",
+        "src.utils",
+        "src.models",
+        "src.models.components",
+        "src.masks",
+        "src.masks.components",
+        "src.data",
+        "src.data.components",
+    ):
+        m = types.ModuleType(name)
+        m.__path__ = [str(root / name.replace(".", "/"))]
+        sys.modules[name] = m
+
+
+def compute_mel_spec(waveform: torch.Tensor) -> torch.Tensor:
+    """Kaldi-fbank mel spectrogram of shape (1, TARGET_TIME_BINS, N_MELS)."""
+    hop_length_ms = (CLIP_LENGTH_S * 1000) / TARGET_TIME_BINS
+    frame_length_ms = 2.5 * hop_length_ms
+    spec = torchaudio.compliance.kaldi.fbank(
+        waveform - waveform.mean(),
+        sample_frequency=SAMPLE_RATE,
+        frame_length=frame_length_ms,
+        frame_shift=hop_length_ms,
+        num_mel_bins=N_MELS,
+        low_freq=20,
+        high_freq=SAMPLE_RATE // 2,
+        use_log_fbank=True,
+        window_type="hanning",
+    )
+    if spec.shape[0] < TARGET_TIME_BINS:
+        pad = TARGET_TIME_BINS - spec.shape[0]
+        spec = torch.cat([spec, torch.zeros(pad, N_MELS)], dim=0)
+    elif spec.shape[0] > TARGET_TIME_BINS:
+        spec = spec[:TARGET_TIME_BINS]
+    return spec.unsqueeze(0)
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description="Audio-JEPA CPU inference example.")
+    parser.add_argument(
+        "--audio-jepa-src",
+        default="./audio-jepa",
+        help="Path to the cloned LudovicTuncay/Audio-JEPA repo (default: ./audio-jepa).",
+    )
+    parser.add_argument(
+        "--wav",
+        default=None,
+        help="Optional WAV file to encode (mono, will be resampled to 32 kHz).",
+    )
+    args = parser.parse_args()
+
+    root = Path(args.audio_jepa_src).resolve()
+    if not (root / "src" / "models" / "components" / "vision_transformer.py").exists():
+        print(f"ERROR: Audio-JEPA source not found at {root}", file=sys.stderr)
+        print("Run: git clone --depth 1 https://github.com/LudovicTuncay/Audio-JEPA.git", file=sys.stderr)
+        return 2
+
+    install_flash_attn_shim()
+    if str(root) not in sys.path:
+        sys.path.insert(0, str(root))
+    stub_upstream_inits(root)
+
+    from src.models.components.vision_transformer import VisionTransformer
+
+    print("Building encoder (input_size=(256, 128), patch=(16, 16))...")
+    encoder = VisionTransformer(
+        input_size=(TARGET_TIME_BINS, N_MELS),
+        patch_size=PATCH_SIZE,
+        in_chans=1,
+        embed_dim=EMBED_DIM,
+        depth=DEPTH,
+        num_heads=NUM_HEADS,
+        mlp_ratio=MLP_RATIO,
+        use_flash_attn=False,
+    )
+
+    print("Downloading checkpoint (JEPA.ckpt, ~350 MB, first run only)...")
+    ckpt_path = hf_hub_download("ltuncay/Audio-JEPA", "JEPA.ckpt")
+    state = torch.load(ckpt_path, map_location="cpu", weights_only=False)
+    raw_sd = state.get("state_dict", state)
+    encoder_sd = {
+        k[len("encoder.") :]: v
+        for k, v in raw_sd.items()
+        if k.startswith("encoder.") and not k.startswith("encoder_")
+    }
+    missing, unexpected = encoder.load_state_dict(encoder_sd, strict=True)
+    print(f"Loaded: {len(encoder_sd)} keys, {len(missing)} missing, {len(unexpected)} unexpected.")
+    encoder.eval()
+
+    if args.wav:
+        waveform, sr = torchaudio.load(args.wav)
+        if waveform.shape[0] > 1:
+            waveform = waveform.mean(dim=0, keepdim=True)
+        if sr != SAMPLE_RATE:
+            waveform = torchaudio.functional.resample(waveform, sr, SAMPLE_RATE)
+    else:
+        n = SAMPLE_RATE * 5
+        t = torch.arange(n).float() / SAMPLE_RATE
+        waveform = (0.3 * torch.sin(2 * torch.pi * 440 * t)).unsqueeze(0)
+        print(f"No --wav given, using a 5 s synthetic tone at {SAMPLE_RATE} Hz.")
+
+    duration_s = waveform.shape[1] / SAMPLE_RATE
+    print(f"Waveform: {waveform.shape[1]} samples ({duration_s:.2f} s)")
+
+    spec = compute_mel_spec(waveform).unsqueeze(0)
+    t0 = time.perf_counter()
+    with torch.inference_mode():
+        emb = encoder(spec)
+    wall = time.perf_counter() - t0
+
+    print(f"Embeddings: shape={tuple(emb.shape)}  wall={wall*1000:.1f} ms  RTF={wall/duration_s:.3f}")
+    print("Note: the encoder always produces 128 patches (8 temporal x 16 frequency).")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())