| 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) |
|
|