training/train_diffusion.py

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