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