scripts/build_heldout_refs.py

"""Build the 15.08 s held-out references from the dataset dry sources.

The a2a render length is 377 frames @ 25 fps = 15.08 s, but every dataset
source is exactly 6.000 s. So the reference is the dry clip repeated with a
short gap and trimmed to length — the same construction used for the eval
references on 08-13, which lived inline in vast_push.sh and was therefore lost
when that box died. It is reproducible, so it belongs in a file.

Held-out on purpose: beat_02 and male_voice_02 are the VALIDATION split, never
seen in training, and both are owned/cleared — which makes them the only
sources usable for public examples.

Reads and writes 24-bit stereo WAV at 48 kHz with no external dependencies
beyond numpy, so it runs on a bare pod before any venv exists.
"""

import argparse
import sys
import wave
from pathlib import Path

import numpy as np

SR = 48000
TARGET_S = 15.08
GAP_S = 1.54


def read24(path: Path) -> np.ndarray:
    with wave.open(str(path), "rb") as f:
        if f.getsampwidth() != 3:
            raise SystemExit(f"{path}: expected 24-bit, got {f.getsampwidth()*8}-bit")
        if f.getframerate() != SR:
            raise SystemExit(f"{path}: expected {SR} Hz, got {f.getframerate()}")
        ch = f.getnchannels()
        raw = f.readframes(f.getnframes())
    b = np.frombuffer(raw, dtype=np.uint8).reshape(-1, ch, 3).astype(np.int32)
    v = b[..., 0] | (b[..., 1] << 8) | (b[..., 2] << 16)
    v = np.where(v >= 1 << 23, v - (1 << 24), v)          # sign-extend
    return v.astype(np.float64) / 8388607.0


def save24(x: np.ndarray, path: Path) -> None:
    i = (np.clip(x, -1.0, 1.0) * 8388607.0).astype(np.int32)
    bb = np.stack([i & 255, (i >> 8) & 255, (i >> 16) & 255], axis=-1).astype(np.uint8)
    with wave.open(str(path), "wb") as w:
        w.setnchannels(x.shape[1])
        w.setsampwidth(3)
        w.setframerate(SR)
        w.writeframes(bb.tobytes())


def build(src: Path, dst: Path) -> None:
    x = read24(src)
    n_target = int(TARGET_S * SR)
    gap = np.zeros((int(GAP_S * SR), x.shape[1]))
    out = x
    while len(out) < n_target:                            # repeat until long enough
        out = np.concatenate([out, gap, x])
    out = out[:n_target]
    save24(out, dst)
    peak = float(np.max(np.abs(out))) or 1e-12
    print(f"  {dst.name}: {len(out)/SR:.2f}s, peak {20*np.log10(peak):.1f} dBFS")


def main() -> int:
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--dry-dir", default="/workspace/AUDIO-LTX-LORA/Dataset/v5-rooms/medium_room_1_9/dry")
    p.add_argument("--out-dir", default="/workspace")
    a = p.parse_args()

    pairs = [("male_voice_02.wav", "v5_eval_voice_15s.wav"),
             ("beat_02.wav", "v5_eval_beat_15s.wav")]
    dry, out = Path(a.dry_dir), Path(a.out_dir)
    missing = [s for s, _ in pairs if not (dry / s).exists()]
    if missing:
        print(f"MISSING held-out sources in {dry}: {missing}", file=sys.stderr)
        print("The payload subset must include the VALIDATION split, not just train rows.", file=sys.stderr)
        return 2
    for src, dest in pairs:
        build(dry / src, out / dest)
    print("REFS_BUILT")
    return 0


if __name__ == "__main__":
    sys.exit(main())