michael-treadgold/speech-nano
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add colab_run.py -- single-script pipeline with graceful checkpointing; proper .ipynb notebook

Browse Source
This commit is contained in:
Michael Treadgold
2026-06-18 20:29:05 +12:00
parent 1bc7d89ce9
commit 0ded3603ce
2 changed files with 217 additions and 281 deletions
Download Patch File Download Diff File Expand all files Collapse all files
colab_smooth_finetune.ipynb
+47 -281
View File
@@ -1,281 +1,47 @@
# 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).
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Inflect-Nano Smoothness Fine-Tuning\n",
"Single-click pipeline. Each cell is safe to re-run \u2014 it skips steps already completed.\n",
"\n",
"**Runtime \u2192 Change runtime type \u2192 T4 GPU**"
]
},
{
"cell_type": "code",
"metadata": {
"id": "runner"
},
"source": [
"# Download the runner script and execute it\n",
"import os, sys, subprocess\n",
"from pathlib import Path\n",
"\n",
"REPO = Path(\"/content/speech-nano\")\n",
"if not REPO.exists():\n",
" subprocess.run([\"git\", \"clone\", \"https://git.quietidiot.com/michael-treadgold/speech-nano.git\", str(REPO)], check=False)\n",
"else:\n",
" subprocess.run([\"git\", \"-C\", str(REPO), \"pull\"], check=False)\n",
"\n",
"sys.path.insert(0, str(REPO))\n",
"os.chdir(REPO)\n",
"\n",
"# Run the pipeline\n",
"exec(open(\"colab_run.py\").read())"
],
"execution_count": null,
"outputs": []
}
],
"metadata": {
"accelerator": "GPU",
"colab": {
"provenance": []
}
},
"nbformat": 4,
"nbformat_minor": 5
}