Modules/codec_decoder_speaker.py

"""
Codec Decoder with Learnable Speaker Embeddings.

This module extends the hybrid temporal codec with:
- Learnable speaker embeddings: nn.Embedding(num_speakers=11, embedding_dim=128)
- Speaker conditioning via AdaIN1d (like ringformer.py)
- Speaker IDs 0-10 for 11 speakers

Usage:
    model = HybridTTSCodecVocoderSpeaker(...)
    output = model(pitch, energy, text_emb, mel, speaker_ids=speaker_ids)
"""

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

# ==============================================================================
# Utility modules
# ==============================================================================

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.
    Follows the ringformer.py implementation.
    
    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)


# ==============================================================================
# Hybrid Prosody Encoder (No style - forces codebook usage)
# ==============================================================================

class HybridProsodyEncoderSpeaker(nn.Module):
    def __init__(
        self,
        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.compression_ratio = int(np.prod(strides))

        self.pitch_down = nn.Sequential(
            weight_norm(nn.Conv1d(1, hidden_dim, 7, stride=2, padding=3)),
            nn.SiLU(),
            weight_norm(nn.Conv1d(hidden_dim, hidden_dim, 5, stride=1, padding=2)),
            nn.SiLU(),
            weight_norm(nn.Conv1d(hidden_dim, hidden_dim, 3, stride=1, padding=1)),
            nn.SiLU(),
        )
        self.energy_down = nn.Sequential(
            weight_norm(nn.Conv1d(1, hidden_dim, 7, stride=2, padding=3)),
            nn.SiLU(),
            weight_norm(nn.Conv1d(hidden_dim, hidden_dim, 5, stride=1, padding=2)),
            nn.SiLU(),
            weight_norm(nn.Conv1d(hidden_dim, hidden_dim, 3, stride=1, padding=1)),
            nn.SiLU(),
        )

        input_dim = hidden_dim * 2
        self.fusion = nn.Sequential(
            weight_norm(nn.Conv1d(input_dim, hidden_dim, 7, padding=3)),
            nn.SiLU(),
            weight_norm(nn.Conv1d(hidden_dim, hidden_dim, 5, padding=2)),
            nn.SiLU(),
            weight_norm(nn.Conv1d(hidden_dim, hidden_dim, 3, padding=1)),
            nn.SiLU(),
        )

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

        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, pitch, energy):
        """Encode pitch + energy into prosody latent. No speaker here - forces codebook usage."""
        pitch_feat = self.pitch_down(pitch.unsqueeze(1))
        energy_feat = self.energy_down(energy.unsqueeze(1))
        min_len = min(pitch_feat.shape[-1], energy_feat.shape[-1])
        pitch_feat = pitch_feat[..., :min_len]
        energy_feat = energy_feat[..., :min_len]
        x = torch.cat([pitch_feat, energy_feat], dim=1)
        x = self.fusion(x)
        x = self.refine(x)
        return self.to_latent(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)


# ==============================================================================
# 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))
        # AdaIN1d with speaker embedding (global, not temporal)
        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 SpeakerFusionModule(nn.Module):
    """
    ResNet-style fusion module with speaker conditioning via AdaIN1d.
    """
    def __init__(self, dim_in, hidden_dim, speaker_dim=128):
        super().__init__()
        self.input_mix = SpeakerFusionResBlock(dim_in, hidden_dim, speaker_dim)
        
        self.decode = nn.ModuleList()
        concat_dim = hidden_dim + dim_in
        
        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, prosody_latent, text_emb, speaker_emb, language_emb=None):
        """
        Args:
            prosody_latent: [B, prosody_dim, T]
            text_emb: [B, text_dim, T]
            speaker_emb: [B, speaker_dim] - global speaker embedding
            language_emb: [B, language_dim] optional
        """
        if language_emb is not None:
            language_emb_expanded = language_emb.unsqueeze(-1).expand(-1, -1, prosody_latent.shape[-1])
            fused = torch.cat([prosody_latent, text_emb, language_emb_expanded], dim=1)
        else:
            fused = torch.cat([prosody_latent, text_emb], dim=1)
             
        x = self.input_mix(fused, speaker_emb)
        
        for block in self.decode:
            x = torch.cat([x, fused], dim=1)
            x = block(x, speaker_emb)
            
        return x


# ==============================================================================
# Waveform Decoder with Speaker Conditioning
# ==============================================================================

class HybridWaveformDecoderSpeaker(nn.Module):
    """
    Waveform decoder conditioned on learnable speaker embeddings via AdaIN1d.
    """
    def __init__(
        self,
        prosody_latent_dim: int = 256,
        text_dim: int = 512,
        speaker_dim: int = 128,
        language_dim: int = 0,
        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
        
        self.prosody_upsampler = nn.Sequential(
            nn.Upsample(scale_factor=2, mode='linear', align_corners=False),
            weight_norm(nn.Conv1d(prosody_latent_dim, prosody_latent_dim, 3, stride=1, padding=1)),
            nn.SiLU(),
        )
        
        self.f0_upsampler = nn.Sequential(
            nn.Upsample(scale_factor=2, mode='linear', align_corners=False),
            weight_norm(nn.Conv1d(1, 1, 3, stride=1, padding=1)),
        )
        
        self.f0_predictor = nn.Sequential(
            weight_norm(nn.Conv1d(prosody_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))
        )
        
        self.m_source = SourceModuleHnNSF(
            sampling_rate=sample_rate,
            upsample_scale=self.source_upsample_rate, 
            harmonic_num=14,
            voiced_threshold=0,
        )
        self.f0_upsamp = nn.Upsample(scale_factor=self.source_upsample_rate)

        self.language_dim = language_dim
        
        fusion_dim = prosody_latent_dim + text_dim + language_dim
        
        # Speaker-conditioned fusion module
        self.pre_decoder = SpeakerFusionModule(
            dim_in=fusion_dim,
            hidden_dim=hidden_dim,
            speaker_dim=speaker_dim
        )
        
        self.conformers = nn.ModuleList()
        for i in range(len(upsample_rates)):
            ch = hidden_dim // (2 ** i)
            self.conformers.append(
                Conformer(
                    dim=ch,
                    depth=2,
                    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,
                )
            )
        
        self.snakes = nn.ModuleList()
        self.snakes.append(Snake1d(hidden_dim))
        
        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))
        
        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))
        
        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))
        
        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, prosody_latent, text_emb, speaker_emb, f0_gt=None, cache=None, language_emb=None):
        """
        Args:
            prosody_latent: [B, prosody_dim, T_comp]
            text_emb: [B, text_dim, T]
            speaker_emb: [B, speaker_dim] - global speaker embedding
            f0_gt: [B, T] optional ground truth F0
            cache: dict for streaming inference
            language_emb: [B, language_dim] optional
        
        Returns:
            wav, spec, phase, f0_pred, (new_cache if streaming)
        """
        B = prosody_latent.shape[0]
        
        f0_pred_latent = self.f0_predictor(prosody_latent)
        f0_pred = self.f0_upsampler(f0_pred_latent)
        
        if f0_gt is not None:
            f0_pred = F.interpolate(f0_pred, size=f0_gt.shape[-1], mode='linear')
        else:
            target_len = int(prosody_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
        f0_log = self.f0_upsamp(f0_to_use[:, None]).transpose(1, 2)
        f0_lin = (10.0 ** f0_log.float()).to(f0_log.dtype)
        
        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 prosody to match text
        prosody_latent = self.prosody_upsampler(prosody_latent)

        if prosody_latent.shape[-1] < text_emb.shape[-1]:
            pad_amount = text_emb.shape[-1] - prosody_latent.shape[-1]
            prosody_latent = F.pad(prosody_latent, (0, pad_amount), mode='replicate')
            
        prosody_latent = prosody_latent[..., :text_emb.shape[-1]]
        text_emb = text_emb[..., :prosody_latent.shape[-1]]
        
        # Speaker-conditioned fusion
        x = self.pre_decoder(prosody_latent, text_emb, speaker_emb, language_emb)
        
        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 Codec with Learnable Speaker Embeddings
# ==============================================================================

class HybridTTSCodecVocoderSpeaker(nn.Module):
    """
    Hybrid TTS Codec with LEARNABLE SPEAKER EMBEDDINGS.
    
    Key features:
    - Learnable speaker embedding: nn.Embedding(num_speakers=11, embedding_dim=128)
    - Speaker IDs: 0-10 for 11 speakers
    - Speaker conditioning via AdaIN1d throughout the decoder
    - No mel-based style encoder - purely speaker ID based
    
    Usage:
        model = HybridTTSCodecVocoderSpeaker(num_speakers=11, speaker_dim=128, ...)
        output = model(pitch, energy, text_emb, speaker_ids=speaker_ids)
        
        # speaker_ids: [B] tensor with values 0-10
    """
    def __init__(
        self,
        num_speakers: int = 11,
        speaker_dim: int = 128,
        text_dim: int = 512,
        prosody_latent_dim: int = 512,
        hidden_dim: int = 512,
        codec_strides: List[int] = [2, 2],
        codebook_size: int = 4096,
        upsample_rates: List[int] = [12, 10],
        gen_istft_n_fft: int = 30,
        gen_istft_hop_size: int = 5,
        sample_rate: int = 44100,
        source_upsample_rate: int = 600,
        fsq_levels: Optional[List[int]] = None,
        language_dim: int = 0,
    ):
        super().__init__()
        self.num_speakers = num_speakers
        self.speaker_dim = speaker_dim
        self.text_dim = text_dim
        self.prosody_latent_dim = prosody_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.language_dim = language_dim
        
        # =====================================================================
        # LEARNABLE SPEAKER EMBEDDING
        # 11 speakers (IDs 0-10), each with 128-dim embedding
        # =====================================================================
        self.speaker_embedding = nn.Embedding(
            num_embeddings=num_speakers,
            embedding_dim=speaker_dim
        )
        # Initialize with normal distribution
        nn.init.normal_(self.speaker_embedding.weight, mean=0, std=0.5)
        
        self.prosody_encoder = HybridProsodyEncoderSpeaker(
            speaker_dim=speaker_dim,
            latent_dim=prosody_latent_dim,
            hidden_dim=hidden_dim,
            strides=codec_strides,
        )
        
        self.quantizer = FiniteScalarQuantization(
            input_dim=prosody_latent_dim,
            levels=self.fsq_levels,
        )
        
        self.decoder = HybridWaveformDecoderSpeaker(
            prosody_latent_dim=prosody_latent_dim,
            text_dim=text_dim,
            speaker_dim=speaker_dim,
            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,
            language_dim=language_dim,
        )
        
    def forward(self, pitch, energy, text_emb, speaker_ids, n_quantizers=None, use_predicted_f0=False, language_emb=None):
        """
        Training forward pass.
        
        Args:
            pitch: [B, T] - pitch contour (log F0)
            energy: [B, T] - energy contour
            text_emb: [B, text_dim, T] - text embeddings
            speaker_ids: [B] - speaker IDs (0-10 for 11 speakers)
            n_quantizers: unused, for compatibility
            use_predicted_f0: bool - whether to use predicted F0
            language_emb: [B, language_dim] optional
            
        Returns:
            dict with wav, tokens, speaker_emb, etc.
        """
        # Get speaker embedding from ID
        speaker_emb = self.speaker_embedding(speaker_ids)  # [B, speaker_dim]
        
        # Prosody encoder (no speaker - forces codebook usage)
        prosody_latent = self.prosody_encoder(pitch, energy)
        
        # Quantize prosody
        quantized_prosody, tokens, commitment_loss = self.quantizer(prosody_latent)
        
        decoder_f0 = None if use_predicted_f0 else pitch

        # Decode with speaker conditioning via AdaIN1d
        wav, mag, phase, f0_pred = self.decoder(
            quantized_prosody,
            text_emb,
            speaker_emb,  # Speaker embedding passed to decoder
            f0_gt=decoder_f0,
            cache=None,
            language_emb=language_emb
        )
        
        return {
            "wav": wav,
            "mag": mag,
            "phase": phase,
            "tokens": tokens,
            "prosody_latent": prosody_latent,
            "quantized_prosody": quantized_prosody,
            "text_down": text_emb,
            "speaker_emb": speaker_emb,
            "commitment_loss": commitment_loss,
            "f0_pred": f0_pred,
            "f0_gt": pitch,
        }
    
    def get_speaker_embedding(self, speaker_ids):
        """Get speaker embedding from IDs."""
        return self.speaker_embedding(speaker_ids)
    
    @torch.no_grad()
    def tokenize(self, pitch, energy, text_emb, speaker_ids, n_quantizers=None):
        """Tokenize prosody."""
        speaker_emb = self.speaker_embedding(speaker_ids)
        prosody_latent = self.prosody_encoder(pitch, energy)
        _, tokens, _ = self.quantizer(prosody_latent)
        return tokens, text_emb, speaker_emb
    
    @torch.no_grad()
    def decode_tokens(self, tokens, text_emb, speaker_ids, language_emb=None):
        """
        Decode tokens with speaker ID.
        
        Args:
            tokens: [B, 1, T_comp] - prosody tokens
            text_emb: [B, text_dim, T] - text embeddings  
            speaker_ids: [B] - speaker IDs (0-10)
            language_emb: optional
        """
        speaker_emb = self.speaker_embedding(speaker_ids)
        quantized_prosody = self.quantizer.decode(tokens)
        
        wav, _, _, _ = self.decoder(
            quantized_prosody,
            text_emb,
            speaker_emb,
            f0_gt=None,
            cache=None,
            language_emb=language_emb
        )
        return wav

    @torch.no_grad()
    def decode_tokens_with_speaker_emb(self, tokens, text_emb, speaker_emb, language_emb=None):
        """
        Decode tokens with pre-computed speaker embedding.
        Useful for speaker interpolation or external speaker embeddings.
        
        Args:
            tokens: [B, 1, T_comp] - prosody tokens
            text_emb: [B, text_dim, T] - text embeddings
            speaker_emb: [B, speaker_dim] - speaker embedding (can be interpolated)
            language_emb: optional
        """
        quantized_prosody = self.quantizer.decode(tokens)
        
        wav, _, _, _ = self.decoder(
            quantized_prosody,
            text_emb,
            speaker_emb,
            f0_gt=None,
            cache=None,
            language_emb=language_emb
        )
        return wav

    @torch.no_grad()
    def decode_chunk(self, tokens, text_emb, speaker_ids, cache=None, language_emb=None):
        """
        Streaming inference by chunk.
        
        Args:
            tokens: Chunk of tokens
            text_emb: Chunk of text embeddings
            speaker_ids: [B] speaker IDs
            cache: Dictionary from previous chunk call
        
        Returns:
            wav_chunk, new_cache
        """
        if cache is None:
            cache = {}
            
        speaker_emb = self.speaker_embedding(speaker_ids)
        quantized_prosody = self.quantizer.decode(tokens)

        wav, _, _, _, new_cache = self.decoder(
            quantized_prosody,
            text_emb,
            speaker_emb,
            f0_gt=None,
            cache=cache,
            language_emb=language_emb
        )
        
        return wav, 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


# ==============================================================================
# Example Usage
# ==============================================================================

if __name__ == "__main__":
    # Test the model
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    
    # Create model
    model = HybridTTSCodecVocoderSpeaker(
        num_speakers=11,
        speaker_dim=128,
        text_dim=512,
        prosody_latent_dim=512,
        hidden_dim=512,
        codec_strides=[2, 2],
        upsample_rates=[12, 10],
        gen_istft_n_fft=30,
        gen_istft_hop_size=5,
        sample_rate=44100,
        source_upsample_rate=600,
        fsq_levels=[4, 4, 4, 4, 4, 4],
        language_dim=0,
    ).to(device)
    
    print(f"Model created with {model.num_speakers} speakers, {model.speaker_dim}-dim embeddings")
    print(f"Speaker embedding shape: {model.speaker_embedding.weight.shape}")
    
    # Test forward pass
    batch_size = 2
    seq_len = 100
    
    pitch = torch.randn(batch_size, seq_len).to(device)
    energy = torch.randn(batch_size, seq_len).to(device)
    text_emb = torch.randn(batch_size, 512, seq_len * 2).to(device)
    speaker_ids = torch.randint(0, 11, (batch_size,)).to(device)  # Random speaker IDs 0-10
    
    print(f"\nTest inputs:")
    print(f"  pitch: {pitch.shape}")
    print(f"  energy: {energy.shape}")
    print(f"  text_emb: {text_emb.shape}")
    print(f"  speaker_ids: {speaker_ids}")
    
    # Forward pass
    output = model(pitch, energy, text_emb, speaker_ids)
    
    print(f"\nOutputs:")
    print(f"  wav: {output['wav'].shape}")
    print(f"  tokens: {output['tokens'].shape}")
    print(f"  speaker_emb: {output['speaker_emb'].shape}")
    print(f"  f0_pred: {output['f0_pred'].shape}")
    
    # Test speaker interpolation
    interp_emb = model.interpolate_speakers(0, 5, alpha=0.5)
    print(f"\nInterpolated speaker embedding (0 <-> 5): {interp_emb.shape}")
    
    # Count parameters
    total_params = sum(p.numel() for p in model.parameters())
    speaker_params = model.speaker_embedding.weight.numel()
    print(f"\nTotal parameters: {total_params:,}")
    print(f"Speaker embedding parameters: {speaker_params:,} ({speaker_params/total_params*100:.2f}%)")