Upload comfy/custom_nodes/ComfyUI-IndexTTS2/render_klaus_episode.py with huggingface_hub

comfy/custom_nodes/ComfyUI-IndexTTS2/render_klaus_episode.py

@@ -0,0 +1,505 @@
+#!/usr/bin/env python3
+"""render_klaus_episode.py — render a Klaus episode .md into one final audio FLAC.
+
+Runs on the server (must have IndexTTS-2 + ffmpeg installed). Pipeline:
+  1. Parse .md: drop H1 heading, strip **emphasis** and non-pause （…） stage cues,
+     split on （停顿） into fragments.
+  2. For each fragment: apply klaus_lexicon (v04 tone-only, locked-in profile);
+     IndexTTS-2 infer with the slowed Hugh ref → per-fragment .wav.
+  3. ffmpeg concat all fragments with 0.6s silence padding between every two.
+  4. ffmpeg atempo=0.95 post-process → final ~0.95× effective speed on the
+     clean Hugh ref (current production). Stacking a pre-slowed atempo=0.85
+     ref on top → ~0.81× (opt-in; ref no longer on server — see CLAUDE.md).
+
+Usage (on server):
+    cd /root/IndexTTS2
+    .venv/bin/python render_klaus_episode.py /path/to/script.md /tmp/klaus_ep01_full.flac
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import queue
+import re
+import subprocess
+import sys
+import threading
+import time
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).parent))
+import klaus_lexicon  # noqa: E402
+
+from indextts.infer_v2 import IndexTTS2  # noqa: E402
+
+# Locked-in Klaus voice config (tuned 2026-05-03 — see ts-tts-cached-flurry.md plan)
+LEXICON_PROFILE = dict(
+    tone_4to1_prob=0.95,
+    tone_3to2_prob=0.95,
+    retroflex_loss_prob=0.0,
+    u_prob=0.0,
+)
+SEED = 42
+PAUSE_SEC = 0.6           # silence between （停顿）-delimited fragments
+
+# Dropout guard (2026-05-11 architecture: persistent verifier daemon + per-frag
+# full-scan audit + 1× respace retry).
+# IndexTTS-2's GPT2 occasionally drops a syllable on tightly-packed phrases —
+# ep01 frag 5 lost `sai` in 下场比赛见, classic `任重道远 → 任重远`. After all
+# fragments render, we spawn verify_klaus_dropout.py in --daemon mode (loads
+# whisperx ONCE, ~15s amortized over the whole sweep) and ASR-check every
+# fragment via pinyin-sequence diff (tone-stripped, digit-normalized). Any
+# fragment with a `delete` opcode on its pinyin sequence is a real dropout —
+# we respace-retry once (`任重道远 → 任 重 道 远` forces GPT2 attention
+# boundaries) and re-verify. Persistent failures ship anyway with stderr
+# warning — never blocks episode rendering. The Layer 1 duration check was
+# removed: structurally blind to 1-syllable drops in moderate-length fragments
+# (49-char frag 5 had < 2% duration delta — drowned in baseline noise).
+DEFAULT_VERIFIER_PYTHON = "/root/miniconda3/envs/whisperx/bin/python"
+VERIFIER_SCRIPT = str(Path(__file__).parent / "verify_klaus_dropout.py")
+# Default --speed is the post-process atempo on the master FLAC. 0.95 alone
+# gives ~0.95× effective speed on the clean Hugh ref (current production).
+# Stacking a pre-slowed atempo=0.85 ref on top → ~0.81× (the original
+# locked-in cadence trick; opt-in, slow ref no longer kept on server).
+# Web UI overrides via per-character sidecar (Voices/<char>/.character.json
+# `speed`, default 1.0). Range 0.6–1.5.
+DEFAULT_SPEED = 0.95
+
+
+def md_to_fragments(md: str) -> list[str]:
+    body = "\n".join(ln for ln in md.splitlines() if not ln.lstrip().startswith("#"))
+    # 任何含「停顿」的括注都当作完整停顿：（停顿）/（停顿，XXX）/（XXX，停顿）/（A，停顿，B）。
+    # 多余表演词随括号被下面的 re.sub(r"（[^）]*）", ...) 丢弃。
+    raw = re.split(r"\n*\s*（[^）]*停顿[^）]*）\s*\n*", body)
+    out: list[str] = []
+    for s in raw:
+        s = re.sub(r"\*\*(.+?)\*\*", r"\1", s)
+        s = re.sub(r"（[^）]*）", "", s)
+        s = re.sub(r"\n{3,}", "\n\n", s).strip()
+        if s:
+            out.append(s)
+    return out
+
+
+def _ffprobe_duration(path: Path) -> float:
+    out = subprocess.check_output(
+        ["ffprobe", "-v", "error", "-show_entries", "format=duration",
+         "-of", "csv=p=0", str(path)]
+    )
+    return float(out.decode().strip())
+
+
+def _respace_han(text: str) -> str:
+    """Insert ASCII space between consecutive Han chars to force GPT2 attention
+    boundaries (the documented 任重道远 → 任 重 道 远 hard-fix). klaus_lexicon
+    is char-by-char and passes spaces through verbatim, so the lexicon profile
+    still applies on top."""
+    out: list[str] = []
+    prev_han = False
+    for ch in text:
+        is_han = "一" <= ch <= "鿿"
+        if is_han and prev_han:
+            out.append(" ")
+        out.append(ch)
+        prev_han = is_han
+    return "".join(out)
+
+
+def _render_one(tts: IndexTTS2, text: str, ref_wav: str, wav_path: Path) -> tuple[str, dict, float]:
+    """Single TTS infer; returns (mangled_text, lexicon_stats, elapsed_sec)."""
+    mangled, stats = klaus_lexicon.transform(text, seed=SEED, **LEXICON_PROFILE)
+    t0 = time.perf_counter()
+    tts.infer(spk_audio_prompt=ref_wav, text=mangled, output_path=str(wav_path), verbose=False)
+    return mangled, stats, time.perf_counter() - t0
+
+
+class _VerifierDaemon:
+    """Long-running [verify_klaus_dropout.py](IndexTTS2/verify_klaus_dropout.py)
+    subprocess in --daemon mode. Spawn once per episode → whisperx loads once
+    (~15 s) instead of cold-spawning per call (cold-spawn racing IndexTTS-2
+    on the GPU previously hit 180 s timeouts in production). Protocol is
+    line-delimited JSON over stdin/stdout.
+
+    Usage:
+        with _VerifierDaemon(py_path, script_path) as v:
+            verdict = v.verify(wav_path, expected_text)  # returns dict
+    """
+
+    _HANDSHAKE_TIMEOUT = 60.0   # whisperx model load ceiling
+    _CALL_TIMEOUT = 60.0        # per-frag ASR inference ceiling
+
+    def __init__(self, python_path: str, script_path: str):
+        self.proc = subprocess.Popen(
+            [python_path, script_path, "--daemon"],
+            stdin=subprocess.PIPE,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.DEVNULL,  # ML libs are chatty on stderr; we don't surface it
+            text=True, bufsize=1,
+            encoding="utf-8",
+        )
+        line = self._readline(self._HANDSHAKE_TIMEOUT)
+        if not line:
+            raise RuntimeError("verifier daemon exited during startup (no handshake)")
+        msg = json.loads(line)
+        if msg.get("status") != "ready":
+            raise RuntimeError(f"verifier daemon unexpected handshake: {msg}")
+
+    def _readline(self, timeout: float) -> str:
+        # subprocess.Popen.stdout.readline has no timeout; rely on the daemon
+        # being responsive (model load already done at handshake time, per-call
+        # is just inference). If it ever wedges, the user kills the parent.
+        #
+        # Defensive: skip lines that aren't JSON objects. whisperx/pyannote
+        # historically leaked INFO log lines to stdout despite our basicConfig
+        # → we belt+suspender it with _stdout_to_stderr() on the daemon side,
+        # and this client-side skip handles any future regression.
+        while True:
+            line = self.proc.stdout.readline()
+            if not line:
+                return ""
+            stripped = line.strip()
+            if stripped.startswith("{"):
+                return line
+            # else: non-JSON line (log spam), ignore and read next
+
+    def verify(self, audio_path: Path, expected: str) -> dict:
+        req = {"audio": str(audio_path), "expected": expected}
+        self.proc.stdin.write(json.dumps(req, ensure_ascii=False) + "\n")
+        self.proc.stdin.flush()
+        line = self._readline(self._CALL_TIMEOUT)
+        if not line:
+            raise RuntimeError("verifier daemon closed stdout mid-stream")
+        return json.loads(line)
+
+    def shutdown(self) -> None:
+        try:
+            if self.proc.stdin and not self.proc.stdin.closed:
+                self.proc.stdin.write("QUIT\n")
+                self.proc.stdin.flush()
+                self.proc.stdin.close()
+            self.proc.wait(timeout=10)
+        except (subprocess.TimeoutExpired, Exception):
+            self.proc.kill()
+            try:
+                self.proc.wait(timeout=5)
+            except subprocess.TimeoutExpired:
+                pass
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc, tb):
+        self.shutdown()
+
+
+class _StreamingAuditor:
+    """Concurrent daemon + worker thread for streaming per-fragment audit.
+
+    Lifecycle (2026-05-11 architecture A+B):
+      auditor = _StreamingAuditor(py, script)
+        ↳ kicks off two threads: daemon loader (~15s whisperx) + queue worker
+        ↳ caller continues with IndexTTS-2 load + Phase 1 render in parallel
+
+      for i, frag in render_loop:
+          render_one(...)
+          auditor.submit(i, frag, wav_path)   # non-blocking; worker picks up
+
+      auditor.drain()      # signal end of submissions, wait for worker
+      auditor.results      # dict[idx -> verdict]; ready for Phase 2 retries
+
+      for idx, verdict in auditor.results.items():
+          if not verdict["ok"]:
+              retry...
+              verdict2 = auditor.verify_direct(wav, frag)  # main thread only, post-drain
+
+      auditor.shutdown()
+
+    Threading model: daemon's stdin/stdout protocol is single-threaded. The
+    worker thread owns the daemon during Phase 1 streaming. After `drain()`,
+    the worker exits and the main thread takes over for retries via
+    `verify_direct()`. No concurrent access ever — no lock needed.
+    """
+
+    def __init__(self, python_path: str, script_path: str):
+        self.queue: queue.Queue = queue.Queue()
+        self.results: dict[int, dict] = {}
+        self.daemon: _VerifierDaemon | None = None
+        self.daemon_error: Exception | None = None
+        self._daemon_ready = threading.Event()
+
+        # Thread 1: load whisperx (~15s). Parallel with IndexTTS-2 load
+        # happening in the main thread.
+        self._loader = threading.Thread(
+            target=self._load, args=(python_path, script_path), daemon=True,
+        )
+        self._loader.start()
+
+        # Thread 2: worker. Waits for daemon to be ready, then drains queue.
+        self._worker = threading.Thread(target=self._run, daemon=True)
+        self._worker.start()
+
+    def _load(self, python_path: str, script_path: str) -> None:
+        try:
+            self.daemon = _VerifierDaemon(python_path, script_path)
+        except Exception as e:
+            self.daemon_error = e
+        finally:
+            self._daemon_ready.set()
+
+    def _run(self) -> None:
+        self._daemon_ready.wait()
+        while True:
+            item = self.queue.get()
+            if item is None:
+                break
+            if self.daemon is None:
+                continue  # daemon failed to load — drop the request
+            idx, frag, wav = item
+            try:
+                self.results[idx] = self.daemon.verify(wav, frag)
+            except Exception as e:
+                print(f"  ⚠ audit frag {idx}: daemon error: {e}",
+                      file=sys.stderr, flush=True)
+
+    def submit(self, idx: int, frag: str, wav: Path) -> None:
+        self.queue.put((idx, frag, wav))
+
+    def drain(self) -> None:
+        """Signal end of submissions and wait for the worker to process them all."""
+        self.queue.put(None)
+        self._worker.join()
+
+    def verify_direct(self, wav: Path, frag: str) -> dict | None:
+        """Synchronous verify on the daemon (no queue). Only safe AFTER
+        `drain()` — used for Phase 2 retries when the worker thread is gone
+        and the main thread has exclusive access to the daemon."""
+        if self.daemon is None:
+            return None
+        return self.daemon.verify(wav, frag)
+
+    def shutdown(self) -> None:
+        if self.daemon is not None:
+            self.daemon.shutdown()
+
+
+def _process_audit_results(
+    auditor: _StreamingAuditor,
+    tts: IndexTTS2,
+    fragments: list[str],
+    frag_wavs: list[Path],
+    ref_wav: str,
+) -> list[tuple[int, list[str]]]:
+    """Iterate the streaming auditor's results; for each fragment flagged DROP,
+    respace-retry once and re-verify. Returns list of `(idx_1based, dropped_pinyin)`
+    for fragments still failing after retry."""
+    failures: list[tuple[int, list[str]]] = []
+
+    for idx in sorted(auditor.results.keys()):
+        verdict = auditor.results[idx]
+        if verdict.get("error"):
+            print(f"  ⚠ frag {idx}: daemon error: {verdict['error']}",
+                  file=sys.stderr, flush=True)
+            continue
+        if verdict.get("ok"):
+            continue
+
+        frag = fragments[idx - 1]
+        wav_path = frag_wavs[idx - 1]
+        dropped = verdict.get("pinyin_dropped", [])
+        py_exp = verdict.get("pinyin_expected_count")
+        py_got = verdict.get("pinyin_got_count")
+        print(f"  ⚠ frag {idx}: dropped pinyin [{' '.join(dropped)}] "
+              f"(exp={py_exp} got={py_got}) — respace retry", flush=True)
+
+        respaced = _respace_han(frag)
+        _, _, elapsed = _render_one(tts, respaced, ref_wav, wav_path)
+        actual_dur = _ffprobe_duration(wav_path)
+        print(f"    retry: {len(respaced):>4} chars → wav "
+              f"({elapsed:.1f}s render, {actual_dur:.2f}s audio)", flush=True)
+
+        verdict2 = auditor.verify_direct(wav_path, frag)
+        if verdict2 is None:
+            print(f"  ⚠ frag {idx}: re-verify daemon unavailable",
+                  file=sys.stderr, flush=True)
+            continue
+        if verdict2.get("ok"):
+            print(f"  ✓ frag {idx}: retry cleared the dropout", flush=True)
+            continue
+
+        dropped2 = verdict2.get("pinyin_dropped", [])
+        print(f"  ⚠⚠⚠ frag {idx}: still dropping after retry: [{' '.join(dropped2)}] — shipping anyway",
+              file=sys.stderr, flush=True)
+        failures.append((idx, dropped2))
+
+    return failures
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
+    ap.add_argument("script", type=Path, help=".md episode script")
+    ap.add_argument("output", type=Path, help="final .flac path")
+    ap.add_argument("--workdir", type=Path, default=Path("/tmp/klaus_render"))
+    ap.add_argument("--ep-id", default=None,
+                    help="ep_id used for the segments.json sidecar filename. "
+                         "Sidecar is written next to <output> as {ep_id}-segments.json. "
+                         "Skip writing sidecar when omitted.")
+    ap.add_argument("--ref-wav", required=True,
+                    help="voice clone reference path (relative to cwd or absolute). "
+                         "Web UI persists per-character refs at refs/web/<voice>.flac; "
+                         "the standalone Klaus pipeline historically used "
+                         "refs/hugh_184_slow_0.85.flac (Hugh Jackman atempo=0.85). "
+                         "No default — pass explicitly so the chosen voice is unambiguous.")
+    ap.add_argument("--speed", type=float, default=DEFAULT_SPEED,
+                    help=f"final ffmpeg atempo factor on the master FLAC + scales "
+                         f"sidecar timings. Default {DEFAULT_SPEED} (Klaus "
+                         f"production cadence trick); Web UI passes the per-character "
+                         f"sidecar value (default 1.0). Range 0.6–1.5.")
+    ap.add_argument("--no-dropout-guard", action="store_true",
+                    help="disable the Layer-1 duration check + Layer-2 ASR confirm + "
+                         "respace retry pass. Default ON; use this only for A/B or when "
+                         "the WhisperX env on the server is broken and you need to ship.")
+    ap.add_argument("--verifier-python", default=DEFAULT_VERIFIER_PYTHON,
+                    help=f"path to the whisperx conda env's python (it loads whisperx, "
+                         f"which the IndexTTS-2 venv doesn't have). "
+                         f"Default {DEFAULT_VERIFIER_PYTHON}.")
+    args = ap.parse_args()
+    if not (0.6 <= args.speed <= 1.5):
+        sys.exit(f"--speed must be in [0.6, 1.5], got {args.speed}")
+
+    ref_wav = args.ref_wav
+    if not Path(ref_wav).exists():
+        sys.exit(f"REF missing: {ref_wav} (cwd={Path.cwd()})")
+    if not args.script.exists():
+        sys.exit(f"script missing: {args.script}")
+    args.workdir.mkdir(parents=True, exist_ok=True)
+    for old in args.workdir.glob("frag*.wav"):
+        old.unlink()
+
+    fragments = md_to_fragments(args.script.read_text())
+    if not fragments:
+        sys.exit("no fragments parsed")
+    print(f">> {len(fragments)} fragments, lengths: {[len(f) for f in fragments]}", flush=True)
+
+    # Early-spawn the streaming auditor: daemon's whisperx load (~15s) runs in
+    # a background thread, parallel with IndexTTS-2 load (~17s) below. By the
+    # time Phase 1 ends, the audit worker has typically consumed the entire
+    # queue, leaving Phase 2 with just retries to do.
+    enable_guard = not args.no_dropout_guard
+    auditor: _StreamingAuditor | None = None
+    if enable_guard:
+        if not Path(args.verifier_python).exists() or not Path(VERIFIER_SCRIPT).exists():
+            print(f">> warning: verifier python ({args.verifier_python}) or script "
+                  f"({VERIFIER_SCRIPT}) missing — skipping streaming audit",
+                  file=sys.stderr, flush=True)
+        else:
+            print(">> spawning verifier daemon (whisperx loads in parallel with IndexTTS-2)…",
+                  flush=True)
+            auditor = _StreamingAuditor(args.verifier_python, VERIFIER_SCRIPT)
+
+    print(">> loading IndexTTS2…", flush=True)
+    t0 = time.perf_counter()
+    tts = IndexTTS2(cfg_path="checkpoints/config.yaml", model_dir="checkpoints", use_fp16=True)
+    print(f">> loaded in {time.perf_counter()-t0:.1f}s on {tts.device}", flush=True)
+
+    # Phase 1: render each fragment, submit to streaming audit queue.
+    frag_wavs: list[Path] = []
+    total_overrides = 0
+    total_han = 0
+    for i, frag in enumerate(fragments, 1):
+        wav_path = args.workdir / f"frag{i:03d}.wav"
+        _, stats, elapsed = _render_one(tts, frag, ref_wav, wav_path)
+        actual_dur = _ffprobe_duration(wav_path)
+        print(f"  frag {i:>3}/{len(fragments)}: {len(frag):>4} chars → wav "
+              f"({elapsed:.1f}s render, {actual_dur:.2f}s audio)", flush=True)
+        total_overrides += stats["overrides_emitted"]
+        total_han += stats["total_han"]
+        frag_wavs.append(wav_path)
+        if auditor is not None:
+            auditor.submit(i, frag, wav_path)
+
+    rate = total_overrides / max(total_han, 1)
+    print(f">> total overrides {total_overrides}/{total_han} ({rate:.0%})", flush=True)
+
+    # Phase 2: drain audit worker (most ASR already done in parallel with
+    # Phase 1) → respace-retry any DROP verdicts.
+    dropout_failures: list[tuple[int, list[str]]] = []
+    if auditor is not None:
+        print(">> draining audit worker (ASR ran in parallel with Phase 1)…", flush=True)
+        auditor.drain()
+        if auditor.daemon_error:
+            print(f">> warning: verifier daemon failed to load: {auditor.daemon_error} — "
+                  f"no audit done", file=sys.stderr, flush=True)
+        elif auditor.daemon is None:
+            print(">> warning: verifier daemon unavailable — no audit done",
+                  file=sys.stderr, flush=True)
+        else:
+            dropout_failures = _process_audit_results(auditor, tts, fragments, frag_wavs, ref_wav)
+        auditor.shutdown()
+
+    if dropout_failures:
+        idxs = [i for i, _ in dropout_failures]
+        print(f"\n⚠⚠⚠ DROPOUTS PERSISTED in fragments: {idxs}", file=sys.stderr, flush=True)
+        for i, dropped in dropout_failures:
+            print(f"    frag {i}: dropped pinyin = {dropped}", file=sys.stderr, flush=True)
+        print(f"    The respaced retries are in the final FLAC; play them and decide whether to ship.",
+              file=sys.stderr, flush=True)
+        print(f"    To hard-fix, edit the .md to insert spaces (任 重 道 远) or （停顿） around the "
+              f"affected chars and re-render.", file=sys.stderr, flush=True)
+
+    # ffmpeg pad-concat: every fragment but the last gets apad=PAUSE_SEC silence trailer
+    raw_concat = args.workdir / "concat_raw.wav"
+    inputs: list[str] = []
+    for w in frag_wavs:
+        inputs += ["-i", str(w)]
+    n = len(frag_wavs)
+    if n == 1:
+        fc = "[0:a]anull[out]"
+    else:
+        pads = ";".join(f"[{i}:a]apad=pad_dur={PAUSE_SEC}[a{i}]" for i in range(n - 1))
+        chain = "".join(f"[a{i}]" for i in range(n - 1)) + f"[{n-1}:a]"
+        fc = f"{pads};{chain}concat=n={n}:v=0:a=1[out]"
+    cmd = ["ffmpeg", "-y", *inputs, "-filter_complex", fc, "-map", "[out]", "-c:a", "flac", str(raw_concat)]
+    print(f">> ffmpeg concat ({n} frags + {PAUSE_SEC}s pads)…", flush=True)
+    subprocess.run(cmd, check=True, capture_output=True)
+
+    # ffmpeg atempo post-process
+    final = args.output.resolve()
+    final.parent.mkdir(parents=True, exist_ok=True)
+    cmd = ["ffmpeg", "-y", "-i", str(raw_concat), "-filter:a", f"atempo={args.speed}", "-c:a", "flac", str(final)]
+    print(f">> ffmpeg atempo={args.speed} → {final}", flush=True)
+    subprocess.run(cmd, check=True, capture_output=True)
+
+    dur = subprocess.check_output(
+        ["ffprobe", "-v", "error", "-show_entries", "format=duration", "-of", "csv=p=0", str(final)]
+    ).decode().strip()
+    print(f">> done: {final} ({dur}s)", flush=True)
+
+    # Sidecar — consumed by Wan2.2-s2v-infinite/render_episode.py for per-group sub-render.
+    # Durations are POST-atempo (master-FLAC timeline), so render_episode's
+    # compute_reset_points lands （停顿） boundaries at the right master-t.
+    if args.ep_id:
+        per_frag_durs: list[float] = []
+        for w in frag_wavs:
+            d = subprocess.check_output(
+                ["ffprobe", "-v", "error", "-show_entries", "format=duration", "-of", "csv=p=0", str(w)]
+            ).decode().strip()
+            per_frag_durs.append(float(d))
+        sidecar_path = final.parent / f"{args.ep_id}-segments.json"
+        import json as _json
+        sidecar_path.write_text(_json.dumps(
+            {
+                "pause_sec": PAUSE_SEC / args.speed,
+                "segments": [
+                    {"index": i + 1, "chars": len(s), "duration_sec": d / args.speed}
+                    for i, (s, d) in enumerate(zip(fragments, per_frag_durs))
+                ],
+            },
+            indent=2, ensure_ascii=False,
+        ))
+        print(f">> sidecar: {sidecar_path}", flush=True)
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())