Delete train_boson.py

train_boson.py

@@ -1,891 +0,0 @@
-#!/usr/bin/env python3
-"""
-Training script for Boson Audio Codec with DAC-inspired losses
-"""
-
-import os
-import json
-import argparse
-import random
-from pathlib import Path
-from datetime import datetime
-import numpy as np
-import pandas as pd
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import torch.distributed as dist
-from torch.nn.parallel import DistributedDataParallel as DDP
-from torch.utils.data import Dataset, DataLoader
-from torch.utils.data.distributed import DistributedSampler
-from torch.utils.tensorboard import SummaryWriter
-import torchaudio
-import librosa
-from tqdm import tqdm
-from audiotools import AudioSignal, STFTParams
-
-# Import from the provided codebase
-from higgs_audio_tokenizer import HiggsAudioTokenizer
-from quantization.distrib import broadcast_tensors, sync_buffer, is_distributed, world_size, rank
-from quantization.ddp_utils import set_random_seed, is_logging_process, get_timestamp
-
-# Import DAC losses and discriminator
-import sys
-sys.path.append('.')  # Add current directory to path
-from loss import L1Loss, MultiScaleSTFTLoss, MelSpectrogramLoss, GANLoss
-from discriminator import Discriminator
-
-
-class CosineWarmupScheduler(torch.optim.lr_scheduler._LRScheduler):
-    """Cosine scheduler with linear warmup"""
-    def __init__(self, optimizer, warmup_steps, total_steps, eta_min=1e-6, last_epoch=-1):
-        self.warmup_steps = warmup_steps
-        self.total_steps = total_steps
-        self.eta_min = eta_min
-        super().__init__(optimizer, last_epoch)
-    
-    def get_lr(self):
-        if self.last_epoch < self.warmup_steps:
-            # Linear warmup
-            warmup_factor = self.last_epoch / self.warmup_steps
-            return [base_lr * warmup_factor for base_lr in self.base_lrs]
-        else:
-            # Cosine annealing
-            progress = (self.last_epoch - self.warmup_steps) / (self.total_steps - self.warmup_steps)
-            cosine_factor = 0.5 * (1 + np.cos(np.pi * progress))
-            return [self.eta_min + (base_lr - self.eta_min) * cosine_factor for base_lr in self.base_lrs]
-
-
-class AudioDataset(Dataset):
-    """Dataset for loading audio files from CSV"""
-    def __init__(self, csv_path, sample_rate=44100, segment_duration=2.0, is_train=True):
-        self.df = pd.read_csv(csv_path)
-        self.sample_rate = sample_rate
-        self.segment_duration = segment_duration
-        self.segment_length = int(sample_rate * segment_duration)
-        self.is_train = is_train
-        
-        # Filter out files that don't exist
-        valid_files = []
-        for idx, row in self.df.iterrows():
-            if os.path.exists(row.iloc[0]):
-                valid_files.append(row.iloc[0])
-        self.audio_paths = valid_files
-        print(f"Found {len(self.audio_paths)} valid audio files")
-    
-    def __len__(self):
-        return len(self.audio_paths)
-    
-    def __getitem__(self, idx):
-        audio_path = self.audio_paths[idx]
-        
-        try:
-            # Load audio using librosa
-            audio, sr = librosa.load(audio_path, sr=self.sample_rate, mono=True)
-            
-            # Random segment extraction for training
-            if len(audio) > self.segment_length:
-                if self.is_train:
-                    start = random.randint(0, len(audio) - self.segment_length)
-                else:
-                    start = 0  # Always use beginning for validation
-                audio = audio[start:start + self.segment_length]
-            else:
-                # Pad if too short
-                audio = np.pad(audio, (0, self.segment_length - len(audio)))
-            
-            # Convert to tensor and add batch dimension
-            audio_tensor = torch.FloatTensor(audio).unsqueeze(0)
-            
-            return audio_tensor, audio_path
-            
-        except Exception as e:
-            print(f"Error loading {audio_path}: {e}")
-            # Return silence if loading fails
-            return torch.zeros(1, self.segment_length), audio_path
-
-
-class BosonTrainer:
-    def __init__(self, args):
-        self.args = args
-        self.distributed = False
-        
-        # Check if we're in a distributed environment
-        if 'WORLD_SIZE' in os.environ and int(os.environ['WORLD_SIZE']) > 1:
-            self.distributed = True
-            self.setup_ddp()
-            self.device = torch.device(f'cuda:{args.local_rank}')
-        else:
-            # Single GPU mode
-            self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-            torch.cuda.set_device(0)
-            set_random_seed(args.seed)
-        
-        # Load config
-        with open(args.config, 'r') as f:
-            self.config = json.load(f)
-        
-        # Initialize models
-        self.model = self.build_model()
-        self.discriminator = self.build_discriminator() if args.use_discriminator else None
-        
-        # Setup data loaders
-        self.train_loader, self.val_loader = self.setup_data_loaders()
-        
-        # Setup optimizers
-        self.optimizer_g = torch.optim.AdamW(
-            self.model.parameters(),
-            lr=args.learning_rate,
-            betas=(0.5, 0.9),
-            weight_decay=args.weight_decay
-        )
-        
-        if self.discriminator is not None:
-            self.optimizer_d = torch.optim.AdamW(
-                self.discriminator.parameters(),
-                lr=args.learning_rate * 2,  # Typically discriminator learns faster
-                betas=(0.5, 0.9),
-                weight_decay=args.weight_decay
-            )
-        
-        # Calculate total training steps
-        self.total_steps = args.num_epochs * len(self.train_loader)
-        
-        # Setup schedulers with warmup
-        self.scheduler_g = CosineWarmupScheduler(
-            self.optimizer_g,
-            warmup_steps=args.warmup_steps,
-            total_steps=self.total_steps,
-            eta_min=1e-6
-        )
-        
-        if self.discriminator is not None:
-            self.scheduler_d = CosineWarmupScheduler(
-                self.optimizer_d,
-                warmup_steps=args.warmup_steps,
-                total_steps=self.total_steps,
-                eta_min=1e-6
-            )
-        
-        # Setup losses
-        self.setup_losses()
-        
-        # Setup tensorboard
-        if not self.distributed or rank() == 0:
-            self.writer = SummaryWriter(
-                log_dir=os.path.join(args.output_dir, 'logs', get_timestamp())
-            )
-        
-        self.global_step = 0
-        self.start_epoch = 0
-        
-        # Load checkpoint if exists
-        if args.resume:
-            self.load_checkpoint()
-    
-    def setup_ddp(self):
-        """Initialize DDP"""
-        if 'LOCAL_RANK' in os.environ:
-            self.args.local_rank = int(os.environ['LOCAL_RANK'])
-        dist.init_process_group(backend='nccl')
-        torch.cuda.set_device(self.args.local_rank)
-        set_random_seed(self.args.seed + rank())
-    
-    def build_model(self):
-        """Build and wrap model with DDP if needed"""
-        
-        print(self.config)
-        model = HiggsAudioTokenizer(
-            n_filters=self.config['n_filters'],
-            D=self.config['D'],
-            target_bandwidths=self.config['target_bandwidths'],
-            ratios=self.config['ratios'],
-            sample_rate=self.config['sample_rate'],
-            bins=self.config['bins'],
-            n_q=self.config['n_q'],
-            codebook_dim=self.config.get('codebook_dim', None),
-            semantic_techer=self.config['semantic_techer'],
-            device=self.device
-        ).to(self.device)
-        
-        if self.distributed:
-            # Broadcast model parameters to ensure all ranks have same initialization
-            broadcast_tensors(model.parameters())
-            # Wrap with DDP
-            model = DDP(model, device_ids=[self.args.local_rank])
-        
-        return model
-    
-    def build_discriminator(self):
-        """Build discriminator with DDP if needed"""
-        # Use sample rate from config
-        discriminator = Discriminator(
-            rates=[],  # No multi-rate discriminator for now
-            periods=[2, 3, 5, 7, 11],
-            fft_sizes=[2048, 1024, 512],
-            sample_rate=self.config['sample_rate'],
-        ).to(self.device)
-        
-        if self.distributed:
-            broadcast_tensors(discriminator.parameters())
-            discriminator = DDP(discriminator, device_ids=[self.args.local_rank])
-        
-        return discriminator
-    
-    def setup_losses(self):
-        """Setup all loss functions"""
-        # Basic losses
-        self.l1_loss = L1Loss()
-        self.stft_loss = MultiScaleSTFTLoss(
-            window_lengths=[2048, 1024, 512, 256, 128],
-            loss_fn=nn.L1Loss(),
-            clamp_eps=1e-5,
-            mag_weight=1.0,
-            log_weight=1.0,
-        )
-        self.mel_loss = MelSpectrogramLoss(
-            n_mels=[150, 80],
-            window_lengths=[2048, 512],
-            mel_fmin=[0.0, 0.0],
-            mel_fmax=[None, None],
-            clamp_eps=1e-5,
-            mag_weight=1.0,
-            log_weight=1.0,
-        )
-        
-        # GAN loss if using discriminator
-        if self.discriminator is not None:
-            self.gan_loss = GANLoss(self.discriminator)
-        
-        # Loss weights (matching DAC's proven configuration)
-        self.loss_weights = {
-            'rec': 1.,  # Waveform L1 loss
-            'stft': 1.,  # Multi-scale STFT loss
-            # 'mel': 15.0,  # Mel-spectrogram loss (ENABLE it after 20-25k steps)
-            'mel': 0.0,  # Mel-spectrogram loss (DISABLED)
-            'commit': 0.25,  # Commitment loss
-            'semantic': 1.,  # Semantic loss
-            'gen': 1.,  # Generator adversarial loss
-            'feat': 1.0,  # Feature matching loss
-        }
-    
-    def setup_data_loaders(self):
-        """Setup data loaders (distributed or single GPU)"""
-        # Split data into train/val
-        df = pd.read_csv(self.args.data_csv)
-        n_total = len(df)
-        n_train = int(n_total * 0.9)
-        
-        # Create temporary CSV files for train/val split
-        train_csv = '/tmp/train_audio.csv'
-        val_csv = '/tmp/val_audio.csv'
-        
-        if not self.distributed or rank() == 0:
-            df[:n_train].to_csv(train_csv, index=False)
-            df[n_train:].to_csv(val_csv, index=False)
-        
-        # Synchronize across processes if distributed
-        if self.distributed:
-            dist.barrier()
-        
-        # Create datasets
-        train_dataset = AudioDataset(
-            train_csv,
-            sample_rate=self.config['sample_rate'],
-            segment_duration=self.args.segment_duration,
-            is_train=True
-        )
-        
-        val_dataset = AudioDataset(
-            val_csv,
-            sample_rate=self.config['sample_rate'],
-            segment_duration=self.args.segment_duration,
-            is_train=False
-        )
-        
-        # Create samplers and loaders
-        if self.distributed:
-            train_sampler = DistributedSampler(train_dataset, shuffle=True)
-            val_sampler = DistributedSampler(val_dataset, shuffle=False)
-        else:
-            train_sampler = None
-            val_sampler = None
-        
-        train_loader = DataLoader(
-            train_dataset,
-            batch_size=self.args.batch_size,
-            sampler=train_sampler,
-            shuffle=(train_sampler is None),
-            num_workers=self.args.num_workers,
-            pin_memory=True,
-            drop_last=True
-        )
-        
-        val_loader = DataLoader(
-            val_dataset,
-            batch_size=self.args.batch_size,
-            sampler=val_sampler,
-            shuffle=False,
-            num_workers=self.args.num_workers,
-            pin_memory=True,
-            drop_last=False
-        )
-        
-        return train_loader, val_loader
-    
-    def is_main_process(self):
-        """Check if this is the main process"""
-        return not self.distributed or rank() == 0
-    
-    def train_epoch(self, epoch):
-        """Train for one epoch"""
-        self.model.train()
-        if self.discriminator is not None:
-            self.discriminator.train()
-        
-        if self.distributed:
-            self.train_loader.sampler.set_epoch(epoch)
-        
-        total_losses = {
-            'total': 0, 'rec': 0, 'stft': 0, 'mel': 0, 
-            'commit': 0, 'semantic': 0, 'gen': 0, 'feat': 0, 'disc': 0
-        }
-        
-        pbar = tqdm(self.train_loader, desc=f'Epoch {epoch}', disable=not self.is_main_process())
-        
-        for batch_idx, (audio, paths) in enumerate(pbar):
-            audio = audio.to(self.device)
-            
-            # Create AudioSignal objects for loss computation
-            audio_signal = AudioSignal(audio, self.config['sample_rate'])
-            
-            # Forward pass with random bandwidth
-            bw_idx = random.randint(0, len(self.config['target_bandwidths']) - 1)
-            bw = self.config['target_bandwidths'][bw_idx]
-            
-            output, commit_loss, semantic_loss, _ = self.model(audio, bw)
-            recons_signal = AudioSignal(output, self.config['sample_rate'])
-            
-            # Check if discriminator should be active (after discriminator_start_step)
-            use_discriminator = (self.discriminator is not None and 
-                                self.global_step >= self.args.discriminator_start_step)
-            
-            # Train discriminator first if using GAN and past the start step
-            if use_discriminator and self.global_step % self.args.disc_interval == 0:
-                self.optimizer_d.zero_grad()
-                disc_loss = self.gan_loss.discriminator_loss(recons_signal, audio_signal)
-                disc_loss.backward()
-                torch.nn.utils.clip_grad_norm_(self.discriminator.parameters(), 10.0)
-                self.optimizer_d.step()
-                self.scheduler_d.step()
-                total_losses['disc'] += disc_loss.item()
-            
-            # Train generator
-            losses = {}
-            
-            # Reconstruction losses
-            losses['rec'] = self.l1_loss(recons_signal, audio_signal)
-            losses['stft'] = self.stft_loss(recons_signal, audio_signal)
-            # losses['mel'] = self.mel_loss(recons_signal, audio_signal)
-            losses['mel'] = torch.tensor(0.0, device=self.device) # 15.
-            losses['commit'] = commit_loss
-            losses['semantic'] = semantic_loss
-            
-            # GAN losses if discriminator is active
-            if use_discriminator:
-                gen_loss, feat_loss = self.gan_loss.generator_loss(recons_signal, audio_signal)
-                losses['gen'] = gen_loss
-                losses['feat'] = feat_loss
-            else:
-                # Set to zero for logging purposes
-                losses['gen'] = torch.tensor(0.0, device=self.device)
-                losses['feat'] = torch.tensor(0.0, device=self.device)
-            
-            # Total weighted loss
-            total_loss = sum(self.loss_weights.get(k, 0) * v for k, v in losses.items() 
-                           if k not in ['gen', 'feat'] or use_discriminator)
-            
-            # Backward pass
-            self.optimizer_g.zero_grad()
-            total_loss.backward()
-            torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
-            self.optimizer_g.step()
-            self.scheduler_g.step()
-            
-            # Update metrics
-            total_losses['total'] += total_loss.item()
-            for k, v in losses.items():
-                total_losses[k] += v.item()
-            
-            # Update progress bar
-            if self.is_main_process():
-                pbar.set_postfix({
-                    'loss': f'{total_loss.item():.4f}',
-                    'rec': f'{losses["rec"].item():.4f}',
-                    'mel': f'{losses["mel"].item():.4f}',
-                    'commit_loss': f'{losses["commit"].item():.4f}',
-                    'semantic_loss': f'{losses["semantic"].item():.4f}',
-                    'lr': f'{self.scheduler_g.get_last_lr()[0]:.9f}',
-                    'disc': 'ON' if use_discriminator else 'OFF',
-                    'step': self.global_step
-                })
-            
-            # Log to tensorboard
-            if self.is_main_process() and self.global_step % self.args.log_interval == 0:
-                for k, v in losses.items():
-                    self.writer.add_scalar(f'train/{k}_loss', v.item(), self.global_step)
-                self.writer.add_scalar('train/total_loss', total_loss.item(), self.global_step)
-                self.writer.add_scalar('train/lr', self.scheduler_g.get_last_lr()[0], self.global_step)
-                self.writer.add_scalar('train/bandwidth', bw, self.global_step)
-                self.writer.add_scalar('train/discriminator_active', float(use_discriminator), self.global_step)
-                if use_discriminator:
-                    self.writer.add_scalar('train/disc_loss', total_losses['disc'] / max(1, batch_idx), self.global_step)
-            
-            # Save checkpoint at step intervals
-            if self.global_step > 0 and self.global_step % self.args.save_step_interval == 0:
-                self.save_checkpoint_step(self.global_step)
-                if self.is_main_process():
-                    print(f"\nSaved checkpoint at step {self.global_step}")
-            
-            self.global_step += 1
-        
-        # Return average losses
-        n_batches = len(self.train_loader)
-        return {k: v / n_batches for k, v in total_losses.items()}
-    
-    @torch.no_grad()
-    def validate(self, epoch):
-        """Validation loop"""
-        self.model.eval()
-        
-        total_losses = {
-            'total': 0, 'rec': 0, 'stft': 0, 'mel': 0, 
-            'commit': 0, 'semantic': 0
-        }
-        
-        # Store audio samples for tensorboard
-        audio_samples = {'train': [], 'val': []}
-        
-        for batch_idx, (audio, paths) in enumerate(tqdm(self.val_loader, desc='Validation', disable=not self.is_main_process())):
-            audio = audio.to(self.device)
-            audio_signal = AudioSignal(audio, self.config['sample_rate'])
-            
-            # Use medium bandwidth for validation
-            bw = self.config['target_bandwidths'][2]
-            
-            output, commit_loss, semantic_loss, _ = self.model(audio, bw)
-            recons_signal = AudioSignal(output, self.config['sample_rate'])
-            
-            # Compute losses
-            losses = {
-                'rec': self.l1_loss(recons_signal, audio_signal),
-                'stft': self.stft_loss(recons_signal, audio_signal),
-                'mel': self.mel_loss(recons_signal, audio_signal),
-                'commit': commit_loss,
-                'semantic': semantic_loss
-            }
-            
-            total_loss = sum(self.loss_weights.get(k, 0) * v for k, v in losses.items())
-            
-            total_losses['total'] += total_loss.item()
-            for k, v in losses.items():
-                total_losses[k] += v.item()
-            
-            # Collect audio samples for tensorboard (first 3 from validation)
-            if self.is_main_process() and len(audio_samples['val']) < 3:
-                audio_samples['val'].append({
-                    'original': audio[0].cpu(),
-                    'reconstructed': output[0].cpu(),
-                    'path': paths[0]
-                })
-        
-        # Get train samples for comparison
-        if self.is_main_process():
-            self.model.eval()
-            for batch_idx, (audio, paths) in enumerate(self.train_loader):
-                if len(audio_samples['train']) >= 3:
-                    break
-                audio = audio.to(self.device)
-                bw = self.config['target_bandwidths'][2]
-                output, _, _, _ = self.model(audio, bw)
-                audio_samples['train'].append({
-                    'original': audio[0].cpu(),
-                    'reconstructed': output[0].cpu(),
-                    'path': paths[0]
-                })
-        
-        # Log audio samples to tensorboard
-        if self.is_main_process():
-            for split in ['train', 'val']:
-                for idx, sample in enumerate(audio_samples[split]):
-                    self.writer.add_audio(
-                        f'{split}/original_{idx}',
-                        sample['original'],
-                        epoch,
-                        sample_rate=self.config['sample_rate']
-                    )
-                    self.writer.add_audio(
-                        f'{split}/reconstructed_{idx}',
-                        sample['reconstructed'],
-                        epoch,
-                        sample_rate=self.config['sample_rate']
-                    )
-        
-        # Average losses
-        n_batches = len(self.val_loader)
-        val_metrics = {k: v / n_batches for k, v in total_losses.items()}
-        
-        # Log validation metrics
-        if self.is_main_process():
-            for key, value in val_metrics.items():
-                self.writer.add_scalar(f'val/{key}_loss', value, epoch)
-        
-        return val_metrics
-        
-    def save_checkpoint(self, epoch, is_best=False):
-        """Save model checkpoint (epoch-based)"""
-        if not self.is_main_process():
-            return
-        
-        model_state = self.model.module.state_dict() if self.distributed else self.model.state_dict()
-        
-        # Get current learning rates for verification
-        current_lr_g = self.scheduler_g.get_last_lr()[0]
-        
-        checkpoint = {
-            'epoch': epoch,
-            'global_step': self.global_step,
-            'model_state_dict': model_state,
-            'optimizer_g_state_dict': self.optimizer_g.state_dict(),
-            'scheduler_g_state_dict': self.scheduler_g.state_dict(),
-            'scheduler_g_last_epoch': self.scheduler_g.last_epoch,  # Explicitly save this
-            'current_lr_g': current_lr_g,  # Save for verification
-            'config': self.config,
-            'args': self.args
-        }
-        
-        if self.discriminator is not None:
-            disc_state = self.discriminator.module.state_dict() if self.distributed else self.discriminator.state_dict()
-            current_lr_d = self.scheduler_d.get_last_lr()[0]
-            checkpoint['discriminator_state_dict'] = disc_state
-            checkpoint['optimizer_d_state_dict'] = self.optimizer_d.state_dict()
-            checkpoint['scheduler_d_state_dict'] = self.scheduler_d.state_dict()
-            checkpoint['scheduler_d_last_epoch'] = self.scheduler_d.last_epoch
-            checkpoint['current_lr_d'] = current_lr_d
-        
-        # Save latest checkpoint
-        checkpoint_path = os.path.join(self.args.output_dir, 'checkpoints', 'latest.pth')
-        os.makedirs(os.path.dirname(checkpoint_path), exist_ok=True)
-        torch.save(checkpoint, checkpoint_path)
-        
-        # Save best checkpoint
-        if is_best:
-            best_path = os.path.join(self.args.output_dir, 'checkpoints', 'best.pth')
-            torch.save(checkpoint, best_path)
-        
-        # Save periodic checkpoint
-        if epoch % self.args.save_interval == 0:
-            epoch_path = os.path.join(self.args.output_dir, 'checkpoints', f'epoch_{epoch}.pth')
-            torch.save(checkpoint, epoch_path)
-
-
-    def save_checkpoint_step(self, step):
-        """Save model checkpoint (step-based)"""
-        if not self.is_main_process():
-            return
-        
-        # Get current epoch from training loop
-        current_epoch = step // len(self.train_loader)
-        
-        model_state = self.model.module.state_dict() if self.distributed else self.model.state_dict()
-        
-        # Get current learning rates for verification
-        current_lr_g = self.scheduler_g.get_last_lr()[0]
-        
-        checkpoint = {
-            'epoch': current_epoch,
-            'global_step': step,
-            'model_state_dict': model_state,
-            'optimizer_g_state_dict': self.optimizer_g.state_dict(),
-            'scheduler_g_state_dict': self.scheduler_g.state_dict(),
-            'scheduler_g_last_epoch': self.scheduler_g.last_epoch,  # Explicitly save this
-            'current_lr_g': current_lr_g,  # Save for verification
-            'config': self.config,
-            'args': self.args
-        }
-        
-        if self.discriminator is not None:
-            disc_state = self.discriminator.module.state_dict() if self.distributed else self.discriminator.state_dict()
-            current_lr_d = self.scheduler_d.get_last_lr()[0]
-            checkpoint['discriminator_state_dict'] = disc_state
-            checkpoint['optimizer_d_state_dict'] = self.optimizer_d.state_dict()
-            checkpoint['scheduler_d_state_dict'] = self.scheduler_d.state_dict()
-            checkpoint['scheduler_d_last_epoch'] = self.scheduler_d.last_epoch
-            checkpoint['current_lr_d'] = current_lr_d
-        
-        # Save step-based checkpoint
-        step_path = os.path.join(self.args.output_dir, 'checkpoints', f'step_{step}.pth')
-        torch.save(checkpoint, step_path)
-        
-        # Also update latest checkpoint
-        latest_path = os.path.join(self.args.output_dir, 'checkpoints', 'latest.pth')
-        torch.save(checkpoint, latest_path)
-        
-        # Keep only the last N step-based checkpoints to save disk space
-        if self.args.keep_last_n_steps > 0:
-            checkpoint_dir = os.path.join(self.args.output_dir, 'checkpoints')
-            step_checkpoints = sorted([f for f in os.listdir(checkpoint_dir) if f.startswith('step_')])
-            if len(step_checkpoints) > self.args.keep_last_n_steps:
-                for old_checkpoint in step_checkpoints[:-self.args.keep_last_n_steps]:
-                    os.remove(os.path.join(checkpoint_dir, old_checkpoint))
-
-
-    def load_checkpoint(self):
-        """Load checkpoint with proper state restoration"""
-        checkpoint_path = os.path.join(self.args.output_dir, 'checkpoints', 'latest.pth')
-        if os.path.exists(checkpoint_path):
-            print(f"Loading checkpoint from {checkpoint_path}")
-            checkpoint = torch.load(checkpoint_path, map_location=self.device, weights_only=False)
-            
-            # Load model state
-            if self.distributed:
-                self.model.module.load_state_dict(checkpoint['model_state_dict'])
-            else:
-                self.model.load_state_dict(checkpoint['model_state_dict'])
-            
-            # Load optimizer state
-            self.optimizer_g.load_state_dict(checkpoint['optimizer_g_state_dict'])
-            
-            # Load scheduler state
-            self.scheduler_g.load_state_dict(checkpoint['scheduler_g_state_dict'])
-            
-            # Restore scheduler's last_epoch from checkpoint
-            if 'scheduler_g_last_epoch' in checkpoint:
-                self.scheduler_g.last_epoch = checkpoint['scheduler_g_last_epoch']
-            else:
-                # Fallback: use global_step if the explicit value wasn't saved
-                self.scheduler_g.last_epoch = checkpoint['global_step']
-            
-            # Force scheduler to recompute its internal state
-            self.scheduler_g._last_lr = self.scheduler_g.get_lr()
-            
-            # Load discriminator if present
-            if self.discriminator is not None and 'discriminator_state_dict' in checkpoint:
-                if self.distributed:
-                    self.discriminator.module.load_state_dict(checkpoint['discriminator_state_dict'])
-                else:
-                    self.discriminator.load_state_dict(checkpoint['discriminator_state_dict'])
-                self.optimizer_d.load_state_dict(checkpoint['optimizer_d_state_dict'])
-                self.scheduler_d.load_state_dict(checkpoint['scheduler_d_state_dict'])
-                
-                # Restore discriminator scheduler's last_epoch
-                if 'scheduler_d_last_epoch' in checkpoint:
-                    self.scheduler_d.last_epoch = checkpoint['scheduler_d_last_epoch']
-                else:
-                    self.scheduler_d.last_epoch = checkpoint['global_step']
-                
-                self.scheduler_d._last_lr = self.scheduler_d.get_lr()
-            
-            # Restore training state
-            self.start_epoch = checkpoint['epoch'] + 1
-            self.global_step = checkpoint['global_step']
-            
-            # Verify learning rate restoration
-            current_lr_g = self.scheduler_g.get_last_lr()[0]
-            saved_lr_g = checkpoint.get('current_lr_g', None)
-            
-            print(f"\n{'='*60}")
-            print(f"CHECKPOINT LOADED SUCCESSFULLY")
-            print(f"{'='*60}")
-            print(f"Resumed from epoch: {checkpoint['epoch']}")
-            print(f"Global step: {self.global_step}")
-            print(f"Scheduler last_epoch: {self.scheduler_g.last_epoch}")
-            print(f"Current learning rate (generator): {current_lr_g:.9f}")
-            if saved_lr_g is not None:
-                print(f"Saved learning rate (generator): {saved_lr_g:.9f}")
-                if abs(current_lr_g - saved_lr_g) > 1e-9:
-                    print("⚠️  WARNING: Learning rate mismatch! This might indicate improper state restoration.")
-            
-            if self.discriminator is not None:
-                current_lr_d = self.scheduler_d.get_last_lr()[0]
-                saved_lr_d = checkpoint.get('current_lr_d', None)
-                print(f"Current learning rate (discriminator): {current_lr_d:.9f}")
-                if saved_lr_d is not None:
-                    print(f"Saved learning rate (discriminator): {saved_lr_d:.9f}")
-                print(f"Discriminator status: {'ACTIVE' if self.global_step >= self.args.discriminator_start_step else f'INACTIVE (starts at step {self.args.discriminator_start_step})'}")
-            
-            print(f"Next epoch: {self.start_epoch}")
-            print(f"Next step checkpoint at: step {((self.global_step // self.args.save_step_interval) + 1) * self.args.save_step_interval}")
-            print(f"{'='*60}\n")
-            
-            # Double-check by creating a fresh scheduler and comparing
-            if self.global_step > 0:
-                temp_scheduler = CosineWarmupScheduler(
-                    self.optimizer_g, 
-                    self.args.warmup_steps, 
-                    self.total_steps,
-                    eta_min=1e-6,
-                    last_epoch=-1
-                )
-                # Step it to the current global step
-                for _ in range(self.global_step):
-                    temp_scheduler.step()
-                expected_lr = temp_scheduler.get_last_lr()[0]
-                if abs(current_lr_g - expected_lr) > 1e-9:
-                    print(f"⚠️  Learning rate verification failed!")
-                    print(f"   Expected: {expected_lr:.9f}")
-                    print(f"   Got: {current_lr_g:.9f}")
-                    print("   The scheduler state might not be properly restored.")
-        else:
-            print(f"No checkpoint found at {checkpoint_path}, starting from scratch")
-        
-    def train(self):
-        """Main training loop"""
-        best_val_loss = float('inf')
-        
-        # Print training configuration
-        if self.is_main_process():
-            print(f"\n{'='*50}")
-            print(f"Training Configuration:")
-            print(f"{'='*50}")
-            print(f"Total epochs: {self.args.num_epochs}")
-            print(f"Steps per epoch: {len(self.train_loader)}")
-            print(f"Total steps: {self.total_steps}")
-            print(f"Warmup steps: {self.args.warmup_steps}")
-            print(f"Discriminator starts at step: {self.args.discriminator_start_step}")
-            print(f"Checkpoint saving:")
-            print(f"  - Every {self.args.save_interval} epochs")
-            print(f"  - Every {self.args.save_step_interval} steps")
-            print(f"  - Keep last {self.args.keep_last_n_steps} step checkpoints")
-            if self.start_epoch > 0:
-                print(f"RESUMING from epoch {self.start_epoch}, step {self.global_step}")
-            print(f"{'='*50}\n")
-        
-        for epoch in range(self.start_epoch, self.args.num_epochs):
-            # IMPORTANT: Set the epoch for distributed sampler when resuming
-            # This ensures proper data shuffling across epochs
-            if self.distributed and hasattr(self.train_loader.sampler, 'set_epoch'):
-                self.train_loader.sampler.set_epoch(epoch)
-            
-            # Train
-            train_metrics = self.train_epoch(epoch)
-            
-            # Validate
-            val_metrics = self.validate(epoch)
-            
-            # Log epoch metrics
-            if self.is_main_process():
-                print(f"\nEpoch {epoch} Summary:")
-                print(f"Train - Total: {train_metrics['total']:.4f}, Rec: {train_metrics['rec']:.4f}, "
-                    f"STFT: {train_metrics['stft']:.4f}, Mel: {train_metrics['mel']:.4f}, "
-                    f"Commit: {train_metrics['commit']:.4f}, Semantic: {train_metrics['semantic']:.4f}")
-                if self.discriminator is not None:
-                    print(f"       Gen: {train_metrics['gen']:.4f}, Feat: {train_metrics['feat']:.4f}, "
-                        f"Disc: {train_metrics['disc']:.4f}")
-                    print(f"       Discriminator Status: {'Active' if self.global_step >= self.args.discriminator_start_step else f'Starting at step {self.args.discriminator_start_step}'}")
-                print(f"Val   - Total: {val_metrics['total']:.4f}, Rec: {val_metrics['rec']:.4f}, "
-                    f"STFT: {val_metrics['stft']:.4f}, Mel: {val_metrics['mel']:.4f}, "
-                    f"Commit: {val_metrics['commit']:.4f}, Semantic: {val_metrics['semantic']:.4f}")
-                print(f"Current Step: {self.global_step}, Next step checkpoint at: {((self.global_step // self.args.save_step_interval) + 1) * self.args.save_step_interval}")
-                print(f"Current LR: {self.scheduler_g.get_last_lr()[0]:.9f}")
-            
-            # Save checkpoint
-            is_best = val_metrics['total'] < best_val_loss
-            if is_best:
-                best_val_loss = val_metrics['total']
-            self.save_checkpoint(epoch, is_best)
-        
-        # Save final model
-        if self.is_main_process():
-            model_state = self.model.module.state_dict() if self.distributed else self.model.state_dict()
-            
-            final_path = os.path.join(self.args.output_dir, 'checkpoints', 'final.pth')
-            torch.save({
-                'model_state_dict': model_state,
-                'config': self.config
-            }, final_path)
-            
-            # Also save just the model weights in the format expected by the original code
-            model_only_path = os.path.join(self.args.output_dir, 'model.pth')
-            torch.save(model_state, model_only_path)
-            
-            # Copy config
-            import shutil
-            shutil.copy(self.args.config, os.path.join(self.args.output_dir, 'config.json'))
-        
-        # Cleanup
-        if self.is_main_process():
-            self.writer.close()
-        if self.distributed:
-            dist.destroy_process_group()
-
-
-def main():
-    parser = argparse.ArgumentParser(description='Train Boson Audio Codec')
-    
-    # Data arguments
-    parser.add_argument('--data_csv', type=str, required=True,
-                        help='Path to CSV file containing audio file paths')
-    parser.add_argument('--config', type=str, default='config.json',
-                        help='Path to config JSON file')
-    
-    # Training argumentssss
-    parser.add_argument('--batch_size', type=int, default=32,
-                        help='Batch size per GPU')
-    parser.add_argument('--num_epochs', type=int, default=100,
-                        help='Number of training epochs')
-    parser.add_argument('--learning_rate', type=float, default=1e-4,
-                        help='Initial learning rate')
-    parser.add_argument('--weight_decay', type=float, default=0.01,
-                        help='Weight decay')
-    parser.add_argument('--segment_duration', type=float, default=2.,
-                        help='Audio segment duration in seconds')
-    
-    # Scheduler arguments
-    parser.add_argument('--warmup_steps', type=int, default=5000,
-                        help='Number of warmup steps for cosine scheduler')
-    
-    # Loss arguments
-    parser.add_argument('--use_discriminator', action='store_true',
-                        help='Use adversarial training with discriminator')
-    parser.add_argument('--discriminator_start_step', type=int, default=24_000,
-                        help='Start training discriminator after N steps')
-    parser.add_argument('--disc_interval', type=int, default=1,
-                        help='Train discriminator every N steps')
-    
-    # System arguments
-    parser.add_argument('--output_dir', type=str, default='outputs',
-                        help='Output directory for checkpoints and logs')
-    parser.add_argument('--num_workers', type=int, default=16,
-                        help='Number of data loading workers')
-    parser.add_argument('--seed', type=int, default=42,
-                        help='Random seed')
-    parser.add_argument('--local_rank', type=int, default=0,
-                        help='Local rank for distributed training')
-    
-    # Logging arguments
-    parser.add_argument('--log_interval', type=int, default=10,
-                        help='Log every N steps')
-    parser.add_argument('--save_interval', type=int, default=1,
-                        help='Save checkpoint every N epochs')
-    parser.add_argument('--save_step_interval', type=int, default=1000,
-                        help='Save checkpoint every N steps')
-    parser.add_argument('--keep_last_n_steps', type=int, default=5,
-                        help='Keep only the last N step-based checkpoints (0 to keep all)')
-    
-    # Resume training
-    parser.add_argument('--resume', action='store_true',
-                        help='Resume training from latest checkpoint')
-    
-    args = parser.parse_args()
-    
-    # Create output directory
-    os.makedirs(args.output_dir, exist_ok=True)
-    
-    # Train
-    trainer = BosonTrainer(args)
-    trainer.train()
-
-
-if __name__ == '__main__':
-    main()