| """Arm B: CosyVoice2-0.5B (Apache-2.0) fine-tune with the Zhang et al. 2019 recipe added to its text->speech-token LM. |
| |
| Additions (new, ours): |
| * per-token language embedding added to the Qwen text-token embeddings (zero-init) |
| * learned speaker embedding inserted as one token AFTER the text, before the task token (zero-init), so text |
| positions cannot see it; speech positions attend to it (anchor-voice conditioning) |
| * domain-adversarial speaker classifier with gradient reversal on the LM hidden states at text positions |
| Flow-matching decoder + HiFT vocoder stay frozen; at inference they are prompted with an anchor-voice clip. |
| |
| python train_b.py train --out ~/runs/b1 --minutes 30 |
| python train_b.py synth --ckpt ~/runs/b1/last.pt --out ~/eval/sys/armB_ft |
| python train_b.py zeroshot --out ~/eval/sys/cosyvoice2_zeroshot (pretrained model, cross-lingual, no fine-tune) |
| """ |
| import argparse, inspect, json, math, os, random, 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"] |
| MODEL_DIR = f"{HOME}/models/CosyVoice2-0.5B" |
| LANGS = ["eng", "ibo", "hau", "yor", "pcm"] |
| L2I = {l: i for i, l in enumerate(LANGS)} |
| sys.path.insert(0, f"{HOME}/da") |
| from da.text import segments |
|
|
| from da.model import grad_reverse |
|
|
|
|
| def _sf_load_wav(wav, target_sr, min_sr=16000): |
| """soundfile replacement for cosyvoice load_wav (torchaudio>=2.9 load needs torchcodec).""" |
| 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(): |
| 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 |
| kw = {k: v for k, v in kw.items() if k in sig} |
| return CosyVoice2(MODEL_DIR, **kw) |
|
|
|
|
| def text_tokens(tok, text, default="eng"): |
| """tagged text -> (qwen ids, per-token lang ids).""" |
| 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 |
|
|
|
|
| class DALM(nn.Module): |
| def __init__(self, qlm, n_spk, d=896): |
| super().__init__() |
| self.q = qlm |
| 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.adv = nn.Sequential(nn.Linear(d, 256), nn.ReLU(), nn.Linear(256, n_spk)) |
| 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 build(self, ids, langs, spk, speech=None): |
| """one sequence: [sos, text+lang, spk, task, speech...]""" |
| sos = self.q.llm_embedding.weight[self.q.sos][None] |
| task = self.q.llm_embedding.weight[self.q.task_id][None] |
| parts = [sos, self.embed_text(ids, langs), self.spk_emb(spk)[None], task] |
| if speech is not None: |
| parts.append(self.q.speech_embedding(speech)) |
| return torch.cat(parts, 0) |
|
|
| def hidden(self, x, mask): |
| out = self.q.llm.model.model(inputs_embeds=x, attention_mask=mask, return_dict=True) |
| return out.last_hidden_state |
|
|
| def forward(self, batch, adv_lambda=1.0): |
| seqs, tgts, tpos = [], [], [] |
| for ids, langs, spk, sp in batch: |
| x = self.build(ids, langs, spk, sp) |
| n_text = len(ids) |
| tgt = torch.full((x.shape[0],), -100, device=x.device, dtype=torch.long) |
| |
| start = 1 + n_text + 1 |
| tgt[start: start + len(sp)] = sp |
| tgt[start + len(sp)] = self.V |
| seqs.append(x); tgts.append(tgt); tpos.append((1, 1 + n_text)) |
| T = max(s.shape[0] for s in seqs) + 1 |
| 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[:len(s)] if len(t) > len(s) else 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,), int(batch[i][2]), 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) |
| adv_acc = (al.argmax(-1) == ys).float().mean() |
| return ce, acc, adv, adv_acc |
|
|
| @torch.no_grad() |
| def generate(self, ids, langs, spk, top_p=0.8, top_k=25, win=10, tau_r=0.1, seed=0): |
| g = torch.Generator(device=ids.device).manual_seed(seed) |
| x = self.build(ids, langs, spk)[None] |
| n = len(ids) |
| min_len, max_len = int(n * 2), int(n * 20) |
| past, out = None, [] |
| inp = 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 |
| cand_p, cand_i = sp[keep], si[keep] |
| tokid = int(cand_i[torch.multinomial(cand_p / cand_p.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(root, datasets, split, max_tok=500): |
| items, mm = [], {} |
| for ds in datasets: |
| if not os.path.exists(f"{root}/{ds}/s3tok.jsonl"): |
| continue |
| mm[ds] = np.memmap(f"{root}/{ds}/s3tok.i32", dtype=np.int32, mode="r") |
| tk = {} |
| for l in open(f"{root}/{ds}/s3tok.jsonl"): |
| r = json.loads(l); tk[r["id"]] = (r["toff"], r["ntok"]) |
| for l in open(f"{root}/{ds}/meta.jsonl", encoding="utf-8"): |
| m = json.loads(l) |
| if m["split"] != split or m["id"] not in tk: |
| continue |
| m["toff"], m["ntok"] = tk[m["id"]] |
| if 5 < m["ntok"] <= max_tok: |
| items.append(m) |
| return items, mm |
|
|
|
|
| def batches(items, budget, weights, rng): |
| idx = [] |
| for i, m in enumerate(items): |
| w = weights.get(m["ds"], 1.0) |
| idx += [i] * int(w) + ([i] if rng.random() < w - int(w) else []) |
| idx.sort(key=lambda i: items[i]["ntok"] + rng.random() * 10) |
| out, cur, mx = [], [], 0 |
| for i in idx: |
| L = items[i]["ntok"] + len(items[i]["text"]) // 2 + 4 |
| 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(a): |
| from transformers import AutoTokenizer |
| tok = AutoTokenizer.from_pretrained(f"{MODEL_DIR}/CosyVoice-BlankEN") |
| cosy = load_cosy() |
| items, mm = load_items(a.data, a.datasets.split(","), "train") |
| val, _ = load_items(a.data, a.datasets.split(","), "val") |
| spks = sorted({m["spk"] for m in items}); spk2id = {s: i for i, s in enumerate(spks)} |
| val = [m for m in val if m["spk"] in spk2id][:200] |
| os.makedirs(a.out, exist_ok=True) |
| json.dump(spk2id, open(f"{a.out}/spk2id.json", "w")) |
| print("armB items", len(items), "hours", round(sum(m["dur"] for m in items) / 3600, 1), "spk", len(spks), flush=True) |
| model = DALM(cosy.model.llm, len(spks)).cuda().train() |
| 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 enc(m): |
| if m["id"] not in cache: |
| cache[m["id"]] = text_tokens(tok, f"<{m['lang']}>{m['text']}</{m['lang']}>") |
| ids, langs = cache[m["id"]] |
| sp = np.asarray(mm[m["ds"]][m["toff"]: m["toff"] + m["ntok"]], dtype=np.int64) |
| return (torch.tensor(ids).cuda(), torch.tensor(langs).cuda(), torch.tensor(spk2id[m["spk"]]).cuda(), |
| torch.from_numpy(sp).cuda()) |
|
|
| rng = random.Random(0); step = 0; nskip = 0; t0 = time.time(); logf = open(f"{a.out}/log.jsonl", "a") |
|
|
| def validate(): |
| model.eval(); L = [] |
| with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16): |
| for i in range(0, len(val), 16): |
| ce, acc, _, _ = model([enc(m) for m in val[i:i + 16]]); L.append(float(ce)) |
| model.train(); return float(np.mean(L)) |
|
|
| v0 = validate(); print(json.dumps({"step": 0, "val_ce": v0}), flush=True); logf.write(json.dumps({"step": 0, "val_ce": v0}) + "\n") |
| done = False |
| while not done: |
| for bi in batches(items, a.tokens, {"waxal": a.waxal_w}, rng): |
| lrm = min(1.0, (step + 1) / 200) |
| for g, b in zip(opt.param_groups, base): |
| g["lr"] = b * lrm |
| with torch.autocast("cuda", dtype=torch.bfloat16): |
| ce, acc, adv, adv_acc = model([enc(items[i]) for i in bi]) |
| 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), "B": len(bi), |
| "maxtok": max(items[i]["ntok"] for i in bi)}), 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), min=round((time.time() - t0) / 60, 2)) |
| print(json.dumps(r), flush=True); logf.write(json.dumps(r) + "\n"); logf.flush() |
| if step % 500 == 0: |
| v = validate(); print(json.dumps({"step": step, "val_ce": v}), flush=True) |
| logf.write(json.dumps({"step": step, "val_ce": v}) + "\n"); logf.flush() |
| if (time.time() - t0) / 60 > a.minutes: |
| done = True; break |
| v = validate(); logf.write(json.dumps({"step": step, "val_ce": v, "final": True}) + "\n") |
| sd = {k: v for k, v in model.state_dict().items()} |
| torch.save(dict(model=sd, step=step, n_spk=len(spks)), f"{a.out}/last.pt") |
| print("TRAIN_B_DONE", step, v, flush=True) |
|
|
|
|
| |
| def prompt_feats(cosy, wav_path): |
| """flow prompt (speech tokens, 24k mel feat, campplus embedding) from an anchor clip.""" |
| fe = cosy.frontend |
| tokn, tokl = fe._extract_speech_token(wav_path) |
| feat, featl = 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.cuda(), feat.cuda(), emb.cuda() |
|
|
|
|
| @torch.no_grad() |
| def token2wav(cosy, tokens, ptok, pfeat, emb): |
| flow = cosy.model.flow |
| t = torch.tensor([tokens], dtype=torch.int32).cuda() |
| kw = dict(token=t, token_len=torch.tensor([t.shape[1]], dtype=torch.int32).cuda(), |
| prompt_token=ptok, 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) |
| 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).cuda() |
| 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).cuda()} if "cache_source" in hs else {})) |
| w = w[0] if isinstance(w, tuple) else w |
| return w.squeeze().float().cpu().numpy() |
|
|
|
|
| def eval_rows(): |
| return [json.loads(l) for l in open(f"{HOME}/eval/eval_set.jsonl", encoding="utf-8")] |
|
|
|
|
| def synth(a): |
| from transformers import AutoTokenizer |
| tok = AutoTokenizer.from_pretrained(f"{MODEL_DIR}/CosyVoice-BlankEN") |
| cosy = load_cosy() |
| ck = torch.load(a.ckpt, map_location="cuda") |
| spk2id = json.load(open(os.path.join(os.path.dirname(a.ckpt), "spk2id.json"))) |
| model = DALM(cosy.model.llm, ck["n_spk"]).cuda().eval() |
| model.load_state_dict(ck["model"]) |
| anc = json.load(open(f"{HOME}/eval/anchor.json")) |
| ptok, pfeat, emb = prompt_feats(cosy, anc["ref_wavs"][0]) |
| os.makedirs(a.out, exist_ok=True) |
| for r in eval_rows(): |
| text = r["text"] if r["lang"] == "cs" else f"<{r['lang']}>{r['text']}</{r['lang']}>" |
| ids, langs = text_tokens(tok, text) |
| with torch.autocast("cuda", dtype=torch.bfloat16): |
| sp = model.generate(torch.tensor(ids).cuda(), torch.tensor(langs).cuda(), torch.tensor(spk2id[anc["spk"]]).cuda()) |
| sf.write(f"{a.out}/{r['id']}.wav", token2wav(cosy, sp, ptok, pfeat, emb), cosy.sample_rate) |
| print("synth_b done", a.out, flush=True) |
|
|
|
|
| def zeroshot(a): |
| """Unmodified pretrained CosyVoice2, cross-lingual mode with the anchor prompt (reference point).""" |
| cosy = load_cosy() |
| anc = json.load(open(f"{HOME}/eval/anchor.json")) |
| from cosyvoice.utils.file_utils import load_wav |
| pw = anc["ref_wavs"][0] |
| sig = inspect.signature(cosy.inference_cross_lingual).parameters |
| os.makedirs(a.out, exist_ok=True) |
| for r in eval_rows(): |
| text = " ".join(s.strip() for _, s in segments(r["text"])) if r["lang"] == "cs" else r["text"] |
| prompt = pw |
| outs = [o["tts_speech"] for o in cosy.inference_cross_lingual(text, prompt, stream=False)] |
| w = torch.cat(outs, 1).squeeze().cpu().numpy() |
| sf.write(f"{a.out}/{r['id']}.wav", w, cosy.sample_rate) |
| print("zeroshot done", flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| ap = argparse.ArgumentParser() |
| ap.add_argument("cmd") |
| ap.add_argument("--data", default=f"{HOME}/data") |
| ap.add_argument("--datasets", default="waxal,fleurs,slr86,slr70,bible_yor") |
| ap.add_argument("--out", default=f"{HOME}/runs/b1") |
| ap.add_argument("--ckpt", default="") |
| ap.add_argument("--minutes", type=float, default=30) |
| ap.add_argument("--lr", type=float, default=1e-5) |
| ap.add_argument("--lr_new", type=float, default=1e-3) |
| ap.add_argument("--tokens", type=int, default=12000) |
| ap.add_argument("--waxal_w", type=float, default=2.0) |
| ap.add_argument("--adv_w", type=float, default=0.02) |
| a = ap.parse_args() |
| {"train": train, "synth": synth, "zeroshot": zeroshot}[a.cmd](a) |
|
|
