Modules/codec_decoder_hybrid_temporal.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

# ==============================================================================
# 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):
    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):
        h = self.fc(s).view(s.size(0), -1, 1)
        gamma, beta = torch.chunk(h, 2, dim=1)
        return (1 + gamma) * self.norm(x) + beta


class TemporalAdaIN1d(nn.Module):
    """AdaIN conditioning with temporal style [B, T_style, style_dim]."""
    def __init__(self, style_dim, num_features):
        super().__init__()
        self.norm = nn.InstanceNorm1d(num_features, affine=False)
        self.fc = weight_norm(nn.Conv1d(style_dim, num_features * 2, 1))

    def forward(self, x, s):
        """
        x: [B, C, T]
        s: [B, T_style, style_dim] or [B, style_dim, T_style]
        """
        # Ensure s is [B, style_dim, T_style]
        if s.dim() == 2:
            s = s.unsqueeze(-1)
        elif s.shape[1] != self.fc.weight.shape[1]:
            # s is [B, T, D], transpose to [B, D, T]
            s = s.transpose(1, 2)
        
        # Interpolate style to match x's temporal resolution
        if s.shape[-1] != x.shape[-1]:
            s = F.interpolate(s, size=x.shape[-1], mode='linear', align_corners=False)
        
        h = self.fc(s)  # [B, C*2, T]
        gamma, beta = torch.chunk(h, 2, dim=1)  # Each [B, C, T]
        
        return (1 + gamma) * self.norm(x) + beta


class AdaINResBlock1(nn.Module):
    """Residual block with AdaIN style conditioning."""
    def __init__(self, channels, kernel_size=3, dilation=(1, 3, 5), style_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(style_dim, channels) for _ in dilation])
        self.adain2 = nn.ModuleList([AdaIN1d(style_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, s):
        for c1, c2, n1, n2, s1, s2 in zip(
            self.convs1, self.convs2, self.adain1, self.adain2, self.snakes1, self.snakes2
        ):
            xt = n1(x, s)
            xt = s1(xt)
            xt = c1(xt)
            xt = n2(xt, s)
            xt = s2(xt)
            xt = c2(xt)
            x = xt + x
        return x


class TemporalAdaINResBlock1(nn.Module):
    """Residual block with temporal AdaIN style conditioning."""
    def __init__(self, channels, kernel_size=3, dilation=(1, 3, 5), style_dim=64):
        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([TemporalAdaIN1d(style_dim, channels) for _ in dilation])
        self.adain2 = nn.ModuleList([TemporalAdaIN1d(style_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, s):
        """
        x: [B, C, T]
        s: [B, T_style, style_dim] temporal style
        """
        for c1, c2, n1, n2, s1, s2 in zip(
            self.convs1, self.convs2, self.adain1, self.adain2, self.snakes1, self.snakes2
        ):
            xt = n1(x, s)
            xt = s1(xt)
            xt = c1(xt)
            xt = n2(xt, s)
            xt = s2(xt)
            xt = c2(xt)
            x = xt + x
        return x


# ==============================================================================
# Harmonic Source Module (Updated for Streaming)
# ==============================================================================

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):
        """
        f0_values: [B, 1, T] (or similar, depends on caller)
        initial_phase: [B, dim, 1] Phase from the end of previous chunk.
        """
        # Calculate phase increments (radians per sample in original time)
        rad_values = (f0_values / self.sampling_rate) % 1
        
        # Add random initial phase offset for the very first chunk only (usually 0th element)
        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

        # Upsample the phase increments to match waveform resolution
        rad_values = F.interpolate(
            rad_values.transpose(1, 2),
            scale_factor=1 / self.upsample_scale,
            mode="linear",
        ).transpose(1, 2)

        # Integrate to get phase: phi[t] = phi[t-1] + omega[t]
        # cumsum calculates the accumulation of phase increments for this chunk
        phase = torch.cumsum(rad_values, dim=1) * 2 * np.pi
        
        # [CHANGE] Add initial phase from previous chunk if it exists (for streaming)
        if initial_phase is not None:
            # initial_phase is expected to be [B, 1, dim]
            phase = phase + initial_phase

        # Interpolate phase to final resolution (often needed due to upsample implementation details)
        phase = F.interpolate(
            phase.transpose(1, 2) * self.upsample_scale,
            scale_factor=self.upsample_scale,
            mode="linear",
        ).transpose(1, 2)
        
        # [CHANGE] Save the last phase value to pass to the next chunk
        # We need the last value of the interpolated phase.
        last_phase = phase[:, -1:, :] # [B, 1, dim]

        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)
        )
        
        # [CHANGE] Pass initial_phase to _f02sine and get updated phase back
        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 with streaming support."""
    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):
        """
        x: F0 [B, 1, T]
        cache: Optional tensor containing phase from previous chunk [B, 1, harmonic_num+1]
        """
        initial_phase = cache # [CHANGE] Phase from previous chunk

        with torch.no_grad():
            # [CHANGE] Receive updated phase for next chunk
            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)


# ==============================================================================
# Encoder Block
# ==============================================================================

class EncoderBlock(nn.Module):
    """Downsampling encoder block."""
    def __init__(self, dim_in: int, dim_out: int, stride: int = 2):
        super().__init__()
        self.residual = nn.Sequential(
            weight_norm(nn.Conv1d(dim_in, dim_in, 7, padding=3)),
            nn.SiLU(),
            weight_norm(nn.Conv1d(dim_in, dim_in, 7, dilation=3, padding=9)),
            nn.SiLU(),
        )
        
        if stride == 1:
            self.downsample = weight_norm(
                nn.Conv1d(dim_in, dim_out, kernel_size=3, stride=1, padding=1)
            )
        else:
            self.downsample = weight_norm(
                nn.Conv1d(dim_in, dim_out, kernel_size=2*stride, stride=stride, padding=stride//2)
            )
        
    def forward(self, x):
        x = x + self.residual(x)
        return self.downsample(x)


# ==============================================================================
# Temporal Style Encoder
# ==============================================================================

class StyleResBlock2d(nn.Module):
    """2D Residual block for style encoder with optional downsampling."""
    def __init__(self, in_ch, out_ch, stride=(1, 1), dilation=1):
        super().__init__()
        self.conv1 = nn.Conv2d(in_ch, out_ch, kernel_size=3, stride=stride, 
                               padding=dilation, dilation=dilation)
        # GroupNorm is standard here, kept as is
        self.gn1 = nn.GroupNorm(min(8, out_ch), out_ch)
        self.conv2 = nn.Conv2d(out_ch, out_ch, kernel_size=3, stride=1, padding=1)
        self.gn2 = nn.GroupNorm(min(8, out_ch), out_ch)
        self.act = nn.LeakyReLU(0.2)
        
        self.skip = nn.Identity()
        if in_ch != out_ch or stride != (1, 1):
            self.skip = nn.Conv2d(in_ch, out_ch, kernel_size=1, stride=stride)
    
    def forward(self, x):
        residual = self.skip(x)
        x = self.act(self.gn1(self.conv1(x)))
        x = self.gn2(self.conv2(x))
        return self.act(x + residual)


class MultiScaleStyleRefine1d(nn.Module):
    """Multi-scale dilated convs to capture fine details at different scales."""
    def __init__(self, channels, dilations=[1, 2, 4]):
        super().__init__()
        self.branches = nn.ModuleList([
            nn.Sequential(
                nn.Conv1d(channels, channels, kernel_size=3, padding=d, dilation=d),
                nn.GroupNorm(min(8, channels), channels),
                nn.LeakyReLU(0.2),
            )
            for d in dilations
        ])
        self.fuse = nn.Conv1d(channels * len(dilations), channels, kernel_size=1)
        self.act = nn.LeakyReLU(0.2)
    
    def forward(self, x):
        outs = [branch(x) for branch in self.branches]
        fused = self.fuse(torch.cat(outs, dim=1))
        return self.act(fused + x)


class TemporalStyleEncoder(nn.Module):
    def __init__(
        self,
        n_mels: int = 40,
        style_dim: int = 64,
        hidden_dims: List[int] = [32, 64, 128, 256, 512],
        downsample_factor: int = 16,
    ):
        super().__init__()
        self.n_mels = n_mels
        self.style_dim = style_dim
        self.downsample_factor = downsample_factor
        
        self.stem = nn.Sequential(
            nn.Conv2d(1, hidden_dims[0], kernel_size=5, stride=1, padding=2),
            nn.GroupNorm(min(8, hidden_dims[0]), hidden_dims[0]),
            nn.LeakyReLU(0.2),
            nn.Conv2d(hidden_dims[0], hidden_dims[0], kernel_size=3, stride=1, padding=1),
            nn.GroupNorm(min(8, hidden_dims[0]), hidden_dims[0]),
            nn.LeakyReLU(0.2),
        )
        
        self.stages = nn.ModuleList()
        in_ch = hidden_dims[0]
        
        for i, out_ch in enumerate(hidden_dims):
            freq_stride = 2 if i < 2 else 1
            
            if self.downsample_factor == 80 and i == 0:
                time_stride = 5
            elif self.downsample_factor == 16:
                time_stride = 2 if i < 4 else 1
            else:
                time_stride = 2
            
            self.stages.append(nn.Sequential(
                StyleResBlock2d(in_ch, out_ch, stride=(freq_stride, time_stride)),
                StyleResBlock2d(out_ch, out_ch, stride=(1, 1)),
            ))
            in_ch = out_ch
        
        self.out_freq = n_mels // 4
        
        self.freq_pool = nn.Sequential(
            nn.Conv2d(hidden_dims[-1], hidden_dims[-1], 
                      kernel_size=(self.out_freq, 1), padding=0),
            nn.GroupNorm(min(8, hidden_dims[-1]), hidden_dims[-1]),
            nn.LeakyReLU(0.2),
        )
        
        self.temporal_refine = MultiScaleStyleRefine1d(hidden_dims[-1], dilations=[1, 2, 4])
        
        self.proj = nn.Sequential(
            nn.Conv1d(hidden_dims[-1], hidden_dims[-1] // 2, kernel_size=3, padding=1),
            nn.GroupNorm(min(8, hidden_dims[-1] // 2), hidden_dims[-1] // 2),
            nn.LeakyReLU(0.2),
            nn.Conv1d(hidden_dims[-1] // 2, style_dim, kernel_size=3, padding=1),
            nn.LeakyReLU(0.2),
            nn.Conv1d(style_dim, style_dim, kernel_size=1),
        )
        
    def forward(self, x):
        if x.dim() == 3:
            x = x.unsqueeze(1)
        elif x.dim() == 4 and x.shape[1] != 1:
            if x.shape[-1] == 1:
                x = x.squeeze(-1).unsqueeze(1)
        
        x = self.stem(x)
        for stage in self.stages:
            x = stage(x)
        
        x = self.freq_pool(x)
        x = x.squeeze(2)
        x = self.temporal_refine(x)
        x = self.proj(x)
        x = x.transpose(1, 2)
        return x


class WindowedTemporalStyleEncoder(nn.Module):
    """
    Windowed temporal style encoder for timbre extraction.
    
    Instead of fine-grained T//16 downsampling (which overfits to phonemes),
    uses large overlapping windows (2-4 seconds) to extract speaker/timbre
    characteristics. Each window is processed through a 2D CNN that collapses
    frequency and time into a single style vector.
    
    Output: [B, num_windows, style_dim]
    
    For a 7.5s clip at 100fps (750 frames) with window_size=300, window_hop=100:
        num_windows = (750 - 300) / 100 + 1 = 5-6 windows
    Each window covers ~3 seconds of audio - enough for timbre, too coarse for phonemes.
    """
    def __init__(
        self,
        n_mels: int = 40,
        style_dim: int = 32,
        hidden_dims: List[int] = [32, 64, 128, 256],
        window_size: int = 300,   # ~3 sec at 100fps (hop=441, sr=44100)
        window_hop: int = 100,    # ~1 sec stride
    ):
        super().__init__()
        self.n_mels = n_mels
        self.style_dim = style_dim
        self.window_size = window_size
        self.window_hop = window_hop
        
        # Per-window 2D CNN
        self.stem = nn.Sequential(
            nn.Conv2d(1, hidden_dims[0], kernel_size=5, stride=1, padding=2),
            nn.GroupNorm(min(8, hidden_dims[0]), hidden_dims[0]),
            nn.LeakyReLU(0.2),
            nn.Conv2d(hidden_dims[0], hidden_dims[0], kernel_size=3, stride=1, padding=1),
            nn.GroupNorm(min(8, hidden_dims[0]), hidden_dims[0]),
            nn.LeakyReLU(0.2),
        )
        
        # Stages progressively downsample freq and time within each window
        # Freq: 40 -> 20 -> 10 (first 2 stages), Time: 300 -> 150 -> 75 -> 37 -> 18
        self.stages = nn.ModuleList()
        in_ch = hidden_dims[0]
        for i, out_ch in enumerate(hidden_dims):
            freq_stride = 2 if i < 2 else 1     # Collapse frequency in first 2 stages
            time_stride = 2                       # Aggressive time downsampling within window
            
            self.stages.append(nn.Sequential(
                StyleResBlock2d(in_ch, out_ch, stride=(freq_stride, time_stride)),
                StyleResBlock2d(out_ch, out_ch, stride=(1, 1)),
            ))
            in_ch = out_ch
        
        # Pool remaining spatial dims to single vector per window
        self.pool = nn.AdaptiveAvgPool2d((1, 1))
        
        # Project to style_dim
        self.proj = nn.Sequential(
            nn.Linear(hidden_dims[-1], hidden_dims[-1] // 2),
            nn.LeakyReLU(0.2),
            nn.Linear(hidden_dims[-1] // 2, style_dim),
        )
    
    def forward(self, x):
        """
        Args:
            x: [B, n_mels, T] or [B, 1, n_mels, T]
        Returns:
            [B, num_windows, style_dim]
        """
        if x.dim() == 4 and x.shape[1] == 1:
            pass  # already [B, 1, F, T]
        elif x.dim() == 3:
            x = x.unsqueeze(1)  # [B, 1, F, T]
        elif x.dim() == 4 and x.shape[-1] == 1:
            x = x.squeeze(-1).unsqueeze(1)
        
        B, _, F_dim, T = x.shape
        
        # Pad if sequence shorter than window
        if T < self.window_size:
            pad_amount = self.window_size - T
            # Replicate last frames to reach window_size
            x = torch.cat([x, x[:, :, :, -1:].expand(-1, -1, -1, pad_amount)], dim=3)
            T = self.window_size
        
        # Extract overlapping windows using unfold
        # x: [B, 1, F, T] -> [B, 1, F, num_windows, window_size]
        windows = x.unfold(3, self.window_size, self.window_hop)
        num_windows = windows.shape[3]
        
        # Reshape: [B, 1, F, num_win, win_size] -> [B*num_win, 1, F, win_size]
        windows = windows.permute(0, 3, 1, 2, 4).reshape(
            B * num_windows, 1, F_dim, self.window_size
        )
        
        # Process all windows through CNN in batch
        h = self.stem(windows)
        for stage in self.stages:
            h = stage(h)
        
        # Pool to single vector per window
        h = self.pool(h).squeeze(-1).squeeze(-1)  # [B*num_win, hidden_dims[-1]]
        h = self.proj(h)  # [B*num_win, style_dim]
        
        # Reshape back: [B, num_windows, style_dim]
        output = h.reshape(B, num_windows, self.style_dim)
        return output


class StyleUpsampleRefine(nn.Module):
    """
    Upsamples coarse windowed style to target temporal resolution via
    interpolation + learned refinement convolutions.
    Works for any upsampling ratio (unlike ConvTranspose which is ratio-specific).
    """
    def __init__(self, style_dim: int):
        super().__init__()
        self.refine = nn.Sequential(
            weight_norm(nn.Conv1d(style_dim, style_dim * 2, 5, padding=2)),
            nn.SiLU(),
            weight_norm(nn.Conv1d(style_dim * 2, style_dim * 2, 5, padding=2)),
            nn.SiLU(),
            weight_norm(nn.Conv1d(style_dim * 2, style_dim, 3, padding=1)),
        )
    
    def forward(self, x, target_len):
        """
        Args:
            x: [B, style_dim, T_style] (coarse windowed style)
            target_len: target temporal length
        Returns:
            [B, style_dim, target_len]
        """
        x_up = F.interpolate(x, size=target_len, mode='linear', align_corners=False)
        return x_up + self.refine(x_up)  # Residual refinement


# ==============================================================================
# Hybrid Prosody Encoder
# ==============================================================================

class EncoderBlock(nn.Module):
    """
    ResNet-style Encoder Block from File A.
    Uses dilated convolutions for better context capturing.
    """
    def __init__(self, dim_in: int, dim_out: int, stride: int = 2):
        super().__init__()
        self.residual = nn.Sequential(
            weight_norm(nn.Conv1d(dim_in, dim_in, 7, padding=3)),
            nn.SiLU(),
            # Dilation=3 increases receptive field for better global context
            weight_norm(nn.Conv1d(dim_in, dim_in, 7, dilation=3, padding=9)),
            nn.SiLU(),
        )
        
        if stride == 1:
            self.downsample = weight_norm(
                nn.Conv1d(dim_in, dim_out, kernel_size=3, stride=1, padding=1)
            )
        else:
            self.downsample = weight_norm(
                nn.Conv1d(dim_in, dim_out, kernel_size=2*stride, stride=stride, padding=stride//2)
            )
        
    def forward(self, x):
        x = x + self.residual(x)
        return self.downsample(x)


# ==============================================================================

# Hybrid Prosody Encoder

# ==============================================================================



class HybridProsodyEncoderTemporal(nn.Module):
    def __init__(
        self,
        style_dim: int = 64,
        latent_dim: int = 256,
        hidden_dim: int = 256,
        strides: List[int] = [2],
    ):
        super().__init__()
        self.latent_dim = latent_dim
        self.style_dim = style_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(),
        )

        # NOTE: Style is intentionally NOT included here.
        # Style goes only to the decoder (via AdaIN). If style is also in the
        # prosody encoder, the model routes info through the un-quantized style
        # bypass path, causing codebook collapse (low FSQ utilization).
        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 style 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)






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)


# ==============================================================================
# Refined Fusion Modules
# ==============================================================================

class FusionResBlock(nn.Module):
    def __init__(
        self,
        dim_in,
        dim_out,
        style_dim=64,
        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 = TemporalAdaIN1d(style_dim, dim_in)
        self.norm2 = TemporalAdaIN1d(style_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, s):
        x = self.norm1(x, s)
        x = self.actv(x)
        x = self.conv1(self.dropout(x))
        x = self.norm2(x, s)
        x = self.actv(x)
        x = self.conv2(self.dropout(x))
        return x

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


class ResNetFusionModule(nn.Module):
    def __init__(self, dim_in, hidden_dim, style_dim):
        super().__init__()
        # Matches logic from user:
        # x = cat[asr, latent] -> fused
        # x = encode(x, s)
        # decode: x = cat[x, asr, latent] -> cat[x, fused]
        
        self.input_mix = FusionResBlock(dim_in, hidden_dim, style_dim)
        
        self.decode = nn.ModuleList()
        
        # Dimensions for concatenation: x (hidden_dim) + fused (dim_in)
        concat_dim = hidden_dim + dim_in
        
        self.decode.append(FusionResBlock(concat_dim, hidden_dim, style_dim))
        self.decode.append(FusionResBlock(concat_dim, hidden_dim, style_dim))
        self.decode.append(FusionResBlock(concat_dim, hidden_dim, style_dim))
        
    def forward(self, prosody_latent, text_emb, style, language_emb=None):
        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)
             
        # Encode
        x = self.input_mix(fused, style)
        
        # Decode loop with re-injection
        for block in self.decode:
            x = torch.cat([x, fused], dim=1)
            x = block(x, style)
            
        return x


# ==============================================================================
# Hybrid Waveform Decoder with Temporal Style (Inference Chunks)
# ==============================================================================

class HybridWaveformDecoderTemporal(nn.Module):
    def __init__(
        self,
        prosody_latent_dim: int = 256,
        text_dim: int = 512,
        style_dim: int = 64,
        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.style_dim = style_dim
        
        total_upsample = int(np.prod(upsample_rates)) * gen_istft_hop_size
        self.source_upsample_rate = source_upsample_rate or total_upsample
        
        self.style_predictor_down = nn.Sequential(
            weight_norm(nn.Conv1d(prosody_latent_dim, prosody_latent_dim, 4, stride=2, padding=1)),
            nn.SiLU(),
            weight_norm(nn.Conv1d(prosody_latent_dim, prosody_latent_dim, 4, stride=2, padding=1)),
            nn.SiLU(),
            weight_norm(nn.Conv1d(prosody_latent_dim, prosody_latent_dim, 4, stride=2, padding=1)),
            nn.SiLU(),
        )
        
        self.style_predictor = nn.Sequential(
            weight_norm(nn.Conv1d(prosody_latent_dim, hidden_dim, 5, padding=2)),
            nn.SiLU(),
            weight_norm(nn.Conv1d(hidden_dim, hidden_dim, 5, padding=2)),
            nn.SiLU(),
            weight_norm(nn.Conv1d(hidden_dim, hidden_dim // 2, 3, padding=1)),
            nn.SiLU(),
            weight_norm(nn.Conv1d(hidden_dim // 2, style_dim, 3, padding=1)),
        )
        
        self.predicted_style_upsampler = nn.Sequential(
            weight_norm(nn.Conv1d(style_dim, style_dim * 2, 5, padding=2)),
            nn.SiLU(),
            weight_norm(nn.Conv1d(style_dim * 2, style_dim * 2, 5, padding=2)),
            nn.SiLU(),
            weight_norm(nn.Conv1d(style_dim * 2, style_dim, 3, padding=1)),
        )
        
        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))
        )
        
        # [CHANGE] Source Module now handles phase caching via forward argument
        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
        
        self.pre_decoder = ResNetFusionModule(
            dim_in=fusion_dim,
            hidden_dim=hidden_dim,
            style_dim=style_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(TemporalAdaINResBlock1(c_cur, 7, [1, 3, 5], style_dim))
            else:
                self.noise_convs.append(
                    weight_norm(nn.Conv1d(gen_istft_n_fft + 2, c_cur, kernel_size=1))
                )
                self.noise_res.append(TemporalAdaINResBlock1(c_cur, 11, [1, 3, 5], style_dim))
        
        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(TemporalAdaINResBlock1(ch, k, d, style_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, style_temporal=None, f0_gt=None, cache=None, language_emb=None):
        """
        [CHANGE] Added cache argument for streaming inference.
        
        Args:
            cache (dict, optional): Dictionary containing 'source_phase' from previous chunk.
                                    If None, assumes full sequence (Training).
            language_emb (Tensor, optional): [B, language_dim]
        
        Returns:
            ...
            new_cache (dict): Updated cache for next chunk (only if cache is not None).
        """
        B = prosody_latent.shape[0]

      
        prosody_down = self.style_predictor_down(prosody_latent)
        style_pred_compressed = self.style_predictor(prosody_down)
        style_pred_compressed = style_pred_compressed.transpose(1, 2)
        
        if style_temporal is not None:
            style_to_use = style_temporal
        else:
            # Predicted style: interpolate from compressed resolution to reasonable temporal resolution, then refine
            pred_t = style_pred_compressed.transpose(1, 2)  # [B, style_dim, T_compressed]
            # Upsample to ~frame-level via interpolation + learned refinement
            target_len = int(prosody_latent.shape[-1] * self.codec_compression)
            pred_up = F.interpolate(pred_t, size=target_len, mode='linear', align_corners=False)
            pred_up = pred_up + self.predicted_style_upsampler(pred_up)  # residual refinement
            style_to_use = pred_up.transpose(1, 2)
            style_to_use = style_to_use.detach()
        
        f0_pred_latent = self.f0_predictor(prosody_latent)
        
        # Always run the upsampler (e.g. ConvTranspose1d) so it gets trained
        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 with Phase Caching
        # =======================================================================
        f0_log = self.f0_upsamp(f0_to_use[:, None]).transpose(1, 2)
        f0_lin = (10.0 ** f0_log.float()).to(f0_log.dtype)
        
        # [CHANGE] Extract phase from cache if available
        source_phase_cache = cache.get("source_phase") if cache is not None else None
        
        # [CHANGE] Pass phase to source module and get new phase back
        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]]
        
        x = self.pre_decoder(prosody_latent, text_emb, style_to_use, 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, style_to_use)
            
            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, style_to_use)
                else:
                    xs += self.resblocks[i * self.num_kernels + j](x, style_to_use)
            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)
        
        # [CHANGE] If inferencing by chunk (cache provided), return the updated cache
        if cache is not None:
            new_cache = {
                "source_phase": next_source_phase
            }
            return out, spec, phase, f0_pred, style_pred_compressed, new_cache
            
        return out, spec, phase, f0_pred, style_pred_compressed


class HybridTTSCodecVocoderTemporal(nn.Module):
    """
    Hybrid TTS Codec with TEMPORAL style encoder and Chunked Inference support.
    """
    def __init__(
        self,
        n_mels: int = 40,
        text_dim: int = 512,
        style_dim: int = 64,
        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.n_mels = n_mels
        self.text_dim = text_dim
        self.style_dim = style_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
        
        self.style_encoder = WindowedTemporalStyleEncoder(
            n_mels=n_mels,
            style_dim=style_dim,
            hidden_dims=[32, 64, 128, 256],
            window_size=300,   # ~3 sec at 100fps -> captures timbre, not phonemes
            window_hop=100,    # ~1 sec stride -> smooth temporal transitions
        )
        
        # Interpolation + refinement upsampler (works for any ratio, unlike ConvTranspose)
        self.style_upsampler = StyleUpsampleRefine(style_dim)
        
        self.prosody_encoder = HybridProsodyEncoderTemporal(
            style_dim=style_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 = HybridWaveformDecoderTemporal(
            prosody_latent_dim=prosody_latent_dim,
            text_dim=text_dim,
            style_dim=style_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, mel, n_quantizers=None, use_predicted_style=False, use_predicted_f0=False, language_emb=None):
        """Training forward pass (Standard, no cache)."""
        # Windowed style encoder: [B, num_windows, style_dim] — coarse timbre
        style_compressed = self.style_encoder(mel)
        
        # Prosody encoder takes ONLY pitch + energy (no style!)
        # This forces the FSQ codebook to carry all pitch/energy info,
        # preventing codebook collapse from style bypass.
        prosody_latent = self.prosody_encoder(pitch, energy)
        
        # Upsample windowed style to match pitch length via interpolation + refinement
        style_t = style_compressed.transpose(1, 2)  # [B, style_dim, num_windows]
        style_up = self.style_upsampler(style_t, target_len=pitch.shape[1])  # [B, style_dim, T]
        style_up = style_up.transpose(1, 2)  # [B, T, style_dim]
        
        quantized_prosody, tokens, commitment_loss = self.quantizer(prosody_latent)
        
        decoder_style = None if use_predicted_style else style_up
        decoder_f0 = None if use_predicted_f0 else pitch

        # Training: cache is None
        wav, mag, phase, f0_pred, style_pred = self.decoder(
            quantized_prosody,
            text_emb,
            decoder_style,
            f0_gt=decoder_f0,
            cache=None, # Explicitly None for training
            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,
            "style_temporal": style_compressed,
            "style_up": style_up,
            "commitment_loss": commitment_loss,
            "f0_pred": f0_pred,
            "f0_gt": pitch,
            "style_pred": style_pred,
        }
    
    @torch.no_grad()
    def encode_style(self, mel, return_upsampled=False, target_len=None):
        style = self.style_encoder(mel)
        if return_upsampled:
            style_t = style.transpose(1, 2)
            tgt = target_len if target_len is not None else style_t.shape[-1] * 16
            style_up = self.style_upsampler(style_t, target_len=tgt).transpose(1, 2)
            return style, style_up
        return style
    
    @torch.no_grad()
    def tokenize(self, pitch, energy, text_emb, mel, n_quantizers=None):
        style_compressed = self.style_encoder(mel)
        prosody_latent = self.prosody_encoder(pitch, energy)
        _, tokens, _ = self.quantizer(prosody_latent)
        return tokens, text_emb, style_compressed
    
    @torch.no_grad()
    def decode_tokens(self, tokens, text_emb, style_temporal=None, language_emb=None):
        """Non-streaming decode."""
        quantized_prosody = self.quantizer.decode(tokens)
        
        if style_temporal is not None:
            if style_temporal.shape[1] < text_emb.shape[2] // 2:
                style_t = style_temporal.transpose(1, 2)
                style_up = self.style_upsampler(style_t, target_len=text_emb.shape[2])
                style_temporal = style_up.transpose(1, 2)
        
        wav, _, _, _, _ = self.decoder(
            quantized_prosody,
            text_emb,
            style_temporal,
            f0_gt=None,
            cache=None,
            language_emb=language_emb
        )
        return wav

    @torch.no_grad()
    def decode_tokens_with_predictions(self, tokens, text_emb, style_temporal=None, language_emb=None):
        """
        Generate waveform and return intermediate predictions.
        
        Args:
            tokens: [B, 1, T_comp] - prosody tokens
            text_emb: [B, text_dim, T] - text embeddings
            style_temporal: [B, T_style, style_dim] or None (if None, style is predicted)
            
        Returns:
            dict with wav, mag, phase, f0_pred, style_pred
        """
        quantized_prosody = self.quantizer.decode(tokens)
        
        # If style is provided, upsample it; otherwise decoder will predict style
        if style_temporal is not None:
            # Upsample windowed style to match text temporal resolution
            if style_temporal.shape[1] < text_emb.shape[2] // 2: 
                style_t = style_temporal.transpose(1, 2)
                style_up = self.style_upsampler(style_t, target_len=text_emb.shape[2])
                style_temporal = style_up.transpose(1, 2)

        wav, mag, phase, f0_pred, style_pred = self.decoder(
            quantized_prosody,
            text_emb,
            style_temporal,  # Can be None, decoder will use predicted style
            f0_gt=None,
            cache=None,
            language_emb=language_emb
        )
        return {
            "wav": wav,
            "mag": mag,
            "phase": phase,
            "f0_pred": f0_pred,
            "style_pred": style_pred,
        }

    @torch.no_grad()
    def decode_chunk(self, tokens, text_emb, style_temporal=None, cache=None, language_emb=None):
        """
        [NEW] Streaming inference by chunk.
        
        Args:
            tokens: Chunk of tokens
            text_emb: Chunk of text embeddings
            style_temporal: Chunk of style (or None)
            cache: Dictionary from previous chunk call (init with {})
        
        Returns:
            wav_chunk, new_cache
        """
        if cache is None:
            cache = {}
            
        quantized_prosody = self.quantizer.decode(tokens)
        
        if style_temporal is not None:
            # Assumes style_temporal is already upsampled/processed for this chunk size
            # or is None to allow prediction
            pass

        # Call decoder with cache
        wav, mag, phase, f0_pred, style_pred, new_cache = self.decoder(
            quantized_prosody,
            text_emb,
            style_temporal,
            f0_gt=None,
            cache=cache,
            language_emb=language_emb
        )
        
        return wav, new_cache