UmbraAudio/scripts/generate_all_audio.py

#!/usr/bin/env python3
"""Background chapter audio generation with Qwen3-TTS.
Run as: nohup python scripts/generate_all_audio.py [--project NAME] > generation.log 2>&1 &
"""

from __future__ import annotations

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

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

from umbraaudio.config import load_config
from umbraaudio.generation.audio_effects import (
    apply_effects_to_file,
    apply_effects_to_file_by_name,
    load_speaker_distortion_effects,
    load_speaker_traits,
)
from umbraaudio.generation.staleness import is_current, write_marker
from umbraaudio.pronunciation import apply_phonetics, load_phonetic_names
from umbraaudio.voices.lmstudio_lifecycle import get_vram_usage_mb
from umbraaudio.voices.qwen_tts import VOICE_REFERENCE_TEXT, QwenTTSEngine

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s %(levelname)s %(message)s",
    handlers=[
        logging.StreamHandler(),
    ],
)
logger = logging.getLogger(__name__)


def load_voice_references(voices_dir: Path) -> dict[str, Path]:
    refs: dict[str, Path] = {}
    if not voices_dir.exists():
        return refs
    for f in voices_dir.rglob("*_ref.wav"):
        name = f.stem.replace("_ref", "").replace("_", " ")
        if name not in refs:
            refs[name] = f
    return refs


def load_speakers(data_dir: Path) -> dict[str, dict]:
    path = data_dir / "speakers.json"
    if not path.exists():
        return {}
    with open(path) as f:
        data = json.load(f)
    return {s["primary_name"]: s for s in data.get("speakers", [])}


def get_speaker_name_map(speakers: dict[str, dict]) -> dict[str, str]:
    alias_map: dict[str, str] = {}
    for name, info in speakers.items():
        alias_map[name.lower()] = name
        for alias in info.get("aliases", []):
            if alias:
                alias_map[alias.lower()] = name
    return alias_map


def resolve_reference(
    speaker_name: str,
    voice_refs: dict[str, Path],
    alias_map: dict[str, str],
    narrator_ref: Path | None = None,
) -> Path | None:
    if speaker_name in voice_refs:
        return voice_refs[speaker_name]
    low = speaker_name.lower()
    for name, path in voice_refs.items():
        if name.lower() == low:
            return path
    canonical = alias_map.get(low)
    if canonical:
        for name, path in voice_refs.items():
            if name.lower() == canonical.lower():
                return path
    first_word = low.split()[0] if low.split() else ""
    if first_word:
        matches = [p for n, p in voice_refs.items() if n.lower().split()[:1] == [first_word]]
        if len(matches) == 1:
            return matches[0]
    best = None
    best_len = -1
    for name, path in voice_refs.items():
        nlow = name.lower()
        if (nlow in low or low in nlow) and len(nlow) > best_len:
            best = path
            best_len = len(nlow)
    if best is not None:
        return best
    return narrator_ref


def generate_chapter(
    chapter_id: int,
    engine: QwenTTSEngine,
    voice_refs: dict[str, Path],
    alias_map: dict[str, str],
    output_base: Path,
    db_path: Path,
    narrator_ref: Path | None = None,
    phonetic: dict[str, str] | None = None,
    speaker_traits: dict[str, list[str]] | None = None,
    distortion_effects: dict[str, list[str]] | None = None,
    utterance_callback=None,
    progress_callback=None,
    progress_context=None,
) -> tuple[int, int]:
    db = sqlite3.connect(str(db_path))
    db.row_factory = sqlite3.Row

    chapter = db.execute("SELECT * FROM chapters WHERE id = ?", (chapter_id,)).fetchone()
    if not chapter:
        return 0, 0

    title = chapter["title"] or f"Chapter {chapter['chapter_index'] + 1}"

    utterances = db.execute(
        "SELECT id, speaker_name, text, utterance_index FROM utterances "
        "WHERE chapter_id = ? ORDER BY utterance_index",
        (chapter_id,),
    ).fetchall()
    db.close()

    chapter_dir = output_base / f"ch{chapter_id}"
    chapter_dir.mkdir(parents=True, exist_ok=True)

    total = len(utterances)
    skipped = 0
    generated = 0
    failed = 0

    for utt in utterances:
        speaker = utt["speaker_name"] or "narrator"
        text = utt["text"]
        idx = utt["utterance_index"]

        if not text or not text.strip():
            skipped += 1
            if utterance_callback:
                utterance_callback()
            continue

        if text.strip() == "<<<>>>":
            skipped += 1
            if utterance_callback:
                utterance_callback()
            continue

        output_path = chapter_dir / f"utt_{idx:04d}.wav"

        ref_path = resolve_reference(speaker, voice_refs, alias_map, narrator_ref)

        spoken_text = apply_phonetics(text, phonetic or {})

        if is_current(spoken_text, output_path, speaker=speaker, voice_ref=ref_path):
            skipped += 1
            if utterance_callback:
                utterance_callback()
            continue

        utt_success = False
        try:
            if ref_path:
                engine.generate(
                    text=spoken_text,
                    output_path=output_path,
                    voice_ref=ref_path,
                    ref_text=VOICE_REFERENCE_TEXT,
                )
            else:
                engine.generate(
                    text=spoken_text,
                    output_path=output_path,
                    instruct="A calm, neutral voice speaking clearly.",
                )

            if output_path.exists():
                explicit = (distortion_effects or {}).get(speaker, [])
                if explicit:
                    apply_effects_to_file_by_name(output_path, explicit)
                elif speaker_traits:
                    traits = speaker_traits.get(speaker, [])
                    if traits:
                        apply_effects_to_file(output_path, traits)
                write_marker(spoken_text, output_path, speaker=speaker, voice_ref=ref_path)
                generated += 1
                utt_success = True
            else:
                failed += 1

        except Exception as e:
            logger.error("  FAILED utt_%04d (%s): %s", idx, speaker, e)
            failed += 1

        if utterance_callback:
            utterance_callback(success=utt_success)

        if progress_callback and progress_context:
            progress_callback(progress_context["completed"], progress_context["total"], title)

        if (generated + failed) % 50 == 0 and generated > 0:
            elapsed_min = (generated + failed) * 6.4 / 60
            logger.info(
                "  Progress: %d/%d done (%d gen, %d fail), ETA %.0f min",
                generated + failed,
                total,
                generated,
                failed,
                elapsed_min,
            )

    return generated, failed


def main(status_callback=None, gen_manager=None, project: str | None = None):
    cfg = load_config(project)
    data_dir = cfg.data_dir

    file_handler = logging.FileHandler(data_dir / "generation.log")
    file_handler.setFormatter(logging.Formatter("%(asctime)s %(levelname)s %(message)s"))
    logger.addHandler(file_handler)

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

    voice_refs = load_voice_references(voices_dir)
    speakers = load_speakers(data_dir)
    alias_map = get_speaker_name_map(speakers)
    phonetic = load_phonetic_names(data_dir)
    speaker_traits = load_speaker_traits(data_dir / "speakers.json")
    distortion_effects = load_speaker_distortion_effects(
        data_dir / "speakers.json"
    )

    narrator_ref = None
    for name, path in voice_refs.items():
        if name.lower() == "narrator":
            narrator_ref = path
            break
    if not narrator_ref:
        for name, path in voice_refs.items():
            if name.lower() == "kainda":
                narrator_ref = path
                break

    logger.info("=== Qwen3-TTS Background Generation ===")
    logger.info("Voice references: %d", len(voice_refs))
    logger.info("Speakers: %d", len(speakers))
    logger.info("Confirmed phonetic overrides: %d", len(phonetic))

    from umbraaudio.voices.model_coordinator import ModelCoordinator, PipelineStage

    coordinator = ModelCoordinator(cfg)
    if not coordinator.switch_to(PipelineStage.UTTERANCE_GENERATION).get("ok"):
        logger.error("Failed to switch models to UTTERANCE_GENERATION stage")
        sys.exit(1)

    vram = get_vram_usage_mb()
    logger.info("VRAM: %dMB used, %dMB free", vram["used_mb"], vram["free_mb"])

    db = sqlite3.connect(str(cfg.db_path))
    db.row_factory = sqlite3.Row
    chapters = db.execute(
        "SELECT id, title, chapter_index FROM chapters WHERE project_id=1 ORDER BY chapter_index"
    ).fetchall()
    total_utterances = db.execute("SELECT COUNT(*) as cnt FROM utterances").fetchone()["cnt"]
    db.close()

    total_chapters = len(chapters)

    engine = QwenTTSEngine(device="cuda", sample_rate=24000, qwen_tts_url=cfg.tts.qwen_tts_url)

    engine.load_remote_model(model_type="base")

    total_gen = 0
    total_fail = 0
    completed_utterances = 0
    start = time.time()

    if status_callback:
        status_callback(
            "generate",
            "running",
            "Starting audio generation...",
            0,
            total_chapters=total_chapters,
            chapters_completed=[],
        )

    def _utterance_done(success=True):
        nonlocal completed_utterances
        completed_utterances += 1
        if gen_manager is not None:
            gen_manager.update_item(f"utt:{completed_utterances}", success=success)

    def _progress_tick(completed, total, ch_title):
        if status_callback:
            progress_pct = int((completed_utterances / max(total_utterances, 1)) * 100)
            status_callback(
                "generate",
                "running",
                f"Chapter {i + 1}/{total_chapters}: {ch_title}",
                progress_pct,
                total_chapters=total_chapters,
                chapters_completed=list(range(i)),
                current_chapter=chapters[i]["id"],
                completed_utterances=completed_utterances,
                total_utterances=total_utterances,
            )

    batch_total = total_utterances
    _gen_ctx = None
    if gen_manager is not None:
        _gen_ctx = gen_manager.serialize("audio", batch_total, "audio:generate")
        _gen_ctx.__enter__()

    try:
        for i, ch in enumerate(chapters):
            ch_id = ch["id"]
            ch_title = ch["title"] or f"Ch{ch['chapter_index'] + 1}"

            if status_callback:
                progress_pct = int((completed_utterances / max(total_utterances, 1)) * 100)
                status_callback(
                    "generate",
                    "running",
                    f"Chapter {i + 1}/{total_chapters}: {ch_title}",
                    progress_pct,
                    total_chapters=total_chapters,
                    chapters_completed=list(range(i)),
                    current_chapter=ch_id,
                    completed_utterances=completed_utterances,
                    total_utterances=total_utterances,
                )

            progress_ctx = {"completed": completed_utterances, "total": total_utterances}
            gen, fail = generate_chapter(
                ch_id,
                engine,
                voice_refs,
                alias_map,
                output_base,
                cfg.db_path,
                narrator_ref,
                phonetic=phonetic,
                speaker_traits=speaker_traits,
                distortion_effects=distortion_effects,
                utterance_callback=_utterance_done,
                progress_callback=_progress_tick,
                progress_context=progress_ctx,
            )
            completed_utterances = progress_ctx["completed"]
            total_gen += gen
            total_fail += fail

            ch_dir = output_base / f"ch{ch_id}"
            full_wav = output_base / f"ch{ch_id}_full.wav"
            utt_files = sorted(ch_dir.glob("utt_*.wav")) if ch_dir.exists() else []
            needs_stitch = True
            if full_wav.exists() and utt_files:
                newest_utt = max(f.stat().st_mtime for f in utt_files)
                needs_stitch = full_wav.stat().st_mtime < newest_utt
            if utt_files and needs_stitch:
                try:
                    from umbraaudio.assembly import AudioStitcher

                    stitcher = AudioStitcher()
                    stitcher.stitch_utterances(utt_files, full_wav)
                    n_files = len(utt_files)
                    logger.info("  Stitched ch%d -> %s (%d files)", ch_id, full_wav.name, n_files)
                except Exception as e:
                    logger.error("  Stitch failed for ch%d: %s", ch_id, e)

            elapsed = (time.time() - start) / 60
            logger.info(
                "Ch %d (%s): +%d gen, +%d fail | Total: %d gen, %d fail | %.1f min elapsed",
                ch_id,
                ch_title,
                gen,
                fail,
                total_gen,
                total_fail,
                elapsed,
            )
    finally:
        if _gen_ctx is not None:
            _gen_ctx.__exit__(None, None, None)

    engine.unload_all()
    coordinator.switch_to(PipelineStage.IDLE)
    elapsed = (time.time() - start) / 60
    logger.info(
        "=== COMPLETE: %d generated, %d failed, %.1f min total ===",
        total_gen,
        total_fail,
        elapsed,
    )

    if status_callback:
        status_callback(
            "generate",
            "completed",
            f"Done: {total_gen} generated, {total_fail} failed in {elapsed:.1f} min",
            100,
            total_chapters=total_chapters,
            chapters_completed=list(range(total_chapters)),
        )


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="Background chapter audio generation")
    parser.add_argument(
        "--project", type=str, default=None, help="Project name or path"
    )
    args = parser.parse_args()
    main(project=args.project)