VoiceStudio/scripts/render_demos_omnivoice.py

#!/usr/bin/env python3
"""Re-render the demo bundle using the real VoiceStudio TTS engine.

This is the production-quality counterpart to scripts/build_demos.sh, which
uses macOS `say` to bootstrap the demo bundle. Run this once on a machine
with VoiceStudio model weights cached (typically your dev box) to replace the
`say`-rendered placeholders with engine output. Commit the resulting WAVs.

Prerequisites:
  * The project's .venv exists and is activated (`uv sync`).
  * VoiceStudio model weights cached under $HF_HUB_CACHE (the first
    `model.generate()` call will download them otherwise — ~5 GB).
  * Run from the repo root: `python3 scripts/render_demos_omnivoice.py`.

What it produces:
  * backend/assets/samples/demo_voice.wav             (clone reference)
  * backend/assets/samples/demo_clone_output.wav      (clone pre-rendered)
  * backend/assets/samples/voice_design/demo_voice_design_<slug>.wav (7)
  * Updated manifest with rendered_by="omnivoice@<git_sha>"

  * backend/assets/samples/dictation/*.wav (3 replay scripts)

Not regenerated by this script:
  * backend/assets/demo/dubbing/*.mp4 — see scripts/build_dub_demo.sh.

Dictation used to be excluded here on the grounds that `say` was good enough
and engine TTS was overkill. That reasoning only held on macOS: `say` does not
exist on Linux or Windows, so on every other platform the samples simply were
never rendered and the replay buttons pointed at 404s. The engine runs
everywhere the app does, which makes it the portable answer.

Reproducibility:
  * seed=42 fixed across all renders so re-running the script regenerates
    the same audio byte-for-byte (modulo torch nondeterminism).
"""

from __future__ import annotations

import argparse
import json
import os
import subprocess
import sys
import time
from pathlib import Path

REPO_ROOT = Path(__file__).resolve().parent.parent
BACKEND_DIR = REPO_ROOT / "backend"
SAMPLES_DIR = BACKEND_DIR / "assets" / "samples"
VOICE_DESIGN_DIR = SAMPLES_DIR / "voice_design"
DICTATION_DIR = SAMPLES_DIR / "dictation"

# Make `backend/` importable so we can pull personalities + the engine.
sys.path.insert(0, str(BACKEND_DIR))


# Cloning demo — must match scripts/build_demos.sh exactly so the manifest
# stays in sync with what the bootstrap script produced.
CLONE_REF_TEXT = (
    "Hey. I'm the VoiceStudio demo voice. I was made right here, on your "
    "machine: private, local, and ready whenever you are."
)
CLONE_OUTPUT_TEXT = (
    "Welcome aboard. I was just a three-second clip a moment ago. Now I can "
    "say anything you'd like, in your voice or mine."
)

# Original demo identity: a warm, smoky cinematic alto expressed only through
# OmniVoice's supported taxonomy. It is deliberately not based on, trained on,
# or named after a real performer.
CLONE_VOICE_INSTRUCT = "female, young adult, low pitch, american accent"
CLONE_RENDER_STEPS = 48

# Dictation replay scripts. `text` MUST match SCRIPTS in
# frontend/src/components/DictationDemo.jsx verbatim — the card shows that
# string as "what you would say" and then shows what the recogniser heard, so
# any drift between the two reads as a transcription error.
DICTATION_SCRIPTS = [
    {
        "id": "en_conversational",
        "language": "English",
        "instruct": "female, young adult, moderate pitch, american accent",
        "text": (
            "Schedule a meeting with Pat for Tuesday at three PM and remind me "
            "to bring the quarterly report."
        ),
    },
    {
        "id": "en_technical",
        "language": "English",
        "instruct": "male, young adult, moderate pitch, american accent",
        "text": (
            "Patch the WebGPU shader in renderer.tsx, then bump pnpm to nine "
            "point fifteen and rerun the Vitest suite."
        ),
    },
    {
        "id": "fr_reservation",
        "language": "French",
        "instruct": "female, young adult, moderate pitch",
        "text": "Bonjour, je voudrais réserver une table pour deux personnes à vingt heures.",
    },
]


def _git_sha() -> str:
    try:
        return subprocess.check_output(
            ["git", "rev-parse", "--short", "HEAD"],
            cwd=REPO_ROOT, text=True,
        ).strip()
    except Exception:
        return "unknown"


def watermark_file(path: Path, sample_rate: int, *, context: str) -> None:
    """Stamp the provenance watermark into an already-written WAV (#1169).

    These clips ship inside the app and play back to users as VoiceStudio
    output, so they are synthetic audio leaving the app exactly like a
    generated clip — they go through `mark_synthetic`, the one chokepoint
    every producing route uses. `force=True` because a render script runs
    outside the request path that carries the user's watermark preference.

    It runs on the FILE, after loudness normalization, rather than on the
    tensor before it: `loudnorm` re-encodes what it is given, and marking
    first would put the watermark through a gain and true-peak limiter on
    its way to disk.
    """
    import torch
    import torchaudio

    from services.watermark import mark_synthetic

    waveform, rate = torchaudio.load(str(path))
    marked = mark_synthetic(waveform, rate, context=context, force=True)
    if marked is waveform and os.environ.get("OMNIVOICE_DEMO_ALLOW_UNMARKED") == "1":
        print(f"  ! {path.name} is NOT watermarked (OMNIVOICE_DEMO_ALLOW_UNMARKED=1)")
        return
    if marked is waveform:
        # `mark_synthetic` never raises — it degrades, so generation can't be
        # broken by watermarking. A RENDER SCRIPT is the one caller where that
        # is wrong: it exists to produce files a human then commits, and a
        # warning on a scrolling console is not a gate. Fail, so the unmarked
        # file cannot be mistaken for a finished asset.
        raise RuntimeError(
            f"{path.name} could not be watermarked, so it must not be committed. "
            "AudioSeal is missing or its weights are not cached on this machine "
            "(`uv sync --all-extras`, then re-run with the model cache warm). "
            "Set OMNIVOICE_DEMO_ALLOW_UNMARKED=1 only for a local listen — never "
            "for a render you intend to commit."
        )
    torchaudio.save(
        str(path),
        marked.to(torch.float32),
        rate,
        encoding="PCM_S", bits_per_sample=16,
    )


def _save_wav(audio_tensor, sample_rate: int, out_path: Path):
    """Save a torch tensor (C, T) or (T,) to a 16-bit PCM WAV."""
    import torch
    import torchaudio

    if audio_tensor.dim() == 1:
        audio_tensor = audio_tensor.unsqueeze(0)
    # Ensure mono — most VoiceStudio outputs are mono already.
    if audio_tensor.shape[0] > 1:
        audio_tensor = audio_tensor.mean(dim=0, keepdim=True)
    out_path.parent.mkdir(parents=True, exist_ok=True)
    # Peak-normalize for headroom only; perceived level is set by the loudness
    # pass below. Peak alone is not enough on its own: diffusion TTS emits the
    # occasional single-sample transient, and one click is all it takes to hold
    # the rest of the clip down — the Helpdesk preset landed at -30 dB RMS
    # against -17 dB for its neighbours that way, so a preview row played at
    # wildly different volumes depending on which preset you clicked.
    peak = audio_tensor.abs().max().item()
    if peak > 0:
        audio_tensor = audio_tensor / peak * 0.97
    torchaudio.save(
        str(out_path),
        audio_tensor.to(torch.float32),
        sample_rate,
        encoding="PCM_S", bits_per_sample=16,
    )
    _normalize_loudness(out_path, sample_rate)
    watermark_file(out_path, sample_rate, context=f"demo:{out_path.stem}")


# Preview clips are played back to back in a picker, so they have to sit at the
# same perceived level — which peak normalization does not give you: a clip
# whose speech is quiet under one loud transient normalizes to the same peak as
# a clip that is loud throughout, and plays 12 dB softer. EBU R128 measures
# loudness rather than amplitude, and its true-peak ceiling keeps the transient
# legal without pulling the body of the clip down with it.
#
# ffmpeg is already a documented prerequisite of the demo tooling (see
# scripts/build_demos.sh). If it is missing, the raw render is still written and
# usable — the clips are simply not level-matched, which is a cosmetic loss, not
# a broken asset.
_LOUDNESS_TARGET_LUFS = -18.0
_TRUE_PEAK_CEILING_DBTP = -1.5


def _normalize_loudness(path: Path, sample_rate: int) -> None:
    import shutil
    import tempfile

    if shutil.which("ffmpeg") is None:
        print(f"  ! ffmpeg not found — {path.name} left un-levelled")
        return
    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as handle:
        tmp = Path(handle.name)
    try:
        result = subprocess.run(
            [
                "ffmpeg", "-y", "-loglevel", "error", "-i", str(path),
                "-af",
                f"loudnorm=I={_LOUDNESS_TARGET_LUFS}:TP={_TRUE_PEAK_CEILING_DBTP}:LRA=11",
                # `loudnorm` resamples to 192 kHz internally for true-peak
                # measurement and will happily WRITE at 192 kHz if the output
                # rate is not pinned — which turned 2.1 MB of previews into
                # 17.5 MB of identical-sounding audio the first time this ran.
                "-ar", str(sample_rate),
                "-c:a", "pcm_s16le", str(tmp),
            ],
            capture_output=True, text=True,
        )
        if result.returncode != 0 or not tmp.exists() or tmp.stat().st_size == 0:
            print(f"  ! loudnorm failed for {path.name}: {result.stderr.strip()[:120]}")
            return
        # os.replace, not shutil.move: `path` already exists, so move delegates
        # to os.rename — which raises FileExistsError on Windows and fails the
        # render there. os.replace overwrites atomically on every platform.
        os.replace(str(tmp), str(path))
    finally:
        tmp.unlink(missing_ok=True)


def render_cloning(model, args):
    """Render the cloning demo: reference clip + pre-rendered output.

    The reference clip is itself synthesized — chicken-and-egg, but the
    VoiceStudio engine in non-zero-shot mode (no ref_audio) accepts a plain
    `instruct=` taxonomy string and produces a clean voice.
    """
    print("── Cloning demo ─────────────────────────────────────")
    sr = getattr(model, "sampling_rate", 24000)
    from omnivoice.utils.common import fix_random_seed

    # 1) Reference clip — synthesized as an original cinematic alto so the
    #    timbre is distinctive while remaining reproducible and rights-safe.
    out_ref = SAMPLES_DIR / "demo_voice.wav"
    if args.skip_existing and out_ref.exists():
        print(f"  · skip (exists): {out_ref.name}")
    else:
        fix_random_seed(42)
        audios = model.generate(
            text=CLONE_REF_TEXT,
            instruct=CLONE_VOICE_INSTRUCT,
            num_step=CLONE_RENDER_STEPS,
        )
        _save_wav(audios[0], sr, out_ref)
        print(f"  ✓ {out_ref.name} ({sr} Hz, omnivoice)")

    # 2) Pre-rendered clone output — same voice, different text. Use the
    #    reference clip we just rendered as the speaker reference so this
    #    actually demonstrates cloning rather than independent synthesis.
    out_clone = SAMPLES_DIR / "demo_clone_output.wav"
    if args.skip_existing and out_clone.exists():
        print(f"  · skip (exists): {out_clone.name}")
    else:
        fix_random_seed(42)
        audios = model.generate(
            text=CLONE_OUTPUT_TEXT,
            ref_audio=str(out_ref),
            ref_text=CLONE_REF_TEXT,
            num_step=CLONE_RENDER_STEPS,
        )
        _save_wav(audios[0], sr, out_clone)
        print(f"  ✓ {out_clone.name} ({sr} Hz, cloned)")


def render_voice_design(model, args):
    """Re-render the 7 voice-design preset previews."""
    print("\n── Voice design presets ─────────────────────────────")
    from core.personalities import PERSONALITIES

    sr = getattr(model, "sampling_rate", 24000)
    demos = [p for p in PERSONALITIES if p.get("is_demo")]
    if not demos:
        print("  No is_demo presets found in personalities.py")
        return

    for preset in demos:
        slug = preset["id"]
        out = VOICE_DESIGN_DIR / f"demo_voice_design_{slug}.wav"
        if args.skip_existing and out.exists():
            print(f"  · skip (exists): {out.name}")
            continue
        audios = model.generate(
            text=preset["script"],
            instruct=preset["instruct"],
            language=preset.get("language"),
            num_step=24,
        )
        _save_wav(audios[0], sr, out)
        print(f"  ✓ {out.name} ({preset['name']})")


def render_dictation(model, args):
    """Render the three dictation replay clips.

    The replay path posts these to /transcribe and shows the recognized text,
    so the demo works without microphone permission — on a VM, in CI, or before
    the user has granted access. That makes the clip content load-bearing: it
    has to be what the card says it is, or the demo shows a mismatch.
    """
    print("\n── Dictation replay clips ───────────────────────────")
    sr = getattr(model, "sampling_rate", 24000)
    for script in DICTATION_SCRIPTS:
        out = DICTATION_DIR / f"{script['id']}.wav"
        if args.skip_existing and out.exists():
            print(f"  · skip (exists): {out.name}")
            continue
        audios = model.generate(
            text=script["text"],
            instruct=script["instruct"],
            language=script["language"],
            num_step=32,
        )
        _save_wav(audios[0], sr, out)
        print(f"  ✓ {out.name} ({script['language']})")


def update_manifest(args):
    """Update the existing manifest with rendered_by + rendered_at."""
    print("\n── Manifest ─────────────────────────────────────────")
    mpath = SAMPLES_DIR / "demo" / "dubbing" / "manifest.json"
    if not mpath.exists():
        print(f"  ! manifest not found at {mpath} — run scripts/build_dub_demo.sh first")
        return
    data = json.loads(mpath.read_text())
    data["rendered_by"] = f"omnivoice@{_git_sha()}"
    data["rendered_at"] = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
    mpath.write_text(json.dumps(data, indent=2, ensure_ascii=False))
    print(f"  ✓ {mpath.name}")


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--skip-existing", action="store_true",
        help="Don't re-render files that already exist on disk.",
    )
    parser.add_argument(
        "--only", choices=["cloning", "design", "dictation", "manifest"],
        help="Render only a subset (default: all).",
    )
    args = parser.parse_args()

    print("Loading VoiceStudio engine (this can take 30-60 s on first run)…")
    try:
        import asyncio
        from services.model_manager import get_model
        try:
            asyncio.get_running_loop()
            raise RuntimeError("Run this script outside an async context.")
        except RuntimeError:
            model = asyncio.run(get_model())
    except Exception as e:
        print(f"\nERROR: Could not load VoiceStudio engine: {e}\n")
        print("Check that:")
        print("  1. You're running inside the project venv (uv sync first).")
        print("  2. The omnivoice package is importable: `python -c 'import omnivoice'`.")
        print("  3. Model weights are downloaded (~5 GB on first synthesis).")
        sys.exit(1)
    print("Engine loaded.\n")

    if args.only in (None, "cloning"):
        render_cloning(model, args)
    if args.only in (None, "design"):
        render_voice_design(model, args)
    if args.only in (None, "dictation"):
        render_dictation(model, args)
    if args.only in (None, "manifest"):
        update_manifest(args)

    print("\nDone. Dubbing videos are built separately — scripts/build_dub_demo.sh.")


if __name__ == "__main__":
    main()