| """Arm B, Phase 2. CosyVoice2-0.5B (Apache-2.0) + language embeddings + speaker token + GRL speaker classifier. |
| |
| New in Phase 2 |
| * data from ~/data2 (24 kHz, forced-aligned, pause-normalised; long WAXAL clips segmented; BibleTTS Hausa; CS pairs) |
| * in-context prompts during LM training (p=0.5): [sos, prompt_text+text, spk, task, prompt_speech, speech] with |
| loss on the target speech only -> at inference the anchor clip steers prosody and pause style |
| * keep-best on validation CE (every --val_every steps), cosine LR, early stop |
| * flow-matching decoder fine-tune on 24 kHz studio sources (tones / pitch / timing detail) |
| * synthesis: anchor prompt at LM and flow, sentence-level generation for multi-sentence text with anchor-matched |
| inter-sentence pauses, N candidates per text (seed + prompt clip vary) |
| |
| python train_b2.py lm --out ~/runs/p2_lm --minutes 90 |
| python train_b2.py flow --out ~/runs/p2_flow --minutes 60 |
| python train_b2.py synth --ckpt ~/runs/p2_lm/best.pt [--flow ~/runs/p2_flow/best.pt] --out DIR --n 4 |
| """ |
| import argparse, inspect, json, math, os, random, re, sys, time |
| import numpy as np, torch, torch.nn as nn, torch.nn.functional as F, soundfile as sf, librosa |
|
|
| HOME = os.path.expanduser("~") |
| sys.path[:0] = [f"{HOME}/CosyVoice", f"{HOME}/CosyVoice/third_party/Matcha-TTS", f"{HOME}/da"] |
| from train_b import load_cosy, text_tokens, DALM, LANGS, L2I, MODEL_DIR, token2wav |
| from da.model import grad_reverse |
| from da.text import segments |
|
|
| D2 = os.path.expanduser(os.environ.get("P2OUT", f"{HOME}/data2")) |
| ANCHOR = "waxal_yor_4" |
|
|
|
|
| |
| class DALM2(DALM): |
| def build2(self, ids, langs, spk, sp=None, pids=None, plangs=None, psp=None): |
| """[sos, (prompt_text) text, spk, task, (prompt_speech) speech] -> embeddings, n_text_total, n_prompt_speech""" |
| 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]) |
| parts = [sos, self.embed_text(ids, langs), self.spk_emb(spk)[None], task] |
| 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(ids), npsp |
|
|
| def forward2(self, batch, adv_lambda=1.0): |
| seqs, tgts, tpos, spks = [], [], [], [] |
| for b in batch: |
| ids, langs, spk, sp = b[:4] |
| pr = b[4] if len(b) > 4 else None |
| x, nt, npsp = self.build2(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 = 1 + nt + 2 + npsp - 1 |
| tgt[start: start + len(sp)] = sp |
| tgt[start + len(sp)] = self.V |
| seqs.append(x); tgts.append(tgt[:x.shape[0]]); tpos.append((1, 1 + nt)); spks.append(int(spk)) |
| T = max(s.shape[0] for s in seqs) |
| D = seqs[0].shape[1] |
| X = seqs[0].new_zeros(len(seqs), T, D) |
| Y = torch.full((len(seqs), T), -100, device=X.device, dtype=torch.long) |
| M = torch.zeros(len(seqs), T, 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 |
| 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, ys = [], [] |
| for i, (a, b) in enumerate(tpos): |
| hs.append(H[i, a:b]); ys.append(torch.full((b - a,), spks[i], device=X.device)) |
| hs = torch.cat(hs); ys = torch.cat(ys) |
| al = self.adv(grad_reverse(hs.float(), adv_lambda)) |
| adv = F.cross_entropy(al, ys) |
| return ce, acc, adv, (al.argmax(-1) == ys).float().mean() |
|
|
| @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): |
| g = torch.Generator(device=ids.device).manual_seed(seed) |
| x, _, _ = self.build2(ids, langs, spk, None, *(prompt if prompt else (None, None, None))) |
| x = x[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)]) |
| if out and sum(1 for t in out[-win:] if t == tokid) >= win * tau_r and tokid != self.V: |
| 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 |
|
|
|
|
| |
| def load_items(datasets, split, max_tok=500, flow=False): |
| items, mm = [], {} |
| for ds in datasets: |
| od = f"{D2}/{ds}" |
| if not os.path.exists(f"{od}/s3tok.jsonl"): |
| print("missing", ds, flush=True); continue |
| mm[ds] = np.memmap(f"{od}/s3tok.i32", dtype=np.int32, mode="r") |
| tk = [json.loads(l) for l in open(f"{od}/s3tok.jsonl")] |
| for k, l in enumerate(open(f"{od}/meta.jsonl", encoding="utf-8")): |
| m = json.loads(l) |
| t = tk[k] |
| assert t["id"] == m["id"] |
| if m["split"] != split: |
| continue |
| if flow and not m.get("flow_ok"): |
| continue |
| m["toff"], m["ntok"], m["row"] = t["toff"], t["ntok"], k |
| if 5 < m["ntok"] <= max_tok: |
| items.append(m) |
| return items, mm |
|
|
|
|
| def tagged(m): |
| return m["text"] if m["lang"] == "cs" else f"<{m['lang']}>{m['text']}</{m['lang']}>" |
|
|
|
|
| DS_W = {"waxal": 2.0, "bible_yor": 0.5, "bible_hau": 0.6, "fleurs": 1.0, "slr86": 1.0, "slr70": 1.0, "cs_pairs": 1.0} |
|
|
|
|
| def batches(items, budget, rng, extra_w=None, cost=None): |
| idx = [] |
| for i, m in enumerate(items): |
| w = DS_W.get(m["ds"], 1.0) * (extra_w(m) if extra_w else 1.0) |
| idx += [i] * int(w) + ([i] if rng.random() < w - int(w) else []) |
| cost = cost or (lambda m: m["ntok"] + len(m["text"]) // 2 + 4) |
| 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 |
|
|
|
|
| |
| def train_lm(a): |
| from transformers import AutoTokenizer |
| tok = AutoTokenizer.from_pretrained(f"{MODEL_DIR}/CosyVoice-BlankEN") |
| cosy = load_cosy() |
| dss = a.datasets.split(",") |
| items, mm = load_items(dss, "train") |
| val_all, _ = load_items(dss, "val") |
| if a.anchor_only_w > 0: |
| print("anchor items", sum(m["spk"] == ANCHOR for m in items), flush=True) |
| spks = sorted({m["spk"] for m in items}) |
| if a.init: |
| ck = torch.load(a.init, map_location="cpu") |
| spk2id = json.load(open(os.path.join(os.path.dirname(a.init), "spk2id.json"))) |
| for s in spks: |
| if s not in spk2id: |
| spk2id[s] = len(spk2id) |
| else: |
| spk2id = {s: i for i, s in enumerate(spks)} |
| n_spk = max(len(spk2id), ck["n_spk"]) if a.init else 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"] <= 160 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")) |
| print("lm items", len(items), "hours", round(sum(m["dur"] for m in items) / 3600, 1), "spk", len(spk2id), |
| "val", len(val), per_l, flush=True) |
| model = DALM2(cosy.model.llm, n_spk).cuda().train() |
| if a.init: |
| sd = ck["model"] |
| if sd["spk_emb.weight"].shape[0] < n_spk: |
| for k in ["spk_emb.weight"]: |
| w = model.state_dict()[k].clone(); w[:sd[k].shape[0]] = sd[k]; sd[k] = w |
| for k in ["adv.2.weight", "adv.2.bias"]: |
| w = model.state_dict()[k].clone(); w[:sd[k].shape[0]] = sd[k]; sd[k] = w |
| print(model.load_state_dict(sd, strict=False), flush=True) |
| for p in cosy.model.flow.parameters(): |
| p.requires_grad = False |
| for p in cosy.model.hift.parameters(): |
| p.requires_grad = False |
| new = list(model.lang_emb.parameters()) + list(model.spk_emb.parameters()) + list(model.adv.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] |
| cache = {} |
|
|
| def tt(text, key): |
| if key not in cache: |
| cache[key] = text_tokens(tok, text) |
| return cache[key] |
|
|
| def enc(m, prng=None, force_prompt=None): |
| ids, langs = tt(tagged(m), m["id"]) |
| sp = np.asarray(mm[m["ds"]][m["toff"]: m["toff"] + m["ntok"]], dtype=np.int64) |
| r = [torch.tensor(ids).cuda(), torch.tensor(langs).cuda(), torch.tensor(spk2id[m["spk"]]).cuda(), |
| 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"] <= 600: |
| p = c[0] |
| pids, plangs = tt(tagged(p), p["id"]) |
| psp = np.asarray(mm[p["ds"]][p["toff"]: p["toff"] + p["ntok"]], dtype=np.int64) |
| r.append((torch.tensor(pids).cuda(), torch.tensor(plangs).cuda(), torch.from_numpy(psp).cuda())) |
| return r |
|
|
| def validate(): |
| model.eval(); out = {} |
| 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) for m in bt]) |
| L.append(float(ce) * len(bt)); W.append(len(bt)) |
| out["p" if mode else "np"] = sum(L) / sum(W) |
| |
| av = [m for m in val if m["spk"] == ANCHOR] |
| if av: |
| vr = random.Random(7) |
| ce, _, _, _ = model.forward2([enc(m, vr, force_prompt=True) for m in av[:32]]) |
| out["anchor_p"] = float(ce) |
| model.train() |
| out["mean"] = (out["np"] + out["p"]) / 2 |
| return out |
|
|
| extra = (lambda m: a.anchor_only_w if m["spk"] == ANCHOR else 1.0) if a.anchor_only_w > 0 else None |
| rng = random.Random(0); prng = random.Random(3) |
| step, nskip, t0 = 0, 0, time.time() |
| logf = open(f"{a.out}/log.jsonl", "a") |
| v = validate(); best = v["mean"] if not a.select_anchor else 0.5 * (v["mean"] + v.get("anchor_p", v["mean"])); bad = 0 |
| logf.write(json.dumps({"step": 0, "val": v}) + "\n"); print(json.dumps({"step": 0, "val": v}), flush=True) |
| torch.save(dict(model=model.state_dict(), step=0, n_spk=n_spk, val=v), f"{a.out}/best.pt") |
| total = a.max_steps |
| done = False |
| while not done: |
| for bi in batches(items, a.tokens, rng, extra, cost=lambda m: m["ntok"] + len(m["text"]) // 2 + 4 + int(a.p_prompt * 200)): |
| 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 g, b in zip(opt.param_groups, base): |
| g["lr"] = b * lrm |
| batch = [enc(items[i], prng) for i in bi] |
| |
| 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 |
| opt.zero_grad(set_to_none=True) |
| if not torch.isfinite(loss): |
| nskip += 1; print(json.dumps({"skip_nonfinite_fwd": step, "ce": float(ce), "adv": float(adv)}), flush=True) |
| continue |
| loss.backward() |
| 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); print(json.dumps({"skip_nonfinite_grad": step}), flush=True) |
| continue |
| opt.step(); step += 1 |
| if step % 25 == 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, 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() |
| score = v["mean"] if not a.select_anchor else 0.5 * (v["mean"] + v.get("anchor_p", v["mean"])) |
| imp = score < 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 = score, 0 |
| torch.save(dict(model=model.state_dict(), 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, flush=True) |
|
|
|
|
| |
| def train_flow(a): |
| cosy = load_cosy() |
| flow = cosy.model.flow |
| if a.flow_init: |
| flow.load_state_dict(torch.load(a.flow_init, map_location="cuda")["flow"]) |
| flow.train() |
| for p in cosy.model.llm.parameters(): |
| p.requires_grad = False |
| fe = cosy.frontend.feat_extractor |
| dss = a.datasets.split(",") |
| items, mm = load_items(dss, "train", max_tok=400, flow=True) |
| val, _ = load_items(dss, "val", max_tok=400, flow=True) |
| rng0 = random.Random(1); rng0.shuffle(val); val = val[:256] |
| pcm = {ds: np.memmap(f"{D2}/{ds}/pcm24k.i16", dtype=np.int16, mode="r") for ds in dss if os.path.exists(f"{D2}/{ds}/pcm24k.i16")} |
| emb = {ds: np.load(f"{D2}/{ds}/campplus.npy", mmap_mode="r") for ds in dss if os.path.exists(f"{D2}/{ds}/campplus.npy")} |
| os.makedirs(a.out, exist_ok=True) |
| print("flow items", len(items), "hours", round(sum(m["dur"] for m in items) / 3600, 1), "val", len(val), flush=True) |
| extra = (lambda m: a.anchor_only_w if m["spk"] == ANCHOR else 1.0) if a.anchor_only_w > 0 else None |
|
|
| def make(bt): |
| toks, feats, es = [], [], [] |
| for m in bt: |
| y = torch.from_numpy(np.asarray(pcm[m["ds"]][m["off24"]: m["off24"] + m["n24"]], dtype=np.float32) / 32768.0) |
| n = int(math.ceil(len(y) / 960)) * 960 |
| y = F.pad(y, (0, n - len(y)))[None].cuda() |
| ft = fe(y).squeeze(0).transpose(0, 1) |
| tk = torch.from_numpy(np.asarray(mm[m["ds"]][m["toff"]: m["toff"] + m["ntok"]], dtype=np.int64)) |
| k = min(ft.shape[0] // 2, len(tk)) |
| toks.append(tk[:k]); feats.append(ft[:2 * k]); es.append(torch.from_numpy(np.array(emb[m["ds"]][m["row"]]))) |
| tl = torch.tensor([len(t) for t in toks], dtype=torch.int32) |
| fl = torch.tensor([len(f) for f in feats], dtype=torch.int32) |
| T = int(tl.max()); Fm = int(fl.max()) |
| tok = torch.zeros(len(bt), T, dtype=torch.long); feat = torch.zeros(len(bt), Fm, 80) |
| for i, (t, f) in enumerate(zip(toks, feats)): |
| tok[i, :len(t)] = t; feat[i, :len(f)] = f.cpu() |
| return dict(speech_token=tok, speech_token_len=tl, speech_feat=feat, speech_feat_len=fl, |
| embedding=torch.stack(es).float()) |
|
|
| params = [p for p in flow.parameters() if p.requires_grad] |
| opt = torch.optim.AdamW(params, lr=a.lr, weight_decay=0.0) |
| logf = open(f"{a.out}/log.jsonl", "a") |
|
|
| def validate(): |
| flow.eval(); L = [] |
| st = torch.random.get_rng_state(); cst = torch.cuda.get_rng_state(); pst = random.getstate() |
| torch.manual_seed(0); random.seed(0) |
| with torch.no_grad(): |
| for i in range(0, len(val), 16): |
| L.append(float(flow(make(val[i:i + 16]), torch.device("cuda"))["loss"])) |
| torch.random.set_rng_state(st); torch.cuda.set_rng_state(cst); random.setstate(pst) |
| flow.train(); return float(np.mean(L)) |
|
|
| v = validate(); best, bad = v, 0 |
| print(json.dumps({"step": 0, "val": v}), flush=True); logf.write(json.dumps({"step": 0, "val": v}) + "\n") |
| torch.save(dict(flow=flow.state_dict(), step=0, val=v), f"{a.out}/best.pt") |
| rng = random.Random(0); step, t0, done = 0, time.time(), False |
| while not done: |
| for bi in batches(items, a.tokens, rng, extra, cost=lambda m: m["ntok"]): |
| lrm = min(1.0, (step + 1) / a.warmup) |
| for g in opt.param_groups: |
| g["lr"] = a.lr * lrm |
| with torch.autocast("cuda", dtype=torch.bfloat16): |
| loss = flow(make([items[i] for i in bi]), torch.device("cuda"))["loss"] |
| opt.zero_grad(set_to_none=True) |
| if not torch.isfinite(loss): |
| continue |
| loss.backward() |
| gn = torch.nn.utils.clip_grad_norm_(params, 5.0) |
| if not torch.isfinite(gn): |
| continue |
| opt.step(); step += 1 |
| if step % 50 == 0: |
| r = dict(step=step, loss=round(float(loss), 4), gn=round(float(gn), 2), B=len(bi), 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 < best - 1e-4 |
| print(json.dumps({"step": step, "val": v, "best": imp}), flush=True) |
| logf.write(json.dumps({"step": step, "val": v, "best": imp}) + "\n"); logf.flush() |
| if imp: |
| best, bad = v, 0 |
| torch.save(dict(flow=flow.state_dict(), step=step, val=v), f"{a.out}/best.pt") |
| else: |
| bad += 1 |
| if bad >= a.patience: |
| done = True; break |
| if (time.time() - t0) / 60 > a.minutes or step >= a.max_steps: |
| done = True; break |
| print("TRAIN_FLOW_DONE", step, "best", best, flush=True) |
|
|
|
|
| |
| _SENT = re.compile(r"(?<=[.?!])\s+") |
|
|
|
|
| def split_sentences(text): |
| """tagged text -> list of tagged sentences. A sentence ends at . ? ! inside a span or at the end of a span; |
| language spans are cut there and re-tagged, so a sentence keeps its in-sentence switches.""" |
| 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()] |
|
|
|
|
| def trim_edges(w, sr, keep_lead=0.12, keep_tail=0.18): |
| from da.p2common import silence_mask |
| 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): |
| """shorten every internal silence longer than cap seconds to cap (the centre of the silent run is removed).""" |
| from da.p2common import silence_mask, runs |
| 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: |
| cut0 = a * hop + int(cap / 2 * sr) |
| cut1 = b * hop - int(cap / 2 * sr) |
| keep.append(w[cur:cut0]); cur = cut1 |
| keep.append(w[cur:]) |
| return np.concatenate(keep) |
|
|
|
|
| class Synth: |
| def __init__(self, ckpt, flow_ckpt=None, sent_pause=0.5): |
| from transformers import AutoTokenizer |
| self.tok = AutoTokenizer.from_pretrained(f"{MODEL_DIR}/CosyVoice-BlankEN") |
| self.cosy = load_cosy() |
| if flow_ckpt: |
| self.cosy.model.flow.load_state_dict(torch.load(flow_ckpt, map_location="cuda")["flow"]) |
| ck = torch.load(ckpt, map_location="cuda") |
| self.spk2id = json.load(open(os.path.join(os.path.dirname(ckpt), "spk2id.json"))) |
| self.m = DALM2(self.cosy.model.llm, ck["n_spk"]).cuda().eval() |
| self.m.load_state_dict(ck["model"]) |
| self.sent_pause = sent_pause |
| self.cap = 0.6 |
| |
| meta = [json.loads(l) for l in open(f"{D2}/waxal/meta.jsonl", encoding="utf-8")] |
| tk = [json.loads(l) for l in open(f"{D2}/waxal/s3tok.jsonl")] |
| pcm = np.memmap(f"{D2}/waxal/pcm24k.i16", dtype=np.int16, mode="r") |
| s3 = np.memmap(f"{D2}/waxal/s3tok.i32", dtype=np.int32, mode="r") |
| cand = [(m, t) for m, t in zip(meta, tk) if m["spk"] == ANCHOR and m["split"] == "val" and 3.0 <= m["dur"] <= 7.0 |
| and not m.get("longform")] |
| cand.sort(key=lambda x: -x[0]["ascore"]) |
| self.prompts = [] |
| os.makedirs(f"{HOME}/eval/prompts", exist_ok=True) |
| for m, t in cand[:4]: |
| y = np.asarray(pcm[m["off24"]: m["off24"] + m["n24"]], dtype=np.float32) / 32768.0 |
| p = f"{HOME}/eval/prompts/{m['id']}.wav"; sf.write(p, y, 24000) |
| from train_b import prompt_feats |
| ptok, pfeat, pemb = prompt_feats(self.cosy, p) |
| ids, langs = text_tokens(self.tok, tagged(m)) |
| lmp = (torch.tensor(ids).cuda(), torch.tensor(langs).cuda(), |
| torch.from_numpy(np.asarray(s3[t["toff"]: t["toff"] + t["ntok"]], dtype=np.int64)).cuda()) |
| self.prompts.append(dict(id=m["id"], text=m["text"], lm=lmp, flow=(ptok, pfeat, pemb))) |
| print("anchor prompts", [p["id"] for p in self.prompts], flush=True) |
|
|
| def one(self, tagged_text, seed=0, prompt_k=0, use_lm_prompt=True): |
| ids, langs = text_tokens(self.tok, tagged_text) |
| P = self.prompts[prompt_k % len(self.prompts)] |
| with torch.autocast("cuda", dtype=torch.bfloat16): |
| sp = self.m.generate2(torch.tensor(ids).cuda(), torch.tensor(langs).cuda(), |
| torch.tensor(self.spk2id[ANCHOR]).cuda(), prompt=P["lm"] if use_lm_prompt else None, |
| seed=seed) |
| if len(sp) < 3: |
| sp = sp + [0] * 3 |
| self.last_tokens.append(sp) |
| return token2wav(self.cosy, sp, *P["flow"]) |
|
|
| def join(self, ws): |
| sr = self.cosy.sample_rate |
| 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) * sr)), np.float32)) |
| out.append(w) |
| return np.concatenate(out) |
|
|
| def redecode(self, toks, prompt_k=0): |
| """same LM tokens, (possibly different) flow -> identical post-processing""" |
| sr = self.cosy.sample_rate |
| P = self.prompts[prompt_k % len(self.prompts)] |
| return self.join([cap_internal(trim_edges(token2wav(self.cosy, t, *P["flow"]), sr), sr, self.cap) for t in toks]) |
|
|
| def say(self, tagged_text, seed=0, prompt_k=0, stitch=True, use_lm_prompt=True): |
| sr = self.cosy.sample_rate |
| self.last_tokens = [] |
| sents = split_sentences(tagged_text) if stitch else [tagged_text] |
| ws = [cap_internal(trim_edges(self.one(s, seed=seed + 1000 * i, prompt_k=prompt_k, use_lm_prompt=use_lm_prompt), sr), sr, self.cap) |
| for i, s in enumerate(sents)] |
| |
| return self.join(ws) |
|
|
|
|
| def redecode(a): |
| """decode the LM tokens saved by synth (--src DIR) with another flow checkpoint -> --out""" |
| S = Synth(a.ckpt, a.flow or None, a.sent_pause) |
| os.makedirs(a.out, exist_ok=True) |
| import glob |
| for f in sorted(glob.glob(f"{a.src}/*.tok.json")): |
| d = json.load(open(f)) |
| i = os.path.basename(f)[:-9] |
| sf.write(f"{a.out}/{i}.wav", S.redecode(d["tokens"], d["prompt_k"]), S.cosy.sample_rate) |
| print("REDECODE_DONE", a.out, flush=True) |
|
|
|
|
| def synth(a): |
| S = Synth(a.ckpt, a.flow or None, a.sent_pause) |
| rows = [json.loads(l) for l in open(a.eval_set, encoding="utf-8")] |
| if a.only: |
| keep = set(a.only.split(",")); rows = [r for r in rows if r["id"] in keep or r["lang"] in keep] |
| if a.nper: |
| cnt, sel = {}, [] |
| for r in rows: |
| if r["lang"] in ("cs", "para") or cnt.get(r["lang"], 0) < a.nper: |
| sel.append(r); cnt[r["lang"]] = cnt.get(r["lang"], 0) + 1 |
| rows = sel |
| ks = [a.k] if a.k >= 0 else list(range(a.n)) |
| for k in ks: |
| od = a.out if len(ks) == 1 else f"{a.out}/cand{k}" |
| os.makedirs(od, exist_ok=True) |
| for r in rows: |
| text = r["text"] if r["lang"] in ("cs", "para") else f"<{r['lang']}>{r['text']}</{r['lang']}>" |
| w = S.say(text, seed=k, prompt_k=k, stitch=not a.no_stitch, use_lm_prompt=not a.no_lm_prompt) |
| sf.write(f"{od}/{r['id']}.wav", w, S.cosy.sample_rate) |
| json.dump(dict(tokens=S.last_tokens, prompt_k=k), open(f"{od}/{r['id']}.tok.json", "w")) |
| print("synth cand", k, "done", od, flush=True) |
| print("SYNTH2_DONE", a.out, flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| ap = argparse.ArgumentParser() |
| ap.add_argument("cmd") |
| ap.add_argument("--datasets", default="waxal,fleurs,slr86,slr70,bible_yor,bible_hau,cs_pairs") |
| ap.add_argument("--out", default=f"{HOME}/runs/p2_lm") |
| ap.add_argument("--init", default="") |
| ap.add_argument("--flow_init", default="") |
| ap.add_argument("--ckpt", default="") |
| ap.add_argument("--flow", default="") |
| ap.add_argument("--minutes", type=float, default=90) |
| ap.add_argument("--max_steps", type=int, default=4000) |
| ap.add_argument("--lr", type=float, default=1e-5) |
| ap.add_argument("--lr_new", type=float, default=5e-4) |
| ap.add_argument("--warmup", type=int, default=300) |
| ap.add_argument("--tokens", type=int, default=16000) |
| ap.add_argument("--adv_w", type=float, default=0.02) |
| ap.add_argument("--p_prompt", type=float, default=0.5) |
| ap.add_argument("--val_every", type=int, default=200) |
| ap.add_argument("--patience", type=int, default=5) |
| ap.add_argument("--anchor_only_w", type=float, default=0.0) |
| ap.add_argument("--select_anchor", action="store_true") |
| ap.add_argument("--eval_set", default=f"{HOME}/eval/eval_set2.jsonl") |
| ap.add_argument("--n", type=int, default=1) |
| ap.add_argument("--nper", type=int, default=0) |
| ap.add_argument("--k", type=int, default=-1) |
| ap.add_argument("--src", default="") |
| ap.add_argument("--only", default="") |
| ap.add_argument("--sent_pause", type=float, default=0.6) |
| ap.add_argument("--no_stitch", action="store_true") |
| ap.add_argument("--no_lm_prompt", action="store_true") |
| a = ap.parse_args() |
| {"lm": train_lm, "flow": train_flow, "synth": synth, "redecode": redecode}[a.cmd](a) |
|
|
