UmbraAudio/scripts/generate_voice_references.py

# Copyright 2026 Umbra Tools (umbratools.com)
# SPDX-License-Identifier: Apache-2.0

from __future__ import annotations

import json
import logging
from pathlib import Path

from umbraaudio.config import load_config
from umbraaudio.voices.model_coordinator import ModelCoordinator, PipelineStage
from umbraaudio.voices.qwen_tts import QwenTTSEngine, build_voice_description

logger = logging.getLogger(__name__)


def _load_speakers(data_dir: Path) -> list[dict]:
    path = data_dir / "speakers.json"
    if not path.exists():
        raise FileNotFoundError(f"speakers.json not found: {path}")
    with open(path) as f:
        data = json.load(f)
    return data.get("speakers", [])


def _load_profiles(data_dir: Path) -> dict[str, dict]:
    """Load voice profiles from voice_profiles.json, falling back to project.db.

    voice_profiles.json is the canonical source; it contains the full
    voice_description for every speaker. The DB characters table may be
    empty for some projects.
    """
    profiles_path = data_dir / "voice_profiles.json"
    if profiles_path.exists():
        with open(profiles_path) as f:
            data = json.load(f)
        if data:
            result = {}
            for name, profile in data.items():
                result[name] = {
                    "profile_json": json.dumps(profile),
                    "voice_description": profile.get("voice_description"),
                }
            return result

    import sqlite3
    db_path = data_dir / "project.db"
    if not db_path.exists():
        return {}
    db = sqlite3.connect(str(db_path))
    db.row_factory = sqlite3.Row
    pid = db.execute("SELECT id FROM project ORDER BY id LIMIT 1").fetchone()["id"]
    rows = db.execute(
        "SELECT name, profile_json, voice_description FROM characters WHERE project_id=?",
        (pid,),
    ).fetchall()
    db.close()
    result = {}
    for r in rows:
        result[r["name"]] = {
            "profile_json": r["profile_json"],
            "voice_description": r["voice_description"],
        }
    return result


def generate_all_voice_references(
    project_dir: Path | None = None,
    force: bool = False,
) -> None:
    project_dir = project_dir or Path(".")
    cfg = load_config(project_dir)

    coordinator = ModelCoordinator(cfg)
    if not coordinator.switch_to(PipelineStage.VOICE_DESIGN).get("ok"):
        raise RuntimeError("Failed to switch models to VOICE_DESIGN stage")

    speakers = _load_speakers(cfg.data_dir)
    db_profiles = _load_profiles(cfg.data_dir)

    voices_dir = cfg.data_dir / "voices"
    voices_dir.mkdir(parents=True, exist_ok=True)

    tts = QwenTTSEngine(
        device=cfg.tts.device,
        sample_rate=cfg.tts.sample_rate,
        voicedesign_model=cfg.tts.voicedesign_model,
        base_model=cfg.tts.base_model,
        qwen_tts_url=cfg.tts.qwen_tts_url,
    )

    narrator_name = None
    for s in speakers:
        if s.get("is_narrator"):
            narrator_name = s["primary_name"]
            break

    generated = 0
    for s in speakers:
        name = s["primary_name"]
        is_narrator = s.get("is_narrator", False)
        ref_name = name.lower().replace(" ", "_") + "_ref.wav"
        ref_path = voices_dir / ref_name

        if ref_path.exists() and not force:
            logger.info("Skipping %s (exists)", name)
            continue

        db_info = db_profiles.get(name, {})
        voice_desc_override = db_info.get("voice_description")

        if voice_desc_override and not force:
            voice_desc = voice_desc_override
        elif is_narrator:
            voice_desc = (
                "A deep, resonant male voice with a calm, authoritative tone. "
                "Rich baritone timbre, slow and deliberate pacing, "
                "like a late-night radio host. Very controlled and even."
            )
        else:
            profile_data = {}
            if db_info.get("profile_json"):
                try:
                    profile_data = json.loads(db_info["profile_json"])
                except (json.JSONDecodeError, TypeError):
                    pass
            voice_desc = build_voice_description(
                character_name=name,
                profile=profile_data if profile_data else None,
                gender=s.get("gender") or (profile_data or {}).get("gender"),
                energy=(profile_data or {}).get("energy"),
                timbre=(profile_data or {}).get("timbre"),
                pitch_register=(profile_data or {}).get("pitch_register"),
            )

        sample_text = _sample_text(name, is_narrator, s.get("demographics", {}))

        logger.info("[%d/%d] %s: %s", generated + 1, len(speakers), name, voice_desc[:70])
        try:
            tts.generate_voice_reference(
                text=sample_text,
                output_path=ref_path,
                voice_description=voice_desc,
            )
            generated += 1
            logger.info("  -> %s", ref_path.name)
        except Exception as e:
            logger.error("  FAILED %s: %s", name, e)

    tts.unload_all()
    coordinator.switch_to(PipelineStage.IDLE)
    logger.info("Generated %d/%d voice references", generated, len(speakers))


# Rainbow Passage - phonetically balanced paragraph for voice samples
# Source: Grant Fairbanks, Voice and Articulation Drillbook (1960)
# Covers all English phonemes including r-colored vowels
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."
)


def _sample_text(name: str, is_narrator: bool, demographics: dict) -> str:
    return RAINBOW_PASSAGE


if __name__ == "__main__":
    import argparse

    logging.basicConfig(level=logging.INFO, format="%(levelname)s %(message)s")
    parser = argparse.ArgumentParser(description="Generate voice references for all speakers")
    parser.add_argument(
        "--project", type=str, default=None, help="Project name or path"
    )
    parser.add_argument("--force", action="store_true", help="Regenerate existing references")
    args = parser.parse_args()
    generate_all_voice_references(args.project, force=args.force)