UmbraAudio/scripts/assign_voices.py

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

"""Assign unique Edge-TTS voices to all speakers."""
from __future__ import annotations

import hashlib
import json
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent.parent))

from umbraaudio.config import load_config

# Edge-TTS voices organized by gender and pitch
VOICES = {
    "male_low": [
        "en-US-GuyNeural",
        "en-GB-ThomasNeural",
        "en-NG-AbeoNeural",
    ],
    "male_mid": [
        "en-US-ChristopherNeural",
        "en-US-JasonNeural",
        "en-US-TonyNeural",
        "en-GB-RyanNeural",
        "en-AU-WilliamMultilingualNeural",
        "en-IN-PrabhatNeural",
        "en-IE-ConnorNeural",
        "en-CA-LiamNeural",
        "en-KE-ChilembaNeural",
    ],
    "female_mid": [
        "en-US-JennyNeural",
        "en-US-AvaNeural",
        "en-US-MichelleNeural",
        "en-GB-SoniaNeural",
        "en-AU-NatashaNeural",
        "en-IN-NeerjaNeural",
        "en-IE-EmilyNeural",
        "en-CA-ClaraNeural",
        "en-NG-EzinneNeural",
        "en-KE-AsiliaNeural",
    ],
    "female_high": [
        "en-US-SaraNeural",
        "en-GB-LibbyNeural",
        "en-HK-YanNeural",
    ],
    "child": [
        "en-US-AnaNeural",
    ],
}

# All voices in a flat list for fallback
ALL_VOICES = []
for voice_list in VOICES.values():
    ALL_VOICES.extend(voice_list)


def get_voice_category(speaker: dict) -> str:
    """Determine which voice category a speaker should use."""
    demos = speaker.get("demographics", {})
    gender = demos.get("gender", "Unknown").lower()
    stype = demos.get("type", "Human").lower()

    # Check for AI/Robot/Computer types
    if stype in ("robot", "ai", "computer", "machine", "system"):
        return "male_mid"  # Robots get a neutral mid voice

    # Check for child
    if "child" in stype.lower():
        return "child"

    # Gender-based assignment
    if gender == "male":
        return "male_mid"
    elif gender == "female":
        return "female_mid"
    else:
        # Unknown gender - use mid voices
        return "male_mid"


def assign_voices(speakers: list[dict]) -> dict[str, str]:
    """Assign unique voices to all speakers."""
    voice_assignments = {}
    used_voices = set()

    # Track available voices by category
    available = {cat: list(voices) for cat, voices in VOICES.items()}

    # Sort speakers by chapter count (more chapters = higher priority)
    sorted_speakers = sorted(
        speakers,
        key=lambda s: len(s.get("chapters_appeared", [])),
        reverse=True
    )

    for speaker in sorted_speakers:
        name = speaker["primary_name"]
        category = get_voice_category(speaker)

        # Try to find an unused voice in the preferred category
        assigned = False
        if available.get(category):
            for voice in available[category]:
                if voice not in used_voices:
                    voice_assignments[name] = voice
                    used_voices.add(voice)
                    available[category].remove(voice)
                    assigned = True
                    break

        # If no voice in preferred category, try other categories
        if not assigned:
            for cat, voices in available.items():
                for voice in voices:
                    if voice not in used_voices:
                        voice_assignments[name] = voice
                        used_voices.add(voice)
                        available[cat].remove(voice)
                        assigned = True
                        break
                if assigned:
                    break

        # If still no voice, use a hash-based fallback with pitch shifting
        if not assigned:
            name_hash = int(hashlib.md5(name.encode()).hexdigest()[:4], 16)
            base_voice = ALL_VOICES[name_hash % len(ALL_VOICES)]
            # Add pitch shift indicator
            pitch_shift = (name_hash % 200) - 100  # -100 to +100
            voice_assignments[name] = f"{base_voice}:{pitch_shift:+d}"

    return voice_assignments


def main():
    import argparse

    parser = argparse.ArgumentParser(description="Assign unique Edge-TTS voices to speakers")
    parser.add_argument(
        "--project", type=str, default=None, help="Project name or path"
    )
    args = parser.parse_args()

    cfg = load_config(args.project)

    # Load speakers
    speakers_path = cfg.data_dir / "speakers.json"
    with open(speakers_path) as f:
        data = json.load(f)

    speakers = data["speakers"]
    print(f"Assigning voices to {len(speakers)} speakers...")

    # Assign voices
    assignments = assign_voices(speakers)

    # Save assignments
    output_path = cfg.data_dir / "voice_assignments.json"
    with open(output_path, "w") as f:
        json.dump(assignments, f, indent=2)

    print(f"\nVoice assignments saved to {output_path}")

    # Print summary
    print("\n=== Voice Assignments ===")
    for name, voice in sorted(assignments.items()):
        print(f"  {name:35} -> {voice}")

    # Count unique base voices
    base_voices = set()
    for voice in assignments.values():
        if ":" in voice:
            base_voices.add(voice.split(":")[0])
        else:
            base_voices.add(voice)

    print("\n=== Summary ===")
    print(f"Total speakers: {len(speakers)}")
    print(f"Unique base voices: {len(base_voices)}")
    print(f"Speakers with pitch shift: {sum(1 for v in assignments.values() if ':' in v)}")


if __name__ == "__main__":
    main()