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_run.py
+170
View File
@@ -0,0 +1,170 @@
"""
Colab runner for Inflect-Nano smoothness fine-tuning.
Single script that handles the full pipeline with graceful checkpointing —
safe to re-run after interruptions. Each step is skipped if its output
already exists.
Usage in Colab:
!wget https://your-server/colab_run.py
!python colab_run.py
"""
from __future__ import annotations
import os
import subprocess
import sys
import time
from pathlib import Path
# ---------------------------------------------------------------------------
# Config — change these
# ---------------------------------------------------------------------------
REPO_URL = "https://git.quietidiot.com/michael-treadgold/speech-nano.git"
REPO_DIR = Path("/content/speech-nano")
DATASET = "MikhailT/cmu-arctic"
SPEAKER = "rms" # "rms", "bdl", "jmk", "awb", "ksp" (CMU ARCTIC male speakers)
TRAIN_STEPS = 5000
BATCH_SIZE = 8
# ---------------------------------------------------------------------------
# Step helpers
# ---------------------------------------------------------------------------
def step(name: str) -> None:
print(f"\n{'='*60}\n {name}\n{'='*60}")
def run(cmd: str, **kwargs) -> bool:
print(f" $ {cmd}")
result = subprocess.run(cmd, shell=True, **kwargs)
return result.returncode == 0
def already(path: Path | str) -> bool:
p = Path(path)
ok = p.exists()
if ok:
print(f" ⏭ Already exists: {p}")
return ok
def pip_installed(pkg: str) -> bool:
result = subprocess.run(
[sys.executable, "-c", f"import {pkg}"],
capture_output=True, text=True, timeout=30,
)
return result.returncode == 0
def main() -> None:
start = time.time()
# ---- Step 1: Clone or pull repo ----
step("1/6 Clone / pull repo")
if REPO_DIR.exists():
print(f" Repo exists, pulling latest...")
subprocess.run(["git", "-C", str(REPO_DIR), "pull"], check=False)
subprocess.run(["git", "-C", str(REPO_DIR), "lfs", "pull"], check=False)
else:
run(f"git clone {REPO_URL} {REPO_DIR}")
sys.path.insert(0, str(REPO_DIR))
sys.path.insert(0, str(REPO_DIR / "third_party" / "tiny_tts_frontend"))
os.chdir(REPO_DIR)
print(f" Working dir: {REPO_DIR}")
# ---- Step 2: Install dependencies ----
step("2/6 Install dependencies")
pkgs = "torch torchaudio soundfile numpy g2p_en transformers gradio numba scipy datasets"
if not pip_installed("g2p_en"):
run(f"{sys.executable} -m pip install -q {pkgs}")
else:
print(" ⏭ Already installed (g2p_en found)")
import nltk
nltk.download("averaged_perceptron_tagger_eng", quiet=True)
nltk.download("cmudict", quiet=True)
print(" NLTK data OK")
# ---- Step 3: Download model weights ----
step("3/6 Download model weights")
WEIGHTS_DIR = REPO_DIR / "weights"
AC_WEIGHTS = WEIGHTS_DIR / "inflect_nano_v1_acoustic.pt"
VO_WEIGHTS = WEIGHTS_DIR / "inflect_nano_v1_vocoder.pt"
if AC_WEIGHTS.exists() and AC_WEIGHTS.stat().st_size > 1000:
print(f" ⏭ Acoustic weights OK ({AC_WEIGHTS.stat().st_size / 1e6:.1f} MB)")
else:
run(f"{sys.executable} download_weights.py")
# ---- Step 4: Preprocess dataset ----
step("4/6 Preprocess dataset")
JSONL = Path(f"/content/durations_{SPEAKER}.jsonl")
if JSONL.exists() and JSONL.stat().st_size > 100:
print(f" ⏭ JSONL exists ({JSONL.stat().st_size / 1e6:.2f} MB)")
else:
ok = run(
f"{sys.executable} preprocess_dataset.py hf "
f"--dataset {DATASET} --split {SPEAKER} "
f"--out {JSONL} --voice-id mark"
)
if not ok:
print(" ❌ Preprocessing failed. Check error above.")
sys.exit(1)
# ---- Step 5: Train ----
step("5/6 Train smoothness model")
OUT_DIR = Path("/content/checkpoints/smooth-v1")
LATEST_CKPT = OUT_DIR / "inflect-smooth-latest.pt"
import torch
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f" Device: {device} ({torch.cuda.get_device_name(0) if device == 'cuda' else 'CPU'})")
if LATEST_CKPT.exists():
print(f" ⏭ Checkpoint exists — adding {TRAIN_STEPS} more steps")
resume_flag = "--resume"
else:
print(f" Starting fresh — {TRAIN_STEPS} steps")
resume_flag = ""
ok = run(
f"{sys.executable} -m inflect_nano.train_smooth "
f"--durations-jsonl {JSONL} "
f"--out-dir {OUT_DIR} "
f"--init-checkpoint {AC_WEIGHTS} "
f"--vocoder-checkpoint {VO_WEIGHTS} "
f"--steps {TRAIN_STEPS} "
f"--batch-size {BATCH_SIZE} "
f"--device {device} "
f"{resume_flag}"
)
if not ok:
print(" ❌ Training failed. Check error above.")
sys.exit(1)
# ---- Step 6: Generate sample ----
step("6/6 Generate sample")
from inference import load_acoustic, load_vocoder, synthesize
import soundfile as sf
import numpy as np
ac_path = LATEST_CKPT if LATEST_CKPT.exists() else AC_WEIGHTS
acoustic, speakers, ap = load_acoustic(ac_path, torch.device(device))
vocoder, vp = load_vocoder(VO_WEIGHTS, torch.device(device))
print(f" Acoustic: {ap:,} params Vocoder: {vp:,} params Total: {ap+vp:,}")
text = "Every man is destined to die, but his work echoes through the ages."
audio = synthesize(
text, acoustic, vocoder, speakers, torch.device(device),
smooth_prosody=True, mel_smooth_sigma=0.8,
)
out = Path("/content/sample_smooth.wav")
sf.write(str(out), audio, 24000, subtype="PCM_16")
elapsed = (time.time() - start) / 60
print(f"\n{'='*60}")
print(f" Done in {elapsed:.0f}m. Sample: {out} ({audio.size/24000:.1f}s)")
print(f" Checkpoints: {OUT_DIR}")
print(f"{'='*60}")
if __name__ == "__main__":
main()
colab_smooth_finetune.ipynb
+47 -281
View File
@@ -1,281 +1,47 @@
# Inflect-Nano Smoothness Fine-Tuning — Google Colab Notebook {
"cells": [
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_type": "markdown",
"metadata": {},
--- "source": [
"# Inflect-Nano Smoothness Fine-Tuning\n",
## Cell 1: Setup — clone repo, install deps "Single-click pipeline. Each cell is safe to re-run \u2014 it skips steps already completed.\n",
"\n",
```python "**Runtime \u2192 Change runtime type \u2192 T4 GPU**"
# @title Setup environment (run once) ]
import os, sys, subprocess },
from pathlib import Path {
"cell_type": "code",
REPO_URL = "https://huggingface.co/owensong/Inflect-Nano-v1" "metadata": {
REPO_DIR = "/content/Inflect-Nano-v1" "id": "runner"
},
# Clone "source": [
if not Path(REPO_DIR).exists(): "# Download the runner script and execute it\n",
!git clone {REPO_URL} {REPO_DIR} "import os, sys, subprocess\n",
else: "from pathlib import Path\n",
%cd {REPO_DIR} "\n",
!git pull "REPO = Path(\"/content/speech-nano\")\n",
"if not REPO.exists():\n",
%cd {REPO_DIR} " subprocess.run([\"git\", \"clone\", \"https://git.quietidiot.com/michael-treadgold/speech-nano.git\", str(REPO)], check=False)\n",
"else:\n",
# Install deps (numba is needed by vendored frontend; scipy for lowpass) " subprocess.run([\"git\", \"-C\", str(REPO), \"pull\"], check=False)\n",
!pip install -q torch torchaudio soundfile numpy g2p_en transformers gradio numba scipy datasets "\n",
"sys.path.insert(0, str(REPO))\n",
# Download NLTK data "os.chdir(REPO)\n",
import nltk "\n",
nltk.download('averaged_perceptron_tagger_eng', quiet=True) "# Run the pipeline\n",
nltk.download('cmudict', quiet=True) "exec(open(\"colab_run.py\").read())"
],
print("✓ Setup complete") "execution_count": null,
print(f" PyTorch {torch.__version__} | GPU: {torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'N/A'}") "outputs": []
``` }
],
--- "metadata": {
"accelerator": "GPU",
## Cell 2: Pick a dataset & preprocess "colab": {
"provenance": []
Choose one: }
},
| Dataset | Speakers | Size | Best for | "nbformat": 4,
|---------|----------|------|----------| "nbformat_minor": 5
| **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).