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

comfy/custom_nodes/ComfyUI-IndexTTS2/audit_klaus_dropouts.py

@@ -0,0 +1,158 @@
+#!/usr/bin/env python3
+"""audit_klaus_dropouts.py — re-ASR all per-fragment wav files from a Klaus
+render and compute the ground-truth dropout report.
+
+Runs in the whisperx conda env (NOT the IndexTTS-2 venv). Reuses md_to_fragments
+from render_klaus_episode.py to derive the expected text per fragment from the
+.md (same processing the render loop uses), then ASRs each frag wav file and
+diffs got vs expected at the **pinyin sequence** level (tone-stripped, digit-
+normalized) — same primitive used by verify_klaus_dropout.py.
+
+Why pinyin diff (v2, 2026-05-11): grapheme-level Levenshtein had ~90% false
+positives on Klaus audio because klaus_lexicon deliberately mangles tones
+and Whisper substitutes homophones. Pinyin sequence diff treats homophone
+substitutions as match (no false positive) and only flags actual missing
+syllables.
+
+Outputs a per-fragment table + summary of which fragments dropped which
+specific syllables. Loads whisperx ONCE — total wall ~30s + 1-2s per fragment.
+
+Usage (on server, after a Klaus render):
+    /root/miniconda3/envs/whisperx/bin/python \
+        /root/IndexTTS2/audit_klaus_dropouts.py \
+        --md /tmp/klaus_ep01_listen.md \
+        --frag-dir /tmp/klaus_render \
+        [--json /tmp/klaus_audit.json]
+"""
+
+from __future__ import annotations
+
+import argparse
+import difflib
+import json
+import re
+import sys
+import time
+import unicodedata
+from pathlib import Path
+
+# Reuse the verifier's pinyin pipeline so audit + production stay in sync.
+sys.path.insert(0, str(Path(__file__).parent))
+from verify_klaus_dropout import (  # noqa: E402
+    text_to_pinyin_sequence,
+    pinyin_diff,
+    _grapheme_normalize,
+)
+
+
+def md_to_fragments(md: str) -> list[str]:
+    """Mirror of render_klaus_episode.py's md_to_fragments — kept in sync."""
+    body = "\n".join(ln for ln in md.splitlines() if not ln.lstrip().startswith("#"))
+    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 main() -> int:
+    ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
+    ap.add_argument("--md", required=True, type=Path, help=".md script used for the render")
+    ap.add_argument("--frag-dir", required=True, type=Path,
+                    help="dir containing frag001.wav ... fragNNN.wav")
+    ap.add_argument("--json", type=Path, default=None,
+                    help="optional path to write the full report as JSON")
+    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()
+
+    fragments = md_to_fragments(args.md.read_text())
+    print(f"## script: {args.md.name}", flush=True)
+    print(f"## fragments parsed: {len(fragments)}", flush=True)
+
+    frag_wavs = sorted(args.frag_dir.glob("frag*.wav"))
+    print(f"## frag wavs found: {len(frag_wavs)}", flush=True)
+    if len(fragments) != len(frag_wavs):
+        print(f"⚠ mismatch: {len(fragments)} fragments vs {len(frag_wavs)} wavs — proceeding "
+              f"with min(len)", file=sys.stderr, flush=True)
+    n = min(len(fragments), len(frag_wavs))
+
+    print(f"## loading whisperx {args.asr_model} …", flush=True)
+    t0 = time.perf_counter()
+    import whisperx
+    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},
+    )
+    print(f"## loaded in {time.perf_counter()-t0:.1f}s", flush=True)
+
+    rows = []
+    print()
+    print(f"{'idx':>3} {'py_exp':>6} {'py_got':>6} {'del':>3} {'rpl':>3} {'ins':>3} {'verdict':>7}  dropped_syllables")
+    print("-" * 80)
+
+    for i in range(n):
+        frag_text = fragments[i]
+        wav_path = frag_wavs[i]
+        t0 = time.perf_counter()
+        audio = whisperx.load_audio(str(wav_path))
+        result = asr.transcribe(audio, batch_size=8)
+        got_text = "".join(seg.get("text", "") for seg in result["segments"]).strip()
+        elapsed = time.perf_counter() - t0
+
+        exp_pinyin = text_to_pinyin_sequence(frag_text)
+        got_pinyin = text_to_pinyin_sequence(got_text)
+        diff = pinyin_diff(exp_pinyin, got_pinyin)
+        counts = diff["counts"]
+        verdict = "DROP" if counts["delete"] >= 1 else "OK"
+
+        # Grapheme diagnostics (informational only)
+        exp_g = _grapheme_normalize(frag_text)
+        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()
+
+        dropped_str = " ".join(diff["dropped"]) if diff["dropped"] else ""
+        print(f"{i+1:>3} {len(exp_pinyin):>6} {len(got_pinyin):>6} "
+              f"{counts['delete']:>3} {counts['replace']:>3} {counts['insert']:>3} "
+              f"{verdict:>7}  {dropped_str}", flush=True)
+
+        rows.append({
+            "idx": i + 1,
+            "wav": wav_path.name,
+            "expected_text": frag_text,
+            "got_text": got_text,
+            "pinyin_expected_count": len(exp_pinyin),
+            "pinyin_got_count": len(got_pinyin),
+            "pinyin_diff_counts": counts,
+            "pinyin_dropped": diff["dropped"],
+            "pinyin_ratio": diff["ratio"],
+            "char_ratio": round(char_ratio, 3),
+            "seq_ratio": round(seq_ratio, 3),
+            "verdict": verdict,
+            "asr_sec": round(elapsed, 2),
+        })
+
+    drops = [r for r in rows if r["verdict"] == "DROP"]
+    print()
+    print(f"## summary: {len(drops)}/{n} fragments flagged DROP (pinyin-level, ground truth)")
+    for r in drops:
+        print(f"   frag {r['idx']:>2}: dropped pinyin = {r['pinyin_dropped']}")
+        print(f"     expected:  {r['expected_text'][:120]}")
+        print(f"     got:       {r['got_text'][:120]}")
+
+    if args.json:
+        args.json.write_text(json.dumps(rows, ensure_ascii=False, indent=2))
+        print(f"## json report: {args.json}")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())