"""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) @torch.no_grad() 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}" 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() @torch.no_grad() 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']}" 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)