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

comfy/custom_nodes/ComfyUI-IndexTTS2/verify_klaus_dropout.py

@@ -0,0 +1,268 @@
+#!/usr/bin/env python3
+"""verify_klaus_dropout.py — ASR-confirm a TTS dropout suspicion.
+
+Sibling tool to render_klaus_episode.py. Lives on the server alongside it
+(/root/IndexTTS2/) but runs in the *whisperx* conda env (NOT the IndexTTS-2
+venv), since whisperx pulls a different torch + faster-whisper stack.
+
+Layer 2 v2 — pinyin-sequence verdict (2026-05-11 rewrite):
+  Klaus uses klaus_lexicon to deliberately mangle tones; Whisper without
+  tone context picks homophone characters (`任重道远 → 任中刀援`,
+  `是→诗`, `像→相`), so grapheme-level Levenshtein had ~90% false
+  positives in audit_klaus_dropouts.py against ep01. The fix: compute
+  pinyin sequences (tone-stripped, Arabic numerals normalized to Chinese
+  via cn2an so `30 → 三十 → san shi`), and use the count of `delete`
+  opcodes in SequenceMatcher as the dropout verdict. Same-syllable
+  substitutions match (no false positive); missing syllables = real
+  dropout.
+
+Output: single JSON verdict line on stdout. `ok=true` ⇒ no missing
+syllables; `ok=false` ⇒ at least one expected pinyin syllable is absent
+from the ASR output. Grapheme char_ratio/seq_ratio stay in the output for
+diagnostics, but they no longer drive the verdict.
+
+Cost: ~10-15s model load + ~1-2s inference per call.
+"""
+
+from __future__ import annotations
+
+import argparse
+import contextlib
+import difflib
+import json
+import re
+import sys
+import unicodedata
+from pathlib import Path
+
+
+@contextlib.contextmanager
+def _stdout_to_stderr():
+    """Redirect Python-level sys.stdout → stderr inside the block. Required
+    around whisperx + pyannote calls: they leak INFO log lines straight to
+    stdout (pyannote in particular, e.g. "Performing voice activity
+    detection..."), which corrupts our line-delimited JSON daemon protocol.
+    `logging.basicConfig(stream=stderr, force=True)` alone doesn't catch
+    these because pyannote's logger gets configured during import, after we
+    set basicConfig."""
+    real = sys.stdout
+    sys.stdout = sys.stderr
+    try:
+        yield
+    finally:
+        sys.stdout = real
+
+
+_HAN = re.compile(r"[一-鿿]")
+
+
+def _to_simplified(text: str) -> str:
+    """Trad → simp via zhconv. WhisperX often returns traditional on Klaus
+    audio (tone errors push the decoder toward Cantonese-prior tokens);
+    without this, simp `远` vs trad `遠` would be treated as a substitution."""
+    try:
+        from zhconv import convert  # type: ignore
+        return convert(text, "zh-cn")
+    except ImportError:
+        return text
+
+
+def _normalize_digits(text: str) -> str:
+    """Replace Arabic digit clusters with Chinese readings (`30 → 三十`,
+    `8000 → 八千`). Klaus scripts mix Arabic and Chinese; ASR may transcribe
+    spoken Chinese numbers as Arabic digits. Without this, every digit
+    cluster collapses to zero pinyin tokens on one side and a syllable run
+    on the other, producing fake `delete` opcodes."""
+    try:
+        import cn2an  # type: ignore
+    except ImportError:
+        return text
+
+    def repl(m: re.Match) -> str:
+        try:
+            return cn2an.an2cn(m.group())
+        except Exception:
+            return m.group()
+    return re.sub(r"\d+", repl, text)
+
+
+def text_to_pinyin_sequence(text: str) -> list[str]:
+    """Convert text → list of tone-stripped lowercase pinyin syllables (Han only).
+
+    Pipeline: NFKC normalize → trad→simp → digits→Chinese → keep only Han
+    chars → pypinyin Style.NORMAL (no tone marks).
+    """
+    from pypinyin import pinyin, Style  # type: ignore
+
+    s = unicodedata.normalize("NFKC", text)
+    s = _to_simplified(s)
+    s = _normalize_digits(s)
+    han_only = "".join(c for c in s if _HAN.match(c))
+    if not han_only:
+        return []
+    # pypinyin returns nested list: [['ren'], ['zhong'], ...]; flatten and lowercase
+    return [p[0].lower() for p in pinyin(han_only, style=Style.NORMAL)]
+
+
+def _grapheme_normalize(text: str) -> str:
+    """Legacy grapheme normalize (alphanumerics only, trad→simp, NFKC) for
+    diagnostic char_ratio / seq_ratio. NOT used in the verdict anymore."""
+    s = unicodedata.normalize("NFKC", _to_simplified(text))
+    return "".join(c for c in s if unicodedata.category(c)[0] in ("L", "N"))
+
+
+def pinyin_diff(exp_seq: list[str], got_seq: list[str]) -> dict:
+    """SequenceMatcher diff on pinyin syllable lists. Returns counts +
+    list of dropped syllables (the things that show up in `delete` opcodes
+    of the expected sequence)."""
+    matcher = difflib.SequenceMatcher(None, exp_seq, got_seq, autojunk=False)
+    counts = {"equal": 0, "delete": 0, "replace": 0, "insert": 0}
+    dropped: list[str] = []
+    replaces: list[tuple[list[str], list[str]]] = []
+    for tag, i1, i2, j1, j2 in matcher.get_opcodes():
+        if tag == "equal":
+            counts["equal"] += i2 - i1
+        elif tag == "delete":
+            counts["delete"] += i2 - i1
+            dropped.extend(exp_seq[i1:i2])
+        elif tag == "replace":
+            counts["replace"] += max(i2 - i1, j2 - j1)
+            replaces.append((exp_seq[i1:i2], got_seq[j1:j2]))
+        elif tag == "insert":
+            counts["insert"] += j2 - j1
+    return {
+        "counts": counts,
+        "dropped": dropped,
+        "replaces": replaces,
+        "ratio": round(matcher.ratio(), 3),
+    }
+
+
+def _compute_verdict(expected: str, got_text: str) -> dict:
+    """Run the pinyin-sequence diff and pack the verdict + diagnostics into the
+    shared response shape (used by both one-shot and daemon modes)."""
+    exp_pinyin = text_to_pinyin_sequence(expected)
+    got_pinyin = text_to_pinyin_sequence(got_text)
+    diff = pinyin_diff(exp_pinyin, got_pinyin)
+    ok = diff["counts"]["delete"] == 0
+
+    exp_g = _grapheme_normalize(expected)
+    got_g = _grapheme_normalize(got_text)
+    char_ratio = len(got_g) / max(len(exp_g), 1)
+    seq_ratio = difflib.SequenceMatcher(None, exp_g, got_g, autojunk=False).ratio()
+
+    return {
+        "ok": ok,
+        "expected": expected,
+        "got": got_text,
+        # Pinyin verdict (primary, drives `ok`)
+        "pinyin_expected_count": len(exp_pinyin),
+        "pinyin_got_count": len(got_pinyin),
+        "pinyin_diff_counts": diff["counts"],
+        "pinyin_dropped": diff["dropped"],
+        "pinyin_ratio": diff["ratio"],
+        # Grapheme diagnostics (informational only)
+        "char_ratio": round(char_ratio, 3),
+        "seq_ratio": round(seq_ratio, 3),
+        # Backward-compat aliases for older render_klaus_episode.py print lines
+        "expected_chars": len(exp_g),
+        "got_chars": len(got_g),
+        "ratio": round(seq_ratio, 3),
+    }
+
+
+def _daemon_loop(args) -> int:
+    """Long-running mode: load model once, then read line-delimited JSON
+    requests `{"audio": "...", "expected": "..."}` from stdin and write a JSON
+    verdict response per line to stdout.
+
+    Special control lines on stdin:
+      QUIT    — graceful shutdown, returns 0.
+      (EOF)   — same as QUIT.
+
+    Designed to be spawned by render_klaus_episode.py for per-fragment audit
+    after all fragments have been rendered. Loading whisperx once means each
+    fragment check is ~1-2s instead of ~15s cold-spawn."""
+    # Force ML libs' logging to stderr BEFORE importing whisperx.
+    import logging
+    logging.basicConfig(stream=sys.stderr, force=True, level=logging.WARNING)
+
+    # whisperx + pyannote leak INFO via direct print/logger-with-stdout-handler
+    # despite our basicConfig — wrap all ML calls in stdout→stderr redirect.
+    with _stdout_to_stderr():
+        import whisperx  # type: ignore
+        asr = whisperx.load_model(
+            args.asr_model, args.device, language=args.language,
+            compute_type=args.compute_type,
+            asr_options={"beam_size": 1, "condition_on_previous_text": False},
+        )
+
+    # Handshake: client blocks on this line until it appears.
+    print(json.dumps({"status": "ready"}), flush=True)
+
+    for raw in sys.stdin:
+        line = raw.strip()
+        if not line or line == "QUIT":
+            break
+        try:
+            req = json.loads(line)
+            audio_path = req["audio"]
+            expected = req["expected"]
+        except (json.JSONDecodeError, KeyError) as e:
+            print(json.dumps({"error": f"bad request: {e}"}, ensure_ascii=False), flush=True)
+            continue
+
+        try:
+            with _stdout_to_stderr():
+                audio = whisperx.load_audio(audio_path)
+                result = asr.transcribe(audio, batch_size=8)
+            got_text = "".join(seg.get("text", "") for seg in result["segments"]).strip()
+            verdict = _compute_verdict(expected, got_text)
+            verdict["audio"] = audio_path
+            print(json.dumps(verdict, ensure_ascii=False), flush=True)
+        except Exception as e:
+            print(json.dumps({"error": str(e), "audio": audio_path}, ensure_ascii=False), flush=True)
+    return 0
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
+    ap.add_argument("--daemon", action="store_true",
+                    help="Long-running mode: read JSON requests from stdin, "
+                         "write JSON verdicts to stdout. Whisperx loads once. "
+                         "Used by render_klaus_episode.py's post-render audit.")
+    ap.add_argument("--audio", help="(one-shot mode) path to .wav/.flac to verify")
+    ap.add_argument("--expected", help="(one-shot mode) the text that *should* "
+                                       "have been spoken")
+    ap.add_argument("--language", default="zh")
+    ap.add_argument("--asr-model", default="large-v3")
+    ap.add_argument("--device", default="cuda")
+    ap.add_argument("--compute-type", default="float16")
+    args = ap.parse_args()
+
+    if args.daemon:
+        return _daemon_loop(args)
+
+    # One-shot mode (kept for ad-hoc debugging + smoke-tests)
+    if not args.audio or not args.expected:
+        sys.exit("one-shot mode requires --audio and --expected (or use --daemon)")
+    if not Path(args.audio).exists():
+        sys.exit(f"audio missing: {args.audio}")
+
+    import whisperx  # type: ignore
+    audio = whisperx.load_audio(args.audio)
+    asr = whisperx.load_model(
+        args.asr_model, args.device, language=args.language,
+        compute_type=args.compute_type,
+        asr_options={"beam_size": 1, "condition_on_previous_text": False},
+    )
+    result = asr.transcribe(audio, batch_size=8)
+    got_text = "".join(seg.get("text", "") for seg in result["segments"]).strip()
+
+    verdict = _compute_verdict(args.expected, got_text)
+    print(json.dumps(verdict, ensure_ascii=False))
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())