Dynamic Accent Flash: add Flash-DPO (CosyVoice3) as --version flash (default), diacritic front-end, Flash results, samples
e4ae2bb verified | """Dynamic Accent Flash: CosyVoice3 (Fun-CosyVoice3-0.5B-2512, Apache-2.0) LM fine-tune + synthesis. | |
| Same additions as Phase 2/3 (per-token language embedding, learned speaker token after the text, rate embedding added to | |
| the speaker vector, GRL speaker classifier on text positions, in-context same-speaker prompt p=0.5), ported to the CV3 LM: | |
| seq = [sos, "You are a helpful assistant.<|endofprompt|>", (prompt text) text, spk(+rate), task, (prompt speech) speech, eos] | |
| sos/task live in speech_embedding (CV3), eos = speech_token_size+1. Targets are CV3 speech tokens (s3v3.*). | |
| Flow (DiT) + HiFT = CV3 base, zero-shot prompted with the voice clip (no flow fine-tune). | |
| python flash_lm.py lm --out ~/flash/runs/flash_lm --max_steps 12000 --minutes 75 | |
| python flash_lm.py gen --lm CKPT --out DIR [--voice anchor|waxal_ibo_5|waxal_yor_5] [--rate normal] [--rows ROWS] | |
| Release note: in this repo use `infer.py --version flash` (it wraps Synth below). Training needs the Phase 3 data | |
| manifest, which is not shipped; the training code is included for reference. Paths come from COSYVOICE_REPO (CosyVoice | |
| clone with CosyVoice3 support) and CV3_DIR (Fun-CosyVoice3-0.5B-2512 snapshot). | |
| """ | |
| import argparse, json, math, os, random, sys, time, inspect | |
| import numpy as np, torch, torch.nn as nn, torch.nn.functional as F, soundfile as sf, librosa | |
| HOME = os.path.expanduser("~") | |
| _HERE = os.path.dirname(os.path.abspath(__file__)) | |
| _CV = os.environ.get("COSYVOICE_REPO") or (f"{_HERE}/CosyVoice" if os.path.isdir(f"{_HERE}/CosyVoice") else f"{HOME}/CosyVoice") | |
| sys.path[:0] = [_CV, f"{_CV}/third_party/Matcha-TTS", f"{_HERE}/code", f"{HOME}/da"] | |
| from da.text import segments # noqa: E402 | |
| from da.model import grad_reverse # noqa: E402 | |
| CV3 = os.environ.get("CV3_DIR") or (f"{_HERE}/pretrained/Fun-CosyVoice3-0.5B-2512" if os.path.isdir(f"{_HERE}/pretrained/Fun-CosyVoice3-0.5B-2512") else f"{HOME}/models/Fun-CosyVoice3-0.5B-2512") | |
| LANGS = ["eng", "ibo", "hau", "yor", "pcm"] | |
| L2I = {l: i for i, l in enumerate(LANGS)} | |
| RATES = [None, "very slow", "slow", "normal", "fast", "very fast"] | |
| R2I = {r: i for i, r in enumerate(RATES)} | |
| TAGS = ("[pause]", "[breath]", "[quick_breath]", "[laughter]") | |
| SIL = {1, 2, 28, 29, 55, 248, 494, 2241, 2242, 2322, 2323} # CV3 FSQ silent/breath tokens (cosyvoice/cli/model.py) | |
| ANCHOR = "waxal_yor_4" | |
| NEW = ["waxal", "igboapi", "open_bible_igbo", "naijas2st_igbo", "igbosyncorp"] | |
| MAN = f"{HOME}/p3t/manifest_p3.jsonl" | |
| PDIR = os.environ.get("DA_PROMPTS") or (f"{_HERE}/prompts" if os.path.isdir(f"{_HERE}/prompts") else f"{HOME}/hf_test/prompts") | |
| VOICES = { | |
| "anchor": dict(wav=f"{PDIR}/clip0_waxal_yor_yor_1258.wav" if os.path.exists(f"{PDIR}/clip0_waxal_yor_yor_1258.wav") | |
| else f"{PDIR}/waxal_yor_yor_1258.wav", spk=ANCHOR), | |
| "waxal_ibo_5": dict(wav=f"{PDIR}/voice_waxal_ibo_5.wav", spk="waxal_ibo_5"), | |
| "waxal_yor_5": dict(wav=f"{PDIR}/voice_waxal_yor_5.wav", spk="waxal_yor_5"), | |
| } | |
| # Flash source re-weighting on top of the Phase 3 manifest weight: less Bible narration and slow phone/dictionary reads | |
| FLASH_W = {"open_bible_igbo": 0.4, "bible_hau": 0.5, "bible_yor": 0.5, "igboapi": 0.6, "naijas2st_igbo": 0.6, | |
| "igbosyn_pseudo": 1.0} | |
| def strip_tags(t): | |
| for g in TAGS: | |
| t = t.replace(g, " ") | |
| return " ".join(t.split()).replace(" ,", ",").replace(" .", ".") | |
| 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_cv3(): | |
| 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 CosyVoice3 | |
| kw = dict(load_jit=False, load_trt=False, fp16=False) | |
| sig = inspect.signature(CosyVoice3.__init__).parameters | |
| return CosyVoice3(CV3, **{k: v for k, v in kw.items() if k in sig}) | |
| class Tok: | |
| def __init__(self, cosy): | |
| self.t = cosy.frontend.tokenizer | |
| self.prefix = self.t.encode("You are a helpful assistant.") + [151646] | |
| self.cache = {} | |
| def __call__(self, text, default="eng"): | |
| if text in self.cache: | |
| return self.cache[text] | |
| ids, langs = [], [] | |
| for l, s in segments(text, default): | |
| t = self.t.encode(s) | |
| ids += t; langs += [L2I[l]] * len(t) | |
| self.cache[text] = (ids, langs) | |
| return ids, langs | |
| class FLM(nn.Module): | |
| def __init__(self, q, n_spk, prefix, d=896): | |
| super().__init__() | |
| self.q = q | |
| self.lang_emb = nn.Embedding(len(LANGS), d); nn.init.zeros_(self.lang_emb.weight) | |
| self.spk_emb = nn.Embedding(n_spk, d); nn.init.normal_(self.spk_emb.weight, std=0.02) | |
| self.rate_emb = nn.Embedding(len(RATES), d, padding_idx=0); nn.init.zeros_(self.rate_emb.weight) | |
| self.adv = nn.Sequential(nn.Linear(d, 256), nn.ReLU(), nn.Linear(256, n_spk)) | |
| self.V = q.speech_token_size | |
| self.EOS = q.eos_token | |
| self.register_buffer("prefix", torch.tensor(prefix, dtype=torch.long), persistent=False) | |
| def embed_text(self, ids, langs): | |
| return self.q.llm.model.model.embed_tokens(ids) + self.lang_emb(langs) | |
| def spkvec(self, spk): | |
| return self.spk_emb(spk[0]) + self.rate_emb(spk[1]) | |
| def build(self, ids, langs, spk, sp=None, pids=None, plangs=None, psp=None): | |
| E = self.q.speech_embedding.weight | |
| if pids is not None: | |
| ids = torch.cat([pids, ids]); langs = torch.cat([plangs, langs]) | |
| pre = self.q.llm.model.model.embed_tokens(self.prefix.to(ids.device)) | |
| parts = [E[self.q.sos][None], pre, self.embed_text(ids, langs), self.spkvec(spk)[None], E[self.q.task_id][None]] | |
| npsp = 0 | |
| if psp is not None: | |
| parts.append(self.q.speech_embedding(psp)); npsp = len(psp) | |
| if sp is not None: | |
| parts.append(self.q.speech_embedding(sp)) | |
| return torch.cat(parts, 0), len(pre), len(ids), npsp | |
| def hidden(self, x, mask): | |
| return self.q.llm.model.model(inputs_embeds=x, attention_mask=mask, return_dict=True).last_hidden_state | |
| def pack(self, batch): | |
| seqs, tgts, tpos, spks = [], [], [], [] | |
| for b in batch: | |
| ids, langs, spk, sp = b[:4] | |
| pr = b[4] if len(b) > 4 else None | |
| x, P, nt, npsp = self.build(ids, langs, spk, sp, *(pr if pr else (None, None, None))) | |
| tgt = torch.full((x.shape[0] + 1,), -100, device=x.device, dtype=torch.long) | |
| start = P + nt + 2 + npsp | |
| tgt[start: start + len(sp)] = sp | |
| tgt[start + len(sp)] = self.EOS | |
| seqs.append(x); tgts.append(tgt[:x.shape[0]]); tpos.append((1 + P, 1 + P + nt)); spks.append(int(spk[0])) | |
| Tm = max(s.shape[0] for s in seqs); D = seqs[0].shape[1] | |
| X = seqs[0].new_zeros(len(seqs), Tm, D) | |
| Y = torch.full((len(seqs), Tm), -100, device=X.device, dtype=torch.long) | |
| M = torch.zeros(len(seqs), Tm, device=X.device, dtype=torch.long) | |
| for i, (s, t) in enumerate(zip(seqs, tgts)): | |
| X[i, :len(s)] = s; Y[i, :len(t)] = t; M[i, :len(s)] = 1 | |
| return X, Y, M, tpos, spks | |
| def forward2(self, batch, adv_lambda=1.0): | |
| X, Y, M, tpos, spks = self.pack(batch) | |
| H = self.hidden(X, M) | |
| logits = self.q.llm_decoder(H) | |
| ce = F.cross_entropy(logits.float().reshape(-1, logits.shape[-1]), Y.reshape(-1), ignore_index=-100) | |
| valid = Y != -100 | |
| acc = (logits.argmax(-1)[valid] == Y[valid]).float().mean() | |
| hs = torch.cat([H[i, a:b] for i, (a, b) in enumerate(tpos)]) | |
| ys = torch.cat([torch.full((b - a,), spks[i], device=X.device) for i, (a, b) in enumerate(tpos)]) | |
| al = self.adv(grad_reverse(hs.float(), adv_lambda)) | |
| return ce, acc, F.cross_entropy(al, ys), (al.argmax(-1) == ys).float().mean() | |
| def seq_logp(self, batch): | |
| """sum log p(speech tokens + eos) per item (for DPO).""" | |
| X, Y, M, _, _ = self.pack(batch) | |
| logits = self.q.llm_decoder(self.hidden(X, M)).float() | |
| lp = logits.log_softmax(-1) | |
| valid = Y != -100 | |
| g = lp.gather(-1, Y.clamp(min=0)[..., None])[..., 0] * valid | |
| return g.sum(-1), valid.sum(-1) | |
| def generate(self, ids, langs, spk, prompt=None, B=1, top_p=0.8, top_k=25, win=10, tau_r=0.1, seed=0, | |
| min_ratio=2.0, max_ratio=30.0, sampler="silexcap", max_tok=900): | |
| g = torch.Generator(device=ids.device).manual_seed(seed) | |
| x, _, _, _ = self.build(ids, langs, spk, None, *(prompt if prompt else (None, None, None))) | |
| inp = x[None].repeat(B, 1, 1) | |
| n = len(ids) | |
| min_len, max_len = int(n * min_ratio), min(int(n * max_ratio), max_tok) | |
| past, outs, done = None, [[] for _ in range(B)], [False] * B | |
| E = self.q.speech_embedding.weight | |
| 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) | |
| eos_lp = logp[:, self.EOS].clone() | |
| logp[:, self.V:] = -float("inf") | |
| if i >= min_len: | |
| logp[:, self.EOS] = eos_lp | |
| P = logp.exp() | |
| nxt = [] | |
| for b in range(B): | |
| if done[b]: | |
| nxt.append(0); continue | |
| p = P[b]; out = outs[b] | |
| sp_, si = p.sort(descending=True) | |
| keep = (sp_.cumsum(0) - sp_) < top_p | |
| keep[top_k:] = False | |
| cp, ci = sp_[keep], si[keep] | |
| t = int(ci[torch.multinomial(cp / cp.sum(), 1, generator=g)]) | |
| exempt = False | |
| if sampler == "silexcap" and t in SIL: | |
| run = 0 | |
| for u in reversed(out): | |
| if u not in SIL: | |
| break | |
| run += 1 | |
| exempt = run < 25 | |
| if not exempt and out and t != self.EOS and sum(1 for u in out[-win:] if u == t) >= win * tau_r: | |
| t = int(torch.multinomial(p, 1, generator=g)) | |
| if t >= self.V: | |
| done[b] = True; nxt.append(0); continue | |
| out.append(t); nxt.append(t) | |
| if all(done): | |
| break | |
| inp = E[torch.tensor(nxt, device=ids.device)][:, None] | |
| return outs | |
| def load_model(cosy, ckpt, tok): | |
| ck = torch.load(ckpt, map_location="cuda", weights_only=False) | |
| m = FLM(cosy.model.llm, ck["n_spk"], tok.prefix).cuda().eval() | |
| print("load", ckpt, m.load_state_dict(ck["model"], strict=False), flush=True) | |
| return m, json.load(open(os.path.join(os.path.dirname(ckpt), "spk2id.json"))) | |
| # ------------------------------------------------------------------ data | |
| def load_data(splits, max_tok=750, extra_man=None, qual=None, qual_drop=0.1): | |
| rows = [json.loads(l) for l in open(MAN, encoding="utf-8")] | |
| if extra_man and os.path.exists(extra_man): | |
| rows += [json.loads(l) for l in open(extra_man, encoding="utf-8")] | |
| want = {r["id"] for r in rows if r["split"] in splits} | |
| drop = set() | |
| if qual and os.path.exists(qual) and "train" in splits: # drop the bottom qual_drop of UTMOS per source | |
| Q = [json.loads(l) for l in open(qual)] | |
| by = {} | |
| for q in Q: | |
| by.setdefault(q["src"], []).append(q) | |
| for src, L in by.items(): | |
| if src in ("cs_pairs",): | |
| continue | |
| thr = float(np.quantile([q["u22"] for q in L], qual_drop)) | |
| drop |= {q["id"] for q in L if q["u22"] < thr} | |
| print("quality filter drops", len(drop), flush=True) | |
| src_idx, mm = {}, {} | |
| for src in sorted({r["src"] for r in rows}): | |
| if src in NEW: | |
| d = f"{HOME}/p3t/{src}" | |
| elif src == "igbosyn_pseudo": | |
| d = f"{HOME}/flash/data/igbosyn_pseudo" | |
| else: | |
| d = f"{HOME}/data3/{src}" | |
| if not os.path.exists(f"{d}/s3v3.done"): | |
| print("missing v3 tokens", d, flush=True); continue | |
| meta = [json.loads(l) for l in open(f"{d}/meta.jsonl", encoding="utf-8")] | |
| tk = [json.loads(l) for l in open(f"{d}/s3v3.jsonl")] | |
| ix = {} | |
| for k, (m, t) in enumerate(zip(meta, tk)): | |
| assert m["id"] == t["id"] | |
| if m["id"] in want: | |
| ix[m["id"]] = (t["toff"], t["ntok"]) | |
| src_idx[src] = ix | |
| mm[src] = np.memmap(f"{d}/s3v3.i32", dtype=np.int32, mode="r") | |
| items = [] | |
| for r in rows: | |
| if r["split"] not in splits or r["src"] not in src_idx or r["id"] in drop and r["speaker"] != ANCHOR: | |
| continue | |
| ix = src_idx[r["src"]] | |
| if r["id"] not in ix: | |
| continue | |
| toff, ntok = ix[r["id"]] | |
| if not (5 < ntok <= max_tok): | |
| continue | |
| w = float(r.get("weight", 1.0)) * FLASH_W.get(r["src"], 1.0) | |
| items.append(dict(id=r["id"], ds=r["src"], src=r["src"], lang=r["lang"], spk=r["speaker"], split=r["split"], | |
| dur=r["dur"], text=r["text_plain"], tags=r.get("text_tags"), rate=R2I.get(r.get("rate_bin"), 0), | |
| w=w, toff=toff, ntok=ntok)) | |
| return items, mm | |
| def tagged(text, lang): | |
| return text if lang == "cs" else f"<{lang}>{text}</{lang}>" | |
| def batches(items, budget, rng, wfn, cost): | |
| idx = [] | |
| for i, m in enumerate(items): | |
| w = wfn(m) | |
| idx += [i] * int(w) + ([i] if rng.random() < w - int(w) else []) | |
| idx.sort(key=lambda i: cost(items[i]) + rng.random() * 10) | |
| out, cur, mx = [], [], 0 | |
| for i in idx: | |
| L = cost(items[i]) | |
| if cur and max(mx, L) * (len(cur) + 1) > budget: | |
| out.append(cur); cur, mx = [], 0 | |
| cur.append(i); mx = max(mx, L) | |
| if cur: | |
| out.append(cur) | |
| rng.shuffle(out) | |
| return out | |
| VAL_SPLITS = {"val", "validation", "dev"} | |
| def train_lm(a): | |
| cosy = load_cv3(); tok = Tok(cosy) | |
| items, mm = load_data({"train"}, a.max_tok, a.extra_man, a.qual, a.qual_drop) | |
| val_all, mmv = load_data(VAL_SPLITS, a.max_tok) | |
| mm.update(mmv) | |
| spk2id = {} | |
| for s in sorted({m["spk"] for m in items}): | |
| spk2id[s] = len(spk2id) | |
| n_spk = len(spk2id) | |
| rng = random.Random(1); rng.shuffle(val_all) | |
| val, per_l, per_s = [], {}, {} | |
| for m in val_all: | |
| if m["spk"] not in spk2id or m["lang"] == "cs": | |
| continue | |
| L = "anchor" if m["spk"] == ANCHOR else m["lang"] | |
| if per_l.get(L, 0) >= 60 or per_s.get(m["spk"], 0) >= 15: | |
| continue | |
| per_l[L] = per_l.get(L, 0) + 1; per_s[m["spk"]] = per_s.get(m["spk"], 0) + 1; val.append(m) | |
| by_spk = {} | |
| for m in items: | |
| if m["ntok"] <= 200 and m["lang"] != "cs": | |
| by_spk.setdefault(m["spk"], []).append(m) | |
| os.makedirs(a.out, exist_ok=True) | |
| json.dump(spk2id, open(f"{a.out}/spk2id.json", "w")); json.dump(vars(a), open(f"{a.out}/args.json", "w")) | |
| hrs = {} | |
| for m in items: | |
| hrs[m["src"]] = hrs.get(m["src"], 0) + m["dur"] / 3600 | |
| print("lm items", len(items), "hours", round(sum(hrs.values()), 1), {k: round(v, 1) for k, v in hrs.items()}, | |
| "spk", n_spk, "val", len(val), per_l, flush=True) | |
| model = FLM(cosy.model.llm, n_spk, tok.prefix).cuda().train() | |
| for p in list(cosy.model.flow.parameters()) + list(cosy.model.hift.parameters()): | |
| p.requires_grad = False | |
| new = list(model.lang_emb.parameters()) + list(model.spk_emb.parameters()) + list(model.adv.parameters()) + list(model.rate_emb.parameters()) | |
| newids = {id(p) for p in new} | |
| old = [p for p in model.parameters() if id(p) not in newids and p.requires_grad] | |
| opt = torch.optim.AdamW([{"params": old, "lr": a.lr}, {"params": new, "lr": a.lr_new}], weight_decay=0.01) | |
| base = [a.lr, a.lr_new] | |
| def text_of(m, r): | |
| if m["tags"] and r is not None and r.random() >= a.tag_drop: | |
| return m["tags"] | |
| return strip_tags(m["text"]) | |
| def spk_t(m, r): | |
| rate = m["rate"] if (r is not None and r.random() >= a.rate_drop) else 0 | |
| return torch.tensor([spk2id[m["spk"]], rate]).cuda() | |
| def enc(m, prng=None, force_prompt=None, val_mode=False): | |
| r = None if val_mode else prng | |
| ids, langs = tok(tagged(text_of(m, r), m["lang"])) | |
| sp = np.asarray(mm[m["ds"]][m["toff"]: m["toff"] + m["ntok"]], dtype=np.int64) | |
| out = [torch.tensor(ids).cuda(), torch.tensor(langs).cuda(), spk_t(m, r), torch.from_numpy(sp).cuda()] | |
| use = force_prompt if force_prompt is not None else (prng is not None and prng.random() < a.p_prompt) | |
| if use: | |
| c = by_spk.get(m["spk"], []) | |
| c = [x for x in (prng or random).sample(c, min(6, len(c))) if x["id"] != m["id"]] if c else [] | |
| if c and c[0]["ntok"] + m["ntok"] <= a.max_tok + 150: | |
| p = c[0] | |
| pids, plangs = tok(tagged(strip_tags(p["text"]), p["lang"])) | |
| psp = np.asarray(mm[p["ds"]][p["toff"]: p["toff"] + p["ntok"]], dtype=np.int64) | |
| out.append((torch.tensor(pids).cuda(), torch.tensor(plangs).cuda(), torch.from_numpy(psp).cuda())) | |
| return out | |
| def validate(): | |
| model.eval(); res = {} | |
| with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16): | |
| for mode in (False, True): | |
| vr = random.Random(7); L, W = [], [] | |
| for i in range(0, len(val), 16): | |
| bt = val[i:i + 16] | |
| ce, _, _, _ = model.forward2([enc(m, vr, force_prompt=mode, val_mode=True) for m in bt]) | |
| L.append(float(ce) * len(bt)); W.append(len(bt)) | |
| res["p" if mode else "np"] = sum(L) / sum(W) | |
| model.train() | |
| res["mean"] = (res["np"] + res["p"]) / 2 | |
| return res | |
| wfn = lambda m: m["w"] * (a.anchor_w if m["spk"] == ANCHOR else 1.0) # noqa | |
| cost = lambda m: (lambda L: int(L * (1 + L / a.quad)))(m["ntok"] + len(m["text"]) // 2 + 20 + int(a.p_prompt * 200)) # noqa | |
| rng = random.Random(0); prng = random.Random(3) | |
| step, nskip, noom, t0 = 0, 0, 0, time.time() | |
| logf = open(f"{a.out}/log.jsonl", "a") | |
| v = validate(); best = v["mean"]; bad = 0 | |
| print(json.dumps({"step": 0, "val": v}), flush=True); logf.write(json.dumps({"step": 0, "val": v}) + "\n") | |
| total, done = a.max_steps, False | |
| while not done: | |
| for bi in batches(items, a.tokens, rng, wfn, cost): | |
| lrm = min(1.0, (step + 1) / a.warmup) * (0.1 + 0.9 * 0.5 * (1 + math.cos(math.pi * min(1.0, step / total)))) | |
| for gp, b in zip(opt.param_groups, base): | |
| gp["lr"] = b * lrm | |
| batch = [enc(items[i], prng) for i in bi] | |
| opt.zero_grad(set_to_none=True) | |
| oom = False | |
| try: | |
| with torch.autocast("cuda", dtype=torch.bfloat16): | |
| ce, acc, adv, adv_acc = model.forward2(batch, adv_lambda=min(1.0, step / 1000)) | |
| loss = ce + a.adv_w * adv | |
| if not torch.isfinite(loss): | |
| nskip += 1; continue | |
| loss.backward() | |
| except torch.cuda.OutOfMemoryError: | |
| oom = True | |
| if oom: | |
| batch = ce = acc = adv = adv_acc = loss = None | |
| opt.zero_grad(set_to_none=True); torch.cuda.empty_cache(); noom += 1; continue | |
| gn = torch.nn.utils.clip_grad_norm_([p for p in model.parameters() if p.requires_grad], 1.0) | |
| if not torch.isfinite(gn): | |
| nskip += 1; opt.zero_grad(set_to_none=True); continue | |
| opt.step(); step += 1 | |
| if step % 50 == 0: | |
| r = dict(step=step, ce=round(float(ce), 4), acc=round(float(acc), 4), adv=round(float(adv), 3), | |
| adv_acc=round(float(adv_acc), 3), gn=round(float(gn), 2), B=len(bi), lr=round(lrm * a.lr, 9), | |
| skips=nskip, oom=noom, min=round((time.time() - t0) / 60, 2)) | |
| print(json.dumps(r), flush=True); logf.write(json.dumps(r) + "\n"); logf.flush() | |
| if step % a.val_every == 0: | |
| v = validate(); imp = v["mean"] < best - 1e-4 | |
| rec = {"step": step, "val": v, "best": imp} | |
| print(json.dumps(rec), flush=True); logf.write(json.dumps(rec) + "\n"); logf.flush() | |
| if imp: | |
| best, bad = v["mean"], 0 | |
| torch.save(dict(model={k: t for k, t in model.state_dict().items()}, step=step, n_spk=n_spk, val=v), f"{a.out}/best.pt") | |
| else: | |
| bad += 1 | |
| if bad >= a.patience: | |
| print("EARLY_STOP", step, flush=True); done = True; break | |
| if (time.time() - t0) / 60 > a.minutes or step >= total: | |
| done = True; break | |
| print("TRAIN_LM_DONE", step, "best", best, "skips", nskip, "oom", noom, flush=True) | |
| # ------------------------------------------------------------------ synthesis | |
| def prompt_feats(cosy, wav): | |
| fe = cosy.frontend | |
| tokn, _ = fe._extract_speech_token(wav) | |
| feat, _ = fe._extract_speech_feat(wav) | |
| n = min(feat.shape[1] // 2, tokn.shape[1]) | |
| return tokn[:, :n].cuda(), feat[:, :2 * n].cuda(), fe._extract_spk_embedding(wav).cuda() | |
| def token2wav(cosy, tokens, ptok, pfeat, emb): | |
| flow = cosy.model.flow | |
| t = torch.tensor([tokens], dtype=torch.int32).cuda() | |
| mel, _ = flow.inference(token=t, token_len=torch.tensor([t.shape[1]], dtype=torch.int32).cuda(), | |
| prompt_token=ptok.int(), prompt_token_len=torch.tensor([ptok.shape[1]], dtype=torch.int32).cuda(), | |
| prompt_feat=pfeat, prompt_feat_len=torch.tensor([pfeat.shape[1]], dtype=torch.int32).cuda(), | |
| embedding=emb, streaming=False, finalize=True) | |
| w, _ = cosy.model.hift.inference(speech_feat=mel, finalize=True) | |
| return w.squeeze().float().cpu().numpy() | |
| def item_text(r): | |
| return r["text"] if r["lang"] in ("cs", "para") else f"<{r['lang']}>{r['text']}</{r['lang']}>" | |
| class Synth: | |
| def __init__(self, lm_ckpt): | |
| self.cosy = load_cv3(); self.tok = Tok(self.cosy) | |
| self.m, self.spk2id = load_model(self.cosy, lm_ckpt, self.tok) | |
| self.vcache = {} | |
| vt = f"{PDIR}/flash_voices_text.json" | |
| self.vtext = json.load(open(vt if os.path.exists(vt) else f"{_HERE}/voices_text.json", encoding="utf-8")) | |
| def voice(self, name): | |
| if name not in self.vcache: | |
| v = VOICES[name] | |
| ptok, pfeat, emb = prompt_feats(self.cosy, v["wav"]) | |
| ids, langs = self.tok(self.vtext[name]) | |
| lmp = (torch.tensor(ids).cuda(), torch.tensor(langs).cuda(), ptok[0].long()) | |
| self.vcache[name] = (ptok, pfeat, emb, lmp, self.spk2id.get(v["spk"], self.spk2id.get(ANCHOR))) | |
| return self.vcache[name] | |
| def tokens(self, text, voice="anchor", rate="normal", seed=0, B=1, sampler="silexcap"): | |
| ptok, pfeat, emb, lmp, sid = self.voice(voice) | |
| ids, langs = self.tok(text) | |
| spk = torch.tensor([sid, R2I.get(None if rate == "none" else rate, 0)]).cuda() | |
| with torch.autocast("cuda", dtype=torch.bfloat16): | |
| return self.m.generate(torch.tensor(ids).cuda(), torch.tensor(langs).cuda(), spk, prompt=lmp, B=B, seed=seed, | |
| sampler=sampler) | |
| def wav(self, toks, voice="anchor"): | |
| ptok, pfeat, emb, _, _ = self.voice(voice) | |
| if len(toks) < 3: | |
| toks = list(toks) + [1, 1, 1] | |
| return token2wav(self.cosy, toks, ptok, pfeat, emb) | |
| def gen(a): | |
| S = Synth(a.lm) | |
| rows = [json.loads(l) for l in open(a.rows, encoding="utf-8")] | |
| if a.shard: | |
| i, n = map(int, a.shard.split("/")); rows = rows[i::n] | |
| os.makedirs(a.out, exist_ok=True) | |
| json.dump(vars(a), open(f"{a.out}/variant.json", "w")) | |
| for r in rows: | |
| if os.path.exists(f"{a.out}/{r['id']}.wav"): | |
| continue | |
| text = r.get("tts_text") or item_text(r) | |
| toks = S.tokens(text, a.voice, a.rate, a.seed)[0] | |
| w = S.wav(toks, a.voice) | |
| sf.write(f"{a.out}/{r['id']}.wav", w, S.cosy.sample_rate) | |
| json.dump(dict(tokens=[toks], ntok=[len(toks)]), open(f"{a.out}/{r['id']}.tok.json", "w")) | |
| print("GEN_DONE", a.out, len(rows), flush=True) | |
| if __name__ == "__main__": | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("cmd") | |
| ap.add_argument("--out", required=True) | |
| ap.add_argument("--lm") | |
| ap.add_argument("--rows", default=f"{HOME}/flash/score/rows_flash.jsonl") | |
| ap.add_argument("--voice", default="anchor") | |
| ap.add_argument("--rate", default="normal") | |
| ap.add_argument("--seed", type=int, default=0) | |
| ap.add_argument("--shard") | |
| ap.add_argument("--minutes", type=float, default=75) | |
| ap.add_argument("--max_steps", type=int, default=12000) | |
| ap.add_argument("--lr", type=float, default=2e-5) | |
| ap.add_argument("--lr_new", type=float, default=1e-3) | |
| ap.add_argument("--warmup", type=int, default=300) | |
| ap.add_argument("--tokens", type=int, default=16000) | |
| ap.add_argument("--max_tok", type=int, default=750) | |
| ap.add_argument("--adv_w", type=float, default=0.02) | |
| ap.add_argument("--p_prompt", type=float, default=0.5) | |
| ap.add_argument("--tag_drop", type=float, default=0.5) | |
| ap.add_argument("--rate_drop", type=float, default=0.5) | |
| ap.add_argument("--anchor_w", type=float, default=1.0) | |
| ap.add_argument("--val_every", type=int, default=1000) | |
| ap.add_argument("--patience", type=int, default=4) | |
| ap.add_argument("--quad", type=float, default=600.0) | |
| ap.add_argument("--extra_man", default=f"{HOME}/flash/data/igbosyn_pseudo/manifest.jsonl") | |
| ap.add_argument("--qual", default=f"{HOME}/flash/qual/utmos_train.jsonl") | |
| ap.add_argument("--qual_drop", type=float, default=0.1) | |
| a = ap.parse_args() | |
| {"lm": train_lm, "gen": gen}[a.cmd](a) | |