EvoTalk / inference.py
akkiisfrommars's picture
Upload inference.py with huggingface_hub
06b6b0f verified
Raw
History Blame Contribute Delete
8.56 kB
import os
import re
import sys
import json
import warnings
warnings.filterwarnings("ignore")
os.environ["PYTHONWARNINGS"] = "ignore"
import torch
import torchaudio
torch.set_warn_always(False)
from model import EvoTalk, EvoTalkConfig
VALID_ARPABET = {
"AA", "AE", "AH", "AO", "AW", "AY",
"B", "CH", "D", "DH", "EH", "ER", "EY",
"F", "G", "HH", "IH", "IY", "JH", "K",
"L", "M", "N", "NG", "OW", "OY", "P",
"R", "S", "SH", "T", "TH", "UH", "UW",
"V", "W", "Y", "Z", "ZH",
}
_G2P_BACKEND = None
def get_g2p_backend():
global _G2P_BACKEND
if _G2P_BACKEND is None:
from g2p_en import G2p
_G2P_BACKEND = G2p()
return _G2P_BACKEND
def text_to_arpabet(text):
try:
g2p = get_g2p_backend()
raw_tokens = g2p(text.strip())
arpabet = []
for tok in raw_tokens:
tok = tok.strip()
if not tok:
continue
phon = re.sub(r"[0-9]", "", tok).upper()
if phon in VALID_ARPABET:
arpabet.append(phon)
return arpabet if arpabet else None
except Exception:
return None
def phonemes_to_ids(phonemes, token2id):
unk = token2id["<unk>"]
return [token2id.get(p, unk) for p in phonemes]
SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
CHECKPOINT = os.path.join(SCRIPT_DIR, "ckpt.pt")
METADATA = os.path.join(SCRIPT_DIR, "metadata.json")
AUDIO_DIR = os.path.join(SCRIPT_DIR, "audio")
MAX_WORDS = 12
PAUSE = {"sentence": 0.40, "clause": 0.20, "wrap": 0.10}
def split_sentences(text):
parts = re.split(r"(?<=[.!?])\s+", text.strip())
return [p.strip() for p in parts if p.strip()]
def split_clauses(sentence):
parts = re.split(r"(?<=[,;:])\s+|\s+--\s+|\s+—\s+", sentence)
return [p.strip() for p in parts if p and p.strip()]
def wrap_words(piece, max_words):
words = piece.split()
if len(words) <= max_words:
return [piece]
out = []
for i in range(0, len(words), max_words):
out.append(" ".join(words[i:i + max_words]))
return out
def chunk_text(text, max_words=MAX_WORDS):
chunks = []
for sentence in split_sentences(text):
clauses = split_clauses(sentence)
for ci, clause in enumerate(clauses):
last_clause = ci == len(clauses) - 1
pieces = wrap_words(clause, max_words)
for pi, piece in enumerate(pieces):
last_piece = pi == len(pieces) - 1
if last_piece and last_clause:
kind = "sentence"
elif last_piece:
kind = "clause"
else:
kind = "wrap"
chunks.append((piece, kind))
return chunks
def load_metadata(path):
with open(path) as f:
return json.load(f)
def load_model(checkpoint_path, device):
ckpt = torch.load(checkpoint_path, map_location=device, weights_only=False)
config = ckpt.get("config", None) or EvoTalkConfig()
model = EvoTalk(config).to(device)
model.load_state_dict(ckpt["model"])
model.eval()
return model
def load_vocos(device):
from vocos import Vocos
vocos = Vocos.from_pretrained("charactr/vocos-mel-24khz").to(device)
vocos.eval()
return vocos
def synth_chunk(model, vocos, text, token2id, speaker_ids, duration_scale, device):
phonemes = text_to_arpabet(text)
if not phonemes:
diagnose_phonemizer(text)
return None
try:
ids_list = phonemes_to_ids(phonemes, token2id)
except Exception as e:
print(f"warning: phonemes_to_ids failed for {phonemes!r}: {e}")
return None
if not ids_list:
print(f"warning: phonemes_to_ids returned empty ids for {phonemes!r}")
return None
ids = torch.tensor([ids_list], dtype=torch.long, device=device)
with torch.no_grad():
mel = model.inference(ids, speaker_ids, src_mask=None,
duration_scale=duration_scale)
wav = vocos.decode(mel.float().transpose(1, 2).contiguous())
return wav.squeeze(0).float().cpu()
def diagnose_phonemizer(text):
print(f"warning: text_to_arpabet produced no phonemes for: {text!r}")
try:
g2p = get_g2p_backend()
raw = g2p(text.strip())
print(f" g2p_en loaded fine, raw output: {raw!r}")
except Exception as e:
print(f" underlying g2p_en error: {type(e).__name__}: {e}")
print(" make sure g2p_en is installed: pip install g2p_en")
def synthesize(model, vocos, meta, text, speaker_id, duration_scale, device,
max_words=MAX_WORDS):
sr = meta["sample_rate"]
n_speakers = meta.get("n_speakers", 1)
speaker_id = max(0, min(int(speaker_id), n_speakers - 1))
spk = torch.tensor([speaker_id], dtype=torch.long, device=device)
chunks = chunk_text(text, max_words)
if not chunks:
return None
pieces = []
for piece, kind in chunks:
wav = synth_chunk(model, vocos, piece, meta["token2id"], spk,
duration_scale, device)
if wav is None:
continue
pieces.append(wav)
pieces.append(torch.zeros(int(PAUSE[kind] * sr)))
if not pieces:
return None
return torch.cat(pieces[:-1])
def save_waveform(waveform, sr, out_path):
os.makedirs(os.path.dirname(out_path) or ".", exist_ok=True)
wav = waveform.unsqueeze(0) if waveform.dim() == 1 else waveform
peak = wav.abs().max()
if peak > 0:
wav = wav / peak
torchaudio.save(out_path, wav, sr)
return out_path
def play_audio(path):
try:
if sys.platform == "darwin":
os.system(f"afplay '{path}'")
elif sys.platform.startswith("linux"):
os.system(f"aplay '{path}' 2>/dev/null || paplay '{path}' 2>/dev/null")
elif sys.platform.startswith("win"):
os.system(f'start "" "{path}"')
except Exception:
pass
BANNER = r"""
_____ _______ _ _
| ____|_ _____|_ _|_ _ | | | __
| _| \ \ / / _ \| |/ _` || | |/ /
| |___ \ V / (_) | | (_| || | <
|_____| \_/ \___/|_|\__,_||_|_|\_\
"""
def run(speaker_id, duration_scale, max_words):
print(BANNER)
print(f"loading model from {CHECKPOINT}")
device = "cuda" if torch.cuda.is_available() else "cpu"
meta = load_metadata(METADATA)
model = load_model(CHECKPOINT, device)
vocos = load_vocos(device)
print(f"device: {device}")
print(f"speaker: {speaker_id} duration_scale: {duration_scale} max_words: {max_words}")
print("type text and press enter to hear it")
print("commands: :speaker N :scale X :words N :quit")
print("-" * 50)
n_speakers = meta.get("n_speakers", 1)
idx = 0
while True:
try:
text = input("\n> ").strip()
except (EOFError, KeyboardInterrupt):
print()
break
if not text:
continue
if text == ":quit":
break
if text.startswith(":speaker"):
parts = text.split()
if len(parts) == 2 and parts[1].isdigit():
speaker_id = max(0, min(int(parts[1]), n_speakers - 1))
print(f"speaker set to {speaker_id}")
else:
print("usage: :speaker N")
continue
if text.startswith(":scale"):
parts = text.split()
try:
duration_scale = float(parts[1])
print(f"duration_scale set to {duration_scale}")
except (IndexError, ValueError):
print("usage: :scale X")
continue
if text.startswith(":words"):
parts = text.split()
if len(parts) == 2 and parts[1].isdigit():
max_words = int(parts[1])
print(f"max_words set to {max_words}")
else:
print("usage: :words N")
continue
wav = synthesize(model, vocos, meta, text, speaker_id, duration_scale,
device, max_words)
if wav is None:
print("could not synthesize that text")
continue
out_path = os.path.join(AUDIO_DIR, f"line_{idx:04d}.wav")
save_waveform(wav, meta["sample_rate"], out_path)
duration = wav.numel() / meta["sample_rate"]
print(f"{duration:.2f}s -> {out_path}")
play_audio(out_path)
idx += 1
print("goodbye")
if __name__ == "__main__":
run(speaker_id=0, duration_scale=1, max_words=MAX_WORDS)