scripts/generate_docdoe_kokoro_audio.py

"""Synthesize Kokoro af_heart audio for DocDoe teaching-system chapters.

Does not touch tuition/ legacy audio trees. Writes measured durations into
docdoe-teaching-system audio manifests and a single concatenated WAV.
"""

from __future__ import annotations

import argparse
import json
import logging
import os
import subprocess
import warnings
from pathlib import Path

warnings.filterwarnings("ignore")
os.environ.setdefault("PYTHONWARNINGS", "ignore")
logging.disable(logging.WARNING)

import soundfile as sf

ROOT = Path(__file__).resolve().parents[2]


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 loudnorm_wav(src: Path, dst: Path) -> None:
    subprocess.run(
        [
            "ffmpeg",
            "-y",
            "-i",
            str(src),
            "-af",
            "loudnorm=I=-16:TP=-1.5:LRA=11",
            "-ar",
            "24000",
            str(dst),
        ],
        check=True,
        capture_output=True,
    )


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--source-manifest",
        default=str(
            ROOT
            / "outputs/video/docdoe-teaching-system/audio/narration-manifests/bio-p1-c1-release/source-manifest.json"
        ),
    )
    parser.add_argument(
        "--output-dir",
        default=str(
            ROOT
            / "outputs/video/docdoe-teaching-system/audio/narration-manifests/bio-p1-c1-release"
        ),
    )
    parser.add_argument("--voice", default="af_heart")
    parser.add_argument("--speed", type=float, default=0.95)
    parser.add_argument("--force", 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
    import numpy as np

    source = json.loads(Path(args.source_manifest).read_text(encoding="utf-8"))
    chunks = source["chunks"]
    out_dir = Path(args.output_dir)
    chunk_dir = out_dir / "chunks"
    chunk_dir.mkdir(parents=True, exist_ok=True)

    pipeline = KPipeline(lang_code="a", repo_id="hexgrad/Kokoro-82M")
    generated = []
    concat_parts = []

    for index, chunk in enumerate(chunks, start=1):
        cid = chunk["id"]
        text = (chunk.get("spokenText") or chunk.get("text") or "").strip()
        if not text:
            raise RuntimeError(f"Empty text for {cid}")
        wav_path = chunk_dir / f"{cid}.wav"
        if args.force or not wav_path.exists():
            audio_parts = []
            for _g, _p, audio in pipeline(text, voice=args.voice, speed=args.speed, split_pattern=r"\n+"):
                audio_parts.append(audio)
            if not audio_parts:
                raise RuntimeError(f"Silent/empty Kokoro output for {cid} — aborting (no silent fallback)")
            joined = np.concatenate(audio_parts)
            sf.write(wav_path, joined, 24_000, subtype="PCM_16")
        duration = probe_duration(wav_path)
        pause = float(chunk.get("pauseAfterSeconds") or 0.28)
        generated.append(
            {
                **chunk,
                "spokenText": text,
                "outputPath": wav_path.relative_to(ROOT).as_posix(),
                "status": "synthesized",
                "durationSeconds": duration,
                "pauseAfterSeconds": pause,
            }
        )
        audio_data, sr = sf.read(str(wav_path))
        concat_parts.append(audio_data)
        if pause > 0:
            concat_parts.append(np.zeros(int(sr * pause), dtype=audio_data.dtype))
        if index == 1 or index % 10 == 0 or index == len(chunks):
            print(f"[docdoe-kokoro] {index}/{len(chunks)} {cid} ({duration:.2f}s)", flush=True)

    # Timeline with measured start times
    t = 0.0
    for item in generated:
        item["startSeconds"] = round(t, 3)
        t += item["durationSeconds"] + item["pauseAfterSeconds"]
        item["endSeconds"] = round(t, 3)

    full_raw = out_dir / "full-narration-raw.wav"
    full_norm = out_dir / "full-narration.wav"
    full_audio = np.concatenate(concat_parts)
    sf.write(full_raw, full_audio, 24_000, subtype="PCM_16")
    try:
        loudnorm_wav(full_raw, full_norm)
        final_audio_path = full_norm
    except Exception as exc:  # noqa: BLE001
        print(f"[docdoe-kokoro] loudnorm failed, using raw: {exc}", flush=True)
        final_audio_path = full_raw

    total = probe_duration(final_audio_path)
    manifest = {
        "provider": "kokoro-local",
        "model": "hexgrad/Kokoro-82M",
        "voice": args.voice,
        "speed": args.speed,
        "sampleRate": 24000,
        "chapterId": source.get("chapterId", "bio-p1-c1"),
        "sourceManifest": Path(args.source_manifest).relative_to(ROOT).as_posix(),
        "fullNarrationPath": final_audio_path.relative_to(ROOT).as_posix(),
        "chunkCount": len(generated),
        "totalDurationSeconds": total,
        "failedChunks": [c["id"] for c in generated if c.get("status") != "synthesized"],
        "chunks": generated,
    }
    if manifest["failedChunks"]:
        raise RuntimeError(f"Failed chunks: {manifest['failedChunks']}")

    manifest_path = out_dir / "kokoro-manifest.json"
    manifest_path.write_text(json.dumps(manifest, indent=2, ensure_ascii=False) + "\n", encoding="utf-8")
    print(f"[docdoe-kokoro] wrote {manifest_path} total={total:.2f}s", flush=True)

    # WebVTT from phrase timings
    vtt_lines = ["WEBVTT", ""]
    for item in generated:
        def ts(sec: float) -> str:
            h = int(sec // 3600)
            m = int((sec % 3600) // 60)
            s = sec % 60
            return f"{h:02d}:{m:02d}:{s:06.3f}"

        start = item["startSeconds"]
        end = item["startSeconds"] + item["durationSeconds"]
        # phrase-based, max ~2 lines by splitting long text lightly
        text = item["spokenText"]
        vtt_lines.append(f"{ts(start)} --> {ts(end)}")
        if len(text) > 90:
            mid = text.rfind(" ", 0, len(text) // 2 + 10)
            if mid < 20:
                mid = len(text) // 2
            vtt_lines.append(text[:mid].strip())
            vtt_lines.append(text[mid:].strip())
        else:
            vtt_lines.append(text)
        vtt_lines.append("")

    vtt_path = (
        ROOT
        / "outputs/video/docdoe-teaching-system/final/subtitles/bio-p1-c1-genetics-of-life.en.vtt"
    )
    vtt_path.parent.mkdir(parents=True, exist_ok=True)
    vtt_path.write_text("\n".join(vtt_lines), encoding="utf-8")
    print(f"[docdoe-kokoro] wrote {vtt_path}", flush=True)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())