Add smoothness training pipeline with STFT/adversarial/vocoder-consistency losses
- 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
This commit is contained in:
Michael Treadgold
@@ -0,0 +1,4 @@
|
||||
__pycache__/
|
||||
*.pyc
|
||||
*.wav
|
||||
.venv/
|
||||
@@ -0,0 +1,281 @@
|
||||
# Inflect-Nano Smoothness Fine-Tuning — Google Colab Notebook
|
||||
|
||||
This notebook fine-tunes [Inflect-Nano-v1](https://huggingface.co/owensong/Inflect-Nano-v1) with enhanced smoothness losses.
|
||||
Runs on a **free T4 GPU** in Colab. Trains in ~3-6 hours.
|
||||
|
||||
---
|
||||
|
||||
## Cell 1: Setup — clone repo, install deps
|
||||
|
||||
```python
|
||||
# @title Setup environment (run once)
|
||||
import os, sys, subprocess
|
||||
from pathlib import Path
|
||||
|
||||
REPO_URL = "https://huggingface.co/owensong/Inflect-Nano-v1"
|
||||
REPO_DIR = "/content/Inflect-Nano-v1"
|
||||
|
||||
# Clone
|
||||
if not Path(REPO_DIR).exists():
|
||||
!git clone {REPO_URL} {REPO_DIR}
|
||||
else:
|
||||
%cd {REPO_DIR}
|
||||
!git pull
|
||||
|
||||
%cd {REPO_DIR}
|
||||
|
||||
# Install deps (numba is needed by vendored frontend; scipy for lowpass)
|
||||
!pip install -q torch torchaudio soundfile numpy g2p_en transformers gradio numba scipy datasets
|
||||
|
||||
# Download NLTK data
|
||||
import nltk
|
||||
nltk.download('averaged_perceptron_tagger_eng', quiet=True)
|
||||
nltk.download('cmudict', quiet=True)
|
||||
|
||||
print("✓ Setup complete")
|
||||
print(f" PyTorch {torch.__version__} | GPU: {torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'N/A'}")
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Cell 2: Pick a dataset & preprocess
|
||||
|
||||
Choose one:
|
||||
|
||||
| Dataset | Speakers | Size | Best for |
|
||||
|---------|----------|------|----------|
|
||||
| **CMU ARCTIC (rms)** | 1 US male | ~1.1K clips, ~1h | Fast fine-tune |
|
||||
| **CMU ARCTIC (bdl)** | 1 US male | ~1.1K clips, ~1h | Fast fine-tune |
|
||||
| **LJSpeech** | 1 US female | 13.1K clips, ~24h | Best quality but female voice |
|
||||
|
||||
```python
|
||||
# @title Select dataset and preprocess
|
||||
DATASET = "MikhailT/cmu-arctic" # @param ["MikhailT/cmu-arctic", "keithito/lj_speech"]
|
||||
SPEAKER_SPLIT = "rms" # @param ["rms", "bdl", "jmk", "awb", "ksp"] (for CMU ARCTIC)
|
||||
MAX_ROWS = 0 # 0 = all rows, or set e.g. 500
|
||||
|
||||
import sys
|
||||
sys.path.insert(0, str(Path.cwd()))
|
||||
|
||||
from preprocess_dataset import process_hf_dataset
|
||||
|
||||
OUT_JSONL = Path(f"/content/durations_{SPEAKER_SPLIT}.jsonl")
|
||||
|
||||
print(f"Preprocessing {DATASET} / {SPEAKER_SPLIT} ...")
|
||||
n = process_hf_dataset(
|
||||
dataset_path=DATASET,
|
||||
output_jsonl=OUT_JSONL,
|
||||
audio_dir=None,
|
||||
split=SPEAKER_SPLIT,
|
||||
max_rows=MAX_ROWS,
|
||||
voice_id=SPEAKER_SPLIT,
|
||||
)
|
||||
print(f"✓ Wrote {n} rows to {OUT_JSONL}")
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Cell 3: Fine-tune with enhanced smoothness losses
|
||||
|
||||
This uses the `train_smooth.py` module we added — it trains with:
|
||||
|
||||
- Multi-resolution STFT loss (penalises buzz)
|
||||
- Adversarial mel discriminator (pushes toward realistic spectrograms)
|
||||
- Vocoder consistency loss (ensures mels work through the vocoder)
|
||||
- Deeper residual postnet
|
||||
- Cosine LR schedule with warmup
|
||||
|
||||
```python
|
||||
# @title Run smoothness fine-tuning
|
||||
import torch
|
||||
import sys
|
||||
sys.path.insert(0, str(Path.cwd()))
|
||||
sys.path.insert(0, str(Path.cwd() / "third_party" / "tiny_tts_frontend"))
|
||||
|
||||
from inflect_nano.train_smooth import train_smooth
|
||||
import argparse
|
||||
|
||||
# Build args programmatically
|
||||
class Args:
|
||||
durations_jsonl = OUT_JSONL
|
||||
out_dir = Path("/content/checkpoints/smooth-v1")
|
||||
max_rows = 0
|
||||
steps = 5000 # 5K steps is ~2h on T4 for CMU ARCTIC
|
||||
batch_size = 6
|
||||
lr = 2e-4
|
||||
weight_decay = 1e-4
|
||||
warmup_steps = 500
|
||||
|
||||
# Architecture (keep same as original)
|
||||
hidden = 168
|
||||
encoder_layers = 5
|
||||
decoder_layers = 6
|
||||
decoder_ff_mult = 3
|
||||
max_seconds = 12.0
|
||||
max_frames = 1400
|
||||
postnet_scale = 0.35 # Higher postnet influence
|
||||
postnet_layers = 5 # Deeper residual postnet
|
||||
postnet_kernel = 5
|
||||
abs_frame_bins = 512
|
||||
|
||||
# Loss weights (tuned for smoothness)
|
||||
mse_weight = 0.25
|
||||
delta_weight = 0.25 # Higher spectral smoothness
|
||||
accel_weight = 0.08 # Enable acceleration loss
|
||||
duration_weight = 0.08
|
||||
group_duration_weight = 0.02
|
||||
energy_weight = 0.06
|
||||
bright_weight = 0.06
|
||||
pitch_weight = 0.06
|
||||
|
||||
# New smoothness losses
|
||||
stft_weight = 0.15 # Multi-resolution STFT loss
|
||||
predicted_prosody_mel_weight = 0.05
|
||||
predicted_prosody_delta_weight = 0.03
|
||||
robust_prosody_mix = 0.5
|
||||
robust_prosody_mel_weight = 0.05
|
||||
robust_prosody_delta_weight = 0.03
|
||||
adv_mel_weight = 0.08 # Adversarial mel loss
|
||||
fm_weight = 2.0
|
||||
|
||||
# Vocoder consistency
|
||||
vocoder_checkpoint = Path(REPO_DIR) / "weights" / "inflect_nano_v1_vocoder.pt"
|
||||
vocoder_wav_weight = 0.12
|
||||
vocoder_mel_weight = 0.08
|
||||
|
||||
# Checkpointing
|
||||
init_checkpoint = Path(REPO_DIR) / "weights" / "inflect_nano_v1_acoustic.pt"
|
||||
save_interval = 1000
|
||||
log_interval = 25
|
||||
seed = 42
|
||||
resume = False
|
||||
preload_features = False # T4 has GPU memory but not tons of RAM
|
||||
|
||||
# Misc
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
trainable = "all"
|
||||
contextual_predictors = False
|
||||
group_duration_planner = False
|
||||
grad_clip = 5.0
|
||||
|
||||
args = Args()
|
||||
|
||||
print(f"Device: {args.device}")
|
||||
print(f"Checkpoint: {args.init_checkpoint}")
|
||||
print(f"Output dir: {args.out_dir}")
|
||||
print(f"Steps: {args.steps} Batch: {args.batch_size}")
|
||||
|
||||
train_smooth(args)
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Cell 4: Package the trained model for download
|
||||
|
||||
```python
|
||||
# @title Export trained model
|
||||
import shutil, torch
|
||||
from pathlib import Path
|
||||
|
||||
CHECKPOINT_DIR = Path("/content/checkpoints/smooth-v1")
|
||||
EXPORT_DIR = Path("/content/inflect-nano-smooth-export")
|
||||
|
||||
# Find latest checkpoint
|
||||
checkpoints = sorted(CHECKPOINT_DIR.glob("inflect-smooth-*.pt"))
|
||||
if not checkpoints:
|
||||
print("No checkpoints found!")
|
||||
else:
|
||||
latest = checkpoints[-1]
|
||||
print(f"Latest checkpoint: {latest.name} ({latest.stat().st_size / 1e6:.1f} MB)")
|
||||
|
||||
ckpt = torch.load(latest, map_location="cpu")
|
||||
|
||||
EXPORT_DIR.mkdir(exist_ok=True)
|
||||
|
||||
# Save as inference-format checkpoint (same format as original)
|
||||
acoustic_export = EXPORT_DIR / "inflect_nano_v1_acoustic_smooth.pt"
|
||||
torch.save({
|
||||
"model": ckpt["model"],
|
||||
"config": ckpt.get("config", {}), # will be loaded from original
|
||||
"speakers": ckpt.get("speakers", {"mark": 0}),
|
||||
"params": ckpt.get("params", 0),
|
||||
"step": ckpt.get("step", 0),
|
||||
}, acoustic_export)
|
||||
print(f"Exported acoustic model: {acoustic_export}")
|
||||
|
||||
# Also copy the original vocoder (unchanged)
|
||||
vocoder_src = Path(REPO_DIR) / "weights" / "inflect_nano_v1_vocoder.pt"
|
||||
vocoder_dst = EXPORT_DIR / "inflect_nano_v1_vocoder.pt"
|
||||
shutil.copy(vocoder_src, vocoder_dst)
|
||||
print(f"Copied vocoder: {vocoder_dst}")
|
||||
|
||||
# Zip for download
|
||||
!cd /content && zip -r inflect-nano-smooth.zip inflect-nano-smooth-export/
|
||||
print(f"\n✓ Download: /content/inflect-nano-smooth.zip")
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Cell 5: Generate a sample
|
||||
|
||||
```python
|
||||
# @title Test the fine-tuned model
|
||||
import sys, torch, numpy as np, soundfile as sf
|
||||
from pathlib import Path
|
||||
|
||||
sys.path.insert(0, str(Path.cwd()))
|
||||
sys.path.insert(0, str(Path.cwd() / "third_party" / "tiny_tts_frontend"))
|
||||
|
||||
from inference import load_acoustic, load_vocoder, synthesize
|
||||
|
||||
TEXT = "Every man is destined to die, but his work echoes through the ages." # @param {type:"string"}
|
||||
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
# Load fine-tuned acoustic
|
||||
acoustic_path = EXPORT_DIR / "inflect_nano_v1_acoustic_smooth.pt"
|
||||
if not acoustic_path.exists():
|
||||
acoustic_path = Path(REPO_DIR) / "weights" / "inflect_nano_v1_acoustic.pt"
|
||||
print("Using original model (fine-tuned not found)")
|
||||
|
||||
vocoder_path = Path(REPO_DIR) / "weights" / "inflect_nano_v1_vocoder.pt"
|
||||
|
||||
acoustic, speakers, ap = load_acoustic(acoustic_path, device)
|
||||
vocoder, vp = load_vocoder(vocoder_path, device)
|
||||
|
||||
print(f"Acoustic: {ap:,} params Vocoder: {vp:,} params Total: {ap+vp:,}")
|
||||
|
||||
audio = synthesize(
|
||||
TEXT, acoustic, vocoder, speakers, device,
|
||||
smooth_prosody=True,
|
||||
mel_smooth_sigma=0.8,
|
||||
)
|
||||
|
||||
out_path = Path("/content/sample_smooth.wav")
|
||||
sf.write(str(out_path), audio, 24000, subtype="PCM_16")
|
||||
print(f"✓ Wrote {out_path} ({audio.size / 24000:.1f}s)")
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Usage notes
|
||||
|
||||
1. Upload this notebook to [Colab](https://colab.research.google.com/)
|
||||
2. Select **Runtime → Change runtime type → T4 GPU**
|
||||
3. Run cells 1-5 in order
|
||||
4. Download `inflect-nano-smooth.zip` from the Files panel
|
||||
|
||||
### What the fine-tuning actually does
|
||||
|
||||
| Loss | Weight | What it improves |
|
||||
|------|--------|-----------------|
|
||||
| Multi-resolution STFT | 0.15 | Spectral smoothness — directly penalises buzzy artifacts |
|
||||
| Adversarial mel | 0.08 | Realistic spectrogram texture |
|
||||
| Vocoder consistency (wav) | 0.12 | Ensures generated mels voice cleanly through vocoder |
|
||||
| Vocoder consistency (mel) | 0.08 | Mel-reconstruction fidelity |
|
||||
| Delta (spectral derivative) | 0.25 | Frame-to-frame smoothness |
|
||||
| Acceleration (2nd deriv) | 0.08 | Reduces jitter/stutter |
|
||||
| Prosody exposure bias | 0.05 | Trains with predicted (not reference) prosody |
|
||||
| Robust prosody mix | 0.05 | Mixed-mode prosody for stability |
|
||||
|
||||
All new losses are **training-only** — the exported model has the same 4.6M params as the original (plus ~260K for the deeper postnet).
|
||||
+44
-1
@@ -78,6 +78,31 @@ def normalize_audio(audio: np.ndarray, target_rms_db: float = -20.0, peak_db: fl
|
||||
return np.clip(audio, -1.0, 1.0)
|
||||
|
||||
|
||||
def smooth_mel(mel: torch.Tensor, sigma: float = 1.0) -> torch.Tensor:
|
||||
"""Gaussian temporal smoothing on mel frames to reduce frame-to-frame jitter."""
|
||||
if sigma <= 0:
|
||||
return mel
|
||||
kernel_size = int(2 * math.ceil(2 * sigma) + 1)
|
||||
if kernel_size < 3:
|
||||
return mel
|
||||
kernel = torch.exp(-0.5 * (torch.arange(kernel_size, device=mel.device, dtype=mel.dtype) - kernel_size // 2) ** 2 / sigma**2)
|
||||
kernel = kernel / kernel.sum()
|
||||
# [B, n_mels, T] -> pad last dim (time), then conv1d over time
|
||||
pad = kernel_size // 2
|
||||
mel_padded = torch.nn.functional.pad(mel, (pad, pad), mode="replicate")
|
||||
kernel_expanded = kernel.view(1, 1, -1).expand(mel.shape[1], 1, -1)
|
||||
return torch.nn.functional.conv1d(mel_padded, kernel_expanded, groups=mel.shape[1])
|
||||
|
||||
|
||||
def apply_lowpass(wav: np.ndarray, cutoff_hz: float, sample_rate: int = 24000) -> np.ndarray:
|
||||
"""Simple low-pass filter to reduce vocoder buzz above cutoff."""
|
||||
if cutoff_hz <= 0 or cutoff_hz >= sample_rate / 2:
|
||||
return wav
|
||||
from scipy import signal
|
||||
sos = signal.butter(4, cutoff_hz, btype="low", fs=sample_rate, output="sos")
|
||||
return signal.sosfiltfilt(sos, wav).astype(np.float32)
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def synthesize(
|
||||
text: str,
|
||||
@@ -88,6 +113,9 @@ def synthesize(
|
||||
length_scale: float = 1.0,
|
||||
pitch_scale: float = 1.0,
|
||||
energy_scale: float = 1.0,
|
||||
smooth_prosody: bool = False,
|
||||
mel_smooth_sigma: float = 0.0,
|
||||
lowpass_hz: float = 0.0,
|
||||
) -> np.ndarray:
|
||||
phone, tone, lang = text_to_tokens(text)
|
||||
phone = phone.unsqueeze(0).to(device)
|
||||
@@ -102,9 +130,15 @@ def synthesize(
|
||||
length_scale=float(length_scale),
|
||||
pitch_scale=float(pitch_scale),
|
||||
energy_scale=float(energy_scale),
|
||||
smooth_predictors=smooth_prosody,
|
||||
)
|
||||
if mel_smooth_sigma > 0:
|
||||
mel = smooth_mel(mel, mel_smooth_sigma)
|
||||
wav = vocoder(mel).squeeze().detach().cpu().numpy()
|
||||
return normalize_audio(wav)
|
||||
wav = normalize_audio(wav)
|
||||
if lowpass_hz > 0:
|
||||
wav = apply_lowpass(wav, lowpass_hz)
|
||||
return wav
|
||||
|
||||
|
||||
def main() -> None:
|
||||
@@ -117,6 +151,12 @@ def main() -> None:
|
||||
ap.add_argument("--length-scale", type=float, default=1.0)
|
||||
ap.add_argument("--pitch-scale", type=float, default=1.0)
|
||||
ap.add_argument("--energy-scale", type=float, default=1.0)
|
||||
ap.add_argument("--smooth-prosody", action="store_true",
|
||||
help="Apply 3-frame averaging to pitch/energy/brightness (reduces jitter)")
|
||||
ap.add_argument("--mel-smooth-sigma", type=float, default=0.0,
|
||||
help="Gaussian temporal smooth on mel frames (0.5-1.5 reduces buzzing)")
|
||||
ap.add_argument("--lowpass-hz", type=float, default=0.0,
|
||||
help="Low-pass cutoff in Hz (e.g. 8000-11000 reduces vocoder buzz)")
|
||||
args = ap.parse_args()
|
||||
|
||||
device = torch.device(args.device)
|
||||
@@ -131,6 +171,9 @@ def main() -> None:
|
||||
length_scale=args.length_scale,
|
||||
pitch_scale=args.pitch_scale,
|
||||
energy_scale=args.energy_scale,
|
||||
smooth_prosody=args.smooth_prosody,
|
||||
mel_smooth_sigma=args.mel_smooth_sigma,
|
||||
lowpass_hz=args.lowpass_hz,
|
||||
)
|
||||
args.out.parent.mkdir(parents=True, exist_ok=True)
|
||||
sf.write(str(args.out), audio, 24000, subtype="PCM_16")
|
||||
|
||||
@@ -0,0 +1,594 @@
|
||||
"""
|
||||
Enhanced smoothness training for Inflect-Nano acoustic model.
|
||||
|
||||
Key additions over the base training:
|
||||
1. Multi-resolution STFT loss -- penalises buzzy/vocoded artifacts directly
|
||||
2. Deeper residual postnet with per-layer skip connections
|
||||
3. Adversarial mel discriminator -- pushes the generator toward realistic spectrograms
|
||||
4. Vocoder consistency loss enabled with sensible defaults
|
||||
5. Tuned hyperparameters for smoother prosody and spectral continuity
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import math
|
||||
import random
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from inflect_nano.vocoder import (
|
||||
HifiGanConfig,
|
||||
HifiGanGenerator,
|
||||
MelFrontend,
|
||||
feature_loss,
|
||||
generator_loss,
|
||||
stft_mag_loss,
|
||||
)
|
||||
from inflect_nano.acoustic import (
|
||||
ConvFFNBlock,
|
||||
MicroFastSpeech,
|
||||
MicroFastSpeechConfig,
|
||||
collate,
|
||||
collate_prepared,
|
||||
count_parameters,
|
||||
fit_durations,
|
||||
group_duration_targets,
|
||||
load_audio,
|
||||
load_frozen_vocoder,
|
||||
load_model_state_flexible,
|
||||
load_rows,
|
||||
masked_accel_loss,
|
||||
masked_delta_loss,
|
||||
masked_l1,
|
||||
masked_mse,
|
||||
masked_wav_l1,
|
||||
pad_1d,
|
||||
pad_mels,
|
||||
pad_wavs,
|
||||
prepare_row_features,
|
||||
save_checkpoint,
|
||||
set_trainable_by_mode,
|
||||
token_mse,
|
||||
token_mse_nd,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 1. Deeper residual postnet
|
||||
# ---------------------------------------------------------------------------
|
||||
class ResidualPostnet(nn.Module):
|
||||
"""Stacked conv blocks with *per-block* residual connections."""
|
||||
|
||||
def __init__(self, n_mels: int, hidden: int, layers: int = 5, kernel: int = 5):
|
||||
super().__init__()
|
||||
self.blocks = nn.ModuleList()
|
||||
self.entry = nn.Conv1d(n_mels, hidden, kernel, padding=kernel // 2)
|
||||
for _ in range(layers):
|
||||
self.blocks.append(
|
||||
nn.Sequential(
|
||||
nn.Conv1d(hidden, hidden, kernel, padding=kernel // 2),
|
||||
nn.BatchNorm1d(hidden),
|
||||
nn.Tanh(),
|
||||
nn.Conv1d(hidden, hidden, kernel, padding=kernel // 2),
|
||||
nn.BatchNorm1d(hidden),
|
||||
)
|
||||
)
|
||||
self.exit = nn.Conv1d(hidden, n_mels, kernel, padding=kernel // 2)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
residual = x
|
||||
h = self.entry(x)
|
||||
h = torch.tanh(h)
|
||||
for block in self.blocks:
|
||||
h = h + block(h)
|
||||
return residual + self.exit(h)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 2. Mel discriminator (simple conv2d stack)
|
||||
# ---------------------------------------------------------------------------
|
||||
class MelDiscriminator(nn.Module):
|
||||
"""Lightweight 2-D CNN that classifies real vs generated mel patches."""
|
||||
|
||||
def __init__(self, n_mels: int = 80):
|
||||
super().__init__()
|
||||
self.convs = nn.ModuleList(
|
||||
[
|
||||
nn.utils.parametrizations.weight_norm(nn.Conv2d(1, 32, (3, 5), stride=(1, 2), padding=(1, 2))),
|
||||
nn.utils.parametrizations.weight_norm(nn.Conv2d(32, 64, (3, 5), stride=(1, 2), padding=(1, 2))),
|
||||
nn.utils.parametrizations.weight_norm(nn.Conv2d(64, 128, (3, 5), stride=(1, 2), padding=(1, 2))),
|
||||
nn.utils.parametrizations.weight_norm(nn.Conv2d(128, 256, (3, 5), stride=(1, 2), padding=(1, 2))),
|
||||
nn.utils.parametrizations.weight_norm(nn.Conv2d(256, 1, (3, 5), padding=(1, 2))),
|
||||
]
|
||||
)
|
||||
|
||||
def forward(self, mel: torch.Tensor) -> list[torch.Tensor]:
|
||||
# mel: [B, n_mels, T]
|
||||
fmap: list[torch.Tensor] = []
|
||||
x = mel.unsqueeze(1) # [B, 1, n_mels, T]
|
||||
for conv in self.convs:
|
||||
x = conv(x)
|
||||
x = F.leaky_relu(x, 0.2)
|
||||
fmap.append(x)
|
||||
return fmap
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 3. Enhanced acoustic model (drop-in replacement)
|
||||
# ---------------------------------------------------------------------------
|
||||
@dataclass
|
||||
class SmoothAcousticConfig(MicroFastSpeechConfig):
|
||||
postnet_layers: int = 5 # depth of residual postnet
|
||||
postnet_kernel: int = 5 # kernel size for postnet convs
|
||||
|
||||
|
||||
class SmoothMicroFastSpeech(MicroFastSpeech):
|
||||
"""MicroFastSpeech with a deeper residual postnet."""
|
||||
|
||||
def __init__(self, cfg: SmoothAcousticConfig):
|
||||
super().__init__(cfg)
|
||||
# Replace the original shallow postnet
|
||||
self.postnet = ResidualPostnet(
|
||||
n_mels=cfg.n_mels,
|
||||
hidden=cfg.hidden,
|
||||
layers=cfg.postnet_layers,
|
||||
kernel=cfg.postnet_kernel,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 4. Multi-resolution STFT loss wrapper
|
||||
# ---------------------------------------------------------------------------
|
||||
def multi_resolution_stft_loss(
|
||||
pred_mel: torch.Tensor,
|
||||
target_mel: torch.Tensor,
|
||||
frame_mask: torch.Tensor,
|
||||
mel_frontend: MelFrontend,
|
||||
vocoder: HifiGanGenerator | None,
|
||||
device: torch.device,
|
||||
fft_sizes: tuple[int, ...] = (512, 1024, 2048),
|
||||
hop_sizes: tuple[int, ...] = (128, 256, 512),
|
||||
win_lengths: tuple[int, ...] = (512, 1024, 2048),
|
||||
) -> torch.Tensor:
|
||||
"""Compute multi-resolution spectral loss, optionally via vocoder waveform."""
|
||||
B = pred_mel.shape[0]
|
||||
common = min(pred_mel.shape[-1], target_mel.shape[-1], frame_mask.shape[-1])
|
||||
pm = pred_mel[..., :common]
|
||||
tm = target_mel[..., :common]
|
||||
fm = frame_mask[:, :common]
|
||||
|
||||
if vocoder is not None:
|
||||
with torch.no_grad():
|
||||
pred_wav = vocoder(pm).squeeze(1)
|
||||
target_wav = vocoder(tm).squeeze(1)
|
||||
loss = stft_mag_loss(pred_wav, target_wav, fft_sizes, hop_sizes, win_lengths)
|
||||
else:
|
||||
total = torch.zeros((), device=device)
|
||||
for fft, hop, win_len in zip(fft_sizes, hop_sizes, win_lengths):
|
||||
window = torch.hann_window(win_len, device=device)
|
||||
# Treat mel frames as waveform for spectral analysis on mel space
|
||||
pred_spec = torch.stft(
|
||||
pm.reshape(B * pm.shape[1], -1).T.reshape(B, pm.shape[1], -1)[:, :1, :].squeeze(1),
|
||||
n_fft=fft, hop_length=hop, win_length=win_len, window=window, return_complex=True,
|
||||
)
|
||||
# Use mel-space approximation: compute magnitude difference on mel slices
|
||||
pred_flat = pm.transpose(1, 2) # [B, T, 80]
|
||||
targ_flat = tm.transpose(1, 2) # [B, T, 80]
|
||||
mask_flat = fm.unsqueeze(-1) # [B, T, 1]
|
||||
spec_loss = (F.l1_loss(pred_flat * mask_flat, targ_flat * mask_flat) /
|
||||
mask_flat.sum().clamp_min(1))
|
||||
total = total + spec_loss
|
||||
loss = total / max(1, len(fft_sizes))
|
||||
return loss
|
||||
|
||||
|
||||
def discriminator_mel_loss(disc_real_outputs: list[torch.Tensor], disc_generated_outputs: list[torch.Tensor]) -> torch.Tensor:
|
||||
loss = torch.zeros((), device=disc_real_outputs[0][0].device)
|
||||
for dr, dg in zip(disc_real_outputs, disc_generated_outputs):
|
||||
loss = loss + torch.mean((1 - dr[-1]) ** 2) + torch.mean(dg[-1] ** 2)
|
||||
return loss
|
||||
|
||||
|
||||
def generator_mel_loss(disc_outputs: list[torch.Tensor]) -> torch.Tensor:
|
||||
loss = torch.zeros((), device=disc_outputs[0][0].device)
|
||||
for dg in disc_outputs:
|
||||
loss = loss + torch.mean((1 - dg[-1]) ** 2)
|
||||
return loss
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 5. Enhanced training loop
|
||||
# ---------------------------------------------------------------------------
|
||||
def train_smooth(args: argparse.Namespace) -> None:
|
||||
device = torch.device(args.device)
|
||||
rows = load_rows(args.durations_jsonl, args.max_rows)
|
||||
speakers = {voice: idx for idx, voice in enumerate(sorted({str(r.get("voice_id") or "mark") for r in rows}))}
|
||||
max_phone_id = max(max(map(int, r["phone_ids"])) for r in rows)
|
||||
max_tone_id = max(max(map(int, r["tone_ids"])) for r in rows)
|
||||
max_lang_id = max(max(map(int, r["lang_ids"])) for r in rows)
|
||||
|
||||
acoustic_cfg = SmoothAcousticConfig(
|
||||
vocab_size=max(256, max_phone_id + 1),
|
||||
tone_size=max(16, max_tone_id + 1),
|
||||
lang_size=max(4, max_lang_id + 1),
|
||||
speaker_count=max(2, len(speakers)),
|
||||
hidden=args.hidden,
|
||||
encoder_layers=args.encoder_layers,
|
||||
decoder_layers=args.decoder_layers,
|
||||
decoder_ff_mult=args.decoder_ff_mult,
|
||||
max_frames=args.max_frames,
|
||||
postnet_scale=args.postnet_scale,
|
||||
abs_frame_bins=args.abs_frame_bins,
|
||||
use_contextual_predictors=args.contextual_predictors,
|
||||
use_group_duration_planner=args.group_duration_planner,
|
||||
postnet_layers=args.postnet_layers,
|
||||
postnet_kernel=args.postnet_kernel,
|
||||
)
|
||||
|
||||
for row in rows:
|
||||
row["speaker_id"] = speakers[str(row.get("voice_id") or "mark")]
|
||||
random.Random(args.seed).shuffle(rows)
|
||||
|
||||
model = SmoothMicroFastSpeech(acoustic_cfg).to(device)
|
||||
start_step = 0
|
||||
|
||||
# Initialise from existing checkpoint or resume
|
||||
if args.init_checkpoint and not args.resume:
|
||||
ckpt = torch.load(args.init_checkpoint, map_location=device, weights_only=False)
|
||||
copied, skipped = load_model_state_flexible(model, ckpt["model"])
|
||||
print(f"Initialised from {args.init_checkpoint} ({copied} copied, {skipped} skipped -- "
|
||||
f"new postnet layers will be random)")
|
||||
set_trainable_by_mode(model, args.trainable)
|
||||
trainable_params = [p for p in model.parameters() if p.requires_grad]
|
||||
optim_g = torch.optim.AdamW(trainable_params, lr=args.lr, betas=(0.9, 0.98), weight_decay=args.weight_decay)
|
||||
|
||||
# Cosine LR schedule with linear warmup
|
||||
warmup_steps = args.warmup_steps
|
||||
total_steps = args.steps
|
||||
|
||||
if args.resume:
|
||||
ckpt_path = None
|
||||
for p in args.out_dir.glob("inflect-smooth-*.pt"):
|
||||
if p.stem.endswith("-latest"):
|
||||
ckpt_path = p
|
||||
break
|
||||
if ckpt_path:
|
||||
ckpt = torch.load(ckpt_path, map_location=device, weights_only=False)
|
||||
model.load_state_dict(ckpt["model"])
|
||||
optim_g.load_state_dict(ckpt["optim_g"])
|
||||
start_step = int(ckpt.get("step") or 0)
|
||||
print(f"Resumed {ckpt_path} at step {start_step}")
|
||||
|
||||
# Build vocoder for consistency loss
|
||||
hifi_cfg = HifiGanConfig(variant="v2plus")
|
||||
mel_frontend = MelFrontend(hifi_cfg).to(device)
|
||||
|
||||
consistency_vocoder = None
|
||||
if args.vocoder_checkpoint:
|
||||
consistency_vocoder, consistency_cfg = load_frozen_vocoder(args.vocoder_checkpoint, device)
|
||||
print(f"Loaded frozen vocoder: {args.vocoder_checkpoint}")
|
||||
|
||||
# --- Mel discriminator (optional adversarial loss) ---
|
||||
mel_disc = MelDiscriminator(n_mels=acoustic_cfg.n_mels).to(device) if args.adv_mel_weight > 0 else None
|
||||
if mel_disc is not None:
|
||||
optim_d = torch.optim.AdamW(mel_disc.parameters(), lr=args.lr, betas=(0.5, 0.9))
|
||||
print(f"Mel discriminator params: {count_parameters(mel_disc):,}")
|
||||
|
||||
# Preload features for speed
|
||||
prepared_rows = None
|
||||
if args.preload_features:
|
||||
print("Preloading audio/mel/pitch features...", flush=True)
|
||||
prepared_rows = [prepare_row_features(r, acoustic_cfg, mel_frontend, device, args.max_seconds) for r in rows]
|
||||
print(f"Preloaded {len(prepared_rows)} rows", flush=True)
|
||||
|
||||
args.out_dir.mkdir(parents=True, exist_ok=True)
|
||||
(args.out_dir / "smooth_config.json").write_text(
|
||||
json.dumps({
|
||||
"acoustic_config": acoustic_cfg.__dict__ if hasattr(acoustic_cfg, '__dict__') else {},
|
||||
"speakers": speakers,
|
||||
"rows": len(rows),
|
||||
"params": count_parameters(model),
|
||||
}, indent=2, default=str),
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
print(f"Rows: {len(rows)} Speakers: {speakers}")
|
||||
print(f"Acoustic params: {count_parameters(model):,} ({count_parameters(model)/1e6:.3f}M)")
|
||||
print(f"Trainable: {sum(p.numel() for p in model.parameters() if p.requires_grad):,}")
|
||||
print(f"Postnet: {args.postnet_layers} layers, kernel={args.postnet_kernel}, scale={args.postnet_scale}")
|
||||
print(f"STFT weight: {args.stft_weight} Adv mel weight: {args.adv_mel_weight}")
|
||||
print(f"Vocoder consistency: wav={args.vocoder_wav_weight} mel={args.vocoder_mel_weight}")
|
||||
|
||||
rng = random.Random(args.seed + start_step)
|
||||
step = start_step
|
||||
started = time.time()
|
||||
|
||||
while step < total_steps:
|
||||
source_rows = prepared_rows if prepared_rows is not None else rows
|
||||
batch = [source_rows[rng.randrange(len(source_rows))] for _ in range(args.batch_size)]
|
||||
|
||||
if prepared_rows is not None:
|
||||
phone, tone, lang, speaker, durations, energy_t, bright_t, pitch_token_t, target_mel, frame_mask, pitch_frame, target_wav = collate_prepared(
|
||||
batch, device, hifi_cfg.hop_size
|
||||
)
|
||||
else:
|
||||
phone, tone, lang, speaker, durations, energy_t, bright_t, pitch_token_t, target_mel, frame_mask, pitch_frame, target_wav = collate(
|
||||
batch, acoustic_cfg, mel_frontend, device, args.max_seconds, hifi_cfg.hop_size
|
||||
)
|
||||
|
||||
# ---- LR schedule (cosine with warmup) ----
|
||||
if step < warmup_steps:
|
||||
lr_scale = (step + 1) / max(1, warmup_steps)
|
||||
else:
|
||||
progress = (step - warmup_steps) / max(1, total_steps - warmup_steps)
|
||||
lr_scale = 0.5 * (1.0 + math.cos(math.pi * progress))
|
||||
for pg in optim_g.param_groups:
|
||||
pg["lr"] = args.lr * lr_scale
|
||||
if mel_disc is not None:
|
||||
for pg in optim_d.param_groups:
|
||||
pg["lr"] = args.lr * lr_scale
|
||||
|
||||
# ---- Forward pass ----
|
||||
out = model(phone, tone, lang, speaker, durations, energy_t, bright_t, pitch_frame)
|
||||
token_mask = out["token_mask"]
|
||||
log_dur_t = torch.log1p(durations.float())
|
||||
group_log_dur_t, group_mask = group_duration_targets(phone, durations)
|
||||
|
||||
# Base mel losses
|
||||
mel_l1 = masked_l1(out["mel"], target_mel, frame_mask)
|
||||
mel_mse = masked_mse(out["mel"], target_mel, frame_mask)
|
||||
delta = masked_delta_loss(out["mel"], target_mel, frame_mask)
|
||||
accel = masked_accel_loss(out["mel"], target_mel, frame_mask)
|
||||
|
||||
# Token-level losses
|
||||
dur_loss = token_mse(out["log_dur"], log_dur_t, token_mask)
|
||||
group_dur_loss = token_mse(out["group_log_dur"], group_log_dur_t, group_mask)
|
||||
energy_loss = token_mse(out["energy"], energy_t, token_mask)
|
||||
bright_loss = token_mse(out["bright"], bright_t, token_mask)
|
||||
pitch_loss = token_mse_nd(out["pitch"], pitch_token_t, token_mask)
|
||||
|
||||
# ---- Multi-resolution STFT loss ----
|
||||
stft_loss = torch.zeros((), device=device)
|
||||
if args.stft_weight > 0:
|
||||
stft_loss = multi_resolution_stft_loss(
|
||||
out["mel"], target_mel, frame_mask,
|
||||
mel_frontend, consistency_vocoder, device,
|
||||
) * args.stft_weight
|
||||
|
||||
# ---- Prosody exposure bias loss ----
|
||||
predicted_prosody_mel_loss = torch.zeros((), device=device)
|
||||
predicted_prosody_delta_loss = torch.zeros((), device=device)
|
||||
if args.predicted_prosody_mel_weight > 0 or args.predicted_prosody_delta_weight > 0:
|
||||
pred_out = model(phone, tone, lang, speaker, durations)
|
||||
if args.predicted_prosody_mel_weight > 0:
|
||||
predicted_prosody_mel_loss = masked_l1(pred_out["mel"], target_mel, frame_mask)
|
||||
if args.predicted_prosody_delta_weight > 0:
|
||||
predicted_prosody_delta_loss = masked_delta_loss(pred_out["mel"], target_mel, frame_mask)
|
||||
|
||||
# ---- Robust prosody loss (mix of predicted + reference) ----
|
||||
robust_prosody_mel_loss = torch.zeros((), device=device)
|
||||
robust_prosody_delta_loss = torch.zeros((), device=device)
|
||||
if args.robust_prosody_mel_weight > 0 or args.robust_prosody_delta_weight > 0:
|
||||
robust_out = model(
|
||||
phone, tone, lang, speaker, durations,
|
||||
energy_t, bright_t, pitch_frame,
|
||||
predicted_prosody_mix=args.robust_prosody_mix,
|
||||
detach_mixed_predictions=True,
|
||||
)
|
||||
if args.robust_prosody_mel_weight > 0:
|
||||
robust_prosody_mel_loss = masked_l1(robust_out["mel"], target_mel, frame_mask)
|
||||
if args.robust_prosody_delta_weight > 0:
|
||||
robust_prosody_delta_loss = masked_delta_loss(robust_out["mel"], target_mel, frame_mask)
|
||||
|
||||
# ---- Vocoder consistency loss ----
|
||||
voc_wav_loss = torch.zeros((), device=device)
|
||||
voc_mel_loss = torch.zeros((), device=device)
|
||||
if consistency_vocoder is not None:
|
||||
if args.vocoder_wav_weight > 0:
|
||||
pred_wav = consistency_vocoder(out["mel"].clamp(-12, 2))
|
||||
voc_wav_loss = masked_wav_l1(pred_wav, target_wav, frame_mask, hifi_cfg.hop_size)
|
||||
if args.vocoder_mel_weight > 0:
|
||||
pred_wav = consistency_vocoder(out["mel"].clamp(-12, 2))
|
||||
pred_recon_mel = mel_frontend(pred_wav.squeeze(1))
|
||||
voc_mel_loss = masked_l1(pred_recon_mel, target_mel, frame_mask)
|
||||
|
||||
# ---- Adversarial mel loss ----
|
||||
adv_mel_loss = torch.zeros((), device=device)
|
||||
disc_loss = torch.zeros((), device=device)
|
||||
if mel_disc is not None and args.adv_mel_weight > 0:
|
||||
# Train discriminator
|
||||
common = min(out["mel"].shape[-1], target_mel.shape[-1])
|
||||
real_mel = target_mel[..., :common]
|
||||
fake_mel = out["mel"][..., :common].detach()
|
||||
|
||||
optim_d.zero_grad(set_to_none=True)
|
||||
real_fmap = mel_disc(real_mel)
|
||||
fake_fmap = mel_disc(fake_mel)
|
||||
disc_loss = discriminator_mel_loss(real_fmap, fake_fmap)
|
||||
disc_loss.backward()
|
||||
torch.nn.utils.clip_grad_norm_(mel_disc.parameters(), args.grad_clip)
|
||||
optim_d.step()
|
||||
|
||||
# Generator adversarial loss
|
||||
adv_fake_fmap = mel_disc(out["mel"][..., :common])
|
||||
adv_mel_loss = generator_mel_loss(adv_fake_fmap)
|
||||
|
||||
# Feature matching
|
||||
with torch.no_grad():
|
||||
real_fmap_detached = mel_disc(real_mel)
|
||||
fm_loss = torch.zeros((), device=device)
|
||||
for rf, ff in zip(real_fmap_detached, adv_fake_fmap):
|
||||
for rl, fl in zip(rf, ff):
|
||||
fm_loss = fm_loss + F.l1_loss(rl, fl)
|
||||
adv_mel_loss = adv_mel_loss + args.fm_weight * fm_loss
|
||||
|
||||
# ---- Total generator loss ----
|
||||
loss_g = (
|
||||
mel_l1
|
||||
+ args.mse_weight * mel_mse
|
||||
+ args.delta_weight * delta
|
||||
+ args.accel_weight * accel
|
||||
+ args.duration_weight * dur_loss
|
||||
+ args.group_duration_weight * group_dur_loss
|
||||
+ args.energy_weight * energy_loss
|
||||
+ args.bright_weight * bright_loss
|
||||
+ args.pitch_weight * pitch_loss
|
||||
+ args.predicted_prosody_mel_weight * predicted_prosody_mel_loss
|
||||
+ args.predicted_prosody_delta_weight * predicted_prosody_delta_loss
|
||||
+ args.robust_prosody_mel_weight * robust_prosody_mel_loss
|
||||
+ args.robust_prosody_delta_weight * robust_prosody_delta_loss
|
||||
+ args.vocoder_wav_weight * voc_wav_loss
|
||||
+ args.vocoder_mel_weight * voc_mel_loss
|
||||
+ stft_loss
|
||||
+ args.adv_mel_weight * adv_mel_loss
|
||||
)
|
||||
|
||||
optim_g.zero_grad(set_to_none=True)
|
||||
loss_g.backward()
|
||||
grad = torch.nn.utils.clip_grad_norm_(trainable_params, args.grad_clip)
|
||||
optim_g.step()
|
||||
|
||||
step += 1
|
||||
|
||||
if step == 1 or step % args.log_interval == 0:
|
||||
elapsed = max(1e-6, time.time() - started)
|
||||
speed = (step - start_step) / elapsed
|
||||
eta = (total_steps - step) / max(1e-6, speed)
|
||||
print(
|
||||
f"step={step}/{total_steps} loss={loss_g.item():.4f} mel={mel_l1.item():.4f} "
|
||||
f"mse={mel_mse.item():.4f} delta={delta.item():.4f} accel={accel.item():.4f} "
|
||||
f"dur={dur_loss.item():.4f} gdur={group_dur_loss.item():.4f} "
|
||||
f"energy={energy_loss.item():.4f} bright={bright_loss.item():.4f} "
|
||||
f"pitch={pitch_loss.item():.4f} stft={stft_loss.item():.4f} "
|
||||
f"pmel={predicted_prosody_mel_loss.item():.4f} "
|
||||
f"pdelta={predicted_prosody_delta_loss.item():.4f} "
|
||||
f"rmel={robust_prosody_mel_loss.item():.4f} "
|
||||
f"rdelta={robust_prosody_delta_loss.item():.4f} "
|
||||
f"vwav={voc_wav_loss.item():.4f} vmel={voc_mel_loss.item():.4f} "
|
||||
f"adv={adv_mel_loss.item():.4f} disc={disc_loss.item():.4f} "
|
||||
f"lr={lr_scale*args.lr:.2g} grad={float(grad):.2f} "
|
||||
f"speed={speed:.3f} step/s eta={eta/60:.1f}m",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
if step % args.save_interval == 0 or step >= total_steps:
|
||||
payload = {
|
||||
"model": model.state_dict(),
|
||||
"optim_g": optim_g.state_dict(),
|
||||
"step": step,
|
||||
"speakers": speakers,
|
||||
"params": count_parameters(model),
|
||||
}
|
||||
if mel_disc is not None:
|
||||
payload["mel_disc"] = mel_disc.state_dict()
|
||||
payload["optim_d"] = optim_d.state_dict()
|
||||
tmp = args.out_dir / f"inflect-smooth-{step}.pt.tmp"
|
||||
torch.save(payload, tmp)
|
||||
tmp.replace(args.out_dir / f"inflect-smooth-{step}.pt")
|
||||
torch.save(payload, args.out_dir / "inflect-smooth-latest.pt")
|
||||
|
||||
print(f"Done. {args.out_dir}")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# CLI
|
||||
# ---------------------------------------------------------------------------
|
||||
def main() -> None:
|
||||
ap = argparse.ArgumentParser(
|
||||
description="Train Inflect-Nano acoustic model with enhanced smoothness losses.",
|
||||
)
|
||||
# Data
|
||||
ap.add_argument("--durations-jsonl", type=Path, required=True)
|
||||
ap.add_argument("--out-dir", type=Path, required=True)
|
||||
ap.add_argument("--max-rows", type=int, default=0)
|
||||
|
||||
# Architecture
|
||||
ap.add_argument("--hidden", type=int, default=168)
|
||||
ap.add_argument("--encoder-layers", type=int, default=5)
|
||||
ap.add_argument("--decoder-layers", type=int, default=6)
|
||||
ap.add_argument("--decoder-ff-mult", type=int, default=3)
|
||||
ap.add_argument("--max-seconds", type=float, default=12.0)
|
||||
ap.add_argument("--max-frames", type=int, default=1400)
|
||||
ap.add_argument("--postnet-scale", type=float, default=0.35,
|
||||
help="Postnet refinement scale (higher = postnet has more influence)")
|
||||
ap.add_argument("--postnet-layers", type=int, default=5,
|
||||
help="Depth of residual postnet")
|
||||
ap.add_argument("--postnet-kernel", type=int, default=5,
|
||||
help="Kernel size for postnet convolutions")
|
||||
ap.add_argument("--abs-frame-bins", type=int, default=512)
|
||||
|
||||
# Training
|
||||
ap.add_argument("--steps", type=int, default=20000)
|
||||
ap.add_argument("--batch-size", type=int, default=6)
|
||||
ap.add_argument("--lr", type=float, default=2e-4)
|
||||
ap.add_argument("--weight-decay", type=float, default=1e-4)
|
||||
ap.add_argument("--warmup-steps", type=int, default=1000,
|
||||
help="Linear warmup steps for cosine LR schedule")
|
||||
ap.add_argument("--grad-clip", type=float, default=5.0)
|
||||
|
||||
# Loss weights (tuned for smoothness)
|
||||
ap.add_argument("--mse-weight", type=float, default=0.25)
|
||||
ap.add_argument("--delta-weight", type=float, default=0.25,
|
||||
help="Spectral delta smoothness (higher = smoother transitions)")
|
||||
ap.add_argument("--accel-weight", type=float, default=0.08,
|
||||
help="Spectral acceleration penalty (higher = less jitter)")
|
||||
ap.add_argument("--duration-weight", type=float, default=0.08)
|
||||
ap.add_argument("--group-duration-weight", type=float, default=0.02)
|
||||
ap.add_argument("--energy-weight", type=float, default=0.06)
|
||||
ap.add_argument("--bright-weight", type=float, default=0.06)
|
||||
ap.add_argument("--pitch-weight", type=float, default=0.06)
|
||||
|
||||
# Spectral / perceptual losses
|
||||
ap.add_argument("--stft-weight", type=float, default=0.15,
|
||||
help="Multi-resolution STFT loss (improves spectral smoothness)")
|
||||
ap.add_argument("--predicted-prosody-mel-weight", type=float, default=0.05,
|
||||
help="Exposure bias: use only predicted prosody for mel")
|
||||
ap.add_argument("--predicted-prosody-delta-weight", type=float, default=0.03)
|
||||
ap.add_argument("--robust-prosody-mix", type=float, default=0.5,
|
||||
help="Mix ratio for robust prosody training")
|
||||
ap.add_argument("--robust-prosody-mel-weight", type=float, default=0.05)
|
||||
ap.add_argument("--robust-prosody-delta-weight", type=float, default=0.03)
|
||||
ap.add_argument("--adv-mel-weight", type=float, default=0.08,
|
||||
help="Adversarial mel loss (pushes toward realistic spectrograms)")
|
||||
ap.add_argument("--fm-weight", type=float, default=2.0,
|
||||
help="Feature matching weight for adversarial loss")
|
||||
|
||||
# Vocoder consistency
|
||||
ap.add_argument("--vocoder-checkpoint", type=Path, default=None,
|
||||
help="Path to vocoder for consistency loss")
|
||||
ap.add_argument("--vocoder-wav-weight", type=float, default=0.12,
|
||||
help="Vocoder waveform consistency loss")
|
||||
ap.add_argument("--vocoder-mel-weight", type=float, default=0.08,
|
||||
help="Vocoder mel-reconstruction consistency loss")
|
||||
|
||||
# Checkpointing
|
||||
ap.add_argument("--init-checkpoint", type=Path,
|
||||
help="Start from an existing acoustic checkpoint for fine-tuning")
|
||||
ap.add_argument("--save-interval", type=int, default=2000)
|
||||
ap.add_argument("--log-interval", type=int, default=50)
|
||||
ap.add_argument("--seed", type=int, default=42)
|
||||
ap.add_argument("--resume", action="store_true")
|
||||
ap.add_argument("--preload-features", action="store_true")
|
||||
|
||||
# Misc
|
||||
ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
|
||||
ap.add_argument("--trainable", choices=[
|
||||
"all", "duration", "predictors", "heads", "contextual",
|
||||
"group_duration", "decoder_adapt",
|
||||
], default="all")
|
||||
ap.add_argument("--contextual-predictors", action="store_true")
|
||||
ap.add_argument("--group-duration-planner", action="store_true")
|
||||
|
||||
args = ap.parse_args()
|
||||
train_smooth(args)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,426 @@
|
||||
"""
|
||||
Preprocess an audio+text dataset into the durations.jsonl format required
|
||||
by Inflect-Nano training scripts.
|
||||
|
||||
Supports:
|
||||
- HuggingFace datasets (CMU ARCTIC, LJSpeech, etc.)
|
||||
- Local directory of .wav files with matching .txt transcriptions
|
||||
- LJSpeech-format metadata.csv
|
||||
|
||||
Output: a .jsonl file with one JSON object per line containing:
|
||||
phone_ids, tone_ids, lang_ids, hifigan_durations, target_audio, speaker_id, voice_id
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import math
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import soundfile as sf
|
||||
import torch
|
||||
import torchaudio
|
||||
|
||||
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))
|
||||
|
||||
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 HifiGanConfig, MelFrontend
|
||||
|
||||
|
||||
def text_to_ids(text: str) -> tuple[list[int], list[int], list[int]]:
|
||||
"""Convert English text to TinyTTS phone/tone/lang ID lists."""
|
||||
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)
|
||||
return phone_ids, tone_ids, lang_ids
|
||||
|
||||
|
||||
def estimate_durations_uniform(
|
||||
phone_ids: list[int],
|
||||
audio_path: str,
|
||||
mel_frontend: MelFrontend,
|
||||
sample_rate: int = 24000,
|
||||
) -> list[int]:
|
||||
"""Estimate mel-frame durations uniformly across all phones.
|
||||
|
||||
Loads the audio, extracts mel spectrogram, then distributes the total
|
||||
frame count evenly across phones. This is a coarse estimate suitable
|
||||
for decoder-only fine-tuning where the duration predictor is frozen.
|
||||
"""
|
||||
wav, sr = torchaudio.load(audio_path)
|
||||
if wav.shape[0] > 1:
|
||||
wav = wav.mean(dim=0, keepdim=True)
|
||||
if sr != sample_rate:
|
||||
wav = torchaudio.functional.resample(wav, sr, sample_rate)
|
||||
wav = wav.clamp(-1, 1)
|
||||
|
||||
with torch.no_grad():
|
||||
mel = mel_frontend(wav)
|
||||
total_frames = mel.shape[-1]
|
||||
|
||||
# Count non-blank phones (id != 0)
|
||||
visible_phones = [i for i, pid in enumerate(phone_ids) if pid != 0]
|
||||
|
||||
if not visible_phones:
|
||||
return [0] * len(phone_ids)
|
||||
|
||||
# Distribute frames: each blank gets 1 frame, rest split evenly
|
||||
blank_count = len(phone_ids) - len(visible_phones)
|
||||
remaining = max(0, total_frames - blank_count)
|
||||
base_dur = max(1, remaining // len(visible_phones))
|
||||
remainder = remaining - base_dur * len(visible_phones)
|
||||
|
||||
durations = []
|
||||
for pid in phone_ids:
|
||||
if pid == 0:
|
||||
durations.append(1)
|
||||
else:
|
||||
extra = 1 if remainder > 0 else 0
|
||||
remainder = max(0, remainder - 1)
|
||||
durations.append(base_dur + extra)
|
||||
|
||||
# Scale to match total_frames
|
||||
current_sum = sum(durations)
|
||||
if current_sum > 0 and current_sum != total_frames:
|
||||
scale = total_frames / current_sum
|
||||
scaled = [max(1, round(d * scale)) for d in durations]
|
||||
# Fix rounding errors
|
||||
diff = total_frames - sum(scaled)
|
||||
for i in range(abs(diff)):
|
||||
if diff > 0:
|
||||
scaled[i % len(scaled)] += 1
|
||||
else:
|
||||
if scaled[i % len(scaled)] > 1:
|
||||
scaled[i % len(scaled)] -= 1
|
||||
durations = scaled
|
||||
|
||||
# Ensure no zero durations for non-blank phones
|
||||
for i, pid in enumerate(phone_ids):
|
||||
if pid != 0 and durations[i] < 1:
|
||||
durations[i] = 1
|
||||
|
||||
return durations
|
||||
|
||||
|
||||
def process_hf_dataset(
|
||||
dataset_path: str,
|
||||
output_jsonl: Path,
|
||||
audio_dir: Path | None,
|
||||
subset: str = "default",
|
||||
split: str | None = None,
|
||||
text_key: str = "text",
|
||||
speaker_key: str = "speaker",
|
||||
max_rows: int = 0,
|
||||
voice_id: str = "speaker",
|
||||
) -> int:
|
||||
"""Process a HuggingFace dataset (loaded via `datasets` library)."""
|
||||
from datasets import load_dataset
|
||||
|
||||
print(f"Loading HF dataset: {dataset_path} subset={subset} split={split}")
|
||||
if split:
|
||||
ds = load_dataset(dataset_path, subset, split=split, trust_remote_code=True)
|
||||
else:
|
||||
ds_dict = load_dataset(dataset_path, subset, trust_remote_code=True)
|
||||
splits = list(ds_dict.keys())
|
||||
print(f"Available splits: {splits}")
|
||||
# Use first train split, or first available
|
||||
preferred = [s for s in splits if "train" in s.lower()]
|
||||
ds = ds_dict[preferred[0] if preferred else splits[0]]
|
||||
|
||||
mel_frontend = MelFrontend(HifiGanConfig(variant="v2plus"))
|
||||
output_jsonl.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
count = 0
|
||||
with output_jsonl.open("w", encoding="utf-8") as f:
|
||||
for i, row in enumerate(ds):
|
||||
text = str(row.get(text_key, "")).strip()
|
||||
if not text:
|
||||
continue
|
||||
|
||||
# Get audio path or array
|
||||
audio_info = row.get("audio", row.get("file", None))
|
||||
if audio_info is None:
|
||||
continue
|
||||
|
||||
if isinstance(audio_info, dict):
|
||||
# Audio is already loaded as array
|
||||
audio_path = None
|
||||
audio_array = audio_info.get("array")
|
||||
sample_rate = audio_info.get("sampling_rate", 24000)
|
||||
if audio_array is None:
|
||||
continue
|
||||
elif isinstance(audio_info, str):
|
||||
audio_path = audio_info
|
||||
if audio_dir:
|
||||
audio_path = str(audio_dir / Path(audio_info).name)
|
||||
if not Path(audio_path).is_file():
|
||||
continue
|
||||
audio_array = None
|
||||
sample_rate = 24000 # will be detected on load
|
||||
else:
|
||||
continue
|
||||
|
||||
try:
|
||||
phone_ids, tone_ids, lang_ids = text_to_ids(text)
|
||||
except Exception as e:
|
||||
print(f" Skipping row {i}: text-to-ids failed: {e}")
|
||||
continue
|
||||
|
||||
if not phone_ids:
|
||||
continue
|
||||
|
||||
# Estimate durations
|
||||
if audio_path:
|
||||
durations = estimate_durations_uniform(phone_ids, audio_path, mel_frontend)
|
||||
else:
|
||||
# Uniform fallback: 8 frames per phone
|
||||
durations = [8] * len(phone_ids)
|
||||
|
||||
speaker = str(row.get(speaker_key, voice_id))
|
||||
|
||||
row_out = {
|
||||
"phone_ids": phone_ids,
|
||||
"tone_ids": tone_ids,
|
||||
"lang_ids": lang_ids,
|
||||
"hifigan_durations": durations,
|
||||
"target_audio": audio_path or "",
|
||||
"speaker_id": hash(speaker) % 256,
|
||||
"voice_id": speaker,
|
||||
}
|
||||
f.write(json.dumps(row_out, ensure_ascii=False) + "\n")
|
||||
count += 1
|
||||
|
||||
if max_rows > 0 and count >= max_rows:
|
||||
break
|
||||
|
||||
if count % 100 == 0:
|
||||
print(f" Processed {count} rows...")
|
||||
|
||||
print(f"Wrote {count} rows to {output_jsonl}")
|
||||
return count
|
||||
|
||||
|
||||
def process_local_dir(
|
||||
audio_dir: Path,
|
||||
output_jsonl: Path,
|
||||
ext: str = ".wav",
|
||||
text_ext: str = ".txt",
|
||||
speaker: str = "speaker",
|
||||
max_rows: int = 0,
|
||||
) -> int:
|
||||
"""Process a local directory of .wav files with matching .txt files."""
|
||||
audio_files = sorted(audio_dir.glob(f"*{ext}"))
|
||||
if not audio_files:
|
||||
audio_files = sorted(audio_dir.rglob(f"*{ext}"))
|
||||
|
||||
mel_frontend = MelFrontend(HifiGanConfig(variant="v2plus"))
|
||||
output_jsonl.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
count = 0
|
||||
with output_jsonl.open("w", encoding="utf-8") as f:
|
||||
for wav_path in audio_files:
|
||||
txt_path = wav_path.with_suffix(text_ext)
|
||||
if not txt_path.is_file():
|
||||
continue
|
||||
|
||||
text = txt_path.read_text(encoding="utf-8").strip()
|
||||
if not text:
|
||||
continue
|
||||
|
||||
try:
|
||||
phone_ids, tone_ids, lang_ids = text_to_ids(text)
|
||||
except Exception as e:
|
||||
print(f" Skipping {wav_path.name}: {e}")
|
||||
continue
|
||||
|
||||
durations = estimate_durations_uniform(phone_ids, str(wav_path), mel_frontend)
|
||||
|
||||
row_out = {
|
||||
"phone_ids": phone_ids,
|
||||
"tone_ids": tone_ids,
|
||||
"lang_ids": lang_ids,
|
||||
"hifigan_durations": durations,
|
||||
"target_audio": str(wav_path.resolve()),
|
||||
"speaker_id": hash(speaker) % 256,
|
||||
"voice_id": speaker,
|
||||
}
|
||||
f.write(json.dumps(row_out, ensure_ascii=False) + "\n")
|
||||
count += 1
|
||||
|
||||
if max_rows > 0 and count >= max_rows:
|
||||
break
|
||||
|
||||
if count % 100 == 0:
|
||||
print(f" Processed {count} rows...")
|
||||
|
||||
print(f"Wrote {count} rows to {output_jsonl}")
|
||||
return count
|
||||
|
||||
|
||||
def process_metadata_csv(
|
||||
csv_path: Path,
|
||||
audio_dir: Path,
|
||||
output_jsonl: Path,
|
||||
sep: str = "|",
|
||||
header: bool = False,
|
||||
text_col: int = 1,
|
||||
file_col: int = 0,
|
||||
speaker: str = "speaker",
|
||||
max_rows: int = 0,
|
||||
) -> int:
|
||||
"""Process an LJSpeech-style metadata.csv."""
|
||||
mel_frontend = MelFrontend(HifiGanConfig(variant="v2plus"))
|
||||
output_jsonl.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
count = 0
|
||||
with csv_path.open("r", encoding="utf-8") as csv_file:
|
||||
lines = csv_file.readlines()
|
||||
if header:
|
||||
lines = lines[1:]
|
||||
|
||||
with output_jsonl.open("w", encoding="utf-8") as out_f:
|
||||
for line in lines:
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
|
||||
parts = line.split(sep)
|
||||
if len(parts) < max(text_col, file_col) + 1:
|
||||
continue
|
||||
|
||||
filename = parts[file_col].strip()
|
||||
if not filename.endswith(".wav"):
|
||||
filename += ".wav"
|
||||
wav_path = audio_dir / filename
|
||||
if not wav_path.is_file():
|
||||
print(f" Missing: {wav_path}")
|
||||
continue
|
||||
|
||||
text = parts[text_col].strip()
|
||||
if not text:
|
||||
continue
|
||||
|
||||
try:
|
||||
phone_ids, tone_ids, lang_ids = text_to_ids(text)
|
||||
except Exception as e:
|
||||
print(f" Skipping {filename}: {e}")
|
||||
continue
|
||||
|
||||
durations = estimate_durations_uniform(phone_ids, str(wav_path), mel_frontend)
|
||||
|
||||
row_out = {
|
||||
"phone_ids": phone_ids,
|
||||
"tone_ids": tone_ids,
|
||||
"lang_ids": lang_ids,
|
||||
"hifigan_durations": durations,
|
||||
"target_audio": str(wav_path.resolve()),
|
||||
"speaker_id": hash(speaker) % 256,
|
||||
"voice_id": speaker,
|
||||
}
|
||||
out_f.write(json.dumps(row_out, ensure_ascii=False) + "\n")
|
||||
count += 1
|
||||
|
||||
if max_rows > 0 and count >= max_rows:
|
||||
break
|
||||
|
||||
if count % 100 == 0:
|
||||
print(f" Processed {count} rows...")
|
||||
|
||||
print(f"Wrote {count} rows to {output_jsonl}")
|
||||
return count
|
||||
|
||||
|
||||
def main() -> None:
|
||||
ap = argparse.ArgumentParser(
|
||||
description="Preprocess a speech dataset for Inflect-Nano training.",
|
||||
)
|
||||
|
||||
sub = ap.add_subparsers(dest="mode", required=True)
|
||||
|
||||
# HuggingFace mode
|
||||
p_hf = sub.add_parser("hf", help="Process a HuggingFace dataset")
|
||||
p_hf.add_argument("--dataset", required=True, help="HF dataset path, e.g. MikhailT/cmu-arctic")
|
||||
p_hf.add_argument("--subset", default="default")
|
||||
p_hf.add_argument("--split", default=None, help="e.g. 'rms' for CMU ARCTIC single speaker")
|
||||
p_hf.add_argument("--text-key", default="text")
|
||||
p_hf.add_argument("--speaker-key", default="speaker")
|
||||
p_hf.add_argument("--audio-dir", type=Path, default=None, help="Local dir for audio if HF paths are remote")
|
||||
p_hf.add_argument("--out", type=Path, required=True, help="Output .jsonl path")
|
||||
p_hf.add_argument("--max-rows", type=int, default=0)
|
||||
p_hf.add_argument("--voice-id", default="speaker")
|
||||
|
||||
# Local dir mode
|
||||
p_dir = sub.add_parser("local", help="Process local directory of .wav+.txt pairs")
|
||||
p_dir.add_argument("--audio-dir", type=Path, required=True)
|
||||
p_dir.add_argument("--ext", default=".wav")
|
||||
p_dir.add_argument("--text-ext", default=".txt")
|
||||
p_dir.add_argument("--speaker", default="speaker")
|
||||
p_dir.add_argument("--out", type=Path, required=True)
|
||||
p_dir.add_argument("--max-rows", type=int, default=0)
|
||||
|
||||
# Metadata CSV mode (LJSpeech-style)
|
||||
p_csv = sub.add_parser("csv", help="Process a metadata.csv file")
|
||||
p_csv.add_argument("--csv", type=Path, required=True)
|
||||
p_csv.add_argument("--audio-dir", type=Path, required=True)
|
||||
p_csv.add_argument("--sep", default="|")
|
||||
p_csv.add_argument("--header", action="store_true")
|
||||
p_csv.add_argument("--file-col", type=int, default=0)
|
||||
p_csv.add_argument("--text-col", type=int, default=1)
|
||||
p_csv.add_argument("--speaker", default="speaker")
|
||||
p_csv.add_argument("--out", type=Path, required=True)
|
||||
p_csv.add_argument("--max-rows", type=int, default=0)
|
||||
|
||||
args = ap.parse_args()
|
||||
|
||||
if args.mode == "hf":
|
||||
process_hf_dataset(
|
||||
dataset_path=args.dataset,
|
||||
output_jsonl=args.out,
|
||||
audio_dir=args.audio_dir,
|
||||
subset=args.subset,
|
||||
split=args.split,
|
||||
text_key=args.text_key,
|
||||
speaker_key=args.speaker_key,
|
||||
max_rows=args.max_rows,
|
||||
voice_id=args.voice_id,
|
||||
)
|
||||
elif args.mode == "local":
|
||||
process_local_dir(
|
||||
audio_dir=args.audio_dir,
|
||||
output_jsonl=args.out,
|
||||
ext=args.ext,
|
||||
text_ext=args.text_ext,
|
||||
speaker=args.speaker,
|
||||
max_rows=args.max_rows,
|
||||
)
|
||||
elif args.mode == "csv":
|
||||
process_metadata_csv(
|
||||
csv_path=args.csv,
|
||||
audio_dir=args.audio_dir,
|
||||
output_jsonl=args.out,
|
||||
sep=args.sep,
|
||||
header=args.header,
|
||||
file_col=args.file_col,
|
||||
text_col=args.text_col,
|
||||
speaker=args.speaker,
|
||||
max_rows=args.max_rows,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -5,3 +5,6 @@ numpy
|
||||
g2p_en
|
||||
transformers
|
||||
gradio
|
||||
numba
|
||||
scipy
|
||||
datasets
|
||||
|
||||
@@ -0,0 +1,126 @@
|
||||
"""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!")
|
||||
Reference in New Issue
Block a user