Upload 9 files

uploaded training folder

training/config.py

@@ -0,0 +1,89 @@
+"""Training hyperparameters."""
+
+from dataclasses import dataclass
+from pathlib import Path
+
+
+@dataclass
+class AutoencoderConfig:
+    # Data
+    data_dir: Path = Path("data/processed")
+    val_split: float = 0.1
+
+    # Model
+    latent_dim: int = 4
+
+    # Training
+    batch_size: int = 32
+    learning_rate: float = 1e-4
+    epochs: int = 100
+    num_workers: int = 4
+
+    # KL annealing: weight ramps from kl_weight_start to kl_weight_end
+    # over the first kl_anneal_epochs epochs
+    kl_weight_start: float = 0.0001
+    kl_weight_end: float = 0.001
+    kl_anneal_epochs: int = 20
+
+    # Mixed precision
+    use_amp: bool = True
+
+    # Checkpointing
+    checkpoint_dir: Path = Path("checkpoints")
+    checkpoint_every: int = 10
+
+    # Logging
+    log_dir: Path = Path("runs/autoencoder")
+
+    def kl_weight_at_epoch(self, epoch: int) -> float:
+        """Linearly anneal KL weight over first N epochs."""
+        if epoch >= self.kl_anneal_epochs:
+            return self.kl_weight_end
+        t = epoch / self.kl_anneal_epochs
+        return self.kl_weight_start + t * (
+            self.kl_weight_end - self.kl_weight_start
+        )
+
+
+@dataclass
+class DiffusionConfig:
+    # Data
+    data_dir: Path = Path("data/processed")
+    latents_dir: Path = Path("data/latents")
+    metadata_csv: Path = Path("data/metadata.csv")
+    val_split: float = 0.1
+
+    # VAE
+    vae_checkpoint: Path = Path("checkpoints/vae_epoch_100.pt")
+    latent_dim: int = 4
+
+    # Noise schedule
+    timesteps: int = 1000
+    beta_start: float = 0.0001
+    beta_end: float = 0.02
+
+    # Model
+    base_channels: int = 64
+    cond_dim: int = 256
+    text_embed_dim: int = 64
+
+    # Training
+    batch_size: int = 16
+    learning_rate: float = 1e-4
+    iterations: int = 100_000
+    num_workers: int = 4
+    gradient_accumulation: int = 2
+    use_amp: bool = True
+
+    # EMA
+    ema_decay: float = 0.9999
+
+    # Classifier-free guidance
+    cfg_dropout: float = 0.15
+
+    # Checkpointing
+    checkpoint_dir: Path = Path("checkpoints")
+    checkpoint_every: int = 5000  # iterations
+
+    # Logging
+    log_dir: Path = Path("runs/diffusion")

training/losses.py

@@ -0,0 +1,54 @@
+"""Loss functions for autoencoder training."""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def reconstruction_loss(
+    recon: torch.Tensor, target: torch.Tensor
+) -> torch.Tensor:
+    """MSE reconstruction loss."""
+    return F.mse_loss(recon, target)
+
+
+def spectral_convergence_loss(
+    recon: torch.Tensor, target: torch.Tensor
+) -> torch.Tensor:
+    """L1 loss on magnitude (encourages spectral fidelity)."""
+    return F.l1_loss(recon, target)
+
+
+def kl_divergence(
+    mu: torch.Tensor, logvar: torch.Tensor
+) -> torch.Tensor:
+    """KL divergence from N(mu, sigma) to N(0, 1)."""
+    return -0.5 * torch.mean(1 + logvar - mu.pow(2) - logvar.exp())
+
+
+def vae_loss(
+    recon: torch.Tensor,
+    target: torch.Tensor,
+    mu: torch.Tensor,
+    logvar: torch.Tensor,
+    kl_weight: float = 0.0001,
+) -> tuple[torch.Tensor, dict[str, float]]:
+    """Combined VAE loss.
+
+    Returns:
+        Total loss tensor and dict of individual loss values for logging.
+    """
+    mse = reconstruction_loss(recon, target)
+    spectral = spectral_convergence_loss(recon, target)
+    kl = kl_divergence(mu, logvar)
+
+    total = mse + spectral + kl_weight * kl
+
+    metrics = {
+        "mse": mse.item(),
+        "spectral": spectral.item(),
+        "kl": kl.item(),
+        "kl_weighted": (kl_weight * kl).item(),
+        "total": total.item(),
+    }
+    return total, metrics

training/train_autoencoder.py

@@ -0,0 +1,169 @@
+"""Training script for the kick drum VAE."""
+
+import sys
+from pathlib import Path
+
+import torch
+from torch.utils.data import DataLoader, Dataset, random_split
+from torch.utils.tensorboard import SummaryWriter
+from tqdm import tqdm
+
+# Add project root to path
+sys.path.insert(0, str(Path(__file__).parent.parent))
+
+from models.autoencoder import KickVAE
+from training.config import AutoencoderConfig
+from training.losses import vae_loss
+
+
+class MelDataset(Dataset):
+    """Dataset of preprocessed mel spectrogram tensors."""
+
+    def __init__(self, data_dir: Path) -> None:
+        self.files = sorted(
+            f for f in data_dir.glob("*.pt") if not f.name.startswith("._")
+        )
+        if not self.files:
+            raise FileNotFoundError(
+                f"No .pt files found in {data_dir}"
+            )
+
+    def __len__(self) -> int:
+        return len(self.files)
+
+    def __getitem__(self, idx: int) -> torch.Tensor:
+        return torch.load(self.files[idx], weights_only=False)
+
+
+def train(cfg: AutoencoderConfig | None = None) -> None:
+    """Run VAE training."""
+    if cfg is None:
+        cfg = AutoencoderConfig()
+
+    device = torch.device(
+        "cuda" if torch.cuda.is_available()
+        else "mps" if torch.backends.mps.is_available()
+        else "cpu"
+    )
+    print(f"Using device: {device}")
+
+    # Data
+    dataset = MelDataset(cfg.data_dir)
+    val_size = int(len(dataset) * cfg.val_split)
+    train_size = len(dataset) - val_size
+    train_set, val_set = random_split(
+        dataset, [train_size, val_size],
+        generator=torch.Generator().manual_seed(42),
+    )
+
+    train_loader = DataLoader(
+        train_set,
+        batch_size=cfg.batch_size,
+        shuffle=True,
+        num_workers=cfg.num_workers,
+        pin_memory=True,
+    )
+    val_loader = DataLoader(
+        val_set,
+        batch_size=cfg.batch_size,
+        shuffle=False,
+        num_workers=cfg.num_workers,
+        pin_memory=True,
+    )
+
+    print(f"Train: {train_size}, Val: {val_size}")
+
+    # Model
+    model = KickVAE(latent_dim=cfg.latent_dim).to(device)
+    optimizer = torch.optim.AdamW(model.parameters(), lr=cfg.learning_rate)
+    scaler = torch.amp.GradScaler(enabled=cfg.use_amp and device.type == "cuda")
+
+    # Logging
+    cfg.log_dir.mkdir(parents=True, exist_ok=True)
+    cfg.checkpoint_dir.mkdir(parents=True, exist_ok=True)
+    writer = SummaryWriter(cfg.log_dir)
+
+    global_step = 0
+
+    for epoch in range(cfg.epochs):
+        model.train()
+        kl_weight = cfg.kl_weight_at_epoch(epoch)
+
+        epoch_metrics: dict[str, float] = {}
+        epoch_count = 0
+
+        pbar = tqdm(train_loader, desc=f"Epoch {epoch+1}/{cfg.epochs}")
+        for batch in pbar:
+            batch = batch.to(device)
+
+            with torch.amp.autocast(
+                device_type=device.type,
+                enabled=cfg.use_amp and device.type == "cuda",
+            ):
+                recon, mu, logvar = model(batch)
+                loss, metrics = vae_loss(
+                    recon, batch, mu, logvar, kl_weight
+                )
+
+            optimizer.zero_grad()
+            scaler.scale(loss).backward()
+            scaler.step(optimizer)
+            scaler.update()
+
+            # Accumulate metrics
+            for k, v in metrics.items():
+                epoch_metrics[k] = epoch_metrics.get(k, 0.0) + v
+            epoch_count += 1
+            global_step += 1
+
+            pbar.set_postfix(loss=f"{metrics['total']:.4f}")
+
+        # Log epoch averages
+        for k, v in epoch_metrics.items():
+            writer.add_scalar(f"train/{k}", v / epoch_count, epoch)
+        writer.add_scalar("train/kl_weight", kl_weight, epoch)
+
+        # Validation
+        model.eval()
+        val_metrics: dict[str, float] = {}
+        val_count = 0
+
+        with torch.no_grad():
+            for batch in val_loader:
+                batch = batch.to(device)
+                recon, mu, logvar = model(batch)
+                _, metrics = vae_loss(
+                    recon, batch, mu, logvar, kl_weight
+                )
+                for k, v in metrics.items():
+                    val_metrics[k] = val_metrics.get(k, 0.0) + v
+                val_count += 1
+
+        avg_val_loss = val_metrics.get("total", 0.0) / max(val_count, 1)
+        for k, v in val_metrics.items():
+            writer.add_scalar(f"val/{k}", v / val_count, epoch)
+
+        print(
+            f"Epoch {epoch+1}: "
+            f"train={epoch_metrics['total']/epoch_count:.4f} "
+            f"val={avg_val_loss:.4f} "
+            f"kl_w={kl_weight:.6f}"
+        )
+
+        # Checkpoint
+        if (epoch + 1) % cfg.checkpoint_every == 0:
+            path = cfg.checkpoint_dir / f"vae_epoch_{epoch+1}.pt"
+            torch.save({
+                "epoch": epoch + 1,
+                "model_state_dict": model.state_dict(),
+                "optimizer_state_dict": optimizer.state_dict(),
+                "config": cfg,
+            }, path)
+            print(f"Saved checkpoint: {path}")
+
+    writer.close()
+    print("Training complete.")
+
+
+if __name__ == "__main__":
+    train()

training/train_diffusion.py

@@ -0,0 +1,314 @@
+"""Training script for the latent diffusion model."""
+
+import copy
+import csv
+import sys
+from pathlib import Path
+
+import torch
+from torch.utils.data import DataLoader, Dataset, random_split
+from torch.utils.tensorboard import SummaryWriter
+from tqdm import tqdm
+
+sys.path.insert(0, str(Path(__file__).parent.parent))
+
+from models.autoencoder import KickVAE
+from models.diffusion import LatentUNet, NoiseScheduler
+from models.text_encoder import KeywordEncoder, build_vocab
+from training.config import DiffusionConfig
+
+
+# ---------------------------------------------------------------------------
+# Dataset
+# ---------------------------------------------------------------------------
+
+class LatentDataset(Dataset):
+    """Dataset of pre-encoded VAE latents with keyword token IDs."""
+
+    def __init__(
+        self,
+        latents_dir: Path,
+        metadata_csv: Path,
+        vocab: list[str],
+    ) -> None:
+        self.latent_files = sorted(
+            f for f in latents_dir.glob("*.pt") if not f.name.startswith("._")
+        )
+        if not self.latent_files:
+            raise FileNotFoundError(f"No .pt files in {latents_dir}")
+
+        # Build keyword lookup: filename_stem -> list of keyword strings
+        self.kw_to_idx = {kw: i for i, kw in enumerate(vocab)}
+        self.keywords: dict[str, list[int]] = {}
+        with open(metadata_csv) as f:
+            reader = csv.DictReader(f)
+            for row in reader:
+                stem = Path(row["filename"]).stem
+                ids = []
+                for kw in row["keywords"].split(","):
+                    kw = kw.strip().lower()
+                    if kw in self.kw_to_idx:
+                        ids.append(self.kw_to_idx[kw])
+                self.keywords[stem] = ids
+
+    def __len__(self) -> int:
+        return len(self.latent_files)
+
+    def __getitem__(self, idx: int) -> tuple[torch.Tensor, list[int]]:
+        path = self.latent_files[idx]
+        latent = torch.load(path, weights_only=False)
+        # latent filename matches mel filename stem
+        stem = path.stem
+        token_ids = self.keywords.get(stem, [])
+        return latent, token_ids
+
+
+def collate_fn(
+    batch: list[tuple[torch.Tensor, list[int]]],
+) -> tuple[torch.Tensor, list[list[int]]]:
+    """Custom collate to handle variable-length keyword lists."""
+    latents = torch.stack([b[0] for b in batch])
+    token_ids = [b[1] for b in batch]
+    return latents, token_ids
+
+
+# ---------------------------------------------------------------------------
+# Pre-encode latents
+# ---------------------------------------------------------------------------
+
+def pre_encode_latents(cfg: DiffusionConfig) -> None:
+    """Encode all mel spectrograms to latents using frozen VAE."""
+    cfg.latents_dir.mkdir(parents=True, exist_ok=True)
+
+    # Check if already done
+    existing = list(cfg.latents_dir.glob("*.pt"))
+    if len(existing) > 100:
+        print(f"Latents dir already has {len(existing)} files, skipping encoding.")
+        return
+
+    device = torch.device(
+        "cuda" if torch.cuda.is_available()
+        else "mps" if torch.backends.mps.is_available()
+        else "cpu"
+    )
+
+    # Load VAE
+    checkpoint = torch.load(cfg.vae_checkpoint, weights_only=False)
+    vae = KickVAE(latent_dim=cfg.latent_dim).to(device)
+    vae.load_state_dict(checkpoint["model_state_dict"])
+    vae.eval()
+
+    mel_files = sorted(
+        f for f in cfg.data_dir.glob("*.pt") if not f.name.startswith("._")
+    )
+    print(f"Encoding {len(mel_files)} mel spectrograms to latents...")
+
+    with torch.no_grad():
+        for f in tqdm(mel_files):
+            out_path = cfg.latents_dir / f.name
+            if out_path.exists():
+                continue
+            mel = torch.load(f, weights_only=False).unsqueeze(0).to(device)
+            latent = vae.encode(mel).squeeze(0).cpu()
+            torch.save(latent, out_path)
+
+    print("Latent encoding complete.")
+
+
+# ---------------------------------------------------------------------------
+# EMA
+# ---------------------------------------------------------------------------
+
+class EMA:
+    """Exponential moving average of model parameters."""
+
+    def __init__(self, model: torch.nn.Module, decay: float = 0.9999) -> None:
+        self.decay = decay
+        self.shadow = copy.deepcopy(model)
+        self.shadow.eval()
+        for p in self.shadow.parameters():
+            p.requires_grad_(False)
+
+    @torch.no_grad()
+    def update(self, model: torch.nn.Module) -> None:
+        for s, p in zip(self.shadow.parameters(), model.parameters()):
+            s.data.mul_(self.decay).add_(p.data, alpha=1 - self.decay)
+
+
+# ---------------------------------------------------------------------------
+# Training
+# ---------------------------------------------------------------------------
+
+def train(cfg: DiffusionConfig | None = None) -> None:
+    if cfg is None:
+        cfg = DiffusionConfig()
+
+    # Pre-encode latents
+    pre_encode_latents(cfg)
+
+    device = torch.device(
+        "cuda" if torch.cuda.is_available()
+        else "mps" if torch.backends.mps.is_available()
+        else "cpu"
+    )
+    print(f"Using device: {device}")
+
+    # Build vocab and dataset
+    vocab = build_vocab(cfg.metadata_csv)
+    print(f"Vocabulary size: {len(vocab)}")
+
+    dataset = LatentDataset(cfg.latents_dir, cfg.metadata_csv, vocab)
+    val_size = int(len(dataset) * cfg.val_split)
+    train_size = len(dataset) - val_size
+    train_set, val_set = random_split(
+        dataset, [train_size, val_size],
+        generator=torch.Generator().manual_seed(42),
+    )
+
+    train_loader = DataLoader(
+        train_set,
+        batch_size=cfg.batch_size,
+        shuffle=True,
+        num_workers=cfg.num_workers,
+        pin_memory=True,
+        collate_fn=collate_fn,
+    )
+    val_loader = DataLoader(
+        val_set,
+        batch_size=cfg.batch_size,
+        shuffle=False,
+        num_workers=cfg.num_workers,
+        pin_memory=True,
+        collate_fn=collate_fn,
+    )
+    print(f"Train: {train_size}, Val: {val_size}")
+
+    # Model
+    model = LatentUNet(
+        latent_dim=cfg.latent_dim,
+        base_channels=cfg.base_channels,
+        cond_dim=cfg.cond_dim,
+    ).to(device)
+    text_enc = KeywordEncoder(
+        vocab_size=len(vocab),
+        embed_dim=cfg.text_embed_dim,
+        cond_dim=cfg.cond_dim,
+    ).to(device)
+
+    scheduler = NoiseScheduler(cfg.timesteps, cfg.beta_start, cfg.beta_end).to(device)
+    ema = EMA(model, cfg.ema_decay)
+
+    optimizer = torch.optim.AdamW(
+        list(model.parameters()) + list(text_enc.parameters()),
+        lr=cfg.learning_rate,
+    )
+    scaler = torch.amp.GradScaler(enabled=cfg.use_amp and device.type == "cuda")
+
+    # Logging
+    cfg.log_dir.mkdir(parents=True, exist_ok=True)
+    cfg.checkpoint_dir.mkdir(parents=True, exist_ok=True)
+    writer = SummaryWriter(cfg.log_dir)
+
+    # Training loop (iteration-based)
+    global_step = 0
+    model.train()
+    text_enc.train()
+
+    print(f"Training for {cfg.iterations} iterations...")
+
+    while global_step < cfg.iterations:
+        for latents, token_ids in train_loader:
+            if global_step >= cfg.iterations:
+                break
+
+            latents = latents.to(device)
+            batch_size = latents.shape[0]
+
+            # Classifier-free guidance dropout: replace keywords with empty list
+            dropped_ids = []
+            for ids in token_ids:
+                if torch.rand(1).item() < cfg.cfg_dropout:
+                    dropped_ids.append([])
+                else:
+                    dropped_ids.append(ids)
+
+            # Sample timesteps and noise
+            t = torch.randint(0, cfg.timesteps, (batch_size,), device=device)
+            noise = torch.randn_like(latents)
+            noisy = scheduler.add_noise(latents, noise, t)
+
+            with torch.amp.autocast(
+                device_type=device.type,
+                enabled=cfg.use_amp and device.type == "cuda",
+            ):
+                cond = text_enc(dropped_ids, device)
+                pred_noise = model(noisy, t, cond)
+                loss = torch.nn.functional.mse_loss(pred_noise, noise)
+                loss = loss / cfg.gradient_accumulation
+
+            scaler.scale(loss).backward()
+
+            if (global_step + 1) % cfg.gradient_accumulation == 0:
+                scaler.step(optimizer)
+                scaler.update()
+                optimizer.zero_grad()
+                ema.update(model)
+
+            # Logging
+            if global_step % 50 == 0:
+                writer.add_scalar(
+                    "train/loss", loss.item() * cfg.gradient_accumulation, global_step
+                )
+
+            if global_step % 500 == 0:
+                print(
+                    f"Step {global_step}/{cfg.iterations} "
+                    f"loss={loss.item() * cfg.gradient_accumulation:.6f}"
+                )
+
+            # Validation
+            if global_step % 1000 == 0 and global_step > 0:
+                model.eval()
+                text_enc.eval()
+                val_loss_sum = 0.0
+                val_count = 0
+                with torch.no_grad():
+                    for vl, vt in val_loader:
+                        vl = vl.to(device)
+                        vt_step = torch.randint(
+                            0, cfg.timesteps, (vl.shape[0],), device=device
+                        )
+                        vn = torch.randn_like(vl)
+                        vnoisy = scheduler.add_noise(vl, vn, vt_step)
+                        vcond = text_enc(vt, device)
+                        vpred = model(vnoisy, vt_step, vcond)
+                        val_loss_sum += torch.nn.functional.mse_loss(vpred, vn).item()
+                        val_count += 1
+                avg_val = val_loss_sum / max(val_count, 1)
+                writer.add_scalar("val/loss", avg_val, global_step)
+                print(f"  val_loss={avg_val:.6f}")
+                model.train()
+                text_enc.train()
+
+            # Checkpoint
+            if (global_step + 1) % cfg.checkpoint_every == 0:
+                path = cfg.checkpoint_dir / f"diffusion_step_{global_step+1}.pt"
+                torch.save({
+                    "step": global_step + 1,
+                    "model_state_dict": model.state_dict(),
+                    "ema_state_dict": ema.shadow.state_dict(),
+                    "text_enc_state_dict": text_enc.state_dict(),
+                    "optimizer_state_dict": optimizer.state_dict(),
+                    "vocab": vocab,
+                    "config": cfg,
+                }, path)
+                print(f"Saved checkpoint: {path}")
+
+            global_step += 1
+
+    writer.close()
+    print("Diffusion training complete.")
+
+
+if __name__ == "__main__":
+    train()

training/train_vocoder.py

@@ -0,0 +1,352 @@
+"""
+HiFi-GAN vocoder training.
+
+Trains a mel-to-waveform generator with multi-period and multi-scale discriminators.
+Designed for 6GB VRAM: lazy-loads audio, uses random 8192-sample segments, small batch size.
+
+Usage:
+    uv run training/train_vocoder.py
+    uv run training/train_vocoder.py --batch-size 4 --segment-size 8192
+"""
+
+import argparse
+import sys
+from pathlib import Path
+
+import numpy as np
+import soundfile as sf
+import torch
+import torch.nn.functional as F
+from torch.utils.data import Dataset, DataLoader
+from torch.utils.tensorboard import SummaryWriter
+from tqdm import tqdm
+
+sys.path.insert(0, str(Path(__file__).parent.parent))
+
+from models.vocoder import HiFiGANGenerator, MultiPeriodDiscriminator, MultiScaleDiscriminator
+
+# Audio params (must match preprocess.py)
+SAMPLE_RATE = 44100
+N_FFT = 2048
+HOP_LENGTH = 512
+N_MELS = 128
+DURATION_SECONDS = 2.0
+TARGET_SAMPLES = int(SAMPLE_RATE * DURATION_SECONDS)
+
+
+# ---------------------------------------------------------------------------
+# Dataset - lazy loads audio, pairs with pre-computed mel
+# ---------------------------------------------------------------------------
+
+class VocoderDataset(Dataset):
+    """Lazy-loading dataset that pairs raw audio with processed mel spectrograms.
+
+    Only stores file paths in memory. Loads audio on-the-fly per __getitem__.
+    Returns random segments of `segment_size` samples for memory efficiency.
+    """
+
+    def __init__(self, raw_dir: Path, processed_dir: Path, segment_size: int = 8192) -> None:
+        self.segment_size = segment_size
+        self.raw_dir = raw_dir
+        self.processed_dir = processed_dir
+
+        # Build list of (raw_path, mel_path) pairs
+        self.pairs: list[tuple[Path, Path]] = []
+        mel_stems = {p.stem for p in processed_dir.glob("*.pt")}
+        for raw_path in sorted(raw_dir.iterdir()):
+            if raw_path.suffix.lower() not in (".wav", ".aif", ".aiff", ".mp3", ".flac"):
+                continue
+            stem = raw_path.stem
+            if stem in mel_stems:
+                self.pairs.append((raw_path, processed_dir / f"{stem}.pt"))
+
+        print(f"VocoderDataset: {len(self.pairs)} paired samples found")
+
+    def __len__(self) -> int:
+        return len(self.pairs)
+
+    def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor]:
+        raw_path, mel_path = self.pairs[idx]
+
+        # Load raw audio
+        audio, sr = sf.read(raw_path, dtype="float32", always_2d=True)
+        audio = audio[:, 0]  # mono
+        if sr != SAMPLE_RATE:
+            # Simple resample via linear interpolation
+            audio = np.interp(
+                np.linspace(0, len(audio) - 1, int(len(audio) * SAMPLE_RATE / sr)),
+                np.arange(len(audio)),
+                audio,
+            ).astype(np.float32)
+
+        # Pad/trim to target length
+        if len(audio) < TARGET_SAMPLES:
+            audio = np.pad(audio, (0, TARGET_SAMPLES - len(audio)))
+        else:
+            audio = audio[:TARGET_SAMPLES]
+
+        # Load pre-computed mel
+        mel = torch.load(mel_path, weights_only=True)  # (1, 128, 173)
+
+        # Pick random segment
+        # segment_size audio samples = segment_size // HOP_LENGTH mel frames
+        mel_frames = self.segment_size // HOP_LENGTH
+        max_mel_start = mel.shape[-1] - mel_frames
+        if max_mel_start > 0:
+            mel_start = torch.randint(0, max_mel_start, (1,)).item()
+        else:
+            mel_start = 0
+
+        audio_start = mel_start * HOP_LENGTH
+        audio_end = audio_start + self.segment_size
+
+        mel_seg = mel[:, :, mel_start:mel_start + mel_frames]  # (1, 128, mel_frames)
+        audio_seg = torch.from_numpy(audio[audio_start:audio_end])  # (segment_size,)
+
+        # Pad if needed (edge cases)
+        if mel_seg.shape[-1] < mel_frames:
+            mel_seg = F.pad(mel_seg, (0, mel_frames - mel_seg.shape[-1]))
+        if audio_seg.shape[-1] < self.segment_size:
+            audio_seg = F.pad(audio_seg, (0, self.segment_size - audio_seg.shape[-1]))
+
+        return mel_seg.squeeze(0), audio_seg.unsqueeze(0)  # (128, mel_frames), (1, segment_size)
+
+
+# ---------------------------------------------------------------------------
+# Loss functions
+# ---------------------------------------------------------------------------
+
+def generator_adversarial_loss(disc_outputs: list[torch.Tensor]) -> torch.Tensor:
+    loss = 0.0
+    for dg in disc_outputs:
+        loss = loss + torch.mean((1 - dg) ** 2)
+    return loss
+
+
+def discriminator_loss(real_outputs: list[torch.Tensor], fake_outputs: list[torch.Tensor]) -> torch.Tensor:
+    loss = 0.0
+    for dr, dg in zip(real_outputs, fake_outputs):
+        loss = loss + torch.mean((1 - dr) ** 2) + torch.mean(dg ** 2)
+    return loss
+
+
+def feature_matching_loss(real_fmaps: list[list[torch.Tensor]], fake_fmaps: list[list[torch.Tensor]]) -> torch.Tensor:
+    loss = 0.0
+    for rf, ff in zip(real_fmaps, fake_fmaps):
+        for r, f in zip(rf, ff):
+            loss = loss + F.l1_loss(f, r.detach())
+    return loss
+
+
+def mel_spectrogram_loss(y: torch.Tensor, y_hat: torch.Tensor) -> torch.Tensor:
+    """L1 loss on mel spectrograms of real vs generated audio."""
+    mel_transform = torch.nn.Sequential(
+        torch.nn.Identity(),  # placeholder
+    )
+    # Compute mel on-the-fly for loss
+    # Use torchaudio for consistency
+    import torchaudio
+    mel_spec = torchaudio.transforms.MelSpectrogram(
+        sample_rate=SAMPLE_RATE, n_fft=N_FFT, hop_length=HOP_LENGTH,
+        n_mels=N_MELS, power=1.0,
+    ).to(y.device)
+
+    mel_real = torch.log(mel_spec(y.squeeze(1)).clamp(min=1e-5))
+    mel_fake = torch.log(mel_spec(y_hat.squeeze(1)).clamp(min=1e-5))
+    return F.l1_loss(mel_real, mel_fake)
+
+
+# ---------------------------------------------------------------------------
+# Training loop
+# ---------------------------------------------------------------------------
+
+def train(args: argparse.Namespace) -> None:
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Device: {device}")
+
+    if device.type == "cuda":
+        vram_gb = torch.cuda.get_device_properties(0).total_memory / 1e9
+        print(f"GPU VRAM: {vram_gb:.1f} GB")
+
+    # Dataset
+    dataset = VocoderDataset(
+        raw_dir=Path(args.raw_dir),
+        processed_dir=Path(args.processed_dir),
+        segment_size=args.segment_size,
+    )
+    loader = DataLoader(
+        dataset,
+        batch_size=args.batch_size,
+        shuffle=True,
+        num_workers=args.num_workers,
+        pin_memory=True,
+        drop_last=True,
+        persistent_workers=args.num_workers > 0,
+    )
+
+    # Models
+    generator = HiFiGANGenerator(in_channels=N_MELS).to(device)
+    mpd = MultiPeriodDiscriminator().to(device)
+    msd = MultiScaleDiscriminator().to(device)
+
+    # Print param counts
+    g_params = sum(p.numel() for p in generator.parameters()) / 1e6
+    d_params = (sum(p.numel() for p in mpd.parameters()) + sum(p.numel() for p in msd.parameters())) / 1e6
+    print(f"Generator: {g_params:.1f}M params | Discriminators: {d_params:.1f}M params")
+
+    # Optimizers
+    optim_g = torch.optim.AdamW(generator.parameters(), lr=args.lr, betas=(0.8, 0.99))
+    optim_d = torch.optim.AdamW(
+        list(mpd.parameters()) + list(msd.parameters()),
+        lr=args.lr, betas=(0.8, 0.99),
+    )
+
+    # Schedulers
+    sched_g = torch.optim.lr_scheduler.ExponentialLR(optim_g, gamma=0.999)
+    sched_d = torch.optim.lr_scheduler.ExponentialLR(optim_d, gamma=0.999)
+
+    # Checkpointing - resume if exists
+    checkpoint_dir = Path(args.checkpoint_dir)
+    checkpoint_dir.mkdir(parents=True, exist_ok=True)
+    start_epoch = 0
+
+    resume_path = checkpoint_dir / "vocoder_latest.pt"
+    if resume_path.exists():
+        print(f"Resuming from {resume_path}")
+        ckpt = torch.load(resume_path, weights_only=False, map_location=device)
+        generator.load_state_dict(ckpt["generator"])
+        mpd.load_state_dict(ckpt["mpd"])
+        msd.load_state_dict(ckpt["msd"])
+        optim_g.load_state_dict(ckpt["optim_g"])
+        optim_d.load_state_dict(ckpt["optim_d"])
+        sched_g.load_state_dict(ckpt["sched_g"])
+        sched_d.load_state_dict(ckpt["sched_d"])
+        start_epoch = ckpt["epoch"] + 1
+        print(f"Resumed at epoch {start_epoch}")
+
+    # Logging
+    log_dir = Path(args.log_dir)
+    log_dir.mkdir(parents=True, exist_ok=True)
+    writer = SummaryWriter(log_dir)
+
+    # Training
+    global_step = start_epoch * len(loader)
+    for epoch in range(start_epoch, args.epochs):
+        generator.train()
+        mpd.train()
+        msd.train()
+
+        pbar = tqdm(loader, desc=f"Epoch {epoch+1}/{args.epochs}")
+        for mel, audio in pbar:
+            mel = mel.to(device)       # (B, 128, mel_frames)
+            audio = audio.to(device)   # (B, 1, segment_size)
+
+            # ---- Discriminator step ----
+            optim_d.zero_grad()
+            with torch.no_grad():
+                audio_fake = generator(mel)
+            # Trim to match lengths
+            min_len = min(audio.shape[-1], audio_fake.shape[-1])
+            audio_t = audio[..., :min_len]
+            audio_f = audio_fake[..., :min_len]
+
+            mpd_real, _ = mpd(audio_t)
+            mpd_fake, _ = mpd(audio_f)
+            msd_real, _ = msd(audio_t)
+            msd_fake, _ = msd(audio_f)
+
+            loss_d = discriminator_loss(mpd_real, mpd_fake) + discriminator_loss(msd_real, msd_fake)
+
+            loss_d.backward()
+            optim_d.step()
+
+            # ---- Generator step ----
+            optim_g.zero_grad()
+            audio_fake = generator(mel)
+            min_len = min(audio.shape[-1], audio_fake.shape[-1])
+            audio_t = audio[..., :min_len]
+            audio_f = audio_fake[..., :min_len]
+
+            mpd_real, mpd_real_fmap = mpd(audio_t)
+            mpd_fake, mpd_fake_fmap = mpd(audio_f)
+            msd_real, msd_real_fmap = msd(audio_t)
+            msd_fake, msd_fake_fmap = msd(audio_f)
+
+            loss_gen = generator_adversarial_loss(mpd_fake) + generator_adversarial_loss(msd_fake)
+            loss_fm = feature_matching_loss(mpd_real_fmap, mpd_fake_fmap) + feature_matching_loss(msd_real_fmap, msd_fake_fmap)
+            loss_mel = mel_spectrogram_loss(audio_t, audio_f)
+
+            loss_g = loss_gen + 2.0 * loss_fm + 45.0 * loss_mel
+
+            loss_g.backward()
+            optim_g.step()
+
+            global_step += 1
+            pbar.set_postfix(loss_g=f"{loss_g.item():.3f}", loss_d=f"{loss_d.item():.3f}")
+
+            if global_step % 100 == 0:
+                writer.add_scalar("loss/generator", loss_g.item(), global_step)
+                writer.add_scalar("loss/discriminator", loss_d.item(), global_step)
+                writer.add_scalar("loss/mel", loss_mel.item(), global_step)
+                writer.add_scalar("loss/feature_matching", loss_fm.item(), global_step)
+
+        # End of epoch
+        sched_g.step()
+        sched_d.step()
+
+        # Save checkpoint
+        if (epoch + 1) % args.checkpoint_every == 0 or epoch == args.epochs - 1:
+            ckpt_path = checkpoint_dir / f"vocoder_epoch_{epoch+1}.pt"
+            torch.save({
+                "generator": generator.state_dict(),
+                "mpd": mpd.state_dict(),
+                "msd": msd.state_dict(),
+                "optim_g": optim_g.state_dict(),
+                "optim_d": optim_d.state_dict(),
+                "sched_g": sched_g.state_dict(),
+                "sched_d": sched_d.state_dict(),
+                "epoch": epoch,
+            }, ckpt_path)
+            print(f"Saved {ckpt_path}")
+
+        # Always save latest for resume
+        torch.save({
+            "generator": generator.state_dict(),
+            "mpd": mpd.state_dict(),
+            "msd": msd.state_dict(),
+            "optim_g": optim_g.state_dict(),
+            "optim_d": optim_d.state_dict(),
+            "sched_g": sched_g.state_dict(),
+            "sched_d": sched_d.state_dict(),
+            "epoch": epoch,
+        }, checkpoint_dir / "vocoder_latest.pt")
+
+    writer.close()
+    print("Training complete.")
+
+    # Save inference-ready checkpoint
+    torch.save({
+        "generator_state_dict": generator.state_dict(),
+    }, checkpoint_dir / "vocoder.pt")
+    print(f"Saved inference checkpoint: {checkpoint_dir / 'vocoder.pt'}")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Train HiFi-GAN vocoder")
+    parser.add_argument("--raw-dir", type=str, default="data/raw")
+    parser.add_argument("--processed-dir", type=str, default="data/processed")
+    parser.add_argument("--batch-size", type=int, default=8)
+    parser.add_argument("--segment-size", type=int, default=8192)
+    parser.add_argument("--lr", type=float, default=2e-4)
+    parser.add_argument("--epochs", type=int, default=50)
+    parser.add_argument("--num-workers", type=int, default=2)
+    parser.add_argument("--use-amp", action="store_true", default=True)
+    parser.add_argument("--checkpoint-dir", type=str, default="checkpoints")
+    parser.add_argument("--checkpoint-every", type=int, default=10)
+    parser.add_argument("--log-dir", type=str, default="runs/vocoder")
+    args = parser.parse_args()
+    train(args)
+
+
+if __name__ == "__main__":
+    main()