Add Gradio app, model code, and deps — checkpoint downloads from dayngerous/whoSampledAST

README.md

@@ -1,13 +1,30 @@
 ---
-title: Sampled
-emoji: 🚀
-colorFrom: pink
-colorTo: pink
+title: Sample Match Verifier
+emoji: 🎵
+colorFrom: blue
+colorTo: purple
 sdk: gradio
-sdk_version: 6.13.0
+sdk_version: "5.0"
 app_file: app.py
 pinned: false
-short_description: Detect if a sample is in another song
+license: mit
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Sample Match Verifier
+
+Upload a track and a possible source sample. Waveforms appear immediately on upload. Click **Verify match** to run the model — it scans beat-aligned windows, scores the best match, and highlights the predicted sampled sections on both the waveform and mel spectrogram. If no confident match is found, the mel spectrogram shows a **No Match** overlay.
+
+## Model checkpoint
+
+Place your checkpoint at `models/best.pt` (committed via Git LFS) or set the `MODEL_CHECKPOINT` environment variable to its path. The app falls back to `checkpoints/best.pt` if `models/best.pt` is not found.
+
+## Environment variables
+
+| Variable | Default | Description |
+|---|---|---|
+| `MODEL_CHECKPOINT` | `models/best.pt` | Path to the `.pt` checkpoint |
+| `MODEL_BACKBONE` | `ast` | Backbone: `ast`, `sslam`, or `cnn` |
+| `AST_MODEL` | `MIT/ast-finetuned-audioset-10-10-0.4593` | HuggingFace AST model ID |
+| `MODEL_BARS` | `4` | Bars per analysis window |
+| `MODEL_N_MELS` | `128` | Mel frequency bins |
+| `APP_SAMPLE_RATE` | `16000` | Audio sample rate |

app.py

@@ -0,0 +1,597 @@
+import json
+import os
+from dataclasses import dataclass
+from functools import lru_cache
+from pathlib import Path
+
+os.environ.setdefault("AST_MODEL", "MIT/ast-finetuned-audioset-10-10-0.4593")
+os.environ.setdefault("SSLAM_MODEL", "ta012/SSLAM_pretrain")
+
+import gradio as gr
+import librosa
+import matplotlib
+import numpy as np
+import torch
+import torchaudio.transforms as T
+from huggingface_hub import hf_hub_download
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+
+from model import CNNSampleDetector, SSLAMSampleDetector, SampleDetector
+
+
+SAMPLE_RATE = int(os.environ.get("APP_SAMPLE_RATE", "16000"))
+MODEL_REPO = os.environ.get("MODEL_REPO", "dayngerous/whoSampledAST")
+
+
+def _resolve_checkpoint() -> str:
+    """Return local checkpoint path, downloading from HF Hub if needed."""
+    env_path = os.environ.get("MODEL_CHECKPOINT", "")
+    for p in [env_path, "models/best.pt", "checkpoints/best.pt", "checkpoints2/best.pt"]:
+        if p and Path(p).exists():
+            return p
+    try:
+        return hf_hub_download(repo_id=MODEL_REPO, filename="models/best.pt")
+    except Exception as exc:
+        raise FileNotFoundError(
+            f"No local checkpoint found and download from {MODEL_REPO} failed: {exc}"
+        )
+
+
+def _resolve_meta() -> str:
+    """Return local test_indices.json path, downloading from HF Hub if needed."""
+    for p in ["models/test_indices.json", "checkpoints2/test_indices.json", "checkpoints/test_indices.json"]:
+        if Path(p).exists():
+            return p
+    try:
+        return hf_hub_download(repo_id=MODEL_REPO, filename="models/test_indices.json")
+    except Exception:
+        return ""
+
+
+DEFAULT_CHECKPOINT = _resolve_checkpoint()
+DEFAULT_META = DEFAULT_META or _resolve_meta()
+TARGET_FRAMES_PER_BEAT = 50
+N_FFT = 1024
+MEL_HOP = 512
+N_MELS_VIZ = 128
+
+
+@dataclass
+class AudioClip:
+    waveform: torch.Tensor
+    sample_rate: int
+    offset_sec: float
+    duration_sec: float
+
+
+@dataclass
+class BeatWindow:
+    waveform: torch.Tensor
+    start_sec: float
+    end_sec: float
+    beat_intervals: list[tuple[float, float]]
+
+
+def _format_time(seconds: float) -> str:
+    seconds = max(0.0, float(seconds))
+    minutes = int(seconds // 60)
+    rem = seconds - minutes * 60
+    return f"{minutes}:{rem:04.1f}"
+
+
+def _format_intervals(intervals: list[tuple[float, float]], limit: int = 4) -> str:
+    if not intervals:
+        return "none"
+    shown = ", ".join(f"{_format_time(a)}-{_format_time(b)}" for a, b in intervals[:limit])
+    if len(intervals) > limit:
+        shown += f", +{len(intervals) - limit} more"
+    return shown
+
+
+def _merge_intervals(intervals: list[tuple[float, float]], gap: float = 0.05) -> list[tuple[float, float]]:
+    if not intervals:
+        return []
+    ordered = sorted((float(a), float(b)) for a, b in intervals if b > a)
+    merged = [ordered[0]]
+    for start, end in ordered[1:]:
+        prev_start, prev_end = merged[-1]
+        if start <= prev_end + gap:
+            merged[-1] = (prev_start, max(prev_end, end))
+        else:
+            merged.append((start, end))
+    return merged
+
+
+def _load_args(checkpoint_path: Path) -> dict:
+    meta_path = Path(DEFAULT_META) if DEFAULT_META else checkpoint_path.parent / "test_indices.json"
+    args = {}
+    if meta_path.exists():
+        with open(meta_path) as f:
+            args = json.load(f).get("args", {})
+
+    args.setdefault("backbone", os.environ.get("MODEL_BACKBONE", "ast"))
+    args.setdefault("ast_model", os.environ.get("AST_MODEL"))
+    args.setdefault("bars", int(os.environ.get("MODEL_BARS", "4")))
+    args.setdefault("n_mels", int(os.environ.get("MODEL_N_MELS", "128")))
+    args.setdefault("sample_rate", SAMPLE_RATE)
+    return args
+
+
+def _build_model(args: dict, device: torch.device):
+    beats_per_window = int(args.get("bars", 4)) * 4
+    n_mels = int(args.get("n_mels", 128))
+    backbone = args.get("backbone", "ast")
+    if backbone == "ast":
+        model = SampleDetector(
+            model_name=args.get("ast_model", os.environ["AST_MODEL"]),
+            freeze_encoder=True,
+            beats_per_window=beats_per_window,
+            n_mels=n_mels,
+        )
+    elif backbone == "sslam":
+        model = SSLAMSampleDetector(
+            freeze_encoder=True,
+            beats_per_window=beats_per_window,
+            n_mels=n_mels,
+        )
+    else:
+        model = CNNSampleDetector(beats_per_window=beats_per_window, n_mels=n_mels)
+    return model.to(device)
+
+
+@lru_cache(maxsize=2)
+def _load_model(checkpoint_path: str):
+    path = Path(checkpoint_path)
+    if not path.exists():
+        raise FileNotFoundError(
+            f"Checkpoint not found: {path}. Set MODEL_CHECKPOINT or place a checkpoint at models/best.pt."
+        )
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    args = _load_args(path)
+    model = _build_model(args, device)
+    ckpt = torch.load(path, map_location=device)
+    state = ckpt.get("model_state", ckpt)
+    pair_head_loaded = any(k.startswith("pair_mask_head.") for k in state)
+    missing, unexpected = model.load_state_dict(state, strict=False)
+    model.eval()
+    return {
+        "model": model,
+        "args": args,
+        "device": device,
+        "epoch": ckpt.get("epoch", "?"),
+        "pair_head_loaded": pair_head_loaded,
+        "missing": missing,
+        "unexpected": unexpected,
+    }
+
+
+def _load_audio(path: str, offset_sec: float, max_seconds: float) -> AudioClip:
+    if not path:
+        raise gr.Error("Upload both audio files before running verification.")
+
+    audio, sr = librosa.load(path, sr=SAMPLE_RATE, mono=True)
+    waveform = torch.from_numpy(audio).float()
+
+    offset_sec = max(0.0, float(offset_sec or 0.0))
+    max_seconds = max(1.0, float(max_seconds or 1.0))
+    start = min(int(offset_sec * sr), max(waveform.numel() - 1, 0))
+    end = min(start + int(max_seconds * sr), waveform.numel())
+    waveform = waveform[start:end].float().contiguous()
+    if waveform.numel() < sr // 4:
+        raise gr.Error("Each upload must contain at least 0.25 seconds of audio after offset trimming.")
+
+    peak = waveform.abs().max().clamp_min(1e-6)
+    waveform = waveform / peak
+    return AudioClip(
+        waveform=waveform,
+        sample_rate=sr,
+        offset_sec=offset_sec,
+        duration_sec=waveform.numel() / sr,
+    )
+
+
+def _estimate_beats(waveform: torch.Tensor, sample_rate: int) -> tuple[float, np.ndarray]:
+    y = waveform.detach().cpu().numpy().astype(np.float32)
+    tempo, beat_frames = librosa.beat.beat_track(y=y, sr=sample_rate, hop_length=512)
+    bpm = float(np.atleast_1d(tempo)[0]) if np.size(tempo) else 120.0
+    if not np.isfinite(bpm) or bpm <= 0:
+        bpm = 120.0
+    bpm = float(np.clip(bpm, 60.0, 200.0))
+
+    beat_samples = librosa.frames_to_samples(beat_frames, hop_length=512)
+    beat_samples = beat_samples[(beat_samples >= 0) & (beat_samples < waveform.numel())]
+    if len(beat_samples) < 2:
+        step = max(1, int(round(sample_rate * 60.0 / bpm)))
+        beat_samples = np.arange(0, waveform.numel(), step, dtype=np.int64)
+    elif beat_samples[0] > sample_rate * 60.0 / bpm:
+        beat_samples = np.insert(beat_samples, 0, 0)
+    return bpm, beat_samples.astype(np.int64)
+
+
+def _to_mel(waveform: torch.Tensor, bpm: float, args: dict) -> torch.Tensor:
+    sample_rate = int(args.get("sample_rate", SAMPLE_RATE))
+    n_mels = int(args.get("n_mels", 128))
+    bars = int(args.get("bars", 4))
+    fixed_frames = bars * 4 * TARGET_FRAMES_PER_BEAT
+    hop = max(1, round(60 * sample_rate / (bpm * TARGET_FRAMES_PER_BEAT)))
+
+    mel_transform = T.MelSpectrogram(
+        sample_rate=sample_rate,
+        n_fft=N_FFT,
+        hop_length=hop,
+        n_mels=n_mels,
+        power=2.0,
+    )
+    amp_to_db = T.AmplitudeToDB(stype="power", top_db=80)
+    mel = amp_to_db(mel_transform(waveform)).T
+    if mel.shape[0] > fixed_frames:
+        mel = mel[:fixed_frames]
+    elif mel.shape[0] < fixed_frames:
+        mel = torch.cat([mel, torch.zeros(fixed_frames - mel.shape[0], mel.shape[1])], dim=0)
+    mel = (mel - mel.mean()) / (mel.std() + 1e-6)
+    return mel.unsqueeze(0)
+
+
+def _make_windows(
+    clip: AudioClip,
+    bpm: float,
+    beat_samples: np.ndarray,
+    args: dict,
+    stride_beats: int,
+    max_windows: int,
+) -> list[BeatWindow]:
+    bars = int(args.get("bars", 4))
+    beats_per_window = bars * 4
+    window_samples = max(1, int(round(beats_per_window * 60.0 / bpm * clip.sample_rate)))
+    beat_seconds = 60.0 / bpm
+    stride_beats = max(1, int(stride_beats))
+    max_windows = max(1, int(max_windows))
+
+    valid = [i for i in range(0, len(beat_samples), stride_beats) if beat_samples[i] < clip.waveform.numel()]
+    if not valid:
+        valid = [0]
+
+    if len(valid) > max_windows:
+        chosen_positions = np.linspace(0, len(valid) - 1, max_windows, dtype=np.int64)
+        valid = [valid[i] for i in sorted(set(chosen_positions.tolist()))]
+
+    windows = []
+    for beat_idx in valid:
+        start_sample = int(beat_samples[beat_idx]) if len(beat_samples) else 0
+        chunk = clip.waveform[start_sample:start_sample + window_samples]
+        if chunk.numel() < window_samples:
+            chunk = torch.nn.functional.pad(chunk, (0, window_samples - chunk.numel()))
+
+        start_sec = clip.offset_sec + start_sample / clip.sample_rate
+        end_sec = start_sec + window_samples / clip.sample_rate
+        beat_intervals = [
+            (start_sec + i * beat_seconds, start_sec + (i + 1) * beat_seconds)
+            for i in range(beats_per_window)
+        ]
+        windows.append(BeatWindow(chunk, start_sec, end_sec, beat_intervals))
+    return windows
+
+
+def _encode(model, mels: torch.Tensor, batch_size: int) -> torch.Tensor:
+    embs = []
+    for start in range(0, mels.shape[0], batch_size):
+        embs.append(model.encoder(mels[start:start + batch_size]))
+    return torch.cat(embs, dim=0)
+
+
+def _score_pairs(model, track_mels: torch.Tensor, source_mels: torch.Tensor, batch_size: int) -> torch.Tensor:
+    track_emb = _encode(model, track_mels, batch_size)
+    source_emb = _encode(model, source_mels, batch_size)
+    n_track, n_source = track_emb.shape[0], source_emb.shape[0]
+    scores = []
+
+    pair_indices = [(i, j) for i in range(n_track) for j in range(n_source)]
+    for start in range(0, len(pair_indices), batch_size):
+        chunk = pair_indices[start:start + batch_size]
+        ti = torch.tensor([p[0] for p in chunk], device=track_emb.device)
+        sj = torch.tensor([p[1] for p in chunk], device=track_emb.device)
+        t = track_emb.index_select(0, ti)
+        s = source_emb.index_select(0, sj)
+        combined = torch.cat([t, s, torch.abs(t - s), t * s], dim=-1)
+        logits = model.head(combined)
+        scores.append(torch.softmax(logits, dim=-1)[:, 1])
+
+    return torch.cat(scores).reshape(n_track, n_source)
+
+
+def _intervals_from_mask(mask: np.ndarray, window: BeatWindow, max_end: float) -> list[tuple[float, float]]:
+    intervals = []
+    for use, (start, end) in zip(mask.tolist(), window.beat_intervals):
+        if use:
+            intervals.append((start, min(end, max_end)))
+    return _merge_intervals(intervals)
+
+
+def _localize_match(
+    model,
+    track_mel: torch.Tensor,
+    source_mel: torch.Tensor,
+    track_window: BeatWindow,
+    source_window: BeatWindow,
+    track_clip: AudioClip,
+    source_clip: AudioClip,
+    threshold: float,
+    pair_head_loaded: bool,
+) -> tuple[list[tuple[float, float]], list[tuple[float, float]], str]:
+    if not pair_head_loaded:
+        return (
+            [(track_window.start_sec, min(track_window.end_sec, track_clip.offset_sec + track_clip.duration_sec))],
+            [(source_window.start_sec, min(source_window.end_sec, source_clip.offset_sec + source_clip.duration_sec))],
+            "The checkpoint does not include a trained pairwise beat head, so the highlight covers the best matching window.",
+        )
+
+    with torch.inference_mode():
+        pair_probs = torch.sigmoid(model.pair_mask_head(track_mel, source_mel))[0].detach().cpu().numpy()
+
+    selected = pair_probs >= float(threshold)
+    if not selected.any():
+        top_k = min(6, pair_probs.size)
+        flat = np.argpartition(pair_probs.reshape(-1), -top_k)[-top_k:]
+        selected = np.zeros_like(pair_probs, dtype=bool)
+        selected.reshape(-1)[flat] = True
+
+    track_mask = selected.any(axis=1)
+    source_mask = selected.any(axis=0)
+    track_regions = _intervals_from_mask(
+        track_mask,
+        track_window,
+        track_clip.offset_sec + track_clip.duration_sec,
+    )
+    source_regions = _intervals_from_mask(
+        source_mask,
+        source_window,
+        source_clip.offset_sec + source_clip.duration_sec,
+    )
+    return track_regions, source_regions, ""
+
+
+def _draw_waveform(ax, clip: AudioClip, regions: list[tuple[float, float]], color: str, title: str):
+    y = clip.waveform.detach().cpu().numpy()
+    n = len(y)
+    points = min(20000, n)
+    idx = np.linspace(0, n - 1, points, dtype=np.int64)
+    x = clip.offset_sec + idx / clip.sample_rate
+    ax.plot(x, y[idx], color="#111827", linewidth=0.55)
+    for start, end in regions:
+        ax.axvspan(start, end, color=color, alpha=0.28)
+    ax.set_title(title, loc="left", fontsize=10)
+    ax.set_ylabel("Amplitude")
+    ax.set_xlim(clip.offset_sec, clip.offset_sec + clip.duration_sec)
+    ax.set_ylim(-1.05, 1.05)
+    ax.grid(True, alpha=0.18)
+
+
+def _draw_mel(ax, clip: AudioClip, regions: list[tuple[float, float]], color: str, title: str, matched: bool):
+    y = clip.waveform.detach().cpu().numpy().astype(np.float32)
+    mel = librosa.feature.melspectrogram(y=y, sr=clip.sample_rate, n_mels=N_MELS_VIZ, hop_length=MEL_HOP)
+    mel_db = librosa.power_to_db(mel, ref=np.max)
+
+    t_start = clip.offset_sec
+    t_end = clip.offset_sec + clip.duration_sec
+    f_max = clip.sample_rate / 2
+
+    ax.imshow(
+        mel_db,
+        aspect="auto",
+        origin="lower",
+        extent=[t_start, t_end, 0, f_max],
+        cmap="magma",
+        interpolation="nearest",
+    )
+    ax.set_title(title, loc="left", fontsize=10)
+    ax.set_ylabel("Frequency (Hz)")
+    ax.set_xlim(t_start, t_end)
+
+    if matched and regions:
+        for start, end in regions:
+            ax.axvspan(start, end, color=color, alpha=0.38, linewidth=0)
+    elif not matched:
+        ax.text(
+            0.5, 0.5, "No Match",
+            transform=ax.transAxes,
+            fontsize=18,
+            color="white",
+            ha="center",
+            va="center",
+            fontweight="bold",
+            bbox=dict(boxstyle="round,pad=0.4", facecolor="#111827", alpha=0.65),
+        )
+
+
+def _plot_waveforms(
+    track_clip: AudioClip,
+    source_clip: AudioClip,
+    track_regions: list[tuple[float, float]],
+    source_regions: list[tuple[float, float]],
+    score: float | None,
+    matched: bool,
+) -> plt.Figure:
+    fig, axes = plt.subplots(2, 1, figsize=(12, 5), sharex=False)
+    color = "#22c55e" if matched else "#f59e0b"
+    title_score = "unavailable" if score is None else f"{score:.3f}"
+    fig.suptitle(f"Best match score: {title_score}" if score is not None else "Waveform preview", fontsize=12)
+
+    _draw_waveform(axes[0], track_clip, track_regions, color, "Track / song audio")
+    _draw_waveform(axes[1], source_clip, source_regions, color, "Source sample audio")
+    axes[1].set_xlabel("Time in uploaded file (seconds)")
+    fig.tight_layout()
+    return fig
+
+
+def _plot_mels(
+    track_clip: AudioClip,
+    source_clip: AudioClip,
+    track_regions: list[tuple[float, float]],
+    source_regions: list[tuple[float, float]],
+    matched: bool,
+) -> plt.Figure:
+    fig, axes = plt.subplots(2, 1, figsize=(12, 6), sharex=False)
+    color = "#22c55e" if matched else "#f59e0b"
+
+    _draw_mel(axes[0], track_clip, track_regions, color, "Track mel spectrogram", matched)
+    _draw_mel(axes[1], source_clip, source_regions, color, "Source mel spectrogram", matched)
+    axes[1].set_xlabel("Time in uploaded file (seconds)")
+    fig.tight_layout()
+    return fig
+
+
+def preview_waveforms(track_audio, source_audio):
+    if not track_audio or not source_audio:
+        return None, None
+    try:
+        track_clip = _load_audio(track_audio, 0.0, 120.0)
+        source_clip = _load_audio(source_audio, 0.0, 120.0)
+        wfig = _plot_waveforms(track_clip, source_clip, [], [], None, False)
+        mfig = _plot_mels(track_clip, source_clip, [], [], False)
+        return wfig, mfig
+    except Exception:
+        return None, None
+
+
+def verify(
+    track_audio,
+    source_audio,
+    checkpoint_path,
+    match_threshold,
+    localization_threshold,
+    track_offset,
+    source_offset,
+    max_seconds,
+    stride_beats,
+    max_windows,
+):
+    try:
+        track_clip = _load_audio(track_audio, track_offset, max_seconds)
+        source_clip = _load_audio(source_audio, source_offset, max_seconds)
+    except Exception as exc:
+        raise gr.Error(str(exc))
+
+    try:
+        loaded = _load_model(checkpoint_path or DEFAULT_CHECKPOINT)
+    except Exception as exc:
+        wfig = _plot_waveforms(track_clip, source_clip, [], [], None, False)
+        mfig = _plot_mels(track_clip, source_clip, [], [], False)
+        return f"Model could not be loaded: {exc}", wfig, mfig
+
+    model = loaded["model"]
+    args = loaded["args"]
+    device = loaded["device"]
+    batch_size = 8 if device.type == "cpu" else 32
+
+    track_bpm, track_beats = _estimate_beats(track_clip.waveform, track_clip.sample_rate)
+    source_bpm, source_beats = _estimate_beats(source_clip.waveform, source_clip.sample_rate)
+    track_windows = _make_windows(track_clip, track_bpm, track_beats, args, stride_beats, max_windows)
+    source_windows = _make_windows(source_clip, source_bpm, source_beats, args, stride_beats, max_windows)
+
+    track_mels = torch.stack([_to_mel(w.waveform, track_bpm, args) for w in track_windows]).to(device)
+    source_mels = torch.stack([_to_mel(w.waveform, source_bpm, args) for w in source_windows]).to(device)
+
+    with torch.inference_mode():
+        score_matrix = _score_pairs(model, track_mels, source_mels, batch_size)
+    best_flat = int(torch.argmax(score_matrix).item())
+    best_track = best_flat // score_matrix.shape[1]
+    best_source = best_flat % score_matrix.shape[1]
+    best_score = float(score_matrix[best_track, best_source].detach().cpu())
+    matched = best_score >= float(match_threshold)
+
+    track_regions, source_regions, note = _localize_match(
+        model,
+        track_mels[best_track:best_track + 1],
+        source_mels[best_source:best_source + 1],
+        track_windows[best_track],
+        source_windows[best_source],
+        track_clip,
+        source_clip,
+        localization_threshold,
+        loaded["pair_head_loaded"],
+    )
+
+    highlight_track = track_regions if matched else []
+    highlight_source = source_regions if matched else []
+
+    wfig = _plot_waveforms(track_clip, source_clip, highlight_track, highlight_source, best_score, matched)
+    mfig = _plot_mels(track_clip, source_clip, highlight_track, highlight_source, matched)
+
+    verdict = "Likely match" if matched else "No confident match"
+    details = [
+        f"**{verdict}**",
+        f"Score: `{best_score:.3f}` with threshold `{float(match_threshold):.2f}`.",
+        f"Estimated BPM: track `{track_bpm:.1f}`, source `{source_bpm:.1f}`.",
+        f"Highlighted track section(s): {_format_intervals(highlight_track)}.",
+        f"Highlighted source section(s): {_format_intervals(highlight_source)}.",
+        f"Model: `{args.get('backbone', 'ast')}` checkpoint epoch `{loaded['epoch']}` on `{device}`.",
+    ]
+    if note:
+        details.append(note)
+    if loaded["missing"]:
+        details.append(f"Missing checkpoint keys initialized at load time: `{len(loaded['missing'])}`.")
+    return "\n\n".join(details), wfig, mfig
+
+
+with gr.Blocks(title="Sample Match Verifier") as demo:
+    gr.Markdown("# Sample Match Verifier")
+    gr.Markdown(
+        "Upload a track and a possible source sample. "
+        "Waveforms appear immediately on upload. "
+        "Click **Verify match** to run the model and highlight sampled sections."
+    )
+
+    with gr.Row():
+        track_audio = gr.Audio(label="Track / song audio", type="filepath", sources=["upload"])
+        source_audio = gr.Audio(label="Source sample audio", type="filepath", sources=["upload"])
+
+    with gr.Accordion("Settings", open=False):
+        checkpoint_path = gr.Textbox(label="Checkpoint path", value=DEFAULT_CHECKPOINT)
+        with gr.Row():
+            match_threshold = gr.Slider(0.0, 1.0, value=0.50, step=0.01, label="Match threshold")
+            localization_threshold = gr.Slider(0.0, 1.0, value=0.55, step=0.01, label="Highlight threshold")
+        with gr.Row():
+            track_offset = gr.Number(value=0.0, label="Track start offset, seconds")
+            source_offset = gr.Number(value=0.0, label="Source start offset, seconds")
+        with gr.Row():
+            max_seconds = gr.Slider(5, 180, value=60, step=5, label="Analyze duration per upload, seconds")
+            stride_beats = gr.Slider(1, 16, value=4, step=1, label="Window stride, beats")
+            max_windows = gr.Slider(4, 64, value=24, step=1, label="Max windows per upload")
+
+    run = gr.Button("Verify match", variant="primary")
+    result = gr.Markdown()
+
+    waveform_plot = gr.Plot(label="Waveforms")
+    mel_plot = gr.Plot(label="Mel Spectrograms")
+
+    # Show waveforms as soon as both files are uploaded
+    for audio_input in [track_audio, source_audio]:
+        audio_input.change(
+            preview_waveforms,
+            inputs=[track_audio, source_audio],
+            outputs=[waveform_plot, mel_plot],
+        )
+
+    run.click(
+        verify,
+        inputs=[
+            track_audio,
+            source_audio,
+            checkpoint_path,
+            match_threshold,
+            localization_threshold,
+            track_offset,
+            source_offset,
+            max_seconds,
+            stride_beats,
+            max_windows,
+        ],
+        outputs=[result, waveform_plot, mel_plot],
+    )
+
+
+if __name__ == "__main__":
+    demo.queue(max_size=8).launch()

model.py

@@ -0,0 +1,316 @@
+import glob
+import importlib
+import os
+
+import torch
+import torch.nn as nn
+from dotenv import load_dotenv
+from transformers import ASTModel, ASTConfig
+
+load_dotenv()
+
+AST_TIME_DIM = 1024
+AST_FREQ_DIM = 128
+SSLAM_HF_REPO = os.environ["SSLAM_MODEL"]
+SSLAM_TIME_DIM = 1024
+SSLAM_FREQ_DIM = 128
+
+
+class ASTEncoder(nn.Module):
+    """Wraps ASTModel and returns the [CLS] token embedding."""
+
+    def __init__(self, model_name: str, freeze: bool = True):
+        super().__init__()
+        self.ast = ASTModel.from_pretrained(model_name, ignore_mismatched_sizes=True)
+        # print(f"AST hidden size: {self.ast.config.hidden_size}")
+        if freeze:
+            for p in self.ast.parameters():
+                p.requires_grad = False
+
+    def unfreeze_last_n(self, n: int = 2):
+        for block in self.ast.encoder.layer[-n:]:
+            for p in block.parameters():
+                p.requires_grad = True
+        for p in self.ast.layernorm.parameters():
+            p.requires_grad = True
+        # trainable = sum(p.numel() for p in self.ast.parameters() if p.requires_grad)
+        # print(f"unfroze {n} blocks, trainable params: {trainable:,}")
+
+
+
+    @staticmethod
+    def _prep(mel: torch.Tensor) -> torch.Tensor:
+        """mel: [B, 1, T, F] => [B, AST_TIME_DIM, AST_FREQ_DIM]"""
+        x = mel.squeeze(1)
+        T = x.shape[1]
+        # print(f"input T={T}, target={AST_TIME_DIM}")
+        if T < AST_TIME_DIM:
+            pad = torch.zeros(x.shape[0], AST_TIME_DIM - T, x.shape[2], device=x.device, dtype=x.dtype)
+            x = torch.cat([x, pad], dim=1)
+        elif T > AST_TIME_DIM:
+            x = x[:, :AST_TIME_DIM, :]
+        return x
+
+    def forward(self, mel: torch.Tensor) -> torch.Tensor:
+        x = self._prep(mel)
+        out = self.ast(input_values=x)
+        # print(f"AST output shape: {out.last_hidden_state.shape}")
+        return out.last_hidden_state[:, 0, :]
+
+
+class PairMaskHead(nn.Module):
+    """Beat-by-beat pair matching head over two mel spectrograms."""
+
+    def __init__(self, beats_per_window: int, n_mels: int, beat_dim: int = 64):
+        super().__init__()
+        self.pool = nn.AdaptiveAvgPool2d((beats_per_window, n_mels))
+        self.beat_proj = nn.Sequential(
+            nn.Linear(n_mels, beat_dim),
+            nn.GELU(),
+            nn.Linear(beat_dim, beat_dim),
+        )
+        self.logit_scale = nn.Parameter(torch.tensor(1.0))
+        self.bias = nn.Parameter(torch.zeros(()))
+
+    def _beats(self, mel: torch.Tensor) -> torch.Tensor:
+        # mel: [B, 1, T, F] -> [B, beats, F] -> [B, beats, beat_dim]
+        x = self.pool(mel).squeeze(1)
+        return torch.nn.functional.normalize(self.beat_proj(x), dim=-1)
+
+    def forward(self, track_mel: torch.Tensor, orig_mel: torch.Tensor) -> torch.Tensor:
+        t = self._beats(track_mel)
+        o = self._beats(orig_mel)
+        return torch.einsum("bih,bjh->bij", t, o) * self.logit_scale.exp() + self.bias
+
+
+class SampleDetector(nn.Module):
+    """Siamese AST encoder + interaction head for binary sample detection."""
+
+    def __init__(
+        self,
+        model_name: str = os.environ["AST_MODEL"],
+        freeze_encoder: bool = True,
+        dropout: float = 0.3,
+        beats_per_window: int = 16,
+        n_mels: int = 128,
+    ):
+        super().__init__()
+        self.encoder = ASTEncoder(model_name, freeze=freeze_encoder)
+        H = self.encoder.ast.config.hidden_size
+        self.head = nn.Sequential(
+            nn.LayerNorm(4 * H),
+            nn.Linear(4 * H, 512),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(512, 128),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(128, 2),
+        )
+        self.pair_mask_head = PairMaskHead(beats_per_window, n_mels)
+
+    def unfreeze_encoder(self, n_blocks: int = 2):
+        self.encoder.unfreeze_last_n(n_blocks)
+
+    def forward(self, track_mel: torch.Tensor, orig_mel: torch.Tensor) -> torch.Tensor:
+        t = self.encoder(track_mel)
+        o = self.encoder(orig_mel)
+        # print(f"embeddings: t={t.shape}, o={o.shape}")
+        combined = torch.cat([t, o, torch.abs(t - o), t * o], dim=-1)
+        # print(f"combined shape: {combined.shape}")
+        return self.head(combined)
+
+
+class ConvBlock(nn.Module):
+    def __init__(self, in_ch: int, out_ch: int, stride: int = 2):
+        super().__init__()
+        self.block = nn.Sequential(
+            nn.Conv2d(in_ch, out_ch, 3, stride=stride, padding=1, bias=False),
+            nn.BatchNorm2d(out_ch),
+            nn.GELU(),
+            nn.Conv2d(out_ch, out_ch, 3, padding=1, bias=False),
+            nn.BatchNorm2d(out_ch),
+            nn.GELU(),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.block(x)
+
+
+class CNNEncoder(nn.Module):
+    def __init__(self, embed_dim: int = 256):
+        super().__init__()
+        self.net = nn.Sequential(
+            ConvBlock(1, 32),
+            ConvBlock(32, 64),
+            ConvBlock(64, 128),
+            ConvBlock(128, 256),
+            nn.AdaptiveAvgPool2d(1),
+            nn.Flatten(),
+            nn.Linear(256, embed_dim),
+        )
+
+    def forward(self, mel: torch.Tensor) -> torch.Tensor:
+        return self.net(mel)
+
+
+class CNNSampleDetector(nn.Module):
+    """Drop-in CNN alternative to SampleDetector."""
+
+    def __init__(self, embed_dim: int = 256, dropout: float = 0.3, beats_per_window: int = 16, n_mels: int = 128):
+        super().__init__()
+        self.encoder = CNNEncoder(embed_dim)
+        self.head = nn.Sequential(
+            nn.LayerNorm(4 * embed_dim),
+            nn.Linear(4 * embed_dim, 256),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(256, 64),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(64, 2),
+        )
+        self.pair_mask_head = PairMaskHead(beats_per_window, n_mels)
+
+    def forward(self, track_mel: torch.Tensor, orig_mel: torch.Tensor) -> torch.Tensor:
+        t = self.encoder(track_mel)
+        o = self.encoder(orig_mel)
+        combined = torch.cat([t, o, torch.abs(t - o), t * o], dim=-1)
+        return self.head(combined)
+
+
+
+class SSLAMEncoder(nn.Module):
+    """Wraps the EAT (SSLAM) model and returns the CLS-like token embedding.
+
+    Bypasses AutoModel.from_pretrained due to a transformers >= 5.5 incompatibility
+    with EATModel's missing all_tied_weights_keys attribute.
+    """
+
+    def __init__(self, freeze: bool = True):
+        super().__init__()
+        from transformers import AutoConfig
+        import safetensors.torch
+        from huggingface_hub import hf_hub_download
+
+        cfg = AutoConfig.from_pretrained(SSLAM_HF_REPO, trust_remote_code=True)
+        self.hidden_size = cfg.embed_dim
+        sha = cfg._commit_hash or self._find_sha()
+        eat_mod = importlib.import_module(
+            f"transformers_modules.ta012.SSLAM_pretrain.{sha}.eat_model"
+        )
+        self.eat = eat_mod.EAT(cfg)
+
+        weights_path = hf_hub_download(SSLAM_HF_REPO, "model.safetensors")
+        raw = safetensors.torch.load_file(weights_path)
+        state = {k.removeprefix("model."): v for k, v in raw.items()}
+        self.eat.load_state_dict(state, strict=True)
+        if freeze:
+            for p in self.eat.parameters():
+                p.requires_grad = False
+
+    @staticmethod
+    def _find_sha() -> str:
+        dirs = glob.glob(
+            os.path.expanduser(
+                f"~/.cache/huggingface/modules/transformers_modules/{SSLAM_HF_REPO}/*"
+            )
+        )
+        dirs = [d for d in dirs if os.path.isdir(d)]
+        if not dirs:
+            raise RuntimeError("SSLAM modules not found in HF cache — run AutoConfig.from_pretrained first")
+        return os.path.basename(sorted(dirs)[-1])
+
+    def unfreeze_last_n(self, n: int):
+        for block in self.eat.blocks[-n:]:
+            for p in block.parameters():
+                p.requires_grad = True
+                
+        for p in self.eat.pre_norm.parameters():
+            p.requires_grad = True
+
+    @staticmethod
+    def _prep(mel: torch.Tensor) -> torch.Tensor:
+        """mel: [B, 1, T, F] => [B, 1, SSLAM_TIME_DIM, SSLAM_FREQ_DIM]"""
+        x = mel.float()
+        T = x.shape[2]
+        if T < SSLAM_TIME_DIM:
+            pad = torch.zeros(x.shape[0], 1, SSLAM_TIME_DIM - T, x.shape[3],
+                              device=x.device, dtype=x.dtype)
+            x = torch.cat([x, pad], dim=2)
+        elif T > SSLAM_TIME_DIM:
+            x = x[:, :, :SSLAM_TIME_DIM, :]
+        return x
+
+    def forward(self, mel: torch.Tensor) -> torch.Tensor:
+        x = self._prep(mel)
+        feats = self.eat.extract_features(x)
+        # print(f"SSLAM features: {feats.shape}")  # should be [B, 1+patches, 768]
+        return feats[:, 0, :]
+
+
+
+class SSLAMSampleDetector(nn.Module):
+    """SampleDetector using SSLAM/EAT encoder instead of AST."""
+
+    def __init__(self, freeze_encoder: bool = True, dropout: float = 0.3, beats_per_window: int = 16, n_mels: int = 128):
+        super().__init__()
+        self.encoder = SSLAMEncoder(freeze=freeze_encoder)
+        H = self.encoder.hidden_size
+        self.head = nn.Sequential(
+            nn.LayerNorm(4 * H),
+            nn.Linear(4 * H, 512),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(512, 128),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(128, 2),
+        )
+        self.pair_mask_head = PairMaskHead(beats_per_window, n_mels)
+
+    def unfreeze_encoder(self, n_blocks: int):
+        self.encoder.unfreeze_last_n(n_blocks)
+
+    def forward(self, track_mel: torch.Tensor, orig_mel: torch.Tensor) -> torch.Tensor:
+        t = self.encoder(track_mel)
+        o = self.encoder(orig_mel)
+        combined = torch.cat([t, o, torch.abs(t - o), t * o], dim=-1)
+        return self.head(combined)
+
+
+class ContrastiveSampleDetector(nn.Module):
+    """Siamese AST encoder + projection head trained with CosineEmbeddingLoss."""
+
+    def __init__(
+        self,
+        model_name: str = os.environ["AST_MODEL"],
+        freeze_encoder: bool = True,
+        proj_dim: int = 256,
+        dropout: float = 0.2,
+    ):
+        super().__init__()
+        self.encoder = ASTEncoder(model_name, freeze=freeze_encoder)
+        H = self.encoder.ast.config.hidden_size
+        self.proj = nn.Sequential(
+            nn.Linear(H, 512),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(512, proj_dim),
+        )
+
+    def embed(self, mel: torch.Tensor) -> torch.Tensor:
+        h = self.encoder(mel)
+        # print(f"encoder output: {h.shape}, norm={h.norm(dim=-1).mean():.3f}")
+        z = self.proj(h)
+        return torch.nn.functional.normalize(z, dim=-1)
+
+    def forward(self, track_mel: torch.Tensor, orig_mel: torch.Tensor) -> tuple:
+        return self.embed(track_mel), self.embed(orig_mel)
+
+    def similarity(self, track_mel: torch.Tensor, orig_mel: torch.Tensor) -> torch.Tensor:
+        t, o = self.embed(track_mel), self.embed(orig_mel)
+        return (t * o).sum(dim=-1)
+
+    def unfreeze_encoder(self, n_blocks: int = 2):
+        self.encoder.unfreeze_last_n(n_blocks)

requirements.txt

@@ -0,0 +1,17 @@
+gradio>=5.0
+matplotlib>=3.8
+torch>=2.5
+torchaudio>=2.5
+accelerate==1.13.0
+python-dotenv==1.2.2
+safetensors==0.7.0
+audiomentations==0.43.1
+av==17.0.0
+huggingface-hub==1.10.1
+librosa==0.11.0
+numpy==2.4.4
+scikit-learn==1.8.0
+scipy==1.17.1
+soundfile==0.13.1
+transformers==5.5.4
+yt-dlp==2026.3.17