UmbraAudio/scripts/generate_voice_samples.py

#!/usr/bin/env python3
# Copyright 2026 Umbra Tools (umbratools.com)
# SPDX-License-Identifier: Apache-2.0

"""Generate Rainbow Passage samples from cached speaker embeddings.

For each speaker in the given chapter, generates a WAV file of the Rainbow Passage
spoken in that character's voice. Used for human review/approval before full generation.

Usage:
    python scripts/generate_voice_samples.py --chapter 15
    python scripts/generate_voice_samples.py --speakers narrator,nanny,kainda
    python scripts/generate_voice_samples.py --all  # all speakers with embeddings
"""
from __future__ import annotations

import argparse
import json
import sqlite3
import sys
import time
from pathlib import Path

import numpy as np
import soundfile as sf

project_root = Path(__file__).resolve().parent.parent
if str(project_root) not in sys.path:
    sys.path.insert(0, str(project_root))

from umbraaudio.config import load_config  # noqa: E402
from umbraaudio.voices.qwen_embeddings import QwenEmbeddingCache  # noqa: E402

RAINBOW_PASSAGE = (
    "When the sunlight strikes raindrops in the air, they act as a prism and form a rainbow. "
    "The rainbow is a division of white light into many beautiful colors. "
    "These take the shape of a long round arch, with its path high above, "
    "and its two ends apparently beyond the horizon. "
    "There is, according to legend, a boiling pot of gold at one end. "
    "People look, but no one ever finds it. "
    "When a man looks for something beyond his reach, "
    "his friends say he is looking for the pot of gold at the end of the rainbow. "
    "When the sunlight strikes raindrops in the air, they act as a prism and form a rainbow. "
    "The rainbow is a division of white light into many beautiful colors. "
    "These take the shape of a long round arch, with its path high above, "
    "and its two ends apparently beyond the horizon. "
    "There is, according to legend, a boiling pot of gold at one end. "
    "People look, but no one ever finds it. "
    "When a man looks for something beyond his reach, "
    "his friends say he is looking for the pot of gold at the end of the rainbow."
)


def name_to_slug(name: str) -> str:
    return name.lower().replace(' ', '_').replace('.', '').replace("'", '').replace('(', '').replace(')', '').replace('-', '_').replace('/', '_')


def get_chapter_speakers(project_name: str, chapter_id: int) -> list[str]:
    config = load_config(project_name)
    db_path = config.data_dir / "project.db"
    db = sqlite3.connect(str(db_path))
    rows = db.execute('''
        SELECT DISTINCT speaker_name FROM utterances 
        WHERE chapter_id=? AND speaker_name IS NOT NULL
    ''', (chapter_id,)).fetchall()
    db.close()
    return [r[0] for r in rows]


def generate_samples(
    project_name: str,
    speaker_names: list[str],
    output_dir: Path,
    temperature: float = 0.5,
    seed: int = 42,
) -> list[Path]:
    from umbraaudio.voices.qwen_tts import QwenTTSEngine

    config = load_config(project_name)
    embeddings_dir = config.data_dir / "voice_embeddings"
    cache = QwenEmbeddingCache(embeddings_dir)

    output_dir.mkdir(parents=True, exist_ok=True)
    engine = QwenTTSEngine(qwen_tts_url="http://localhost:8881")

    generated = []
    for name in speaker_names:
        slug = name_to_slug(name)
        if not cache.has(slug):
            print(f"  SKIP {name}: no embedding cached (slug={slug})")
            continue

        out_path = output_dir / f"{slug}_rainbow.wav"
        t0 = time.time()
        # Use a dummy ref path — the engine will find the cached embedding
        dummy_ref = config.data_dir / "voices" / f"{slug}_ref.wav"
        if not dummy_ref.exists():
            dummy_ref = config.data_dir / "voices" / "narrator_ref.wav"
        engine.generate(
            text=RAINBOW_PASSAGE,
            output_path=out_path,
            voice_ref=dummy_ref,
            temperature=temperature,
            seed=seed,
        )
        elapsed = time.time() - t0
        size_kb = out_path.stat().st_size / 1024
        print(f"  OK {name:40s} -> {out_path.name:40s} ({elapsed:.1f}s, {size_kb:.0f}KB)")
        generated.append(out_path)

    return generated


def main() -> None:
    parser = argparse.ArgumentParser(description="Generate voice samples for review")
    parser.add_argument("--project", default="DealHunterShip_4.0e_whole-to-audio")
    parser.add_argument("--chapter", type=int, help="Generate samples for all speakers in this chapter")
    parser.add_argument("--speakers", help="Comma-separated speaker names")
    parser.add_argument("--all", action="store_true", help="All speakers with embeddings")
    parser.add_argument("--temperature", type=float, default=0.5)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--output-dir", default=None, help="Output directory")
    args = parser.parse_args()

    if args.chapter:
        speakers = get_chapter_speakers(args.project, args.chapter)
        out_dir = Path(args.output_dir or f"/tmp/voice_samples_ch{args.chapter}")
        print(f"Chapter {args.chapter}: {len(speakers)} speakers")
    elif args.speakers:
        speakers = [s.strip() for s in args.speakers.split(",")]
        out_dir = Path(args.output_dir or "/tmp/voice_samples")
    elif args.all:
        config = load_config(args.project)
        cache = QwenEmbeddingCache(config.data_dir / "voice_embeddings")
        speakers = [s.replace("_", " ").title() for s in cache.list_cached()]
        out_dir = Path(args.output_dir or "/tmp/voice_samples_all")
    else:
        parser.error("Specify --chapter, --speakers, or --all")

    print(f"Output: {out_dir}")
    t0 = time.time()
    files = generate_samples(args.project, speakers, out_dir, args.temperature, args.seed)
    elapsed = time.time() - t0
    print(f"\nDone: {len(files)} samples in {elapsed:.1f}s")


if __name__ == "__main__":
    main()