UmbraAudio/scripts/generate_test_samples.py

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

"""Generate test samples for all speakers using Edge-TTS with speaker attributes."""
from __future__ import annotations

import asyncio
import json
import sys
import time
from pathlib import Path

# Add project root to path
sys.path.insert(0, str(Path(__file__).parent.parent))

from umbraaudio.config import load_config
from umbraaudio.voices.edge_tts import EdgeTTSEngine

# Rainbow Passage - phonetically balanced paragraph for voice samples
# Source: Grant Fairbanks, Voice and Articulation Drillbook (1960)
# Covers all English phonemes including r-colored vowels
VERIFICATION_PHRASE = (
    "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 _map_traits_to_energy(traits: list[str]) -> str:
    """Map vocal traits to energy level."""
    traits_lower = [t.lower() for t in traits]
    if any(t in traits_lower for t in ["explosive", "thunderous", "shouting"]):
        return "explosive"
    if any(t in traits_lower for t in ["volatile", "intense", "agitated", "emotional"]):
        return "volatile"
    if any(t in traits_lower for t in ["commanding", "authoritative", "powerful", "dominant"]):
        return "commanding"
    if any(t in traits_lower for t in ["assertive", "confident", "firm", "strong"]):
        return "assertive"
    if any(t in traits_lower for t in ["measured", "even", "controlled", "steady"]):
        return "measured"
    if any(t in traits_lower for t in ["soft", "gentle", "soft_spoken", "quiet"]):
        return "soft_spoken"
    if any(t in traits_lower for t in ["subdued", "muted", "restrained", "hushed"]):
        return "subdued"
    if any(t in traits_lower for t in ["whispered", "whisper", "barely_audible"]):
        return "whispered"
    return "measured"


def _map_traits_to_pitch(traits: list[str], gender: str) -> str:
    """Map vocal traits to pitch register."""
    traits_lower = [t.lower() for t in traits]
    if any(t in traits_lower for t in ["very_high", "ultra_high", "squeaky", "falsetto"]):
        return "very_high" if gender == "Male" else "ultra_high"
    if any(t in traits_lower for t in ["high", "high_pitched", "soprano"]):
        return "high"
    if any(t in traits_lower for t in ["mid_high", "slightly_high"]):
        return "mid_high"
    if any(t in traits_lower for t in ["mid", "medium", "moderate", "average"]):
        return "mid"
    if any(t in traits_lower for t in ["mid_low", "slightly_low"]):
        return "mid_low"
    if any(t in traits_lower for t in ["low", "deep", "bass", "baritone"]):
        return "low"
    if any(t in traits_lower for t in ["very_low", "very_deep"]):
        return "very_low"
    # Default based on gender
    return "mid" if gender == "Male" else "mid_high"


def _extract_accent(traits: list[str], accent_field: str | None) -> str:
    """Extract accent from traits or accent field. Default to Canadian."""
    if accent_field and accent_field.lower() not in ("none", "unknown", ""):
        return accent_field.lower()
    traits_lower = [t.lower() for t in traits]
    accent_keywords = {
        "british": ["british", "uk", "english", "cockney", "rp"],
        "american": ["american", "us", "general_american"],
        "australian": ["australian", "aussie"],
        "indian": ["indian", "south_asian"],
        "irish": ["irish", "irish_accent"],
        "canadian": ["canadian", "canada"],
        "nigerian": ["nigerian", "west_african"],
        "kenyan": ["kenyan", "east_african"],
        "hong_kong": ["hong_kong", "hk", "cantonese"],
    }
    for accent, keywords in accent_keywords.items():
        if any(k in " ".join(traits_lower) for k in keywords):
            return accent
    return "canadian"  # Default accent


def _generate_with_retry(engine: EdgeTTSEngine, text: str, output_path: Path, **kwargs) -> bool:
    """Generate audio with retry logic for transient Edge-TTS failures."""
    max_retries = 3
    base_delay = 2.0
    
    for attempt in range(max_retries):
        try:
            engine.generate_voice_reference(
                text=text,
                output_path=output_path,
                **kwargs,
            )
            if output_path.exists():
                return True
        except Exception as e:
            if attempt == max_retries - 1:
                raise
            delay = base_delay * (2 ** attempt)
            print(f"  Retry {attempt + 1}/{max_retries} after {delay:.1f}s: {e}")
            time.sleep(delay)
    return False


async def generate_test_samples(project: str | None = None):
    """Generate test samples for all speakers using their vocal traits."""
    cfg = load_config(project)

    # Load voice assignments
    assignments_path = cfg.data_dir / "voice_assignments.json"
    with open(assignments_path) as f:
        assignments = json.load(f)

    # Load speakers for vocal traits
    speakers_path = cfg.data_dir / "speakers.json"
    with open(speakers_path) as f:
        speakers_data = json.load(f)
    speakers = {s["primary_name"]: s for s in speakers_data["speakers"]}

    print(f"Generating test samples for {len(assignments)} speakers...")
    print(f"Verification phrase: \"{VERIFICATION_PHRASE}\"")
    print()

    # Create output directory
    output_dir = cfg.data_dir / "voice_tests"
    output_dir.mkdir(exist_ok=True)

    # Initialize Edge-TTS engine
    engine = EdgeTTSEngine()

    # Track results
    success_count = 0
    fail_count = 0

    for speaker_name, voice_assignment in sorted(assignments.items()):
        speaker = speakers.get(speaker_name, {})
        demographics = speaker.get("demographics", {})
        gender = demographics.get("gender", "Unknown")
        stype = demographics.get("type", "Human")
        traits = speaker.get("vocal_traits", [])
        accent_field = speaker.get("accent")

        # Map traits to Edge-TTS parameters
        energy = _map_traits_to_energy(traits)
        pitch = _map_traits_to_pitch(traits, gender)
        accent = _extract_accent(traits, accent_field)

        # Parse voice assignment for pitch shift (fallback)
        if ":" in voice_assignment:
            base_voice, pitch_shift = voice_assignment.split(":")
            pitch_shift = int(pitch_shift)
        else:
            base_voice = voice_assignment
            pitch_shift = 0

        # Create safe filename
        safe_name = speaker_name.replace(" ", "_").replace("/", "_")
        output_path = output_dir / f"{safe_name}.wav"

        # Skip if already exists
        if output_path.exists():
            success_count += 1
            print(f"⊙ {speaker_name:35} -> already exists")
            continue

        try:
            # Generate audio using speaker attributes (accent defaults to Canadian)
            success = _generate_with_retry(
                engine=engine,
                text=VERIFICATION_PHRASE,
                output_path=output_path,
                energy=energy,
                pitch=pitch,
                gender=gender,
                character_name=speaker_name,
                accent=accent,
            )

            if success and output_path.exists():
                success_count += 1
                print(f"✓ {speaker_name:35} -> energy={energy}, pitch={pitch}, gender={gender}, accent={accent}")
            else:
                fail_count += 1
                print(f"✗ {speaker_name:35} -> (file not created)")

        except Exception as e:
            fail_count += 1
            print(f"✗ {speaker_name:35} -> (error: {e})")

        # Small delay between requests to avoid rate limiting
        time.sleep(0.5)

    print("\n=== Summary ===")
    print(f"Success: {success_count}")
    print(f"Failed: {fail_count}")
    print(f"Total: {len(assignments)}")
    print(f"\nTest samples saved to: {output_dir}")


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="Generate test samples for all speakers")
    parser.add_argument(
        "--project", type=str, default=None, help="Project name or path"
    )
    args = parser.parse_args()
    asyncio.run(generate_test_samples(project=args.project))