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

comfy/custom_nodes/ComfyUI-IndexTTS2/klaus_lexicon.py

@@ -0,0 +1,176 @@
+#!/usr/bin/env python3
+"""klaus_lexicon.py — transform Chinese text into IndexTTS-2 mangled-pinyin format
+simulating a 3-month-Mandarin Australian learner (Klaus).
+
+IndexTTS-2 accepts inline UPPERCASE pinyin tokens directly in the input string,
+e.g. '今天HEN3累' forces 很 to be read as "HEN3" instead of the standard pinyin.
+Constraint (checkpoints/pinyin.vocab, 1728 entries): tokens MUST be valid
+Mandarin syllables — `WOR`/`DUH`/English-vowel reshapes do NOT work; we can
+only swap one legal syllable for another (wrong tone, wrong reading, plain
+for retroflex, u for ü).
+
+This module applies four classes of beginner-level error, weighted by frequency
+in real Anglo L2 Mandarin learners (Cantonese/non-tonal-L1 priors are similar):
+
+  - Tone collapse: 4→1 (loss of falling), 3→2 (loss of dipping)
+  - Retroflex loss: ZH/CH/SH → Z/C/S (no curl-tongue training)
+  - ü→u (V→U): English has no /y/, defaults to /u/
+
+Each candidate token is validated against the IndexTTS-2 pinyin.vocab set;
+invalid candidates fall back to emitting the Chinese character verbatim (so the
+model uses its built-in standard pronunciation).
+
+Output: mixed Chinese + uppercase pinyin tokens, ready to feed
+`indextts.infer_v2.IndexTTS2.infer(text=...)`.
+"""
+
+from __future__ import annotations
+
+import argparse
+import random
+import re
+import sys
+from pathlib import Path
+
+try:
+    import pypinyin
+except ImportError:
+    sys.exit("klaus_lexicon.py: pypinyin required. `pip install pypinyin`")
+
+VOCAB_PATH = Path(__file__).parent / "pinyin.vocab"
+VOCAB: set[str] = set(VOCAB_PATH.read_text().split())  # 1728 legal syllables
+
+# Default probability tuning — calibrated for "3-month learner" perceptual signal.
+# Override per-call via transform(... tone_4to1_prob=, tone_3to2_prob=, retroflex_loss_prob=, u_prob=).
+# Hitting 1.0 across the board produces robotic-not-human output; keep at least one rule under 0.8.
+TONE_4TO1_PROB = 0.4   # loss of falling tone
+TONE_3TO2_PROB = 0.5   # loss of dipping tone (the hardest tone for English speakers)
+RETROFLEX_LOSS_PROB = 0.6  # zh/ch/sh → z/c/s
+U_PROB = 0.7           # ü → u (V → U)
+
+_HAN = re.compile(r"[一-鿿]")
+
+
+def _is_han(ch: str) -> bool:
+    return bool(_HAN.match(ch))
+
+
+def _standard_pinyin(ch: str) -> str | None:
+    """Return e.g. 'ZHONG1' for 中, or None if not a known Chinese char."""
+    py = pypinyin.pinyin(ch, style=pypinyin.Style.TONE3, errors="ignore", heteronym=False)
+    if not py or not py[0]:
+        return None
+    raw = py[0][0]  # 'zhong1' or 'ma' (neutral tone, no digit)
+    if not raw or not raw[-1].isdigit():
+        raw = (raw or "") + "5"  # neutral tone
+    return raw.upper()
+
+
+def _retroflex_loss(token: str) -> str:
+    if token.startswith("ZH"):
+        return "Z" + token[2:]
+    if token.startswith("CH"):
+        return "C" + token[2:]
+    if token.startswith("SH"):
+        return "S" + token[2:]
+    return token
+
+
+def transform(
+    text: str,
+    seed: int = 42,
+    tone_4to1_prob: float | None = None,
+    tone_3to2_prob: float | None = None,
+    retroflex_loss_prob: float | None = None,
+    u_prob: float | None = None,
+) -> tuple[str, dict]:
+    """Apply Klaus lexicon to `text`. Returns (transformed_string, stats).
+
+    Per-call probability overrides default to the module constants when None,
+    so existing callers (smoke_test, segment_klaus_tts) get the calibrated
+    behavior; sweeps pass explicit values.
+
+    stats keys: total_han, overrides_emitted, fallbacks_invalid, fallbacks_unchanged.
+    Use stats to sanity-check that the lexicon is hitting reasonable error rates
+    (e.g. overrides_emitted / total_han should land around 0.4–0.7 for a heavy accent).
+    """
+    p_4to1 = TONE_4TO1_PROB if tone_4to1_prob is None else tone_4to1_prob
+    p_3to2 = TONE_3TO2_PROB if tone_3to2_prob is None else tone_3to2_prob
+    p_retro = RETROFLEX_LOSS_PROB if retroflex_loss_prob is None else retroflex_loss_prob
+    p_u = U_PROB if u_prob is None else u_prob
+
+    rng = random.Random(seed)
+    out: list[str] = []
+    stats = dict(total_han=0, overrides_emitted=0, fallbacks_invalid=0, fallbacks_unchanged=0)
+
+    for ch in text:
+        if not _is_han(ch):
+            out.append(ch)
+            continue
+        stats["total_han"] += 1
+
+        std = _standard_pinyin(ch)
+        if std is None or std not in VOCAB:
+            # Don't know how to mangle it; let the model handle the char.
+            out.append(ch)
+            stats["fallbacks_invalid"] += 1
+            continue
+
+        candidate = std
+
+        # Retroflex loss
+        if rng.random() < p_retro:
+            candidate = _retroflex_loss(candidate)
+
+        # Tone collapse
+        body, digit = candidate[:-1], candidate[-1]
+        if digit == "4" and rng.random() < p_4to1:
+            candidate = body + "1"
+        elif digit == "3" and rng.random() < p_3to2:
+            candidate = body + "2"
+
+        # ü → u
+        if "V" in candidate and rng.random() < p_u:
+            candidate = candidate.replace("V", "U")
+
+        if candidate == std:
+            # No mutation rolled true — keep Chinese, model uses standard pronunciation.
+            out.append(ch)
+            stats["fallbacks_unchanged"] += 1
+            continue
+
+        if candidate not in VOCAB:
+            # Mutation produced an illegal syllable (rare, e.g. some V→U combos).
+            # Fall back to the standard pronunciation rather than risk an invalid token.
+            out.append(ch)
+            stats["fallbacks_invalid"] += 1
+            continue
+
+        out.append(candidate)
+        stats["overrides_emitted"] += 1
+
+    return "".join(out), stats
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
+    ap.add_argument("text", nargs="?", default="大家好我叫Klaus，我学中文学了三个月，普通话真的很难，特别是声调")
+    ap.add_argument("--seed", type=int, default=42)
+    ap.add_argument("--quiet", action="store_true", help="suppress stats line on stderr")
+    args = ap.parse_args()
+
+    transformed, stats = transform(args.text, seed=args.seed)
+    print(transformed)
+    if not args.quiet:
+        rate = stats["overrides_emitted"] / max(stats["total_han"], 1)
+        print(
+            f"[stats] han={stats['total_han']} overrides={stats['overrides_emitted']} "
+            f"({rate:.0%}) unchanged={stats['fallbacks_unchanged']} "
+            f"invalid={stats['fallbacks_invalid']}",
+            file=sys.stderr,
+        )
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())