#!/usr/bin/env python3 """Batch-synthesize the Egyptian eval set with s2-pro (+optional LoRA ckpt). Loads the text2semantic model + codec ONCE, then synthesizes every sentence in --sentences (jsonl: {id, text}) using a fixed reference voice, writing /.wav plus timing.jsonl with generation stats (RTF etc). Usage: python synth_eval.py --outdir /opt/work/eval/baseline python synth_eval.py --outdir /opt/work/eval/step_2000 \ --lora-ckpt results/s2pro_egy_lora/checkpoints/step_000002000.ckpt \ --lora-config r_32_egy """ import argparse import json import time from pathlib import Path import numpy as np import soundfile as sf import torch from loguru import logger from fish_speech.models.text2semantic.inference import ( encode_audio, generate_long, init_model, load_codec_model, ) CKPT = Path("/opt/work/checkpoints/s2-pro") def load_lora(model, ckpt_path: str, lora_config_name: str, lora_filter: str = "all"): from hydra import compose, initialize_config_dir from hydra.utils import instantiate from fish_speech.models.text2semantic.lora import setup_lora cfg_dir = str( Path("/opt/work/fish-speech/fish_speech/configs/lora").resolve() ) with initialize_config_dir(version_base="1.3", config_dir=cfg_dir): lora_cfg = instantiate(compose(config_name=lora_config_name)) setup_lora(model, lora_cfg) sd = torch.load(ckpt_path, map_location="cpu", weights_only=False) if "state_dict" in sd: sd = sd["state_dict"] sd = {k.removeprefix("model."): v for k, v in sd.items()} if lora_filter == "slow": sd = {k: v for k, v in sd.items() if ".fast_" not in k and not k.startswith("fast_")} logger.info(f"ablation slow-only: {len(sd)} tensors kept") elif lora_filter == "fast": sd = {k: v for k, v in sd.items() if ".fast_" in k or k.startswith("fast_")} logger.info(f"ablation fast-only: {len(sd)} tensors kept") err = model.load_state_dict(sd, strict=False) n_lora = sum(1 for k in sd if "lora" in k) assert n_lora > 0, "no lora keys in checkpoint!" logger.info(f"Loaded {n_lora} LoRA tensors; missing={len(err.missing_keys)} (expected: base weights)") if err.unexpected_keys: logger.warning(f"Unexpected keys: {err.unexpected_keys[:5]}") return model def main(): ap = argparse.ArgumentParser() ap.add_argument("--sentences", default="/opt/work/scripts/eval_sentences.jsonl") ap.add_argument("--outdir", required=True) ap.add_argument("--lora-ckpt", default=None) ap.add_argument("--lora-config", default="r_32_egy") ap.add_argument("--ref-audio", default="/opt/work/eval/ref_voice.wav") ap.add_argument("--ref-text-file", default="/opt/work/eval/ref_voice.txt") ap.add_argument("--temperature", type=float, default=0.8) ap.add_argument("--top-p", type=float, default=0.8) ap.add_argument("--seed", type=int, default=42) ap.add_argument("--skip-existing", action="store_true", help="only synthesize ids whose wav is missing; merge timing") ap.add_argument("--lora-filter", choices=["all", "slow", "fast"], default="all", help="ablation: load only slow-transformer or only fast-transformer LoRA weights") ap.add_argument("--limit", type=int, default=0, help="synth only first N sentences") args = ap.parse_args() outdir = Path(args.outdir) outdir.mkdir(parents=True, exist_ok=True) device = "cuda" precision = torch.bfloat16 sentences = [ json.loads(l) for l in Path(args.sentences).read_text(encoding="utf-8").splitlines() if l.strip() ] old_timing = {} tj = outdir / "timing.jsonl" if args.skip_existing and tj.exists(): for l in tj.read_text(encoding="utf-8").splitlines(): if l.strip(): r = json.loads(l) old_timing[r["id"]] = r before = len(sentences) sentences = [s for s in sentences if not (outdir / f"{s['id']}.wav").exists()] logger.info(f"skip-existing: {before - len(sentences)} kept, {len(sentences)} to synth") if not sentences: logger.info("nothing to do") return if args.limit: sentences = sentences[: args.limit] logger.info("Loading text2semantic model...") model, decode_one_token = init_model(CKPT, device, precision, compile=False) if args.lora_ckpt: model = load_lora(model, args.lora_ckpt, args.lora_config, args.lora_filter) model = model.to(device=device, dtype=precision).eval() with torch.device(device): model.setup_caches( max_batch_size=1, max_seq_len=model.config.max_seq_len, dtype=next(model.parameters()).dtype, ) logger.info("Loading codec...") codec = load_codec_model(CKPT / "codec.pth", device, precision) ref_text = Path(args.ref_text_file).read_text(encoding="utf-8").strip() ref_tokens = encode_audio(args.ref_audio, codec, device).cpu() results = [] for s in sentences: torch.manual_seed(args.seed) torch.cuda.manual_seed(args.seed) t0 = time.time() gen = generate_long( model=model, device=device, decode_one_token=decode_one_token, text=s["text"], num_samples=1, max_new_tokens=0, top_p=args.top_p, top_k=30, temperature=args.temperature, compile=False, iterative_prompt=True, chunk_length=300, prompt_text=[ref_text], prompt_tokens=[ref_tokens], ) codes = [] for r in gen: if r.action == "sample": codes.append(r.codes) gen_s = time.time() - t0 if not codes: logger.error(f"{s['id']}: NO CODES GENERATED") results.append({"id": s["id"], "error": "no_codes"}) continue merged = torch.cat(codes, dim=1).to(device) with torch.no_grad(): fake = codec.from_indices(merged.unsqueeze(0)) wav = fake[0, 0].float().cpu().numpy() dur = len(wav) / codec.sample_rate sf.write(outdir / f"{s['id']}.wav", wav, codec.sample_rate) rtf = gen_s / max(dur, 1e-6) results.append( {"id": s["id"], "text": s["text"], "gen_s": round(gen_s, 2), "dur_s": round(dur, 2), "rtf": round(rtf, 2)} ) logger.info(f"{s['id']}: {dur:.1f}s audio in {gen_s:.1f}s (RTF {rtf:.2f})") merged = {**old_timing, **{r["id"]: r for r in results}} with open(outdir / "timing.jsonl", "w", encoding="utf-8") as f: for r in merged.values(): f.write(json.dumps(r, ensure_ascii=False) + "\n") ok = [r for r in results if "rtf" in r] if ok: logger.info( f"DONE {len(ok)}/{len(results)} ok; mean RTF " f"{sum(r['rtf'] for r in ok)/len(ok):.2f}" ) if __name__ == "__main__": main()