| |
|
|
| import argparse |
| import contextlib |
| import io |
| import json |
| import os |
| import sys |
| import time |
| from pathlib import Path |
|
|
| os.environ.setdefault("TQDM_DISABLE", "1") |
|
|
| _REPO_ROOT = Path(__file__).resolve().parent.parent |
| sys.path.insert(0, str(_REPO_ROOT)) |
| sys.path.insert(0, str(_REPO_ROOT / "scripts")) |
|
|
| from fish_codec_util import fish_repo_path |
| from gen_s1mini_synth import ( |
| EMOTION_TAG, |
| S1MiniTiktokenTokenizer, |
| load_ref_specs, |
| resolve_checkpoint, |
| ) |
|
|
| SEM_MAX = 4095 |
| RES_MAX = 1023 |
| FULL_VOCAB_MIN = 100_000 |
|
|
|
|
| class ForcedSemantic: |
| def __init__(self, real_sample, sem_begin, im_end_id, fast_temps=None): |
| self.real = real_sample |
| self.sem_begin = sem_begin |
| self.im_end_id = im_end_id |
| self.queue = None |
| self.calls = 0 |
| self.fast_temps = fast_temps |
| self.res = None |
| self.fast_calls = 0 |
|
|
| def arm(self, q0, residuals=None): |
| self.queue = list(q0) |
| self.calls = 0 |
| self.res = residuals |
| self.fast_calls = 0 |
|
|
| @property |
| def frames_forced(self): |
| return (self.calls + 1) // 2 |
|
|
| def __call__(self, logits, temperature, top_p, top_k): |
| import torch |
| if logits.shape[-1] >= FULL_VOCAB_MIN: |
| idx = self.calls // 2 |
| self.calls += 1 |
| if self.queue is not None and idx < len(self.queue): |
| tid = self.sem_begin + int(self.queue[idx]) |
| else: |
| tid = self.im_end_id |
| return (torch.tensor([tid], device=logits.device, |
| dtype=torch.long), None) |
| frame, level = divmod(self.fast_calls, 9) |
| self.fast_calls += 1 |
| if self.res is not None and frame < len(self.res[0]): |
| return (torch.tensor([int(self.res[level][frame])], |
| device=logits.device, |
| dtype=torch.long), None) |
| if self.fast_temps is not None: |
| temperature = torch.as_tensor(self.fast_temps[level], |
| device=logits.device, |
| dtype=logits.dtype) |
| if logits.shape[-1] > RES_MAX + 1: |
| logits = logits.clone() |
| logits[..., RES_MAX + 1:] = float("-inf") |
| return self.real(logits, temperature, top_p, top_k) |
|
|
|
|
| def load_rows(path): |
| rows = [] |
| for line in open(path, encoding="utf-8"): |
| if line.strip(): |
| rows.append(json.loads(line)) |
| return rows |
|
|
|
|
| def resume_state(path): |
| done, kept = set(), [] |
| if not Path(path).exists(): |
| return done, kept |
| for line in open(path, encoding="utf-8"): |
| try: |
| j = json.loads(line)["renderer"]["row"] |
| except (ValueError, KeyError, TypeError): |
| continue |
| if j not in done: |
| done.add(j) |
| kept.append(line.rstrip("\n")) |
| return done, kept |
|
|
|
|
| def parse_voice_map(arg, voices): |
| vm = {} |
| if not arg: |
| return vm |
| for part in arg.split(","): |
| k, v = part.split("=", 1) |
| if v not in voices: |
| raise SystemExit(f"--voice_map: unknown voice {v!r}; " |
| f"available: {sorted(voices)[:12]}...") |
| vm[int(k)] = v |
| return vm |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser(description=__doc__.splitlines()[0]) |
| ap.add_argument("--codes", required=True, |
| help="cell jsonl with speech_tokens (q0) + prompt per row") |
| ap.add_argument("--out", required=True) |
| ap.add_argument("--ref_dir", default="data/s1mini_refs") |
| ap.add_argument("--ref_manifest", default=None) |
| ap.add_argument("--voice_map", default=None, |
| help="'4200=<ref>,4201=<ref>' — speaker_id to reference " |
| "voice. Unmapped speakers cycle the ref list " |
| "DETERMINISTICALLY by speaker id, so a speaker keeps " |
| "one voice across cells (the flip contrast needs " |
| "identical voices per row across cells).") |
| ap.add_argument("--emotion_tag", action="store_true", default=True) |
| ap.add_argument("--no_emotion_tag", dest="emotion_tag", |
| action="store_false", |
| help="drop the '(emotion)' inline marker from the text " |
| "conditioning. Same tag across a row's cells either " |
| "way, so the flip contrast is never confounded.") |
| ap.add_argument("--checkpoint", default=None) |
| ap.add_argument("--device", default="cuda", choices=["cuda", "mps", "cpu"]) |
| ap.add_argument("--limit", type=int, default=0) |
| ap.add_argument("--force_residuals", action="store_true", |
| help="PLUMBING SELF-TEST: force q1..q9 from the input " |
| "rows too. Output must reproduce the GT roundtrip; " |
| "anything else means the harness is broken.") |
| ap.add_argument("--temp", type=float, default=1.0, |
| help="fast-stage sampling (repo default; the " |
| "GRPO-polished setting). Semantic is forced.") |
| ap.add_argument("--fast_temps", type=float, nargs=9, default=None, |
| help="PER-LEVEL fast-stage temps for q1..q9 (9 values), " |
| "overriding the flat --temp. q0 is forced, so ALL " |
| "articulation loss lives in the residuals; cooling " |
| "the coarse ones (q1-q2) sharpens phonetic detail " |
| "while leaving the fine ones hot preserves texture.") |
| ap.add_argument("--top_p", type=float, default=0.9) |
| ap.add_argument("--top_k", type=int, default=30) |
| ap.add_argument("--seed", type=int, default=20260824, |
| help="row j is sampled with seed + j, so a row renders the same whatever else ran before it") |
| ap.add_argument("--resume", action="store_true", |
| help="keep the rows already in --out (those carrying renderer.row) and render only the rest; " |
| "a torn last line from an interrupted write is dropped") |
| ap.add_argument("--max_seq_len", type=int, default=4096) |
| args = ap.parse_args() |
|
|
| repo = fish_repo_path() |
| if not Path(repo).exists(): |
| raise SystemExit(f"fish-speech repo not found at {repo}") |
| if repo not in sys.path: |
| sys.path.insert(0, repo) |
|
|
| import soundfile as sf |
| import torch |
| from loguru import logger |
| logger.remove() |
| logger.add(sys.stderr, level="WARNING") |
|
|
| import fish_speech.models.text2semantic.inference as inf_mod |
| from fish_speech.models.text2semantic.inference import ( |
| generate_long, |
| init_model, |
| ) |
| try: |
| from fish_speech.models.text2semantic.inference import load_codec_model |
| except ImportError: |
| from fish_codec_util import load_fish_codec |
|
|
| def load_codec_model(checkpoint_path, device, precision=None): |
| os.environ.setdefault("FISH_CODEC_PTH", str(checkpoint_path)) |
| return load_fish_codec(device) |
| from fish_codec_util import encode_np |
|
|
| import inspect |
| gl_accepts_top_k = "top_k" in inspect.signature(generate_long).parameters |
|
|
| torch.manual_seed(args.seed) |
| if torch.cuda.is_available(): |
| torch.cuda.manual_seed(args.seed) |
|
|
| precision = torch.bfloat16 if args.device == "cuda" else torch.float32 |
| ckpt = resolve_checkpoint(args.checkpoint) |
| print(f"[render] checkpoint {ckpt} device={args.device}") |
| model, decode_one_token = init_model(ckpt, args.device, precision, |
| compile=False) |
| if (getattr(model, "tokenizer", None) is None |
| or getattr(model.config, "semantic_begin_id", 0) == 0): |
| tok = S1MiniTiktokenTokenizer(ckpt) |
| model.tokenizer = tok |
| model.config.semantic_begin_id = tok.semantic_begin_id |
| model.config.semantic_end_id = tok.semantic_end_id |
| print("[render] tiktoken adapter attached") |
| if args.max_seq_len < model.config.max_seq_len: |
| model.config.max_seq_len = args.max_seq_len |
| codec = load_codec_model(ckpt / "codec.pth", args.device, torch.float32) |
|
|
| im_end_id = model.tokenizer.get_token_id("<|im_end|>") |
| forced = ForcedSemantic(inf_mod.sample, model.config.semantic_begin_id, |
| im_end_id, fast_temps=args.fast_temps) |
| if args.fast_temps: |
| print(f"[render] per-level fast temps q1..q9: {args.fast_temps}") |
| inf_mod.sample = forced |
| print(f"[render] sample() patched (semantic_begin " |
| f"{model.config.semantic_begin_id}, im_end {im_end_id})") |
|
|
| ref_specs = load_ref_specs(args.ref_dir, args.ref_manifest) |
| voices = {r["name"]: r for r in ref_specs} |
| vmap = parse_voice_map(args.voice_map, voices) |
| names = sorted(voices) |
|
|
| ref_cache = {} |
|
|
| def get_ref(name): |
| if name not in ref_cache: |
| spec = voices[name] |
| wav, sr = sf.read(spec["wav"], dtype="float32") |
| codes = encode_np(codec, wav, sr, args.device) |
| if codes is None: |
| raise RuntimeError(f"reference encode failed: {spec['wav']}") |
| ref_cache[name] = (spec["text"], |
| torch.tensor(codes, dtype=torch.long)) |
| return ref_cache[name] |
|
|
| def voice_for(speaker_id): |
| if speaker_id in vmap: |
| return vmap[speaker_id] |
| return names[int(speaker_id) % len(names)] |
|
|
| rows = load_rows(args.codes) |
| if args.limit: |
| rows = rows[: args.limit] |
| out_path = Path(args.out) |
| out_path.parent.mkdir(parents=True, exist_ok=True) |
|
|
| done, kept = resume_state(out_path) if args.resume else (set(), []) |
| if args.resume: |
| print(f"[render] resume: {len(done)} rows already rendered, {len(rows) - len(done)} to go") |
| n_ok = n_mismatch = n_short = 0 |
| t0 = time.time() |
| with open(out_path, "w", encoding="utf-8") as fout: |
| for line in kept: |
| fout.write(line + "\n") |
| fout.flush() |
| for j, r in enumerate(rows): |
| if j in done: |
| continue |
| torch.manual_seed(args.seed + j) |
| if torch.cuda.is_available(): |
| torch.cuda.manual_seed(args.seed + j) |
| q0 = r["speech_tokens"] |
| if not q0 or max(q0) > SEM_MAX or min(q0) < 0: |
| raise SystemExit(f"row {j}: q0 out of the fish semantic range " |
| f"— wrong cell file?") |
| spk = r.get("speaker_id", 0) |
| vname = voice_for(spk) |
| ref_text, ref_tokens = get_ref(vname) |
| emo = r.get("emotion") or r.get("emotion_label") |
| text = r["prompt"] |
| if args.emotion_tag and emo and emo != "neutral": |
| text = f"({EMOTION_TAG.get(emo, emo)}) {text}" |
|
|
| forced.arm(q0, residuals=(r["residual_codes"] |
| if args.force_residuals else None)) |
| gen_kwargs = dict( |
| model=model, device=args.device, |
| decode_one_token=decode_one_token, |
| text=text, num_samples=1, |
| max_new_tokens=len(q0) + 8, |
| top_p=args.top_p, repetition_penalty=1.1, |
| temperature=args.temp, compile=False, |
| prompt_text=[ref_text], prompt_tokens=[ref_tokens], |
| ) |
| if gl_accepts_top_k: |
| gen_kwargs["top_k"] = args.top_k |
| parts = [] |
| t_row = time.time() |
| with contextlib.redirect_stdout(io.StringIO()): |
| for resp in generate_long(**gen_kwargs): |
| if resp.action == "sample" and resp.codes is not None: |
| parts.append(resp.codes.cpu()) |
| if args.device == "cuda": |
| torch.cuda.synchronize() |
| render_s = time.time() - t_row |
| if not parts: |
| print(f"[render] row {j}: EMPTY generation — skipped") |
| n_short += 1 |
| continue |
| codes = torch.cat(parts, dim=1)[:, : len(q0)] |
| sem = codes[0].tolist() |
| if sem != q0[: len(sem)]: |
| bad = next(i for i, (a, b) in enumerate(zip(sem, q0)) |
| if a != b) |
| print(f"[render] row {j}: FORCING MISMATCH at frame {bad} " |
| f"(got {sem[bad]}, forced {q0[bad]}) — row dropped") |
| n_mismatch += 1 |
| continue |
| if len(sem) < len(q0): |
| print(f"[render] row {j}: short render {len(sem)}/{len(q0)} " |
| f"frames — kept, truncated") |
| n_short += 1 |
| res = codes[1:] |
| if res.min() < 0 or res.max() > RES_MAX: |
| raise SystemExit(f"row {j}: residual out of range — " |
| f"codec/vocab drift, aborting") |
| out = dict(r) |
| out["speech_tokens"] = sem |
| out["residual_codes"] = [res[k].tolist() |
| for k in range(res.shape[0])] |
| out["n_frames"] = len(sem) |
| out["renderer"] = {"model": "fishaudio/s1-mini", "row": j, "seed": args.seed + j, |
| "voice": vname, "text_tagged": text != r["prompt"], |
| "temp": args.temp, "top_p": args.top_p, |
| "seconds": round(render_s, 3), |
| "rtf": round(render_s / max(len(sem) / 21.533, 1e-6), 4), |
| "device": (torch.cuda.get_device_name(0) if args.device == "cuda" else args.device)} |
| fout.write(json.dumps(out, ensure_ascii=False) + "\n") |
| fout.flush() |
| n_ok += 1 |
| if (j + 1) % 5 == 0: |
| print(f" {j + 1}/{len(rows)} rows " |
| f"({(time.time() - t0) / 60:.1f} min)", flush=True) |
|
|
| print(f"[render] DONE: {n_ok} ok, {n_mismatch} forcing-mismatch, " |
| f"{n_short} short/empty -> {out_path}") |
| if n_mismatch: |
| print("[render] WARNING: forcing mismatches mean the shim missed a " |
| "sampling path — do NOT use this output until explained.") |
| print("S1MINI_RENDER_DONE") |
|
|
| if __name__ == "__main__": |
| main() |
|
|
