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,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).
|
||||
Reference in New Issue
Block a user