scripts/generate_kokoro_chapter_audio.py

"""Generate locally hosted, teacher-paced Kokoro narration for SSLC chapters.

This intentionally writes a parallel ``kokoro-full-chapter`` manifest. Existing
Deepgram audio remains untouched so every generation run is reversible.
"""

from __future__ import annotations

import argparse
import json
import logging
import os
import re
import subprocess
import sys
import warnings
from collections import OrderedDict
from pathlib import Path

# Silence torch/kokoro warnings so nothing reaches stderr. Windows PowerShell 5.1
# with $ErrorActionPreference='Stop' promotes ANY native-command stderr line to a
# terminating NativeCommandError, which made the render-*-teacher-v3-all.ps1
# runners abort at Kokoro model load (~30s in) and record the torch UserWarning as
# a bogus "failure". The warnings are harmless; suppressing them keeps stderr clean.
warnings.filterwarnings("ignore")
os.environ.setdefault("PYTHONWARNINGS", "ignore")
logging.disable(logging.WARNING)

import soundfile as sf


ROOT = Path(__file__).resolve().parents[2]
DEFAULT_IDS = [f"phy-p1-c{index}" for index in range(1, 8)]
MIN_WORDS = 18
MAX_WORDS = 82


def words(text: str) -> int:
    return len(re.findall(r"\S+", text))


def spoken_text(text: str) -> str:
    """Turn compact board notation into narration that sounds like a teacher."""
    substitutions = (
        (r"\bv\s*=\s*f\s*(?:lambda|λ)\b", "v equals f multiplied by lambda"),
        (r"\bflambda\b", "f multiplied by lambda"),
        (r"\bλ\b", "lambda"),
        (r"\bμ\b", "mu"),
        (r"\bΔ\b", "delta"),
        (r"\b([0-9]+)\s*Hz\b", r"\1 hertz"),
        (r"\b([0-9]+)\s*kHz\b", r"\1 kilohertz"),
        (r"\bm/s\b", "metres per second"),
        (r"\bcm\b", "centimetres"),
        (r"\bmm\b", "millimetres"),
    )
    result = text
    for pattern, replacement in substitutions:
        result = re.sub(pattern, replacement, result, flags=re.IGNORECASE)
    result = re.sub(r"\s+", " ", result).strip()
    return result


def sentence_parts(text: str) -> list[str]:
    parts = re.split(r"(?<=[.!?])\s+", text.strip())
    return [part.strip() for part in parts if part.strip()]


def split_long_text(text: str, max_words: int = MAX_WORDS) -> list[str]:
    sentences = sentence_parts(text)
    output: list[str] = []
    current: list[str] = []
    current_words = 0
    for sentence in sentences:
        sentence_words = words(sentence)
        if current and current_words + sentence_words > max_words:
            output.append(" ".join(current))
            current = []
            current_words = 0
        if sentence_words > max_words:
            tokens = sentence.split()
            for start in range(0, len(tokens), max_words):
                if current:
                    output.append(" ".join(current))
                    current = []
                    current_words = 0
                output.append(" ".join(tokens[start : start + max_words]))
        else:
            current.append(sentence)
            current_words += sentence_words
    if current:
        output.append(" ".join(current))
    return output


def merge_scene_chunks(chunks: list[dict]) -> list[dict]:
    """Merge tiny TTS calls while preserving scene boundaries and caption text."""
    scenes: OrderedDict[str, list[dict]] = OrderedDict()
    for chunk in chunks:
        scenes.setdefault(chunk["sceneId"], []).append(chunk)

    merged: list[dict] = []
    for scene_id, scene_chunks in scenes.items():
        scene_text = " ".join(chunk["text"].strip() for chunk in scene_chunks if chunk.get("text", "").strip())
        pieces = split_long_text(scene_text)

        # Avoid a clipped-sounding final phrase by folding it into the prior piece.
        if len(pieces) > 1 and words(pieces[-1]) < MIN_WORDS and words(pieces[-2]) + words(pieces[-1]) <= MAX_WORDS + 12:
            pieces[-2] = f"{pieces[-2]} {pieces[-1]}"
            pieces.pop()

        for index, text in enumerate(pieces, start=1):
            source = scene_chunks[min(index - 1, len(scene_chunks) - 1)]
            merged.append(
                {
                    "id": f"{scene_id}-voice-{index}",
                    "sceneId": scene_id,
                    "segmentKind": source.get("segmentKind", "concept"),
                    "text": text,
                    "spokenText": spoken_text(text),
                    "characterCount": len(text),
                    "pauseAfterSeconds": 0.45 if text.rstrip().endswith("?") else 0.22,
                }
            )
    return merged


def preserve_authored_chunks(chunks: list[dict]) -> list[dict]:
    """Keep Biology V4 cue IDs intact so measured audio drives exact visual beats."""
    prepared: list[dict] = []
    for chunk in chunks:
        text = chunk.get("text", "").strip()
        if not text:
            continue
        prepared.append(
            {
                **chunk,
                "spokenText": chunk.get("spokenText") or spoken_text(text),
                "characterCount": chunk.get("characterCount") or len(text),
                "pauseAfterSeconds": float(chunk.get("pauseAfterSeconds") or 0.22),
            }
        )
    return prepared


def probe_duration(path: Path) -> float:
    completed = subprocess.run(
        ["ffprobe", "-v", "error", "-show_entries", "format=duration", "-of", "default=nw=1:nk=1", str(path)],
        capture_output=True,
        check=True,
        text=True,
    )
    return round(float(completed.stdout.strip()), 6)


def synthesize_chapter(
    pipeline,
    chapter_id: str,
    voice: str,
    speed: float,
    force: bool,
    source_audio_name: str = "full-chapter",
    output_audio_name: str = "kokoro-full-chapter",
    preserve_source_chunks: bool = False,
) -> dict:
    source_dir = ROOT / "outputs" / "video" / "tuition" / chapter_id / "audio" / source_audio_name
    source_manifest_path = source_dir / "manifest.json"
    if not source_manifest_path.exists():
        raise FileNotFoundError(f"Missing source narration manifest: {source_manifest_path}")

    source_manifest = json.loads(source_manifest_path.read_text(encoding="utf-8"))
    chunks = preserve_authored_chunks(source_manifest["chunks"]) if preserve_source_chunks else merge_scene_chunks(source_manifest["chunks"])
    output_dir = source_dir.parent / output_audio_name
    chunk_dir = output_dir / "chunks"
    chunk_dir.mkdir(parents=True, exist_ok=True)
    manifest_path = output_dir / "manifest.json"
    existing: dict[str, dict] = {}
    if manifest_path.exists() and not force:
        existing_manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
        settings_match = (
            existing_manifest.get("model") == "hexgrad/Kokoro-82M"
            and existing_manifest.get("voice") == voice
            and float(existing_manifest.get("speed", -1)) == speed
        )
        if settings_match:
            existing = {item["id"]: item for item in existing_manifest["chunks"]}
        else:
            print(
                f"[{chapter_id}] voice settings changed; regenerating all phrases "
                f"with {voice} at {speed}x",
                flush=True,
            )

    generated: list[dict] = []
    for index, chunk in enumerate(chunks, start=1):
        output_path = chunk_dir / f"{chunk['id']}.wav"
        cached = existing.get(chunk["id"])
        stale = cached is not None and cached.get("spokenText") != chunk["spokenText"]
        if force or not output_path.exists() or stale:
            audio_parts = []
            for _graphemes, _phonemes, audio in pipeline(chunk["spokenText"], voice=voice, speed=speed, split_pattern=r"\n+"):
                audio_parts.append(audio)
            if not audio_parts:
                raise RuntimeError(f"Kokoro produced no audio for {chunk['id']}")
            import numpy as np

            joined = np.concatenate(audio_parts)
            sf.write(output_path, joined, 24_000, subtype="PCM_16")

        if cached and not stale and output_path.exists() and not force:
            # Timing metadata belongs to the current authored cue contract even
            # when the waveform itself is safely reused from cache.
            generated.append({**cached, "pauseAfterSeconds": chunk["pauseAfterSeconds"]})
        else:
            relative_path = output_path.relative_to(ROOT).as_posix()
            generated.append(
                {
                    **chunk,
                    "outputPath": relative_path,
                    "status": "synthesized",
                    "durationSeconds": probe_duration(output_path),
                }
            )
        if index == 1 or index % 25 == 0 or index == len(chunks):
            print(f"[{chapter_id}] {index}/{len(chunks)} teacher phrases", flush=True)

    manifest = {
        "provider": "kokoro-local",
        "model": "hexgrad/Kokoro-82M",
        "voice": voice,
        "speed": speed,
        "sampleRate": 24_000,
        "configured": True,
        "dryRun": False,
        "sourceManifest": source_manifest_path.relative_to(ROOT).as_posix(),
        "sourceChunkCount": len(source_manifest["chunks"]),
        "naturalPhraseCount": len(generated),
        "preserveSourceChunks": preserve_source_chunks,
        "totalCharacters": sum(chunk["characterCount"] for chunk in generated),
        "totalDurationSeconds": round(sum(chunk["durationSeconds"] for chunk in generated), 3),
        "chunks": generated,
    }
    output_dir.mkdir(parents=True, exist_ok=True)
    manifest_path.write_text(json.dumps(manifest, indent=2, ensure_ascii=False) + "\n", encoding="utf-8")
    print(f"[{chapter_id}] wrote {manifest_path}", flush=True)
    return manifest


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--chapter-id", action="append", dest="chapter_ids", help="Repeat to select chapters; defaults to all seven Physics chapters.")
    # Heart is the clearest natural Kokoro voice in the current student proof.
    # Scene selection uses measured audio duration, so the calmer pace does not
    # push the final production package outside its 30-40 minute target.
    parser.add_argument("--voice", default="af_heart")
    parser.add_argument("--speed", type=float, default=1.0)
    parser.add_argument("--force", action="store_true")
    parser.add_argument("--source-audio-name", default="full-chapter")
    parser.add_argument("--output-audio-name", default="kokoro-full-chapter")
    parser.add_argument("--preserve-source-chunks", action="store_true")
    args = parser.parse_args()

    os.environ.setdefault("USE_TF", "0")
    os.environ.setdefault("TRANSFORMERS_NO_TF", "1")
    os.environ.setdefault("HF_HUB_DISABLE_XET", "1")
    from kokoro import KPipeline

    pipeline = KPipeline(lang_code="a", repo_id="hexgrad/Kokoro-82M")
    summaries = []
    for chapter_id in args.chapter_ids or DEFAULT_IDS:
        manifest = synthesize_chapter(
            pipeline,
            chapter_id,
            args.voice,
            args.speed,
            args.force,
            args.source_audio_name,
            args.output_audio_name,
            args.preserve_source_chunks,
        )
        summaries.append(
            {
                "chapterId": chapter_id,
                "phrases": manifest["naturalPhraseCount"],
                "durationMinutes": round(manifest["totalDurationSeconds"] / 60, 2),
            }
        )
    print(json.dumps(summaries, indent=2))
    return 0


if __name__ == "__main__":
    sys.exit(main())