Spaces:
Running on Zero
Running on Zero
| """Serving controls: the decode-side fixes for stop behaviour, long-form drift and output level. | |
| Everything here is decode-side; the checkpoint is untouched. Three controls, each switchable so | |
| ablations can be run with and without each one: | |
| --stop {sampled,controlled} | |
| sampled = the trained behaviour: <aeos> is one class in the multinomial cb0 draw (a | |
| well-fit model spreads the stop hazard over trailing-silence frames, so a | |
| non-aeos draw keeps talking = babble tails). | |
| controlled = <aeos> is masked out of the sampled draw; the stop fires when, after a | |
| text-conditioned minimum length, argmax == <aeos>, or p(<aeos>) >= --tau, or the | |
| cumulative stop hazard 1 - prod(1 - p_t) >= --hazard. The cap is the same | |
| per-language text -> frames prior at its p99 residual instead of a flat 1500. | |
| --chunk-s S | |
| chunked generation with RE-PROMPTING (against long-rollout drift): the text is split at sentence terminals into pieces | |
| of about S seconds (prior-estimated). Chunk 1 uses the requested arm. Chunk k > 1 is generated | |
| in a FRESH backbone context whose in-context reference is the previous chunk's text plus the | |
| last --ref-tail-frames of its generated codes (the same speaker memory throughout). Drift | |
| accumulated over long rollouts resets at each chunk boundary. 0 disables. | |
| --level {none,peaknorm,limiter} | |
| post-decode level control (against loudness ramps): peaknorm = the corpus's own consumer rule (attenuate so the | |
| peak sits at --ceiling-db, never boost); limiter = look-ahead peak limiter at --ceiling-db | |
| (5 ms look-ahead, 80 ms release) so a loudness ramp cannot flat-top. | |
| Rails are computed on the RAW decode (so the ramp is visible) and on the processed output: clip rate | |
| (|x| >= 0.999), flat-top runs (>= 3 consecutive pinned samples), pinned fraction, RMS dBFS, per-second | |
| RMS slope (dB/s) and end-minus-start (last 3 s vs first 3 s), HF (> 6 kHz) energy fraction, silent | |
| (>= 3 s of 1-s windows under RMS 0.005), frames / cap / aeos, duration vs the prior's expectation. | |
| Targets: by default a LONG text per row is built by concatenating the target texts of other eval-set | |
| rows of the same language until --long-chars characters (the corpus targets are ~10 s; drift and | |
| loudness ramps need 25-40 s to show). --targets orig uses each row's own text instead. | |
| python scripts/s2_serve_controls.py --config <eval_config.yaml> --set <eval_set.pt> \ | |
| --ckpt <run>/ckpt/<step> --out-dir <out_dir> --langs hi,te --per-lang 2 \ | |
| --views spk_only,icl_spk --configs baseline,controlled [--device cuda:0] [--q-serve 16] | |
| The checkpoint is loaded by the SHARED strict loader tts/model/load_s3.py, so an | |
| S3 checkpoint builds and loads its speaker anchor and an S2 checkpoint does not — a disagreement | |
| raises instead of silently serving a different model. --model-config <yaml> names the run config | |
| that built the checkpoint (default: auto-discovered next to it, else --config); | |
| --require-anchor {yes,no} asserts the arm before anything is built. The loaded model is echoed as | |
| `CKPT_LOADED {...}` and stored in <out-dir>/run.json under "ckpt_provenance". | |
| Configs (--configs): baseline = sampled stop, no chunking, no level control (the plain trained decode); | |
| controlled = controlled stop + chunk-s (re-prompt on the ORIGINAL reference) + limiter; stop_only / | |
| chunk_only / level_only ablations. | |
| Further controls (added after listening tests reported "static / double voice from the start, worse by | |
| the end, pure static once the transcript is done"): | |
| --greedy-from L depth levels >= L are taken by argmax instead of sampled (the fine codebooks | |
| sampled at temperature decode as broadband hiss on EVERY frame, and that hiss is | |
| fed back into the model's context). Configs greedy9 / greedy17 isolate it. | |
| --deep-temp T levels >= 9 sampled at temperature T (levels 1-8 keep --temp). | |
| --reprompt {orig,own} chunk k > 1 re-prompts on the row's ORIGINAL reference block (orig, default: | |
| clean human audio; for spk_only a fresh embedding-only context) or on the | |
| model's own previous chunk (own, the original behaviour, which carries the drift | |
| forward — te icl went 0.5% -> 8.4% HF by the last chunk in a short test). | |
| --static-stop H every --static-check frames after the minimum length, decode the last 3 s and | |
| stop when the HF fraction of the last 2 s exceeds H twice in a row (the | |
| "transcript is over, only static remains" state), or when it is silent | |
| (RMS < -45 dBFS). 0 disables. | |
| config "full" = controlled stop + chunking (orig) + limiter + greedy-from 9 + static-stop 0.05. | |
| """ | |
| import argparse | |
| import collections | |
| import json | |
| import os | |
| import re | |
| import sys | |
| import time | |
| import wave | |
| sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) | |
| import numpy as np | |
| import torch | |
| import yaml | |
| KIND_TEXT, KIND_AUDIO = 0, 1 | |
| SEED0 = 20260901 | |
| SR = 24000 | |
| FPS = 12.5 | |
| # per-language text -> frames prior (frames = a + b * chars), fitted on 40 shards of the | |
| # training corpus (the --prior file, configs/text_frames_prior.json by default). resid_p99 = 99th percentile of the | |
| # residual, used as the cap margin. Fallback when the JSON is not present. | |
| PRIOR_DEFAULT = {"a": 2.0, "b": 0.72, "resid_p99": 90.0} | |
| TERMINALS = "।.!?;।॥" | |
| def load_prior(path): | |
| prior = {} | |
| if path and os.path.exists(path): | |
| for lang, d in json.load(open(path)).items(): | |
| prior[lang] = {"a": float(d["affine_a"]), "b": float(d["affine_b"]), | |
| "resid_p99": float(d["resid_p99"])} | |
| return prior | |
| def expected_frames(prior, lang, nchar): | |
| p = prior.get(lang, PRIOR_DEFAULT) | |
| return p["a"] + p["b"] * nchar, p["resid_p99"] | |
| # ----------------------------------------------------------------------------- text | |
| def split_text(text, lang, prior, chunk_s): | |
| """Split at sentence terminals into pieces of about chunk_s seconds (prior-estimated). | |
| Falls back to commas, then whitespace, for run-on text.""" | |
| if chunk_s <= 0: | |
| return [text] | |
| target = chunk_s * FPS | |
| # sentence pieces keep their terminal punctuation | |
| pieces = [p.strip() for p in re.split(f"(?<=[{re.escape(TERMINALS)}])\\s+", text) if p.strip()] | |
| out, cur = [], "" | |
| for p in pieces: | |
| cand = (cur + " " + p).strip() if cur else p | |
| if cur and expected_frames(prior, lang, len(cand))[0] > target: | |
| out.append(cur) | |
| cur = p | |
| else: | |
| cur = cand | |
| if cur: | |
| out.append(cur) | |
| # a single piece longer than the budget: split at commas / whitespace | |
| final = [] | |
| for piece in out: | |
| if expected_frames(prior, lang, len(piece))[0] <= target * 1.3: | |
| final.append(piece) | |
| continue | |
| words = re.split(r"(?<=[,،])\s+|\s+", piece) | |
| cur = "" | |
| for w in words: | |
| cand = (cur + " " + w).strip() if cur else w | |
| if cur and expected_frames(prior, lang, len(cand))[0] > target: | |
| final.append(cur) | |
| cur = w | |
| else: | |
| cur = cand | |
| if cur: | |
| final.append(cur) | |
| return final | |
| def row_text(ex, tok, ctl_ids): | |
| """Target-span text of an eval-set view, decoded from its text tokens without control tokens.""" | |
| t, k = np.asarray(ex["tokens"]), np.asarray(ex["kind"]) | |
| ts = int(ex.get("target_start", 0)) | |
| ids = [int(x) for x in t[ts:][k[ts:] == KIND_TEXT] if int(x) not in ctl_ids] | |
| return tok.decode(ids).strip() | |
| # ----------------------------------------------------------------------------- prompts | |
| def make_example(packer, prefix, text, lang, spk_refs, spk_null): | |
| """Pack a prompt: an optional reference prefix (tokens/kind/codes/loss_mask arrays, already | |
| packed) followed by the target span for `text` (one dummy audio frame; the prompt is truncated | |
| at it).""" | |
| from tts.tokens import N_Q | |
| cut = {"text_final": text, "lang_final": lang, "codes": np.zeros((N_Q, 1), np.int64), "row": 0} | |
| first = prefix is None | |
| t, k, c, m = packer._span(cut, first=first, cont=False) | |
| if prefix is not None: | |
| t = np.concatenate([prefix["tokens"], t]) | |
| k = np.concatenate([prefix["kind"], k]) | |
| c = np.concatenate([prefix["codes"], c]) | |
| m = np.concatenate([prefix["loss_mask"], m]) | |
| ts = int(prefix["tokens"].shape[0]) | |
| else: | |
| ts = 0 | |
| return {"tokens": t, "kind": k, "codes": c, "loss_mask": m, "spk_refs": spk_refs, | |
| "spk_null": spk_null, "mode": "serve", "n_cuts": 1, "target_rows": [0], | |
| "ref_rows": [], "target_start": ts, "lang": lang} | |
| def ref_prefix_from_view(ex): | |
| """The packed reference block of an icl view (everything before target_start).""" | |
| ts = int(ex.get("target_start", 0)) | |
| if ts == 0: | |
| return None | |
| return {key: np.asarray(ex[key])[:ts] for key in ("tokens", "kind", "codes", "loss_mask")} | |
| def ref_prefix_from_generated(packer, text, lang, codes_qF, q_serve, tail_frames): | |
| """Re-prompting reference: the previous chunk's text + the tail of its generated codes, | |
| packed as a first (lang-tagged) reference span. Levels >= q_serve are zero; the prompt is | |
| composed with q_use = q_serve so they are never read.""" | |
| from tts.tokens import N_Q | |
| codes = codes_qF[:, -tail_frames:] if tail_frames > 0 else codes_qF | |
| full = np.zeros((N_Q, codes.shape[1]), np.int64) | |
| full[: codes.shape[0]] = codes.cpu().numpy() if torch.is_tensor(codes) else codes | |
| cut = {"text_final": text, "lang_final": lang, "codes": full, "row": 0} | |
| t, k, c, m = packer._span(cut, first=True, cont=False) | |
| return {"tokens": t, "kind": k, "codes": c, "loss_mask": m} | |
| def prompt_from_example(ex, model, dev, dt, q_serve): | |
| """Collate one example with q_use = q_serve, cut at the target's first audio position and | |
| compose the backbone inputs. Returns (prompt_emb, n_prompt_tokens). | |
| The prompt goes through collate_v2 (PackerV2._span -> collate_v2), so the S3 per-span | |
| geometry is built by the SAME code as in training and is only sliced here — never | |
| hand-made. The speaker vector the anchor reads is collate_v2's `spk_refs`/`spk_ref_mask` | |
| /`spk_null` block, i.e. the same mean-pooled ECAPA-192 input as in training | |
| (tts/model/anchors.py:52-65).""" | |
| from tts.data.packing_v2 import collate_v2 | |
| from tts.model.load_s3 import truncate_prompt_batch | |
| b = collate_v2([ex], q_choices=(q_serve,), q_rng=np.random.default_rng(0)) | |
| kind = b["kind"][0] | |
| ts = int(ex.get("target_start", 0)) | |
| a = (kind == KIND_AUDIO).nonzero().flatten() | |
| a = a[a >= ts] | |
| assert a.numel(), "prompt has no target audio position" | |
| a0 = int(a[0]) | |
| # slice the span tables with the columns. The target span's <|audio|> | |
| # sentinel is at a0-1 (the last prompt position, whose hidden state generates frame 0) | |
| # and MUST survive the cut, or the anchor never reaches the generating position. | |
| pb = truncate_prompt_batch(b, a0, | |
| require_target_span=model.speaker_anchor is not None) | |
| pb = {k: (v.to(dev, dtype=dt) if torch.is_tensor(v) and v.is_floating_point() | |
| else (v.to(dev) if torch.is_tensor(v) else v)) for k, v in pb.items()} | |
| return model.compose_inputs(pb), a0 | |
| # ----------------------------------------------------------------------------- depth sampling | |
| def depth_generate(depth, h_frame, cb0, spk_mem, q_serve, temp, top_k, cfg_scale, ref_mask, | |
| greedy_from=None, deep_temp=None, deep_from=9): | |
| """DepthV2.generate with per-level control: levels >= greedy_from by argmax; levels >= deep_from | |
| at deep_temp (else temp). Same prefix construction and CFG as the model's own generate.""" | |
| from tts.model.depth_v2 import GROUP_OF_LEVEL | |
| from tts.tokens import AEOS_CLASS | |
| if torch.any(cb0 >= AEOS_CLASS) or torch.any(cb0 < 0): | |
| raise ValueError("aeos/invalid cb0 must be handled by the caller") | |
| B = h_frame.shape[0] | |
| null_mem = None | |
| if cfg_scale > 0: | |
| null_mem = depth.null_mem.to(spk_mem.dtype).expand(B, spk_mem.shape[1], -1) | |
| outs = [cb0] | |
| for level in range(1, q_serve): | |
| codes_prev = torch.stack(outs, 1) | |
| x = depth._prefix(h_frame, codes_prev, level) | |
| h = depth._run(x, GROUP_OF_LEVEL[level], spk_mem, ref_mask) | |
| logits = depth.heads[level - 1](h).float() | |
| if cfg_scale > 0: | |
| hn = depth._run(x, GROUP_OF_LEVEL[level], null_mem) | |
| ln = depth.heads[level - 1](hn).float() | |
| logits = ln + cfg_scale * (logits - ln) | |
| if greedy_from is not None and level >= greedy_from: | |
| outs.append(logits.argmax(-1)) | |
| continue | |
| t = deep_temp if (deep_temp is not None and level >= deep_from) else temp | |
| logits = logits / t | |
| if top_k: | |
| kth = logits.topk(top_k, -1).values[:, -1:] | |
| logits = logits.masked_fill(logits < kth, -float("inf")) | |
| outs.append(torch.multinomial(logits.softmax(-1), 1).squeeze(-1)) | |
| return torch.stack(outs, -1) | |
| def hf_of(w, sr=SR): | |
| if w.size < 256: | |
| return 0.0 | |
| sp = np.abs(np.fft.rfft(w)) ** 2 | |
| f = np.fft.rfftfreq(len(w), 1 / sr) | |
| return float(sp[f >= 6000].sum() / (sp.sum() + 1e-12)) | |
| # ----------------------------------------------------------------------------- session | |
| class ServeSession: | |
| """Backbone KV session + DepthV2 frame decoding with speaker memory, plus the stop policy.""" | |
| def __init__(self, model, spk_mem, ref_mask, q_serve=16, temp=0.7, top_k=64, cfg_scale=0.0, | |
| device="cuda:0", greedy_from=None, deep_temp=None, mimi=None, static_stop=0.0, | |
| static_check=25): | |
| self.m, self.dev = model, device | |
| self.mem, self.rmask = spk_mem, ref_mask | |
| self.q, self.temp, self.top_k, self.cfg = q_serve, temp, top_k, cfg_scale | |
| self.greedy_from, self.deep_temp = greedy_from, deep_temp | |
| self.mimi, self.static_stop, self.static_check = mimi, static_stop, static_check | |
| self.past = None | |
| def _static_state(self, out): | |
| """Decode the last 3 s of generated codes; return (hf of the last 2 s, rms dBFS).""" | |
| codes = torch.stack(out[-38:], -1)[None].to(self.dev).long() | |
| w = self.mimi.decode(codes)[0][0, 0].float().cpu().numpy() | |
| w = w[-2 * SR:] | |
| rms = 20 * np.log10(np.sqrt((w ** 2).mean()) + 1e-9) if w.size else -99.0 | |
| return hf_of(w), float(rms) | |
| def _step(self, emb): | |
| h, self.past = self.m.backbone(emb, past_key_values=self.past, use_cache=True) | |
| return h[:, -1] | |
| def _cb0_probs(self, h): | |
| logits = self.m.backbone.cb0_logits(h).float() / self.temp | |
| if self.top_k: | |
| kth = logits.topk(self.top_k, -1).values[:, -1:] | |
| logits = logits.masked_fill(logits < kth, -float("inf")) | |
| return logits.softmax(-1)[0] | |
| def speak(self, prompt_emb, cap, stop="sampled", min_frames=0, tau=0.35, hazard=0.6): | |
| from tts.tokens import AEOS_CLASS | |
| h = self._step(prompt_emb.to(self.dev)) | |
| out, p_trace, H = [], [], 0.0 | |
| reason = "cap" | |
| static_hits = 0 | |
| for t in range(cap): | |
| if (self.static_stop > 0 and self.mimi is not None and t >= max(min_frames, 38) | |
| and t % self.static_check == 0): | |
| hf, rms = self._static_state(out) | |
| if rms < -45.0: | |
| reason = "silence" | |
| break | |
| static_hits = static_hits + 1 if hf >= self.static_stop else 0 | |
| if static_hits >= 2: | |
| reason = "static" | |
| break | |
| p = self._cb0_probs(h) | |
| p_aeos = float(p[AEOS_CLASS]) | |
| p_trace.append(p_aeos) | |
| if stop == "sampled": | |
| cb0 = int(torch.multinomial(p, 1)) | |
| if cb0 == AEOS_CLASS: | |
| reason = "aeos" | |
| break | |
| else: | |
| if t >= min_frames: | |
| H = 1.0 - (1.0 - H) * (1.0 - p_aeos) | |
| if int(p.argmax()) == AEOS_CLASS: | |
| reason = "aeos_argmax" | |
| break | |
| if p_aeos >= tau: | |
| reason = "aeos_tau" | |
| break | |
| if H >= hazard: | |
| reason = "aeos_hazard" | |
| break | |
| q = p.clone() | |
| q[AEOS_CLASS] = 0.0 | |
| q = q / q.sum() | |
| cb0 = int(torch.multinomial(q, 1)) | |
| frame = depth_generate( | |
| self.m.depth, h.float(), torch.tensor([cb0], device=self.dev), self.mem, self.q, | |
| self.temp, self.top_k, self.cfg, self.rmask, greedy_from=self.greedy_from, | |
| deep_temp=self.deep_temp)[0] | |
| out.append(frame) | |
| fb = self.m.audio_in.frames(frame[None, :], frame.shape[0]) | |
| h = self._step(fb[None].to(prompt_emb.dtype)) | |
| if reason.startswith("aeos"): | |
| self._step(self.m.audio_in.aeos(1)[None].to(prompt_emb.dtype)) | |
| codes = (torch.stack(out, -1) if out | |
| else torch.zeros(self.q, 0, dtype=torch.long, device=self.dev)) | |
| return codes, reason, p_trace | |
| # ----------------------------------------------------------------------------- level + rails | |
| def peak_norm(x, ceiling_db=-1.0): | |
| ceil = 10 ** (ceiling_db / 20) | |
| peak = float(np.abs(x).max()) if x.size else 0.0 | |
| g = min(1.0, ceil / peak) if peak > 0 else 1.0 | |
| return (x * g).astype(np.float32), g | |
| def limiter(x, ceiling_db=-1.0, lookahead_s=0.005, release_s=0.08, sr=SR): | |
| """Look-ahead peak limiter: gain drops instantly (within the look-ahead window) to keep the | |
| windowed peak at the ceiling and recovers with an exponential release.""" | |
| from scipy.ndimage import maximum_filter1d | |
| if x.size == 0: | |
| return x.astype(np.float32), 1.0 | |
| ceil = 10 ** (ceiling_db / 20) | |
| L = max(1, int(lookahead_s * sr)) | |
| env = maximum_filter1d(np.abs(x), size=2 * L + 1, mode="nearest") | |
| g = np.minimum(1.0, ceil / np.maximum(env, 1e-9)) | |
| rel = float(np.exp(-1.0 / (release_s * sr))) | |
| out_g = np.empty_like(g) | |
| cur = 1.0 | |
| for n in range(len(g)): | |
| gn = g[n] | |
| cur = gn if gn < cur else gn + (cur - gn) * rel | |
| out_g[n] = cur | |
| y = np.clip(x * out_g, -1.0, 1.0).astype(np.float32) | |
| return y, float(out_g.min()) | |
| def rails(w, sr=SR): | |
| a = np.abs(w) | |
| n = len(w) | |
| pinned = a >= 0.999 | |
| runs = 0 | |
| if n: | |
| d = np.diff(np.concatenate([[0], pinned.astype(np.int8), [0]])) | |
| starts, ends = np.where(d == 1)[0], np.where(d == -1)[0] | |
| runs = int(((ends - starts) >= 3).sum()) | |
| sec = w[: (n // sr) * sr].reshape(-1, sr) if n >= sr else np.zeros((0, sr)) | |
| rms_s = 20 * np.log10(np.sqrt((sec ** 2).mean(1)) + 1e-9) if len(sec) else np.array([]) | |
| slope = float(np.polyfit(np.arange(len(rms_s)), rms_s, 1)[0]) if len(rms_s) >= 4 else None | |
| end_start = (float(rms_s[-3:].mean() - rms_s[:3].mean()) if len(rms_s) >= 6 else None) | |
| if n: | |
| sp = np.abs(np.fft.rfft(w)) ** 2 | |
| f = np.fft.rfftfreq(n, 1 / sr) | |
| hf = float(sp[f >= 6000].sum() / (sp.sum() + 1e-12)) | |
| else: | |
| hf = None | |
| thirds = None | |
| if len(sec) >= 3: | |
| hfs = [hf_of(x) for x in sec] | |
| t3 = len(hfs) // 3 | |
| thirds = [round(float(np.mean(hfs[:t3])), 4), round(float(np.mean(hfs[t3:2 * t3])), 4), | |
| round(float(np.mean(hfs[2 * t3:])), 4)] | |
| silent_run = 0 | |
| best = 0 | |
| for r in rms_s: | |
| silent_run = silent_run + 1 if r < 20 * np.log10(0.005) else 0 | |
| best = max(best, silent_run) | |
| return { | |
| "dur_s": round(n / sr, 2), | |
| "peak_dbfs": round(float(20 * np.log10(a.max() + 1e-9)), 2) if n else None, | |
| "rms_dbfs": round(float(20 * np.log10(np.sqrt((w ** 2).mean()) + 1e-9)), 2) if n else None, | |
| "clip_frac": round(float(pinned.mean()), 5) if n else None, | |
| "flat_top_runs": runs, | |
| "rms_slope_db_s": None if slope is None else round(slope, 3), | |
| "rms_end_minus_start_db": None if end_start is None else round(end_start, 2), | |
| "hf_frac_gt6k": None if hf is None else round(hf, 4), | |
| "hf_thirds": thirds, | |
| "silent_ge3s": bool(best >= 3), | |
| } | |
| def save_wav(p, w): | |
| with wave.open(p, "wb") as fh: | |
| fh.setnchannels(1) | |
| fh.setsampwidth(2) | |
| fh.setframerate(SR) | |
| fh.writeframes((np.clip(w, -1, 1) * 32767).astype("<i2").tobytes()) | |
| CONFIGS = { | |
| # stop / chunk / level are the base controls; greedy_from / deep_temp / reprompt / static_stop the further ones | |
| "baseline": {"stop": "sampled", "chunk": False, "level": "none"}, | |
| "controlled": {"stop": "controlled", "chunk": True, "level": "limiter", "reprompt": "orig"}, | |
| "stop_only": {"stop": "controlled", "chunk": False, "level": "none"}, | |
| "chunk_only": {"stop": "sampled", "chunk": True, "level": "none", "reprompt": "orig"}, | |
| "chunk_own": {"stop": "sampled", "chunk": True, "level": "none", "reprompt": "own"}, | |
| "level_only": {"stop": "sampled", "chunk": False, "level": "limiter"}, | |
| "greedy17": {"stop": "sampled", "chunk": False, "level": "none", "greedy_from": 17}, | |
| "greedy9": {"stop": "sampled", "chunk": False, "level": "none", "greedy_from": 9}, | |
| "deep05": {"stop": "sampled", "chunk": False, "level": "none", "deep_temp": 0.5}, | |
| "full": {"stop": "controlled", "chunk": True, "level": "limiter", "reprompt": "orig", | |
| "greedy_from": 9, "static_stop": 0.05}, | |
| } | |
| def main(): | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("--config", required=True) | |
| ap.add_argument("--set", dest="set_path", required=True) | |
| ap.add_argument("--ckpt", required=True) | |
| ap.add_argument("--out-dir", required=True) | |
| ap.add_argument("--langs", default="hi,te") | |
| ap.add_argument("--per-lang", type=int, default=2) | |
| ap.add_argument("--views", default="spk_only,icl_spk") | |
| ap.add_argument("--configs", default="baseline,controlled") | |
| ap.add_argument("--targets", choices=["long", "orig"], default="long") | |
| ap.add_argument("--long-chars", type=int, default=450) | |
| ap.add_argument("--temp", type=float, default=0.7) | |
| ap.add_argument("--top-k", type=int, default=64) | |
| ap.add_argument("--q-serve", type=int, default=16) | |
| ap.add_argument("--cfg", type=float, default=0.0) | |
| ap.add_argument("--hard-cap", type=int, default=1500) | |
| ap.add_argument("--tau", type=float, default=0.35) | |
| ap.add_argument("--hazard", type=float, default=0.6) | |
| ap.add_argument("--min-frac", type=float, default=0.5, | |
| help="minimum length as a fraction of the prior's expected frames") | |
| ap.add_argument("--chunk-s", type=float, default=16.0) | |
| ap.add_argument("--ref-tail-frames", type=int, default=150) | |
| ap.add_argument("--level", default=None, help="override the config's level control") | |
| ap.add_argument("--ceiling-db", type=float, default=-1.0) | |
| ap.add_argument("--prior", default=os.path.join(os.path.dirname(__file__), "..", "configs", | |
| "text_frames_prior.json")) | |
| ap.add_argument("--greedy-from", type=int, default=None, help="override: levels >= L by argmax") | |
| ap.add_argument("--deep-temp", type=float, default=None, help="override: levels >= 9 at this temp") | |
| ap.add_argument("--reprompt", choices=["orig", "own"], default=None, help="override chunk re-prompt source") | |
| ap.add_argument("--static-stop", type=float, default=None, help="override: HF fraction that stops (0 off)") | |
| ap.add_argument("--static-check", type=int, default=25) | |
| ap.add_argument("--device", default="cuda:0") | |
| ap.add_argument("--threads", type=int, default=8) | |
| ap.add_argument("--seed", type=int, default=SEED0) | |
| # the run config that BUILT the checkpoint (speaker anchor etc). Default = | |
| # auto-discovered next to the checkpoint, else --config. --require-anchor asserts the arm. | |
| ap.add_argument("--model-config", default=None) | |
| ap.add_argument("--require-anchor", choices=["yes", "no"], default=None) | |
| A = ap.parse_args() | |
| torch.set_num_threads(A.threads) | |
| cfg = yaml.safe_load(open(A.config)) | |
| dev = A.device | |
| dt = torch.bfloat16 if dev.startswith("cuda") else torch.float32 | |
| os.makedirs(A.out_dir + "/wav", exist_ok=True) | |
| prior = load_prior(A.prior) | |
| if not prior: | |
| print(f"WARNING: no prior at {A.prior}; using the default {PRIOR_DEFAULT}", flush=True) | |
| views = A.views.split(",") | |
| configs = A.configs.split(",") | |
| langs = A.langs.split(",") | |
| from transformers import AutoTokenizer | |
| from tts.data.packing_v2 import PackerV2 | |
| from tts.model.load_s3 import load_model_for_eval | |
| from tts.tokens import extend_tokenizer | |
| tok = AutoTokenizer.from_pretrained(cfg["model"]["backbone"]) | |
| tok, ids = extend_tokenizer(tok) | |
| ctl_ids = set(int(v) for v in ids.values()) | |
| packer = PackerV2(tok, ids, max_tokens=int(cfg["data"].get("max_tokens", 8192))) | |
| if dev.startswith("cuda"): | |
| torch.backends.cuda.enable_cudnn_sdp(False) | |
| t0 = time.time() | |
| # shared STRICT loader (tts/model/load_s3.py). A loader that always built an anchor-free | |
| # model and never asked for the checkpoint's extra tensors would silently serve a different | |
| # model from an anchored checkpoint. | |
| model, prov = load_model_for_eval( | |
| A.ckpt, device=dev, dtype=dt, cfg=cfg, config=A.model_config, | |
| amortize_frac=1.0, # unused here (generation never calls forward()) | |
| require_anchor=None if A.require_anchor is None | |
| else A.require_anchor == "yes") | |
| model.depth.float() | |
| print(f"model loaded in {time.time() - t0:.0f}s on {dev} " | |
| f"(anchor={prov['anchor']}, update={prov['update']})", flush=True) | |
| from transformers import MimiModel | |
| mimi = MimiModel.from_pretrained("kyutai/mimi").to(dev).eval() | |
| data = torch.load(A.set_path, weights_only=False) | |
| by_lang = collections.defaultdict(list) | |
| for i, row in enumerate(data["rows"]): | |
| by_lang[row["lang"]].append((i, row)) | |
| picked = [] | |
| for lang in langs: | |
| rows = by_lang.get(lang, []) | |
| picked += [(lang, i, r) for i, r in rows[: A.per_lang]] | |
| # long targets: concatenate OTHER rows' target texts of the same language | |
| texts = {lang: [row_text(r["views"]["spk_only"], tok, ctl_ids) for _, r in by_lang[lang]] | |
| for lang in langs} | |
| recs = [] | |
| for lang, i, row in picked: | |
| own = row_text(row["views"]["spk_only"], tok, ctl_ids) | |
| if A.targets == "long": | |
| pool = [t for t in texts[lang] if t != own] or [own] | |
| text, k = own, 0 | |
| while len(text) < A.long_chars: | |
| text = text + " " + pool[k % len(pool)] | |
| k += 1 | |
| else: | |
| text = own | |
| exp_frames, _resid = expected_frames(prior, lang, len(text)) | |
| for view in views: | |
| ex0 = row["views"][view] | |
| spk_refs = np.asarray(ex0["spk_refs"], np.float32) | |
| spk_null = bool(ex0.get("spk_null", True)) | |
| refs_t = torch.as_tensor(spk_refs, dtype=torch.float32, device=dev)[None] | |
| null_t = torch.tensor([spk_null], device=dev) | |
| mem = model.depth.spk_memory(refs_t, "ecapa", null_mask=null_t) | |
| rmask = torch.ones(1, refs_t.shape[1], dtype=torch.bool, device=dev) | |
| icl_prefix = ref_prefix_from_view(ex0) if view.startswith("icl") else None | |
| for cname in configs: | |
| c = dict(CONFIGS[cname]) | |
| if A.level is not None: | |
| c["level"] = A.level | |
| greedy_from = A.greedy_from if A.greedy_from is not None else c.get("greedy_from") | |
| deep_temp = A.deep_temp if A.deep_temp is not None else c.get("deep_temp") | |
| reprompt = A.reprompt or c.get("reprompt", "orig") | |
| static_stop = A.static_stop if A.static_stop is not None else c.get("static_stop", 0.0) | |
| chunks = split_text(text, lang, prior, A.chunk_s if c["chunk"] else 0) | |
| torch.manual_seed(A.seed + 7919 * i + 101 * views.index(view)) | |
| gen_codes, reasons, frames_per_chunk, caps = [], [], [], [] | |
| prev_text, prev_codes = None, None | |
| tg = time.time() | |
| for ci, ctext in enumerate(chunks): | |
| if ci == 0 or reprompt == "orig": | |
| prefix = icl_prefix | |
| else: | |
| prefix = ref_prefix_from_generated(packer, prev_text, lang, prev_codes, | |
| A.q_serve, A.ref_tail_frames) | |
| ex = make_example(packer, prefix, ctext, lang, spk_refs, spk_null) | |
| emb, _ = prompt_from_example(ex, model, dev, dt, A.q_serve) | |
| ef, er = expected_frames(prior, lang, len(ctext)) | |
| if c["stop"] == "controlled": | |
| cap = min(A.hard_cap, int(ef + er + 10)) | |
| min_frames = max(4, int(A.min_frac * ef)) | |
| else: | |
| cap = A.hard_cap | |
| min_frames = 0 | |
| sess = ServeSession(model, mem, rmask, q_serve=A.q_serve, temp=A.temp, | |
| top_k=A.top_k, cfg_scale=A.cfg if not spk_null else 0.0, | |
| device=dev, greedy_from=greedy_from, deep_temp=deep_temp, | |
| mimi=mimi, static_stop=static_stop, | |
| static_check=A.static_check) | |
| codes, reason, _ = sess.speak(emb.to(dt), cap, stop=c["stop"], | |
| min_frames=min_frames, tau=A.tau, | |
| hazard=A.hazard) | |
| gen_codes.append(codes) | |
| reasons.append(reason) | |
| frames_per_chunk.append(int(codes.shape[-1])) | |
| caps.append(cap) | |
| prev_text, prev_codes = ctext, codes | |
| print(f" [{lang} r{i} {view} {cname}] chunk {ci + 1}/{len(chunks)} " | |
| f"{codes.shape[-1]}f cap {cap} stop={reason} " | |
| f"({time.time() - tg:.0f}s)", flush=True) | |
| codes = torch.cat(gen_codes, -1) if gen_codes else None | |
| recs.append({"lang": lang, "row": i, "view": view, "config": cname, "text": text, | |
| "nchar": len(text), "exp_frames": round(exp_frames, 1), | |
| "chunks": len(chunks), "frames": frames_per_chunk, "caps": caps, | |
| "stops": reasons, "gen_frames": int(codes.shape[-1]), | |
| "level": c["level"], "greedy_from": greedy_from, | |
| "deep_temp": deep_temp, "reprompt": reprompt if c["chunk"] else None, | |
| "static_stop": static_stop, "codes": codes.cpu()}) | |
| del model | |
| if dev.startswith("cuda"): | |
| torch.cuda.empty_cache() | |
| summary = collections.defaultdict(list) | |
| with torch.inference_mode(), open(A.out_dir + "/rails.jsonl", "w") as fh: | |
| for n, rec in enumerate(recs): | |
| codes = rec.pop("codes") | |
| if codes.shape[-1] == 0: | |
| w = np.zeros(0, np.float32) | |
| else: | |
| w = mimi.decode(codes[None].to(dev).long())[0][0, 0].float().cpu().numpy() | |
| raw = rails(w) | |
| if rec["level"] == "peaknorm": | |
| y, g = peak_norm(w, A.ceiling_db) | |
| elif rec["level"] == "limiter": | |
| y, g = limiter(w, A.ceiling_db) | |
| else: | |
| y, g = np.clip(w, -1, 1).astype(np.float32), 1.0 | |
| proc = rails(y) | |
| base = f"{A.out_dir}/wav/{n:03d}_{rec['lang']}_r{rec['row']}_{rec['view']}_{rec['config']}" | |
| save_wav(base + "_raw.wav", w) | |
| if rec["level"] != "none": | |
| save_wav(base + "_proc.wav", y) | |
| rec.update({"raw": raw, "proc": proc, "min_gain": round(float(g), 3), | |
| "dur_ratio_vs_prior": round(rec["gen_frames"] / max(1.0, rec["exp_frames"]), 3)}) | |
| fh.write(json.dumps(rec, ensure_ascii=False) + "\n") | |
| key = (rec["view"], rec["config"]) | |
| summary[key].append(rec) | |
| print(f"{n:03d} {rec['lang']} r{rec['row']} {rec['view']:8s} {rec['config']:10s} " | |
| f"{rec['gen_frames']:4d}f/{rec['exp_frames']:.0f} exp, chunks {rec['chunks']}, " | |
| f"stops {','.join(rec['stops'])} | RAW peak {raw['peak_dbfs']} rms {raw['rms_dbfs']} " | |
| f"clip {raw['clip_frac']} runs {raw['flat_top_runs']} slope {raw['rms_slope_db_s']} " | |
| f"hf {raw['hf_frac_gt6k']} | PROC peak {proc['peak_dbfs']} clip {proc['clip_frac']} " | |
| f"runs {proc['flat_top_runs']}", flush=True) | |
| print("\n=== SUMMARY (means over rows; RAW decode unless marked proc) ===") | |
| print(f"{'view':9s}{'config':11s}{'n':>3s}{'frames':>8s}{'dur/exp':>8s}{'peak':>7s}{'rms':>7s}" | |
| f"{'clip%':>7s}{'runs':>6s}{'slope':>7s}{'end-st':>7s}{'hf%':>6s}{'procclip%':>10s}{'procruns':>9s}") | |
| for (view, cname), rs in summary.items(): | |
| def mean(key, sub="raw", rs=rs): | |
| vals = [r[sub][key] for r in rs if r[sub][key] is not None] | |
| return sum(vals) / len(vals) if vals else float("nan") | |
| print(f"{view:9s}{cname:11s}{len(rs):3d}{sum(r['gen_frames'] for r in rs) / len(rs):8.0f}" | |
| f"{sum(r['dur_ratio_vs_prior'] for r in rs) / len(rs):8.2f}{mean('peak_dbfs'):7.2f}" | |
| f"{mean('rms_dbfs'):7.2f}{100 * mean('clip_frac'):7.3f}{mean('flat_top_runs'):6.0f}" | |
| f"{mean('rms_slope_db_s'):7.3f}{mean('rms_end_minus_start_db'):7.2f}" | |
| f"{100 * mean('hf_frac_gt6k'):6.2f}{100 * mean('clip_frac', 'proc'):10.3f}" | |
| f"{mean('flat_top_runs', 'proc'):9.0f}") | |
| json.dump({"args": vars(A), "n": len(recs), "ckpt_provenance": prov}, | |
| open(A.out_dir + "/run.json", "w"), indent=1) | |
| if __name__ == "__main__": | |
| main() | |