inference_example.py

"""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())