train.py

import os
import time
import math
import json
import argparse
import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader, Dataset
from torch.amp import GradScaler, autocast
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

from model import EvoTalk, EvoTalkConfig


LEARNING_RATE = 1e-4
WEIGHT_DECAY = 0.01
BETAS = (0.9, 0.98)
GRAD_CLIP = 1.0
WARMUP_STEPS = 4000

# The ~80M model uses more memory than the old ~20M one. If you hit CUDA OOM,
# drop BATCH_SIZE (e.g. 8) and/or raise GRAD_ACCUM-equivalent by training longer.
BATCH_SIZE = 32
NUM_WORKERS = 4

MAX_STEPS = 150000
SAVE_EVERY = 5000
LOG_EVERY = 50
EVAL_EVERY = 500
EVAL_STEPS = 50

# When True, the decoder is fed the predictor's own (detached) pitch/energy
# during training instead of ground-truth. This removes the teacher-forcing leak
# that otherwise makes inference collapse. Set via --predicted_variance (used for
# fine-tuning from a checkpoint trained with teacher forcing).
EMBED_PREDICTED = False

GRAPH_DIR = "graphs"
CHECKPOINT_DIR = "checkpoints"
RESUME_FROM = None


def smooth(values, window=50):
    if len(values) < 2:
        return values
    w = min(window, len(values))
    kernel = torch.ones(w) / w
    t = torch.tensor(values, dtype=torch.float32)
    pad = w // 2
    t_padded = torch.cat([t[:pad].flip(0), t, t[-pad:].flip(0)])
    smoothed = torch.nn.functional.conv1d(
        t_padded.view(1, 1, -1), kernel.view(1, 1, -1), padding=0
    ).squeeze()
    return smoothed[:len(values)].tolist()


def save_graphs(history, graph_dir):
    os.makedirs(graph_dir, exist_ok=True)
    steps = history["steps"]
    if len(steps) < 2:
        return

    style = dict(linewidth=1.2, alpha=0.9)
    raw_style = dict(linewidth=0.5, alpha=0.25)

    fig, ax = plt.subplots(figsize=(10, 4))
    raw = history["loss"]
    s = smooth(raw)
    ax.plot(steps, raw, color="#4C8BE8", **raw_style)
    ax.plot(steps, s, color="#4C8BE8", label="train", **style)
    if history["val_steps"]:
        ax.plot(history["val_steps"], history["val_loss"], "o--",
                color="#E8744C", label="val", linewidth=1.2, markersize=3)
    ax.set_title("Total Loss")
    ax.set_xlabel("step")
    ax.set_ylabel("loss")
    ax.legend()
    ax.grid(True, alpha=0.3)
    fig.tight_layout()
    fig.savefig(os.path.join(graph_dir, "loss_total.png"), dpi=120)
    plt.close(fig)

    fig, axes = plt.subplots(1, 3, figsize=(15, 4))
    components = [
        ("mel", "#4C8BE8", "Mel Loss"),
        ("dur", "#50C87A", "Duration Loss"),
        ("energy", "#C850A8", "Energy Loss"),
    ]
    for ax, (key, color, title) in zip(axes, components):
        raw = history[key]
        s = smooth(raw)
        ax.plot(steps, raw, color=color, **raw_style)
        ax.plot(steps, s, color=color, **style)
        ax.set_title(title)
        ax.set_xlabel("step")
        ax.grid(True, alpha=0.3)
    fig.tight_layout()
    fig.savefig(os.path.join(graph_dir, "loss_components.png"), dpi=120)
    plt.close(fig)

    fig, ax = plt.subplots(figsize=(10, 4))
    raw = history["pitch"]
    s = smooth(raw)
    ax.plot(steps, raw, color="#E8C84C", **raw_style)
    ax.plot(steps, s, color="#E8C84C", **style)
    ax.set_title("Pitch Loss")
    ax.set_xlabel("step")
    ax.set_ylabel("loss")
    ax.grid(True, alpha=0.3)
    fig.tight_layout()
    fig.savefig(os.path.join(graph_dir, "loss_pitch.png"), dpi=120)
    plt.close(fig)

    fig, ax = plt.subplots(figsize=(10, 4))
    ax.plot(steps, history["lr"], color="#888888", **style)
    ax.set_title("Learning Rate")
    ax.set_xlabel("step")
    ax.set_ylabel("lr")
    ax.grid(True, alpha=0.3)
    fig.tight_layout()
    fig.savefig(os.path.join(graph_dir, "learning_rate.png"), dpi=120)
    plt.close(fig)


class LibriTTSDataset(Dataset):
    """Single-speaker dataset. Pitch and energy targets are normalized here with
    the GLOBAL training statistics from metadata.json, so the exact same normalized
    representation is used for (a) the variance-predictor loss target and (b) the
    pitch/energy embedding at both training and inference time."""

    def __init__(self, data_dir, split, pitch_mean, pitch_std, energy_mean, energy_std):
        self.data_dir = data_dir
        self.split = split
        self.pitch_mean = pitch_mean
        self.pitch_std = pitch_std
        self.energy_mean = energy_mean
        self.energy_std = energy_std
        self.samples = self._load_samples()

    def _load_samples(self):
        manifest = os.path.join(self.data_dir, self.split, "manifest.txt")
        samples = []
        with open(manifest, "r") as f:
            for line in f:
                line = line.strip()
                if line:
                    samples.append(line)
        return samples

    def __len__(self):
        return len(self.samples)

    def _norm_pitch(self, pitch):
        return (torch.log1p(pitch.clamp(min=0.0)) - self.pitch_mean) / self.pitch_std

    def _norm_energy(self, energy):
        return (energy - self.energy_mean) / self.energy_std

    @staticmethod
    def _to_phoneme_level(frame_values, durations):
        """Average a per-frame contour into one value per phoneme using the
        (integer) durations. sum(durations) == n_frames by construction, so the
        segmentation is exact. This is what makes pitch/energy phoneme-level."""
        d = durations.long()
        T = int(d.sum().item())
        fv = frame_values[:T]
        out = torch.zeros(d.size(0), dtype=torch.float32)
        idx = 0
        for i, di in enumerate(d.tolist()):
            if di > 0:
                seg = fv[idx:idx + di]
                if seg.numel() > 0:
                    out[i] = seg.mean()
            idx += di
        return out

    def __getitem__(self, idx):
        path = self.samples[idx]
        data = torch.load(path, weights_only=True)
        durations = data["durations"]
        # Per-frame pitch/energy (stored by prepare.py) -> phoneme-level averages.
        phon_pitch = self._to_phoneme_level(data["pitch"], durations)
        phon_energy = self._to_phoneme_level(data["energy"], durations)
        return (
            data["phonemes"],
            data["speaker_id"],
            durations,
            self._norm_pitch(phon_pitch),
            self._norm_energy(phon_energy),
            data["mel"],
        )


def collate_fn(batch):
    phonemes, speaker_ids, durations, pitches, energies, mels = zip(*batch)

    phoneme_lengths = [p.size(0) for p in phonemes]
    mel_lengths = [m.size(0) for m in mels]

    max_phoneme_len = max(phoneme_lengths)
    max_mel_len = max(mel_lengths)
    n_mels = mels[0].size(1)

    phonemes_padded = torch.zeros(len(batch), max_phoneme_len, dtype=torch.long)
    durations_padded = torch.zeros(len(batch), max_phoneme_len)
    mels_padded = torch.zeros(len(batch), max_mel_len, n_mels)
    # pitch/energy are now PHONEME-LEVEL, so they pad to phoneme length.
    pitches_padded = torch.zeros(len(batch), max_phoneme_len)
    energies_padded = torch.zeros(len(batch), max_phoneme_len)
    src_mask = torch.zeros(len(batch), max_phoneme_len, dtype=torch.bool)

    for i, (p, d, pi, e, m) in enumerate(zip(phonemes, durations, pitches, energies, mels)):
        phonemes_padded[i, :p.size(0)] = p
        durations_padded[i, :d.size(0)] = d
        mels_padded[i, :m.size(0)] = m
        pitches_padded[i, :pi.size(0)] = pi
        energies_padded[i, :e.size(0)] = e
        src_mask[i, p.size(0):] = True

    speaker_ids = torch.stack(speaker_ids)
    mel_lengths = torch.tensor(mel_lengths, dtype=torch.long)

    return {
        "phonemes": phonemes_padded,
        "speaker_ids": speaker_ids,
        "durations": durations_padded,
        "pitch": pitches_padded,
        "energy": energies_padded,
        "mels": mels_padded,
        "src_mask": src_mask,
        "mel_lengths": mel_lengths,
        "max_mel_len": max_mel_len,
    }


def get_lr(step):
    if step < WARMUP_STEPS:
        return LEARNING_RATE * step / WARMUP_STEPS
    progress = (step - WARMUP_STEPS) / (MAX_STEPS - WARMUP_STEPS)
    progress = min(max(progress, 0.0), 1.0)
    return LEARNING_RATE * 0.5 * (1.0 + math.cos(math.pi * progress))


def frame_mask_from_lengths(lengths, max_len, device):
    """(B, max_len) bool: True where the frame is valid (t < length)."""
    positions = torch.arange(max_len, device=device).unsqueeze(0)
    return positions < lengths.to(device).unsqueeze(1)


def masked_mse(pred, target, valid_mask):
    """Mean squared error over valid entries only. valid_mask is (B, T) with True
    marking valid frames; it is broadcast over any trailing feature dimension."""
    if valid_mask.dim() == pred.dim() - 1:
        valid_mask = valid_mask.unsqueeze(-1)
    m = valid_mask.to(pred.dtype)
    diff = (pred - target) ** 2 * m
    denom = m.expand_as(pred).sum().clamp(min=1.0)
    return diff.sum() / denom


def masked_l1(pred, target, valid_mask):
    """Mean absolute error over valid entries only. L1 on the mel produces sharper,
    less over-smoothed spectrograms than MSE (standard in FastSpeech2-style models)."""
    if valid_mask.dim() == pred.dim() - 1:
        valid_mask = valid_mask.unsqueeze(-1)
    m = valid_mask.to(pred.dtype)
    diff = (pred - target).abs() * m
    denom = m.expand_as(pred).sum().clamp(min=1.0)
    return diff.sum() / denom


def compute_loss(mel_out, duration_preds, pitch_preds, energy_preds, batch):
    mels = batch["mels"]
    durations = batch["durations"]
    pitch = batch["pitch"]        # already normalized in the dataset
    energy = batch["energy"]      # already normalized in the dataset
    src_mask = batch["src_mask"]  # True = padded phoneme
    mel_lengths = batch["mel_lengths"]
    device = mel_out.device

    # --- Mel loss (L1, masked over valid frames) ---
    T_mel = min(mel_out.size(1), mels.size(1))
    valid_frames = frame_mask_from_lengths(mel_lengths, T_mel, device)
    mel_loss = masked_l1(mel_out[:, :T_mel], mels[:, :T_mel], valid_frames)

    # --- Duration loss (masked over valid phonemes) ---
    phon_valid = ~src_mask
    dur_targets = torch.clamp(durations.float(), min=1.0)
    log_dur_targets = torch.log(dur_targets + 1.0)
    dur_loss = masked_mse(duration_preds, log_dur_targets, phon_valid)

    # --- Pitch / energy loss (now PHONEME-LEVEL, masked over valid phonemes) ---
    T_var = min(pitch_preds.size(1), pitch.size(1))
    valid_var = phon_valid[:, :T_var]
    pitch_loss = masked_mse(pitch_preds[:, :T_var], pitch[:, :T_var], valid_var)
    energy_loss = masked_mse(energy_preds[:, :T_var], energy[:, :T_var], valid_var)

    total = mel_loss + dur_loss + pitch_loss + energy_loss
    return total, mel_loss, dur_loss, pitch_loss, energy_loss


def save_checkpoint(model, optimizer, scaler, step, loss, path):
    os.makedirs(os.path.dirname(path), exist_ok=True)
    torch.save({
        "step": step,
        "model": model.state_dict(),
        "optimizer": optimizer.state_dict(),
        "scaler": scaler.state_dict(),
        "loss": loss,
        "config": model.config,
    }, path)
    print(f"saved checkpoint at step {step} -> {path}")


def load_checkpoint(path, model, optimizer, scaler, device):
    ckpt = torch.load(path, map_location=device, weights_only=False)
    model.load_state_dict(ckpt["model"])
    optimizer.load_state_dict(ckpt["optimizer"])
    scaler.load_state_dict(ckpt["scaler"])
    print(f"resumed from step {ckpt['step']} (loss {ckpt['loss']:.4f})")
    return ckpt["step"]


@torch.no_grad()
def evaluate(model, val_loader, device, max_steps):
    model.eval()
    total_loss = 0.0
    steps = 0
    for batch in val_loader:
        if steps >= max_steps:
            break
        batch = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in batch.items()}
        mel_out, duration_preds, pitch_preds, energy_preds, _ = model(
            phonemes=batch["phonemes"],
            speaker_ids=batch["speaker_ids"],
            durations=batch["durations"],
            pitch_targets=batch["pitch"],
            energy_targets=batch["energy"],
            mel_targets=batch["mels"],
            src_mask=batch["src_mask"],
            max_mel_len=batch["max_mel_len"],
            embed_predicted=EMBED_PREDICTED,
        )
        loss, *_ = compute_loss(mel_out, duration_preds, pitch_preds, energy_preds, batch)
        total_loss += loss.item()
        steps += 1
    model.train()
    return total_loss / max(steps, 1)


def train(data_dir, init_from=None):
    device = "cuda" if torch.cuda.is_available() else "cpu"
    use_amp = device == "cuda"

    with open(os.path.join(data_dir, "metadata.json")) as f:
        meta = json.load(f)

    config = EvoTalkConfig(
        phoneme_vocab_size=meta["phoneme_vocab_size"],
        n_speakers=meta["n_speakers"],
        n_mels=meta["n_mels"],
    )
    model = EvoTalk(config).to(device)

    if init_from is not None:
        ckpt = torch.load(init_from, map_location=device, weights_only=False)
        model.load_state_dict(ckpt["model"])
        print(f"warm-started weights from {init_from}")

    optimizer = model.configure_optimizers(WEIGHT_DECAY, LEARNING_RATE, BETAS, device)
    scaler = GradScaler("cuda", enabled=use_amp)

    stat_kwargs = dict(
        pitch_mean=meta["pitch_mean"], pitch_std=meta["pitch_std"],
        energy_mean=meta["energy_mean"], energy_std=meta["energy_std"],
    )
    train_dataset = LibriTTSDataset(data_dir, "train", **stat_kwargs)
    val_dataset = LibriTTSDataset(data_dir, "val", **stat_kwargs)
    print(f"train samples: {len(train_dataset)} | val samples: {len(val_dataset)}")

    train_loader = DataLoader(
        train_dataset,
        batch_size=BATCH_SIZE,
        shuffle=True,
        num_workers=NUM_WORKERS,
        collate_fn=collate_fn,
        pin_memory=True,
        drop_last=True,
    )
    val_loader = DataLoader(
        val_dataset,
        batch_size=BATCH_SIZE,
        shuffle=False,
        num_workers=NUM_WORKERS,
        collate_fn=collate_fn,
        pin_memory=True,
    )

    step = 0
    if RESUME_FROM is not None:
        step = load_checkpoint(RESUME_FROM, model, optimizer, scaler, device)

    history = {
        "steps": [], "loss": [], "mel": [], "dur": [],
        "pitch": [], "energy": [], "lr": [],
        "val_steps": [], "val_loss": [],
    }

    model.train()
    train_iter = iter(train_loader)
    t0 = time.time()

    while step < MAX_STEPS:
        try:
            batch = next(train_iter)
        except StopIteration:
            train_iter = iter(train_loader)
            batch = next(train_iter)

        lr = get_lr(step)
        for param_group in optimizer.param_groups:
            param_group["lr"] = lr

        batch = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in batch.items()}

        with autocast("cuda", enabled=use_amp):
            mel_out, duration_preds, pitch_preds, energy_preds, _ = model(
                phonemes=batch["phonemes"],
                speaker_ids=batch["speaker_ids"],
                durations=batch["durations"],
                pitch_targets=batch["pitch"],
                energy_targets=batch["energy"],
                mel_targets=batch["mels"],
                src_mask=batch["src_mask"],
                max_mel_len=batch["max_mel_len"],
                embed_predicted=EMBED_PREDICTED,
            )
            loss, mel_loss, dur_loss, pitch_loss, energy_loss = compute_loss(
                mel_out, duration_preds, pitch_preds, energy_preds, batch
            )

        optimizer.zero_grad(set_to_none=True)
        scaler.scale(loss).backward()
        scaler.unscale_(optimizer)
        torch.nn.utils.clip_grad_norm_(model.parameters(), GRAD_CLIP)
        scaler.step(optimizer)
        scaler.update()

        step += 1

        history["steps"].append(step)
        history["loss"].append(loss.item())
        history["mel"].append(mel_loss.item())
        history["dur"].append(dur_loss.item())
        history["pitch"].append(pitch_loss.item())
        history["energy"].append(energy_loss.item())
        history["lr"].append(lr)

        if step % LOG_EVERY == 0:
            t1 = time.time()
            dt = (t1 - t0) / LOG_EVERY
            t0 = t1
            print(
                f"step {step:>6} | loss {loss.item():.4f} | mel {mel_loss.item():.4f} "
                f"| dur {dur_loss.item():.4f} | pitch {pitch_loss.item():.4f} "
                f"| energy {energy_loss.item():.4f} | lr {lr:.2e} | {dt*1000:.0f}ms/step"
            )
            save_graphs(history, GRAPH_DIR)

        if step % EVAL_EVERY == 0 and len(val_dataset) > 0:
            val_loss = evaluate(model, val_loader, device, EVAL_STEPS)
            print(f"step {step:>6} | val loss {val_loss:.4f}")
            history["val_steps"].append(step)
            history["val_loss"].append(val_loss)
            save_graphs(history, GRAPH_DIR)

        if step % SAVE_EVERY == 0:
            ckpt_path = os.path.join(CHECKPOINT_DIR, f"evotalk_{step:06d}.pt")
            save_checkpoint(model, optimizer, scaler, step, loss.item(), ckpt_path)

    save_checkpoint(model, optimizer, scaler, step, loss.item(),
                    os.path.join(CHECKPOINT_DIR, "evotalk_final.pt"))
    print("training complete")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--data_dir", type=str, required=True)
    parser.add_argument("--init_from", type=str, default=None,
                        help="checkpoint to warm-start weights from (for fine-tuning)")
    parser.add_argument("--predicted_variance", action="store_true",
                        help="feed the decoder predicted (detached) pitch/energy instead of "
                             "ground-truth; removes the teacher-forcing leak")
    parser.add_argument("--lr", type=float, default=None, help="override learning rate")
    parser.add_argument("--max_steps", type=int, default=None, help="override max steps")
    parser.add_argument("--warmup", type=int, default=None, help="override warmup steps")
    args = parser.parse_args()

    if args.predicted_variance:
        EMBED_PREDICTED = True
    if args.lr is not None:
        LEARNING_RATE = args.lr
    if args.max_steps is not None:
        MAX_STEPS = args.max_steps
    if args.warmup is not None:
        WARMUP_STEPS = args.warmup

    train(args.data_dir, init_from=args.init_from)