|
|
| """Run a Navoiy TTS LLM checkpoint with the upstream CosyVoice2 runtime."""
|
|
|
| from __future__ import annotations
|
|
|
| import argparse
|
| import json
|
| import random
|
| import sys
|
| from pathlib import Path
|
|
|
|
|
| ROOT = Path(__file__).resolve().parent
|
| DEFAULT_EMOTIONS = ROOT / "emotions_40h.json"
|
|
|
|
|
| def parse_args() -> argparse.Namespace:
|
| parser = argparse.ArgumentParser(description=__doc__)
|
| parser.add_argument("--cosyvoice-dir", type=Path, help="CosyVoice source checkout")
|
| parser.add_argument("--base-model-dir", type=Path, help="Downloaded CosyVoice2-0.5B directory")
|
| parser.add_argument("--checkpoint", type=Path, help="Navoiy .pt LLM checkpoint")
|
| parser.add_argument("--reference", type=Path, help="Consented reference-speaker WAV")
|
| parser.add_argument("--text", help="Uzbek text to synthesize")
|
| parser.add_argument("--emotion", default="calm", help="Emotion preset name or tag")
|
| parser.add_argument("--emotions-file", type=Path, default=DEFAULT_EMOTIONS)
|
| parser.add_argument("--output", type=Path, default=Path("output.wav"))
|
| parser.add_argument("--speed", type=float, default=1.0)
|
| parser.add_argument("--seed", type=int, default=1986)
|
| parser.add_argument("--list-emotions", action="store_true")
|
| args = parser.parse_args()
|
|
|
| if not args.list_emotions:
|
| required = ("cosyvoice_dir", "base_model_dir", "checkpoint", "reference", "text")
|
| missing = [name.replace("_", "-") for name in required if not getattr(args, name)]
|
| if missing:
|
| parser.error("missing required arguments: " + ", ".join(f"--{name}" for name in missing))
|
| if args.speed <= 0:
|
| parser.error("--speed must be greater than zero")
|
| return args
|
|
|
|
|
| def load_emotions(path: Path) -> tuple[list[dict], dict[str, dict]]:
|
| entries = json.loads(path.read_text(encoding="utf-8"))
|
| lookup: dict[str, dict] = {}
|
| for entry in entries:
|
| lookup[entry["uz"].lower()] = entry
|
| for tag in entry["tag"].replace("[", " ").replace("]", " ").split():
|
| lookup[tag.lower()] = entry
|
| return entries, lookup
|
|
|
|
|
| def unwrap_state_dict(value):
|
| if not isinstance(value, dict):
|
| return value
|
| for key in ("state_dict", "model", "llm"):
|
| nested = value.get(key)
|
| if isinstance(nested, dict):
|
| return unwrap_state_dict(nested)
|
| return value
|
|
|
|
|
| def main() -> None:
|
| args = parse_args()
|
| entries, emotions = load_emotions(args.emotions_file)
|
|
|
| if args.list_emotions:
|
| for entry in entries:
|
| aliases = ", ".join(
|
| part.strip() for part in entry["tag"].replace("[", "").split("]") if part.strip()
|
| )
|
| print(f"{entry['uz']}: {aliases}")
|
| return
|
|
|
| emotion = emotions.get(args.emotion.lower().strip("[]"))
|
| if emotion is None:
|
| valid = ", ".join(sorted(emotions))
|
| raise SystemExit(f"Unknown emotion {args.emotion!r}. Available names/tags: {valid}")
|
|
|
| cosyvoice_dir = args.cosyvoice_dir.resolve()
|
| sys.path.insert(0, str(cosyvoice_dir))
|
| sys.path.insert(0, str(cosyvoice_dir / "third_party" / "Matcha-TTS"))
|
|
|
| import torch
|
| import torchaudio
|
| from cosyvoice.cli.cosyvoice import CosyVoice2
|
| from uztts.normalize import normalize
|
|
|
| if not torch.cuda.is_available():
|
| raise SystemExit("CUDA GPU is required for this inference script.")
|
| for path, label in (
|
| (args.base_model_dir, "base model"),
|
| (args.checkpoint, "checkpoint"),
|
| (args.reference, "reference WAV"),
|
| ):
|
| if not path.exists():
|
| raise SystemExit(f"{label} not found: {path}")
|
|
|
| random.seed(args.seed)
|
| torch.manual_seed(args.seed)
|
| torch.cuda.manual_seed_all(args.seed)
|
|
|
| model = CosyVoice2(
|
| str(args.base_model_dir.resolve()),
|
| load_jit=False,
|
| load_trt=False,
|
| fp16=True,
|
| )
|
| state = torch.load(args.checkpoint, map_location="cpu", weights_only=True)
|
| state = unwrap_state_dict(state)
|
| incompatible = model.model.llm.load_state_dict(state, strict=False)
|
| if incompatible.missing_keys:
|
| print(f"Warning: {len(incompatible.missing_keys)} missing checkpoint keys", file=sys.stderr)
|
| if incompatible.unexpected_keys:
|
| print(f"Warning: {len(incompatible.unexpected_keys)} unexpected checkpoint keys", file=sys.stderr)
|
| model.model.llm.eval()
|
|
|
| text = normalize(args.text, mode="infer")
|
| instruction = emotion["instruct"].strip() + "<|endofprompt|>"
|
| chunks = []
|
| with torch.inference_mode():
|
| for result in model.inference_instruct2(
|
| text,
|
| instruction,
|
| str(args.reference.resolve()),
|
| stream=False,
|
| speed=args.speed,
|
| ):
|
| chunks.append(result["tts_speech"].detach().cpu())
|
| if not chunks:
|
| raise SystemExit("The model returned no audio.")
|
|
|
| audio = torch.cat(chunks, dim=1)
|
| args.output.parent.mkdir(parents=True, exist_ok=True)
|
| torchaudio.save(str(args.output), audio, 24000)
|
| duration = audio.shape[-1] / 24000
|
| print(f"Wrote {args.output} ({duration:.2f}s, emotion={args.emotion}, seed={args.seed})")
|
|
|
|
|
| if __name__ == "__main__":
|
| main()
|
|
|
