File size: 5,324 Bytes
018434c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 | #!/usr/bin/env python3
"""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()
|