#!/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()