clap_embed.py

"""CLAP audio + text embedding for the live matcher Space.

The tiling here is a deliberate, line-by-line reimplementation of
audio-brief's `clap_worker.py::_embed` (checkpoint `default-630k`,
HTSAT-tiny, enable_fusion=False, librosa.load(sr=48000, mono=True), 10 s
consecutive tiles, mean, L2). That worker produced every vector in the
exported corpus. If this file drifts from it, every similarity in the app
becomes quietly wrong while still looking plausible — so `cli verify-live`
re-embeds a known corpus track through this code and compares against its
stored database vector before anything ships.

Two things beyond the worker:

* **A tile cache.** A 30 s window starting on a 5 s grid decomposes into
  exactly three 10 s tiles that also lie on that grid, so the whole snippet
  sweep reuses the tiles the full-track vector already paid for. A 3-minute
  track costs ~37 tile embeddings for the full track *and* ~31 windows,
  instead of ~130. Same numbers, a third of the CPU — which is what makes
  this affordable on a free CPU Space.
* **The text tower.** `get_text_embedding` rides in the same checkpoint, so
  descriptor tags are scored in the same 512-dim space as the audio, with no
  extra model.
"""

from __future__ import annotations

import numpy as np

CLAP_SR = 48_000              # CLAP is trained at 48 kHz
CLAP_WINDOW = CLAP_SR * 10    # laion-clap's fixed input length (non-fusion)
GRID_S = 5                    # snippet hop; also the tile-cache grid
DIM = 512
CKPT_ID = "default-630k"


def l2(v: np.ndarray) -> np.ndarray:
    n = float(np.linalg.norm(v))
    return v / n if n > 1e-12 else v


def load_audio(path: str) -> np.ndarray:
    import librosa

    audio, _ = librosa.load(path, sr=CLAP_SR, mono=True)
    if audio.size == 0:
        raise ValueError("empty audio")
    return audio


class Embedder:
    """One loaded CLAP model. Holds the tile cache for the current track."""

    ckpt = CKPT_ID

    def __init__(self):
        import laion_clap

        # HTSAT-tiny is the arch of the auto-downloaded 630k-audioset
        # checkpoint. $LAION_CLAP_MUSIC_CKPT is NOT honoured here: that
        # variable switches to the HTSAT-base music checkpoint, a different
        # vector space, and would silently invalidate every score.
        self._model = laion_clap.CLAP_Module(enable_fusion=False,
                                             amodel="HTSAT-tiny")
        self._model.load_ckpt(ckpt=_local_ckpt())
        try:
            self._model.model.eval()   # no dropout — reproducibility
        except Exception:
            pass
        self._cache: dict[tuple[int, int], np.ndarray] = {}

    # -- raw calls ---------------------------------------------------------
    def _raw(self, audio: np.ndarray) -> np.ndarray:
        x = np.ascontiguousarray(audio[None, :], dtype=np.float32)
        emb = self._model.get_audio_embedding_from_data(x=x, use_tensor=False)[0]
        return np.asarray(emb, dtype=np.float32)

    def _tile(self, audio: np.ndarray, s: int, e: int) -> np.ndarray:
        key = (s, e - s)
        hit = self._cache.get(key)
        if hit is not None:
            return hit
        w = audio[s:e]
        if len(w) < CLAP_WINDOW:          # pad the last kept tile, as the worker does
            w = np.pad(w, (0, CLAP_WINDOW - len(w)))
        out = self._tile_uncached(w)
        self._cache[key] = out
        return out

    def _tile_uncached(self, w: np.ndarray) -> np.ndarray:
        return self._raw(w)

    # -- clap_worker._embed, over an arbitrary span ------------------------
    def embed_span(self, audio: np.ndarray, lo: int, hi: int) -> np.ndarray:
        """L2-normed vector for audio[lo:hi], identical to what clap_worker
        would return for that span written out as its own file."""
        W = CLAP_WINDOW
        n = hi - lo
        if n <= 0:
            raise ValueError("empty span")
        if n <= W:
            # Short input is never randomly truncated → embed as-is,
            # unpadded, exactly as the worker does.
            return l2(self._raw(audio[lo:hi]))
        tiles = []
        for s in range(lo, hi, W):
            e = min(s + W, hi)
            if e - s < W // 2:            # drop a remainder under half a window
                break
            tiles.append(self._tile(audio, s, e))
        return l2(np.mean(tiles, axis=0).astype(np.float32))

    def embed_track(self, audio: np.ndarray) -> np.ndarray:
        return self.embed_span(audio, 0, len(audio))

    def embed_windows(self, audio: np.ndarray, window_s: float = 30.0,
                      hop_s: float = GRID_S):
        """-> (starts_seconds, matrix). Only windows that fit entirely inside
        the track are emitted: a short tail window would be embedded under
        different padding rules than the corpus, and a snippet you cannot
        actually deliver is not a recommendation. A track shorter than the
        window yields none at all — claiming a 0:00-0:30 clip of a
        20-second file would be a straightforwardly false recommendation."""
        n = len(audio)
        win = int(round(window_s * CLAP_SR))
        hop = int(round(hop_s * CLAP_SR))
        if n < win:
            return [], np.zeros((0, DIM), dtype=np.float32)
        starts = list(range(0, n - win + 1, hop))
        vecs = [self.embed_span(audio, s, s + win) for s in starts]
        return [s / CLAP_SR for s in starts], np.asarray(vecs, dtype=np.float32)

    def tile_matrix(self):
        """-> (starts_seconds, matrix) for every FULL-length tile in the cache.

        The snippet sweep has already embedded the whole track as overlapping
        10-second tiles on a 5-second grid, so a self-similarity read of the
        record's sections costs one dot product and no new audio work. Short
        tail tiles are excluded: they are zero-padded, which shifts their
        vector, and comparing them against full ones would read as a section
        change that is really an artefact of padding.

        Must be called BEFORE `reset_cache()`.
        """
        keys = sorted(k for k in self._cache if k[1] == CLAP_WINDOW)
        if not keys:
            return np.zeros(0), np.zeros((0, DIM), dtype=np.float32)
        starts = np.asarray([k[0] / CLAP_SR for k in keys], dtype=np.float64)
        mat = np.stack([l2(self._cache[k]) for k in keys]).astype(np.float32)
        return starts, mat

    def reset_cache(self) -> None:
        self._cache.clear()

    # -- text tower --------------------------------------------------------
    def embed_texts(self, texts: list[str]) -> np.ndarray:
        vecs = self._model.get_text_embedding(texts, use_tensor=False)
        arr = np.asarray(vecs, dtype=np.float32)
        return np.stack([l2(v) for v in arr])


def _local_ckpt():
    """Fetch 630k-audioset-best.pt from the Hub once and cache it.

    laion_clap's own `load_ckpt()` wgets into its site-packages directory,
    which a Space container discards on every rebuild. hf_hub_download puts
    it in HF_HOME instead, where it survives restarts. Returns None (i.e.
    let laion_clap do its default download) if the Hub is unreachable, so a
    local run without network still works off an already-downloaded copy.
    """
    try:
        from huggingface_hub import hf_hub_download

        return hf_hub_download("lukewys/laion_clap", "630k-audioset-best.pt")
    except Exception:
        return None