Dynamic Accent Flash: add Flash-DPO (CosyVoice3) as --version flash (default), diacritic front-end, Flash results, samples
e4ae2bb verified | #!/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) | |
| 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) | |
| 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() | |