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

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()