# 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).