""" Codec Decoder with Learnable Speaker Embeddings. This module extends the hybrid temporal codec with: - Learnable speaker embeddings: nn.Embedding(num_speakers=11, embedding_dim=128) - Speaker conditioning via AdaIN1d (like ringformer.py) - Speaker IDs 0-10 for 11 speakers Usage: model = HybridTTSCodecVocoderSpeaker(...) output = model(pitch, energy, text_emb, mel, speaker_ids=speaker_ids) """ import math import random import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from torch.nn.utils import weight_norm, remove_weight_norm from scipy.signal import get_window from einops import rearrange from typing import Tuple, Optional, List, Dict, Union from .conformer import Conformer from .utils import init_weights, get_padding # ============================================================================== # Utility modules # ============================================================================== class TorchSTFT(nn.Module): def __init__(self, filter_length=800, hop_length=200, win_length=800, window="hann"): super().__init__() self.filter_length = filter_length self.hop_length = hop_length self.win_length = win_length self.window = torch.from_numpy( get_window(window, win_length, fftbins=True).astype(np.float32) ) def transform(self, input_data): forward_transform = torch.stft( input_data, self.filter_length, self.hop_length, self.win_length, window=self.window.to(input_data.device), return_complex=True, ) return torch.abs(forward_transform), torch.angle(forward_transform) def inverse(self, magnitude, phase): inverse_transform = torch.istft( magnitude * torch.exp(phase * 1j), self.filter_length, self.hop_length, self.win_length, window=self.window.to(magnitude.device), ) return inverse_transform.unsqueeze(-2) class Snake1d(nn.Module): """Learned periodic activation from BigVGAN.""" def __init__(self, in_features): super().__init__() self.alpha = nn.Parameter(torch.ones(1, in_features, 1)) def forward(self, x): return x + (1.0 / (self.alpha + 1e-9)) * (torch.sin(self.alpha * x) ** 2) class AdaIN1d(nn.Module): """ Adaptive Instance Normalization for 1D signals. Follows the ringformer.py implementation. Takes a style vector [B, style_dim] and applies affine transformation to normalized features [B, C, T]. """ def __init__(self, style_dim, num_features): super().__init__() self.norm = nn.InstanceNorm1d(num_features, affine=False) self.fc = nn.Linear(style_dim, num_features * 2) def forward(self, x, s): """ Args: x: [B, C, T] input features s: [B, style_dim] style/speaker embedding Returns: [B, C, T] AdaIN-transformed features """ h = self.fc(s) h = h.view(h.size(0), h.size(1), 1) gamma, beta = torch.chunk(h, chunks=2, dim=1) return (1 + gamma) * self.norm(x) + beta class SpeakerAdaINResBlock1(nn.Module): """ Residual block with AdaIN speaker conditioning. Uses global speaker embedding [B, speaker_dim] for style. """ def __init__(self, channels, kernel_size=3, dilation=(1, 3, 5), speaker_dim=128): super().__init__() self.convs1 = nn.ModuleList([ weight_norm(nn.Conv1d(channels, channels, kernel_size, 1, dilation=d, padding=get_padding(kernel_size, d))) for d in dilation ]) self.convs1.apply(init_weights) self.convs2 = nn.ModuleList([ weight_norm(nn.Conv1d(channels, channels, kernel_size, 1, dilation=1, padding=get_padding(kernel_size, 1))) for _ in dilation ]) self.convs2.apply(init_weights) self.adain1 = nn.ModuleList([AdaIN1d(speaker_dim, channels) for _ in dilation]) self.adain2 = nn.ModuleList([AdaIN1d(speaker_dim, channels) for _ in dilation]) self.snakes1 = nn.ModuleList([Snake1d(channels) for _ in dilation]) self.snakes2 = nn.ModuleList([Snake1d(channels) for _ in dilation]) def forward(self, x, speaker_emb): """ Args: x: [B, C, T] input features speaker_emb: [B, speaker_dim] speaker embedding """ for c1, c2, n1, n2, s1, s2 in zip( self.convs1, self.convs2, self.adain1, self.adain2, self.snakes1, self.snakes2 ): xt = n1(x, speaker_emb) xt = s1(xt) xt = c1(xt) xt = n2(xt, speaker_emb) xt = s2(xt) xt = c2(xt) x = xt + x return x # ============================================================================== # Harmonic Source Module # ============================================================================== class SineGen(nn.Module): """Sine generator for F0-based harmonic source with phase caching.""" def __init__(self, samp_rate, upsample_scale, harmonic_num=0, sine_amp=0.1, noise_std=0.003, voiced_threshold=0, flag_for_pulse=False): super().__init__() self.sine_amp = sine_amp self.noise_std = noise_std self.harmonic_num = harmonic_num self.dim = harmonic_num + 1 self.sampling_rate = samp_rate self.voiced_threshold = voiced_threshold self.upsample_scale = upsample_scale self.flag_for_pulse = flag_for_pulse def _f02uv(self, f0): return (f0 > self.voiced_threshold).float() def _f02sine(self, f0_values, initial_phase=None): rad_values = (f0_values / self.sampling_rate) % 1 rand_ini = torch.rand(f0_values.shape[0], f0_values.shape[2], device=f0_values.device) rand_ini[:, 0] = 0 rad_values[:, 0, :] = rad_values[:, 0, :] + rand_ini rad_values = F.interpolate( rad_values.transpose(1, 2), scale_factor=1 / self.upsample_scale, mode="linear", ).transpose(1, 2) phase = torch.cumsum(rad_values, dim=1) * 2 * np.pi if initial_phase is not None: phase = phase + initial_phase phase = F.interpolate( phase.transpose(1, 2) * self.upsample_scale, scale_factor=self.upsample_scale, mode="linear", ).transpose(1, 2) last_phase = phase[:, -1:, :] if self.flag_for_pulse: sines = torch.cos(phase) else: sines = torch.sin(phase) return sines, last_phase def forward(self, f0, initial_phase=None): f0_buf = torch.zeros(f0.shape[0], f0.shape[1], self.dim, device=f0.device) fn = torch.multiply( f0, torch.FloatTensor([[range(1, self.harmonic_num + 2)]]).to(f0.device) ) sine_waves, next_phase = self._f02sine(fn, initial_phase) sine_waves = sine_waves * self.sine_amp uv = self._f02uv(f0) noise_amp = uv * self.noise_std + (1 - uv) * self.sine_amp / 3 noise = noise_amp * torch.randn_like(sine_waves) sine_waves = sine_waves * uv + noise return sine_waves, uv, noise, next_phase class SourceModuleHnNSF(nn.Module): """Source module for harmonic-plus-noise synthesis.""" def __init__(self, sampling_rate, upsample_scale, harmonic_num=0, sine_amp=0.1, add_noise_std=0.003, voiced_threshold=0): super().__init__() self.sine_amp = sine_amp self.noise_std = add_noise_std self.l_sin_gen = SineGen( sampling_rate, upsample_scale, harmonic_num, sine_amp, add_noise_std, voiced_threshold, flag_for_pulse=False ) self.l_linear = nn.Linear(harmonic_num + 1, 1) self.l_tanh = nn.Tanh() def forward(self, x, cache=None): initial_phase = cache with torch.no_grad(): sine_wavs, uv, _, next_phase = self.l_sin_gen(x, initial_phase=initial_phase) sine_merge = self.l_tanh(self.l_linear(sine_wavs)) noise = torch.randn_like(uv) * self.sine_amp / 3 return sine_merge, noise, uv, next_phase # ============================================================================== # Pixel Shuffle Upsampling # ============================================================================== def pixel_shuffle_1d(x: torch.Tensor, r: int) -> torch.Tensor: B, Cr, L = x.size() C = Cr // r x = x.view(B, C, r, L).permute(0, 1, 3, 2) return x.reshape(B, C, L * r) class UpsamplePixelShuffle1D(nn.Module): def __init__(self, in_ch: int, out_ch: int, kernel_size: int, r: int): super().__init__() self.r = r pad_l, pad_r = (kernel_size - 1) // 2, kernel_size // 2 self.pad = nn.ReflectionPad1d((pad_l, pad_r)) self.conv = weight_norm(nn.Conv1d(in_ch, out_ch * r, kernel_size, padding=0)) self._init_icnr(in_ch, out_ch, r, kernel_size) def _init_icnr(self, in_ch, out_ch, r, kernel_size): """ICNR initialization for smooth upsampling.""" weight = self.conv.weight.data kernel = torch.zeros(out_ch, in_ch, kernel_size) nn.init.kaiming_normal_(kernel) weight.copy_(kernel.repeat(r, 1, 1)) if self.conv.bias is not None: self.conv.bias.data.fill_(0) def forward(self, x): x = self.pad(x) x = self.conv(x) return pixel_shuffle_1d(x, self.r) # ============================================================================== # Hybrid Prosody Encoder (No style - forces codebook usage) # ============================================================================== class HybridProsodyEncoderSpeaker(nn.Module): def __init__( self, speaker_dim: int = 128, latent_dim: int = 256, hidden_dim: int = 256, strides: List[int] = [2], ): super().__init__() self.latent_dim = latent_dim self.speaker_dim = speaker_dim self.compression_ratio = int(np.prod(strides)) self.pitch_down = nn.Sequential( weight_norm(nn.Conv1d(1, hidden_dim, 7, stride=2, padding=3)), nn.SiLU(), weight_norm(nn.Conv1d(hidden_dim, hidden_dim, 5, stride=1, padding=2)), nn.SiLU(), weight_norm(nn.Conv1d(hidden_dim, hidden_dim, 3, stride=1, padding=1)), nn.SiLU(), ) self.energy_down = nn.Sequential( weight_norm(nn.Conv1d(1, hidden_dim, 7, stride=2, padding=3)), nn.SiLU(), weight_norm(nn.Conv1d(hidden_dim, hidden_dim, 5, stride=1, padding=2)), nn.SiLU(), weight_norm(nn.Conv1d(hidden_dim, hidden_dim, 3, stride=1, padding=1)), nn.SiLU(), ) input_dim = hidden_dim * 2 self.fusion = nn.Sequential( weight_norm(nn.Conv1d(input_dim, hidden_dim, 7, padding=3)), nn.SiLU(), weight_norm(nn.Conv1d(hidden_dim, hidden_dim, 5, padding=2)), nn.SiLU(), weight_norm(nn.Conv1d(hidden_dim, hidden_dim, 3, padding=1)), nn.SiLU(), ) self.refine = nn.Sequential( weight_norm(nn.Conv1d(hidden_dim, hidden_dim, 7, padding=3)), nn.SiLU(), weight_norm(nn.Conv1d(hidden_dim, hidden_dim * 2, 5, padding=2)), nn.SiLU(), weight_norm(nn.Conv1d(hidden_dim * 2, hidden_dim * 2, 3, padding=1)), nn.SiLU(), ) self.to_latent = nn.Sequential( weight_norm(nn.Conv1d(hidden_dim * 2, hidden_dim * 2, 5, padding=2)), nn.SiLU(), weight_norm(nn.Conv1d(hidden_dim * 2, latent_dim, 1)), ) def forward(self, pitch, energy): """Encode pitch + energy into prosody latent. No speaker here - forces codebook usage.""" pitch_feat = self.pitch_down(pitch.unsqueeze(1)) energy_feat = self.energy_down(energy.unsqueeze(1)) min_len = min(pitch_feat.shape[-1], energy_feat.shape[-1]) pitch_feat = pitch_feat[..., :min_len] energy_feat = energy_feat[..., :min_len] x = torch.cat([pitch_feat, energy_feat], dim=1) x = self.fusion(x) x = self.refine(x) return self.to_latent(x) # ============================================================================== # Finite Scalar Quantization # ============================================================================== class FiniteScalarQuantization(nn.Module): def __init__(self, input_dim=256, levels: List[int] = [4]*6): super().__init__() self.input_dim = input_dim self.levels = levels self.dims = len(levels) self.codebook_size = math.prod(levels) self.in_proj = nn.Sequential( nn.Linear(input_dim, input_dim // 2), nn.SiLU(), nn.Linear(input_dim // 2, self.dims), ) self.out_proj = nn.Sequential( nn.Linear(self.dims, input_dim // 2), nn.SiLU(), nn.Linear(input_dim // 2, input_dim), ) self.scale = nn.Parameter(torch.ones(self.dims) * 1.5) self.bias = nn.Parameter(torch.zeros(self.dims)) for m in self.in_proj.modules(): if isinstance(m, nn.Linear): nn.init.xavier_uniform_(m.weight, gain=2.0) if m.bias is not None: nn.init.zeros_(m.bias) for m in self.out_proj.modules(): if isinstance(m, nn.Linear): nn.init.xavier_uniform_(m.weight, gain=1.0) if m.bias is not None: nn.init.zeros_(m.bias) self.register_buffer('levels_tensor', torch.tensor(levels, dtype=torch.float32)) _basis = torch.cumprod(torch.tensor([1] + levels[:-1]), dim=0) self.register_buffer('basis', _basis) self.register_buffer('num_steps', torch.tensor(0)) self.warmup_steps = 5000 def forward(self, x, n_quantizers=None): x = x.transpose(1, 2) z = self.in_proj(x) z = z * self.scale + self.bias z_bound = torch.tanh(z) if self.training: self.num_steps += 1 noise_scale = max(0.3 * (1 - self.num_steps.float() / self.warmup_steps), 0.05) noise = (torch.rand_like(z_bound) - 0.5) * 2 * noise_scale z_bound_noisy = z_bound + noise z_bound_noisy = torch.clamp(z_bound_noisy, -1, 1) else: z_bound_noisy = z_bound levels = self.levels_tensor.to(z.device) half_l = (levels - 1) / 2 z_scaled = z_bound_noisy * half_l z_shifted = z_scaled + half_l z_ind = z_shifted.round() z_ind = torch.clamp(z_ind, torch.zeros_like(levels), levels - 1) z_q_target = z_ind - half_l z_q = z_scaled + (z_q_target - z_scaled).detach() out = self.out_proj(z_q) z_ind_long = z_ind.long() indices = (z_ind_long * self.basis).sum(dim=-1) out = out.transpose(1, 2) aux_loss = self._entropy_loss(z_shifted, levels) return out, indices.unsqueeze(1), aux_loss def _entropy_loss(self, z_shifted, levels): B, T, D = z_shifted.shape total_entropy_loss = torch.tensor(0.0, device=z_shifted.device) for d in range(D): vals = z_shifted[..., d].reshape(-1) num_levels = int(levels[d].item()) centers = torch.arange(num_levels, device=z_shifted.device, dtype=torch.float32) dist = (vals.unsqueeze(1) - centers.unsqueeze(0)).pow(2) probs = F.softmax(-dist / 0.5, dim=1) avg_probs = probs.mean(dim=0) uniform = torch.ones_like(avg_probs) / num_levels kl_div = (avg_probs * (torch.log(avg_probs + 1e-7) - torch.log(uniform + 1e-7))).sum() total_entropy_loss = total_entropy_loss + kl_div return 0.1 * total_entropy_loss / D def decode(self, indices): if indices.dim() == 3: indices = indices.squeeze(1) z_q = [] remainder = indices for i in range(self.dims): val = remainder % self.levels[i] remainder = remainder // self.levels[i] z_q.append(val) z_q = torch.stack(z_q, dim=-1).float().to(indices.device) levels = self.levels_tensor.to(indices.device) half_l = (levels - 1) / 2 z_q = z_q - half_l out = self.out_proj(z_q) return out.transpose(1, 2) # ============================================================================== # Speaker-Conditioned Fusion Module with AdaIN1d # ============================================================================== class SpeakerFusionResBlock(nn.Module): """ Fusion ResBlock conditioned on speaker embedding via AdaIN1d. """ def __init__( self, dim_in, dim_out, speaker_dim=128, actv=nn.LeakyReLU(0.2), dropout_p=0.0, ): super().__init__() self.actv = actv self.learned_sc = dim_in != dim_out self.dropout = nn.Dropout(dropout_p) self.conv1 = weight_norm(nn.Conv1d(dim_in, dim_out, 3, 1, 1)) self.conv2 = weight_norm(nn.Conv1d(dim_out, dim_out, 3, 1, 1)) # AdaIN1d with speaker embedding (global, not temporal) self.norm1 = AdaIN1d(speaker_dim, dim_in) self.norm2 = AdaIN1d(speaker_dim, dim_out) if self.learned_sc: self.conv1x1 = weight_norm(nn.Conv1d(dim_in, dim_out, 1, 1, 0, bias=False)) def _shortcut(self, x): if self.learned_sc: x = self.conv1x1(x) return x def _residual(self, x, speaker_emb): x = self.norm1(x, speaker_emb) x = self.actv(x) x = self.conv1(self.dropout(x)) x = self.norm2(x, speaker_emb) x = self.actv(x) x = self.conv2(self.dropout(x)) return x def forward(self, x, speaker_emb): out = self._residual(x, speaker_emb) out = (out + self._shortcut(x)) / math.sqrt(2) return out class SpeakerFusionModule(nn.Module): """ ResNet-style fusion module with speaker conditioning via AdaIN1d. """ def __init__(self, dim_in, hidden_dim, speaker_dim=128): super().__init__() self.input_mix = SpeakerFusionResBlock(dim_in, hidden_dim, speaker_dim) self.decode = nn.ModuleList() concat_dim = hidden_dim + dim_in self.decode.append(SpeakerFusionResBlock(concat_dim, hidden_dim, speaker_dim)) self.decode.append(SpeakerFusionResBlock(concat_dim, hidden_dim, speaker_dim)) self.decode.append(SpeakerFusionResBlock(concat_dim, hidden_dim, speaker_dim)) def forward(self, prosody_latent, text_emb, speaker_emb, language_emb=None): """ Args: prosody_latent: [B, prosody_dim, T] text_emb: [B, text_dim, T] speaker_emb: [B, speaker_dim] - global speaker embedding language_emb: [B, language_dim] optional """ if language_emb is not None: language_emb_expanded = language_emb.unsqueeze(-1).expand(-1, -1, prosody_latent.shape[-1]) fused = torch.cat([prosody_latent, text_emb, language_emb_expanded], dim=1) else: fused = torch.cat([prosody_latent, text_emb], dim=1) x = self.input_mix(fused, speaker_emb) for block in self.decode: x = torch.cat([x, fused], dim=1) x = block(x, speaker_emb) return x # ============================================================================== # Waveform Decoder with Speaker Conditioning # ============================================================================== class HybridWaveformDecoderSpeaker(nn.Module): """ Waveform decoder conditioned on learnable speaker embeddings via AdaIN1d. """ def __init__( self, prosody_latent_dim: int = 256, text_dim: int = 512, speaker_dim: int = 128, language_dim: int = 0, hidden_dim: int = 512, upsample_rates: List[int] = [12, 10], resblock_kernel_sizes: List[int] = [3, 7, 11], resblock_dilation_sizes: List[List[int]] = [[1, 3, 5], [1, 3, 5], [1, 3, 5]], gen_istft_n_fft: int = 30, gen_istft_hop_size: int = 5, sample_rate: int = 44100, source_upsample_rate: Optional[int] = None, codec_strides: Optional[List[int]] = None, ): super().__init__() self.num_upsamples = len(upsample_rates) self.num_kernels = len(resblock_kernel_sizes) self.gen_istft_n_fft = gen_istft_n_fft self.gen_istft_hop_size = gen_istft_hop_size self.codec_strides = codec_strides or [1] self.codec_compression = int(np.prod(self.codec_strides)) self.speaker_dim = speaker_dim total_upsample = int(np.prod(upsample_rates)) * gen_istft_hop_size self.source_upsample_rate = source_upsample_rate or total_upsample self.prosody_upsampler = nn.Sequential( nn.Upsample(scale_factor=2, mode='linear', align_corners=False), weight_norm(nn.Conv1d(prosody_latent_dim, prosody_latent_dim, 3, stride=1, padding=1)), nn.SiLU(), ) self.f0_upsampler = nn.Sequential( nn.Upsample(scale_factor=2, mode='linear', align_corners=False), weight_norm(nn.Conv1d(1, 1, 3, stride=1, padding=1)), ) self.f0_predictor = nn.Sequential( weight_norm(nn.Conv1d(prosody_latent_dim, hidden_dim, 3, padding=1)), nn.SiLU(), weight_norm(nn.Conv1d(hidden_dim, hidden_dim, 3, padding=1)), nn.SiLU(), weight_norm(nn.Conv1d(hidden_dim, hidden_dim // 2, 3, padding=1)), nn.SiLU(), weight_norm(nn.Conv1d(hidden_dim // 2, hidden_dim // 4, 3, padding=1)), nn.SiLU(), weight_norm(nn.Conv1d(hidden_dim // 4, 1, 3, padding=1)) ) self.m_source = SourceModuleHnNSF( sampling_rate=sample_rate, upsample_scale=self.source_upsample_rate, harmonic_num=14, voiced_threshold=0, ) self.f0_upsamp = nn.Upsample(scale_factor=self.source_upsample_rate) self.language_dim = language_dim fusion_dim = prosody_latent_dim + text_dim + language_dim # Speaker-conditioned fusion module self.pre_decoder = SpeakerFusionModule( dim_in=fusion_dim, hidden_dim=hidden_dim, speaker_dim=speaker_dim ) self.conformers = nn.ModuleList() for i in range(len(upsample_rates)): ch = hidden_dim // (2 ** i) self.conformers.append( Conformer( dim=ch, depth=2, dim_head=64, heads=8, ff_mult=4, conv_expansion_factor=2, conv_kernel_size=31, attn_dropout=0.1, ff_dropout=0.1, conv_dropout=0.1, ) ) self.snakes = nn.ModuleList() self.snakes.append(Snake1d(hidden_dim)) self.ups = nn.ModuleList() upsample_kernel_sizes = [2 * u for u in upsample_rates] for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)): in_ch = hidden_dim // (2 ** i) out_ch = hidden_dim // (2 ** (i + 1)) self.ups.append(UpsamplePixelShuffle1D(in_ch, out_ch, kernel_size=k, r=u)) self.snakes.append(Snake1d(out_ch)) self.noise_convs = nn.ModuleList() self.noise_res = nn.ModuleList() for i in range(len(upsample_rates)): c_cur = hidden_dim // (2 ** (i + 1)) if i + 1 < len(upsample_rates): stride_f0 = int(np.prod(upsample_rates[i + 1:])) self.noise_convs.append( weight_norm(nn.Conv1d( gen_istft_n_fft + 2, c_cur, kernel_size=stride_f0 * 2, stride=stride_f0, padding=(stride_f0 + 1) // 2, )) ) self.noise_res.append(SpeakerAdaINResBlock1(c_cur, 7, [1, 3, 5], speaker_dim)) else: self.noise_convs.append( weight_norm(nn.Conv1d(gen_istft_n_fft + 2, c_cur, kernel_size=1)) ) self.noise_res.append(SpeakerAdaINResBlock1(c_cur, 11, [1, 3, 5], speaker_dim)) # ResBlocks with speaker AdaIN conditioning self.resblocks = nn.ModuleList() for i in range(len(upsample_rates)): ch = hidden_dim // (2 ** (i + 1)) for k, d in zip(resblock_kernel_sizes, resblock_dilation_sizes): self.resblocks.append(SpeakerAdaINResBlock1(ch, k, d, speaker_dim)) self.post_n_fft = gen_istft_n_fft final_ch = hidden_dim // (2 ** len(upsample_rates)) self.conv_post = weight_norm(nn.Conv1d(final_ch, self.post_n_fft + 2, 7, padding=3)) self.stft = TorchSTFT( filter_length=gen_istft_n_fft, hop_length=gen_istft_hop_size, win_length=gen_istft_n_fft, ) self.reflection_pad = nn.ReflectionPad1d((1, 0)) def forward(self, prosody_latent, text_emb, speaker_emb, f0_gt=None, cache=None, language_emb=None): """ Args: prosody_latent: [B, prosody_dim, T_comp] text_emb: [B, text_dim, T] speaker_emb: [B, speaker_dim] - global speaker embedding f0_gt: [B, T] optional ground truth F0 cache: dict for streaming inference language_emb: [B, language_dim] optional Returns: wav, spec, phase, f0_pred, (new_cache if streaming) """ B = prosody_latent.shape[0] f0_pred_latent = self.f0_predictor(prosody_latent) f0_pred = self.f0_upsampler(f0_pred_latent) if f0_gt is not None: f0_pred = F.interpolate(f0_pred, size=f0_gt.shape[-1], mode='linear') else: target_len = int(prosody_latent.shape[-1] * self.codec_compression) if f0_pred.shape[-1] != target_len: f0_pred = F.interpolate(f0_pred, size=target_len, mode='linear') f0_pred = f0_pred.squeeze(1) f0_to_use = f0_gt if f0_gt is not None else f0_pred.detach() # Generate harmonic source f0_log = self.f0_upsamp(f0_to_use[:, None]).transpose(1, 2) f0_lin = (10.0 ** f0_log.float()).to(f0_log.dtype) source_phase_cache = cache.get("source_phase") if cache is not None else None har_source, noi_source, uv, next_source_phase = self.m_source(f0_lin, cache=source_phase_cache) har_source = har_source.transpose(1, 2).squeeze(1) har_spec, har_phase = self.stft.transform(har_source) har = torch.cat([har_spec, har_phase], dim=1) # Upsample prosody to match text prosody_latent = self.prosody_upsampler(prosody_latent) if prosody_latent.shape[-1] < text_emb.shape[-1]: pad_amount = text_emb.shape[-1] - prosody_latent.shape[-1] prosody_latent = F.pad(prosody_latent, (0, pad_amount), mode='replicate') prosody_latent = prosody_latent[..., :text_emb.shape[-1]] text_emb = text_emb[..., :prosody_latent.shape[-1]] # Speaker-conditioned fusion x = self.pre_decoder(prosody_latent, text_emb, speaker_emb, language_emb) for i in range(self.num_upsamples): x = self.snakes[i](x) x = rearrange(x, "b f t -> b t f") x = self.conformers[i](x) x = rearrange(x, "b t f -> b f t") x = self.ups[i](x) x_source = self.noise_convs[i](har) x_source = self.noise_res[i](x_source, speaker_emb) if i == self.num_upsamples - 1: x = self.reflection_pad(x) if x.shape[-1] != x_source.shape[-1]: min_len_add = min(x.shape[-1], x_source.shape[-1]) x = x[..., :min_len_add] x_source = x_source[..., :min_len_add] x = x + x_source xs = None for j in range(self.num_kernels): if xs is None: xs = self.resblocks[i * self.num_kernels + j](x, speaker_emb) else: xs += self.resblocks[i * self.num_kernels + j](x, speaker_emb) x = xs / self.num_kernels x = self.snakes[-1](x) x = self.conv_post(x) spec = torch.exp(x[:, :self.post_n_fft // 2 + 1, :]) phase = torch.sin(x[:, self.post_n_fft // 2 + 1:, :]) out = self.stft.inverse(spec, phase) if cache is not None: new_cache = { "source_phase": next_source_phase } return out, spec, phase, f0_pred, new_cache return out, spec, phase, f0_pred # ============================================================================== # Main Codec with Learnable Speaker Embeddings # ============================================================================== class HybridTTSCodecVocoderSpeaker(nn.Module): """ Hybrid TTS Codec with LEARNABLE SPEAKER EMBEDDINGS. Key features: - Learnable speaker embedding: nn.Embedding(num_speakers=11, embedding_dim=128) - Speaker IDs: 0-10 for 11 speakers - Speaker conditioning via AdaIN1d throughout the decoder - No mel-based style encoder - purely speaker ID based Usage: model = HybridTTSCodecVocoderSpeaker(num_speakers=11, speaker_dim=128, ...) output = model(pitch, energy, text_emb, speaker_ids=speaker_ids) # speaker_ids: [B] tensor with values 0-10 """ def __init__( self, num_speakers: int = 11, speaker_dim: int = 128, text_dim: int = 512, prosody_latent_dim: int = 512, hidden_dim: int = 512, codec_strides: List[int] = [2, 2], codebook_size: int = 4096, upsample_rates: List[int] = [12, 10], gen_istft_n_fft: int = 30, gen_istft_hop_size: int = 5, sample_rate: int = 44100, source_upsample_rate: int = 600, fsq_levels: Optional[List[int]] = None, language_dim: int = 0, ): super().__init__() self.num_speakers = num_speakers self.speaker_dim = speaker_dim self.text_dim = text_dim self.prosody_latent_dim = prosody_latent_dim self.codec_compression = math.prod(codec_strides) self.use_fsq = fsq_levels is not None self.fsq_levels = fsq_levels or [4] * 6 self.language_dim = language_dim # ===================================================================== # LEARNABLE SPEAKER EMBEDDING # 11 speakers (IDs 0-10), each with 128-dim embedding # ===================================================================== self.speaker_embedding = nn.Embedding( num_embeddings=num_speakers, embedding_dim=speaker_dim ) # Initialize with normal distribution nn.init.normal_(self.speaker_embedding.weight, mean=0, std=0.5) self.prosody_encoder = HybridProsodyEncoderSpeaker( speaker_dim=speaker_dim, latent_dim=prosody_latent_dim, hidden_dim=hidden_dim, strides=codec_strides, ) self.quantizer = FiniteScalarQuantization( input_dim=prosody_latent_dim, levels=self.fsq_levels, ) self.decoder = HybridWaveformDecoderSpeaker( prosody_latent_dim=prosody_latent_dim, text_dim=text_dim, speaker_dim=speaker_dim, hidden_dim=hidden_dim, upsample_rates=upsample_rates, gen_istft_n_fft=gen_istft_n_fft, gen_istft_hop_size=gen_istft_hop_size, sample_rate=sample_rate, source_upsample_rate=source_upsample_rate, codec_strides=codec_strides, language_dim=language_dim, ) def forward(self, pitch, energy, text_emb, speaker_ids, n_quantizers=None, use_predicted_f0=False, language_emb=None): """ Training forward pass. Args: pitch: [B, T] - pitch contour (log F0) energy: [B, T] - energy contour text_emb: [B, text_dim, T] - text embeddings speaker_ids: [B] - speaker IDs (0-10 for 11 speakers) n_quantizers: unused, for compatibility use_predicted_f0: bool - whether to use predicted F0 language_emb: [B, language_dim] optional Returns: dict with wav, tokens, speaker_emb, etc. """ # Get speaker embedding from ID speaker_emb = self.speaker_embedding(speaker_ids) # [B, speaker_dim] # Prosody encoder (no speaker - forces codebook usage) prosody_latent = self.prosody_encoder(pitch, energy) # Quantize prosody quantized_prosody, tokens, commitment_loss = self.quantizer(prosody_latent) decoder_f0 = None if use_predicted_f0 else pitch # Decode with speaker conditioning via AdaIN1d wav, mag, phase, f0_pred = self.decoder( quantized_prosody, text_emb, speaker_emb, # Speaker embedding passed to decoder f0_gt=decoder_f0, cache=None, language_emb=language_emb ) return { "wav": wav, "mag": mag, "phase": phase, "tokens": tokens, "prosody_latent": prosody_latent, "quantized_prosody": quantized_prosody, "text_down": text_emb, "speaker_emb": speaker_emb, "commitment_loss": commitment_loss, "f0_pred": f0_pred, "f0_gt": pitch, } def get_speaker_embedding(self, speaker_ids): """Get speaker embedding from IDs.""" return self.speaker_embedding(speaker_ids) @torch.no_grad() def tokenize(self, pitch, energy, text_emb, speaker_ids, n_quantizers=None): """Tokenize prosody.""" speaker_emb = self.speaker_embedding(speaker_ids) prosody_latent = self.prosody_encoder(pitch, energy) _, tokens, _ = self.quantizer(prosody_latent) return tokens, text_emb, speaker_emb @torch.no_grad() def decode_tokens(self, tokens, text_emb, speaker_ids, language_emb=None): """ Decode tokens with speaker ID. Args: tokens: [B, 1, T_comp] - prosody tokens text_emb: [B, text_dim, T] - text embeddings speaker_ids: [B] - speaker IDs (0-10) language_emb: optional """ speaker_emb = self.speaker_embedding(speaker_ids) quantized_prosody = self.quantizer.decode(tokens) wav, _, _, _ = self.decoder( quantized_prosody, text_emb, speaker_emb, f0_gt=None, cache=None, language_emb=language_emb ) return wav @torch.no_grad() def decode_tokens_with_speaker_emb(self, tokens, text_emb, speaker_emb, language_emb=None): """ Decode tokens with pre-computed speaker embedding. Useful for speaker interpolation or external speaker embeddings. Args: tokens: [B, 1, T_comp] - prosody tokens text_emb: [B, text_dim, T] - text embeddings speaker_emb: [B, speaker_dim] - speaker embedding (can be interpolated) language_emb: optional """ quantized_prosody = self.quantizer.decode(tokens) wav, _, _, _ = self.decoder( quantized_prosody, text_emb, speaker_emb, f0_gt=None, cache=None, language_emb=language_emb ) return wav @torch.no_grad() def decode_chunk(self, tokens, text_emb, speaker_ids, cache=None, language_emb=None): """ Streaming inference by chunk. Args: tokens: Chunk of tokens text_emb: Chunk of text embeddings speaker_ids: [B] speaker IDs cache: Dictionary from previous chunk call Returns: wav_chunk, new_cache """ if cache is None: cache = {} speaker_emb = self.speaker_embedding(speaker_ids) quantized_prosody = self.quantizer.decode(tokens) wav, _, _, _, new_cache = self.decoder( quantized_prosody, text_emb, speaker_emb, f0_gt=None, cache=cache, language_emb=language_emb ) return wav, new_cache @torch.no_grad() def interpolate_speakers(self, speaker_id_1, speaker_id_2, alpha=0.5): """ Interpolate between two speaker embeddings. Args: speaker_id_1: int - first speaker ID speaker_id_2: int - second speaker ID alpha: float - interpolation weight (0 = speaker_1, 1 = speaker_2) Returns: [1, speaker_dim] interpolated embedding """ emb1 = self.speaker_embedding(torch.tensor([speaker_id_1], device=self.speaker_embedding.weight.device)) emb2 = self.speaker_embedding(torch.tensor([speaker_id_2], device=self.speaker_embedding.weight.device)) return (1 - alpha) * emb1 + alpha * emb2 # ============================================================================== # Example Usage # ============================================================================== if __name__ == "__main__": # Test the model device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # Create model model = HybridTTSCodecVocoderSpeaker( num_speakers=11, speaker_dim=128, text_dim=512, prosody_latent_dim=512, hidden_dim=512, codec_strides=[2, 2], upsample_rates=[12, 10], gen_istft_n_fft=30, gen_istft_hop_size=5, sample_rate=44100, source_upsample_rate=600, fsq_levels=[4, 4, 4, 4, 4, 4], language_dim=0, ).to(device) print(f"Model created with {model.num_speakers} speakers, {model.speaker_dim}-dim embeddings") print(f"Speaker embedding shape: {model.speaker_embedding.weight.shape}") # Test forward pass batch_size = 2 seq_len = 100 pitch = torch.randn(batch_size, seq_len).to(device) energy = torch.randn(batch_size, seq_len).to(device) text_emb = torch.randn(batch_size, 512, seq_len * 2).to(device) speaker_ids = torch.randint(0, 11, (batch_size,)).to(device) # Random speaker IDs 0-10 print(f"\nTest inputs:") print(f" pitch: {pitch.shape}") print(f" energy: {energy.shape}") print(f" text_emb: {text_emb.shape}") print(f" speaker_ids: {speaker_ids}") # Forward pass output = model(pitch, energy, text_emb, speaker_ids) print(f"\nOutputs:") print(f" wav: {output['wav'].shape}") print(f" tokens: {output['tokens'].shape}") print(f" speaker_emb: {output['speaker_emb'].shape}") print(f" f0_pred: {output['f0_pred'].shape}") # Test speaker interpolation interp_emb = model.interpolate_speakers(0, 5, alpha=0.5) print(f"\nInterpolated speaker embedding (0 <-> 5): {interp_emb.shape}") # Count parameters total_params = sum(p.numel() for p in model.parameters()) speaker_params = model.speaker_embedding.weight.numel() print(f"\nTotal parameters: {total_params:,}") print(f"Speaker embedding parameters: {speaker_params:,} ({speaker_params/total_params*100:.2f}%)")