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
282 lines
8.6 KiB
Plaintext
282 lines
8.6 KiB
Plaintext
# 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).
|