scripts/train.py

#!/usr/bin/env python3
"""VAE decoder knowledge distillation v3: research-grounded two-phase training.

Phase 1 (steps 1-500K): Reconstruction convergence
  - L1 waveform loss
  - Log-magnitude multi-resolution STFT loss
  - Multi-scale mel spectrogram loss (7 scales, following DAC)
  - Feature-level distillation L1

Phase 2 (steps 500K-800K): Adversarial refinement
  - All Phase 1 losses
  - Multi-scale STFT discriminator (following EnCodec/DAC)
  - Feature matching loss from discriminator

Grounded in: DAC, EnCodec, APCodec, StreamCodec2, RAVE, Turbo-VAED research.

Usage on vast.ai:
    uv pip install torch diffusers transformers accelerate safetensors soundfile
    python distill_vae_decoder.py

Fully resumable from any checkpoint. Crashes if real audio unavailable.
"""

import logging
import math
import os
import sys
import time
from pathlib import Path

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.utils import weight_norm
from torch.nn.utils.parametrizations import weight_norm as weight_norm_v2

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
)
logger = logging.getLogger(__name__)

DEVICE = "cuda"
DTYPE = torch.float32

# ---------------------------------------------------------------------------
# Hyperparameters (grounded in published audio codec training)
# ---------------------------------------------------------------------------
TOTAL_STEPS = 800_000
PHASE2_START = 500_000          # adversarial kicks in here (RAVE/Turbo-VAED pattern)
BATCH_SIZE = 2
GRAD_ACCUM = 4                  # effective batch = 8 (close to DAC ablation batch 12)
CLIP_FRAMES = 100               # 4 seconds of latent (100 * 1920 / 48000 = 4s)
LATENT_CHANNELS = 64
LR = 3e-4                       # EnCodec uses 3e-4
LR_MIN = 1e-6
WEIGHT_DECAY = 1e-4
GRAD_CLIP = 1.0
LOG_EVERY = 100
SAVE_EVERY = 5000               # checkpoint every 5K steps, fully resumable

# Loss weights (grounded in DAC/EnCodec/APCodec)
W_L1 = 1.0                     # waveform L1
W_STFT = 1.0                   # multi-res STFT (spectral convergence + log mag)
W_MEL = 2.0                    # multi-scale mel (DAC weights mel at 15.0 but that's their primary)
W_FEAT = 0.1                   # feature distillation
W_ADV = 1.0                    # adversarial (phase 2, DAC uses 1.0)
W_FM = 2.0                     # feature matching from discriminator (DAC uses 2.0)

# STFT window sizes for multi-resolution loss
STFT_SIZES = [256, 512, 1024, 2048]

# Mel spectrogram scales (following DAC: 7 scales)
MEL_SIZES = [32, 64, 128, 256, 512, 1024, 2048]
MEL_BINS = 80
SAMPLE_RATE = 48000

# Discriminator LR (EnCodec/DAC use same or slightly different)
D_LR = 3e-4

# Latent dataset
LATENT_CLIP_SECONDS = 8         # longer clips for more diversity per track
MAX_CLIPS_PER_TRACK = 3         # multiple clips from long tracks

# Paths
OUTPUT_DIR = Path("./checkpoints/fast_decoder_v3")
ONNX_PATH = Path("./exports/vae_decode_fast_v3.onnx")

# Teacher config (ACE-Step VAE)
HF_REPO = "ACE-Step/Ace-Step1.5"
HF_SUBFOLDER = "vae"

# Teacher architecture
TEACHER_CHANNELS = 128
TEACHER_CHANNEL_MULTIPLES = [1, 2, 4, 8, 16]
TEACHER_DOWNSAMPLING_RATIOS = [2, 4, 4, 6, 10]

# Student architecture
STUDENT_CHANNELS = 128
STUDENT_CHANNEL_MULTIPLES = [1, 2, 4, 8, 8]
STUDENT_UPSAMPLING_RATIOS = [10, 6, 4, 4, 2]

# Audio directory on vast.ai
AUDIO_DIR = "/workspace/audio"
LATENT_CACHE_DIR = "/workspace/latent_cache"


# =========================================================================
# Student model definition (self-contained)
# =========================================================================

class Snake1d(nn.Module):
    """Snake activation from the DAC paper."""

    def __init__(self, hidden_dim, logscale=True):
        super().__init__()
        self.alpha = nn.Parameter(torch.zeros(1, hidden_dim, 1))
        self.beta = nn.Parameter(torch.zeros(1, hidden_dim, 1))
        self.alpha.requires_grad = True
        self.beta.requires_grad = True
        self.logscale = logscale

    def forward(self, hidden_states):
        shape = hidden_states.shape
        alpha = self.alpha if not self.logscale else torch.exp(self.alpha)
        beta = self.beta if not self.logscale else torch.exp(self.beta)
        hidden_states = hidden_states.reshape(shape[0], shape[1], -1)
        hidden_states = hidden_states + (beta + 1e-9).reciprocal() * torch.sin(alpha * hidden_states).pow(2)
        hidden_states = hidden_states.reshape(shape)
        return hidden_states


class FastResidualUnit(nn.Module):
    def __init__(self, dim: int, dilation: int = 1):
        super().__init__()
        pad = ((7 - 1) * dilation) // 2
        self.snake1 = Snake1d(dim)
        self.conv1 = weight_norm(
            nn.Conv1d(dim, dim, kernel_size=7, dilation=dilation, padding=pad)
        )
        self.snake2 = Snake1d(dim)
        self.conv2 = weight_norm(nn.Conv1d(dim, dim, kernel_size=1))

    def forward(self, x):
        h = self.conv1(self.snake1(x))
        h = self.conv2(self.snake2(h))
        pad = (x.shape[-1] - h.shape[-1]) // 2
        if pad > 0:
            x = x[..., pad:-pad]
        return x + h


class FastDecoderBlock(nn.Module):
    def __init__(self, in_dim: int, out_dim: int, stride: int = 1):
        super().__init__()
        self.snake1 = Snake1d(in_dim)
        self.conv_t = weight_norm(
            nn.ConvTranspose1d(
                in_dim, out_dim,
                kernel_size=2 * stride, stride=stride,
                padding=math.ceil(stride / 2),
            )
        )
        self.res1 = FastResidualUnit(out_dim, dilation=1)
        self.res2 = FastResidualUnit(out_dim, dilation=3)

    def forward(self, x):
        x = self.snake1(x)
        x = self.conv_t(x)
        x = self.res1(x)
        x = self.res2(x)
        return x


class FastOobleckDecoder(nn.Module):
    def __init__(
        self,
        channels: int = 128,
        input_channels: int = 64,
        audio_channels: int = 2,
        upsampling_ratios: list = None,
        channel_multiples: list = None,
    ):
        super().__init__()
        if upsampling_ratios is None:
            upsampling_ratios = [10, 6, 4, 4, 2]
        if channel_multiples is None:
            channel_multiples = [1, 2, 4, 8, 8]

        strides = upsampling_ratios
        cm = [1] + channel_multiples

        self.conv1 = weight_norm(
            nn.Conv1d(input_channels, channels * cm[-1], kernel_size=7, padding=3)
        )

        blocks = []
        for i, stride in enumerate(strides):
            in_dim = channels * cm[len(strides) - i]
            out_dim = channels * cm[len(strides) - i - 1]
            blocks.append(FastDecoderBlock(in_dim, out_dim, stride=stride))
        self.blocks = nn.ModuleList(blocks)

        self.final_snake = Snake1d(channels)
        self.conv2 = weight_norm(
            nn.Conv1d(channels, audio_channels, kernel_size=7, padding=3, bias=False)
        )

    def forward(self, latents: torch.Tensor) -> torch.Tensor:
        x = self.conv1(latents)
        for block in self.blocks:
            x = block(x)
        x = self.final_snake(x)
        x = self.conv2(x)
        return x

    def forward_with_features(self, latents: torch.Tensor):
        features = []
        x = self.conv1(latents)
        for block in self.blocks:
            x = block(x)
            features.append(x)
        x = self.final_snake(x)
        x = self.conv2(x)
        return x, features


# =========================================================================
# Teacher feature extraction wrapper
# =========================================================================

class TeacherWithFeatures(nn.Module):
    def __init__(self, teacher_decoder):
        super().__init__()
        self.teacher = teacher_decoder

    @torch.no_grad()
    def forward(self, hidden_state):
        features = []
        hidden_state = self.teacher.conv1(hidden_state)
        for layer in self.teacher.block:
            hidden_state = layer(hidden_state)
            features.append(hidden_state)
        hidden_state = self.teacher.snake1(hidden_state)
        hidden_state = self.teacher.conv2(hidden_state)
        return hidden_state, features


# =========================================================================
# Feature distillation projections
# =========================================================================

class FeatureProjectors(nn.Module):
    def __init__(self, student_dims, teacher_dims):
        super().__init__()
        projectors = []
        for s_dim, t_dim in zip(student_dims, teacher_dims):
            if s_dim != t_dim:
                projectors.append(nn.Conv1d(s_dim, t_dim, kernel_size=1))
            else:
                projectors.append(nn.Identity())
        self.projectors = nn.ModuleList(projectors)

    def forward(self, student_features, teacher_features):
        loss = torch.tensor(0.0, device=student_features[0].device)
        n = 0
        for proj, s_feat, t_feat in zip(self.projectors, student_features, teacher_features):
            s_proj = proj(s_feat)
            min_t = min(s_proj.shape[-1], t_feat.shape[-1])
            s_proj = s_proj[..., :min_t]
            t_feat = t_feat[..., :min_t]
            loss = loss + F.l1_loss(s_proj, t_feat.detach())
            n += 1
        return loss / max(n, 1)


# =========================================================================
# Multi-scale STFT discriminator (following EnCodec/DAC)
# =========================================================================

class STFTDiscriminatorBlock(nn.Module):
    """Single-scale complex STFT discriminator (EnCodec style)."""

    def __init__(self, n_fft: int, hop_length: int):
        super().__init__()
        self.n_fft = n_fft
        self.hop_length = hop_length
        # Input: real + imag = 2 channels, freq_bins = n_fft//2+1
        freq_bins = n_fft // 2 + 1

        self.layers = nn.ModuleList([
            nn.Sequential(
                nn.Conv2d(2, 32, kernel_size=(3, 9), padding=(1, 4)),
                nn.LeakyReLU(0.2),
            ),
            nn.Sequential(
                nn.Conv2d(32, 32, kernel_size=(3, 9), stride=(1, 2), padding=(1, 4)),
                nn.LeakyReLU(0.2),
            ),
            nn.Sequential(
                nn.Conv2d(32, 32, kernel_size=(3, 9), stride=(1, 2), padding=(1, 4)),
                nn.LeakyReLU(0.2),
            ),
            nn.Sequential(
                nn.Conv2d(32, 32, kernel_size=(3, 3), padding=(1, 1)),
                nn.LeakyReLU(0.2),
            ),
            nn.Conv2d(32, 1, kernel_size=(3, 3), padding=(1, 1)),
        ])

    def forward(self, x):
        """x: [B, C, T] audio. Returns (logits, features_list)."""
        # Mix to mono for discriminator
        if x.shape[1] > 1:
            x = x.mean(dim=1, keepdim=False)  # [B, T]
        else:
            x = x.squeeze(1)

        window = torch.hann_window(self.n_fft, device=x.device)
        stft = torch.stft(
            x, self.n_fft, self.hop_length, window=window,
            return_complex=True, normalized=True,
        )
        # stft: [B, freq, time] complex -> [B, 2, freq, time]
        x = torch.stack([stft.real, stft.imag], dim=1)

        features = []
        for layer in self.layers:
            x = layer(x)
            features.append(x)

        return x, features[:-1]  # logits, intermediate features


class MultiScaleSTFTDiscriminator(nn.Module):
    """Multi-scale STFT discriminator with 5 scales (following EnCodec)."""

    def __init__(self):
        super().__init__()
        # EnCodec uses windows: 2048, 1024, 512, 256, 128
        configs = [
            (2048, 512),
            (1024, 256),
            (512, 128),
            (256, 64),
            (128, 32),
        ]
        self.discriminators = nn.ModuleList([
            STFTDiscriminatorBlock(n_fft, hop) for n_fft, hop in configs
        ])

    def forward(self, x):
        """Returns list of (logits, features) per scale."""
        results = []
        for disc in self.discriminators:
            logits, feats = disc(x)
            results.append((logits, feats))
        return results


# =========================================================================
# Loss functions
# =========================================================================

def log_stft_loss(pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
    """Multi-resolution STFT loss: spectral convergence + log magnitude L1."""
    eps = 1e-5
    B, C, T = pred.shape
    pred_flat = pred.reshape(B * C, T)
    target_flat = target.reshape(B * C, T)

    loss = torch.tensor(0.0, device=pred.device)

    for n_fft in STFT_SIZES:
        hop = n_fft // 4
        window = torch.hann_window(n_fft, device=pred.device)

        pred_stft = torch.stft(
            pred_flat, n_fft, hop_length=hop, window=window,
            return_complex=True, normalized=True,
        )
        tgt_stft = torch.stft(
            target_flat, n_fft, hop_length=hop, window=window,
            return_complex=True, normalized=True,
        )

        pred_mag = pred_stft.abs()
        tgt_mag = tgt_stft.abs()

        sc = (tgt_mag - pred_mag).norm(p="fro") / (tgt_mag.norm(p="fro") + eps)
        loss = loss + sc

        log_mag_loss = F.l1_loss(
            torch.log(pred_mag + eps),
            torch.log(tgt_mag + eps),
        )
        loss = loss + log_mag_loss

    return loss / len(STFT_SIZES)


def multi_scale_mel_loss(pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
    """Multi-scale mel spectrogram loss (7 scales, following DAC)."""
    eps = 1e-5
    B, C, T = pred.shape
    pred_flat = pred.reshape(B * C, T)
    target_flat = target.reshape(B * C, T)

    loss = torch.tensor(0.0, device=pred.device)

    for n_fft in MEL_SIZES:
        hop = n_fft // 4
        window = torch.hann_window(n_fft, device=pred.device)
        n_mels = min(MEL_BINS, n_fft // 2)

        # Create mel filterbank
        mel_fb = _mel_filterbank(n_fft, n_mels, SAMPLE_RATE, pred.device)

        pred_stft = torch.stft(
            pred_flat, n_fft, hop_length=hop, window=window,
            return_complex=True, normalized=True,
        )
        tgt_stft = torch.stft(
            target_flat, n_fft, hop_length=hop, window=window,
            return_complex=True, normalized=True,
        )

        pred_mel = torch.matmul(mel_fb, pred_stft.abs().pow(2)).clamp(min=eps).log()
        tgt_mel = torch.matmul(mel_fb, tgt_stft.abs().pow(2)).clamp(min=eps).log()

        loss = loss + F.l1_loss(pred_mel, tgt_mel)

    return loss / len(MEL_SIZES)


def _mel_filterbank(n_fft: int, n_mels: int, sr: int, device) -> torch.Tensor:
    """Create a mel filterbank matrix [n_mels, n_fft//2+1]."""
    f_min, f_max = 0.0, sr / 2.0
    freq_bins = n_fft // 2 + 1

    def hz_to_mel(f):
        return 2595.0 * math.log10(1.0 + f / 700.0)

    def mel_to_hz(m):
        return 700.0 * (10.0 ** (m / 2595.0) - 1.0)

    mel_min = hz_to_mel(f_min)
    mel_max = hz_to_mel(f_max)
    mel_points = torch.linspace(mel_min, mel_max, n_mels + 2, device=device)
    hz_points = mel_to_hz(mel_points)
    bin_points = (hz_points * n_fft / sr).long().clamp(0, freq_bins - 1)

    fb = torch.zeros(n_mels, freq_bins, device=device)
    for i in range(n_mels):
        left, center, right = bin_points[i], bin_points[i + 1], bin_points[i + 2]
        if center > left:
            fb[i, left:center] = torch.linspace(0, 1, center - left, device=device)
        if right > center:
            fb[i, center:right] = torch.linspace(1, 0, right - center, device=device)

    return fb


def adversarial_g_loss(disc_outputs):
    """Hinge generator loss across all discriminator scales."""
    loss = torch.tensor(0.0, device=disc_outputs[0][0].device)
    for logits, _ in disc_outputs:
        loss = loss + torch.mean(F.relu(1.0 - logits))
    return loss / len(disc_outputs)


def adversarial_d_loss(real_outputs, fake_outputs):
    """Hinge discriminator loss across all scales."""
    loss = torch.tensor(0.0, device=real_outputs[0][0].device)
    for (real_logits, _), (fake_logits, _) in zip(real_outputs, fake_outputs):
        loss = loss + torch.mean(F.relu(1.0 - real_logits))
        loss = loss + torch.mean(F.relu(1.0 + fake_logits))
    return loss / len(real_outputs)


def feature_matching_loss(real_outputs, fake_outputs):
    """L1 feature matching across discriminator layers (DAC/EnCodec)."""
    loss = torch.tensor(0.0, device=real_outputs[0][0].device)
    n = 0
    for (_, real_feats), (_, fake_feats) in zip(real_outputs, fake_outputs):
        for rf, ff in zip(real_feats, fake_feats):
            loss = loss + F.l1_loss(ff, rf.detach())
            n += 1
    return loss / max(n, 1)


# =========================================================================
# Latent dataset generation
# =========================================================================

def generate_latent_dataset(vae, audio_dir: str, cache_dir: str):
    """Encode ALL available audio through the VAE encoder.

    Encodes every track, taking multiple clips from long tracks.
    Caches to disk so re-runs skip encoding. Crashes if no audio found.
    """
    import glob
    import subprocess
    import tempfile
    import random

    os.makedirs(cache_dir, exist_ok=True)
    cache_file = os.path.join(cache_dir, "latents_all.pt")

    # Check cache first
    if os.path.exists(cache_file):
        logger.info("Loading cached latents from %s", cache_file)
        data = torch.load(cache_file, map_location="cpu", weights_only=True)
        logger.info("Loaded %d cached latents", len(data))
        return data

    # Must have soundfile for reading wav
    import soundfile as sf

    mp3_files = sorted(glob.glob(os.path.join(audio_dir, "**", "*.mp3"), recursive=True))
    if not mp3_files:
        raise RuntimeError(
            f"No MP3 files found in {audio_dir}. "
            f"Download FMA dataset first: wget https://os.unil.cloud.switch.ch/fma/fma_small.zip"
        )

    logger.info("Found %d MP3 files in %s", len(mp3_files), audio_dir)
    random.shuffle(mp3_files)

    target_sr = SAMPLE_RATE
    target_samples = int(LATENT_CLIP_SECONDS * target_sr)
    latents = []
    errors = 0
    tracks_used = 0

    vae.eval()
    with torch.no_grad():
        for mp3_path in mp3_files:
            try:
                # Decode MP3 to 48kHz stereo WAV via ffmpeg
                tmp = tempfile.mktemp(suffix=".wav")
                result = subprocess.run(
                    ["ffmpeg", "-y", "-i", mp3_path,
                     "-ar", str(target_sr), "-ac", "2", "-f", "wav", tmp],
                    capture_output=True, timeout=30,
                )
                if result.returncode != 0:
                    errors += 1
                    continue

                data, sr = sf.read(tmp, dtype="float32")
                os.unlink(tmp)

                waveform = torch.tensor(data, dtype=torch.float32).T  # [2, samples]

                if waveform.shape[-1] < target_samples:
                    # Still use short clips, just pad
                    if waveform.shape[-1] < target_sr:  # skip < 1 second
                        continue
                    waveform = F.pad(waveform, (0, target_samples - waveform.shape[-1]))
                    clips_from_track = 1
                else:
                    clips_from_track = min(
                        MAX_CLIPS_PER_TRACK,
                        waveform.shape[-1] // target_samples
                    )

                tracks_used += 1

                for clip_idx in range(clips_from_track):
                    if clips_from_track == 1 and waveform.shape[-1] >= target_samples:
                        start = torch.randint(0, waveform.shape[-1] - target_samples, (1,)).item()
                    else:
                        start = clip_idx * (waveform.shape[-1] - target_samples) // max(clips_from_track - 1, 1)

                    clip = waveform[:, start:start + target_samples]

                    # Normalize to [-1, 1]
                    peak = clip.abs().max()
                    if peak > 1e-6:
                        clip = clip / peak

                    # Encode through VAE
                    clip_gpu = clip.unsqueeze(0).to(DEVICE, dtype=DTYPE)
                    enc_out = vae.encode(clip_gpu)
                    latent = enc_out.latent_dist.sample()
                    latents.append(latent.cpu())

            except Exception:
                errors += 1
                continue

            if tracks_used % 200 == 0:
                logger.info("  Encoded %d tracks -> %d latent clips (%d errors)",
                            tracks_used, len(latents), errors)

    if len(latents) < 100:
        raise RuntimeError(
            f"Only encoded {len(latents)} latents from {tracks_used} tracks (need >= 100). "
            f"{errors} files failed. Check ffmpeg and audio files."
        )

    logger.info("Encoded %d latent clips from %d tracks (%d errors)",
                len(latents), tracks_used, errors)
    sample = latents[0]
    logger.info("  Latent shape: %s, mean=%.3f, std=%.3f",
                list(sample.shape), sample.mean().item(), sample.std().item())

    # Cache to disk
    torch.save(latents, cache_file)
    logger.info("Cached latents to %s", cache_file)

    return latents


def sample_from_dataset(latent_dataset, batch_size, clip_frames):
    """Sample a batch of random clips from the pre-generated latent dataset."""
    batch = []
    for _ in range(batch_size):
        idx = torch.randint(0, len(latent_dataset), (1,)).item()
        lat = latent_dataset[idx]  # [1, 64, T]
        T = lat.shape[-1]
        if T > clip_frames:
            start = torch.randint(0, T - clip_frames, (1,)).item()
            lat = lat[:, :, start:start + clip_frames]
        elif T < clip_frames:
            lat = F.pad(lat, (0, clip_frames - T))
        batch.append(lat)
    return torch.cat(batch, dim=0).to(DEVICE, dtype=DTYPE)


# =========================================================================
# Utilities
# =========================================================================

def remove_weight_norm_recursive(module):
    for name, child in module.named_children():
        try:
            torch.nn.utils.remove_weight_norm(child)
        except ValueError:
            pass
        remove_weight_norm_recursive(child)


def get_block_output_dims(channels, channel_multiples, upsampling_ratios):
    cm = [1] + channel_multiples
    strides = upsampling_ratios
    dims = []
    for i in range(len(strides)):
        out_dim = channels * cm[len(strides) - i - 1]
        dims.append(out_dim)
    return dims


# =========================================================================
# Checkpoint save/load (fully resumable)
# =========================================================================

def save_checkpoint(path, step, student, feat_projectors, optimizer_g, scheduler_g,
                    discriminator=None, optimizer_d=None, scheduler_d=None):
    """Save a fully resumable checkpoint."""
    ckpt = {
        "step": step,
        "student_state_dict": student.state_dict(),
        "feat_proj_state_dict": feat_projectors.state_dict(),
        "optimizer_g_state_dict": optimizer_g.state_dict(),
        "scheduler_g_state_dict": scheduler_g.state_dict(),
    }
    if discriminator is not None:
        ckpt["discriminator_state_dict"] = discriminator.state_dict()
    if optimizer_d is not None:
        ckpt["optimizer_d_state_dict"] = optimizer_d.state_dict()
    if scheduler_d is not None:
        ckpt["scheduler_d_state_dict"] = scheduler_d.state_dict()

    torch.save(ckpt, path)
    logger.info("Checkpoint saved: %s (step %d)", path, step)


def find_latest_checkpoint(output_dir):
    """Find the latest checkpoint in the output directory."""
    ckpts = sorted(output_dir.glob("student_step*.pt"))
    if not ckpts:
        return None
    # Sort by step number
    def step_from_path(p):
        name = p.stem  # e.g. "student_step5000"
        return int(name.replace("student_step", ""))
    ckpts.sort(key=step_from_path)
    return ckpts[-1]


# =========================================================================
# Main training loop
# =========================================================================

def main():
    os.makedirs(OUTPUT_DIR, exist_ok=True)

    logger.info("=" * 60)
    logger.info("VAE Decoder Knowledge Distillation v3")
    logger.info("=" * 60)
    logger.info("Device: %s", DEVICE)
    logger.info("Total steps: %d (phase 2 at %d)", TOTAL_STEPS, PHASE2_START)
    logger.info("Batch: %d x %d accum = %d effective", BATCH_SIZE, GRAD_ACCUM, BATCH_SIZE * GRAD_ACCUM)
    logger.info("Clip frames: %d (%.1fs)", CLIP_FRAMES, CLIP_FRAMES * 1920 / SAMPLE_RATE)
    logger.info("LR: %s -> %s", LR, LR_MIN)
    logger.info("Loss weights: L1=%.1f STFT=%.1f Mel=%.1f Feat=%.2f Adv=%.1f FM=%.1f",
                W_L1, W_STFT, W_MEL, W_FEAT, W_ADV, W_FM)

    # ==================================================================
    # 1. Load teacher VAE
    # ==================================================================
    logger.info("Loading teacher VAE from %s...", HF_REPO)
    from diffusers import AutoencoderOobleck

    vae = AutoencoderOobleck.from_pretrained(HF_REPO, subfolder=HF_SUBFOLDER)
    vae = vae.to(DEVICE, dtype=DTYPE)
    vae.eval()

    logger.info("Teacher loaded. hop_length=%d", vae.hop_length)

    # ==================================================================
    # 2. Generate real latent dataset from ALL audio
    # ==================================================================
    latent_dataset = generate_latent_dataset(vae, AUDIO_DIR, LATENT_CACHE_DIR)
    logger.info("Dataset: %d latent clips", len(latent_dataset))

    # ==================================================================
    # 3. Set up teacher decoder with feature extraction
    # ==================================================================
    teacher = TeacherWithFeatures(vae.decoder).eval().to(DEVICE)
    for p in teacher.parameters():
        p.requires_grad_(False)

    del vae
    torch.cuda.empty_cache()

    teacher_params = sum(p.numel() for p in teacher.parameters())
    logger.info("Teacher decoder: %.2fM params", teacher_params / 1e6)

    # ==================================================================
    # 4. Create student
    # ==================================================================
    student = FastOobleckDecoder(
        channels=STUDENT_CHANNELS,
        input_channels=LATENT_CHANNELS,
        audio_channels=2,
        upsampling_ratios=STUDENT_UPSAMPLING_RATIOS,
        channel_multiples=STUDENT_CHANNEL_MULTIPLES,
    ).to(DEVICE, dtype=DTYPE)

    student_params = sum(p.numel() for p in student.parameters())
    logger.info("Student decoder: %.2fM params (%.0f%% of teacher)",
                student_params / 1e6, 100 * student_params / teacher_params)

    # ==================================================================
    # 5. Feature projectors
    # ==================================================================
    teacher_block_dims = get_block_output_dims(
        TEACHER_CHANNELS, TEACHER_CHANNEL_MULTIPLES, STUDENT_UPSAMPLING_RATIOS
    )
    student_block_dims = get_block_output_dims(
        STUDENT_CHANNELS, STUDENT_CHANNEL_MULTIPLES, STUDENT_UPSAMPLING_RATIOS
    )

    for i, (s, t) in enumerate(zip(student_block_dims, teacher_block_dims)):
        logger.info("  Block %d: student=%d, teacher=%d %s",
                     i, s, t, "(proj)" if s != t else "")

    feat_projectors = FeatureProjectors(student_block_dims, teacher_block_dims).to(DEVICE, dtype=DTYPE)

    # ==================================================================
    # 6. Discriminator (created now, used in phase 2)
    # ==================================================================
    discriminator = MultiScaleSTFTDiscriminator().to(DEVICE, dtype=DTYPE)
    d_params = sum(p.numel() for p in discriminator.parameters())
    logger.info("Discriminator: %.2fM params", d_params / 1e6)

    # ==================================================================
    # 7. Optimizers and schedulers
    # ==================================================================
    g_params = list(student.parameters()) + list(feat_projectors.parameters())
    optimizer_g = torch.optim.AdamW(g_params, lr=LR, weight_decay=WEIGHT_DECAY, betas=(0.8, 0.99))
    scheduler_g = torch.optim.lr_scheduler.CosineAnnealingLR(
        optimizer_g, T_max=TOTAL_STEPS, eta_min=LR_MIN
    )

    optimizer_d = torch.optim.AdamW(discriminator.parameters(), lr=D_LR, weight_decay=WEIGHT_DECAY, betas=(0.8, 0.99))
    scheduler_d = torch.optim.lr_scheduler.CosineAnnealingLR(
        optimizer_d, T_max=TOTAL_STEPS - PHASE2_START, eta_min=LR_MIN
    )

    # ==================================================================
    # 8. Resume from checkpoint if available
    # ==================================================================
    start_step = 0
    latest_ckpt = find_latest_checkpoint(OUTPUT_DIR)
    if latest_ckpt is not None:
        logger.info("Resuming from checkpoint: %s", latest_ckpt)
        ckpt = torch.load(latest_ckpt, map_location=DEVICE, weights_only=False)
        start_step = ckpt["step"]
        student.load_state_dict(ckpt["student_state_dict"])
        feat_projectors.load_state_dict(ckpt["feat_proj_state_dict"])
        optimizer_g.load_state_dict(ckpt["optimizer_g_state_dict"])
        scheduler_g.load_state_dict(ckpt["scheduler_g_state_dict"])
        if "discriminator_state_dict" in ckpt:
            discriminator.load_state_dict(ckpt["discriminator_state_dict"])
        if "optimizer_d_state_dict" in ckpt:
            optimizer_d.load_state_dict(ckpt["optimizer_d_state_dict"])
        if "scheduler_d_state_dict" in ckpt:
            scheduler_d.load_state_dict(ckpt["scheduler_d_state_dict"])
        logger.info("Resumed from step %d", start_step)
        del ckpt
        torch.cuda.empty_cache()

    # ==================================================================
    # 9. Training loop
    # ==================================================================
    logger.info("Training from step %d to %d...", start_step + 1, TOTAL_STEPS)
    student.train()
    feat_projectors.train()

    running = {"l1": 0, "stft": 0, "mel": 0, "feat": 0, "adv_g": 0, "fm": 0, "adv_d": 0, "total": 0}
    t_start = time.time()

    optimizer_g.zero_grad(set_to_none=True)
    optimizer_d.zero_grad(set_to_none=True)

    for step in range(start_step + 1, TOTAL_STEPS + 1):
        in_phase2 = step >= PHASE2_START

        # --- Sample latent batch ---
        latents = sample_from_dataset(latent_dataset, BATCH_SIZE, CLIP_FRAMES)

        # --- Teacher forward (no grad) ---
        with torch.no_grad():
            teacher_audio, teacher_feats = teacher(latents)

        # --- Student forward ---
        student_audio, student_feats = student.forward_with_features(latents)

        # --- Trim to matching lengths ---
        min_len = min(student_audio.shape[-1], teacher_audio.shape[-1])
        student_audio_trimmed = student_audio[..., :min_len]
        teacher_audio_trimmed = teacher_audio[..., :min_len]

        # --- Phase 1 losses (always active) ---
        l1_loss = F.l1_loss(student_audio_trimmed, teacher_audio_trimmed)
        stft_loss = log_stft_loss(student_audio_trimmed, teacher_audio_trimmed)
        mel_loss = multi_scale_mel_loss(student_audio_trimmed, teacher_audio_trimmed)
        feat_loss = feat_projectors(student_feats, teacher_feats)

        g_loss = (W_L1 * l1_loss + W_STFT * stft_loss + W_MEL * mel_loss + W_FEAT * feat_loss)

        # --- Phase 2: adversarial losses ---
        adv_g_loss_val = torch.tensor(0.0, device=DEVICE)
        fm_loss_val = torch.tensor(0.0, device=DEVICE)
        d_loss_val = torch.tensor(0.0, device=DEVICE)

        if in_phase2:
            discriminator.train()

            # Discriminator step: detach student output
            with torch.no_grad():
                fake_audio_d = student_audio_trimmed.detach()
            real_out = discriminator(teacher_audio_trimmed.detach())
            fake_out = discriminator(fake_audio_d)
            d_loss_val = adversarial_d_loss(real_out, fake_out)

            # Scale by grad accum
            (d_loss_val / GRAD_ACCUM).backward()

            # Generator adversarial + feature matching
            fake_out_g = discriminator(student_audio_trimmed)
            real_out_g = discriminator(teacher_audio_trimmed.detach())
            adv_g_loss_val = adversarial_g_loss(fake_out_g)
            fm_loss_val = feature_matching_loss(real_out_g, fake_out_g)

            g_loss = g_loss + W_ADV * adv_g_loss_val + W_FM * fm_loss_val

        # --- Generator backward ---
        (g_loss / GRAD_ACCUM).backward()

        # --- Accumulate stats ---
        running["l1"] += l1_loss.item()
        running["stft"] += stft_loss.item()
        running["mel"] += mel_loss.item()
        running["feat"] += feat_loss.item()
        running["adv_g"] += adv_g_loss_val.item()
        running["fm"] += fm_loss_val.item()
        running["adv_d"] += d_loss_val.item()
        running["total"] += g_loss.item()

        # --- Optimizer step (with gradient accumulation) ---
        if step % GRAD_ACCUM == 0:
            torch.nn.utils.clip_grad_norm_(g_params, GRAD_CLIP)
            optimizer_g.step()
            optimizer_g.zero_grad(set_to_none=True)

            if in_phase2:
                torch.nn.utils.clip_grad_norm_(discriminator.parameters(), GRAD_CLIP)
                optimizer_d.step()
                optimizer_d.zero_grad(set_to_none=True)
                scheduler_d.step()

            scheduler_g.step()

        # --- Logging ---
        if step % LOG_EVERY == 0:
            n = LOG_EVERY
            elapsed = time.time() - t_start
            steps_done = step - start_step
            sps = steps_done / elapsed
            eta_min = (TOTAL_STEPS - step) / sps / 60
            lr_now = scheduler_g.get_last_lr()[0]
            phase = "P2" if in_phase2 else "P1"

            logger.info(
                "[%s] step %6d/%d  total=%.4f  l1=%.4f  stft=%.4f  mel=%.4f  "
                "feat=%.4f  adv_g=%.4f  fm=%.4f  d=%.4f  lr=%.1e  %.1f it/s  ETA %.0fm",
                phase, step, TOTAL_STEPS,
                running["total"] / n, running["l1"] / n, running["stft"] / n,
                running["mel"] / n, running["feat"] / n, running["adv_g"] / n,
                running["fm"] / n, running["adv_d"] / n,
                lr_now, sps, eta_min,
            )
            running = {k: 0.0 for k in running}

        # --- Save checkpoint ---
        if step % SAVE_EVERY == 0:
            ckpt_path = OUTPUT_DIR / f"student_step{step}.pt"
            save_checkpoint(
                ckpt_path, step, student, feat_projectors, optimizer_g, scheduler_g,
                discriminator if in_phase2 else None,
                optimizer_d if in_phase2 else None,
                scheduler_d if in_phase2 else None,
            )

    # ==================================================================
    # 10. Final save
    # ==================================================================
    total_time = time.time() - t_start
    logger.info("Training complete in %.1f hours", total_time / 3600)

    final_path = OUTPUT_DIR / "student_final.pt"
    torch.save({
        "step": TOTAL_STEPS,
        "student_state_dict": student.state_dict(),
        "config": {
            "channels": STUDENT_CHANNELS,
            "input_channels": LATENT_CHANNELS,
            "audio_channels": 2,
            "upsampling_ratios": STUDENT_UPSAMPLING_RATIOS,
            "channel_multiples": STUDENT_CHANNEL_MULTIPLES,
        },
    }, final_path)
    logger.info("Final model saved: %s", final_path)

    del teacher, feat_projectors, discriminator, latent_dataset
    torch.cuda.empty_cache()

    # ==================================================================
    # 11. ONNX export
    # ==================================================================
    logger.info("Preparing ONNX export...")
    student.eval()
    remove_weight_norm_recursive(student)

    test_input = torch.randn(1, LATENT_CHANNELS, 150, device=DEVICE, dtype=DTYPE)
    with torch.no_grad():
        test_output = student(test_input)
    logger.info("Post weight_norm-removal: %s -> %s",
                list(test_input.shape), list(test_output.shape))

    example = torch.randn(1, LATENT_CHANNELS, 750, device=DEVICE, dtype=DTYPE)
    os.makedirs(ONNX_PATH.parent, exist_ok=True)
    with torch.no_grad():
        torch.onnx.export(
            student, (example,), str(ONNX_PATH),
            input_names=["latents"], output_names=["audio"],
            dynamic_axes={
                "latents": {0: "batch", 2: "latent_frames"},
                "audio": {0: "batch", 2: "samples"},
            },
            opset_version=18, do_constant_folding=True,
        )
    logger.info("ONNX saved: %s (%.1f MB)", ONNX_PATH, ONNX_PATH.stat().st_size / 1e6)

    logger.info("=" * 60)
    logger.info("DONE. Checkpoint: %s | ONNX: %s", final_path, ONNX_PATH)
    logger.info("=" * 60)


if __name__ == "__main__":
    main()