flash_lm.py

"""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}</{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()


@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']}</{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)