File size: 2,242 Bytes
778f04c | 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 | # -*- coding: utf-8 -*-
"""Minimal inference for the Ewe Spark-TTS model."""
import re, sys, torch, numpy as np, soundfile as sf
from transformers import AutoModelForCausalLM, AutoTokenizer
MODEL = "." # this repo
SPARK_SRC = "./Spark-TTS" # git clone https://github.com/SparkAudio/Spark-TTS
sys.path.append(SPARK_SRC)
from sparktts.models.audio_tokenizer import BiCodecTokenizer
from ewe_text import normalize_ewe
dev = "cuda" if torch.cuda.is_available() else "cpu"
tok = AutoTokenizer.from_pretrained(f"{MODEL}/LLM")
llm = AutoModelForCausalLM.from_pretrained(
f"{MODEL}/LLM", torch_dtype=torch.float32).to(dev).eval()
audio = BiCodecTokenizer(MODEL, dev)
stops = {tok.convert_tokens_to_ids(t)
for t in ["<|im_end|>", "<|end_semantic_token|>", "<|endoftext|>"]}
stops = sorted(i for i in stops if isinstance(i, int) and i >= 0)
@torch.inference_mode()
def say(text, temperature=0.55, max_seconds=20):
# Normalisation is required — see the model card.
clean = normalize_ewe(text, digit_policy="verbalize",
strip_verse_numbers=False)
prompt = ("<|task_tts|><|start_content|>" + clean +
"<|end_content|><|start_global_token|>")
ins = tok([prompt], return_tensors="pt").to(dev)
out = llm.generate(**ins, max_new_tokens=int(max_seconds * 50) + 128,
do_sample=True, temperature=temperature,
top_k=50, top_p=0.95, eos_token_id=stops,
pad_token_id=tok.pad_token_id or tok.eos_token_id)
txt = tok.batch_decode(out[:, ins.input_ids.shape[1]:],
skip_special_tokens=False)[0]
sem = re.findall(r"<\|bicodec_semantic_(\d+)\|>", txt)
glo = re.findall(r"<\|bicodec_global_(\d+)\|>", txt)
if not sem:
raise RuntimeError("no semantic tokens generated")
return audio.detokenize(
torch.tensor([int(i) for i in glo or ["0"]]).long().unsqueeze(0).to(dev),
torch.tensor([int(i) for i in sem]).long().unsqueeze(0).to(dev))
if __name__ == "__main__":
wav = say("Ŋdi na mi. Nye ŋkɔe nye Kofi, eye medzea Eʋegbe nyuie.")
sf.write("out.wav", wav, audio.config.get("sample_rate", 16000))
print("wrote out.wav")
|