Modules/codec_decoder_mel_speaker.py

import math
import random
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.utils import weight_norm, remove_weight_norm
from scipy.signal import get_window
from einops import rearrange
from typing import Tuple, Optional, List, Dict, Union

from .conformer import Conformer
from .utils import init_weights, get_padding

class TorchSTFT(nn.Module):
    def __init__(self, filter_length=800, hop_length=200, win_length=800, window="hann"):
        super().__init__()
        self.filter_length = filter_length
        self.hop_length = hop_length
        self.win_length = win_length
        self.window = torch.from_numpy(
            get_window(window, win_length, fftbins=True).astype(np.float32)
        )

    def transform(self, input_data):
        forward_transform = torch.stft(
            input_data,
            self.filter_length,
            self.hop_length,
            self.win_length,
            window=self.window.to(input_data.device),
            return_complex=True,
        )
        return torch.abs(forward_transform), torch.angle(forward_transform)

    def inverse(self, magnitude, phase):
        inverse_transform = torch.istft(
            magnitude * torch.exp(phase * 1j),
            self.filter_length,
            self.hop_length,
            self.win_length,
            window=self.window.to(magnitude.device),
        )
        return inverse_transform.unsqueeze(-2)


class Snake1d(nn.Module):
    """Learned periodic activation from BigVGAN."""
    def __init__(self, in_features):
        super().__init__()
        self.alpha = nn.Parameter(torch.ones(1, in_features, 1))

    def forward(self, x):
        return x + (1.0 / (self.alpha + 1e-9)) * (torch.sin(self.alpha * x) ** 2)


class AdaIN1d(nn.Module):
    """
    Adaptive Instance Normalization for 1D signals.
    Takes a style vector [B, style_dim] and applies affine transformation
    to normalized features [B, C, T].
    """
    def __init__(self, style_dim, num_features):
        super().__init__()
        self.norm = nn.InstanceNorm1d(num_features, affine=False)
        self.fc = nn.Linear(style_dim, num_features * 2)

    def forward(self, x, s):
        """
        Args:
            x: [B, C, T] input features
            s: [B, style_dim] style/speaker embedding
        Returns:
            [B, C, T] AdaIN-transformed features
        """
        h = self.fc(s)
        h = h.view(h.size(0), h.size(1), 1)
        gamma, beta = torch.chunk(h, chunks=2, dim=1)
        return (1 + gamma) * self.norm(x) + beta


class SpeakerAdaINResBlock1(nn.Module):
    """
    Residual block with AdaIN speaker conditioning.
    Uses global speaker embedding [B, speaker_dim] for style.
    """
    def __init__(self, channels, kernel_size=3, dilation=(1, 3, 5), speaker_dim=128):
        super().__init__()
        self.convs1 = nn.ModuleList([
            weight_norm(nn.Conv1d(channels, channels, kernel_size, 1,
                                  dilation=d, padding=get_padding(kernel_size, d)))
            for d in dilation
        ])
        self.convs1.apply(init_weights)
        
        self.convs2 = nn.ModuleList([
            weight_norm(nn.Conv1d(channels, channels, kernel_size, 1,
                                  dilation=1, padding=get_padding(kernel_size, 1)))
            for _ in dilation
        ])
        self.convs2.apply(init_weights)
        
        self.adain1 = nn.ModuleList([AdaIN1d(speaker_dim, channels) for _ in dilation])
        self.adain2 = nn.ModuleList([AdaIN1d(speaker_dim, channels) for _ in dilation])
        self.snakes1 = nn.ModuleList([Snake1d(channels) for _ in dilation])
        self.snakes2 = nn.ModuleList([Snake1d(channels) for _ in dilation])

    def forward(self, x, speaker_emb):
        """
        Args:
            x: [B, C, T] input features
            speaker_emb: [B, speaker_dim] speaker embedding
        """
        for c1, c2, n1, n2, s1, s2 in zip(
            self.convs1, self.convs2, self.adain1, self.adain2, self.snakes1, self.snakes2
        ):
            xt = n1(x, speaker_emb)
            xt = s1(xt)
            xt = c1(xt)
            xt = n2(xt, speaker_emb)
            xt = s2(xt)
            xt = c2(xt)
            x = xt + x
        return x


# ==============================================================================
# Harmonic Source Module 
# ==============================================================================

class SineGen(nn.Module):
    """Sine generator for F0-based harmonic source with phase caching."""
    def __init__(self, samp_rate, upsample_scale, harmonic_num=0,
                 sine_amp=0.1, noise_std=0.003, voiced_threshold=0,
                 flag_for_pulse=False):
        super().__init__()
        self.sine_amp = sine_amp
        self.noise_std = noise_std
        self.harmonic_num = harmonic_num
        self.dim = harmonic_num + 1
        self.sampling_rate = samp_rate
        self.voiced_threshold = voiced_threshold
        self.upsample_scale = upsample_scale
        self.flag_for_pulse = flag_for_pulse

    def _f02uv(self, f0):
        return (f0 > self.voiced_threshold).float()

    def _f02sine(self, f0_values, initial_phase=None):
        rad_values = (f0_values / self.sampling_rate) % 1
        
        rand_ini = torch.rand(f0_values.shape[0], f0_values.shape[2], device=f0_values.device)
        rand_ini[:, 0] = 0
        rad_values[:, 0, :] = rad_values[:, 0, :] + rand_ini

        rad_values = F.interpolate(
            rad_values.transpose(1, 2),
            scale_factor=1 / self.upsample_scale,
            mode="linear",
        ).transpose(1, 2)

        phase = torch.cumsum(rad_values, dim=1) * 2 * np.pi
        
        if initial_phase is not None:
            phase = phase + initial_phase

        phase = F.interpolate(
            phase.transpose(1, 2) * self.upsample_scale,
            scale_factor=self.upsample_scale,
            mode="linear",
        ).transpose(1, 2)
        
        last_phase = phase[:, -1:, :]

        if self.flag_for_pulse:
            sines = torch.cos(phase)
        else:
            sines = torch.sin(phase)
            
        return sines, last_phase

    def forward(self, f0, initial_phase=None):
        f0_buf = torch.zeros(f0.shape[0], f0.shape[1], self.dim, device=f0.device)
        fn = torch.multiply(
            f0, torch.FloatTensor([[range(1, self.harmonic_num + 2)]]).to(f0.device)
        )
        
        sine_waves, next_phase = self._f02sine(fn, initial_phase)
        sine_waves = sine_waves * self.sine_amp
        
        uv = self._f02uv(f0)
        noise_amp = uv * self.noise_std + (1 - uv) * self.sine_amp / 3
        noise = noise_amp * torch.randn_like(sine_waves)
        sine_waves = sine_waves * uv + noise
        
        return sine_waves, uv, noise, next_phase


class SourceModuleHnNSF(nn.Module):
    """Source module for harmonic-plus-noise synthesis."""
    def __init__(self, sampling_rate, upsample_scale, harmonic_num=0,
                 sine_amp=0.1, add_noise_std=0.003, voiced_threshold=0):
        super().__init__()
        self.sine_amp = sine_amp
        self.noise_std = add_noise_std
        self.l_sin_gen = SineGen(
            sampling_rate, upsample_scale, harmonic_num,
            sine_amp, add_noise_std, voiced_threshold,
            flag_for_pulse=False
        )
        self.l_linear = nn.Linear(harmonic_num + 1, 1)
        self.l_tanh = nn.Tanh()

    def forward(self, x, cache=None):
        initial_phase = cache

        with torch.no_grad():
            sine_wavs, uv, _, next_phase = self.l_sin_gen(x, initial_phase=initial_phase)
            
        sine_merge = self.l_tanh(self.l_linear(sine_wavs))
        noise = torch.randn_like(uv) * self.sine_amp / 3
        
        return sine_merge, noise, uv, next_phase


# ==============================================================================
# Pixel Shuffle Upsampling
# ==============================================================================

def pixel_shuffle_1d(x: torch.Tensor, r: int) -> torch.Tensor:
    B, Cr, L = x.size()
    C = Cr // r
    x = x.view(B, C, r, L).permute(0, 1, 3, 2)
    return x.reshape(B, C, L * r)


class UpsamplePixelShuffle1D(nn.Module):
    def __init__(self, in_ch: int, out_ch: int, kernel_size: int, r: int):
        super().__init__()
        self.r = r
        pad_l, pad_r = (kernel_size - 1) // 2, kernel_size // 2
        self.pad = nn.ReflectionPad1d((pad_l, pad_r))
        self.conv = weight_norm(nn.Conv1d(in_ch, out_ch * r, kernel_size, padding=0))
        self._init_icnr(in_ch, out_ch, r, kernel_size)

    def _init_icnr(self, in_ch, out_ch, r, kernel_size):
        """ICNR initialization for smooth upsampling."""
        weight = self.conv.weight.data
        kernel = torch.zeros(out_ch, in_ch, kernel_size)
        nn.init.kaiming_normal_(kernel)
        weight.copy_(kernel.repeat(r, 1, 1))
        if self.conv.bias is not None:
            self.conv.bias.data.fill_(0)

    def forward(self, x):
        x = self.pad(x)
        x = self.conv(x)
        return pixel_shuffle_1d(x, self.r)


class EncoderResBlock1d(nn.Module):
    """Lightweight residual block used in the mel/pitch encoder."""
    def __init__(self, channels: int, kernel_size: int = 3, dilation: int = 1):
        super().__init__()
        padding = get_padding(kernel_size, dilation)
        self.block = nn.Sequential(
            weight_norm(nn.Conv1d(channels, channels, kernel_size, padding=padding, dilation=dilation)),
            nn.SiLU(),
            weight_norm(nn.Conv1d(channels, channels, kernel_size, padding=padding, dilation=dilation)),
        )
        self.act = nn.SiLU()

    def forward(self, x):
        return self.act(x + self.block(x))


# ==============================================================================
# Mel + Pitch Encoder (replaces text_emb + energy encoder)
# ==============================================================================

class MelPitchEncoderSpeaker(nn.Module):
    """
    Encoder that takes mel spectrogram and pitch as input.
    Produces a latent representation suitable for quantization.
    
    Supports variable compression ratios via strides parameter:
    - strides=[2]: 2x compression (~50Hz tokens)
    - strides=[4]: 4x compression (~25Hz tokens)
    - strides=[2,2]: 4x compression (~25Hz tokens)
    """
    def __init__(
        self,
        n_mels: int = 80,
        speaker_dim: int = 128,
        latent_dim: int = 256,
        hidden_dim: int = 256,
        strides: List[int] = [2],
    ):
        super().__init__()
        self.latent_dim = latent_dim
        self.speaker_dim = speaker_dim
        self.n_mels = n_mels
        self.compression_ratio = int(np.prod(strides))
        self.strides = strides

        # Build mel encoder from EXACT stride list for predictable token rate.
        mel_layers = []
        in_ch = n_mels
        for stage_idx, stride in enumerate(self.strides):
            k = 2 * stride + 1
            p = stride
            mel_layers.extend([
                weight_norm(nn.Conv1d(in_ch, hidden_dim, kernel_size=k, stride=stride, padding=p)),
                nn.SiLU(),
                EncoderResBlock1d(hidden_dim, kernel_size=3, dilation=1),
            ])
            in_ch = hidden_dim

        mel_layers.extend([
            weight_norm(nn.Conv1d(hidden_dim, hidden_dim, 3, stride=1, padding=1)),
            nn.SiLU(),
            EncoderResBlock1d(hidden_dim, kernel_size=3, dilation=1),
            weight_norm(nn.Conv1d(hidden_dim, hidden_dim * 2, 3, stride=1, padding=1)),
            nn.SiLU(),
        ])
        self.mel_encoder = nn.Sequential(*mel_layers)

        # Build pitch encoder from EXACT stride list.
        # Input is [logf0, uv] -> 2 channels.
        pitch_layers = []
        in_ch = 2
        pitch_hidden = hidden_dim // 2
        for stage_idx, stride in enumerate(self.strides):
            k = 2 * stride + 1
            p = stride
            pitch_layers.extend([
                weight_norm(nn.Conv1d(in_ch, pitch_hidden, kernel_size=k, stride=stride, padding=p)),
                nn.SiLU(),
                EncoderResBlock1d(pitch_hidden, kernel_size=3, dilation=1),
            ])
            in_ch = pitch_hidden

        pitch_layers.extend([
            weight_norm(nn.Conv1d(pitch_hidden, hidden_dim, 3, stride=1, padding=1)),
            nn.SiLU(),
            EncoderResBlock1d(hidden_dim, kernel_size=3, dilation=1),
        ])
        self.pitch_encoder = nn.Sequential(*pitch_layers)

        # Fusion of mel and pitch features
        input_dim = hidden_dim * 2 + hidden_dim  # mel features + pitch features
        self.fusion = nn.Sequential(
            weight_norm(nn.Conv1d(input_dim, hidden_dim * 2, 7, padding=3)),
            nn.SiLU(),
            weight_norm(nn.Conv1d(hidden_dim * 2, hidden_dim * 2, 5, padding=2)),
            nn.SiLU(),
            weight_norm(nn.Conv1d(hidden_dim * 2, hidden_dim * 2, 3, padding=1)),
            nn.SiLU(),
        )

        # Refinement layers
        self.refine = nn.Sequential(
            weight_norm(nn.Conv1d(hidden_dim * 2, hidden_dim * 2, 7, padding=3)),
            nn.SiLU(),
            weight_norm(nn.Conv1d(hidden_dim * 2, hidden_dim * 2, 5, padding=2)),
            nn.SiLU(),
            weight_norm(nn.Conv1d(hidden_dim * 2, hidden_dim * 2, 3, padding=1)),
            nn.SiLU(),
        )

        # Project to latent dimension
        self.to_latent = nn.Sequential(
            weight_norm(nn.Conv1d(hidden_dim * 2, hidden_dim * 2, 5, padding=2)),
            nn.SiLU(),
            weight_norm(nn.Conv1d(hidden_dim * 2, latent_dim, 1)),
        )

    def forward(self, mel, pitch):
        """
        Encode mel spectrogram + pitch into latent.
        
        Args:
            mel: [B, n_mels, T] - mel spectrogram
            pitch: [B, T] - pitch contour (log F0)
            
        Returns:
            latent: [B, latent_dim, T / compression_ratio]
        """
        # Encode mel spectrogram
        mel_feat = self.mel_encoder(mel)  # [B, hidden_dim * 2, T/compression]
        
        # Encode pitch + voiced/unvoiced cue
        pitch_in = pitch.unsqueeze(1)
        uv = (pitch > 0).float().unsqueeze(1)
        pitch_feat = self.pitch_encoder(torch.cat([pitch_in, uv], dim=1))  # [B, hidden_dim, T/compression]
        
        # Align lengths
        min_len = min(mel_feat.shape[-1], pitch_feat.shape[-1])
        mel_feat = mel_feat[..., :min_len]
        pitch_feat = pitch_feat[..., :min_len]
        
        # Concatenate and fuse
        x = torch.cat([mel_feat, pitch_feat], dim=1)
        x = self.fusion(x)
        x = self.refine(x)
        
        return self.to_latent(x)


# ==============================================================================
# Learnable Upsampling with Anti-Aliasing
# ==============================================================================

class TokenRatePhonemePredictor(nn.Module):
    """
    Predict phoneme logits at TOKEN RATE directly from quantized_latent.
    No temporal conv (kernel_size=1 only) to prevent the head from learning
    phonetic structure by itself.

    Inputs:
      z_q: [B, latent_dim, T_tok]
      lang_emb (optional): [B, lang_dim]
    Output:
      logits: [B, n_phonemes, T_tok]
    """

    def __init__(
        self,
        latent_dim: int,
        n_phonemes: int,
        language_emb_dim: int = 64,
        hidden_dim: int = 256,
        dropout_p: float = 0.1,
    ):
        super().__init__()
        self.n_phonemes = n_phonemes
        self.language_emb_dim = language_emb_dim

        in_dim = latent_dim + (language_emb_dim if language_emb_dim > 0 else 0)

        self.net = nn.Sequential(
            weight_norm(nn.Conv1d(in_dim, hidden_dim, kernel_size=1)),
            nn.SiLU(),
            nn.Dropout(dropout_p),
            nn.Conv1d(hidden_dim, n_phonemes, kernel_size=1),
        )

    @staticmethod
    def durations_frames_to_tokens(
        durations_frames: torch.Tensor,
        compression: int,
        t_tok: int,
    ) -> torch.Tensor:
        """
        Convert mel-frame durations -> token durations using boundary rounding.
        durations_frames: [B, T_ph] (mel frames)
        returns durations_tokens: [B, T_ph] (tokens)
        """
        d = durations_frames.long().clamp(min=0)
        end_f = torch.cumsum(d, dim=1)  # [B, T_ph] in frames

        # Round boundaries to token grid
        end_t = torch.round(end_f.float() / float(compression)).long()
        end_t = torch.clamp(end_t, 0, t_tok)

        start_t = torch.cat(
            [torch.zeros_like(end_t[:, :1]), end_t[:, :-1]],
            dim=1,
        )
        d_tok = (end_t - start_t).clamp(min=0)
        return d_tok

    @staticmethod
    def durations_frames_to_tokens_by_length(durations_frames, T_mel, t_tok):
        """
        Convert mel-frame durations -> token durations using mapping by length.
        Aligns boundaries proportionally between T_mel and t_tok.
        """
        d = durations_frames.long().clamp(min=0)
        end_f = torch.cumsum(d, dim=1)  # [B, T_ph] frame boundaries

        # Map frame boundary positions -> token boundary positions
        # end_t in [0, t_tok]
        end_t = torch.round(end_f.float() * float(t_tok) / float(T_mel)).long()
        end_t = end_t.clamp(0, t_tok)

        start_t = torch.cat([torch.zeros_like(end_t[:, :1]), end_t[:, :-1]], dim=1)
        return (end_t - start_t).clamp(min=0)

    @staticmethod
    def expand_phonemes_by_durations(
        phonemes: torch.Tensor,
        durations: torch.Tensor,
        target_length: int,
    ):
        """
        Expand phoneme ids by durations to token-level targets.
        phonemes:   [B, T_ph]
        durations:  [B, T_ph] in tokens (int)
        returns:
          expanded:   [B, target_length]
          valid_mask: [B, target_length] bool
        """
        durations = durations.long().clamp(min=0)

        end_idxs = torch.cumsum(durations, dim=1)
        start_idxs = end_idxs - durations

        t = torch.arange(target_length, device=phonemes.device).view(1, 1, -1)

        starts = start_idxs.unsqueeze(2)
        ends = end_idxs.unsqueeze(2)

        mask = (t >= starts) & (t < ends)  # [B, T_ph, T_target]

        expanded = (phonemes.unsqueeze(2) * mask.long()).sum(dim=1)
        valid_mask = mask.sum(dim=1) > 0

        return expanded, valid_mask

    def forward(self, z_q: torch.Tensor, lang_emb: Optional[torch.Tensor] = None):
        if self.language_emb_dim > 0:
            if lang_emb is None:
                raise ValueError("lang_emb is required (language_emb_dim > 0).")
            # [B, D] -> [B, D, T]
            lang = lang_emb.unsqueeze(-1).expand(-1, -1, z_q.size(-1))
            x = torch.cat([z_q, lang], dim=1)
        else:
            x = z_q

        return self.net(x)


class LearnableUpsample1d(nn.Module):
    """
    Learnable upsampling using transposed convolution with anti-aliasing.
    Better than nn.Upsample for preserving high-frequency details.
    """
    def __init__(self, in_channels: int, out_channels: int, scale_factor: int = 2, kernel_size: int = None):
        super().__init__()
        self.scale_factor = scale_factor
        # Kernel size should be 2x scale factor for good coverage
        kernel_size = kernel_size or scale_factor * 4
        padding = (kernel_size - scale_factor) // 2
        
        # Main upsampling via transposed conv
        self.upsample = weight_norm(nn.ConvTranspose1d(
            in_channels, out_channels,
            kernel_size=kernel_size,
            stride=scale_factor,
            padding=padding,
        ))
        
        # Anti-aliasing low-pass filter (learnable)
        self.antialiasing = nn.Sequential(
            weight_norm(nn.Conv1d(out_channels, out_channels, kernel_size=5, padding=2, groups=out_channels)),
            nn.SiLU(),
            weight_norm(nn.Conv1d(out_channels, out_channels, kernel_size=3, padding=1)),
        )
        
        # Initialize for smooth upsampling
        nn.init.kaiming_normal_(self.upsample.weight)
        if self.upsample.bias is not None:
            nn.init.zeros_(self.upsample.bias)
    
    def forward(self, x):
        x = self.upsample(x)
        x = self.antialiasing(x)
        return x


class LearnableUpsampleBlock(nn.Module):
    """
    Multi-stage learnable upsampling block.
    Replaces nn.Upsample with learnable transposed convolutions.
    """
    def __init__(self, channels: int, total_upsample: int):
        super().__init__()
        self.total_upsample = total_upsample
        
        # Decompose into 2x upsamples
        layers = []
        remaining = total_upsample
        while remaining > 1:
            factor = min(2, remaining)
            layers.append(LearnableUpsample1d(channels, channels, scale_factor=factor))
            remaining //= factor
        
        self.layers = nn.ModuleList(layers)
    
    def forward(self, x):
        for layer in self.layers:
            x = layer(x)
        return x


# ==============================================================================
# Finite Scalar Quantization
# ==============================================================================

class FiniteScalarQuantization(nn.Module):
    def __init__(self, input_dim=256, levels: List[int] = [4]*6):
        super().__init__()
        self.input_dim = input_dim
        self.levels = levels
        self.dims = len(levels)
        self.codebook_size = math.prod(levels)
        
        self.in_proj = nn.Sequential(
            nn.Linear(input_dim, input_dim // 2),
            nn.SiLU(),
            nn.Linear(input_dim // 2, self.dims),
        )
        self.out_proj = nn.Sequential(
            nn.Linear(self.dims, input_dim // 2),
            nn.SiLU(),
            nn.Linear(input_dim // 2, input_dim),
        )
        
        self.scale = nn.Parameter(torch.ones(self.dims) * 1.5)
        self.bias = nn.Parameter(torch.zeros(self.dims))
        
        for m in self.in_proj.modules():
            if isinstance(m, nn.Linear):
                nn.init.xavier_uniform_(m.weight, gain=2.0)
                if m.bias is not None:
                    nn.init.zeros_(m.bias)
        for m in self.out_proj.modules():
            if isinstance(m, nn.Linear):
                nn.init.xavier_uniform_(m.weight, gain=1.0)
                if m.bias is not None:
                    nn.init.zeros_(m.bias)
        
        self.register_buffer('levels_tensor', torch.tensor(levels, dtype=torch.float32))
        _basis = torch.cumprod(torch.tensor([1] + levels[:-1]), dim=0)
        self.register_buffer('basis', _basis)
        
        self.register_buffer('num_steps', torch.tensor(0))
        self.warmup_steps = 5000
        
    def forward(self, x, n_quantizers=None):
        x = x.transpose(1, 2)
        z = self.in_proj(x)
        
        z = z * self.scale + self.bias
        z_bound = torch.tanh(z)
        
        if self.training:
            self.num_steps += 1
            noise_scale = max(0.3 * (1 - self.num_steps.float() / self.warmup_steps), 0.05)
            noise = (torch.rand_like(z_bound) - 0.5) * 2 * noise_scale
            z_bound_noisy = z_bound + noise
            z_bound_noisy = torch.clamp(z_bound_noisy, -1, 1)
        else:
            z_bound_noisy = z_bound
        
        levels = self.levels_tensor.to(z.device)
        half_l = (levels - 1) / 2
        z_scaled = z_bound_noisy * half_l
        
        z_shifted = z_scaled + half_l
        z_ind = z_shifted.round()
        z_ind = torch.clamp(z_ind, torch.zeros_like(levels), levels - 1)
        
        z_q_target = z_ind - half_l
        z_q = z_scaled + (z_q_target - z_scaled).detach()
        
        out = self.out_proj(z_q)
        
        z_ind_long = z_ind.long()
        indices = (z_ind_long * self.basis).sum(dim=-1)
        
        out = out.transpose(1, 2)
        aux_loss = self._entropy_loss(z_shifted, levels)
        
        return out, indices.unsqueeze(1), aux_loss
    
    def _entropy_loss(self, z_shifted, levels):
        B, T, D = z_shifted.shape
        total_entropy_loss = torch.tensor(0.0, device=z_shifted.device)
        
        for d in range(D):
            vals = z_shifted[..., d].reshape(-1)
            num_levels = int(levels[d].item())
            centers = torch.arange(num_levels, device=z_shifted.device, dtype=torch.float32)
            dist = (vals.unsqueeze(1) - centers.unsqueeze(0)).pow(2)
            probs = F.softmax(-dist / 0.5, dim=1)
            avg_probs = probs.mean(dim=0)
            uniform = torch.ones_like(avg_probs) / num_levels
            kl_div = (avg_probs * (torch.log(avg_probs + 1e-7) - torch.log(uniform + 1e-7))).sum()
            total_entropy_loss = total_entropy_loss + kl_div
        
        return 0.1 * total_entropy_loss / D
        
    def decode(self, indices):
        if indices.dim() == 3: 
            indices = indices.squeeze(1)
        z_q = []
        remainder = indices
        for i in range(self.dims):
            val = remainder % self.levels[i]
            remainder = remainder // self.levels[i]
            z_q.append(val)
            
        z_q = torch.stack(z_q, dim=-1).float().to(indices.device)
        levels = self.levels_tensor.to(indices.device)
        half_l = (levels - 1) / 2
        z_q = z_q - half_l
        out = self.out_proj(z_q)
        return out.transpose(1, 2)

    # =========================================================================
    # FLOW MATCHING SUPPORT METHODS
    # =========================================================================
    
    @property
    def embed_dim(self) -> int:
        """Dimension of continuous FSQ embeddings (for flow matching)."""
        return self.dims  # 6 for [4,4,4,4,4,4]
    
    def encode_continuous(self, x: torch.Tensor) -> torch.Tensor:
        """
        Encode latent to continuous FSQ space (pre-quantization).
        This is the TARGET for flow matching training.
        
        Args:
            x: [B, input_dim, T] - latent from encoder
            
        Returns:
            z_continuous: [B, T, dims] - continuous embeddings in [-half_l, half_l]
                         For levels=[4,4,4,4,4,4], range is [-1.5, 1.5] per dim
        """
        x = x.transpose(1, 2)  # [B, T, input_dim]
        z = self.in_proj(x)    # [B, T, dims]
        
        z = z * self.scale + self.bias
        z_bound = torch.tanh(z)  # [-1, 1]
        
        levels = self.levels_tensor.to(z.device)
        half_l = (levels - 1) / 2
        z_continuous = z_bound * half_l  # [-half_l, half_l] per dim
        
        return z_continuous  # [B, T, dims]
    
    def quantize_continuous(self, z_continuous: torch.Tensor) -> torch.Tensor:
        """
        Quantize continuous FSQ embeddings to token indices.
        Use this after flow matching generates z_continuous.
        
        Args:
            z_continuous: [B, T, dims] - continuous in [-half_l, half_l]
            
        Returns:
            indices: [B, T] - token indices
        """
        levels = self.levels_tensor.to(z_continuous.device)
        half_l = (levels - 1) / 2
        
        # Shift to [0, L-1] range and round
        z_shifted = z_continuous + half_l
        z_ind = z_shifted.round()
        z_ind = torch.clamp(z_ind, torch.zeros_like(levels), levels - 1)
        
        # Convert to single index
        z_ind_long = z_ind.long()
        indices = (z_ind_long * self.basis).sum(dim=-1)
        
        return indices  # [B, T]
    
    def continuous_to_latent(self, z_continuous: torch.Tensor) -> torch.Tensor:
        """
        Convert continuous FSQ embeddings to decoder-ready latent.
        Quantizes and projects back to input_dim.
        
        Args:
            z_continuous: [B, T, dims] - from flow matching prediction
            
        Returns:
            latent: [B, input_dim, T] - ready for decoder
        """
        levels = self.levels_tensor.to(z_continuous.device)
        half_l = (levels - 1) / 2
        
        # Quantize (round to nearest level)
        z_shifted = z_continuous + half_l
        z_ind = z_shifted.round()
        z_ind = torch.clamp(z_ind, torch.zeros_like(levels), levels - 1)
        z_q = z_ind - half_l
        
        # Project back to latent dim
        out = self.out_proj(z_q)
        return out.transpose(1, 2)  # [B, input_dim, T]
    
    def indices_to_continuous(self, indices: torch.Tensor) -> torch.Tensor:
        """
        Convert token indices to continuous FSQ embeddings.
        Useful for getting GT targets from precomputed tokens.
        
        Args:
            indices: [B, T] or [B, 1, T] - token indices
            
        Returns:
            z_continuous: [B, T, dims] - continuous embeddings
        """
        if indices.dim() == 3:
            indices = indices.squeeze(1)
            
        # Decompose index into per-dimension values
        z_q = []
        remainder = indices
        for i in range(self.dims):
            val = remainder % self.levels[i]
            remainder = remainder // self.levels[i]
            z_q.append(val)
        
        z_q = torch.stack(z_q, dim=-1).float()  # [B, T, dims]
        
        levels = self.levels_tensor.to(indices.device)
        half_l = (levels - 1) / 2
        z_continuous = z_q - half_l  # Center around 0
        
        return z_continuous  # [B, T, dims]


# ==============================================================================
# Speaker-Conditioned Fusion Module with AdaIN1d
# ==============================================================================

class SpeakerFusionResBlock(nn.Module):
    """
    Fusion ResBlock conditioned on speaker embedding via AdaIN1d.
    """
    def __init__(
        self,
        dim_in,
        dim_out,
        speaker_dim=128,
        actv=nn.LeakyReLU(0.2),
        dropout_p=0.0,
    ):
        super().__init__()
        self.actv = actv
        self.learned_sc = dim_in != dim_out
        self.dropout = nn.Dropout(dropout_p)

        self.conv1 = weight_norm(nn.Conv1d(dim_in, dim_out, 3, 1, 1))
        self.conv2 = weight_norm(nn.Conv1d(dim_out, dim_out, 3, 1, 1))
        self.norm1 = AdaIN1d(speaker_dim, dim_in)
        self.norm2 = AdaIN1d(speaker_dim, dim_out)
        if self.learned_sc:
            self.conv1x1 = weight_norm(nn.Conv1d(dim_in, dim_out, 1, 1, 0, bias=False))

    def _shortcut(self, x):
        if self.learned_sc:
            x = self.conv1x1(x)
        return x

    def _residual(self, x, speaker_emb):
        x = self.norm1(x, speaker_emb)
        x = self.actv(x)
        x = self.conv1(self.dropout(x))
        x = self.norm2(x, speaker_emb)
        x = self.actv(x)
        x = self.conv2(self.dropout(x))
        return x

    def forward(self, x, speaker_emb):
        out = self._residual(x, speaker_emb)
        out = (out + self._shortcut(x)) / math.sqrt(2)
        return out


class SpeakerLatentFusionModule(nn.Module):
    """
    ResNet-style fusion module for latent only (no text), with speaker conditioning.
    Takes quantized latent and processes it with speaker AdaIN conditioning.
    """
    def __init__(self, latent_dim, hidden_dim, speaker_dim=128):
        super().__init__()
        self.input_mix = SpeakerFusionResBlock(latent_dim, hidden_dim, speaker_dim)
        
        self.decode = nn.ModuleList()
        concat_dim = hidden_dim + latent_dim
        
        self.decode.append(SpeakerFusionResBlock(concat_dim, hidden_dim, speaker_dim))
        self.decode.append(SpeakerFusionResBlock(concat_dim, hidden_dim, speaker_dim))
        self.decode.append(SpeakerFusionResBlock(concat_dim, hidden_dim, speaker_dim))
        
    def forward(self, latent, speaker_emb):
        """
        Args:
            latent: [B, latent_dim, T]
            speaker_emb: [B, speaker_dim] - global speaker embedding
        """
        x = self.input_mix(latent, speaker_emb)
        
        for block in self.decode:
            x = torch.cat([x, latent], dim=1)
            x = block(x, speaker_emb)
            
        return x


# ==============================================================================
# Waveform Decoder with Speaker Conditioning (for Mel codec)
# ==============================================================================

class MelWaveformDecoderSpeaker(nn.Module):
    """
    Waveform decoder for mel codec, conditioned on learnable speaker embeddings.
    Takes quantized latent and decodes to waveform.
    """
    def __init__(
        self,
        latent_dim: int = 512,
        speaker_dim: int = 128,
        hidden_dim: int = 512,
        upsample_rates: List[int] = [12, 10],
        resblock_kernel_sizes: List[int] = [3, 7, 11],
        resblock_dilation_sizes: List[List[int]] = [[1, 3, 5], [1, 3, 5], [1, 3, 5]],
        gen_istft_n_fft: int = 30,
        gen_istft_hop_size: int = 5,
        sample_rate: int = 44100,
        source_upsample_rate: Optional[int] = None,
        codec_strides: Optional[List[int]] = None,
    ):
        super().__init__()
        self.num_upsamples = len(upsample_rates)
        self.num_kernels = len(resblock_kernel_sizes)
        self.gen_istft_n_fft = gen_istft_n_fft
        self.gen_istft_hop_size = gen_istft_hop_size
        self.codec_strides = codec_strides or [1]
        self.codec_compression = int(np.prod(self.codec_strides))
        self.speaker_dim = speaker_dim
        
        total_upsample = int(np.prod(upsample_rates)) * gen_istft_hop_size
        self.source_upsample_rate = source_upsample_rate or total_upsample
        
        # Learnable upsampler for latent (replaces nn.Upsample)
        # Uses transposed convolutions with anti-aliasing for better quality
        self.latent_upsampler = LearnableUpsampleBlock(latent_dim, self.codec_compression)
        
        # Simple linear F0 upsampler - just interpolate + smooth with conv1d
        # No hidden dims, just direct linear processing
        self.f0_upsample_factor = self.codec_compression
        self.f0_smooth = nn.Sequential(
            weight_norm(nn.Conv1d(1, 1, kernel_size=5, padding=2)),  # Smooth after interpolation
            weight_norm(nn.Conv1d(1, 1, kernel_size=3, padding=1)),  # Final refinement
        )
        
        # F0 predictor from latent
        self.f0_predictor = nn.Sequential(
            weight_norm(nn.Conv1d(latent_dim, hidden_dim, 3, padding=1)),
            nn.SiLU(),
            weight_norm(nn.Conv1d(hidden_dim, hidden_dim, 3, padding=1)),
            nn.SiLU(),
            weight_norm(nn.Conv1d(hidden_dim, hidden_dim // 2, 3, padding=1)),
            nn.SiLU(),
            weight_norm(nn.Conv1d(hidden_dim // 2, hidden_dim // 4, 3, padding=1)),
            nn.SiLU(),
            weight_norm(nn.Conv1d(hidden_dim // 4, 1, 3, padding=1))
        )
        
        # Harmonic source module
        self.m_source = SourceModuleHnNSF(
            sampling_rate=sample_rate,
            upsample_scale=self.source_upsample_rate, 
            harmonic_num=14,
            voiced_threshold=1,
        )
        self.f0_upsamp = nn.Upsample(scale_factor=self.source_upsample_rate)

        # Speaker-conditioned pre-decoder
        self.pre_decoder = SpeakerLatentFusionModule(
            latent_dim=latent_dim,
            hidden_dim=hidden_dim,
            speaker_dim=speaker_dim
        )
        
        # Conformer layers
        self.conformers = nn.ModuleList()
        for i in range(len(upsample_rates)):
            ch = hidden_dim // (2 ** i)
            self.conformers.append(
                Conformer(
                    dim=ch,
                    depth=4,
                    dim_head=64,
                    heads=8,
                    ff_mult=4,
                    conv_expansion_factor=2,
                    conv_kernel_size=31,
                    attn_dropout=0.1,
                    ff_dropout=0.1,
                    conv_dropout=0.1,
                )
            )
        
        # Snake activations
        self.snakes = nn.ModuleList()
        self.snakes.append(Snake1d(hidden_dim))
        
        # Upsampling layers
        self.ups = nn.ModuleList()
        upsample_kernel_sizes = [2 * u for u in upsample_rates]
        for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):
            in_ch = hidden_dim // (2 ** i)
            out_ch = hidden_dim // (2 ** (i + 1))
            self.ups.append(UpsamplePixelShuffle1D(in_ch, out_ch, kernel_size=k, r=u))
            self.snakes.append(Snake1d(out_ch))
        
        # Noise injection layers
        self.noise_convs = nn.ModuleList()
        self.noise_res = nn.ModuleList()
        for i in range(len(upsample_rates)):
            c_cur = hidden_dim // (2 ** (i + 1))
            if i + 1 < len(upsample_rates):
                stride_f0 = int(np.prod(upsample_rates[i + 1:]))
                self.noise_convs.append(
                    weight_norm(nn.Conv1d(
                        gen_istft_n_fft + 2, c_cur,
                        kernel_size=stride_f0 * 2,
                        stride=stride_f0,
                        padding=(stride_f0 + 1) // 2,
                    ))
                )
                self.noise_res.append(SpeakerAdaINResBlock1(c_cur, 7, [1, 3, 5], speaker_dim))
            else:
                self.noise_convs.append(
                    weight_norm(nn.Conv1d(gen_istft_n_fft + 2, c_cur, kernel_size=1))
                )
                self.noise_res.append(SpeakerAdaINResBlock1(c_cur, 11, [1, 3, 5], speaker_dim))
        
        # ResBlocks with speaker AdaIN conditioning
        self.resblocks = nn.ModuleList()
        for i in range(len(upsample_rates)):
            ch = hidden_dim // (2 ** (i + 1))
            for k, d in zip(resblock_kernel_sizes, resblock_dilation_sizes):
                self.resblocks.append(SpeakerAdaINResBlock1(ch, k, d, speaker_dim))
        
        # Post convolution for STFT
        self.post_n_fft = gen_istft_n_fft
        final_ch = hidden_dim // (2 ** len(upsample_rates))
        self.conv_post = weight_norm(nn.Conv1d(final_ch, self.post_n_fft + 2, 7, padding=3))
        
        # STFT for inverse transform
        self.stft = TorchSTFT(
            filter_length=gen_istft_n_fft,
            hop_length=gen_istft_hop_size,
            win_length=gen_istft_n_fft,
        )
        
        self.reflection_pad = nn.ReflectionPad1d((1, 0))
        
    def forward(self, latent, speaker_emb, f0_gt=None, cache=None):
        """
        Args:
            latent: [B, latent_dim, T_comp] - quantized latent
            speaker_emb: [B, speaker_dim] - global speaker embedding
            f0_gt: [B, T] optional ground truth F0
            cache: dict for streaming inference
        
        Returns:
            wav, spec, phase, f0_pred, (new_cache if streaming)
        """
        B = latent.shape[0]
        
        # Predict F0 from latent
        f0_pred_latent = self.f0_predictor(latent)  # [B, 1, T_comp]
        
        # Simple linear upsample - just interpolate and smooth
        f0_pred = F.interpolate(f0_pred_latent, scale_factor=self.f0_upsample_factor, mode='linear', align_corners=False)
        f0_pred = self.f0_smooth(f0_pred)  # Smooth with conv1d
        
        # Match F0 to target length
        if f0_gt is not None:
            f0_pred = F.interpolate(f0_pred, size=f0_gt.shape[-1], mode='linear')
        else:
            target_len = int(latent.shape[-1] * self.codec_compression)
            if f0_pred.shape[-1] != target_len:
                f0_pred = F.interpolate(f0_pred, size=target_len, mode='linear')
            
        f0_pred = f0_pred.squeeze(1)
        f0_to_use = f0_gt if f0_gt is not None else f0_pred.detach()
        
        # Generate harmonic source
        uv = (f0_to_use > 0).float()  # f0_to_use is log10(F0), 0 = unvoiced

        f0_log_up = self.f0_upsamp(f0_to_use[:, None]).transpose(1, 2)
        uv_up = self.f0_upsamp(uv[:, None]).transpose(1, 2)
        uv_up = (uv_up > 0.5).to(f0_log_up.dtype)

        f0_lin = (10.0 ** f0_log_up.float()).to(f0_log_up.dtype) * uv_up
        
        source_phase_cache = cache.get("source_phase") if cache is not None else None
        har_source, noi_source, uv, next_source_phase = self.m_source(f0_lin, cache=source_phase_cache)
        
        har_source = har_source.transpose(1, 2).squeeze(1)
        har_spec, har_phase = self.stft.transform(har_source)
        har = torch.cat([har_spec, har_phase], dim=1)
        
        # Upsample latent
        latent_up = self.latent_upsampler(latent)
        
        # Speaker-conditioned fusion
        x = self.pre_decoder(latent_up, speaker_emb)
        
        # Upsampling with conformers and residual blocks
        for i in range(self.num_upsamples):
            x = self.snakes[i](x)
            
            x = rearrange(x, "b f t -> b t f")
            x = self.conformers[i](x)
            x = rearrange(x, "b t f -> b f t")
            
            x = self.ups[i](x)
            
            x_source = self.noise_convs[i](har)
            x_source = self.noise_res[i](x_source, speaker_emb)
            
            if i == self.num_upsamples - 1:
                x = self.reflection_pad(x)
            
            if x.shape[-1] != x_source.shape[-1]:
                min_len_add = min(x.shape[-1], x_source.shape[-1])
                x = x[..., :min_len_add]
                x_source = x_source[..., :min_len_add]
            
            x = x + x_source
            
            xs = None
            for j in range(self.num_kernels):
                if xs is None:
                    xs = self.resblocks[i * self.num_kernels + j](x, speaker_emb)
                else:
                    xs += self.resblocks[i * self.num_kernels + j](x, speaker_emb)
            x = xs / self.num_kernels
        
        x = self.snakes[-1](x)
        x = self.conv_post(x)
        
        spec = torch.exp(x[:, :self.post_n_fft // 2 + 1, :])
        phase = torch.sin(x[:, self.post_n_fft // 2 + 1:, :])
        
        out = self.stft.inverse(spec, phase)
        
        if cache is not None:
            new_cache = {
                "source_phase": next_source_phase
            }
            return out, spec, phase, f0_pred, new_cache
            
        return out, spec, phase, f0_pred


# ==============================================================================
# Main Mel Codec with Learnable Speaker Embeddings
# ==============================================================================

class MelCodecVocoderSpeaker(nn.Module):
    """
    Mel Spectrogram Codec with LEARNABLE SPEAKER EMBEDDINGS.
    
    Key features:
    - Input: mel spectrogram + pitch (instead of text_emb + energy)
    - Learnable speaker embedding: nn.Embedding(num_speakers, speaker_dim)
    - Speaker conditioning via AdaIN1d throughout the decoder
    - F0 prediction from latent
    
    Usage:
        model = MelCodecVocoderSpeaker(num_speakers=11, speaker_dim=128, ...)
        output = model(mel, pitch, speaker_ids=speaker_ids)
    """
    def __init__(
        self,
        num_speakers: int = 11,
        speaker_dim: int = 128,
        n_mels: int = 80,
        latent_dim: int = 512,
        hidden_dim: int = 512,
        codec_strides: List[int] = [2, 2],
        codebook_size: int = 4096,
        upsample_rates: List[int] = [9, 7],
        gen_istft_n_fft: int = 30,
        gen_istft_hop_size: int = 5,
        sample_rate: int = 44100,
        source_upsample_rate: int = 441,
        fsq_levels: Optional[List[int]] = None,
        # Phoneme predictor settings
        n_phonemes: int = 178,
        num_languages: int = 10,
        language_dim: int = 64,
    ):
        super().__init__()
        self.num_speakers = num_speakers
        self.speaker_dim = speaker_dim
        self.n_mels = n_mels
        self.latent_dim = latent_dim
        self.codec_compression = math.prod(codec_strides)
        self.use_fsq = fsq_levels is not None
        self.fsq_levels = fsq_levels or [4] * 6
        

        self.speaker_emb_dim = 64
        self.language_emb_dim = 64
        
        self.speaker_embedding = nn.Embedding(
            num_embeddings=num_speakers,
            embedding_dim=self.speaker_emb_dim
        )
        nn.init.normal_(self.speaker_embedding.weight, mean=0, std=0.5)

        self.language_embedding = nn.Embedding(
            num_embeddings=num_languages,
            embedding_dim=self.language_emb_dim
        )
        nn.init.normal_(self.language_embedding.weight, mean=0, std=0.5)
        
        # Mel + Pitch encoder
        self.encoder = MelPitchEncoderSpeaker(
            n_mels=n_mels,
            speaker_dim=self.speaker_emb_dim + self.language_emb_dim, # Combined dim (128)
            latent_dim=latent_dim,
            hidden_dim=hidden_dim,
            strides=codec_strides,
        )
        
        # FSQ Quantizer
        self.quantizer = FiniteScalarQuantization(
            input_dim=latent_dim,
            levels=self.fsq_levels,
        )
        
        # Decoder
        self.decoder = MelWaveformDecoderSpeaker(
            latent_dim=latent_dim,
            speaker_dim=self.speaker_emb_dim + self.language_emb_dim, # Combined dim (128)
            hidden_dim=hidden_dim,
            upsample_rates=upsample_rates,
            gen_istft_n_fft=gen_istft_n_fft,
            gen_istft_hop_size=gen_istft_hop_size,
            sample_rate=sample_rate,
            source_upsample_rate=source_upsample_rate,
            codec_strides=codec_strides,
        )
        
        # =====================================================================
        # PHONEME PREDICTOR for linguistic supervision
        # =====================================================================
        self.phoneme_predictor = TokenRatePhonemePredictor(
            latent_dim=latent_dim,
            n_phonemes=n_phonemes,
            language_emb_dim=self.language_emb_dim,  # 64
            hidden_dim=hidden_dim // 2,
            dropout_p=0.1,
        )
        
    def forward(self, mel, pitch, speaker_ids, language_ids=None, n_quantizers=None, use_predicted_f0=False,
                target_phonemes=None, target_durations=None):
        """
        Training forward pass.
        
        Args:
            mel: [B, n_mels, T] - mel spectrogram
            pitch: [B, T] - pitch contour (log F0)
            speaker_ids: [B] - speaker IDs
            language_ids: [B] - language IDs (for phoneme prediction)
            n_quantizers: unused, for compatibility
            use_predicted_f0: bool - whether to use predicted F0
            target_phonemes: [B, T_ph] - optional, target phoneme IDs for loss
            target_durations: [B, T_ph] - optional, target durations for loss
            
        Returns:
            dict with wav, tokens, speaker_emb, f0_pred, phoneme_logits, etc.
        """
        # Get speaker embedding from ID
        s_emb = self.speaker_embedding(speaker_ids)  # [B, 64]
        
        if language_ids is None:
            # Try to infer from device/shape if possible or default to 0? 
            # Ideally we raise error, but to be robust:
            # raise ValueError("language_ids must be provided for style construction")
             # Assuming language_ids are required.
             raise ValueError("language_ids are required for concatenation.")

        l_emb = self.language_embedding(language_ids) # [B, 64]
        
        # Concatenate for Style [B, 128]
        speaker_emb = torch.cat([s_emb, l_emb], dim=1)
        
        # Encode mel + pitch
        latent = self.encoder(mel, pitch)
        
        # Quantize
        quantized_latent, tokens, commitment_loss = self.quantizer(latent)
        
        # Decide whether to use GT or predicted F0
        decoder_f0 = None if use_predicted_f0 else pitch

        # Decode with speaker conditioning
        wav, mag, phase, f0_pred = self.decoder(
            quantized_latent,
            speaker_emb,
            f0_gt=decoder_f0,
            cache=None,
        )
        
        # Phoneme prediction from latent (for linguistic supervision)
        phoneme_logits = None
        loss_phoneme = None
        
        if language_ids is not None:
            # Predict at frame rate (upsampled latent)
            # quantized_latent: [B, C, T_tok]
            T_mel = pitch.size(-1)
            
            # Upsample latent to frame resolution (nearest neighbor - non-learnable)
            quantized_latent_upsampled = F.interpolate(quantized_latent, size=T_mel, mode='nearest')

            phoneme_logits = self.phoneme_predictor(quantized_latent_upsampled, l_emb.detach())  # [B, P, T_mel]

            if target_phonemes is not None and target_durations is not None:
                # Expand phonemes to frame-level targets using frame durations directly
                targets_frames, valid_mask_frames = TokenRatePhonemePredictor.expand_phonemes_by_durations(
                    phonemes=target_phonemes.long(),
                    durations=target_durations,
                    target_length=T_mel,
                )

                # Masked CE over valid frames
                # NOTE: do NOT ignore_index=0 because 0 is a real phoneme in your vocab.
                logits_flat = phoneme_logits.transpose(1, 2).reshape(-1, phoneme_logits.size(1))
                targets_flat = targets_frames.reshape(-1)
                mask_flat = valid_mask_frames.reshape(-1)

                if mask_flat.any():
                    loss_phoneme = F.cross_entropy(logits_flat[mask_flat], targets_flat[mask_flat])
                else:
                    loss_phoneme = torch.tensor(0.0, device=mel.device)
        
        return {
            "wav": wav,
            "mag": mag,
            "phase": phase,
            "tokens": tokens,
            "latent": latent,
            "quantized_latent": quantized_latent,
            "speaker_emb": speaker_emb,
            "commitment_loss": commitment_loss,
            "f0_pred": f0_pred,
            "f0_gt": pitch,
            "phoneme_logits": phoneme_logits,
            "phoneme_loss": loss_phoneme,
        }
    
    def get_speaker_embedding(self, speaker_ids):
        """Get speaker embedding from IDs (only speaker part)."""
        return self.speaker_embedding(speaker_ids)
        
    def get_style_embedding(self, speaker_ids, language_ids):
        """Get full style embedding (speaker + language)."""
        s_emb = self.speaker_embedding(speaker_ids)
        l_emb = self.language_embedding(language_ids)
        return torch.cat([s_emb, l_emb], dim=1)
    
    @torch.no_grad()
    def tokenize(self, mel, pitch, speaker_ids, language_ids=None, n_quantizers=None):
        """Tokenize mel + pitch."""
        s_emb = self.speaker_embedding(speaker_ids)
        
        if language_ids is not None:
            l_emb = self.language_embedding(language_ids)
            style_emb = torch.cat([s_emb, l_emb], dim=1)
        else:
            # Return partial if language not provided (might break if used for decoding)
            style_emb = s_emb 
            
        latent = self.encoder(mel, pitch)
        _, tokens, _ = self.quantizer(latent)
        return tokens, style_emb
    
    # =========================================================================
    # FLOW MATCHING SUPPORT
    # =========================================================================
    
    @torch.no_grad()
    def encode_for_flow_matching(self, mel, pitch):
        """
        Encode mel+pitch to continuous FSQ embeddings for flow matching training.
        
        Args:
            mel: [B, n_mels, T] - mel spectrogram
            pitch: [B, T] - pitch contour (log F0)
            
        Returns:
            z_continuous: [B, T_codec, fsq_dims] - continuous FSQ embeddings
                         These are the TARGETS for flow matching.
                         fsq_dims = 6 for levels=[4,4,4,4,4,4]
                         Range: [-1.5, 1.5] per dimension
        """
        latent = self.encoder(mel, pitch)  # [B, latent_dim, T_codec]
        z_continuous = self.quantizer.encode_continuous(latent)  # [B, T_codec, fsq_dims]
        return z_continuous
    
    @torch.no_grad()
    def tokenize_with_continuous(self, mel, pitch, speaker_ids, language_ids=None):
        """
        Tokenize and also return continuous FSQ embeddings.
        
        Returns:
            tokens: [B, 1, T_codec] - discrete token indices
            z_continuous: [B, T_codec, fsq_dims] - continuous FSQ embeddings
            style_emb: [B, style_dim] - combined speaker+language embedding
        """
        s_emb = self.speaker_embedding(speaker_ids)
        
        if language_ids is not None:
            l_emb = self.language_embedding(language_ids)
            style_emb = torch.cat([s_emb, l_emb], dim=1)
        else:
            style_emb = s_emb
            
        latent = self.encoder(mel, pitch)
        z_continuous = self.quantizer.encode_continuous(latent)
        _, tokens, _ = self.quantizer(latent)
        
        return tokens, z_continuous, style_emb
    
    def decode_from_continuous(self, z_continuous, speaker_ids, language_ids, f0=None):
        """
        Decode from continuous FSQ embeddings (flow matching output).
        
        Args:
            z_continuous: [B, T_codec, fsq_dims] - from flow matching
            speaker_ids: [B] - speaker IDs
            language_ids: [B] - language IDs
            f0: [B, T] optional F0
            
        Returns:
            wav: [B, 1, T_audio] - waveform
            tokens: [B, T_codec] - quantized token indices
            f0_pred: [B, T] - predicted F0
        """
        s_emb = self.speaker_embedding(speaker_ids)
        l_emb = self.language_embedding(language_ids)
        speaker_emb = torch.cat([s_emb, l_emb], dim=1)
        
        # Quantize continuous to tokens
        tokens = self.quantizer.quantize_continuous(z_continuous)
        
        # Convert to decoder latent
        quantized_latent = self.quantizer.continuous_to_latent(z_continuous)
        
        wav, _, _, f0_pred = self.decoder(
            quantized_latent,
            speaker_emb,
            f0_gt=f0,
            cache=None,
        )
        return wav, tokens, f0_pred
    
    @property
    def fsq_embed_dim(self) -> int:
        """Dimension of continuous FSQ embeddings (for flow matching)."""
        return self.quantizer.embed_dim

    @torch.no_grad()
    def decode_tokens(self, tokens, speaker_ids, language_ids, f0=None):
        """
        Decode tokens with speaker ID and Language ID.
        
        Args:
            tokens: [B, 1, T_comp] - tokens
            speaker_ids: [B] - speaker IDs
            language_ids: [B] - language IDs
            f0: [B, T] optional F0 to condition on
        """
        s_emb = self.speaker_embedding(speaker_ids)
        l_emb = self.language_embedding(language_ids)
        speaker_emb = torch.cat([s_emb, l_emb], dim=1)
        
        quantized_latent = self.quantizer.decode(tokens)
        
        wav, _, _, f0_pred = self.decoder(
            quantized_latent,
            speaker_emb,
            f0_gt=f0,
            cache=None,
        )
        return wav, f0_pred

    @torch.no_grad()
    def decode_tokens_with_speaker_emb(self, tokens, speaker_emb, f0=None):
        """
        Decode tokens with pre-computed speaker embedding.
        Useful for speaker interpolation.
        
        Args:
            tokens: [B, 1, T_comp] - tokens
            speaker_emb: [B, speaker_dim] - speaker embedding
            f0: [B, T] optional F0
        """
        quantized_latent = self.quantizer.decode(tokens)
        
        wav, _, _, f0_pred = self.decoder(
            quantized_latent,
            speaker_emb,
            f0_gt=f0,
            cache=None,
        )
        return wav, f0_pred

    @torch.no_grad()
    def decode_chunk(self, tokens, speaker_ids, language_ids, cache=None, f0=None):
        """
        Streaming inference by chunk.
        
        Args:
            tokens: Chunk of tokens
            speaker_ids: [B] speaker IDs
            language_ids: [B] language IDs
            cache: Dictionary from previous chunk call
            f0: [B, T] optional F0
        
        Returns:
            wav_chunk, f0_pred, new_cache
        """
        if cache is None:
            cache = {}
            
        s_emb = self.speaker_embedding(speaker_ids)
        l_emb = self.language_embedding(language_ids)
        speaker_emb = torch.cat([s_emb, l_emb], dim=1)
        
        quantized_latent = self.quantizer.decode(tokens)

        wav, _, _, f0_pred, new_cache = self.decoder(
            quantized_latent,
            speaker_emb,
            f0_gt=f0,
            cache=cache,
        )
        
        return wav, f0_pred, new_cache
    
    @torch.no_grad()
    def interpolate_speakers(self, speaker_id_1, speaker_id_2, alpha=0.5):
        """
        Interpolate between two speaker embeddings.
        
        Args:
            speaker_id_1: int - first speaker ID
            speaker_id_2: int - second speaker ID
            alpha: float - interpolation weight (0 = speaker_1, 1 = speaker_2)
            
        Returns:
            [1, speaker_dim] interpolated embedding
        """
        emb1 = self.speaker_embedding(torch.tensor([speaker_id_1], device=self.speaker_embedding.weight.device))
        emb2 = self.speaker_embedding(torch.tensor([speaker_id_2], device=self.speaker_embedding.weight.device))
        return (1 - alpha) * emb1 + alpha * emb2
    
    @torch.no_grad()
    def encode_and_reconstruct(self, mel, pitch, speaker_ids, language_ids, use_predicted_f0=False):
        """
        Encode mel+pitch and reconstruct waveform.
        Useful for testing reconstruction quality.
        """
        output = self.forward(mel, pitch, speaker_ids, language_ids, use_predicted_f0=use_predicted_f0)
        return output['wav'], output['f0_pred'], output['tokens']