diff --git "a/modeling_spark_tts.py" "b/modeling_spark_tts.py" new file mode 100644--- /dev/null +++ "b/modeling_spark_tts.py" @@ -0,0 +1,3438 @@ +# coding=utf-8 +# Copyright 2024 The SparkAudio Authors and The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" PyTorch SparkTTS model.""" + +import torch +import torch.nn as nn +import numpy as np +import os +import warnings +from pathlib import Path +from typing import Dict, Any, Tuple, Optional, Union + +from transformers import PreTrainedModel, AutoModelForCausalLM, Wav2Vec2FeatureExtractor, Wav2Vec2Model +from transformers.utils import logging, requires_backends +from transformers.generation.utils import GenerationMixin +from transformers.configuration_utils import PretrainedConfig +from safetensors.torch import load_file +import torchaudio.transforms as TT # Directly use torchaudio + +# # Import necessary components from the original codebase structure +# # These are now defined in _modeling_bicodec_components.py +# from ._modeling_bicodec_components import ( +# SpeakerEncoder, +# Encoder, +# Decoder, +# WaveGenerator, +# FactorizedVectorQuantize, +# # Include Snake1d or other base classes if BiCodec.__init__ needs them directly +# ) + + +""" Utility functions for SparkTTS """ + +import random +import soxr +import soundfile +import torch +import torchaudio +import numpy as np + +from pathlib import Path +from typing import Tuple, Dict, Any +from numpy.lib.stride_tricks import sliding_window_view +from omegaconf import OmegaConf # Keep if BiCodec config loading needs it + + +# --- Token Maps (from sparktts/utils/token_parser.py) --- +TASK_TOKEN_MAP = { + "vc": "<|task_vc|>", + "tts": "<|task_tts|>", + "asr": "<|task_asr|>", + "s2s": "<|task_s2s|>", + "t2s": "<|task_t2s|>", + "understand": "<|task_understand|>", + "caption": "<|task_cap|>", + "controllable_tts": "<|task_controllable_tts|>", + "prompt_tts": "<|task_prompt_tts|>", + "speech_edit": "<|task_edit|>", +} + +LEVELS_MAP = { + "very_low": 0, + "low": 1, + "moderate": 2, + "high": 3, + "very_high": 4, +} + +LEVELS_MAP_UI = { + 1: 'very_low', + 2: 'low', + 3: 'moderate', + 4: 'high', + 5: 'very_high' +} + +GENDER_MAP = { + "female": 0, + "male": 1, +} + +# --- Audio Utils (from sparktts/utils/audio.py) --- +def audio_volume_normalize(audio: np.ndarray, coeff: float = 0.2) -> np.ndarray: + temp = np.sort(np.abs(audio)) + if len(temp) == 0: # Handle empty audio case + return audio + if temp[-1] < 0.1: + scaling_factor = max(temp[-1], 1e-3) + audio = audio / scaling_factor * 0.1 + temp = temp[temp > 0.01] + L = temp.shape[0] + if L <= 10: + return audio + volume = np.mean(temp[int(0.9 * L) : int(0.99 * L)]) + if volume == 0: # Avoid division by zero if volume is effectively zero + return audio + audio = audio * np.clip(coeff / volume, a_min=0.1, a_max=10) + max_value = np.max(np.abs(audio)) if len(audio) > 0 else 0 + if max_value > 1: + audio = audio / max_value + return audio + +def load_audio( + adfile: Path, + sampling_rate: int = None, + length: int = None, + volume_normalize: bool = False, + segment_duration: int = None, +) -> np.ndarray: + try: + audio, sr = soundfile.read(adfile, dtype='float32') # Ensure float32 + except Exception as e: + raise IOError(f"Could not read audio file {adfile}: {e}") + + if audio is None or len(audio) == 0: + raise ValueError(f"Audio file {adfile} is empty or invalid.") + + if len(audio.shape) > 1: + audio = audio[:, 0] + + if sampling_rate is not None and sr != sampling_rate: + try: + # Ensure input is float64 for soxr + audio = audio.astype(np.float64) + audio = soxr.resample(audio, sr, sampling_rate, quality="VHQ") + # Convert back to float32 + audio = audio.astype(np.float32) + sr = sampling_rate + except Exception as e: + raise RuntimeError(f"Failed to resample audio from {sr}Hz to {sampling_rate}Hz: {e}") + + if segment_duration is not None: + seg_length = int(sr * segment_duration) + audio = random_select_audio_segment(audio, seg_length) + + if volume_normalize: + audio = audio_volume_normalize(audio) + + if length is not None: + if audio.shape[0] > length: + audio = audio[:length] + else: + audio = np.pad(audio, (0, int(length - audio.shape[0])), mode='constant') + return audio + +def random_select_audio_segment(audio: np.ndarray, length: int) -> np.ndarray: + if audio.shape[0] < length: + audio = np.pad(audio, (0, int(length - audio.shape[0])), mode='constant') + start_index = 0 # If padded, start from beginning + elif audio.shape[0] == length: + start_index = 0 # If exact length, start from beginning + else: + start_index = random.randint(0, audio.shape[0] - length) + + end_index = int(start_index + length) + return audio[start_index:end_index] + +# --- File Utils (Minimal required) --- +def load_config_yaml(config_path: Path) -> Dict: + """Loads a YAML configuration file using OmegaConf.""" + # Check if path exists + if not Path(config_path).is_file(): + raise FileNotFoundError(f"YAML Config file not found: {config_path}") + try: + config = OmegaConf.load(config_path) + # Convert OmegaConf DictConfig to standard Python dict + return OmegaConf.to_container(config, resolve=True) + except Exception as e: + raise IOError(f"Error loading YAML config file {config_path}: {e}") + + +""" PyTorch SparkTTS BiCodec sub-module definitions.""" + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.distributed as dist +import random + +from torch.nn.utils import weight_norm, remove_weight_norm +from torch import Tensor, int32 +from torch.amp import autocast + +from typing import Any, Dict, List, Tuple, Optional +from collections import namedtuple +from functools import wraps, partial +from contextlib import nullcontext +from packaging import version + +from einops import rearrange, repeat, reduce, pack, unpack +from einops.layers.torch import Rearrange +from einx import get_at # Ensure einx is installed: pip install einx + +# =============================================================== +# Start: Content from sparktts/modules/blocks/layers.py +# =============================================================== +def WNConv1d(*args, **kwargs): + return weight_norm(nn.Conv1d(*args, **kwargs)) + + +def WNConvTranspose1d(*args, **kwargs): + return weight_norm(nn.ConvTranspose1d(*args, **kwargs)) + + +# Scripting this brings model speed up 1.4x +@torch.jit.script +def snake(x, alpha): + shape = x.shape + x = x.reshape(shape[0], shape[1], -1) + x = x + (alpha + 1e-9).reciprocal() * torch.sin(alpha * x).pow(2) + x = x.reshape(shape) + return x + + +class Snake1d(nn.Module): + def __init__(self, channels): + super().__init__() + self.alpha = nn.Parameter(torch.ones(1, channels, 1)) + + def forward(self, x): + return snake(x, self.alpha) + + +class ResidualUnit(nn.Module): + def __init__(self, dim: int = 16, dilation: int = 1): + super().__init__() + pad = ((7 - 1) * dilation) // 2 + self.block = nn.Sequential( + Snake1d(dim), + WNConv1d(dim, dim, kernel_size=7, dilation=dilation, padding=pad), + Snake1d(dim), + WNConv1d(dim, dim, kernel_size=1), + ) + + def forward(self, x): + y = self.block(x) + # Adjust padding handling if input and output shapes differ + diff = x.shape[-1] - y.shape[-1] + if diff > 0: + pad = diff // 2 + x = x[..., pad:pad + y.shape[-1]] # Ensure shapes match for residual connection + elif diff < 0: + pad = -diff // 2 + y = y[..., pad:pad + x.shape[-1]] + + return x + y + + +def init_weights(m): + if isinstance(m, nn.Conv1d): + nn.init.trunc_normal_(m.weight, std=0.02) + if m.bias is not None: + nn.init.constant_(m.bias, 0) +# =============================================================== +# End: Content from sparktts/modules/blocks/layers.py +# =============================================================== + + +# =============================================================== +# Start: Content from sparktts/modules/blocks/samper.py +# =============================================================== +class SamplingBlock(nn.Module): + """Sampling block for upsampling or downsampling""" + + def __init__( + self, + dim: int, + groups: int = 1, + upsample_scale: int = 1, + downsample_scale: int = 1, + ) -> None: + """ + Args: + dim: input dimension + groups: number of groups + upsample_scale: upsampling scale + downsample_scale: downsampling scale + """ + super(SamplingBlock, self).__init__() + + self.upsample_scale = upsample_scale + self.downsample_scale = downsample_scale + + if self.upsample_scale > 1: + self.de_conv_upsampler = nn.Sequential( + nn.LeakyReLU(0.2), + nn.ConvTranspose1d( + dim, + dim, + kernel_size=upsample_scale * 2, + stride=upsample_scale, + padding=upsample_scale // 2 + upsample_scale % 2, + output_padding=upsample_scale % 2, + groups=groups, + ), + ) + + if self.downsample_scale > 1: + self.conv_downsampler = nn.Sequential( + nn.LeakyReLU(0.2), + nn.Conv1d( + dim, + dim, + kernel_size=2 * downsample_scale, + stride=downsample_scale, + padding=downsample_scale // 2 + downsample_scale % 2, + groups=groups, + ), + ) + + @staticmethod + def repeat_upsampler(x, upsample_scale): + return x.repeat_interleave(upsample_scale, dim=2) + + @staticmethod + def skip_downsampler(x, downsample_scale): + return F.avg_pool1d(x, kernel_size=downsample_scale, stride=downsample_scale) + + def forward(self, x): + # Input expected as (B, D, T) from VocosBackbone output (B, T, D) + # x = x.transpose(1, 2) # Remove this transpose, input should be (B, D, T) + if self.upsample_scale > 1: + repeat_res = self.repeat_upsampler(x, self.upsample_scale) + deconv_res = self.de_conv_upsampler(x) + # Ensure shapes match for addition + if deconv_res.shape[-1] > repeat_res.shape[-1]: + deconv_res = deconv_res[..., :repeat_res.shape[-1]] + elif repeat_res.shape[-1] > deconv_res.shape[-1]: + repeat_res = repeat_res[..., :deconv_res.shape[-1]] + upmerge_res = repeat_res + deconv_res + else: + upmerge_res = x + repeat_res = x + + if self.downsample_scale > 1: + conv_res = self.conv_downsampler(upmerge_res) + skip2_res = self.skip_downsampler(upmerge_res, self.downsample_scale) + skip1_res = self.skip_downsampler(repeat_res, self.downsample_scale) + # Ensure shapes match + min_len = min(conv_res.shape[-1], skip1_res.shape[-1], skip2_res.shape[-1]) + conv_res = conv_res[..., :min_len] + skip1_res = skip1_res[..., :min_len] + skip2_res = skip2_res[..., :min_len] + else: + conv_res = upmerge_res + skip2_res = upmerge_res + skip1_res = repeat_res + + final_res = conv_res + skip1_res + skip2_res + + # Return (B, D, T) for next VocosBackbone + # return final_res.transpose(1, 2) # Remove this, keep (B, D, T) + return final_res + +# =============================================================== +# End: Content from sparktts/modules/blocks/samper.py +# =============================================================== + + +# =============================================================== +# Start: Content from sparktts/modules/speaker/pooling_layers.py +# =============================================================== +class TAP(nn.Module): + """ + Temporal average pooling, only first-order mean is considered + """ + + def __init__(self, in_dim=0, **kwargs): + super(TAP, self).__init__() + self.in_dim = in_dim + + def forward(self, x): + pooling_mean = x.mean(dim=-1) + # To be compatable with 2D input + pooling_mean = pooling_mean.flatten(start_dim=1) + return pooling_mean + + def get_out_dim(self): + # This method seems specific to the original usage, might not be needed by HF + # self.out_dim = self.in_dim + # return self.out_dim + return self.in_dim + + +class TSDP(nn.Module): + """ + Temporal standard deviation pooling, only second-order std is considered + """ + + def __init__(self, in_dim=0, **kwargs): + super(TSDP, self).__init__() + self.in_dim = in_dim + + def forward(self, x): + # The last dimension is the temporal axis + pooling_std = torch.sqrt(torch.var(x, dim=-1) + 1e-7) + pooling_std = pooling_std.flatten(start_dim=1) + return pooling_std + + def get_out_dim(self): + # self.out_dim = self.in_dim + # return self.out_dim + return self.in_dim + + +class TSTP(nn.Module): + """ + Temporal statistics pooling, concatenate mean and std, which is used in + x-vector + Comment: simple concatenation can not make full use of both statistics + """ + + def __init__(self, in_dim=0, **kwargs): + super(TSTP, self).__init__() + self.in_dim = in_dim + + def forward(self, x): + # The last dimension is the temporal axis + pooling_mean = x.mean(dim=-1) + pooling_std = torch.sqrt(torch.var(x, dim=-1) + 1e-7) + pooling_mean = pooling_mean.flatten(start_dim=1) + pooling_std = pooling_std.flatten(start_dim=1) + stats = torch.cat((pooling_mean, pooling_std), 1) + return stats + + def get_out_dim(self): + # self.out_dim = self.in_dim * 2 + # return self.out_dim + return self.in_dim * 2 + + +class ASTP(nn.Module): + """ Attentive statistics pooling: Channel- and context-dependent + statistics pooling, first used in ECAPA_TDNN. + """ + + def __init__(self, + in_dim, + bottleneck_dim=128, + global_context_att=False, + **kwargs): + super(ASTP, self).__init__() + self.in_dim = in_dim + self.global_context_att = global_context_att + + # Use Conv1d with stride == 1 rather than Linear, then we don't + # need to transpose inputs. + if global_context_att: + self.linear1 = nn.Conv1d( + in_dim * 3, bottleneck_dim, + kernel_size=1) # equals W and b in the paper + else: + self.linear1 = nn.Conv1d( + in_dim, bottleneck_dim, + kernel_size=1) # equals W and b in the paper + self.linear2 = nn.Conv1d(bottleneck_dim, in_dim, + kernel_size=1) # equals V and k in the paper + + def forward(self, x): + """ + x: a 3-dimensional tensor in tdnn-based architecture (B,F,T) + or a 4-dimensional tensor in resnet architecture (B,C,F,T) + 0-dim: batch-dimension, last-dim: time-dimension (frame-dimension) + """ + if len(x.shape) == 4: + x = x.reshape(x.shape[0], x.shape[1] * x.shape[2], x.shape[3]) + assert len(x.shape) == 3 + + if self.global_context_att: + context_mean = torch.mean(x, dim=-1, keepdim=True).expand_as(x) + context_std = torch.sqrt( + torch.var(x, dim=-1, keepdim=True) + 1e-7).expand_as(x) + x_in = torch.cat((x, context_mean, context_std), dim=1) + else: + x_in = x + + # DON'T use ReLU here! ReLU may be hard to converge. + alpha = torch.tanh( + self.linear1(x_in)) # alpha = F.relu(self.linear1(x_in)) + alpha = torch.softmax(self.linear2(alpha), dim=2) + mean = torch.sum(alpha * x, dim=2) + var = torch.sum(alpha * (x**2), dim=2) - mean**2 + std = torch.sqrt(var.clamp(min=1e-7)) + return torch.cat([mean, std], dim=1) + + def get_out_dim(self): + # self.out_dim = 2 * self.in_dim + # return self.out_dim + return self.in_dim * 2 + + +class MHASTP(torch.nn.Module): + """ Multi head attentive statistics pooling + Reference: + Self Multi-Head Attention for Speaker Recognition + https://arxiv.org/pdf/1906.09890.pdf + """ + + def __init__(self, + in_dim, + layer_num=2, + head_num=2, + d_s=1, + bottleneck_dim=64, + **kwargs): + super(MHASTP, self).__init__() + assert (in_dim % head_num + ) == 0 # make sure that head num can be divided by input_dim + self.in_dim = in_dim + self.head_num = head_num + d_model = int(in_dim / head_num) + channel_dims = [bottleneck_dim for i in range(layer_num + 1)] + if d_s > 1: + d_s = d_model + else: + d_s = 1 + self.d_s = d_s + channel_dims[0], channel_dims[-1] = d_model, d_s + heads_att_trans = [] + for i in range(self.head_num): + att_trans = nn.Sequential() + for j in range(layer_num - 1): # Use different loop variable + att_trans.add_module( + 'att_' + str(j), + nn.Conv1d(channel_dims[j], channel_dims[j + 1], 1, 1)) + att_trans.add_module('tanh' + str(j), nn.Tanh()) + att_trans.add_module( + 'att_' + str(layer_num - 1), + nn.Conv1d(channel_dims[layer_num - 1], channel_dims[layer_num], + 1, 1)) + heads_att_trans.append(att_trans) + self.heads_att_trans = nn.ModuleList(heads_att_trans) + + def forward(self, input): + """ + input: a 3-dimensional tensor in xvector architecture + or a 4-dimensional tensor in resnet architecture + 0-dim: batch-dimension, last-dim: time-dimension (frame-dimension) + """ + if len(input.shape) == 4: # B x C x F x T + input = input.reshape(input.shape[0], + input.shape[1] * input.shape[2], + input.shape[3]) # B x (C*F) x T + assert len(input.shape) == 3 + bs, f_dim, t_dim = input.shape + chunks = torch.chunk(input, self.head_num, 1) + # split + chunks_out = [] + for i, layer in enumerate(self.heads_att_trans): + att_score = layer(chunks[i]) + alpha = F.softmax(att_score, dim=-1) + mean = torch.sum(alpha * chunks[i], dim=2) + var = torch.sum(alpha * chunks[i]**2, dim=2) - mean**2 + std = torch.sqrt(var.clamp(min=1e-7)) + chunks_out.append(torch.cat((mean, std), dim=1)) + out = torch.cat(chunks_out, dim=1) + return out + + def get_out_dim(self): + # self.out_dim = 2 * self.in_dim + # return self.out_dim + return self.in_dim * 2 + + +class MQMHASTP(torch.nn.Module): + """ An attentive pooling + Reference: + multi query multi head attentive statistics pooling + https://arxiv.org/pdf/2110.05042.pdf + Args: + in_dim: the feature dimension of input + layer_num: the number of layer in the pooling layer + query_num: the number of querys + head_num: the number of heads + bottleneck_dim: the bottleneck dimension + + SA (H = 1, Q = 1, n = 2, d_s = 1) ref: + https://www.danielpovey.com/files/2018_interspeech_xvector_attention.pdf + MHA (H > 1, Q = 1, n = 1, d_s = 1) ref: + https://arxiv.org/pdf/1906.09890.pdf + AS (H = 1, Q > 1, n = 2, d_s = 1) ref: + https://arxiv.org/pdf/1803.10963.pdf + VSA (H = 1, Q > 1, n = 2, d_s = d_h) ref: + http://www.interspeech2020.org/uploadfile/pdf/Mon-2-10-5.pdf + """ + + def __init__(self, + in_dim, + layer_num=2, + query_num=2, + head_num=8, + d_s=2, + bottleneck_dim=64, + **kwargs): + super(MQMHASTP, self).__init__() + self.n_query = nn.ModuleList([ + MHASTP(in_dim, + layer_num=layer_num, + head_num=head_num, + d_s=d_s, + bottleneck_dim=bottleneck_dim) for i in range(query_num) + ]) + self.query_num = query_num + self.in_dim = in_dim + + def forward(self, input): + """ + input: a 3-dimensional tensor in xvector architecture + or a 4-dimensional tensor in resnet architecture + 0-dim: batch-dimension, last-dim: time-dimension (frame-dimension) + """ + if len(input.shape) == 4: # B x C x F x T + input = input.reshape(input.shape[0], + input.shape[1] * input.shape[2], + input.shape[3]) # B x (C*F) x T + assert len(input.shape) == 3 + res = [] + for i, layer in enumerate(self.n_query): + res.append(layer(input)) + out = torch.cat(res, dim=-1) + return out + + def get_out_dim(self): + # self.out_dim = self.in_dim * 2 * self.query_num + # return self.out_dim + return self.in_dim * 2 * self.query_num + +# =============================================================== +# End: Content from sparktts/modules/speaker/pooling_layers.py +# =============================================================== + + +# =============================================================== +# Start: Content from sparktts/modules/blocks/vocos.py +# =============================================================== +# Helper functions needed by VocosBackbone etc. +def exists(val): + return val is not None + +def default(val, d): + return val if exists(val) else d() if callable(d) else d + +class AdaLayerNorm(nn.Module): + """ + Adaptive Layer Normalization module with learnable embeddings per `num_embeddings` classes + + Args: + condition_dim (int): Dimension of the condition. + embedding_dim (int): Dimension of the embeddings. + """ + + def __init__(self, condition_dim: int, embedding_dim: int, eps: float = 1e-6): + super().__init__() + self.eps = eps + self.dim = embedding_dim + self.scale = nn.Linear(condition_dim, embedding_dim) + self.shift = nn.Linear(condition_dim, embedding_dim) + # Initialize weights similar to original implementation if needed + # torch.nn.init.ones_(self.scale.weight) # Might be default + # torch.nn.init.zeros_(self.shift.weight) # Might be default + if self.scale.bias is not None: nn.init.zeros_(self.scale.bias) + if self.shift.bias is not None: nn.init.zeros_(self.shift.bias) + + def forward(self, x: torch.Tensor, cond_embedding: torch.Tensor) -> torch.Tensor: + scale = self.scale(cond_embedding) + shift = self.shift(cond_embedding) + x = nn.functional.layer_norm(x, (self.dim,), eps=self.eps) + x = x * scale.unsqueeze(1) + shift.unsqueeze(1) + return x + + +class ConvNeXtBlock(nn.Module): + """ConvNeXt Block adapted from https://github.com/facebookresearch/ConvNeXt to 1D audio signal. + + Args: + dim (int): Number of input channels. + intermediate_dim (int): Dimensionality of the intermediate layer. + layer_scale_init_value (float, optional): Initial value for the layer scale. None means no scaling. + Defaults to None. + condition_dim (int, optional): Dimension for AdaLayerNorm. + None means non-conditional LayerNorm. Defaults to None. + """ + + def __init__( + self, + dim: int, + intermediate_dim: int, + layer_scale_init_value: float, + condition_dim: Optional[int] = None, + ): + super().__init__() + self.dwconv = nn.Conv1d( + dim, dim, kernel_size=7, padding=3, groups=dim + ) # depthwise conv + self.adanorm = condition_dim is not None + if self.adanorm: + self.norm = AdaLayerNorm(condition_dim, dim, eps=1e-6) + else: + self.norm = nn.LayerNorm(dim, eps=1e-6) + self.pwconv1 = nn.Linear( + dim, intermediate_dim + ) # pointwise/1x1 convs, implemented with linear layers + self.act = nn.GELU() + self.pwconv2 = nn.Linear(intermediate_dim, dim) + self.gamma = ( + nn.Parameter(layer_scale_init_value * torch.ones(dim), requires_grad=True) + if layer_scale_init_value is not None and layer_scale_init_value > 0 + else None + ) + + def forward( + self, x: torch.Tensor, cond_embedding: Optional[torch.Tensor] = None + ) -> torch.Tensor: + residual = x + x = self.dwconv(x) + x = x.transpose(1, 2) # (B, C, T) -> (B, T, C) + if self.adanorm: + assert cond_embedding is not None, "Conditioning embedding required for AdaLayerNorm" + x = self.norm(x, cond_embedding) + else: + x = self.norm(x) + x = self.pwconv1(x) + x = self.act(x) + x = self.pwconv2(x) + if self.gamma is not None: + x = self.gamma * x + x = x.transpose(1, 2) # (B, T, C) -> (B, C, T) + + x = residual + x + return x + + +class ResBlock1(nn.Module): + """ + ResBlock adapted from HiFi-GAN V1 (https://github.com/jik876/hifi-gan) with dilated 1D convolutions, + but without upsampling layers. + + Args: + dim (int): Number of input channels. + kernel_size (int, optional): Size of the convolutional kernel. Defaults to 3. + dilation (tuple[int], optional): Dilation factors for the dilated convolutions. + Defaults to (1, 3, 5). + lrelu_slope (float, optional): Negative slope of the LeakyReLU activation function. + Defaults to 0.1. + layer_scale_init_value (float, optional): Initial value for the layer scale. None means no scaling. + Defaults to None. + """ + + def __init__( + self, + dim: int, + kernel_size: int = 3, + dilation: Tuple[int, int, int] = (1, 3, 5), + lrelu_slope: float = 0.1, + layer_scale_init_value: Optional[float] = None, + ): + super().__init__() + self.lrelu_slope = lrelu_slope + self.convs1 = nn.ModuleList( + [ + weight_norm( + nn.Conv1d( + dim, + dim, + kernel_size, + 1, + dilation=dilation[0], + padding=self.get_padding(kernel_size, dilation[0]), + ) + ), + weight_norm( + nn.Conv1d( + dim, + dim, + kernel_size, + 1, + dilation=dilation[1], + padding=self.get_padding(kernel_size, dilation[1]), + ) + ), + weight_norm( + nn.Conv1d( + dim, + dim, + kernel_size, + 1, + dilation=dilation[2], + padding=self.get_padding(kernel_size, dilation[2]), + ) + ), + ] + ) + + self.convs2 = nn.ModuleList( + [ + weight_norm( + nn.Conv1d( + dim, + dim, + kernel_size, + 1, + dilation=1, + padding=self.get_padding(kernel_size, 1), + ) + ), + weight_norm( + nn.Conv1d( + dim, + dim, + kernel_size, + 1, + dilation=1, + padding=self.get_padding(kernel_size, 1), + ) + ), + weight_norm( + nn.Conv1d( + dim, + dim, + kernel_size, + 1, + dilation=1, + padding=self.get_padding(kernel_size, 1), + ) + ), + ] + ) + + self.gamma = nn.ParameterList( + [ + ( + nn.Parameter( + layer_scale_init_value * torch.ones(dim, 1), requires_grad=True + ) + if layer_scale_init_value is not None + else None + ), + ( + nn.Parameter( + layer_scale_init_value * torch.ones(dim, 1), requires_grad=True + ) + if layer_scale_init_value is not None + else None + ), + ( + nn.Parameter( + layer_scale_init_value * torch.ones(dim, 1), requires_grad=True + ) + if layer_scale_init_value is not None + else None + ), + ] + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + for c1, c2, gamma in zip(self.convs1, self.convs2, self.gamma): + xt = torch.nn.functional.leaky_relu(x, negative_slope=self.lrelu_slope) + xt = c1(xt) + xt = torch.nn.functional.leaky_relu(xt, negative_slope=self.lrelu_slope) + xt = c2(xt) + if gamma is not None: + xt = gamma * xt + x = xt + x + return x + + def remove_weight_norm(self): + for l in self.convs1: + remove_weight_norm(l) + for l in self.convs2: + remove_weight_norm(l) + + @staticmethod + def get_padding(kernel_size: int, dilation: int = 1) -> int: + return int((kernel_size * dilation - dilation) / 2) + + +class Backbone(nn.Module): + """Base class for the generator's backbone. It preserves the same temporal resolution across all layers.""" + + def forward(self, x: torch.Tensor, **kwargs) -> torch.Tensor: + """ + Args: + x (Tensor): Input tensor of shape (B, C, L), where B is the batch size, + C denotes input features, and L is the sequence length. + + Returns: + Tensor: Output of shape (B, L, H), where B is the batch size, L is the sequence length, + and H denotes the model dimension. + """ + raise NotImplementedError("Subclasses must implement the forward method.") + + +class VocosBackbone(Backbone): + """ + Vocos backbone module built with ConvNeXt blocks. Supports additional conditioning with Adaptive Layer Normalization + + Args: + input_channels (int): Number of input features channels. + dim (int): Hidden dimension of the model. + intermediate_dim (int): Intermediate dimension used in ConvNeXtBlock. + num_layers (int): Number of ConvNeXtBlock layers. + layer_scale_init_value (float, optional): Initial value for layer scaling. Defaults to `1 / num_layers`. + condition_dim (int, optional): Dimension for AdaLayerNorm. + None means non-conditional model. Defaults to None. + """ + + def __init__( + self, + input_channels: int, + dim: int, + intermediate_dim: int, + num_layers: int, + layer_scale_init_value: Optional[float] = None, + condition_dim: Optional[int] = None, + ): + super().__init__() + self.input_channels = input_channels + self.embed = nn.Conv1d(input_channels, dim, kernel_size=7, padding=3) + self.adanorm = condition_dim is not None + if self.adanorm: + self.norm = AdaLayerNorm(condition_dim, dim, eps=1e-6) + else: + self.norm = nn.LayerNorm(dim, eps=1e-6) + layer_scale_init_value = layer_scale_init_value or 1 / num_layers if num_layers > 0 else None # Handle num_layers=0 + self.convnext = nn.ModuleList( + [ + ConvNeXtBlock( + dim=dim, + intermediate_dim=intermediate_dim, + layer_scale_init_value=layer_scale_init_value, + condition_dim=condition_dim, + ) + for _ in range(num_layers) + ] + ) + self.final_layer_norm = nn.LayerNorm(dim, eps=1e-6) + self.apply(self._init_weights) + + def _init_weights(self, m): + if isinstance(m, (nn.Conv1d, nn.Linear)): + nn.init.trunc_normal_(m.weight, std=0.02) + if m.bias is not None: + nn.init.constant_(m.bias, 0) + + def forward(self, x: torch.Tensor, condition: Optional[torch.Tensor] = None) -> torch.Tensor: + # Input x: (B, C, L) + x = self.embed(x) + # After embed: (B, dim, L) + x_transposed = x.transpose(1, 2) # (B, L, dim) + if self.adanorm: + assert condition is not None + norm_out = self.norm(x_transposed, condition) + else: + norm_out = self.norm(x_transposed) + # After norm: (B, L, dim) + x = norm_out.transpose(1, 2) # (B, dim, L) + for conv_block in self.convnext: + x = conv_block(x, condition) + # After convnext blocks: (B, dim, L) + x = self.final_layer_norm(x.transpose(1, 2)) # (B, L, dim) + return x + + +class VocosResNetBackbone(Backbone): + """ + Vocos backbone module built with ResBlocks. + + Args: + input_channels (int): Number of input features channels. + dim (int): Hidden dimension of the model. + num_blocks (int): Number of ResBlock1 blocks. + layer_scale_init_value (float, optional): Initial value for layer scaling. Defaults to None. + """ + + def __init__( + self, + input_channels, + dim, + num_blocks, + layer_scale_init_value=None, + ): + super().__init__() + self.input_channels = input_channels + self.embed = weight_norm( + nn.Conv1d(input_channels, dim, kernel_size=3, padding=1) + ) + layer_scale_init_value = layer_scale_init_value or 1 / num_blocks / 3 if num_blocks > 0 else None # Handle num_blocks=0 + self.resnet = nn.Sequential( + *[ + ResBlock1(dim=dim, layer_scale_init_value=layer_scale_init_value) + for _ in range(num_blocks) + ] + ) + + def forward(self, x: torch.Tensor, **kwargs) -> torch.Tensor: + # Input x: (B, C, L) + x = self.embed(x) + # After embed: (B, dim, L) + x = self.resnet(x) + # After resnet: (B, dim, L) + x = x.transpose(1, 2) # (B, L, dim) + return x + +# =============================================================== +# End: Content from sparktts/modules/blocks/vocos.py +# =============================================================== + + +# =============================================================== +# Start: Content from sparktts/modules/encoder_decoder/feat_decoder.py +# =============================================================== +class Decoder(nn.Module): + """Decoder module with convnext and upsampling blocks + + Args: + sample_ratios (List[int]): sample ratios + example: [2, 2] means upsample by 2x and then upsample by 2x + """ + + def __init__( + self, + input_channels: int, + vocos_dim: int, + vocos_intermediate_dim: int, + vocos_num_layers: int, + out_channels: int, + condition_dim: int = None, + sample_ratios: List[int] = [1, 1], + use_tanh_at_final: bool = False, + ): + super().__init__() + + self.linear_pre = nn.Linear(input_channels, vocos_dim) + + upsample_modules = [] + current_dim = vocos_dim + for i, ratio in enumerate(sample_ratios): + upsample_modules.append( + nn.Sequential( + SamplingBlock( + dim=current_dim, + groups=current_dim, # Maybe use 1 or fewer groups if dim is high? Check original intent. Using current_dim for now. + upsample_scale=ratio, + ), + # Note: The original code used VocosBackbone here, but it changes dims B,T,D -> B,D,T. + # SamplingBlock output is B,D,T, so VocosBackbone input matches. + # However, the VocosBackbone output is B,T,D, which doesn't fit the next SamplingBlock. + # Assuming the intent was to keep B,D,T format between sampling blocks. + # Replacing intermediate VocosBackbone with a simple Conv1d block to maintain format & refine. + nn.Conv1d(current_dim, current_dim, kernel_size=3, padding=1) # Simple refinement layer + # VocosBackbone( + # input_channels=current_dim, + # dim=current_dim, + # intermediate_dim=vocos_intermediate_dim // 2, # Smaller intermediate for efficiency? + # num_layers=2, # Fewer layers + # condition_dim=None, + # ) + ) + ) + # No dimension change expected here if using Conv1d refinement + # If using VocosBackbone, need transpose logic + + self.upsample = nn.Sequential(*upsample_modules) + + # Final Backbone processes the fully upsampled features + self.vocos_backbone = VocosBackbone( + input_channels=current_dim, # Use the dim after upsampling + dim=vocos_dim, # Map back to main vocos_dim or keep current_dim? Using vocos_dim + intermediate_dim=vocos_intermediate_dim, + num_layers=vocos_num_layers, + condition_dim=condition_dim, + ) + self.linear_post = nn.Linear(vocos_dim, out_channels) + self.use_tanh_at_final = use_tanh_at_final + + def forward(self, x: torch.Tensor, c: torch.Tensor = None): + """decoder forward. + + Args: + x (torch.Tensor): (batch_size, input_channels, length) + c (torch.Tensor): (batch_size, condition_dim) - Optional condition + + Returns: + x (torch.Tensor): (batch_size, out_channels, length_upsampled) + """ + # x: (B, C_in, T) + x = self.linear_pre(x.transpose(1, 2)) # (B, T, vocos_dim) + x = x.transpose(1, 2) # (B, vocos_dim, T) + + # Apply upsampling blocks + x = self.upsample(x) # (B, vocos_dim, T_upsampled) + + # Apply final backbone + x = self.vocos_backbone(x, condition=c) # (B, T_upsampled, vocos_dim) + + x = self.linear_post(x) # (B, T_upsampled, C_out) + x = x.transpose(1, 2) # (B, C_out, T_upsampled) + + if self.use_tanh_at_final: + x = torch.tanh(x) + + return x + +# =============================================================== +# End: Content from sparktts/modules/encoder_decoder/feat_decoder.py +# =============================================================== + + +# =============================================================== +# Start: Content from sparktts/modules/encoder_decoder/feat_encoder.py +# =============================================================== +class Encoder(nn.Module): + """Encoder module with convnext and downsampling blocks""" + + def __init__( + self, + input_channels: int, + vocos_dim: int, + vocos_intermediate_dim: int, + vocos_num_layers: int, + out_channels: int, + sample_ratios: List[int] = [1, 1], + ): + super().__init__() + """ + Encoder module with VocosBackbone and sampling blocks. + + Args: + sample_ratios (List[int]): sample ratios + example: [2, 2] means downsample by 2x and then downsample by 2x + """ + # Initial Backbone processing + self.encoder_backbone = VocosBackbone( + input_channels=input_channels, + dim=vocos_dim, + intermediate_dim=vocos_intermediate_dim, + num_layers=vocos_num_layers, # Use main num_layers here + condition_dim=None, + ) + + downsample_modules = [] + current_dim = vocos_dim + for i, ratio in enumerate(sample_ratios): + downsample_modules.append( + nn.Sequential( + SamplingBlock( + dim=current_dim, + groups=current_dim, # Again, check group size. Using current_dim. + downsample_scale=ratio, + ), + # Add refinement layer (optional, similar to Decoder logic) + nn.Conv1d(current_dim, current_dim, kernel_size=3, padding=1) + # VocosBackbone( # Or a lighter VocosBackbone + # input_channels=current_dim, + # dim=current_dim, + # intermediate_dim=vocos_intermediate_dim // 2, + # num_layers=2, + # condition_dim=None, + # ) + ) + ) + # No dimension change expected here + + self.downsample = nn.Sequential(*downsample_modules) + + self.project = nn.Linear(current_dim, out_channels) # Project from the final dimension + + def forward(self, x: torch.Tensor, *args): + """ + Args: + x (torch.Tensor): (batch_size, input_channels, length) + + Returns: + x (torch.Tensor): (batch_size, out_channels, length_downsampled) + """ + # x: (B, C_in, T) + x = self.encoder_backbone(x) # (B, T, vocos_dim) + x = x.transpose(1, 2) # (B, vocos_dim, T) + + # Apply downsampling blocks + x = self.downsample(x) # (B, vocos_dim, T_downsampled) + + x = x.transpose(1, 2) # (B, T_downsampled, vocos_dim) + x = self.project(x) # (B, T_downsampled, C_out) + return x.transpose(1, 2) # (B, C_out, T_downsampled) + +# =============================================================== +# End: Content from sparktts/modules/encoder_decoder/feat_encoder.py +# =============================================================== + + +# =============================================================== +# Start: Content from sparktts/modules/encoder_decoder/wave_generator.py +# =============================================================== +class DecoderBlock(nn.Module): + def __init__( + self, + input_dim: int = 16, + output_dim: int = 8, + kernel_size: int = 2, + stride: int = 1, + ): + super().__init__() + # Ensure stride is at least 1 + stride = max(1, stride) + # Ensure kernel_size is valid for ConvTranspose1d + if kernel_size < stride: + kernel_size = stride # Or handle differently + + padding = (kernel_size - stride) // 2 + output_padding = stride % 2 if kernel_size % 2 == 0 else 0 # Basic calculation, might need adjustment based on desired output length + + # print(f"DecoderBlock - Input: {input_dim}, Output: {output_dim}, Kernel: {kernel_size}, Stride: {stride}, Padding: {padding}, OutputPadding: {output_padding}") + + + self.block = nn.Sequential( + Snake1d(input_dim), + WNConvTranspose1d( + input_dim, + output_dim, + kernel_size=kernel_size, + stride=stride, + padding=padding, + output_padding=output_padding, # Add output_padding + ), + ResidualUnit(output_dim, dilation=1), + ResidualUnit(output_dim, dilation=3), + ResidualUnit(output_dim, dilation=9), + ) + + def forward(self, x): + return self.block(x) + + +class WaveGenerator(nn.Module): + def __init__( + self, + input_channel, + channels, + rates, + kernel_sizes, + d_out: int = 1, + ): + super().__init__() + + # Add first conv layer + layers = [WNConv1d(input_channel, channels, kernel_size=7, padding=3)] + + # Add upsampling + MRF blocks + current_channels = channels + for i, (kernel_size, stride) in enumerate(zip(kernel_sizes, rates)): + input_dim = current_channels + # Ensure output_dim doesn't go below 1 + output_dim = max(1, channels // (2 ** (i + 1))) + layers += [DecoderBlock(input_dim, output_dim, kernel_size, stride)] + current_channels = output_dim # Update for the next block's input + + # Add final conv layer + layers += [ + Snake1d(current_channels), # Use the final output_dim + WNConv1d(current_channels, d_out, kernel_size=7, padding=3), + nn.Tanh(), + ] + + self.model = nn.Sequential(*layers) + + self.apply(init_weights) # Apply weight initialization + + def forward(self, x): + return self.model(x) + +# =============================================================== +# End: Content from sparktts/modules/encoder_decoder/wave_generator.py +# =============================================================== + + +# =============================================================== +# Start: Content from sparktts/modules/fsq/finite_scalar_quantization.py +# =============================================================== +# helper functions moved earlier +def round_ste(z: Tensor) -> Tensor: + """Round with straight through gradients.""" + zhat = z.round() + return z + (zhat - z).detach() + + +class FSQ(nn.Module): + def __init__( + self, + levels: List[int], + dim: int | None = None, + num_codebooks=1, + keep_num_codebooks_dim: bool | None = None, + scale: float | None = None, + allowed_dtypes: Tuple[torch.dtype, ...] = (torch.float32, torch.float64), + channel_first: bool = False, # Added based on usage in ResidualFSQ + projection_has_bias: bool = True, + return_indices=True, + force_quantization_f32=True, + ): + super().__init__() + _levels = torch.tensor(levels, dtype=int32) + self.register_buffer("_levels", _levels, persistent=False) + + _basis = torch.cumprod(torch.tensor([1] + levels[:-1]), dim=0, dtype=int32) + self.register_buffer("_basis", _basis, persistent=False) + + self.scale = scale # Not used in current implementation, but kept + + codebook_dim = len(levels) + self.codebook_dim = codebook_dim + + effective_codebook_dim = codebook_dim * num_codebooks + self.num_codebooks = num_codebooks + self.effective_codebook_dim = effective_codebook_dim + + # keep_num_codebooks_dim = default(keep_num_codebooks_dim, num_codebooks > 1) + # Force keep_num_codebooks_dim to False if num_codebooks is 1 + if num_codebooks == 1: + keep_num_codebooks_dim = False + else: + keep_num_codebooks_dim = default(keep_num_codebooks_dim, True) + + # Original assert was checking if num_codebooks > 1 and keep_num_codebooks_dim is False. Let's refine. + # If num_codebooks > 1, keep_num_codebooks_dim must be True based on how rearrange is used. + if num_codebooks > 1 and not keep_num_codebooks_dim: + raise ValueError("If num_codebooks > 1, keep_num_codebooks_dim must be True or None (defaults to True).") + self.keep_num_codebooks_dim = keep_num_codebooks_dim + + + self.dim = default(dim, len(_levels) * num_codebooks) + + self.channel_first = channel_first # Store channel_first setting + + has_projections = self.dim != effective_codebook_dim + self.project_in = ( + nn.Linear(self.dim, effective_codebook_dim, bias=projection_has_bias) + if has_projections + else nn.Identity() + ) + self.project_out = ( + nn.Linear(effective_codebook_dim, self.dim, bias=projection_has_bias) + if has_projections + else nn.Identity() + ) + + self.has_projections = has_projections + + self.return_indices = return_indices + if return_indices: + self.codebook_size = self._levels.prod().item() + # Calculate implicit codebook based on current device during forward pass if needed + # For now, calculate assuming CPU and move later if necessary + # implicit_codebook = self._indices_to_codes(torch.arange(self.codebook_size, device=self._levels.device)) # Calculate on device + # self.register_buffer("implicit_codebook", implicit_codebook, persistent=False) + + self.allowed_dtypes = allowed_dtypes + self.force_quantization_f32 = force_quantization_f32 + + @property + def implicit_codebook(self): + # Calculate implicit codebook on the fly using the device of _levels + device = self._levels.device + indices = torch.arange(self.codebook_size, device=device) + return self._indices_to_codes(indices) + + + def bound(self, z, eps: float = 1e-3): + """Bound `z`, an array of shape (..., d).""" + levels = self._levels.to(z.device) # Ensure levels are on same device + half_l = (levels - 1) * (1 + eps) / 2 + offset = torch.where(levels % 2 == 0, 0.5, 0.0) + shift = (offset / half_l).atanh() if torch.any(half_l != 0) else torch.zeros_like(offset) # Avoid div by zero + # Ensure shift is compatible shape for broadcasting + shift = shift.view(1, 1, -1) if z.ndim == 3 else shift # Adjust based on z dims + half_l = half_l.view(1, 1, -1) if z.ndim == 3 else half_l + + # Clamp input to avoid inf/-inf in atanh + z_clipped = torch.clamp(z, min=-1.0 + eps, max=1.0 - eps) # Assuming input z is somewhat normalized? + + # Original formula might be sensitive, let's try direct clamping. + # return (z + shift).tanh() * half_l - offset + + # Alternative clamping approach (from original Jax version logic): + upper_bound = (levels - 1) / 2 + lower_bound = -upper_bound + upper_bound = upper_bound.view(1, 1, -1) if z.ndim == 3 else upper_bound + lower_bound = lower_bound.view(1, 1, -1) if z.ndim == 3 else lower_bound + + return torch.clamp(z, min=lower_bound, max=upper_bound) + + + def quantize(self, z): + """Quantizes z, returns quantized zhat, same shape as z.""" + quantized = round_ste(self.bound(z)) + levels = self._levels.to(z.device) + half_width = levels // 2 # Renormalize to [-1, 1]. + # Avoid division by zero if level is 1 + half_width = torch.where(half_width == 0, torch.tensor(1.0, device=z.device), half_width.float()) + half_width_view = half_width.view(1, 1, -1) if quantized.ndim == 3 else half_width + return quantized / half_width_view + + def _scale_and_shift(self, zhat_normalized): + levels = self._levels.to(zhat_normalized.device) + half_width = levels // 2 + half_width_view = half_width.view(1, 1, -1) if zhat_normalized.ndim == 3 else half_width + return (zhat_normalized * half_width_view) + half_width_view + + def _scale_and_shift_inverse(self, zhat): + levels = self._levels.to(zhat.device) + half_width = levels // 2 + # Avoid division by zero if level is 1 + half_width = torch.where(half_width == 0, torch.tensor(1.0, device=zhat.device), half_width.float()) + half_width_view = half_width.view(1, 1, -1) if zhat.ndim == 3 else half_width + return (zhat - half_width_view) / half_width_view + + def _indices_to_codes(self, indices): + level_indices = self.indices_to_level_indices(indices) + codes = self._scale_and_shift_inverse(level_indices.float()) # Convert level indices to float + return codes + + def codes_to_indices(self, zhat): + """Converts a `code` to an index in the codebook.""" + assert zhat.shape[-1] == self.codebook_dim + zhat_scaled = self._scale_and_shift(zhat) + # Ensure basis is on the correct device and dtype, handle potential shape mismatch + basis = self._basis.to(zhat.device, dtype=int32) + basis_view = basis.view(1, 1, -1) if zhat_scaled.ndim == 3 else basis # Match ndim + # Ensure zhat_scaled is integer type for multiplication with basis + product = (zhat_scaled * basis_view).round().int() + return product.sum(dim=-1).to(int32) + + def indices_to_level_indices(self, indices): + """Converts indices to indices at each level, perhaps needed for a transformer with factorized embeddings""" + indices_reshaped = rearrange(indices, "... -> ... 1") + basis = self._basis.to(indices.device) + levels = self._levels.to(indices.device) + # Ensure basis and levels match the device and potentially ndim of indices + basis_view = basis.view(*([1] * (indices_reshaped.ndim - 1)), -1) + levels_view = levels.view(*([1] * (indices_reshaped.ndim - 1)), -1) + + codes_non_centered = (indices_reshaped // basis_view) % levels_view + return codes_non_centered + + # indices_to_codes is now handled by implicit_codebook property + project_out if needed + + def forward(self, z): + """ + einstein notation + b - batch + ... - sequence, spatial dimensions + d - feature dimension + c - number of codebook dim (within a single quantizer) + g - number of quantizers (groups) - handled by ResidualFSQ/GroupedResidualFSQ + """ + + # Input z can be (b d ...) or (b ... d) + # self.channel_first determines the expected input format for projection + + if self.channel_first: + # Expects (b d ...) + if z.ndim > 2: # Has spatial/temporal dims + z = rearrange(z, "b d ... -> b ... d") + z, ps = pack([z], "b * d") + # else: z is (b d) -> processed directly by linear + else: + # Expects (b ... d) + if z.ndim > 2: + z, ps = pack([z], "b * d") + # else: z is (b d) -> processed directly by linear + + + assert ( + z.shape[-1] == self.dim + ), f"expected dimension of {self.dim} but found dimension of {z.shape[-1]}" + + # Project in + z_projected = self.project_in(z) # (b ... effective_codebook_dim) + + # Reshape for codebooks if num_codebooks > 1 + if self.num_codebooks > 1: + z_reshaped = rearrange(z_projected, "b ... (c d) -> b ... c d", c=self.num_codebooks) + else: + # Add a dummy codebook dim for consistent processing + z_reshaped = rearrange(z_projected, "b ... d -> b ... 1 d") + + # Force quantization step to be full precision or not + force_f32 = self.force_quantization_f32 + quantization_context = ( + partial(autocast, "cuda", enabled=False) if force_f32 else nullcontext + ) + + codes = None + indices = None + + with quantization_context(): + orig_dtype = z_reshaped.dtype + + if force_f32 and orig_dtype not in self.allowed_dtypes: + z_for_quant = z_reshaped.float() + else: + z_for_quant = z_reshaped + + codes = self.quantize(z_for_quant) # (b ... c d) + + if self.return_indices: + indices = self.codes_to_indices(codes) # (b ... c) + + # Convert codes back to original dtype if changed + codes = codes.type(orig_dtype) + + + # Reshape codes back and project out + if self.num_codebooks > 1: + codes_reshaped = rearrange(codes, "b ... c d -> b ... (c d)") + else: + codes_reshaped = rearrange(codes, "b ... 1 d -> b ... d") + + out = self.project_out(codes_reshaped) # (b ... dim) + + # Restore original spatial/temporal dimensions + if z.ndim > 2: # If we packed dimensions + out = unpack(out, ps, "b * d")[0] + if self.return_indices: + indices = unpack(indices, ps, "b * c")[0] + + # Restore channel dimension if needed + if self.channel_first and out.ndim > 2: + out = rearrange(out, "b ... d -> b d ...") + if self.return_indices and indices.ndim > 1: # Check indices ndim + # Indices shape (b ... c), need to decide how to handle channel dim + # Often indices might not need channel dim, depends on usage + # If indices are e.g. (b H W c), permuting might be complex. + # Keeping indices as (b ... c) for now. + pass + + + # Remove the dummy codebook dim from indices if num_codebooks was 1 + if self.return_indices and self.num_codebooks == 1 and not self.keep_num_codebooks_dim: + indices = indices.squeeze(-1) + + + return out, indices +# =============================================================== +# End: Content from sparktts/modules/fsq/finite_scalar_quantization.py +# =============================================================== + + +# =============================================================== +# Start: Content from sparktts/modules/fsq/residual_fsq.py +# =============================================================== +# Helper functions needed by ResidualFSQ +def is_distributed(): + return dist.is_initialized() and dist.get_world_size() > 1 + +def get_maybe_sync_seed(device, max_size=10_000): + rand_int = torch.randint(0, max_size, (), device=device) + if is_distributed(): + # Ensure rand_int is on the correct device for all_reduce + if rand_int.device != device: + rand_int = rand_int.to(device) + dist.all_reduce(rand_int) + return rand_int.item() + +def round_up_multiple(num, mult): + # Ensure mult is positive + if mult <= 0: + return num + # Use ceiling division + return (num + mult - 1) // mult * mult + + +class ResidualFSQ(nn.Module): + """Follows Algorithm 1. in https://arxiv.org/pdf/2107.03312.pdf""" + + def __init__( + self, + *, + levels: List[int], + num_quantizers, + dim=None, + # is_channel_first=False, # Handled inside FSQ now + quantize_dropout=False, + quantize_dropout_cutoff_index=0, + quantize_dropout_multiple_of=1, + channel_first: bool = False, # Pass channel_first to FSQ + **kwargs, # Pass remaining kwargs to FSQ + ): + super().__init__() + codebook_dim = len(levels) + dim = default(dim, codebook_dim) + + requires_projection = codebook_dim != dim + self.project_in = ( + nn.Linear(dim, codebook_dim) if requires_projection else nn.Identity() + ) + self.project_out = ( + nn.Linear(codebook_dim, dim) if requires_projection else nn.Identity() + ) + self.has_projections = requires_projection + + self.channel_first = channel_first # Store for potential shape adjustments if needed later + self.num_quantizers = num_quantizers + + self.levels = levels + self.layers = nn.ModuleList([]) + + levels_tensor = torch.Tensor(levels) + + scales = [] + + for ind in range(num_quantizers): + # Calculate scale: (levels - 1) is max value range (- (l-1)/2 to +(l-1)/2) + # Residual is divided by scale before quantization + # Effective scale for quantizer 'ind' is (levels - 1)^ind ? Needs check. + # Original paper scale seems different. Let's stick to FSQ handling scale internally if needed. + # Using scale = 1.0 for now, assuming FSQ handles normalization. + scale_value = 1.0 # ((levels_tensor - 1)**-ind) - Check this logic + scales.append(scale_value) + + # Pass channel_first to FSQ + fsq = FSQ(levels=levels, dim=codebook_dim, channel_first=channel_first, **kwargs) + + self.layers.append(fsq) + + # Check if FSQ layers have projections internally. ResidualFSQ should handle overall projection. + assert all([not fsq.has_projections for fsq in self.layers]), "FSQ layers within ResidualFSQ should not have internal projections." + + self.codebook_size = self.layers[0].codebook_size + + # Using scale = 1.0, so register_buffer might not be needed, or store 1.0s + # self.register_buffer("scales", torch.Tensor(scales), persistent=False) + # If scales are needed, they should likely be parameters or calculated differently. + # For now, assuming FSQ normalizes correctly and scale is 1.0 here. + + self.quantize_dropout = quantize_dropout and num_quantizers > 1 + + assert quantize_dropout_cutoff_index >= 0 + + self.quantize_dropout_cutoff_index = quantize_dropout_cutoff_index + self.quantize_dropout_multiple_of = quantize_dropout_multiple_of # encodec paper proposes structured dropout, believe this was set to 4 + + @property + def codebooks(self): + # Codebooks are implicit in FSQ, access via property + codebooks = [layer.implicit_codebook for layer in self.layers] + codebooks = torch.stack(codebooks, dim=0) + return codebooks + + def get_codes_from_indices(self, indices): + # indices shape: (b ... q) or (b q ...) depending on usage + num_dims = indices.ndim + q_dim = -1 # Assume last dim is quantizer dim by default + + # Find the quantizer dimension (q) + for i in range(num_dims): + if indices.shape[i] == self.num_quantizers: + q_dim = i + break + if q_dim == -1 and self.num_quantizers == 1 and indices.shape[-1] != 1: + # If only 1 quantizer, indices might not have the quantizer dim explicitly + indices = indices.unsqueeze(-1) # Add the quantizer dim + q_dim = -1 + elif q_dim == -1: + raise ValueError(f"Could not find quantizer dimension ({self.num_quantizers}) in indices shape {indices.shape}") + + # Ensure q_dim is the last dimension for processing + if q_dim != num_dims - 1: + permute_dims = list(range(num_dims)) + permute_dims.pop(q_dim) + permute_dims.append(q_dim) + indices = indices.permute(*permute_dims) + + + batch_shape = indices.shape[:-1] # Shape before the quantizer dim + indices = indices.reshape(-1, self.num_quantizers) # Flatten batch/spatial dims + + # Handle dropout indices (-1) + if indices.max() >= self.codebook_size: + raise ValueError(f"Invalid index found in indices: {indices.max()}. Max allowed is {self.codebook_size - 1}.") + if indices.min() < -1: + raise ValueError(f"Invalid index found in indices: {indices.min()}. Min allowed is -1 (dropout).") + + mask = indices == -1 + effective_indices = indices.masked_fill(mask, 0) # Use 0 for dropout indices temporarily + + all_codes = [] + # Iterate through each quantizer layer + for i in range(self.num_quantizers): + layer_indices = effective_indices[:, i] + # Use the FSQ layer's method to convert indices to codes (handles normalization) + # Need to ensure indices_to_codes exists and works correctly in FSQ + # Assuming FSQ.indices_to_codes takes (batch,) indices and returns (batch, codebook_dim) codes + layer_codes = self.layers[i].indices_to_codes(layer_indices) # This needs correct FSQ method + all_codes.append(layer_codes) + + all_codes_tensor = torch.stack(all_codes, dim=0) # (q, b_flat, d) + + # Mask out dropout codes + mask_expanded = mask.permute(1, 0).unsqueeze(-1) # (q, b_flat, 1) + all_codes_tensor = all_codes_tensor.masked_fill(mask_expanded, 0.0) + + # Reshape back to original batch/spatial shape + all_codes_tensor = all_codes_tensor.reshape(self.num_quantizers, *batch_shape, -1) # (q, b ... d) + + # Restore original q_dim position if it was changed + if q_dim != num_dims - 1: + # Need inverse permutation + inv_permute_dims = list(range(num_dims)) # Start with 0, 1, ..., num_dims-1 + inv_permute_dims.insert(q_dim, num_dims) # Insert the last dim (q) at the original position + inv_permute_dims.pop() # Remove the last element + # Permute from (q, b ... d) -> (b ... q ... d) - careful with dims + # Example: Input (b h w q), processed to (q, b*h*w), output (q, b*h*w, d) + # Reshaped to (q, b, h, w, d) + # Want (b, h, w, q, d) -> Need to confirm this logic + # Let's assume output shape (q, b, ..., d) is desired for summation later. + pass # Keep as (q, b ... d) for now + + return all_codes_tensor + + + def get_output_from_indices(self, indices): + # indices shape: (b ... q) + codes = self.get_codes_from_indices(indices) # Output: (q, b ... d) + codes_summed = reduce(codes, "q b ... d -> b ... d", "sum") + # Project back to original dimension + output = self.project_out(codes_summed) + + # Handle channel first if necessary for the final output + if self.channel_first and output.ndim > 2: + # Assumes input was (b d ...), so output should be too + output = rearrange(output, "b ... d -> b d ...") + + return output + + def forward(self, x, return_all_codes=False, rand_quantize_dropout_fixed_seed=None): + num_quant, quant_dropout_multiple_of, device = ( + self.num_quantizers, + self.quantize_dropout_multiple_of, + x.device, + ) + + # handle channel first input if necessary for projection + original_shape = x.shape + if self.channel_first: + if x.ndim > 2: # Has spatial/temporal dims + x = rearrange(x, "b d ... -> b ... d") + x, ps = pack([x], "b * d") + # else: x is (b d), processed directly + else: + # Input is (b ... d) + if x.ndim > 2: + x, ps = pack([x], "b * d") + # else: x is (b d), processed directly + + + # maybe project in + projected_x = self.project_in(x) # (b ... codebook_dim) + + quantized_out = 0.0 + residual = projected_x # Start residual from projected input + + all_indices = [] + + should_quantize_dropout = self.training and self.quantize_dropout + + # sample a layer index at which to dropout further residual quantization + # also prepare null indices + rand_quantize_dropout_index = num_quant # Default to no dropout + + if should_quantize_dropout: + if not exists(rand_quantize_dropout_fixed_seed): + rand_quantize_dropout_fixed_seed = get_maybe_sync_seed(device) + + rand = random.Random(rand_quantize_dropout_fixed_seed) + # Ensure cutoff index is valid + valid_cutoff = max(0, self.quantize_dropout_cutoff_index) + rand_quantize_dropout_index = rand.randrange(valid_cutoff, num_quant) + + if quant_dropout_multiple_of != 1: + rand_quantize_dropout_index = ( + round_up_multiple(rand_quantize_dropout_index + 1, quant_dropout_multiple_of) - 1 + ) + # Clamp index to be within valid range + rand_quantize_dropout_index = min(rand_quantize_dropout_index, num_quant - 1) + + # Null indices shape should match the batch/spatial dims of x before pack + null_indices_shape = list(x.shape[:-1]) # All dims except last feature dim + null_indices = torch.full(null_indices_shape, -1, device=device, dtype=torch.long) + + + # go through the layers + # Assuming scale is handled within FSQ or is 1.0 here + # scales = self.scales.to(device) + + for quantizer_index, layer in enumerate(self.layers): + # scale = scales[quantizer_index] # If using external scales + + if quantizer_index > rand_quantize_dropout_index: + # Append null indices matching the shape of valid indices from FSQ + # FSQ returns indices shape (b ...) or (b ... c) -> need (b ...) + # Use the pre-calculated null_indices + all_indices.append(null_indices) + continue + + # Pass residual to the quantizer layer + # Assume FSQ takes (b ... d) or (b d ...) based on its channel_first setting + # Here, residual is (b ... codebook_dim) + quantized, indices = layer(residual) # layer should handle channel_first internally + + # residual = residual - quantized.detach() # Update residual BEFORE summing output + # quantized_out = quantized_out + quantized # Sum the quantized part + + # Algorithm 1 from paper: + # Input: x + # residual = x + # codes = [] + # for q in quantizers: + # x_q, indices = q(residual) # Quantize + # residual = residual - x_q # Update residual (use x_q directly, not detached?) - Check paper/encodec. Using detached version. + # codes.append(indices) + # x_hat = sum(x_q for each layer?) - No, final quantized output is reconstructed from indices. + + # Let's follow common implementation: sum quantized outputs, update residual with detached quantized + quantized_detached = quantized.detach() + residual = residual - quantized_detached + quantized_out = quantized_out + quantized # Sum quantized outputs from each layer + + # Store indices + if indices is None: + raise ValueError(f"FSQ layer {quantizer_index} did not return indices.") + all_indices.append(indices) + + # project out the summed quantized output + final_quantized_out = self.project_out(quantized_out) # (b ... dim) + + # stack all indices + all_indices = torch.stack(all_indices, dim=-1) # (b ... q) + + # Restore original shape if packed + if x.ndim > 2: # If we packed dimensions + final_quantized_out = unpack(final_quantized_out, ps, "b * d")[0] + all_indices = unpack(all_indices, ps, "b * q")[0] + + # Restore channel dimension if needed + if self.channel_first and final_quantized_out.ndim > 2: + final_quantized_out = rearrange(final_quantized_out, "b ... d -> b d ...") + # Decide how to handle indices shape. Keep as (b ... q) or (b q ...)? + # Keeping as (b ... q) seems more common. + # all_indices = rearrange(all_indices, "b ... q -> b q ...") # Optional rearrange + + # return + ret = (final_quantized_out, all_indices) + + if not return_all_codes: + return ret + + # Return all codes (reconstructed from indices) + # Input to get_codes_from_indices should be (b ... q) + all_codes = self.get_codes_from_indices(all_indices) # Output (q, b ... d) + + # Maybe reshape all_codes to match input shape conventions? + # If input was channel_first (b d ...), maybe output codes as (q b d ...)? + if self.channel_first and all_codes.ndim > 3: + all_codes = rearrange(all_codes, "q b ... d -> q b d ...") + + return (*ret, all_codes) + +# =============================================================== +# End: Content from sparktts/modules/fsq/residual_fsq.py +# =============================================================== + + +# =============================================================== +# Start: Content from sparktts/modules/speaker/ecapa_tdnn.py +# =============================================================== +class Res2Conv1dReluBn(nn.Module): + """ + in_channels == out_channels == channels + """ + + def __init__( + self, + channels, + kernel_size=1, + stride=1, + padding=0, + dilation=1, + bias=True, + scale=4, + ): + super().__init__() + assert channels % scale == 0, "{} % {} != 0".format(channels, scale) + self.scale = scale + self.width = channels // scale + self.nums = scale if scale == 1 else scale - 1 + + self.convs = [] + self.bns = [] + for i in range(self.nums): + self.convs.append( + nn.Conv1d( + self.width, + self.width, + kernel_size, + stride, + padding, + dilation, + bias=bias, + ) + ) + self.bns.append(nn.BatchNorm1d(self.width)) + self.convs = nn.ModuleList(self.convs) + self.bns = nn.ModuleList(self.bns) + + def forward(self, x): + out = [] + spx = torch.split(x, self.width, 1) + sp = spx[0] + # Enumerate starts from 0, matching list indices + for i, (conv, bn) in enumerate(zip(self.convs, self.bns)): + # Order: conv -> relu -> bn + if i >= 1: + sp = sp + spx[i] # Residual connection within block parts + sp = conv(sp) + sp = bn(F.relu(sp)) # Apply ReLU before BatchNorm + out.append(sp) + if self.scale != 1: + # Append the last chunk without processing if scale > 1 + out.append(spx[self.nums]) + out = torch.cat(out, dim=1) + + return out + + +""" Conv1d + BatchNorm1d + ReLU """ +class Conv1dReluBn(nn.Module): + def __init__( + self, + in_channels, + out_channels, + kernel_size=1, + stride=1, + padding=0, + dilation=1, + bias=True, + ): + super().__init__() + self.conv = nn.Conv1d( + in_channels, out_channels, kernel_size, stride, padding, dilation, bias=bias + ) + self.bn = nn.BatchNorm1d(out_channels) + + def forward(self, x): + # Original: bn(relu(conv(x))) + # ECAPA Paper Figure/Desc seems to suggest conv -> bn -> relu ? Check Res2Net paper/ECAPA details. + # Sticking to original code's bn(relu(conv(x))) for now. + return self.bn(F.relu(self.conv(x))) + + +""" The SE connection of 1D case. """ +class SE_Connect(nn.Module): + def __init__(self, channels, se_bottleneck_dim=128): + super().__init__() + self.linear1 = nn.Linear(channels, se_bottleneck_dim) + self.linear2 = nn.Linear(se_bottleneck_dim, channels) + + def forward(self, x): + # x shape: (B, C, T) + out = x.mean(dim=2) # Global average pooling over time -> (B, C) + out = F.relu(self.linear1(out)) + out = torch.sigmoid(self.linear2(out)) + out = x * out.unsqueeze(2) # (B, C, T) * (B, C, 1) -> (B, C, T) + + return out + + +""" SE-Res2Block of the ECAPA-TDNN architecture. """ +class SE_Res2Block(nn.Module): + def __init__(self, channels, kernel_size, stride, padding, dilation, scale): + super().__init__() + self.se_res2block = nn.Sequential( + Conv1dReluBn(channels, channels, kernel_size=1, stride=1, padding=0), + Res2Conv1dReluBn( + channels, kernel_size, stride, padding, dilation, scale=scale + ), + Conv1dReluBn(channels, channels, kernel_size=1, stride=1, padding=0), + SE_Connect(channels), + ) + + def forward(self, x): + return x + self.se_res2block(x) + + +class ECAPA_TDNN(nn.Module): + def __init__( + self, + channels=512, + feat_dim=80, + embed_dim=192, + pooling_func="ASTP", + global_context_att=False, + emb_bn=False, + ): + super().__init__() + + self.layer1 = Conv1dReluBn(feat_dim, channels, kernel_size=5, padding=2) + self.layer2 = SE_Res2Block( + channels, kernel_size=3, stride=1, padding=2, dilation=2, scale=8 + ) + self.layer3 = SE_Res2Block( + channels, kernel_size=3, stride=1, padding=3, dilation=3, scale=8 + ) + self.layer4 = SE_Res2Block( + channels, kernel_size=3, stride=1, padding=4, dilation=4, scale=8 + ) + + cat_channels = channels * 3 + # The output channels after conv depends on the pooling layer input expectation + # Original paper uses 1536. Let's assume pooling expects 1536. + self.conv = nn.Conv1d(cat_channels, cat_channels, kernel_size=1) # Keep channels same for pooling + + # Dynamically get pooling class based on string name from pooling_layers (defined earlier) + if pooling_func == "TAP": pooling_layer = TAP + elif pooling_func == "TSDP": pooling_layer = TSDP + elif pooling_func == "TSTP": pooling_layer = TSTP + elif pooling_func == "ASTP": pooling_layer = ASTP + elif pooling_func == "MHASTP": pooling_layer = MHASTP + elif pooling_func == "MQMHASTP": pooling_layer = MQMHASTP + else: raise ValueError(f"Unsupported pooling function: {pooling_func}") + + self.pool = pooling_layer( + in_dim=cat_channels, # Pooling operates on the output of self.conv + global_context_att=global_context_att # Pass context flag if relevant (ASTP) + # Add other necessary kwargs for specific pooling layers if needed + ) + # self.pool_out_dim = self.pool.get_out_dim() # Get output dim from pooling layer + # Use standard way to get output dim if get_out_dim not standard + # For TSTP/ASTP etc., it's usually 2 * in_dim + if hasattr(self.pool, 'get_out_dim'): + self.pool_out_dim = self.pool.get_out_dim() + elif isinstance(self.pool, (TSTP, ASTP, MHASTP, MQMHASTP)): + # Assuming these double the input dimension + self.pool_out_dim = cat_channels * (2 * getattr(self.pool, 'query_num', 1) if isinstance(self.pool, MQMHASTP) else 2) + else: # TAP, TSDP + self.pool_out_dim = cat_channels + + self.bn = nn.BatchNorm1d(self.pool_out_dim) + self.linear = nn.Linear(self.pool_out_dim, embed_dim) + self.emb_bn = emb_bn + if emb_bn: # better in SSL for SV + self.bn2 = nn.BatchNorm1d(embed_dim) + else: + self.bn2 = nn.Identity() + + def forward(self, x, return_latent=False): + # Input x expected as (B, T, F) e.g., mels + x = x.permute(0, 2, 1) # (B, T, F) -> (B, F, T) + + out1 = self.layer1(x) + out2 = self.layer2(out1) + out3 = self.layer3(out2) + out4 = self.layer4(out3) + + # Concat features from layers 2, 3, 4 + out = torch.cat([out2, out3, out4], dim=1) # (B, 3*channels, T) + latent = F.relu(self.conv(out)) # (B, 3*channels, T) + + # Pooling expects (B, F, T) + pooled_out = self.pool(latent) # (B, pool_out_dim) + bn_out = self.bn(pooled_out) + embedding = self.linear(bn_out) # (B, embed_dim) + + if self.emb_bn: + embedding = self.bn2(embedding) + + if return_latent: + # Return the embedding and the features before pooling + return embedding, latent # latent shape (B, 3*channels, T) + return embedding # Return only the final embedding + + +# Factory functions (optional, but keep if used elsewhere) +def ECAPA_TDNN_c1024(feat_dim, embed_dim, pooling_func="ASTP", emb_bn=False): + return ECAPA_TDNN( + channels=1024, + feat_dim=feat_dim, + embed_dim=embed_dim, + pooling_func=pooling_func, + emb_bn=emb_bn, + ) + +def ECAPA_TDNN_GLOB_c1024(feat_dim, embed_dim, pooling_func="ASTP", emb_bn=False): + return ECAPA_TDNN( + channels=1024, + feat_dim=feat_dim, + embed_dim=embed_dim, + pooling_func=pooling_func, + global_context_att=True, + emb_bn=emb_bn, + ) + +def ECAPA_TDNN_c512(feat_dim, embed_dim, pooling_func="ASTP", emb_bn=False): + return ECAPA_TDNN( + channels=512, + feat_dim=feat_dim, + embed_dim=embed_dim, + pooling_func=pooling_func, + emb_bn=emb_bn, + ) + +def ECAPA_TDNN_GLOB_c512(feat_dim, embed_dim, pooling_func="ASTP", emb_bn=False): + return ECAPA_TDNN( + channels=512, + feat_dim=feat_dim, + embed_dim=embed_dim, + pooling_func=pooling_func, + global_context_att=True, + emb_bn=emb_bn, + ) + +# =============================================================== +# End: Content from sparktts/modules/speaker/ecapa_tdnn.py +# =============================================================== + + +# =============================================================== +# Start: Content from sparktts/modules/speaker/perceiver_encoder.py +# =============================================================== +# Helper functions for Perceiver/Attention +def exists(val): # Redefined earlier + return val is not None + +def once(fn): # Redefined earlier + called = False + @wraps(fn) + def inner(x): + nonlocal called + if called: return + called = True + return fn(x) + return inner + +print_once = once(print) + +class Attend(nn.Module): + def __init__(self, dropout=0.0, causal=False, use_flash=False): + super().__init__() + self.dropout = dropout + self.attn_dropout = nn.Dropout(dropout) + + self.causal = causal + self.register_buffer("mask", None, persistent=False) + + self.use_flash = use_flash + can_use_flash = hasattr(F, 'scaled_dot_product_attention') and use_flash + if can_use_flash: + print_once("Using Flash Attention for Perceiver.") + else: + if use_flash: print_once("Flash Attention requested but not available/enabled.") + self.use_flash = False # Disable if not available + + # Flash attention config (simplified) + self.efficient_config = namedtuple("EfficientAttentionConfig", ["enable_flash", "enable_math", "enable_mem_efficient"]) + # Set default configs, actual backend selection happens in F.scaled_dot_product_attention + self.cpu_config = self.efficient_config(True, True, True) # Default for CPU + self.cuda_config = self.efficient_config(True, True, True) # Default for CUDA + + + def get_mask(self, n, device): + if exists(self.mask) and self.mask.shape[-1] >= n and self.mask.device == device: + return self.mask[:n, :n] + + mask = torch.ones((n, n), device=device, dtype=torch.bool).triu(1) + self.register_buffer("mask", mask, persistent=False) + return mask + + def flash_attn(self, q, k, v, mask=None): + _, heads, q_len, _, k_len, is_cuda = *q.shape, k.shape[-2], q.is_cuda + + # Expand KV if needed (for multi-query attention, though Perceiver might use standard MHA) + if k.ndim == 3: # (b n_kv d) -> (b h n_kv d) ? No, needs (b h n_kv d_head) + # Assume k/v are already (b h n d_head) or need different handling + pass + if v.ndim == 3: + pass + + # Format mask for flash attention (B, N_q, N_kv) or (B, H, N_q, N_kv) + flash_mask = None + if exists(mask): + # mask shape (b, n_kv) -> needs (b, 1, n_q, n_kv) or (b, h, n_q, n_kv) ? + # Check documentation. For key padding mask, usually (B, N_kv). + # Needs expansion. Let's assume (B, H, N_q, N_kv) for safety. + if mask.ndim == 2: # (b, n_kv) + flash_mask = rearrange(mask, "b j -> b 1 1 j") + # Flash attention expects additive mask (-inf for masked) not boolean? Check. + # F.scaled_dot_product_attention takes boolean mask with attn_mask arg. + flash_mask = flash_mask.expand(-1, heads, q_len, -1) # (b h n_q n_kv) + # Use ~mask because True means *mask out* in flash attn's attn_mask. + flash_mask = ~flash_mask + elif mask.ndim == 4 and mask.shape[1] == 1: # Maybe already expanded (b 1 1 n_kv) + flash_mask = mask.expand(-1, heads, q_len, -1) + flash_mask = ~flash_mask + else: + # Assuming mask might already be correctly shaped (e.g., B, H, Nq, Nkv boolean) + flash_mask = ~mask # Invert mask if boolean + + # pytorch 2.0 flash attn: q, k, v, attn_mask, dropout_p, is_causal + # attn_mask should be boolean where True indicates masking. + out = F.scaled_dot_product_attention( + q, + k, + v, + attn_mask=flash_mask if exists(flash_mask) else None, # Pass boolean mask + dropout_p=self.dropout if self.training else 0.0, + is_causal=self.causal # Pass causal flag directly + ) + return out + + def forward(self, q, k, v, mask=None): + """ + einstein notation + b - batch + h - heads + n, i, j - sequence length (query, key/value) + d - feature dimension (d_head) + """ + n, device = q.shape[-2], q.device + scale = q.shape[-1] ** -0.5 + + if self.use_flash: + return self.flash_attn(q, k, v, mask=mask) + + # Manual Attention Calculation + kv_einsum_eq = "b h j d" # Assuming k, v are always (b h n d) + + # similarity + sim = einsum(f"b h i d, {kv_einsum_eq} -> b h i j", q, k) * scale + + # key padding mask + if exists(mask): + # mask shape (b, j) -> (b, 1, 1, j) + mask_value = -torch.finfo(sim.dtype).max + mask = rearrange(mask, "b j -> b 1 1 j") + sim = sim.masked_fill(~mask, mask_value) # Mask where mask is False + + # causal mask (Not typically used in Perceiver cross-attention) + if self.causal: + causal_mask = self.get_mask(n, device) # (i, j) + sim = sim.masked_fill(causal_mask, mask_value) + + # attention + attn = sim.softmax(dim=-1) + attn = self.attn_dropout(attn) + + # aggregate values + out = einsum(f"b h i j, {kv_einsum_eq} -> b h i d", attn, v) + + return out + + +# Need Sequential, default, RMSNorm, GEGLU, FeedForward, Attention for PerceiverResampler +def Sequential(*mods): # Redefined earlier + return nn.Sequential(*filter(exists, mods)) + +class RMSNorm(nn.Module): + def __init__(self, dim, scale=True, dim_cond=None): + super().__init__() + self.cond = exists(dim_cond) + # Conditional LayerNorm not used in PerceiverResampler, simplify + # self.to_gamma_beta = nn.Linear(dim_cond, dim * 2) if self.cond else None + + self.scale = dim**0.5 + self.gamma = nn.Parameter(torch.ones(dim)) if scale else None + + def forward(self, x, cond=None): # Remove cond argument if not used + gamma = default(self.gamma, torch.tensor(1.0, device=x.device)) # Ensure gamma is tensor + # Note: F.normalize normalizes across the *last* dimension by default + normed_x = F.normalize(x, dim=-1) + return normed_x * self.scale * gamma + + +class CausalConv1d(nn.Conv1d): # Already defined earlier + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + kernel_size = self.kernel_size[0] + dilation = self.dilation[0] + stride = self.stride[0] + assert stride == 1 + self.causal_padding = dilation * (kernel_size - 1) + + def forward(self, x): + # Input x: (B, C, T) + causal_padded_x = F.pad(x, (self.causal_padding, 0), value=0.0) + return super().forward(causal_padded_x) + +class GEGLU(nn.Module): # Already defined earlier + def forward(self, x): + x, gate = x.chunk(2, dim=-1) + return F.gelu(gate) * x + +def FeedForward(dim, mult=4, causal_conv=False): # Already defined earlier + dim_inner = int(dim * mult * 2 / 3) + + conv = None + if causal_conv: + conv = nn.Sequential( + Rearrange("b n d -> b d n"), + CausalConv1d(dim_inner, dim_inner, 3), + Rearrange("b d n -> b n d"), + ) + + return Sequential( + nn.Linear(dim, dim_inner * 2, bias=False), # Bias False often used in transformers + GEGLU(), + conv, + nn.Linear(dim_inner, dim, bias=False) # Bias False + ) + + +class Attention(nn.Module): + def __init__( + self, + dim, + *, + dim_context=None, + causal=False, + dim_head=64, + heads=8, + dropout=0.0, + use_flash=False, + cross_attn_include_queries=False, + ): + super().__init__() + # self.scale = dim_head**-0.5 # scale is handled by Attend or flash attn + self.heads = heads + self.cross_attn_include_queries = cross_attn_include_queries + + dim_inner = dim_head * heads + dim_context = default(dim_context, dim) + + self.attend = Attend(causal=causal, dropout=dropout, use_flash=use_flash) + self.to_q = nn.Linear(dim, dim_inner, bias=False) + # Combine K and V projection for efficiency + self.to_kv = nn.Linear(dim_context, dim_inner * 2, bias=False) + self.to_out = nn.Linear(dim_inner, dim, bias=False) + + def forward(self, x, context=None, mask=None): + h, has_context = self.heads, exists(context) + # x shape: (b, n_q, d) + # context shape: (b, n_kv, d_ctx) + + context = default(context, x) # Use self if context not provided + + if has_context and self.cross_attn_include_queries: + # Prepend queries to context for attention calculation + context = torch.cat((x, context), dim=-2) # (b, n_q + n_kv, d_ctx) - ensure dims match + + # Project q, k, v + q = self.to_q(x) + k, v = self.to_kv(context).chunk(2, dim=-1) + + # Reshape for multi-head attention + q, k, v = map(lambda t: rearrange(t, "b n (h d) -> b h n d", h=h), (q, k, v)) + + # Attend + out = self.attend(q, k, v, mask=mask) # mask should be (b, n_kv) + + # Combine heads and project out + out = rearrange(out, "b h n d -> b n (h d)") + return self.to_out(out) + + +class PerceiverResampler(nn.Module): + def __init__( + self, + *, + dim, + depth=2, + dim_context=None, + num_latents=32, + dim_head=64, + heads=8, + ff_mult=4, + use_flash_attn=False, + ): + super().__init__() + dim_context = default(dim_context, dim) + + # Project context to query dimension if different + self.proj_context = ( + nn.Linear(dim_context, dim) if dim_context != dim else nn.Identity() + ) + + # Learnable latent queries + self.latents = nn.Parameter(torch.randn(num_latents, dim)) + nn.init.normal_(self.latents, std=0.02) # Initialize latents + + self.layers = nn.ModuleList([]) + for _ in range(depth): + self.layers.append( + nn.ModuleList( + [ + # Cross-Attention from latents (queries) to context (keys/values) + Attention( + dim=dim, + dim_context=dim, # Context is projected to dim + dim_head=dim_head, + heads=heads, + use_flash=use_flash_attn, + cross_attn_include_queries=False, # Standard Perceiver cross-attn + ), + # Self-Attention within latents + # Optional: Add self-attention block here if needed + # Attention( + # dim=dim, dim_head=dim_head, heads=heads, use_flash=use_flash_attn + # ), + + # FeedForward block + FeedForward(dim=dim, mult=ff_mult), + ] + ) + ) + # Add LayerNorms (typically before attention and ff blocks) + # self.layers[-1].insert(0, RMSNorm(dim)) # Pre-Attention Norm + # self.layers[-1].insert(2, RMSNorm(dim)) # Pre-FF Norm + # Using Post-Norm structure as in original reference: + self.layers[-1].insert(1, RMSNorm(dim)) # After Attention + self.layers[-1].append(RMSNorm(dim)) # After FeedForward + + + # Final normalization of latents + # self.norm = RMSNorm(dim) # Final norm applied inside loop in original? Let's apply at end. + + def forward(self, x, mask=None): + # x shape: (b, n_ctx, d_ctx) + batch = x.shape[0] + + # Project context + x = self.proj_context(x) # (b, n_ctx, d) + + # Repeat latents for batch + latents = repeat(self.latents, "n d -> b n d", b=batch) # (b, n_lat, d) + + # Apply layers + # Original structure had norm inside loop, adapting: Attn -> Norm -> FF -> Norm + for attn, norm1, ff, norm2 in self.layers: + # Cross-Attention + Residual + latents_attn = attn(latents, x, mask=mask) # Query: latents, Context: x + latents = norm1(latents_attn + latents) + + # FeedForward + Residual + latents_ff = ff(latents) + latents = norm2(latents_ff + latents) + + # return self.norm(latents) # Apply final norm if defined outside loop + return latents # Return latents after last block norm + +# =============================================================== +# End: Content from sparktts/modules/speaker/perceiver_encoder.py +# =============================================================== + + +# =============================================================== +# Start: Content from sparktts/modules/speaker/speaker_encoder.py +# =============================================================== +class SpeakerEncoder(nn.Module): + """ + Speaker Encoder using ECAPA-TDNN, Perceiver Resampler, and Residual FSQ. + + Args: + input_dim (int): acoustic feature dimension (e.g., mel bins) + out_dim (int): output dimension of the final d-vector + latent_dim (int): latent dimension for perceiver and quantization + token_num (int): number of latent tokens from perceiver + fsq_levels (List[int]): levels for finite scalar quantization + fsq_num_quantizers (int): number of residual quantizers in FSQ + ecapa_embed_dim (int): embedding dimension from ECAPA-TDNN (before projection) + """ + + def __init__( + self, + input_dim: int = 80, # Default mel bins + out_dim: int = 1024, # Target d-vector dim from config + latent_dim: int = 128, # Latent dim for perceiver/quantizer + token_num: int = 32, # Number of speaker tokens + fsq_levels: List[int] = [4, 4, 4, 4, 4, 4], + fsq_num_quantizers: int = 1, + # Add ECAPA config params if needed, or use defaults + ecapa_channels: int = 512, + ecapa_embed_dim: int = 192, # Default ECAPA embed dim + ): + super(SpeakerEncoder, self).__init__() + + # ECAPA-TDNN for initial feature extraction and x-vector (optional) + # Using the GLOB variant as in the original __main__ test + self.speaker_encoder_base = ECAPA_TDNN_GLOB_c512( + feat_dim=input_dim, + embed_dim=ecapa_embed_dim # Use specific ECAPA embed dim + ) + # Dimension of features extracted by ECAPA (latent before pooling) + ecapa_feature_dim = ecapa_channels * 3 # From concatenation in ECAPA + + # Perceiver Resampler to get fixed-length sequence from variable-length ECAPA features + self.perceiver_sampler = PerceiverResampler( + dim=latent_dim, # Output dim of perceiver latents + dim_context=ecapa_feature_dim, # Input dim from ECAPA features + num_latents=token_num, + depth=2, # Default depth, adjust if needed + dim_head=64, heads=8, ff_mult=4, # Default attention/ff params + use_flash_attn=True # Enable flash attention if available + ) + + # Residual Finite Scalar Quantizer + self.quantizer = ResidualFSQ( + levels=fsq_levels, + num_quantizers=fsq_num_quantizers, + dim=latent_dim, # Quantizer operates on perceiver output dim + channel_first=False, # Perceiver output is (B, T, D), so channel_first=False + quantize_dropout=False, # No dropout specified in config + ) + + # Final projection from flattened quantized tokens to the target output dimension + self.project = nn.Linear(latent_dim * token_num, out_dim) + + def get_codes_from_indices(self, indices: torch.Tensor) -> torch.Tensor: + """Reconstruct quantized vectors from indices.""" + # indices shape: (B, T_token, Q) or (B, Q, T_token)? Check ResidualFSQ output. + # Assuming (B, T_token, Q) from forward pass. + # get_output_from_indices expects (indices_chunk1, indices_chunk2, ...) if grouped. + # If not grouped, expects (B, ... Q). Let's assume (B, T_token, Q). + zq = self.quantizer.get_output_from_indices(indices) + # Output zq shape should be (B, T_token, latent_dim) + return zq + + def get_indices(self, mels: torch.Tensor) -> torch.Tensor: + """Get FSQ indices directly from mel spectrograms.""" + # mels: (B, T_mel, D_mel) + _, features = self.speaker_encoder_base(mels, return_latent=True) # features: (B, ecapa_feat_dim, T_feat) + x = self.perceiver_sampler(features.transpose(1, 2)) # Input: (B, T_feat, ecapa_feat_dim), Output: (B, token_num, latent_dim) + _, indices = self.quantizer(x) # Input: (B, token_num, latent_dim), indices: (B, token_num, Q) + return indices + + def forward(self, mels: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Args: + mels: (B, T_mel, D_mel) - Mel spectrogram input + + Return: + x_vector: (B, ecapa_embed_dim) - Global speaker embedding from ECAPA + d_vector: (B, out_dim) - Speaker embedding derived from quantized tokens + """ + # Get base speaker embedding (x-vector) and intermediate features from ECAPA + x_vector, features = self.speaker_encoder_base(mels, return_latent=True) + # features shape: (B, ecapa_feat_dim, T_feat) + + # Resample features using Perceiver + # Perceiver expects (B, T, D), so transpose features + perceiver_latents = self.perceiver_sampler(features.transpose(1, 2)) + # perceiver_latents shape: (B, token_num, latent_dim) + + # Quantize the perceiver latents + # Quantizer expects (B, T, D) if channel_first=False + zq, indices = self.quantizer(perceiver_latents) + # zq shape: (B, token_num, latent_dim), indices shape: (B, token_num, Q) + + # Flatten quantized tokens and project to final d-vector dimension + zq_flat = rearrange(zq, 'b t d -> b (t d)') # (B, token_num * latent_dim) + d_vector = self.project(zq_flat) # (B, out_dim) + + return x_vector, d_vector + + def tokenize(self, mels: torch.Tensor) -> torch.Tensor: + """Tokenize the input mel spectrogram to get FSQ indices.""" + # Same logic as get_indices + _, features = self.speaker_encoder_base(mels, return_latent=True) # features: (B, ecapa_feat_dim, T_feat) + x = self.perceiver_sampler(features.transpose(1, 2)) # (B, token_num, latent_dim) + _, indices = self.quantizer(x) # indices: (B, token_num, Q) + return indices + + def detokenize(self, indices: torch.Tensor) -> torch.Tensor: + """Detokenize FSQ indices to get the final d-vector.""" + # indices shape: (B, token_num, Q) + # Reconstruct quantized vectors from indices + zq = self.get_codes_from_indices(indices) # (B, token_num, latent_dim) + + # Flatten and project + zq_flat = rearrange(zq, 'b t d -> b (t d)') + d_vector = self.project(zq_flat) + return d_vector + +# =============================================================== +# End: Content from sparktts/modules/speaker/speaker_encoder.py +# =============================================================== + + +# =============================================================== +# Start: Content from sparktts/modules/vq/factorized_vector_quantize.py +# =============================================================== +# Helper function from layers.py (already defined) +# def WNConv1d(*args, **kwargs): +# return weight_norm(nn.Conv1d(*args, **kwargs)) + +def ema_inplace(moving_avg, new, decay): + moving_avg.data.mul_(decay).add_(new, alpha=(1 - decay)) + + +class FactorizedVectorQuantize(nn.Module): + def __init__( + self, + input_dim: int, + codebook_size: int, + codebook_dim: int, + commitment: float, + codebook_loss_weight: float = 1.0, + decay: float = 0.99, + threshold_ema_dead_code: float = 2.0, # Changed default from 2 based on config + momentum: float = 0.99, # Not used in current implementation? + use_l2_normlize: bool = True, # Added from config + **kwargs, + ): + super().__init__() + self.input_dim = input_dim + self.codebook_size = codebook_size + self.codebook_dim = codebook_dim + self.commitment = commitment + self.codebook_loss_weight = codebook_loss_weight + self.decay = decay + self.threshold_ema_dead_code = threshold_ema_dead_code + # self.momentum = momentum # Store if needed later + self.use_l2_normlize = use_l2_normlize + + if input_dim != self.codebook_dim: + self.in_project = WNConv1d(input_dim, self.codebook_dim, kernel_size=1) + self.out_project = WNConv1d(self.codebook_dim, input_dim, kernel_size=1) + else: + self.in_project = nn.Identity() + self.out_project = nn.Identity() + + # Codebook embedding layer + self.codebook = nn.Embedding(self.codebook_size, self.codebook_dim) + # Initialize codebook? Often random init is fine. + + # Buffers for EMA updates (cluster size and maybe embeddings) + self.register_buffer("cluster_size", torch.zeros(self.codebook_size)) + # EMA average embeddings (optional, can use self.codebook.weight directly for loss) + # self.register_buffer("ema_embed", self.codebook.weight.clone()) + + def forward(self, z: torch.Tensor) -> Dict[str, Any]: + """Quantizes the input tensor using a fixed codebook and returns + the corresponding codebook vectors and losses. + + Parameters + ---------- + z : Tensor[B x D_in x T] + + Returns + ------- + Dict containing: + z_q (Tensor[B x D_in x T]): Quantized continuous representation (passed through out_project) + indices (Tensor[B x T]): Codebook indices + vq_loss (Tensor[1]): Combined VQ loss (codebook + commitment) + perplexity (Tensor[1]): Codebook perplexity metric + active_num (Tensor[1]): Number of active codebook entries + """ + # z: (B, D_in, T) + B, _, T = z.shape + + # Project input to codebook dimension if necessary + z_e = self.in_project(z) # (B, D_code, T) + + # Find nearest neighbors and get quantized vectors + indices + z_q, indices, dists = self.decode_latents(z_e) # z_q: (B, D_code, T), indices: (B, T) + + # Calculate statistics for perplexity and active codes + with torch.no_grad(): # Stats should not contribute to gradient + embed_onehot = F.one_hot(indices, self.codebook_size).type(z_e.dtype) # (B, T, C) + # Flatten batch and time dims for stats + embed_onehot_flat = rearrange(embed_onehot, 'b t c -> (b t) c') + avg_probs = torch.mean(embed_onehot_flat, dim=0) # (C,) + perplexity = torch.exp(-torch.sum(avg_probs * torch.log(avg_probs + 1e-10))) + + # EMA update for cluster size (only in training) + active_num_tensor = (embed_onehot_flat.sum(0) > 0).sum() # Before EMA + if self.training: + # Perform EMA update in place + ema_inplace(self.cluster_size, embed_onehot_flat.sum(0), self.decay) + # Calculate active codes based on EMA threshold + active_num_tensor = (self.cluster_size > self.threshold_ema_dead_code).sum() + + + # Calculate losses (only in training) + commit_loss = torch.tensor(0.0, device=z.device) + codebook_loss = torch.tensor(0.0, device=z.device) + vq_loss = torch.tensor(0.0, device=z.device) + + if self.training: + # Commitment loss (encourage encoder output E(x) to be close to codebook z_q) + # Use z_e (projected encoder output) and z_q.detach() + commit_loss = F.mse_loss(z_e, z_q.detach()) * self.commitment + + # Codebook loss (encourage codebook entries z_q to be close to encoder output E(x)) + # Use z_q and z_e.detach() + codebook_loss = F.mse_loss(z_q, z_e.detach()) * self.codebook_loss_weight + + vq_loss = commit_loss + codebook_loss + + + # Straight-through estimator: copy gradient from z_q to z_e + z_q_st = z_e + (z_q - z_e).detach() + + # Project quantized vectors back to input dimension if necessary + z_q_out = self.out_project(z_q_st) # (B, D_in, T) + + return { + "z_q": z_q_out, + "indices": indices, + # "dists": dists, # Dists might be large, exclude unless needed + "vq_loss": vq_loss, + "perplexity": perplexity, + "active_num": active_num_tensor.float(), + } + + def embed_code(self, embed_id): + """Retrieve codebook vectors for given indices.""" + return F.embedding(embed_id, self.codebook.weight) + + def decode_code(self, embed_id): + """Retrieve codebook vectors and transpose to (B, D, T) format.""" + # embed_id: (B, T) + # Embedding: (B, T, D_code) + # Transpose: (B, D_code, T) + return self.embed_code(embed_id).transpose(1, 2) + + def decode_latents(self, latents): + """Find nearest codebook entries for latent vectors.""" + # latents: (B, D_code, T) + B, D_code, T = latents.shape + encodings = rearrange(latents, "b d t -> (b t) d") # ((B*T), D_code) + codebook = self.codebook.weight # (C, D_code) + + # Normalize if required + if self.use_l2_normlize: + encodings = F.normalize(encodings, p=2, dim=-1) + codebook = F.normalize(codebook, p=2, dim=-1) + + # Compute distances (squared Euclidean or Cosine depending on normalization) + # dist = torch.cdist(encodings, codebook, p=2)**2 # Squared Euclidean + # Faster calculation using matrix multiplication if normalized: + # dist = 2 - 2 * (encodings @ codebook.t()) + # Or full squared Euclidean: + dist = ( + encodings.pow(2).sum(1, keepdim=True) # (B*T, 1) + - 2 * (encodings @ codebook.t()) # (B*T, C) + + codebook.pow(2).sum(1, keepdim=True).t() # (1, C) + ) # Result shape: (B*T, C) + + # Find nearest neighbors + indices = torch.argmin(dist, dim=-1) # (B*T) + indices = rearrange(indices, "(b t) -> b t", b=B) # (B, T) + + # Get the quantized vectors + z_q = self.decode_code(indices) # (B, D_code, T) + + return z_q, indices, dist # Return dist if needed, e.g., for debugging + + # --- Methods for inference/tokenization --- + def tokenize(self, z: torch.Tensor) -> torch.Tensor: + """Tokenize the input tensor without loss calculation.""" + # z: (B, D_in, T) + z_e = self.in_project(z) # (B, D_code, T) + _, indices, _ = self.decode_latents(z_e) # indices: (B, T) + return indices + + def detokenize(self, indices: torch.Tensor) -> torch.Tensor: + """Detokenize indices to quantized vectors in input dimension.""" + # indices: (B, T) + z_q_code_dim = self.decode_code(indices) # (B, D_code, T) + z_q_out = self.out_project(z_q_code_dim) # (B, D_in, T) + return z_q_out + +# =============================================================== +# End: Content from sparktts/modules/vq/factorized_vector_quantize.py +# =============================================================== + + + +# --- BiCodec Model Definition (Adapted from sparktts/models/bicodec.py) --- +class BiCodec(nn.Module): + """ + BiCodec model for speech synthesis, incorporating a speaker encoder, feature encoder/decoder, + quantizer, and wave generator. + """ + + def __init__( + self, + mel_params: Dict[str, Any], + encoder: nn.Module, + decoder: nn.Module, + quantizer: nn.Module, + speaker_encoder: nn.Module, + prenet: nn.Module, + postnet: nn.Module, + **kwargs + ) -> None: + """ + Initializes the BiCodec model with the required components. + + Args: + mel_params (dict): Parameters for the mel-spectrogram transformer. + encoder (nn.Module): Encoder module. + decoder (nn.Module): Decoder module. + quantizer (nn.Module): Quantizer module. + speaker_encoder (nn.Module): Speaker encoder module. + prenet (nn.Module): Prenet network. + postnet (nn.Module): Postnet network. + """ + super().__init__() + self.encoder = encoder + self.decoder = decoder + self.quantizer = quantizer + self.speaker_encoder = speaker_encoder + self.prenet = prenet + self.postnet = postnet + self._init_mel_transformer(mel_params) + + @classmethod + def load_from_config_and_checkpoint(cls, model_dir: Path, config_dict: Dict[str, Any], **kwargs) -> "BiCodec": + """Loads the model from a config dictionary and checkpoint file.""" + ckpt_path = model_dir / 'model.safetensors' + if not ckpt_path.is_file(): + raise FileNotFoundError(f"BiCodec checkpoint not found at {ckpt_path}") + + audio_tokenizer_config = config_dict # Assuming config_dict holds the relevant sub-config + + # Instantiate components using classes from _modeling_bicodec_components + mel_params = audio_tokenizer_config.get("mel_params", {}) + encoder_cfg = audio_tokenizer_config.get("encoder", {}) + quantizer_cfg = audio_tokenizer_config.get("quantizer", {}) + prenet_cfg = audio_tokenizer_config.get("prenet", {}) + postnet_cfg = audio_tokenizer_config.get("postnet", {}) + decoder_cfg = audio_tokenizer_config.get("decoder", {}) # This corresponds to WaveGenerator + speaker_encoder_cfg = audio_tokenizer_config.get("speaker_encoder", {}) + + # --- Input Validation --- + required_keys = { + "encoder": ["input_channels", "vocos_dim", "vocos_intermediate_dim", "vocos_num_layers", "out_channels"], + "quantizer": ["input_dim", "codebook_size", "codebook_dim", "commitment"], + "prenet": ["input_channels", "vocos_dim", "vocos_intermediate_dim", "vocos_num_layers", "out_channels"], + "postnet": ["input_channels", "vocos_dim", "vocos_intermediate_dim", "vocos_num_layers", "out_channels"], + "decoder": ["input_channel", "channels", "rates", "kernel_sizes"], # WaveGenerator keys + "speaker_encoder": ["input_dim", "out_dim", "latent_dim", "token_num"], + "mel_params": ["sample_rate", "n_fft", "win_length", "hop_length", "num_mels"] + } + for comp, keys in required_keys.items(): + cfg = audio_tokenizer_config.get(comp, {}) + if not cfg: logging.get_logger(__name__).warning(f"BiCodec config missing section: '{comp}'") + for key in keys: + if key not in cfg: + logging.get_logger(__name__).warning(f"BiCodec config missing key '{key}' in section '{comp}'") + # --- End Validation --- + + + # Instantiate modules + encoder = Encoder(**encoder_cfg) if encoder_cfg else None + quantizer = FactorizedVectorQuantize(**quantizer_cfg) if quantizer_cfg else None + prenet = Decoder(**prenet_cfg) if prenet_cfg else None + postnet = Decoder(**postnet_cfg) if postnet_cfg else None + decoder = WaveGenerator(**decoder_cfg) if decoder_cfg else None # WaveGenerator instance + speaker_encoder = SpeakerEncoder(**speaker_encoder_cfg) if speaker_encoder_cfg else None + + # Check if all components were successfully created + if not all([encoder, quantizer, prenet, postnet, decoder, speaker_encoder, mel_params]): + raise ValueError("Failed to initialize one or more BiCodec components due to missing configuration.") + + # Create the BiCodec instance + model = cls( + mel_params=mel_params, + encoder=encoder, + decoder=decoder, # Pass WaveGenerator instance as decoder + quantizer=quantizer, + speaker_encoder=speaker_encoder, + prenet=prenet, + postnet=postnet, + ) + + # Load state dict + try: + state_dict = load_file(ckpt_path, device="cpu") # Load to CPU first + missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False) + if missing_keys: + print(f"BiCodec missing keys: {missing_keys}") + if unexpected_keys: + print(f"BiCodec unexpected keys: {unexpected_keys}") + except Exception as e: + raise IOError(f"Error loading BiCodec state dict from {ckpt_path}: {e}") + + model.eval() + # model.remove_weight_norm() # Assuming this method exists in components + + return model + + def _init_mel_transformer(self, config: Dict[str, Any]): + # Ensure required keys exist with defaults + sr = config.get("sample_rate", 16000) + n_fft = config.get("n_fft", 1024) + win_length = config.get("win_length", n_fft) + hop_length = config.get("hop_length", n_fft // 4) + fmin = config.get("mel_fmin", 0) + fmax = config.get("mel_fmax", None) + n_mels = config.get("num_mels", 80) + power = config.get("power", 2.0) # Typically 2.0 for power spectrogram + norm = config.get("norm", "slaney") + mel_scale = config.get("mel_scale", "htk") # htk or slaney + + self.mel_transformer = TT.MelSpectrogram( + sample_rate=sr, + n_fft=n_fft, + win_length=win_length, + hop_length=hop_length, + f_min=fmin, + f_max=fmax, + n_mels=n_mels, + power=power, + norm=norm, + mel_scale=mel_scale, + ).eval() # Set to eval mode + + def remove_weight_norm(self): + """Removes weight normalization from components that support it.""" + def _remove_wn(m): + if hasattr(m, 'remove_weight_norm'): + m.remove_weight_norm() + elif isinstance(m, (nn.Conv1d, nn.ConvTranspose1d)): + try: + remove_weight_norm(m) + except ValueError: + pass # Module might not have weight norm applied + + self.apply(_remove_wn) + + + @torch.no_grad() + def tokenize(self, feat: torch.Tensor, ref_wav: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + """ Tokenizes input features and reference wav into semantic and global tokens. """ + # Ensure models are on the correct device + device = feat.device + self.mel_transformer.to(device) + self.encoder.to(device) + self.quantizer.to(device) + self.speaker_encoder.to(device) + + # feat: (B, D_feat, T_feat), ref_wav: (B, T_wav) + mel = self.mel_transformer(ref_wav) # (B, D_mel, T_mel) + + # Encode features to get latents for semantic tokens + z = self.encoder(feat) # (B, D_latent, T_latent) - Assuming Encoder output matches quantizer input dim + + # Quantize latents to get semantic tokens (indices) + semantic_tokens = self.quantizer.tokenize(z) # (B, T_latent) + + # Encode mel spectrogram to get global tokens (indices) + # SpeakerEncoder.tokenize expects (B, T_mel, D_mel) + global_tokens = self.speaker_encoder.tokenize(mel.transpose(1, 2)) # (B, T_token, Q) - Check shape + + # Note: Original BiCodecTokenizer returned (global_tokens, semantic_tokens) + # Let's stick to that order for consistency with original SparkTTS usage. + return global_tokens, semantic_tokens + + + @torch.no_grad() + def detokenize(self, global_tokens: torch.Tensor, semantic_tokens: torch.Tensor) -> torch.Tensor: + """ Detokenizes semantic and global tokens into a waveform. """ + # Ensure models are on the correct device + device = semantic_tokens.device # Assume tokens are on target device + self.quantizer.to(device) + self.speaker_encoder.to(device) + self.prenet.to(device) + self.decoder.to(device) # WaveGenerator + + # semantic_tokens: (B, T_latent) or (B, T_latent, Q)? Check quantizer.tokenize output shape. Assuming (B, T_latent). + # global_tokens: (B, T_token, Q) - Check speaker_encoder.tokenize output shape. + + # Reconstruct quantized vectors from semantic tokens + z_q = self.quantizer.detokenize(semantic_tokens) # (B, D_latent, T_latent) + + # Reconstruct d-vector (condition) from global tokens + # SpeakerEncoder.detokenize expects (B, T_token, Q) + d_vector = self.speaker_encoder.detokenize(global_tokens) # (B, D_dvector) + + # Apply prenet conditioned on d-vector + # Prenet (Decoder class) expects input (B, D_latent, T_latent) and condition (B, D_dvector) + x = self.prenet(z_q, d_vector) # (B, D_prenet_out, T_latent) - Assuming prenet maintains time dim + + # Add condition (broadcasted) before wave generation - Check original logic + # Ensure d_vector has correct shape for broadcasting + if d_vector.ndim == 2: + d_vector_unsqueezed = d_vector.unsqueeze(-1) # (B, D_dvector, 1) + else: # Should not happen if speaker_encoder outputs (B, D) + d_vector_unsqueezed = d_vector + + # Ensure dimensions match for addition + if x.shape[1] == d_vector_unsqueezed.shape[1]: + # Broadcast d_vector across time dimension T_latent + x = x + d_vector_unsqueezed + else: + # Maybe project d_vector or x? Log a warning or adapt based on expected dims. + logging.get_logger(__name__).warning(f"Prenet output dim {x.shape[1]} != d-vector dim {d_vector_unsqueezed.shape[1]}. Skipping residual connection.") + + + # Generate waveform using the decoder (WaveGenerator) + # WaveGenerator expects (B, D_input, T_input) + wav_recon = self.decoder(x) # (B, 1, T_wav) + + return wav_recon + + +# --- Main SparkTTS Model --- +from .configuration_spark_tts import SparkTTSConfig +# from ._utils import load_audio # Use utils from _utils.py + +logger = logging.get_logger(__name__) + +class SparkTTSModel(PreTrainedModel, GenerationMixin): + """ + SparkTTS model integrating LLM, BiCodec, and Wav2Vec2 for text-to-speech. + """ + config_class = SparkTTSConfig + base_model_prefix = "spark_tts" + _supports_load_fast = False + + def __init__(self, config: SparkTTSConfig, llm=None, wav2vec2_model=None, wav2vec2_processor=None, bicodec=None): + super().__init__(config) + self.config = config + self.llm = llm + self.wav2vec2_model = wav2vec2_model + self.wav2vec2_processor = wav2vec2_processor + self.bicodec = bicodec + + # Wav2Vec2 specific config adjustment (needs to happen after loading) + if self.wav2vec2_model and hasattr(self.wav2vec2_model.config, 'output_hidden_states'): + self.wav2vec2_model.config.output_hidden_states = True + + + @classmethod + def from_pretrained( + cls, + pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], + *model_args, + config: Optional[Union[PretrainedConfig, str, os.PathLike]] = None, + cache_dir: Optional[Union[str, os.PathLike]] = None, + ignore_mismatched_sizes: bool = False, + force_download: bool = False, + local_files_only: bool = False, + token: Optional[Union[str, bool]] = None, + revision: str = "main", + use_safetensors: bool = None, + **kwargs, + ): + # 1. Load Config + if config is None: + config, model_kwargs = cls.config_class.from_pretrained( + pretrained_model_name_or_path, + *model_args, + cache_dir=cache_dir, + force_download=force_download, + local_files_only=local_files_only, + token=token, + revision=revision, + return_unused_kwargs=True, + **kwargs, + ) + else: + model_kwargs = kwargs + + # Pop device map info - will handle placement later + device_map = model_kwargs.pop("device_map", None) + torch_dtype = model_kwargs.pop("torch_dtype", "auto") # Use config's or auto + + # Check for trust_remote_code - needed for config loading if custom code involved there too + trust_remote_code = model_kwargs.pop("trust_remote_code", False) # Important + + # Determine actual model directory (could be cache path) + if pretrained_model_name_or_path is not None: + resolved_model_path = Path(pretrained_model_name_or_path) + if not resolved_model_path.is_dir(): + # Attempt to download and resolve cache path if it's an ID + # This requires internet connection if not cached + try: + resolved_model_path = Path(cached_file( + pretrained_model_name_or_path, + filename=cls.config_class.config_files[0], # e.g., "config.json" + cache_dir=cache_dir, + force_download=force_download, + local_files_only=local_files_only, + token=token, + revision=revision, + _raise_exceptions_for_missing_entries=False, + _raise_exceptions_for_connection_errors=False, + )).parent + except Exception as e: + logger.warning(f"Could not resolve cache path for {pretrained_model_name_or_path}: {e}. Assuming it's a local path.") + resolved_model_path = Path(pretrained_model_name_or_path) # Fallback + if not resolved_model_path.is_dir(): + raise EnvironmentError(f"Cannot find model directory at {resolved_model_path}") + else: + raise ValueError("pretrained_model_name_or_path must be provided.") + + + # Helper function to resolve paths relative to the main model directory + def _resolve_path(sub_path): + p = Path(sub_path) + if p.is_absolute(): + return str(p) + else: + # Resolve relative to the potentially cached main model path + return str(resolved_model_path / p) + + # --- Load LLM --- + llm_path = _resolve_path(config.llm_model_name_or_path) + logger.info(f"Loading LLM from resolved path: {llm_path}") + try: + llm = AutoModelForCausalLM.from_pretrained( + llm_path, + torch_dtype=torch_dtype if torch_dtype != "auto" else config.torch_dtype, # Prioritize explicit dtype + trust_remote_code=trust_remote_code, # Pass down trust_remote_code + **model_kwargs # Pass remaining kwargs + ) + except Exception as e: + raise OSError(f"Failed to load LLM from {llm_path}: {e}") + + # --- Load Wav2Vec2 --- + w2v_path = _resolve_path(config.wav2vec2_model_name_or_path) + logger.info(f"Loading Wav2Vec2 from resolved path: {w2v_path}") + try: + wav2vec2_processor = Wav2Vec2FeatureExtractor.from_pretrained(w2v_path, trust_remote_code=trust_remote_code) + wav2vec2_model = Wav2Vec2Model.from_pretrained(w2v_path, trust_remote_code=trust_remote_code) + except Exception as e: + raise OSError(f"Failed to load Wav2Vec2 components from {w2v_path}: {e}") + + # --- Load BiCodec --- + bicodec_path = _resolve_path(config.bicodec_model_name_or_path) + logger.info(f"Loading BiCodec from resolved path: {bicodec_path}") + # print(f"Loading BiCodec from resolved path: {bicodec_path}, {config}") + if not config.bicodec_config or "audio_tokenizer" not in config.bicodec_config: + raise ValueError("BiCodec configuration ('bicodec_config' with 'audio_tokenizer' key) not found in SparkTTSConfig.") + try: + # Assuming BiCodec class is defined above in this file + bicodec = BiCodec.load_from_config_and_checkpoint( + model_dir=Path(bicodec_path), + config_dict=config.bicodec_config["audio_tokenizer"] + ) + except Exception as e: + raise OSError(f"Failed to load BiCodec from {bicodec_path}: {e}") + + + # Instantiate the main model wrapper, passing the loaded components + model = cls(config, llm=llm, wav2vec2_model=wav2vec2_model, wav2vec2_processor=wav2vec2_processor, bicodec=bicodec) + + # --- Handle device placement --- + # Note: device_map is complex; simple .to(device) is easier if not using accelerate + # Determine target device + if torch.cuda.is_available(): + current_device = torch.cuda.current_device() + device = torch.device(f"cuda:{current_device}") + else: + device = torch.device("cpu") + logger.info(f"Placing SparkTTSModel and components on device: {device}") + model.to(device) # This should move all registered nn.Module attributes + + return model + + + # --- Embedding getters/setters (delegate to LLM if loaded) --- + def get_input_embeddings(self): + if self.llm: + return self.llm.get_input_embeddings() + return None # Or raise error + + def set_input_embeddings(self, value): + if self.llm: + self.llm.set_input_embeddings(value) + else: + logger.warning("LLM not loaded, cannot set input embeddings.") + + def get_output_embeddings(self): + if self.llm: + # For causal LM, output embeddings are usually tied to lm_head + return self.llm.get_output_embeddings() + return None # Or raise error + + def set_output_embeddings(self, new_embeddings): + if self.llm and hasattr(self.llm, 'set_output_embeddings'): + self.llm.set_output_embeddings(new_embeddings) + else: + logger.warning("LLM not loaded or does not support set_output_embeddings.") + # --- End Embedding methods --- + + # post_init is less critical now as loading happens in from_pretrained, + # but can be used for final checks or setup. + def post_init(self): + # Ensure wav2vec2 config has output_hidden_states=True + if self.wav2vec2_model and hasattr(self.wav2vec2_model.config, 'output_hidden_states'): + if not self.wav2vec2_model.config.output_hidden_states: + self.wav2vec2_model.config.output_hidden_states = True + logger.info("Set wav2vec2_model.config.output_hidden_states=True") + + @property + def device(self) -> torch.device: + """ Override device property to report the LLM's device as representative """ + if self.llm: + return self.llm.device + else: + # Fallback or default if LLM not loaded yet + # This might be called by pipeline before full init? Be cautious. + try: + return next(self.parameters()).device + except StopIteration: + # If no parameters, default to CPU + return torch.device("cpu") + + @torch.no_grad() + def _extract_wav2vec2_features(self, wavs: torch.Tensor) -> torch.Tensor: + """Extract wav2vec2 features. Input wavs: (B, T_wav)""" + if not self.wav2vec2_model or not self.wav2vec2_processor: + raise RuntimeError("Wav2Vec2 components not loaded.") + + # Use component's device + target_device = self.wav2vec2_model.device + wavs_on_device = wavs.to(target_device) # Expected shape [B, T_wav] e.g., [1, 61120] + + # Process audio using the Wav2Vec2FeatureExtractor + processor_output = self.wav2vec2_processor( + wavs_on_device, + sampling_rate=self.config.sample_rate, + return_tensors="pt", + padding=True, # Ensure padding is handled correctly + ) + inputs = processor_output.input_values # Should be shape [B, T_processed] + # --- START DEBUG & FIX --- + print(f"Shape returned by processor: {inputs.shape}") + + # Reshape if processor added extra dimensions + if inputs.ndim == 4 and inputs.shape[1] == 1 and inputs.shape[2] == 1: + print(f"Reshaping input from {inputs.shape} to 2D.") + inputs = inputs.squeeze(1).squeeze(1) # Remove the two middle dimensions + elif inputs.ndim == 3 and inputs.shape[1] == 1: + print(f"Reshaping input from {inputs.shape} to 2D.") + inputs = inputs.squeeze(1) # Remove the channel dimension + + # Ensure final shape is 2D: (batch_size, sequence_length) + if inputs.ndim != 2: + raise ValueError(f"Unexpected shape after processing/reshaping: {inputs.shape}. Expected 2D input for Wav2Vec2Model.") + + print(f"Shape BEFORE Wav2Vec2Model: {inputs.shape}") + # --- END DEBUG & FIX --- + + inputs = inputs.to(target_device) + # Ensure output_hidden_states=True during call if not set reliably in config + outputs = self.wav2vec2_model(inputs, output_hidden_states=True) + + if outputs.hidden_states is None: + raise ValueError("Wav2Vec2 model did not return hidden states. Ensure config.output_hidden_states=True.") + + # Mix specific layers + num_layers = len(outputs.hidden_states) + indices_to_mix = [11, 14, 16] + valid_indices = [i for i in indices_to_mix if i < num_layers] + + if len(valid_indices) != len(indices_to_mix): + logger.warning(f"Requested Wav2Vec2 hidden state indices {indices_to_mix} out of range (0-{num_layers-1}). Using available valid indices: {valid_indices}.") + if not valid_indices: # If no valid indices, use last hidden state + logger.warning("No valid hidden state indices for mixing. Using last hidden state.") + feats_mix = outputs.last_hidden_state + else: + # Mix available valid indices + feats_mix = torch.stack([outputs.hidden_states[i] for i in valid_indices]).mean(dim=0) + else: + # Original mixing logic + feats_mix = (outputs.hidden_states[11] + outputs.hidden_states[14] + outputs.hidden_states[16]) / 3 + + # Output shape: (B, T_feat, D_feat) - Transpose needed for BiCodec Encoder + return feats_mix.transpose(1, 2) # (B, D_feat, T_feat) + + def _get_ref_clip(self, wav: np.ndarray) -> np.ndarray: + """Get reference audio clip for speaker embedding.""" + ref_samples = int(self.config.sample_rate * self.config.ref_segment_duration) + latent_hop_length = self.config.latent_hop_length + # Ensure length is multiple of hop_length for potential downstream processing + ref_segment_length = max(latent_hop_length, (ref_samples // latent_hop_length) * latent_hop_length) # Ensure at least one hop + + wav_length = len(wav) + + if wav_length == 0: # Handle empty input + return np.zeros(ref_segment_length, dtype=np.float32) + if ref_segment_length > wav_length: + num_repeats = (ref_segment_length // wav_length) + 1 + wav = np.tile(wav, num_repeats) + + return wav[:ref_segment_length].astype(np.float32) # Ensure float32 + + + @torch.no_grad() + def _tokenize_audio(self, audio_path: str) -> Tuple[torch.Tensor, torch.Tensor]: + """Load audio, extract features, and tokenize using BiCodec.""" + wav_np = load_audio( + audio_path, + sampling_rate=self.config.sample_rate, + volume_normalize=self.config.volume_normalize, + ) + wav_ref_np = self._get_ref_clip(wav_np) + + # Convert to tensors, add batch dim, move to device + wav = torch.from_numpy(wav_np).unsqueeze(0).float().to(self.device) + ref_wav = torch.from_numpy(wav_ref_np).unsqueeze(0).float().to(self.device) + + # Extract Wav2Vec2 features -> (B, D_feat, T_feat) + feat = self._extract_wav2vec2_features(wav) + + # Tokenize using BiCodec -> returns (global_tokens, semantic_tokens) + # BiCodec.tokenize expects feat: (B, D_feat, T_feat), ref_wav: (B, T_wav) + global_tokens, semantic_tokens = self.bicodec.tokenize(feat, ref_wav) + # global_tokens: (B, T_token, Q), semantic_tokens: (B, T_latent) + + return global_tokens, semantic_tokens + + @torch.no_grad() + def _detokenize_audio(self, global_tokens: torch.Tensor, semantic_tokens: torch.Tensor) -> np.ndarray: + """Detokenize using BiCodec to get waveform.""" + global_tokens = global_tokens.to(self.device) + semantic_tokens = semantic_tokens.to(self.device) + self.bicodec.to(self.device) # Ensure BiCodec is on device + + # BiCodec.detokenize expects global_tokens: (B, T_token, Q), semantic_tokens: (B, T_latent) + wav_rec = self.bicodec.detokenize(global_tokens, semantic_tokens) # (B, 1, T_wav) + # Remove channel dim and batch dim, convert to numpy + return wav_rec.detach().squeeze(0).squeeze(0).cpu().numpy() + + + def forward(self, *args, **kwargs): + """ Forward pass delegates to the LLM for generation compatibility, but direct use is not intended for TTS. """ + # return self.llm(*args, **kwargs) # Option 1: Delegate fully + logger.warning("Direct forward pass on SparkTTSModel is not the intended use for TTS. Use the generate method or pipeline.") + # Option 2: Minimal implementation for compatibility if needed + if 'input_ids' in kwargs: + return self.llm(input_ids=kwargs['input_ids']) + else: + raise NotImplementedError("SparkTTSModel's forward pass requires 'input_ids' or should not be called directly for TTS.") + + + + # Use GenerationMixin's forward method by default if needed. + # Define prepare_inputs_for_generation if LLM needs specific handling. + def prepare_inputs_for_generation(self, input_ids, **kwargs): + """ Prepares inputs for the LLM's generate method. """ + if not self.llm: + raise RuntimeError("LLM component not loaded.") + + # --- START REVISED IMPLEMENTATION --- + # Delegate to the LLM's prepare_inputs_for_generation method directly. + # This ensures we use the exact logic defined for the specific LLM architecture (Qwen2). + # It should handle past_key_values, attention_mask, use_cache etc. correctly. + try: + # Pass all relevant kwargs received by the top-level generate call + # The LLM's method will select what it needs. + model_inputs = self.llm.prepare_inputs_for_generation(input_ids, **kwargs) + return model_inputs + except AttributeError: + # Fallback if the LLM doesn't have this method (unlikely for recent models) + logger.warning("LLM does not have 'prepare_inputs_for_generation'. Using basic fallback.") + model_kwargs = {} + model_kwargs["past_key_values"] = kwargs.get("past_key_values", None) + model_kwargs["use_cache"] = kwargs.get("use_cache", None) + # Ensure attention_mask is included if present in kwargs + if "attention_mask" in kwargs: + model_kwargs["attention_mask"] = kwargs["attention_mask"] + return {"input_ids": input_ids, **model_kwargs} + # --- END REVISED IMPLEMENTATION --- + + + # We need a minimal forward method compatible with GenerationMixin + # It should accept the output of prepare_inputs_for_generation + def forward( + self, + input_ids: Optional[torch.Tensor] = None, + past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + **kwargs # Accept other potential kwargs from prepare_inputs + ) -> Any: # Return type depends on the LLM, usually CausalLMOutputWithPast + """ + Minimal forward pass that delegates to the underlying LLM. + Required for compatibility with GenerationMixin. + Accepts arguments typically returned by prepare_inputs_for_generation. + """ + if not self.llm: + raise RuntimeError("LLM component not loaded.") + + # Filter arguments for the LLM's forward method + # (Some LLMs might not accept position_ids directly in forward when using past_key_values) + llm_kwargs = { + "past_key_values": past_key_values, + "attention_mask": attention_mask, + **kwargs # Pass through any other relevant kwargs + } + # Only pass position_ids if the LLM's forward signature accepts it + # This requires inspecting the LLM's forward signature or knowing its behavior. + # For simplicity, we might omit it if it causes issues, or handle it more dynamically. + # Let's assume the LLM forward can handle it for now if prepare_inputs included it. + if position_ids is not None: + llm_kwargs["position_ids"] = position_ids + + return self.llm(input_ids=input_ids, **llm_kwargs) + + # Add generate method to use GenerationMixin capabilities directly on SparkTTSModel if desired + # This will internally call prepare_inputs_for_generation and forward (which might need defining/adjusting) + # However, the pipeline calls self.model.llm.generate, so this might not be strictly needed unless you want `model.generate(...)` + # @torch.no_grad() + # def generate(self, *args, **kwargs): + # if not self.llm: + # raise RuntimeError("LLM component not loaded.") + # # This might need adjustments based on how GenerationMixin interacts with the overridden forward + # # return super().generate(*args, **kwargs) # Calls self.prepare_inputs + self.forward loop + # # Or directly call the LLM's generate if forward is problematic: + # return self.llm.generate(*args, **kwargs) +