Michael Treadgold
6e1da4ddfb
- inflect_nano/train_smooth.py: enhanced training with multi-res STFT loss, adversarial mel discriminator, vocoder consistency loss, deeper residual postnet, and cosine LR schedule - preprocess_dataset.py: convert HF datasets, local dirs, or LJSpeech CSVs to the durations.jsonl format needed by training - inference.py: add --smooth-prosody, --mel-smooth-sigma, --lowpass-hz flags for zero-cost inference-time quality improvements - test_inference.py: smoke tests for model loading and synthesis - colab_smooth_finetune.ipynb: Colab notebook for T4 GPU fine-tuning - requirements.txt: add numba, scipy, datasets
127 lines
4.9 KiB
Python
127 lines
4.9 KiB
Python
"""Basic smoke test for Inflect-Nano-v1: model loading, text-to-tokens, and synthesis."""
|
|
from __future__ import annotations
|
|
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
# Ensure repo root and vendored frontend are on sys.path
|
|
REPO_ROOT = Path(__file__).resolve().parent
|
|
VENDORED_FRONTEND = REPO_ROOT / "third_party" / "tiny_tts_frontend"
|
|
sys.path.insert(0, str(REPO_ROOT))
|
|
sys.path.insert(0, str(VENDORED_FRONTEND))
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
from tiny_tts.nn import commons
|
|
from tiny_tts.text import phonemes_to_ids
|
|
from tiny_tts.text.english import grapheme_to_phoneme, normalize_text
|
|
from tiny_tts.utils import ADD_BLANK
|
|
|
|
from inflect_nano.text_cleaning import clean_tinytts_text
|
|
from inflect_nano.vocoder import HifiGanGenerator, make_config
|
|
from inflect_nano.acoustic import MicroFastSpeech, MicroFastSpeechConfig
|
|
|
|
|
|
def test_text_to_tokens():
|
|
"""Verify the text -> token pipeline runs end-to-end."""
|
|
text = "Hello world, this is a test."
|
|
cleaned = clean_tinytts_text(text)
|
|
normalized = normalize_text(cleaned)
|
|
phones, tones, _ = grapheme_to_phoneme(normalized)
|
|
phone_ids, tone_ids, lang_ids = phonemes_to_ids(phones, tones, "EN")
|
|
if ADD_BLANK:
|
|
phone_ids = commons.insert_blanks(phone_ids, 0)
|
|
tone_ids = commons.insert_blanks(tone_ids, 0)
|
|
lang_ids = commons.insert_blanks(lang_ids, 0)
|
|
|
|
assert len(phone_ids) > 0, "Should produce non-empty phone IDs"
|
|
assert len(phone_ids) == len(tone_ids) == len(lang_ids), "All token sequences should have same length"
|
|
print(f" [PASS] text_to_tokens: {len(phone_ids)} tokens from '{text}'")
|
|
|
|
|
|
def test_load_acoustic():
|
|
"""Verify the acoustic model loads without error."""
|
|
device = torch.device("cpu")
|
|
ckpt_path = REPO_ROOT / "weights" / "inflect_nano_v1_acoustic.pt"
|
|
ckpt = torch.load(ckpt_path, map_location=device, weights_only=False)
|
|
cfg = MicroFastSpeechConfig(**ckpt["config"])
|
|
model = MicroFastSpeech(cfg).to(device)
|
|
model.load_state_dict(ckpt["model"])
|
|
model.eval()
|
|
|
|
params = sum(p.numel() for p in model.parameters())
|
|
print(f" [PASS] load_acoustic: {params:,} parameters loaded")
|
|
|
|
|
|
def test_load_vocoder():
|
|
"""Verify the vocoder loads without error."""
|
|
device = torch.device("cpu")
|
|
ckpt_path = REPO_ROOT / "weights" / "inflect_nano_v1_vocoder.pt"
|
|
ckpt = torch.load(ckpt_path, map_location=device, weights_only=False)
|
|
cfg = make_config((ckpt.get("config") or {}).get("variant", "snake_v2mid"))
|
|
model = HifiGanGenerator(cfg).to(device)
|
|
model.load_state_dict(ckpt["generator"])
|
|
model.remove_weight_norm()
|
|
model.eval()
|
|
|
|
params = sum(p.numel() for p in model.parameters())
|
|
print(f" [PASS] load_vocoder: {params:,} parameters loaded")
|
|
|
|
|
|
def test_synthesize():
|
|
"""End-to-end: produce a waveform from text."""
|
|
device = torch.device("cpu")
|
|
|
|
# Load acoustic
|
|
acoustic_ckpt = torch.load(REPO_ROOT / "weights" / "inflect_nano_v1_acoustic.pt", map_location=device, weights_only=False)
|
|
acoustic_cfg = MicroFastSpeechConfig(**acoustic_ckpt["config"])
|
|
acoustic = MicroFastSpeech(acoustic_cfg).to(device)
|
|
acoustic.load_state_dict(acoustic_ckpt["model"])
|
|
acoustic.eval()
|
|
speakers = acoustic_ckpt.get("speakers") or {"mark": 0}
|
|
|
|
# Load vocoder
|
|
vocoder_ckpt = torch.load(REPO_ROOT / "weights" / "inflect_nano_v1_vocoder.pt", map_location=device, weights_only=False)
|
|
vocoder_cfg = make_config((vocoder_ckpt.get("config") or {}).get("variant", "snake_v2mid"))
|
|
vocoder = HifiGanGenerator(vocoder_cfg).to(device)
|
|
vocoder.load_state_dict(vocoder_ckpt["generator"])
|
|
vocoder.remove_weight_norm()
|
|
vocoder.eval()
|
|
|
|
# Tokenize
|
|
text = "This is a quick test."
|
|
cleaned = clean_tinytts_text(text)
|
|
normalized = normalize_text(cleaned)
|
|
phones, tones, _ = grapheme_to_phoneme(normalized)
|
|
phone_ids, tone_ids, lang_ids = phonemes_to_ids(phones, tones, "EN")
|
|
if ADD_BLANK:
|
|
phone_ids = commons.insert_blanks(phone_ids, 0)
|
|
tone_ids = commons.insert_blanks(tone_ids, 0)
|
|
lang_ids = commons.insert_blanks(lang_ids, 0)
|
|
|
|
phone = torch.LongTensor(phone_ids).unsqueeze(0).to(device)
|
|
tone = torch.LongTensor(tone_ids).unsqueeze(0).to(device)
|
|
lang = torch.LongTensor(lang_ids).unsqueeze(0).to(device)
|
|
speaker = torch.LongTensor([int(speakers.get("mark", 0))]).to(device)
|
|
|
|
# Synthesize
|
|
with torch.inference_mode():
|
|
mel = acoustic.infer(phone, tone, lang, speaker)
|
|
wav = vocoder(mel).squeeze().cpu().numpy()
|
|
|
|
assert isinstance(wav, np.ndarray), "Output should be a numpy array"
|
|
assert wav.size > 0, "Waveform should not be empty"
|
|
assert np.abs(wav).max() <= 1.0, "Waveform should be normalized to [-1, 1]"
|
|
|
|
print(f" [PASS] synthesize: {wav.size} samples ({wav.size / 24000:.2f}s) from '{text}'")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
print("Inflect-Nano-v1 smoke tests\n")
|
|
test_text_to_tokens()
|
|
test_load_acoustic()
|
|
test_load_vocoder()
|
|
test_synthesize()
|
|
print("\nAll tests passed!")
|