train_codec_mel_speaker.py

# train_codec_mel_speaker.py
# ==============================================================================
# Training script for Mel-based Codec with LEARNABLE SPEAKER EMBEDDINGS
#
# KEY FEATURES:
#   - Input: mel spectrogram + pitch (NO text embeddings, NO energy)
#   - Uses learnable speaker embeddings: nn.Embedding(num_speakers, speaker_dim)
#   - Speaker conditioning via AdaIN1d throughout decoder
#   - F0 prediction from latent representation
# Architecture:
#   MelPitchEncoder: (mel, pitch) -> latent
#   FSQ: latent -> tokens (spectral + prosody quantized)
#   Decoder: (quantized_latent, speaker_emb) -> waveform
#   F0 Predictor: latent -> pitch prediction
# Usage:
#   python train_codec_mel_speaker.py -p Configs/config_codec_mel_speaker.yml
# ==============================================================================

import os
import os.path as osp
import sys
import yaml
import shutil
import numpy as np
import torch
import click
import warnings
warnings.simplefilter('ignore')

import random
from munch import Munch
from torch import nn
import torch.nn.functional as F
import torchaudio
from torch.nn.utils.rnn import pad_sequence

from models_mel_speaker import build_model_mel_speaker, load_checkpoint
from meldataset import build_dataloader
from utils import get_data_path_list, log_print, length_to_mask, recursive_munch
from losses import MultiResolutionSTFTLoss, MagPhaseLoss, GeneratorLoss, DiscriminatorLoss, WavLMLoss
from optimizers import build_optimizer

import time
from accelerate import Accelerator
from accelerate.utils import LoggerType
from accelerate import DistributedDataParallelKwargs
from torch.utils.tensorboard import SummaryWriter

import logging
from accelerate.logging import get_logger
logger = get_logger(__name__, log_level="DEBUG")

# Fix "Too many open files" error with DataLoader workers
import torch.multiprocessing
torch.multiprocessing.set_sharing_strategy('file_system')

# Increase file descriptor limit
import resource
soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE)
resource.setrlimit(resource.RLIMIT_NOFILE, (min(65536, hard), hard))

import gc


# ==============================================================================
# Helper function to expand phoneme durations to frame-level targets
# ==============================================================================
def expand_durations_to_frames(phonemes, durations, max_frames):
    """
    Expand phoneme IDs by their durations to get frame-level phoneme targets.
    
    Args:
        phonemes: [B, T_phonemes] - phoneme IDs (with padding)
        durations: [B, T_phonemes] - duration for each phoneme (in frames at hop_length=441)
        max_frames: int - maximum number of frames to output
        
    Returns:
        expanded: [B, max_frames] - phoneme ID at each frame
        mask: [B, max_frames] - valid frames mask (1 where valid, 0 where padded)
    """
    B = phonemes.shape[0]
    device = phonemes.device
    
    expanded = torch.zeros(B, max_frames, dtype=torch.long, device=device)
    mask = torch.zeros(B, max_frames, dtype=torch.float, device=device)
    
    for b in range(B):
        frame_idx = 0
        for p_idx in range(phonemes.shape[1]):
            dur = int(durations[b, p_idx].item())
            if dur <= 0:
                continue
            phoneme_id = phonemes[b, p_idx].item()
            
            end_idx = min(frame_idx + dur, max_frames)
            if frame_idx >= max_frames:
                break
                
            expanded[b, frame_idx:end_idx] = phoneme_id
            mask[b, frame_idx:end_idx] = 1.0
            frame_idx = end_idx
            
    return expanded, mask


def crop_phonemes_and_durations(phonemes, durations, start_frame, end_frame):
    """
    Crop phonemes and durations to match a cropped mel region.
    
    Args:
        phonemes: [T_phonemes] - phoneme IDs (with padding at start and end)
        durations: [T_phonemes] - duration for each phoneme in frames
        start_frame: int - start frame of the crop
        end_frame: int - end frame of the crop (exclusive)
        
    Returns:
        cropped_phonemes: [T_cropped] - phoneme IDs in the cropped region
        cropped_durations: [T_cropped] - durations in the cropped region
    """
    # Find which phonemes fall in the cropped region
    frame_pos = 0
    cropped_phonemes = []
    cropped_durations = []
    
    for i in range(len(phonemes)):
        dur = int(durations[i].item())
        if dur <= 0:
            continue
            
        phone_start = frame_pos
        phone_end = frame_pos + dur
        
        # Check if this phoneme overlaps with our crop region
        if phone_end > start_frame and phone_start < end_frame:
            # Calculate how much of this phoneme is in the crop region
            overlap_start = max(phone_start, start_frame)
            overlap_end = min(phone_end, end_frame)
            cropped_dur = overlap_end - overlap_start
            
            if cropped_dur > 0:
                cropped_phonemes.append(phonemes[i].item())
                cropped_durations.append(cropped_dur)
        
        frame_pos = phone_end
        
        # Early exit if we've passed the crop region
        if frame_pos >= end_frame:
            break
            
    return cropped_phonemes, cropped_durations




@click.command()
@click.option('-p', '--config_path', default='Configs/config_codec_mel_speaker.yml', type=str)
def main(config_path):
    config = yaml.safe_load(open(config_path))
    
    log_dir = config['log_dir']
    if not osp.exists(log_dir):
        os.makedirs(log_dir, exist_ok=True)
    shutil.copy(config_path, osp.join(log_dir, osp.basename(config_path)))
    
    ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True)
    accelerator = Accelerator(
        project_dir=log_dir, 
        split_batches=True, 
        kwargs_handlers=[ddp_kwargs], 
        mixed_precision='bf16'
    )
    
    if accelerator.is_main_process:
        writer = SummaryWriter(log_dir + "/tensorboard")
        
    file_handler = logging.FileHandler(osp.join(log_dir, 'train_codec_mel_speaker.log'))
    file_handler.setLevel(logging.DEBUG)
    file_handler.setFormatter(logging.Formatter('%(levelname)s:%(asctime)s: %(message)s'))
    logger.logger.addHandler(file_handler)
    
    # Config
    batch_size = config.get('batch_size', 10)
    device = accelerator.device
    epochs = config.get('epochs_codec', 100)
    save_freq = config.get('save_freq', 2)
    log_interval = config.get('log_interval', 10)
    
    # n_mels - default to 80
    n_mels = config.get('n_mels', 80)
    
    data_params = config.get('data_params', None)
    sr = config['preprocess_params'].get('sr', 44100)
    train_path = data_params['train_data']
    val_path = data_params['val_data']
    root_path = data_params['root_path']
    min_length = data_params['min_length']
    max_len = config.get('max_len', 7500)
    
    train_list, val_list = get_data_path_list(train_path, val_path)
    
    # Determine dump directory based on n_mels
    dump_dir = config.get('dump_dir', f"./dump_{n_mels}")
    
    # Dataloaders
    train_dataloader = build_dataloader(
        train_list,
        root_path,
        alignment_path="../data_preparation_v2/alignments_train.safetensors",
        feature_dir=f"{dump_dir}/train",
        stats_path=f"{dump_dir}/speaker_stats.json",
        min_length=min_length,
        batch_size=batch_size,
        num_workers=16,
        dataset_config={},
        device=device
    )
    
    val_dataloader = build_dataloader(
        val_list,
        root_path,
        alignment_path="../data_preparation_v2/alignments_val.safetensors",
        feature_dir=f"{dump_dir}/val",
        stats_path=f"{dump_dir}/speaker_stats.json",
        min_length=min_length,
        batch_size=20,
        validation=True,
        num_workers=0,
        device=device,
        dataset_config={}
    )
    
    # Model params
    model_params = recursive_munch(config['model_params'])
    
    # Get hop_length from config
    hop_length = config['preprocess_params']['spect_params'].get('hop_length', 441)
    
    # Scheduler params
    scheduler_params = {
        "max_lr": float(config['optimizer_params'].get('lr', 1e-5)),
        "min_lr": float(config['optimizer_params'].get('min_lr', 1e-5)),
        "pct_start": float(config['optimizer_params'].get('pct_start', 0.33)),
        "epochs": epochs,
        "steps_per_epoch": len(train_dataloader),
    }
    
    decay_epochs = int(epochs * scheduler_params["pct_start"])
    print(f"[LR Schedule] Pre-TMA (epochs 0-{decay_epochs}): {scheduler_params['max_lr']:.2e} -> {scheduler_params['min_lr']:.2e}")
    print(f"[LR Schedule] TMA (epochs {decay_epochs}-{epochs}): constant {scheduler_params['min_lr']:.2e}")
    
    # Build model with speaker embeddings
    model = build_model_mel_speaker(model_params)
    
    # Log speaker embedding info
    num_speakers = model_params.get('num_speakers', 11)
    speaker_dim = model_params.get('speaker_dim', 128)
    print(f"[Mel Speaker Codec] Number of speakers: {num_speakers} (IDs 0-{num_speakers-1})")
    print(f"[Mel Speaker Codec] Speaker embedding dim: {speaker_dim}")
    print(f"[Mel Speaker Codec] Speaker embedding shape: {model.codec.speaker_embedding.weight.shape}")
    print(f"[Mel Speaker Codec] Input n_mels: {n_mels}")
    
    scheduler_params_dict = {key: scheduler_params.copy() for key in model}
    
    # Prepare models
    for k in model:
        model[k] = accelerator.prepare(model[k])
        model[k] = model[k].to(device)
        
    train_dataloader, val_dataloader = accelerator.prepare(train_dataloader, val_dataloader)

    optimizer = build_optimizer(
        {key: model[key].parameters() for key in model},
        scheduler_params_dict={key: scheduler_params.copy() for key in model},
        lr=float(config['optimizer_params'].get('lr', 1e-5))
    )
    
    for k, v in optimizer.optimizers.items():
        optimizer.optimizers[k] = accelerator.prepare(optimizer.optimizers[k])
        optimizer.schedulers[k] = accelerator.prepare(optimizer.schedulers[k])
        
    # Load pretrained
    start_epoch = 0
    iters = 0
    
    with accelerator.main_process_first():
        if config.get('pretrained_model', '') != '':
            model, optimizer, start_epoch, iters = load_checkpoint(
                model, optimizer, config['pretrained_model'],
                load_only_params=config.get('load_only_params', False)
            )
            print(f"Loaded pretrained model from {config['pretrained_model']}")
            
            desired_lr = float(config['optimizer_params'].get('lr', 1e-5))
            for key in optimizer.optimizers:
                for param_group in optimizer.optimizers[key].param_groups:
                    param_group['lr'] = desired_lr
            print(f"Reset learning rate to {desired_lr:.2e}")
            
    # Losses
    loss_params = Munch(config.get('loss_params', {}))
    TMA_epoch = loss_params.get('TMA_epoch', 5)
    
    stft_loss = MultiResolutionSTFTLoss().to(device)
    mag_phase_loss = MagPhaseLoss(
        n_fft=model_params.decoder.gen_istft_n_fft,
        hop_length=model_params.decoder.gen_istft_hop_size
    ).to(device)
    gl = GeneratorLoss(model.mpd, model.msd).to(device)
    dl = DiscriminatorLoss(model.mpd, model.msd).to(device)
    wl = WavLMLoss(model_sr=sr, slm_sr=16000).to(device)
    

    
    # Loss weights
    lambda_mel = loss_params.get('lambda_mel', 45.0)
    lambda_gen = loss_params.get('lambda_gen', 1.0)
    lambda_slm = loss_params.get('lambda_slm', 1.0)
    lambda_mag = loss_params.get('lambda_mag', 0.1)
    lambda_entropy = loss_params.get('lambda_entropy', 10.0)
    lambda_f0 = loss_params.get('lambda_f0',1.0)
    lambda_phoneme = loss_params.get('lambda_phoneme', 0.5)  # Phoneme prediction loss weight
    
    best_loss = float('inf')

    # Compression level
    codec_strides = model_params.get('codec_strides', [2])
    upsample_rates = model_params.decoder.upsample_rates
    
    codec_compression = 1
    for s in codec_strides:
        codec_compression *= s
    
    istft_hop = model_params.decoder.gen_istft_hop_size
    
    print(f"[Mel Speaker Codec] Compression: {codec_compression}x")
    print(f"[Mel Speaker Codec] Upsample Rates: {upsample_rates} (Total: {np.prod(upsample_rates)})")
    print(f"[Mel Speaker Codec] Token Rate: {44100 / (hop_length * codec_compression):.2f} Hz")
    print(f"[Mel Speaker Codec] max_len: {max_len}")
    
    # FSQ levels
    fsq_levels = model_params.get('fsq_levels', [4] * 6)
    codebook_size = np.prod(fsq_levels)
    print(f"[Mel Speaker Codec] FSQ Levels: {fsq_levels} ({codebook_size} codes)")
    
    for epoch in range(start_epoch, epochs):
        running_loss = 0
        start_time = time.time()
        _ = [model[key].train() for key in model]
        
        for i, batch in enumerate(train_dataloader):
            try:
                # ===============================================================
                # START OF TRAINING STEP
                # ===============================================================
                waves = batch[1]
                tensors = [b.to(device) for b in batch[2:13] if isinstance(b, torch.Tensor)]
                (
                    speakers_ids,  # [B] - speaker IDs
                    languages_ids,
                    mels,          # [B, n_mels, T] - mel spectrogram
                    mel_input_length,
                    _target_ph,
                    _target_ph_l,
                    context_ph,
                    context_lens,
                    durations_tg,
                    pitches,       # [B, T] - pitch (log F0)
                    energies,      # [B, T] - energy (not used in this model)
                ) = tensors
                
                # Gather lengths
                mel_input_length_all = accelerator.gather(mel_input_length)
                mel_len = min([int(mel_input_length_all.min().item() / 2 - 1), max_len // 2])
                mel_len_st = int(mel_input_length.min().item() / 2 - 1)
                
                # Prepare batches (random crop for training)
                mel_list, wav_list, f0_list, spk_list = [], [], [], []
                phoneme_list, phoneme_len_list = [], []
                lang_list = []
                dur_list = []
                cropped_phoneme_list = []
                cropped_dur_list = []
                
                for bib in range(len(mel_input_length)):
                    mel_length = int(mel_input_length[bib].item() / 2)
                    random_start = np.random.randint(0, mel_length - mel_len)
                    
                    # Crop mel spectrogram
                    mel_start_frame = random_start * 2
                    mel_end_frame = (random_start + mel_len) * 2
                    mel_crop = mels[bib, :, mel_start_frame:mel_end_frame]
                    mel_list.append(mel_crop)
                    
                    # Crop waveform
                    y = waves[bib][mel_start_frame * hop_length:mel_end_frame * hop_length]
                    wav_list.append(y.to(device))
                    
                    # Crop pitch
                    f0_list.append(pitches[bib, mel_start_frame:mel_end_frame])
                    
                    # Speaker ID
                    spk_list.append(speakers_ids[bib])
                    
                    # Language ID
                    lang_list.append(languages_ids[bib])
                    
                    # Phonemes and durations for this sample (full sequence for reference)
                    phoneme_list.append(_target_ph[bib])
                    dur_list.append(durations_tg[bib])
                    
                    # Crop phonemes and durations to match the cropped mel region
                    cropped_ph, cropped_dur = crop_phonemes_and_durations(
                        _target_ph[bib], durations_tg[bib], mel_start_frame, mel_end_frame
                    )
                    cropped_phoneme_list.append(cropped_ph)
                    cropped_dur_list.append(cropped_dur)
                    
                mel_batch = torch.stack(mel_list).detach()  # [B, n_mels, T]
                F0 = torch.stack(f0_list).float().detach()  # [B, T]
                wav = torch.stack(wav_list).float().detach()  # [B, T_audio]
                speaker_ids_batch = torch.stack(spk_list)  # [B] - speaker IDs
                language_ids_batch = torch.stack(lang_list)  # [B] - language IDs

                if mel_batch.shape[-1] < 80:
                    continue

  

                current_prob_pitch = min(0.8, 0.3 + (epoch / 20))
                use_predicted_f0 = random.random() < current_prob_pitch

                # ===============================================================
                # CODEC FORWARD WITH SPEAKER IDS AND LANGUAGE IDS
                # ===============================================================
                # Convert to padded tensors (handling potential length mismatches if simply stacking)
                if len(cropped_phoneme_list) > 0:
                    # Find max length in this batch
                    max_len_ph = max(len(x) for x in cropped_phoneme_list)
                    padded_phonemes = torch.zeros((len(cropped_phoneme_list), max_len_ph), dtype=torch.long, device=device)
                    padded_durations = torch.zeros((len(cropped_dur_list), max_len_ph), dtype=torch.float, device=device)
                    
                    for b in range(len(cropped_phoneme_list)):
                        L = len(cropped_phoneme_list[b])
                        if L > 0:
                            padded_phonemes[b, :L] = torch.tensor(cropped_phoneme_list[b], device=device)
                            padded_durations[b, :L] = torch.tensor(cropped_dur_list[b], dtype=torch.float, device=device)
                else:
                    padded_phonemes = None
                    padded_durations = None

                codec_out = model.codec(
                    mel=mel_batch,
                    pitch=F0,
                    speaker_ids=speaker_ids_batch,
                    language_ids=language_ids_batch,
                    use_predicted_f0=use_predicted_f0,
                    target_phonemes=padded_phonemes,
                    target_durations=padded_durations,
                )
                
                y_rec = codec_out['wav']
                mag = codec_out['mag']
                phase = codec_out['phase']
                commitment_loss = codec_out['commitment_loss']
                tokens = codec_out['tokens']
                f0_pred = codec_out['f0_pred']
                speaker_emb = codec_out['speaker_emb']  # [B, speaker_dim] learned speaker embedding
                
                # Get phoneme loss directly from model (calculated at token rate)
                loss_phoneme = codec_out.get('phoneme_loss')
                if loss_phoneme is None:
                    loss_phoneme = torch.tensor(0.0, device=device)
                
                # Match lengths for reconstruction
                min_len = min(y_rec.shape[-1], wav.shape[-1])
                y_rec = y_rec[..., :min_len]
                wav = wav[..., :min_len]
                
                # Match F0 lengths
                min_len_f0 = min(f0_pred.shape[-1], F0.shape[-1])
                f0_pred = f0_pred[..., :min_len_f0]
                F0_target = F0[..., :min_len_f0]
                
                # F0 Loss
                loss_f0 = F.mse_loss(f0_pred, F0_target)
                
                # ===============================================================
                # DISCRIMINATOR STEP
                # ===============================================================
                if epoch >= TMA_epoch:
                    optimizer.zero_grad()
                    d_loss = dl(wav.detach().unsqueeze(1).float(), y_rec.detach()).mean()
                    accelerator.backward(d_loss)
                    optimizer.step('msd')
                    optimizer.step('mpd')
                else:
                    d_loss = torch.tensor(0.0, device=device)
                    
                # === Generator Step ===
                optimizer.zero_grad()
                
                # STFT loss
                loss_mel = stft_loss(y_rec.squeeze(), wav.detach())
                
                # Mag/Phase loss
                loss_mag_phase = mag_phase_loss(mag, phase, wav.detach())
                
                # Codebook utilization
                with torch.no_grad():
                    B, n_q, T = tokens.shape
                    flat_tokens = tokens.reshape(-1)
                    unique_codes = len(torch.unique(flat_tokens))
                    codebook_size = np.prod(model.codec.fsq_levels)
                    perplexity = unique_codes / codebook_size * 100
                
                if epoch >= TMA_epoch:
                    loss_gen = gl(wav.detach().unsqueeze(1).float(), y_rec).mean()
                    loss_slm = wl(wav.detach(), y_rec).mean()
                    
                    g_loss = (
                        lambda_mel * loss_mel +
                        lambda_gen * loss_gen +
                        lambda_slm * loss_slm +
                        lambda_mag * loss_mag_phase +
                        lambda_entropy * commitment_loss +
                        lambda_f0 * loss_f0 +
                        lambda_phoneme * loss_phoneme
                    )
                else:
                    loss_gen = torch.tensor(0.0, device=device)
                    loss_slm = torch.tensor(0.0, device=device)
                    g_loss = (
                        lambda_mel * loss_mel +
                        lambda_mag * loss_mag_phase +
                        lambda_entropy * commitment_loss +
                        lambda_f0 * loss_f0 +
                        lambda_phoneme * loss_phoneme
                    )
                    
                running_loss += accelerator.gather(loss_mel).mean().item()
                
                accelerator.backward(g_loss)

                optimizer.step('codec')
                
                iters += 1

   
                entropy_loss = commitment_loss.item() if torch.is_tensor(commitment_loss) else commitment_loss
                phoneme_loss_val = loss_phoneme.item() if torch.is_tensor(loss_phoneme) else loss_phoneme

                if (i + 1) % log_interval == 0 and accelerator.is_main_process:
                    current_lr = optimizer.optimizers['codec'].param_groups[0]['lr']
                    
                    # Log unique speakers in batch
                    unique_speakers = torch.unique(speaker_ids_batch).tolist()
                    
                    log_print(
                        f'Epoch [{epoch+1}/{epochs}], Step [{i+1}/{len(train_dataloader)}], '
                        f'LR: {current_lr:.2e}, '
                        f'Mel: {running_loss / log_interval:.5f}, '
                        f'Gen: {loss_gen.item():.5f}, '
                        f'Disc: {d_loss.item():.5f}, '
                        f'F0: {loss_f0.item():.5f}, '
                        f'Ph: {phoneme_loss_val:.5f}, '
                        f'SLM: {loss_slm.item():.5f}, '
                        f'Entropy: {entropy_loss:.5f}, '
                        f'FSQ Util: {perplexity:.1f}% ({unique_codes}/{codebook_size}), '
                        f'Speakers: {unique_speakers}',
                        logger
                    )
                    
                    writer.add_scalar('train/learning_rate', current_lr, iters)
                    writer.add_scalar('train/mel_loss', running_loss / log_interval, iters)
                    writer.add_scalar('train/gen_loss', loss_gen.item(), iters)
                    writer.add_scalar('train/d_loss', d_loss.item(), iters)
                    writer.add_scalar('train/f0_loss', loss_f0.item(), iters)
                    writer.add_scalar('train/phoneme_loss', phoneme_loss_val, iters)
                    writer.add_scalar('train/slm_loss', loss_slm.item(), iters)
                    writer.add_scalar('train/entropy_loss', entropy_loss, iters)
                    writer.add_scalar('train/fsq_utilization', perplexity, iters)
                    writer.add_scalar('train/unique_codes', unique_codes, iters)
                    writer.add_scalar('train/mag_phase_loss', loss_mag_phase.item(), iters)
                    
                    running_loss = 0

            except torch.cuda.OutOfMemoryError:
                if accelerator.is_main_process:
                    print(f"[Warning] OOM at epoch {epoch}, step {i}. Flushing cache.")
                
                for key in model: 
                    model[key].zero_grad()
                    
                torch.cuda.empty_cache()
                gc.collect()
                continue
                
        # === Validation ===
        loss_test = 0
        _ = [model[key].eval() for key in model]
        eval_samples = []  # Store (mel_len, mel, f0, wav, speaker_id)
        
        with torch.no_grad():
            iters_test = 0
            for batch_idx, batch in enumerate(val_dataloader):
                try:
                    waves = batch[1]
                    tensors = [b.to(device) for b in batch[2:13] if isinstance(b, torch.Tensor)]
                    (
                        speakers_ids,
                        languages_ids,
                        mels,
                        mel_input_length,
                        _target_ph,
                        _target_ph_l,
                        context_ph,
                        context_lens,
                        durations_tg,
                        pitches,
                        energies,
                    ) = tensors
                    
                    # Collect samples for evaluation
                    if len(eval_samples) < 20:
                        for bib in range(len(mel_input_length)):
                            if len(eval_samples) >= 20: break
                            eval_samples.append((
                                mel_input_length[bib].item(),
                                mels[bib],
                                pitches[bib],
                                waves[bib],
                                speakers_ids[bib],
                                languages_ids[bib],
                            ))

                    mel_len = min([int(mel_input_length.min().item() / 2 - 1), max_len // 2])
                    mel_list, wav_list, f0_list, spk_list, lang_list = [], [], [], [], []
                    
                    for bib in range(len(mel_input_length)):
                        mel_length = int(mel_input_length[bib].item() / 2)
                        random_start = np.random.randint(0, mel_length - mel_len)
                        
                        mel_list.append(mels[bib, :, (random_start * 2):((random_start + mel_len) * 2)])
                        y = waves[bib][(random_start * 2) * hop_length:((random_start + mel_len) * 2) * hop_length]
                        wav_list.append(y.to(device))
                        f0_list.append(pitches[bib, (random_start * 2):((random_start + mel_len) * 2)])
                        spk_list.append(speakers_ids[bib])
                        lang_list.append(languages_ids[bib])
                        
                    wav = torch.stack(wav_list).float().detach()
                    F0_curve = torch.stack(f0_list).detach()
                    mel_batch = torch.stack(mel_list).detach()
                    speaker_ids_batch = torch.stack(spk_list)
                    language_ids_batch = torch.stack(lang_list)
                    
                    # Forward with speaker IDs
                    codec_out = model.codec(
                        mel=mel_batch,
                        pitch=F0_curve,
                        speaker_ids=speaker_ids_batch,
                        language_ids=language_ids_batch,
                    )
                    y_rec = codec_out['wav']
                    
                    min_len = min(y_rec.shape[-1], wav.shape[-1])
                    y_rec = y_rec[..., :min_len]
                    wav = wav[..., :min_len]
                    
                    loss_mel = stft_loss(y_rec.squeeze(), wav.detach())
                    loss_test += accelerator.gather(loss_mel).mean().item()
                    iters_test += 1
                except Exception as e:
                    print(f"Error in validation loop: {e}")
                    continue
                
        if accelerator.is_main_process:
            print(f'Epoch: {epoch + 1}')
            current_val_loss = loss_test / iters_test if iters_test > 0 else 0
            log_print(f'Validation loss: {current_val_loss:.3f}\n', logger)
            writer.add_scalar('eval/mel_loss', current_val_loss, epoch + 1)
            
            # Generate audio samples
            with torch.no_grad():
                for bib, sample in enumerate(eval_samples):  #
                    mel_len_val, gt_mel, f0_sample, gt_wav, speaker_id, language_id = sample
                    
                    mel_length = int(mel_len_val)
                    gt_mel = gt_mel[:, :mel_length].unsqueeze(0).to(device)
                    f0_sample = f0_sample[:mel_length].unsqueeze(0).float().to(device)
                    speaker_id_sample = speaker_id.unsqueeze(0).to(device)
                    language_id_sample = language_id.unsqueeze(0).to(device)
                    
                    # Forward with speaker ID (use GT f0)
                    codec_out = model.codec(
                        gt_mel, f0_sample, speaker_id_sample,
                        language_ids=language_id_sample,
                        use_predicted_f0=False
                    )
                    y_rec = codec_out['wav']
                    speaker_emb = codec_out['speaker_emb']
                    f0_pred = codec_out['f0_pred']
                    
                    writer.add_audio(
                        f'eval/mel_codec_{bib}_spk{speaker_id.item()}', 
                        y_rec.cpu().numpy().squeeze(), epoch, sample_rate=sr
                    )
                    
                    # Test with predicted F0
                    codec_out_pred_f0 = model.codec(
                        gt_mel, f0_sample, speaker_id_sample,
                        language_ids=language_id_sample,
                        use_predicted_f0=True
                    )
                    y_rec_pred_f0 = codec_out_pred_f0['wav']
                    
                    writer.add_audio(
                        f'eval/mel_codec_pred_f0_{bib}_spk{speaker_id.item()}', 
                        y_rec_pred_f0.cpu().numpy().squeeze(), epoch, sample_rate=sr
                    )
                    
                    print(f"  Sample {bib}: speaker_id={speaker_id.item()}, f0_pred shape={f0_pred.shape}")
                    
                    # Test tokenize + decode path
                    tokens, spk_emb = model.codec.tokenize(
                        gt_mel, f0_sample, speaker_id_sample, 
                        language_ids=language_id_sample
                    )
                    y_from_tokens, _ = model.codec.decode_tokens(
                        tokens, speaker_id_sample, language_id_sample, f0=f0_sample
                    )
                    writer.add_audio(
                        f'eval/from_tokens_{bib}_spk{speaker_id.item()}', 
                        y_from_tokens.cpu().numpy().squeeze(), epoch, sample_rate=sr
                    )
                    
                    # Save GT audio
                    gt_wav_len = int(mel_length * hop_length)
                    if len(gt_wav) > gt_wav_len:
                        gt_wav = gt_wav[:gt_wav_len]
                    writer.add_audio(f'eval/gt_{bib}', gt_wav.cpu().numpy(), epoch, sample_rate=sr)
        # Save checkpoint
        if epoch % save_freq == 0:
            if accelerator.is_main_process:
                print('Saving checkpoint...')
 
                state = {
                    'net': {key: model[key].state_dict() for key in model},
                    'optimizer': optimizer.state_dict(),
                    'iters': iters,
                    'val_loss': loss_test / iters_test if iters_test > 0 else 0,
                    'epoch': epoch,
                }
                save_path = osp.join(log_dir, f'epoch_mel_speaker_codec_{epoch:05d}.pth')
                torch.save(state, save_path)
        
        gc.collect()
        torch.cuda.empty_cache()
                
    # Final save
    if accelerator.is_main_process:
        print('Saving final checkpoint...')
        
        net_state = {
            key: accelerator.unwrap_model(model[key]).state_dict() 
            for key in model
        }
        
        state = {
            'net': net_state,
            'optimizer': optimizer.state_dict(),
            'iters': iters,
            'val_loss': loss_test / iters_test if iters_test > 0 else 0,
            'epoch': epoch,
        }
        torch.save(state, osp.join(log_dir, 'mel_speaker_codec_final.pth'))


if __name__ == "__main__":
    main()