infer.py

#!/usr/bin/env python
"""Dynamic Accent: Nigerian multilingual, code-switching TTS, with selectable model versions.

Versions (`--version`):
  flash   (default) LM `checkpoints/flash/flash_dpo` on Fun-CosyVoice3-0.5B-2512 (Apache-2.0): Flash SFT + DPO preference
          tuning (naturalness + CER + pace + pause reward), licensed data only. CV3 base flow/vocoder, zero-shot from the
          voice clip. Text front-end: numbers verbalised, Yoruba/Igbo diacritics restored when missing (ByT5 model in
          checkpoints/flash/diacritics_byt5). One pass per paragraph (<= ~60 words), model pauses, no post-processing.
          Needs a CUDA GPU and a CosyVoice clone with CosyVoice3 support.
  phase2  LM `p2_lm_B` + anchor-SFT flow `p2_flow_sft`. Phase 2 sampler, sentence stitching, edge trim, internal
          silences capped at 0.6 s, 0.6 s between sentences (exactly the Phase 2 final samples).
  phase3  LM `p3_lm_full` (re-cut long-form data with natural pauses, ~102 h of Igbo) + the shared Phase 2 flow (see
          settings.json and checkpoints/phase3/p3_lm_full/config.json). Silence-aware sampler (silence tokens exempt from the repetition check up
          to a 1 s run), the whole paragraph in ONE pass (up to ~60 words per pass), no audio post-processing: every
          pause and the pace come from the model.

Voices (`--voice`): `anchor_male` (default; WAXAL Yoruba speaker 4), `waxal_ibo_5_female`, `waxal_yor_5_female`
(see --list_voices; recommended: phase3 with a female voice, phase2 with anchor_male). A non-anchor voice is zero-shot: its clip is the flow prompt and the LM prompt, and its
speaker token is used when the LM was trained on that speaker.

  python infer.py --text "<eng>this is dynamic accent,</eng> <pcm>e dey sound nice abi?</pcm>" --version phase3 --out out.wav
  python infer.py --text "..." --compare --out cmp.wav       # -> cmp_phase2/phase3/flash.wav, same text, seed, voice
  python infer.py --list_voices
  python infer.py --auto_tag --tag_only --text "Omo, this traffic no be here o. Ẹ ṣé púpọ̀."   # inspect / edit the tags

Text: wrap each language span in <eng>, <pcm>, <yor>, <ibo> or <hau> tags. Untagged text is treated as English, unless
--auto_tag (automatic word/span language ID, see tagger.py), --force_lang or --lang_map is given.
`--pause_mode clause` / `--speed` (prosody v2: fixed gaps + WSOLA) are kept for phase2 only as an experiment; a
listening test preferred the model's own pauses.
"""

import argparse, copy, inspect, json, os, re, sys, unicodedata

import numpy as np
import soundfile as sf
import librosa
import torch
import torch.nn as nn
import contextlib

HERE = os.path.dirname(os.path.abspath(__file__))
LANGS = ["eng", "ibo", "hau", "yor", "pcm"]
L2I = {l: i for i, l in enumerate(LANGS)}
ANCHOR = "waxal_yor_4"
DEV = torch.device("cuda" if torch.cuda.is_available() else "cpu")   # set_device() / --device to override
DEFAULT_SPEED = 1.0    # user listening test 2026-10-21: original model pauses preferred; clause/1.05 kept as experimental


def set_device(d):
    global DEV
    DEV = torch.device(d if d != "auto" else ("cuda" if torch.cuda.is_available() else "cpu"))
    return DEV


def autocast(dtype):
    if DEV.type == "cuda" and dtype is not None:
        return torch.autocast("cuda", dtype=dtype)
    return contextlib.nullcontext()            # CPU: plain fp32


# ------------------------------------------------------------------ text front-end (tags -> spans)
_TAG = re.compile(r"<(eng|ibo|hau|yor|pcm)>(.*?)</\1>", re.S)
_SENT = re.compile(r"(?<=[.?!])\s+")


def normalize(t):
    t = unicodedata.normalize("NFC", t)
    t = t.replace("\u2019", "'").replace("\u2018", "'").replace("\u201c", '"').replace("\u201d", '"')
    return re.sub(r"\s+", " ", t).strip()


def segments(text, default_lang="eng"):
    out, pos = [], 0
    for m in _TAG.finditer(text):
        pre = text[pos:m.start()]
        if pre.strip():
            out.append((default_lang, pre))
        out.append((m.group(1), m.group(2)))
        pos = m.end()
    if text[pos:].strip():
        out.append((default_lang, text[pos:]))
    res = []
    for l, s in out:
        s = normalize(s)
        if not s:
            continue
        if res:
            s = " " + s
        res.append((l, s))
    return res


def text_tokens(tok, text, default="eng"):
    ids, langs = [], []
    for l, s in segments(text, default):
        t = tok(s, add_special_tokens=False)["input_ids"]
        ids += t; langs += [L2I[l]] * len(t)
    return ids, langs


def split_sentences(text):
    out, cur = [], []
    for l, s in segments(text):
        parts = _SENT.split(s.strip())
        for k, p in enumerate(parts):
            p = p.strip()
            if p:
                cur.append(f"<{l}>{p}</{l}>")
            if (k < len(parts) - 1 or p[-1:] in ".?!") and cur:
                out.append(" ".join(cur)); cur = []
    if cur:
        out.append(" ".join(cur))
    return [o for o in out if o.strip()]


# ------------------------------------------------------------------ silence post-processing
def frame_db(y, sr, hop_s=0.01):
    hop = int(sr * hop_s)
    n = max(1, len(y) // hop)
    fr = y[: n * hop].reshape(n, hop)
    return 10 * np.log10((fr.astype(np.float64) ** 2).mean(1) + 1e-10)


def silence_mask(y, sr, rel_db=-40.0, hop_s=0.01):
    e = frame_db(y, sr, hop_s)
    return e < np.percentile(e, 98) + rel_db


def runs(mask):
    out, i, n = [], 0, len(mask)
    while i < n:
        if mask[i]:
            j = i
            while j < n and mask[j]:
                j += 1
            out.append((i, j)); i = j
        else:
            i += 1
    return out


def trim_edges(w, sr, keep_lead=0.12, keep_tail=0.18):
    sil = silence_mask(w, sr)
    sp = np.where(~sil)[0]
    if not len(sp):
        return w
    a, b = sp[0] * 0.01, (sp[-1] + 1) * 0.01
    out = w[int(max(0, a - keep_lead) * sr): int(min(len(w) / sr, b + keep_tail) * sr)]
    lead = a - max(0, a - keep_lead)
    if lead < keep_lead:
        out = np.concatenate([np.zeros(int((keep_lead - lead) * sr), out.dtype), out])
    return out


def cap_internal(w, sr, cap=0.6):
    sil = silence_mask(w, sr)
    hop = int(0.01 * sr)
    R = [(a, b) for a, b in runs(sil) if a > 0 and b < len(sil) and (b - a) * 0.01 > cap]
    if not R:
        return w
    keep, cur = [], 0
    for a, b in R:
        keep.append(w[cur: a * hop + int(cap / 2 * sr)]); cur = b * hop - int(cap / 2 * sr)
    keep.append(w[cur:])
    return np.concatenate(keep)



# ------------------------------------------------------------------ prosody v2: clause chunks, set gaps, pitch-preserving speed
_PUNCT = re.compile(r"(?<=[,;:.?!])\s+")
GAPS = {"comma": 0.40, "sent": 0.60, "switch": 0.12}
JITTER = {"comma": 0.05, "sent": 0.03, "switch": 0.02}


def clause_chunks(text):
    """tagged text -> [(lang, chunk_text, boundary_after)] with boundary in comma|sent|switch|None (None = end).
    Chunks are cut after , ; : . ? ! and at language switches; a one-word chunk is merged into the next chunk of the
    same language (one-word LM calls are unstable), keeping its comma inside."""
    raw = []
    for l, s in segments(text):
        for p in _PUNCT.split(s.strip()):
            if p.strip():
                raw.append([l, p.strip()])
    merged = []
    for c in raw:
        if merged and merged[-1][0] == c[0] and len(merged[-1][1].split()) < 2:
            merged[-1][1] = merged[-1][1] + " " + c[1]
        else:
            merged.append(c)
    out = []
    for i, (l, t) in enumerate(merged):
        if i == len(merged) - 1:
            b = None
        elif t[-1] in ".?!":
            b = "sent"
        elif t[-1] in ",;:":
            b = "comma"
        else:
            b = "switch"
        out.append((l, t, b))
    return out


def trim_tight(w, sr, keep=0.03):
    sil = silence_mask(w, sr)
    sp = np.where(~sil)[0]
    if not len(sp):
        return w
    a, b = sp[0] * 0.01, (sp[-1] + 1) * 0.01
    return w[int(max(0.0, a - keep) * sr): int(min(len(w) / sr, b + keep) * sr)]


def time_stretch(w, sr, speed):
    """pitch-preserving tempo change (speed > 1 = faster). WSOLA via audiotsm, else librosa's phase vocoder."""
    if abs(speed - 1.0) < 1e-3:
        return w
    try:
        from audiotsm import wsola
        from audiotsm.io.array import ArrayReader, ArrayWriter
        r, wr = ArrayReader(w[None].astype(np.float32)), ArrayWriter(1)
        wsola(1, speed=speed, frame_length=int(0.040 * sr), synthesis_hop=int(0.010 * sr)).run(r, wr)
        return wr.data[0].astype(np.float32)
    except ImportError:
        return librosa.effects.time_stretch(w.astype(np.float32), rate=speed)


def fade(w, sr, ms=10):
    n = min(len(w) // 2, int(sr * ms / 1000))
    if n > 0:
        w = w.copy(); r = np.linspace(0, 1, n, dtype=np.float32)
        w[:n] *= r; w[-n:] *= r[::-1]
    return w

# ------------------------------------------------------------------ CosyVoice2 base
def _sf_load_wav(wav, target_sr, min_sr=16000):
    y, sr = sf.read(wav, dtype="float32", always_2d=True)
    y = y.mean(1)
    if sr != target_sr:
        y = librosa.resample(y, orig_sr=sr, target_sr=target_sr)
    return torch.from_numpy(y)[None]


def load_cosy(model_dir, cosyvoice_dir):
    sys.path[:0] = [cosyvoice_dir, os.path.join(cosyvoice_dir, "third_party", "Matcha-TTS")]
    try:                                   # pyworld is only needed for CosyVoice training (F0 features); it often fails
        import pyworld  # noqa: F401       # to build, so inference runs without it
    except ImportError:
        import types
        sys.modules["pyworld"] = types.ModuleType("pyworld")
    import cosyvoice.utils.file_utils as fu
    fu.load_wav = _sf_load_wav
    import cosyvoice.cli.frontend as cf
    cf.load_wav = _sf_load_wav
    from cosyvoice.cli.cosyvoice import CosyVoice2
    kw = dict(load_jit=False, load_trt=False, fp16=False)
    sig = inspect.signature(CosyVoice2.__init__).parameters
    return CosyVoice2(model_dir, **{k: v for k, v in kw.items() if k in sig})


def prompt_feats(cosy, wav_path):
    fe = cosy.frontend
    tokn, _ = fe._extract_speech_token(wav_path)
    feat, _ = fe._extract_speech_feat(wav_path)
    n = min(feat.shape[1] // 2, tokn.shape[1])
    feat, tokn = feat[:, :2 * n], tokn[:, :n]
    emb = fe._extract_spk_embedding(wav_path)
    return tokn.to(DEV), feat.to(DEV), emb.to(DEV)


@torch.no_grad()
def token2wav(cosy, tokens, ptok, pfeat, emb):
    flow = cosy.model.flow
    t = torch.tensor([tokens], dtype=torch.int32).to(DEV)
    kw = dict(token=t, token_len=torch.tensor([t.shape[1]], dtype=torch.int32).to(DEV),
              prompt_token=ptok, prompt_token_len=torch.tensor([ptok.shape[1]], dtype=torch.int32).to(DEV),
              prompt_feat=pfeat, prompt_feat_len=torch.tensor([pfeat.shape[1]], dtype=torch.int32).to(DEV),
              embedding=emb)
    sig = inspect.signature(flow.inference).parameters
    if "streaming" in sig:
        kw["streaming"] = False
    if "finalize" in sig:
        kw["finalize"] = True
    if "flow_cache" in sig:
        kw["flow_cache"] = torch.zeros(1, 80, 0, 2).to(DEV)
    mel = flow.inference(**kw)
    mel = mel[0] if isinstance(mel, tuple) else mel
    hs = inspect.signature(cosy.model.hift.inference).parameters
    w = cosy.model.hift.inference(speech_feat=mel, **({"cache_source": torch.zeros(1, 1, 0).to(DEV)} if "cache_source" in hs else {}))
    w = w[0] if isinstance(w, tuple) else w
    return w.squeeze().float().cpu().numpy()


# ------------------------------------------------------------------ Dynamic Accent LM (Qwen2LM + language + speaker embeddings)
# S3 tokens that fill silent frames in the anchor's recordings (Phase 3 Stage 0); used by the silence-aware sampler
SIL_TOKENS = {4137, 4218, 1950, 2031, 3894, 6324, 1707, 3677, 1701, 6405, 1947, 3645, 1944, 1704, 3891, 4134, 3680, 5838, 3651}
RATES = [None, "very slow", "slow", "normal", "fast", "very fast"]


class DALM2(nn.Module):
    def __init__(self, qlm, n_spk, d=896):
        super().__init__()
        self.q = qlm
        self.lang_emb = nn.Embedding(len(LANGS), d)
        self.spk_emb = nn.Embedding(n_spk, d)
        self.rate_emb = nn.Embedding(len(RATES), d, padding_idx=0)   # phase3 speed instruct; zero = none (phase2)
        nn.init.zeros_(self.rate_emb.weight)
        self.adv = nn.Sequential(nn.Linear(d, 256), nn.ReLU(), nn.Linear(256, n_spk))  # training-only head
        self.V = qlm.speech_token_size

    def embed_text(self, ids, langs):
        return self.q.llm.model.model.embed_tokens(ids) + self.lang_emb(langs)

    def build2(self, ids, langs, spk, rate=0, pids=None, plangs=None, psp=None):
        sos = self.q.llm_embedding.weight[self.q.sos][None]
        task = self.q.llm_embedding.weight[self.q.task_id][None]
        if pids is not None:
            ids = torch.cat([pids, ids]); langs = torch.cat([plangs, langs])
        sv = self.spk_emb(spk)
        if rate:
            sv = sv + self.rate_emb(torch.tensor(rate, device=sv.device))
        parts = [sos, self.embed_text(ids, langs), sv[None], task]
        if psp is not None:
            parts.append(self.q.speech_embedding(psp))
        return torch.cat(parts, 0)

    @torch.no_grad()
    def generate2(self, ids, langs, spk, prompt=None, top_p=0.8, top_k=25, win=10, tau_r=0.1, seed=0,
                  min_ratio=2.0, max_ratio=20.0, sampler="ras", rate=0, sil_cap=25):
        """sampler 'ras' = Phase 2 repetition-aware sampling; 'silexcap' = silence tokens exempt from the repetition
        check while the current silence run is shorter than sil_cap frames (25 = 1.0 s), then the normal check."""
        g = torch.Generator(device=ids.device).manual_seed(seed)
        x = self.build2(ids, langs, spk, rate, *(prompt if prompt else (None, None, None)))[None]
        n = len(ids)
        min_len, max_len = int(n * min_ratio), int(n * max_ratio)
        past, out, inp = None, [], x
        for i in range(max_len):
            o = self.q.llm.model.model(inputs_embeds=inp, past_key_values=past, use_cache=True, return_dict=True)
            past = o.past_key_values
            logp = self.q.llm_decoder(o.last_hidden_state[:, -1]).float().log_softmax(-1)[0]
            if i < min_len:
                logp[self.V] = -float("inf")
            logp[self.V + 1:] = -float("inf")
            p = logp.exp()
            sp, si = p.sort(descending=True)
            keep = (sp.cumsum(0) - sp) < top_p
            keep[top_k:] = False
            cp, ci = sp[keep], si[keep]
            tokid = int(ci[torch.multinomial(cp / cp.sum(), 1, generator=g)])
            exempt = False
            if sampler == "silexcap" and tokid in SIL_TOKENS:
                run = 0
                for t in reversed(out):
                    if t not in SIL_TOKENS:
                        break
                    run += 1
                exempt = run < sil_cap
            if not exempt and out and sum(1 for t in out[-win:] if t == tokid) >= win * tau_r and tokid != self.V:  # RAS
                tokid = int(torch.multinomial(p, 1, generator=g))
            if tokid >= self.V:
                break
            out.append(tokid)
            inp = self.q.speech_embedding.weight[tokid][None, None]
        return out


SETTINGS_PATH = os.path.join(HERE, "settings.json")
SETTINGS = json.load(open(SETTINGS_PATH)) if os.path.exists(SETTINGS_PATH) else {}
VERSIONS = {k: {kk: vv for kk, vv in v.items() if kk != "desc"} for k, v in SETTINGS.get("versions", {}).items()} or {
    "phase2": dict(lm="checkpoints/phase2/p2_lm_B/best.pt", flow="checkpoints/phase2/p2_flow_sft/best.pt", sampler="ras",
                   one_pass=False, postprocess=True, rate=None),
    "phase3": dict(lm="checkpoints/phase3/p3_lm_full/best.pt", flow="checkpoints/phase2/p2_flow_sft/best.pt",
                   sampler="silexcap", one_pass=True, postprocess=False, rate="normal"),
}
RECOMMENDED = SETTINGS.get("recommended", {})


def recommendation(version, voice):
    """one-line hint when a version/voice pairing is not the recommended one (None when it is)."""
    rec = RECOMMENDED.get(version)
    if rec and voice not in rec:
        return f"note: {version} is recommended with {' or '.join(rec)} ({RECOMMENDED.get('note', '')})"
    return None


def version_cfg(name):
    cfg = dict(VERSIONS[name])
    cj = os.path.join(HERE, os.path.dirname(cfg["lm"]), "config.json")
    if os.path.exists(cj):                      # a checkpoint may override flow / rate / sampler
        cfg.update({k: v for k, v in json.load(open(cj)).items() if k in cfg})
    for k in ("lm", "flow", "diacritics"):
        if cfg.get(k):
            cfg[k] = os.path.join(HERE, cfg[k])
    return cfg


# ------------------------------------------------------------------ Flash (CosyVoice3 engine)
FLASH_VOICE = {"anchor_male": "anchor", "waxal_ibo_5_female": "waxal_ibo_5", "waxal_yor_5_female": "waxal_yor_5"}


def canonical_voice(voice, voices_json=None):
    try:
        al = json.load(open(voices_json, encoding="utf-8")).get("aliases", {}) if voices_json else {}
    except OSError:
        al = {}
    return al.get(voice, voice)


class FlashTTS:
    """--version flash: CosyVoice3 base + Flash-DPO LM (flash_lm.Synth) behind the same say()/sr interface."""

    def __init__(self, a, cfg):
        if not torch.cuda.is_available():
            raise SystemExit("--version flash needs a CUDA GPU (phase2/phase3 also run on CPU)")
        cv3 = getattr(a, "cv3_dir", "") or os.environ.get("CV3_DIR") or os.path.join(HERE, "pretrained", "Fun-CosyVoice3-0.5B-2512")
        if not os.path.isdir(cv3):
            from huggingface_hub import snapshot_download
            snapshot_download("FunAudioLLM/Fun-CosyVoice3-0.5B-2512", local_dir=cv3)
        os.environ["CV3_DIR"] = cv3
        os.environ.setdefault("COSYVOICE_REPO", a.cosyvoice_dir)
        os.environ.setdefault("DA_PROMPTS", os.path.join(HERE, "prompts"))
        import flash_lm
        self.S = flash_lm.Synth(a.lm or cfg["lm"])
        self.cosy = self.tok = None                  # not shareable with the CosyVoice2 versions
        self.m = self.S.m
        self.sr = self.S.cosy.sample_rate
        r = a.rate if a.rate != "default" else cfg.get("rate", "normal")
        self.rate = RATES.index(None if r == "none" else r) if (None if r == "none" else r) in RATES else 0
        self.fe_mode = getattr(a, "frontend", "codeswitch")
        self.fe_lang = getattr(a, "frontend_lang", "eng")
        if self.fe_mode == "translate" and not getattr(a, "experimental_translate", False):
            raise SystemExit("--frontend translate is EXPERIMENTAL and off by default (MADLAD-400 loops on long inputs and is "
                             "weak on Pidgin). Add --experimental_translate to use it anyway.")
        from frontend import FrontEnd
        diac = cfg.get("diacritics") if not getattr(a, "no_diacritics", False) else None
        self.fe = FrontEnd(diac if diac and os.path.isdir(diac) else None, mt=self.fe_mode == "translate")
        self.set_voice(a.voice, a.voices)

    def set_voice(self, voice, voices_json=None):
        v = canonical_voice(voice, voices_json)
        if v not in FLASH_VOICE:
            raise SystemExit(f"unknown voice '{voice}' for flash; available: {', '.join(FLASH_VOICE)}")
        self.voice = FLASH_VOICE[v]

    def frontend(self, text):
        return self.fe(text, self.fe_mode, self.fe_lang)

    def say(self, tagged_text, seed=0, return_info=False, **_):
        t = self.frontend(tagged_text)
        print(f"[flash front-end] {t}")
        rate = RATES[self.rate]
        ws = []
        for i, ch in enumerate(word_chunks(t)):
            toks = self.S.tokens(ch, self.voice, rate or "none", seed + 1000 * i)[0]
            ws.append(self.S.wav(toks, self.voice))
        w = np.concatenate(ws).astype(np.float32)
        return (w, None) if return_info else w


def word_chunks(text, max_words=60):
    """one pass per <= max_words, split only at sentence ends (a whole paragraph usually fits in one pass)."""
    sents = split_sentences(text)
    out, cur, n = [], [], 0
    for s in sents:
        w = len(re.sub(r"<[^>]+>", " ", s).split())
        if cur and n + w > max_words:
            out.append(" ".join(cur)); cur, n = [], 0
        cur.append(s); n += w
    if cur:
        out.append(" ".join(cur))
    return out or [text]


class DynamicAccent:
    def __init__(self, model_dir, cosyvoice_dir, lm_ckpt, flow_ckpt, prompts_json, dtype="auto", sent_pause=0.6, cap=0.6,
                 sampler="ras", one_pass=False, postprocess=True, rate=None, cosy=None, tok=None, copy_llm=False):
        from transformers import AutoTokenizer
        self.tok = tok or AutoTokenizer.from_pretrained(os.path.join(model_dir, "CosyVoice-BlankEN"))
        self.cosy = cosy or load_cosy(model_dir, cosyvoice_dir)
        self.sr = self.cosy.sample_rate
        self.flow_sd = torch.load(flow_ckpt, map_location=DEV, weights_only=False)["flow"]
        self.cosy.model.flow.load_state_dict(self.flow_sd)
        ck = torch.load(lm_ckpt, map_location=DEV, weights_only=False)
        self.spk2id = json.load(open(os.path.join(os.path.dirname(lm_ckpt), "spk2id.json")))
        # two versions resident at once (app.py) must not share LM weights: copy_llm=True gives each its own copy.
        # The CLI runs versions one after the other and reloads the full LM state in place (saves ~2 GB RAM).
        qlm = copy.deepcopy(self.cosy.model.llm) if copy_llm else self.cosy.model.llm
        self.m = DALM2(qlm, ck["n_spk"]).to(DEV).eval()
        miss = self.m.load_state_dict(ck["model"], strict=False)
        assert not [k for k in miss.missing_keys if not k.startswith("rate_emb")], miss
        del ck
        if DEV.type != "cuda":
            dtype = "fp32"
        elif dtype == "auto":
            dtype = "bf16" if torch.cuda.is_bf16_supported() else "fp16"
        self.dtype = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": None}[dtype]
        self.sent_pause, self.cap = sent_pause, cap
        self.sampler, self.one_pass, self.postprocess = sampler, one_pass, postprocess
        self.rate = RATES.index(rate) if rate in RATES else 0
        meta = json.load(open(prompts_json, encoding="utf-8"))
        root = os.path.dirname(os.path.dirname(os.path.abspath(prompts_json)))
        self.prompts = []
        for p in meta["prompts"]:
            ids, langs = text_tokens(self.tok, p["tagged"])
            lmp = (torch.tensor(ids).to(DEV), torch.tensor(langs).to(DEV), torch.tensor(p["s3_tokens"], dtype=torch.int64).to(DEV))
            self.prompts.append(dict(id=p["id"], lm=lmp, flow=prompt_feats(self.cosy, os.path.join(root, p["wav"])), spk=ANCHOR))
        self.voice = None

    def use_flow(self):
        self.cosy.model.flow.load_state_dict(self.flow_sd)       # two versions may share one CosyVoice2 instance

    def set_voice(self, voice, voices_json=None):
        """'anchor' or an id from voices.json: its clip is the flow + LM prompt (zero-shot)."""
        meta = json.load(open(voices_json, encoding="utf-8")) if voices_json else {"voices": {}}
        voice = meta.get("aliases", {}).get(voice, voice)
        if voice in (None, "anchor", "anchor_male"):
            self.voice = None; return
        if voice not in meta["voices"]:
            raise SystemExit(f"unknown voice '{voice}'; available: {', '.join(meta['voices'])}")
        V = meta["voices"][voice]
        root = os.path.dirname(os.path.dirname(os.path.abspath(voices_json)))
        ptok, pfeat, pemb = prompt_feats(self.cosy, os.path.join(root, V["wav"]))
        ids, langs = text_tokens(self.tok, V["tagged"])
        lmp = (torch.tensor(ids).to(DEV), torch.tensor(langs).to(DEV), ptok[0].long())
        self.voice = dict(id=voice, lm=lmp, flow=(ptok, pfeat, pemb), spk=V["spk"] if V["spk"] in self.spk2id else ANCHOR)

    def _prompt(self, prompt_k):
        return self.voice or self.prompts[prompt_k % len(self.prompts)]

    def one(self, tagged_text, seed, prompt_k, use_lm_prompt=True, lm_prompt=None):
        ids, langs = text_tokens(self.tok, tagged_text)
        P = self._prompt(prompt_k)
        with autocast(self.dtype):
            sp = self.m.generate2(torch.tensor(ids).to(DEV), torch.tensor(langs).to(DEV), torch.tensor(self.spk2id[P["spk"]]).to(DEV),
                                  prompt=(lm_prompt or P["lm"]) if use_lm_prompt else None, seed=seed, sampler=self.sampler,
                                  rate=self.rate, max_ratio=30.0 if self.one_pass else 20.0)
        if len(sp) < 3:
            sp = sp + [0] * 3
        self.last_tokens = sp
        return token2wav(self.cosy, sp, *P["flow"])

    def say(self, tagged_text, seed=0, prompt_k=0, stitch=True, pause_mode="model", speed=1.0, return_info=False):
        self.use_flow()
        if self.one_pass:                      # phase3: the model makes every pause; nothing is trimmed or inserted
            ws, prev = [], None
            for i, ch in enumerate(word_chunks(tagged_text)):
                w = self.one(ch, seed + 1000 * i, prompt_k, lm_prompt=prev)
                ws.append(w)
                ids, langs = text_tokens(self.tok, ch)       # next chunk continues from this one (model sets the gap)
                if len(self.last_tokens) <= 400:
                    prev = (torch.tensor(ids).to(DEV), torch.tensor(langs).to(DEV), torch.tensor(self.last_tokens).to(DEV))
            w = np.concatenate(ws)
            return (w, None) if return_info else w
        if pause_mode == "clause":
            return self.say_clause(tagged_text, seed, prompt_k, speed, return_info)
        w = self.say_model(tagged_text, seed, prompt_k, stitch)
        w = time_stretch(w, self.sr, speed)
        return (w, None) if return_info else w

    def say_clause(self, tagged_text, seed=0, prompt_k=0, speed=1.10, return_info=False):
        rng = np.random.default_rng(seed)
        chunks = clause_chunks(tagged_text)
        out, info, lead = [], [], 0.03
        for i, (l, t, b) in enumerate(chunks):
            w = self.one(f"<{l}>{t}</{l}>", seed + 1000 * i, prompt_k)
            w = cap_internal(trim_tight(w, self.sr), self.sr, self.cap)
            w = fade(time_stretch(w, self.sr, speed), self.sr)
            if i == 0:
                out.append(np.zeros(int(0.12 * self.sr), np.float32))
            out.append(w)
            if b is not None:
                g = GAPS[b] + rng.uniform(-JITTER[b], JITTER[b])
                out.append(np.zeros(max(0, int((g - 2 * lead) * self.sr)), np.float32))
            else:
                out.append(np.zeros(int(0.15 * self.sr), np.float32))
            info.append(dict(lang=l, text=t, boundary=b))
        w = np.concatenate(out)
        return (w, info) if return_info else w

    def say_model(self, tagged_text, seed=0, prompt_k=0, stitch=True):
        sents = split_sentences(tagged_text) if stitch else [tagged_text]
        ws = [cap_internal(trim_edges(self.one(s, seed + 1000 * i, prompt_k), self.sr), self.sr, self.cap)
              for i, s in enumerate(sents)]
        if len(ws) == 1:
            return ws[0]
        out = []
        for i, w in enumerate(ws):
            if i:
                out.append(np.zeros(max(0, int((self.sent_pause - 0.30) * self.sr)), np.float32))
            out.append(w)
        return np.concatenate(out)


def build(version, a, cosy=None, tok=None, copy_llm=False):
    cfg = version_cfg(version)
    if cfg.get("engine") == "cosyvoice3":
        return FlashTTS(a, cfg)
    lm = a.lm or cfg["lm"]; flow = a.flow or cfg["flow"]
    S = DynamicAccent(a.model_dir, a.cosyvoice_dir, lm, flow, a.prompts, a.dtype, sampler=cfg["sampler"],
                      one_pass=cfg["one_pass"] and not a.stitch, postprocess=cfg["postprocess"],
                      rate=a.rate if a.rate != "default" else cfg["rate"], cosy=cosy, tok=tok, copy_llm=copy_llm)
    S.set_voice(a.voice, a.voices)
    return S


def list_voices(voices_json):
    V = json.load(open(voices_json, encoding="utf-8"))["voices"] if os.path.exists(voices_json) else {}
    print("voices (--voice):")
    for k, v in V.items():
        print(f"  {k:20s} {v.get('desc', '')}")
    print("\nversions (--version):")
    for k, v in SETTINGS.get("versions", {}).items():
        print(f"  {k:20s} {v.get('desc', '')}")
    if RECOMMENDED:
        print("\nrecommended pairing: " + "; ".join(f"{k} -> {', '.join(v)}" for k, v in RECOMMENDED.items() if isinstance(v, list)))
        print("  " + RECOMMENDED.get("note", ""))


def prepare_text(a):
    """apply --auto_tag / --force_lang / --lang_map; returns the tagged text that will be synthesised."""
    from tagger import Tagger, parse_map, TAGGED
    text = a.text
    if a.force_lang:
        from tagger import strip_tags
        return f"<{a.force_lang}>{strip_tags(text)}</{a.force_lang}>"
    if a.auto_tag or a.lang_map:
        if TAGGED.search(text) and not a.auto_tag:
            return text
        from tagger import strip_tags
        T = Tagger(None if a.lid == "none" else a.lid, parse_map(a.lang_map))
        return T.tag(strip_tags(text))
    return text


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--text", help="text; tagged ('<eng>Good morning,</eng> <pcm>how body?</pcm>') or plain with --auto_tag")
    ap.add_argument("--version", default=SETTINGS.get("default_version", "phase3"), choices=sorted(VERSIONS),
                    help=f"model version (default {SETTINGS.get('default_version', 'phase3')})")
    ap.add_argument("--compare", action="store_true",
                    help="render the same text, seed and voice with every version in --compare_versions -> <out>_<version>.wav")
    ap.add_argument("--compare_versions", default="phase2,phase3,flash", help="versions for --compare (comma-separated)")
    ap.add_argument("--voice", default=SETTINGS.get("default_voice", "anchor_male"),
                    help="anchor_male | waxal_ibo_5_female | waxal_yor_5_female (see --list_voices)")
    ap.add_argument("--list_voices", action="store_true", help="list voices, versions and the recommended pairing")
    ap.add_argument("--rate", default="default", choices=["default", "none", "very slow", "slow", "normal", "fast", "very fast"],
                    help="phase3/flash speed instruction to the LM (default: the version's setting, 'normal'; 'none' = no instruction)")
    # flash text front-end
    ap.add_argument("--frontend", default="codeswitch", choices=["codeswitch", "native", "accent", "translate"],
                    help="flash only: codeswitch (tagged/one-language text), native (whole text in --frontend_lang), accent "
                         "(words as written, read with the --frontend_lang accent), translate (EXPERIMENTAL, off by default)")
    ap.add_argument("--frontend_lang", default="eng", choices=LANGS, help="flash front-end language for untagged text / native / accent / translate target")
    ap.add_argument("--no_diacritics", action="store_true", help="flash: do not restore missing Yoruba/Igbo diacritics")
    ap.add_argument("--experimental_translate", action="store_true",
                    help="flash: allow --frontend translate (MADLAD-400-3B-MT, ~12 GB download; loops on long inputs, weak on Pidgin)")
    ap.add_argument("--cv3_dir", default=os.environ.get("CV3_DIR", os.path.join(HERE, "pretrained", "Fun-CosyVoice3-0.5B-2512")),
                    help="flash: stock FunAudioLLM/Fun-CosyVoice3-0.5B-2512 snapshot (downloaded on first run)")
    # code-switching control
    ap.add_argument("--auto_tag", action="store_true",
                    help="tag untagged mixed text automatically (word/span language ID), print it, then synthesise")
    ap.add_argument("--tag_only", action="store_true", help="with --auto_tag: print the tagged text and exit (edit it, then pass it back as --text)")
    ap.add_argument("--force_lang", choices=LANGS, help="read the whole text in one language (tags are ignored)")
    ap.add_argument("--lang_map", default="", help='per-word language overrides for the tagger, e.g. "omo=pcm,sha=pcm"')
    ap.add_argument("--lid", default="none", choices=["none", "glotlid"],
                    help="optional LID model for --auto_tag (glotlid: cis-lmu/glotlid, Apache-2.0, 1.7 GB, pip install fasttext)")
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--out", default="out.wav")
    ap.add_argument("--prompt_clip", type=int, default=0, help="anchor prompt clip 0-3 (final systems: 0)")
    ap.add_argument("--model_dir", default=os.environ.get("COSYVOICE2_DIR", os.path.join(HERE, "pretrained", "CosyVoice2-0.5B")),
                    help="stock FunAudioLLM/CosyVoice2-0.5B snapshot (downloaded on first run)")
    ap.add_argument("--cosyvoice_dir", default=os.environ.get("COSYVOICE_REPO", os.path.join(HERE, "CosyVoice")),
                    help="clone of github.com/FunAudioLLM/CosyVoice (with submodules)")
    ap.add_argument("--lm", default="", help="override the version's LM checkpoint")
    ap.add_argument("--flow", default="", help="override the version's flow checkpoint")
    ap.add_argument("--prompts", default=os.path.join(HERE, "prompts", "prompts.json"))
    ap.add_argument("--voices", default=os.path.join(HERE, "prompts", "voices.json"))
    ap.add_argument("--dtype", default="auto", choices=["auto", "bf16", "fp16", "fp32"],
                    help="LM autocast dtype on GPU (samples: bf16; auto picks fp16 on GPUs without bf16). CPU always runs fp32")
    ap.add_argument("--stitch", action="store_true", help="phase3: generate sentence by sentence instead of one pass")
    ap.add_argument("--no_stitch", action="store_true", help="phase2: generate multi-sentence text in one pass")
    ap.add_argument("--pause_mode", default="model", choices=["clause", "model"], help="phase2 only; clause = experimental prosody v2 (not recommended)")
    ap.add_argument("--speed", type=float, default=DEFAULT_SPEED, help="phase2 only; pitch-preserving tempo factor (1.0 = as trained; not recommended)")
    a = ap.parse_args()
    if a.list_voices:
        list_voices(a.voices); return
    if not a.text:
        ap.error("--text is required")
    text = prepare_text(a)
    if text != a.text or a.auto_tag:
        print(f"[tagged] {text}")
    if a.tag_only:
        return
    versions = [v.strip() for v in a.compare_versions.split(",") if v.strip()] if a.compare else [a.version]
    bad = [v for v in versions if v not in VERSIONS]
    if bad:
        ap.error(f"unknown version(s) {bad}; available: {sorted(VERSIONS)}")
    if any(VERSIONS[v].get("engine") != "cosyvoice3" for v in versions) and not os.path.isdir(a.model_dir):
        from huggingface_hub import snapshot_download
        snapshot_download("FunAudioLLM/CosyVoice2-0.5B", local_dir=a.model_dir)
    cosy = tok = None
    base, ext = os.path.splitext(a.out)
    for v in versions:
        hint = recommendation(v, json.load(open(a.voices)).get("aliases", {}).get(a.voice, a.voice) if os.path.exists(a.voices) else a.voice)
        if hint:
            print(hint)
        S = build(v, a, cosy, tok)
        if S.cosy is not None:
            cosy, tok = S.cosy, S.tok
        w = S.say(text, seed=a.seed, prompt_k=a.prompt_clip, stitch=not a.no_stitch, pause_mode=a.pause_mode,
                  speed=a.speed if v == "phase2" else 1.0)
        out = f"{base}_{v}{ext or '.wav'}" if a.compare else a.out
        sf.write(out, w, S.sr)
        print(f"[{v}] wrote {out} ({len(w) / S.sr:.2f} s, {S.sr} Hz, voice={a.voice}, device={DEV.type})")
        del S.m
        if isinstance(S, FlashTTS):
            del S.S
        del S; torch.cuda.empty_cache()


if __name__ == "__main__":
    main()