diff --git a/colab_run.py b/colab_run.py new file mode 100644 index 0000000..3723582 --- /dev/null +++ b/colab_run.py @@ -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() diff --git a/colab_smooth_finetune.ipynb b/colab_smooth_finetune.ipynb index b106dc7..31b4dfd 100644 --- a/colab_smooth_finetune.ipynb +++ b/colab_smooth_finetune.ipynb @@ -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 +} \ No newline at end of file