scripts/compute_reference_stats.py

#!/usr/bin/env python3
"""Compute SR-FD reference statistics from an audio manifest.

This runs the *frozen* SR-FD extractors directly on reference waveforms and
stores only the first- and second-order moments (mean + covariance) of the
resulting feature vectors. The reference audio is never used again at training
time; only the stored moments are loaded.

Each of the paper's three targets is one invocation of this script:

    # 1. low-step Whisper anchor (ASR-verified four-step generations)
    python scripts/compute_reference_stats.py \
        --manifest data/ref/asr_true4_good.jsonl \
        --config configs/srfd_compact3.yaml --reps whisper_anchor8_p64 \
        --out stats/ref_whisper_anchor_asr_true4_good.pt

    # 2. teacher CTC target (ten-step teacher generations)
    python scripts/compute_reference_stats.py \
        --manifest data/ref/teacher_t10.jsonl \
        --config configs/srfd_compact3.yaml --reps ctc_content_p64 \
        --out stats/ref_ctc_content_teacher_t10.pt

    # 3. real-speech CTC target (real LibriTTS voice-cloning speech)
    python scripts/compute_reference_stats.py \
        --manifest data/ref/real_voiceclone.jsonl \
        --config configs/srfd_compact3.yaml --reps ctc_content_p64 \
        --out stats/ref_ctc_content_real_voiceclone.pt

The manifest is JSONL with one object per line containing at least an
``audio`` field (path to a wav/flac file). Any other fields are ignored.
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path
from typing import Dict, List, Optional

PROJECT_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(PROJECT_ROOT))

import torch

from srfd.extractors import build_srfd_extractors
from srfd.moments import accumulate_moments, finalize_accumulated_moments
from srfd.stats_io import save_stats


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description="Compute SR-FD reference statistics.")
    p.add_argument("--manifest", required=True, help="JSONL manifest with an 'audio' field.")
    p.add_argument("--config", required=True, help="YAML config; reads srfd.reps for extractors.")
    p.add_argument("--out", required=True, help="Output .pt path.")
    p.add_argument(
        "--reps",
        nargs="*",
        default=None,
        help="Subset of extractor names to compute (default: all enabled in the config).",
    )
    p.add_argument("--audio_field", default="audio", help="Manifest field with the audio path.")
    p.add_argument("--input_sample_rate", type=int, default=0,
                   help="Override input sample rate (0 = use each file's native rate).")
    p.add_argument("--batch_size", type=int, default=4)
    p.add_argument("--max_samples", type=int, default=0, help="0 = process all rows.")
    p.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
    p.add_argument("--dry_run", action="store_true", help="Process 2 batches and print shapes only.")
    return p.parse_args()


def _load_reps(config_path: str, keep: Optional[List[str]]) -> List[Dict]:
    import yaml

    with open(config_path, "r", encoding="utf-8") as f:
        data = yaml.safe_load(f)
    reps = (data.get("srfd", {}) or {}).get("reps", []) if isinstance(data, dict) else []
    if not reps:
        raise ValueError(f"No srfd.reps found in {config_path}")
    if keep:
        keep_set = set(keep)
        reps = [r for r in reps if r.get("name") in keep_set]
        if not reps:
            raise ValueError(f"None of {keep} matched srfd.reps names in {config_path}")
    return reps


def _read_manifest(path: str, audio_field: str, limit: int) -> List[str]:
    paths: List[str] = []
    with open(path, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            row = json.loads(line)
            audio = row[audio_field]
            if isinstance(audio, dict):  # {"path": ...} style
                audio = audio.get("path") or audio.get("array")
            paths.append(str(audio))
            if limit and len(paths) >= limit:
                break
    return paths


def _load_audio(path: str) -> tuple[torch.Tensor, int]:
    import soundfile as sf

    wav, sr = sf.read(path, dtype="float32", always_2d=True)
    # Downmix to mono: [T, C] -> [T]
    mono = torch.from_numpy(wav).mean(dim=1)
    return mono, int(sr)


def _collate(batch_paths: List[str], override_sr: int) -> Dict[str, torch.Tensor]:
    wavs: List[torch.Tensor] = []
    lengths: List[int] = []
    sr_seen: Optional[int] = None
    for p in batch_paths:
        wav, sr = _load_audio(p)
        if override_sr:
            sr = override_sr
        if sr_seen is None:
            sr_seen = sr
        elif sr_seen != sr:
            raise ValueError(
                f"Mixed sample rates in one batch ({sr_seen} vs {sr}); "
                "pre-resample the manifest or use --batch_size 1."
            )
        wavs.append(wav)
        lengths.append(wav.numel())
    max_len = max(lengths)
    B = len(wavs)
    waveform = torch.zeros(B, max_len, dtype=torch.float32)
    mask = torch.zeros(B, max_len, dtype=torch.bool)
    for i, (wav, n) in enumerate(zip(wavs, lengths)):
        waveform[i, :n] = wav
        mask[i, :n] = True
    return {
        "waveform": waveform,
        "waveform_mask": mask,
        "waveform_sample_rate": int(sr_seen or 16000),
    }


def main() -> int:
    args = parse_args()

    reps_config = _load_reps(args.config, args.reps)
    extractors = build_srfd_extractors(reps_config)
    if not extractors:
        raise ValueError("No enabled extractors after filtering.")
    extractors = [e.to(args.device).eval() for e in extractors]
    print(f"Extractors: {[e.name for e in extractors]}", file=sys.stderr)

    paths = _read_manifest(args.manifest, args.audio_field, args.max_samples)
    print(f"Manifest rows: {len(paths)}", file=sys.stderr)

    accumulators: Dict[str, dict] = {ext.name: {} for ext in extractors}
    n_batches = 0
    with torch.no_grad():
        for start in range(0, len(paths), args.batch_size):
            batch_paths = paths[start : start + args.batch_size]
            batch = _collate(batch_paths, args.input_sample_rate)
            batch = {
                k: (v.to(args.device) if isinstance(v, torch.Tensor) else v)
                for k, v in batch.items()
            }
            for ext in extractors:
                rep = ext(batch).to(torch.float32)
                if rep.dim() != 2:
                    raise RuntimeError(f"{ext.name} returned {tuple(rep.shape)}; expected [B, C]")
                accumulate_moments(accumulators[ext.name], rep)
                if args.dry_run and n_batches < 2:
                    print(f"[dry_run] {ext.name}: rep {tuple(rep.shape)}", file=sys.stderr)
            n_batches += 1
            if args.dry_run and n_batches >= 2:
                print("[dry_run] stopping after 2 batches.", file=sys.stderr)
                return 0

    reps_out = {}
    for ext in extractors:
        finalized = finalize_accumulated_moments(accumulators[ext.name])
        reps_out[ext.name] = finalized
        print(
            f"[stats] {ext.name}: n={finalized['n']}, dim={finalized['feature_dim']}, "
            f"||mu||={finalized['mu'].norm().item():.4f}",
            file=sys.stderr,
        )

    metadata = {
        "manifest": str(args.manifest),
        "config": str(args.config),
        "n_samples_processed": int(min(len(paths), args.max_samples) if args.max_samples else len(paths)),
        "extractor_names": [e.name for e in extractors],
    }
    out_path = Path(args.out)
    out_path.parent.mkdir(parents=True, exist_ok=True)
    save_stats(str(out_path), {"reps": reps_out, "metadata": metadata})
    print(f"Saved SR-FD reference stats to {out_path}", file=sys.stderr)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())