import math import typing as tp from functools import partial from dataclasses import dataclass, field from typing import Dict, List, Optional, Tuple, Union import copy import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from transformers.models.auto import AutoModel from transformers.configuration_utils import PretrainedConfig from transformers.utils import logging from transformers.modeling_utils import PreTrainedModel from transformers.activations import ACT2FN from .configuration_vibevoice import VibeVoiceAcousticTokenizerConfig, VibeVoiceSemanticTokenizerConfig logger = logging.get_logger(__name__) import os # Try to import APEX FusedRMSNorm try: from apex.normalization.fused_layer_norm import fused_rms_norm_affine APEX_AVAILABLE = True logger.info("APEX FusedRMSNorm is available and will be used for optimization") if int(os.getenv("OPTIMIZE_FOR_SPEED", "0")) == 0: APEX_AVAILABLE = False logger.warning("APEX FusedRMSNorm is disabled by environment variable OPTIMIZE_FOR_SPEED=0") except ImportError: APEX_AVAILABLE = False logger.warning("APEX FusedRMSNorm not available, using native implementation") # APEX_AVAILABLE=False # Normalization modules class ConvLayerNorm(nn.LayerNorm): """ Convolution-friendly LayerNorm that moves channels to last dimensions before running the normalization and moves them back to original position right after. """ def __init__(self, normalized_shape: tp.Union[int, tp.List[int], torch.Size], **kwargs): super().__init__(normalized_shape, **kwargs) def forward(self, x): x = x.transpose(1, 2) # b ... t -> b t ... x = nn.functional.layer_norm(x.float(), self.normalized_shape, self.weight.float(), self.bias.float(), self.eps).type_as(x) x = x.transpose(1, 2) # b t ... -> b ... t return x class RMSNorm(nn.Module): def __init__(self, dim: int, eps: float = 1e-5, elementwise_affine=True, weight_shape=None): super().__init__() self.dim = dim self.eps = eps self.elementwise_affine = elementwise_affine if self.elementwise_affine: weight_shape = (dim,) if weight_shape is None else weight_shape self.weight = nn.Parameter(torch.ones(weight_shape)) else: self.register_parameter('weight', None) def _norm(self, x): return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) def forward(self, x): output = self._norm(x.float()).type_as(x) if self.weight is not None: output = output * self.weight return output def extra_repr(self) -> str: return f'dim={self.dim}, eps={self.eps}, elementwise_affine={self.elementwise_affine}' class ConvRMSNorm(RMSNorm): def __init__(self, dim: int, eps: float = 1e-5, elementwise_affine=True, weight_shape=None): super().__init__(dim, eps, elementwise_affine, weight_shape) def forward(self, x): x = x.transpose(1, 2) # b ... t -> b t ... if (not APEX_AVAILABLE) or (not self.elementwise_affine): # Fallback to native implementation output = self._norm(x.float()).type_as(x) if self.weight is not None: output = output * self.weight else: output = fused_rms_norm_affine(x, self.weight, self.weight.shape, self.eps) output = output.transpose(1, 2) # b t ... -> b ... t return output # Convolutional layers and utilities CONV_NORMALIZATIONS = frozenset(['none', 'weight_norm', 'spectral_norm', 'time_layer_norm', 'layer_norm', 'time_group_norm']) def apply_parametrization_norm(module: nn.Module, norm: str = 'none') -> nn.Module: assert norm in CONV_NORMALIZATIONS if norm == 'weight_norm': return nn.utils.weight_norm(module) elif norm == 'spectral_norm': return nn.utils.spectral_norm(module) else: # We already check was in CONV_NORMALIZATION, so any other choice # doesn't need reparametrization. return module def get_norm_module(module: nn.Module, causal: bool = False, norm: str = 'none', **norm_kwargs) -> nn.Module: """Return the proper normalization module. If causal is True, this will ensure the returned module is causal, or return an error if the normalization doesn't support causal evaluation. """ assert norm in CONV_NORMALIZATIONS if norm == 'layer_norm': assert isinstance(module, nn.modules.conv._ConvNd) return ConvLayerNorm(module.out_channels, **norm_kwargs) elif norm == 'time_group_norm': if causal: raise ValueError("GroupNorm doesn't support causal evaluation.") assert isinstance(module, nn.modules.conv._ConvNd) return nn.GroupNorm(1, module.out_channels, **norm_kwargs) else: return nn.Identity() def get_extra_padding_for_conv1d(x: torch.Tensor, kernel_size: int, stride: int, padding_total: int = 0) -> int: """Calculate extra padding needed for convolution to have the same output length""" length = x.shape[-1] n_frames = (length - kernel_size + padding_total) / stride + 1 ideal_length = (math.ceil(n_frames) - 1) * stride + (kernel_size - padding_total) return ideal_length - length def pad1d(x: torch.Tensor, paddings: tp.Tuple[int, int], mode: str = 'zero', value: float = 0.): """Pad 1D input with handling for small inputs in reflect mode""" length = x.shape[-1] padding_left, padding_right = paddings assert padding_left >= 0 and padding_right >= 0, (padding_left, padding_right) if mode == 'reflect': max_pad = max(padding_left, padding_right) extra_pad = 0 if length <= max_pad: extra_pad = max_pad - length + 1 x = F.pad(x, (0, extra_pad)) padded = F.pad(x, paddings, mode, value) end = padded.shape[-1] - extra_pad return padded[..., :end] else: return F.pad(x, paddings, mode, value) def unpad1d(x: torch.Tensor, paddings: tp.Tuple[int, int]): """Remove padding from x, handling properly zero padding. Only for 1d!""" padding_left, padding_right = paddings assert padding_left >= 0 and padding_right >= 0, (padding_left, padding_right) assert (padding_left + padding_right) <= x.shape[-1] end = x.shape[-1] - padding_right return x[..., padding_left: end] class NormConv1d(nn.Module): """Wrapper around Conv1d and normalization applied to this conv""" def __init__(self, *args, causal: bool = False, norm: str = 'none', norm_kwargs: tp.Dict[str, tp.Any] = {}, **kwargs): super().__init__() self.conv = apply_parametrization_norm(nn.Conv1d(*args, **kwargs), norm) self.norm = get_norm_module(self.conv, causal, norm, **norm_kwargs) self.norm_type = norm def forward(self, x): x = self.conv(x) x = self.norm(x) return x class NormConvTranspose1d(nn.Module): """Wrapper around ConvTranspose1d and normalization applied to this conv""" def __init__(self, *args, causal: bool = False, norm: str = 'none', norm_kwargs: tp.Dict[str, tp.Any] = {}, **kwargs): super().__init__() self.convtr = apply_parametrization_norm(nn.ConvTranspose1d(*args, **kwargs), norm) self.norm = get_norm_module(self.convtr, causal, norm, **norm_kwargs) self.norm_type = norm def forward(self, x): x = self.convtr(x) x = self.norm(x) return x class VibeVoiceTokenizerStreamingCache: """Cache for streaming convolution, similar to KV cache in attention""" def __init__(self): self.cache = {} # Dict mapping (layer_id, sample_idx) to state tensor def get(self, layer_id: str, sample_indices: torch.Tensor) -> Optional[torch.Tensor]: """Get cached states for given layer and sample indices""" states = [] max_length = 0 # First pass: collect states and find max length for idx in sample_indices.tolist(): key = (layer_id, idx) if key not in self.cache: return None # If any sample is missing, return None state = self.cache[key] states.append(state) max_length = max(max_length, state.shape[-1]) # Second pass: pad states to max length if needed if len(states) > 0 and states[0].dim() >= 2: padded_states = [] for state in states: if state.shape[-1] < max_length: # Pad on the time dimension (last dimension) pad_size = max_length - state.shape[-1] # Pad with zeros on the LEFT to align the most recent samples padded_state = F.pad(state, (pad_size, 0), mode='constant', value=0) padded_states.append(padded_state) else: padded_states.append(state) return torch.stack(padded_states, dim=0) else: return torch.stack(states, dim=0) def set(self, layer_id: str, sample_indices: torch.Tensor, states: torch.Tensor): """Set cached states for given layer and sample indices""" for i, idx in enumerate(sample_indices.tolist()): key = (layer_id, idx) self.cache[key] = states[i].detach() def set_to_zero(self, sample_indices: torch.Tensor): """Set all cached states to zero for given sample indices""" for key in list(self.cache.keys()): layer_id, sample_idx = key if sample_idx in sample_indices.tolist(): # Create zero tensor with same shape and dtype as cached tensor cached_tensor = self.cache[key] self.cache[key] = torch.zeros_like(cached_tensor) def clear(self, layer_id: Optional[str] = None, sample_indices: Optional[torch.Tensor] = None): """Clear cache for specific layer/samples or everything""" if layer_id is None and sample_indices is None: self.cache.clear() elif layer_id is not None and sample_indices is None: # Clear all samples for a specific layer keys_to_remove = [k for k in self.cache.keys() if k[0] == layer_id] for k in keys_to_remove: del self.cache[k] elif layer_id is not None and sample_indices is not None: # Clear specific samples for a specific layer for idx in sample_indices.tolist(): key = (layer_id, idx) self.cache.pop(key, None) class SConv1d(nn.Module): """Conv1d with built-in handling of asymmetric or causal padding and normalization.""" def __init__(self, in_channels: int, out_channels: int, kernel_size: int, stride: int = 1, dilation: int = 1, groups: int = 1, bias: bool = True, causal: bool = False, norm: str = 'none', norm_kwargs: tp.Dict[str, tp.Any] = {}, pad_mode: str = 'reflect'): super().__init__() self.conv = NormConv1d(in_channels, out_channels, kernel_size, stride, dilation=dilation, groups=groups, bias=bias, causal=causal, norm=norm, norm_kwargs=norm_kwargs) self.causal = causal self.pad_mode = pad_mode # Store configuration self.kernel_size = kernel_size self.dilation = dilation self.stride = stride self.in_channels = in_channels self.out_channels = out_channels # For causal convolution, we need to maintain kernel_size - 1 samples as context # need to check use which context_size is more suitable # self.context_size = (kernel_size - 1) * dilation self.context_size = (kernel_size - 1) * dilation - (stride - 1) # For non-streaming mode, calculate padding self.padding_total = (kernel_size - 1) * dilation - (stride - 1) # Create a unique layer ID for cache management self._layer_id = None @property def layer_id(self): if self._layer_id is None: self._layer_id = f"sconv1d_{id(self)}" return self._layer_id def forward(self, x: torch.Tensor, cache: Optional[VibeVoiceTokenizerStreamingCache] = None, sample_indices: Optional[torch.Tensor] = None, use_cache: bool = False, debug: bool = False) -> torch.Tensor: """ Forward pass with optional streaming support via cache. Args: x: Input tensor [batch_size, channels, time] cache: VibeVoiceTokenizerStreamingCache object for maintaining states sample_indices: Indices identifying each sample for cache management use_cache: Whether to use cached states for streaming debug: Whether to print debug information Returns: Output tensor """ B, C, T = x.shape # Non-streaming mode if not use_cache or cache is None: return self._forward_non_streaming(x, debug=debug) # Streaming mode assert self.causal, "Streaming mode is only supported for causal convolutions" assert sample_indices is not None, "sample_indices must be provided for streaming mode" assert len(sample_indices) == B, "sample_indices must match batch size" return self._forward_streaming(x, cache, sample_indices, debug) def _forward_streaming(self, x: torch.Tensor, cache: VibeVoiceTokenizerStreamingCache, sample_indices: torch.Tensor, debug: bool = False) -> torch.Tensor: """Streaming forward pass with cache operations kept separate from compiled code""" B, C, T = x.shape # Cache operations (not compiled) cached_states = cache.get(self.layer_id, sample_indices) if cached_states is None: # First chunk - initialize with zeros for context if self.context_size > 0: cached_states = torch.zeros(B, C, self.context_size, device=x.device, dtype=x.dtype) if debug: print(f"[DEBUG] Initialized cache with shape: {cached_states.shape}, context_size={self.context_size}") else: cached_states = torch.zeros(B, C, 0, device=x.device, dtype=x.dtype) if debug: print(f"[DEBUG] No context needed (kernel_size=stride)") # Concatenate cached states with input if cached_states.shape[2] > 0: input_with_context = torch.cat([cached_states, x], dim=2) else: input_with_context = x if debug: print(f"[DEBUG] Input shape: {x.shape}, Cache shape: {cached_states.shape}, Combined: {input_with_context.shape}") # Apply convolution directly - no extra padding in streaming mode # The conv layer will handle its own padding internally output = self.conv(input_with_context) if debug: print(f"[DEBUG] Output shape: {output.shape}") # Update cache for next chunk if self.context_size > 0: # Calculate how many samples to keep total_input_length = input_with_context.shape[2] # Keep the last context_size samples if total_input_length >= self.context_size: new_cache_start = total_input_length - self.context_size new_cache = input_with_context[:, :, new_cache_start:] else: # If we have less than context_size samples, keep everything new_cache = input_with_context if debug: print(f"[DEBUG] New cache shape: {new_cache.shape}") cache.set(self.layer_id, sample_indices, new_cache) return output def _forward_non_streaming(self, x: torch.Tensor, debug: bool = False) -> torch.Tensor: """Standard forward pass without streaming""" B, C, T = x.shape kernel_size = self.kernel_size stride = self.stride dilation = self.dilation padding_total = self.padding_total # Compute extra padding for stride alignment extra_padding = get_extra_padding_for_conv1d(x, kernel_size, stride, padding_total) if debug: print(f"[DEBUG NON-STREAMING] Input shape: {x.shape}, padding_total={padding_total}, extra_padding={extra_padding}") if self.causal: # Left padding for causal if self.pad_mode == 'constant': x = pad1d(x, (padding_total, extra_padding), mode=self.pad_mode, value=0) else: x = pad1d(x, (padding_total, extra_padding), mode=self.pad_mode) else: # Symmetric padding for non-causal padding_right = padding_total // 2 padding_left = padding_total - padding_right x = pad1d(x, (padding_left, padding_right + extra_padding), mode=self.pad_mode) if debug: print(f"[DEBUG NON-STREAMING] After padding: {x.shape}") output = self.conv(x) if debug: print(f"[DEBUG NON-STREAMING] Output shape: {output.shape}") return output class SConvTranspose1d(nn.Module): """ConvTranspose1d with built-in handling of asymmetric or causal padding and normalization.""" def __init__(self, in_channels: int, out_channels: int, kernel_size: int, stride: int = 1, causal: bool = False, norm: str = 'none', trim_right_ratio: float = 1., norm_kwargs: tp.Dict[str, tp.Any] = {}, bias: bool = True): super().__init__() self.convtr = NormConvTranspose1d(in_channels, out_channels, kernel_size, stride, causal=causal, norm=norm, norm_kwargs=norm_kwargs, bias=bias) self.causal = causal self.trim_right_ratio = trim_right_ratio assert self.causal or self.trim_right_ratio == 1., \ "`trim_right_ratio` != 1.0 only makes sense for causal convolutions" assert self.trim_right_ratio >= 0. and self.trim_right_ratio <= 1. # Store configuration self.kernel_size = kernel_size self.stride = stride self.in_channels = in_channels self.out_channels = out_channels # For transposed convolution, padding calculation is different self.padding_total = kernel_size - stride # For streaming, we need to keep track of input history # Transposed conv needs to see multiple input samples to produce correct output self.context_size = kernel_size - 1 # Create a unique layer ID for cache management self._layer_id = None @property def layer_id(self): if self._layer_id is None: self._layer_id = f"sconvtr1d_{id(self)}" return self._layer_id def forward(self, x: torch.Tensor, cache: Optional[VibeVoiceTokenizerStreamingCache] = None, sample_indices: Optional[torch.Tensor] = None, use_cache: bool = False, debug: bool = False) -> torch.Tensor: """ Forward pass with optional streaming support via cache. """ B, C, T = x.shape # Non-streaming mode if not use_cache or cache is None: return self._forward_non_streaming(x, debug=debug) # Streaming mode assert sample_indices is not None, "sample_indices must be provided for streaming mode" assert len(sample_indices) == B, "sample_indices must match batch size" return self._forward_streaming(x, cache, sample_indices, debug) def _forward_streaming(self, x: torch.Tensor, cache: VibeVoiceTokenizerStreamingCache, sample_indices: torch.Tensor, debug: bool = False) -> torch.Tensor: """Streaming forward pass with cache operations kept separate from compiled code""" B, C, T = x.shape # Cache operations (not compiled) cached_input = cache.get(self.layer_id, sample_indices) if cached_input is None: # First chunk - no history yet cached_input = torch.zeros(B, C, 0, device=x.device, dtype=x.dtype) if debug: print(f"[DEBUG] Initialized empty cache for transposed conv") # Concatenate cached input with new input full_input = torch.cat([cached_input, x], dim=2) if debug: print(f"[DEBUG] Input shape: {x.shape}, Cache shape: {cached_input.shape}, Combined: {full_input.shape}") # First chunk or debug mode - use uncompiled version full_output = self.convtr(full_input) if debug: print(f"[DEBUG] Full transposed conv output shape: {full_output.shape}") # Calculate padding to remove if self.causal: padding_right = math.ceil(self.padding_total * self.trim_right_ratio) padding_left = self.padding_total - padding_right else: padding_right = self.padding_total // 2 padding_left = self.padding_total - padding_right # Remove padding if padding_left + padding_right > 0: full_output = unpad1d(full_output, (padding_left, padding_right)) if debug: print(f"[DEBUG] After unpadding: {full_output.shape}") # Determine which part of the output corresponds to the new input if cached_input.shape[2] == 0: # First chunk - return all output output = full_output else: # Subsequent chunks - return only the new output expected_new_output = T * self.stride # Take the last expected_new_output samples if full_output.shape[2] >= expected_new_output: output = full_output[:, :, -expected_new_output:] else: output = full_output if debug: print(f"[DEBUG] Final streaming output shape: {output.shape}") # Update cache if full_input.shape[2] > self.context_size: new_cache = full_input[:, :, -self.context_size:] else: new_cache = full_input if debug: print(f"[DEBUG] New cache shape: {new_cache.shape}") cache.set(self.layer_id, sample_indices, new_cache) return output def _forward_non_streaming(self, x: torch.Tensor, debug: bool = False) -> torch.Tensor: """Standard forward pass without streaming""" if debug: print(f"[DEBUG NON-STREAMING] Input shape: {x.shape}") # Apply transposed convolution y = self.convtr(x) if debug: print(f"[DEBUG NON-STREAMING] After transposed conv: {y.shape}") # Calculate and remove padding if self.causal: padding_right = math.ceil(self.padding_total * self.trim_right_ratio) padding_left = self.padding_total - padding_right else: padding_right = self.padding_total // 2 padding_left = self.padding_total - padding_right if padding_left + padding_right > 0: y = unpad1d(y, (padding_left, padding_right)) if debug: print(f"[DEBUG NON-STREAMING] Final output shape: {y.shape}") return y # FFN class FFN(nn.Module): def __init__( self, embed_dim, ffn_dim, bias=False, ): super().__init__() self.embed_dim = embed_dim self.linear1 = nn.Linear(self.embed_dim, ffn_dim, bias=bias) self.gelu = ACT2FN["gelu"] self.linear2 = nn.Linear(ffn_dim, self.embed_dim, bias=bias) def forward(self, x): x = self.linear1(x) x = self.gelu(x) x = self.linear2(x) return x class Convlayer(nn.Module): def __init__( self, in_channels, out_channels, kernel_size, stride=1, dilation=1, groups=1, bias=True, pad_mode='zeros', norm='weight_norm', causal=True, ): super().__init__() self.conv = SConv1d(in_channels, out_channels, kernel_size, stride=stride, dilation=dilation, groups=groups, bias=bias, pad_mode=pad_mode, norm=norm, causal=causal) def forward(self, x): return self.conv(x) class Block1D(nn.Module): def __init__(self, dim, kernel_size=7, drop_path=0., mixer_layer='conv', layer_scale_init_value=1e-6, **kwargs): super().__init__() if kwargs.get('layernorm', 'LN') == 'LN': self.norm = ConvLayerNorm(dim, eps=kwargs.get('eps', 1e-6)) self.ffn_norm = ConvLayerNorm(dim, eps=kwargs.get('eps', 1e-6)) elif kwargs.get('layernorm', 'RMSNorm') == 'RMSNorm': self.norm = ConvRMSNorm(dim, eps=kwargs.get('eps', 1e-6)) self.ffn_norm = ConvRMSNorm(dim, eps=kwargs.get('eps', 1e-6)) if mixer_layer == 'conv': self.mixer = Convlayer(dim, dim, groups=kwargs.get('groups', 1), kernel_size=kernel_size, pad_mode=kwargs.get('pad_mode', 'reflect'), norm=kwargs.get('norm', 'none'), causal=kwargs.get('causal', True), bias=kwargs.get('bias', True), ) elif mixer_layer == 'depthwise_conv': self.mixer = Convlayer(dim, dim, groups=dim, kernel_size=kernel_size, pad_mode=kwargs.get('pad_mode', 'reflect'), norm=kwargs.get('norm', 'none'), causal=kwargs.get('causal', True), bias=kwargs.get('bias', True), ) else: raise ValueError(f"Unsupported mixer layer: {mixer_layer}") self.ffn = FFN( dim, kwargs.get('ffn_expansion', 4) * dim, bias=kwargs.get('bias', False), ) self.drop_path = nn.Identity() if drop_path <= 0. else nn.modules.DropPath(drop_path) if layer_scale_init_value > 0: self.gamma = nn.Parameter(layer_scale_init_value * torch.ones((dim)), requires_grad=True) self.ffn_gamma = nn.Parameter(layer_scale_init_value * torch.ones((dim)), requires_grad=True) else: self.gamma = None self.ffn_gamma = None def forward(self, x): # mixer residual = x x = self.norm(x) x = self.mixer(x) if self.gamma is not None: x = x * self.gamma.unsqueeze(-1) x = residual + self.drop_path(x) # ffn residual = x x = self.ffn_norm(x) x = x.permute(0, 2, 1) x = self.ffn(x) x = x.permute(0, 2, 1) if self.ffn_gamma is not None: x = x * self.ffn_gamma.unsqueeze(-1) x = residual + self.drop_path(x) return x class TokenizerEncoder(nn.Module): """ Encoder component for the VibeVoice tokenizer that converts audio to latent representations. Args: config: Configuration object with model parameters """ def __init__(self, config): super().__init__() # Extract parameters from config self.channels = config.channels self.dimension = config.dimension self.n_filters = config.n_filters self.ratios = list(reversed(config.ratios)) self.depths = config.depths self.n_residual_layers = getattr(config, "n_residual_layers", 1) self.hop_length = np.prod(self.ratios) self.causal = config.causal # Additional config parameters with defaults kernel_size = getattr(config, "kernel_size", 7) last_kernel_size = getattr(config, "last_kernel_size", 7) norm = getattr(config, "norm", "none") norm_params = getattr(config, "norm_params", {}) pad_mode = getattr(config, "pad_mode", "reflect") bias = getattr(config, "bias", True) layernorm = getattr(config, "layernorm", "LN") layernorm_eps = getattr(config, "layernorm_eps", 1e-6) layernorm_elementwise_affine = getattr(config, "layernorm_elementwise_affine", True) drop_path_rate = getattr(config, "drop_path_rate", 0.0) mixer_layer = getattr(config, "mixer_layer", "conv") layer_scale_init_value = getattr(config, "layer_scale_init_value", 0) disable_last_norm = getattr(config, "disable_last_norm", False) # determine the norm type based on layernorm if layernorm == 'LN': norm_type = ConvLayerNorm elif layernorm == 'RMSNorm': norm_type = partial(ConvRMSNorm, elementwise_affine=layernorm_elementwise_affine) else: raise ValueError(f"Unsupported norm type: {layernorm}") # stem and intermediate downsampling conv layers stem = nn.Sequential( SConv1d(self.channels, self.n_filters, kernel_size, norm=norm, norm_kwargs=norm_params, causal=self.causal, pad_mode=pad_mode, bias=bias), ) self.downsample_layers = nn.ModuleList() self.downsample_layers.append(stem) for i in range(len(self.ratios)): in_ch = self.n_filters * (2 ** i) out_ch = self.n_filters * (2 ** (i + 1)) downsample_layer = nn.Sequential( SConv1d(in_ch, out_ch, kernel_size=self.ratios[i] * 2, stride=self.ratios[i], causal=self.causal, pad_mode=pad_mode, norm=norm, bias=bias) ) self.downsample_layers.append(downsample_layer) # configure the transformer blocks layer_type = partial( Block1D, mixer_layer=mixer_layer, layernorm=layernorm, eps=layernorm_eps, causal=self.causal, pad_mode=pad_mode, norm=norm, bias=bias, layer_scale_init_value=layer_scale_init_value, ) self.stages = nn.ModuleList() dp_rates = [x.item() for x in torch.linspace(0, drop_path_rate, sum(self.depths))] cur = 0 for i in range(len(self.depths)): in_ch = self.n_filters * (2 ** i) stage = nn.Sequential( *[layer_type(dim=in_ch, drop_path=dp_rates[cur + j]) for j in range(self.depths[i])] ) self.stages.append(stage) cur += self.depths[i] if not disable_last_norm: self.norm = norm_type(in_ch, eps=layernorm_eps) else: self.norm = nn.Identity() self.head = SConv1d(in_ch, self.dimension, kernel_size=last_kernel_size, causal=self.causal, pad_mode=pad_mode, norm=norm, bias=bias) def forward_features(self, x, cache=None, sample_indices=None, use_cache=False, debug=False): for i in range(len(self.depths)): # Apply downsampling for layer in self.downsample_layers[i]: if isinstance(layer, SConv1d): x = layer(x, cache=cache, sample_indices=sample_indices, use_cache=use_cache, debug=debug) else: x = layer(x) # Apply stage (Block1D contains Convlayer which contains SConv1d) for block in self.stages[i]: if hasattr(block, 'mixer') and hasattr(block.mixer, 'conv') and isinstance(block.mixer.conv, SConv1d): # Block1D forward with cache support residual = x x = block.norm(x) x = block.mixer.conv(x, cache=cache, sample_indices=sample_indices, use_cache=use_cache, debug=debug) if block.gamma is not None: x = x * block.gamma.unsqueeze(-1) x = residual + x # FFN part residual = x x = block.ffn_norm(x) x = x.permute(0, 2, 1) x = block.ffn(x) x = x.permute(0, 2, 1) if block.ffn_gamma is not None: x = x * block.ffn_gamma.unsqueeze(-1) x = residual + x else: x = block(x) return self.norm(x) def forward(self, x, cache=None, sample_indices=None, use_cache=False, debug=False): x = self.forward_features(x, cache=cache, sample_indices=sample_indices, use_cache=use_cache, debug=debug) x = self.head(x, cache=cache, sample_indices=sample_indices, use_cache=use_cache, debug=debug) return x class TokenizerDecoder(nn.Module): """ Decoder component for the VibeVoice tokenizer that converts latent representations back to audio. Args: config: Configuration object with model parameters """ def __init__(self, config): super().__init__() # Extract parameters from config self.dimension = config.dimension self.channels = config.channels self.n_filters = config.n_filters self.ratios = config.ratios # IMPORTANT CHANGE: Don't reverse depths again since they're already reversed in VibeVoiceAcousticTokenizerModel self.depths = config.depths # Changed from list(reversed(config.depths)) self.n_residual_layers = getattr(config, "n_residual_layers", 1) self.hop_length = np.prod(self.ratios) self.causal = config.causal # Additional config parameters with defaults kernel_size = getattr(config, "kernel_size", 7) last_kernel_size = getattr(config, "last_kernel_size", 7) norm = getattr(config, "norm", "none") norm_params = getattr(config, "norm_params", {}) pad_mode = getattr(config, "pad_mode", "reflect") bias = getattr(config, "bias", True) layernorm = getattr(config, "layernorm", "LN") layernorm_eps = getattr(config, "layernorm_eps", 1e-6) trim_right_ratio = getattr(config, "trim_right_ratio", 1.0) layernorm_elementwise_affine = getattr(config, "layernorm_elementwise_affine", True) drop_path_rate = getattr(config, "drop_path_rate", 0.0) mixer_layer = getattr(config, "mixer_layer", "conv") layer_scale_init_value = getattr(config, "layer_scale_init_value", 0) disable_last_norm = getattr(config, "disable_last_norm", False) # determine the norm type based on layernorm if layernorm == 'LN': norm_type = ConvLayerNorm elif layernorm == 'RMSNorm': norm_type = partial(ConvRMSNorm, elementwise_affine=layernorm_elementwise_affine) else: raise ValueError(f"Unsupported norm type: {layernorm}") # stem and upsampling layers stem = nn.Sequential( SConv1d(self.dimension, self.n_filters * 2 ** (len(self.depths) - 1), kernel_size, norm=norm, norm_kwargs=norm_params, causal=self.causal, pad_mode=pad_mode, bias=bias), ) self.upsample_layers = nn.ModuleList() self.upsample_layers.append(stem) for i in range(len(self.ratios)): in_ch = self.n_filters * (2 ** (len(self.depths) - 1 - i)) out_ch = self.n_filters * (2 ** (len(self.depths) - 1 - i - 1)) upsample_layer = nn.Sequential( SConvTranspose1d(in_ch, out_ch, kernel_size=self.ratios[i] * 2, stride=self.ratios[i], norm=norm, norm_kwargs=norm_params, bias=bias, causal=self.causal, trim_right_ratio=trim_right_ratio), ) self.upsample_layers.append(upsample_layer) # configure transformer blocks layer_type = partial( Block1D, mixer_layer=mixer_layer, layernorm=layernorm, eps=layernorm_eps, causal=self.causal, pad_mode=pad_mode, norm=norm, bias=bias, layer_scale_init_value=layer_scale_init_value, ) self.stages = nn.ModuleList() dp_rates = [x.item() for x in torch.linspace(0, drop_path_rate, sum(self.depths))] cur = 0 # Create stages in the same order as the original model for i in range(len(self.depths)): in_ch = self.n_filters * (2 ** (len(self.depths) - 1 - i)) stage = nn.Sequential( *[layer_type(dim=in_ch, drop_path=dp_rates[cur + j]) for j in range(self.depths[i])] ) self.stages.append(stage) cur += self.depths[i] if not disable_last_norm: self.norm = norm_type(in_ch, eps=layernorm_eps) else: self.norm = nn.Identity() self.head = SConv1d(in_ch, self.channels, kernel_size=last_kernel_size, causal=self.causal, pad_mode=pad_mode, norm=norm, bias=bias) def forward_features(self, x, cache=None, sample_indices=None, use_cache=False, debug=False): for i in range(len(self.depths)): # Apply upsampling for layer in self.upsample_layers[i]: if isinstance(layer, (SConv1d, SConvTranspose1d)): x = layer(x, cache=cache, sample_indices=sample_indices, use_cache=use_cache, debug=debug) else: x = layer(x) # Apply stage (Block1D contains Convlayer which contains SConv1d) for block in self.stages[i]: if hasattr(block, 'mixer') and hasattr(block.mixer, 'conv') and isinstance(block.mixer.conv, SConv1d): # Block1D forward with cache support residual = x x = block.norm(x) x = block.mixer.conv(x, cache=cache, sample_indices=sample_indices, use_cache=use_cache, debug=debug) if block.gamma is not None: x = x * block.gamma.unsqueeze(-1) x = residual + x # FFN part residual = x x = block.ffn_norm(x) x = x.permute(0, 2, 1) x = block.ffn(x) x = x.permute(0, 2, 1) if block.ffn_gamma is not None: x = x * block.ffn_gamma.unsqueeze(-1) x = residual + x else: x = block(x) return self.norm(x) def forward(self, x, cache=None, sample_indices=None, use_cache=False, debug=False): x = self.forward_features(x, cache=cache, sample_indices=sample_indices, use_cache=use_cache, debug=debug) x = self.head(x, cache=cache, sample_indices=sample_indices, use_cache=use_cache, debug=debug) return x @dataclass class VibeVoiceTokenizerEncoderOutput: """ Output of VibeVoice tokenizer encoder, representing a Gaussian distribution with fixed variance. Args: mean (`torch.FloatTensor`): The mean parameters of the distribution. std (`float` or `torch.FloatTensor`): Fixed standard deviation value. """ mean: torch.Tensor std: Optional[Union[float, torch.Tensor]] = None def sample(self, dist_type='fix'): """ Sample from the distribution. Args: dist_type (`str`): Sampling method, either 'fix' or 'gaussian'. Returns: `torch.FloatTensor`: Sampled values. `torch.FloatTensor` (optional): Standard deviation used (only when dist_type='gaussian'). """ if dist_type == 'fix': x = self.mean + self.std * torch.randn_like(self.mean) return x, self.std elif dist_type == 'gaussian': batch_size = self.mean.size(0) value = self.std / 0.8 std = torch.randn(batch_size, device=self.mean.device, dtype=self.mean.dtype) * value while std.dim() < self.mean.dim(): std = std.unsqueeze(-1) x = self.mean + std * torch.randn_like(self.mean) return x, std else: return self.mean, self.std def kl(self): """Compute KL divergence between this distribution and a standard normal.""" target = torch.zeros_like(self.mean) return F.mse_loss(self.mean, target, reduction='none') def mode(self): """Return the distribution mode (which is the mean for Gaussian).""" return self.mean class VibeVoiceAcousticTokenizerModel(PreTrainedModel): """VibeVoice speech tokenizer model combining encoder and decoder for acoustic tokens""" config_class = VibeVoiceAcousticTokenizerConfig base_model_prefix = "vibevoice_acoustic_tokenizer" _supports_flash_attn_2 = True _supports_sdpa = True _no_split_modules = ["TokenizerEncoder", "TokenizerDecoder"] def __init__(self, config): super().__init__(config) self.register_buffer('fix_std', torch.tensor(config.fix_std), persistent=False) self.std_dist_type = getattr(config, "std_dist_type", "fix") # Parse encoder depths if isinstance(config.encoder_depths, str): encoder_depths = [int(d) for d in config.encoder_depths.split('-')] else: encoder_depths = config.encoder_depths # Parse decoder depths if provided if config.decoder_depths is not None and isinstance(config.decoder_depths, str): decoder_depths = [int(d) for d in config.decoder_depths.split('-')] else: # Default: use reversed encoder depths if decoder_depths is None decoder_depths = list(reversed(encoder_depths)) # Create encoder config encoder_config = copy.deepcopy(config) encoder_config.dimension = config.vae_dim encoder_config.n_filters = config.encoder_n_filters encoder_config.ratios = config.encoder_ratios encoder_config.depths = encoder_depths encoder_config.norm = config.conv_norm encoder_config.pad_mode = config.pad_mode encoder_config.bias = config.conv_bias encoder_config.layernorm_eps = config.layernorm_eps encoder_config.layernorm_elementwise_affine = config.layernorm_elementwise_affine encoder_config.mixer_layer = config.mixer_layer encoder_config.layer_scale_init_value = config.layer_scale_init_value encoder_config.disable_last_norm = config.disable_last_norm # Create decoder config decoder_config = copy.deepcopy(config) decoder_config.dimension = config.vae_dim decoder_config.n_filters = config.decoder_n_filters decoder_config.ratios = config.decoder_ratios decoder_config.depths = decoder_depths decoder_config.norm = config.conv_norm decoder_config.pad_mode = config.pad_mode decoder_config.bias = config.conv_bias decoder_config.layernorm_eps = config.layernorm_eps decoder_config.layernorm_elementwise_affine = config.layernorm_elementwise_affine decoder_config.mixer_layer = config.mixer_layer decoder_config.layer_scale_init_value = config.layer_scale_init_value decoder_config.disable_last_norm = config.disable_last_norm # Initialize encoder and decoder self.encoder = TokenizerEncoder(encoder_config) self.decoder = TokenizerDecoder(decoder_config) # Initialize weights self.apply(self._init_weights) def _init_weights(self, module): """Initialize weights for the model""" if isinstance(module, nn.Linear): nn.init.normal_(module.weight, std=self.config.weight_init_value) if module.bias is not None: nn.init.zeros_(module.bias) elif isinstance(module, nn.LayerNorm): nn.init.ones_(module.weight) nn.init.zeros_(module.bias) elif isinstance(module, nn.Conv1d): nn.init.normal_(module.weight, std=self.config.weight_init_value) if module.bias is not None: nn.init.zeros_(module.bias) @torch.no_grad() def encode(self, audio, cache=None, sample_indices=None, use_cache=False, debug=False): """Convert audio to latent representations""" latents = self.encoder(audio, cache=cache, sample_indices=sample_indices, use_cache=use_cache, debug=debug) return VibeVoiceTokenizerEncoderOutput(mean=latents.permute(0, 2, 1), std=self.fix_std) @torch.no_grad() def sampling(self, encoder_output, dist_type=None): """Sample from the encoder output distribution""" dist_type = dist_type or self.std_dist_type if dist_type == 'fix': return encoder_output.sample(dist_type='fix') elif dist_type == 'gaussian': return encoder_output.sample(dist_type='gaussian') else: raise ValueError(f"Unsupported dist_type: {dist_type}, expected 'fix' or 'gaussian'") @torch.no_grad() def decode(self, latents, cache=None, sample_indices=None, use_cache=False, debug=False): """Convert latent representations back to audio""" if latents.shape[1] == self.config.vae_dim: pass else: latents = latents.permute(0, 2, 1) audio = self.decoder(latents, cache=cache, sample_indices=sample_indices, use_cache=use_cache, debug=debug) return audio def forward(self, audio, cache=None, sample_indices=None, use_cache=False, debug=False): """Full forward pass: encode audio to latents, then decode back to audio""" encoder_output = self.encode(audio, cache=cache, sample_indices=sample_indices, use_cache=use_cache, debug=debug) sampled_latents, _ = self.sampling(encoder_output) reconstructed = self.decode(sampled_latents, cache=cache, sample_indices=sample_indices, use_cache=use_cache, debug=debug) return reconstructed, sampled_latents class VibeVoiceSemanticTokenizerModel(PreTrainedModel): """VibeVoice speech tokenizer model with only encoder for semantic tokens""" config_class = VibeVoiceSemanticTokenizerConfig base_model_prefix = "vibevoice_semantic_tokenizer" _supports_flash_attn_2 = True _supports_sdpa = True _no_split_modules = ["TokenizerEncoder"] def __init__(self, config): super().__init__(config) # Parse encoder depths if isinstance(config.encoder_depths, str): encoder_depths = [int(d) for d in config.encoder_depths.split('-')] else: encoder_depths = config.encoder_depths # Create encoder config encoder_config = copy.deepcopy(config) encoder_config.dimension = config.vae_dim encoder_config.n_filters = config.encoder_n_filters encoder_config.ratios = config.encoder_ratios encoder_config.depths = encoder_depths encoder_config.norm = config.conv_norm encoder_config.pad_mode = config.pad_mode encoder_config.bias = config.conv_bias encoder_config.layernorm_eps = config.layernorm_eps encoder_config.layernorm_elementwise_affine = config.layernorm_elementwise_affine encoder_config.mixer_layer = config.mixer_layer encoder_config.layer_scale_init_value = config.layer_scale_init_value encoder_config.disable_last_norm = config.disable_last_norm # Initialize encoder and decoder self.encoder = TokenizerEncoder(encoder_config) # Initialize weights self.apply(self._init_weights) def _init_weights(self, module): """Initialize weights for the model""" if isinstance(module, nn.Linear): nn.init.normal_(module.weight, std=self.config.weight_init_value) if module.bias is not None: nn.init.zeros_(module.bias) elif isinstance(module, nn.LayerNorm): nn.init.ones_(module.weight) nn.init.zeros_(module.bias) elif isinstance(module, nn.Conv1d): nn.init.normal_(module.weight, std=self.config.weight_init_value) if module.bias is not None: nn.init.zeros_(module.bias) @torch.no_grad() def encode(self, audio, cache=None, sample_indices=None, use_cache=False, debug=False): """Convert audio to latent representations""" latents = self.encoder(audio, cache=cache, sample_indices=sample_indices, use_cache=use_cache, debug=debug) return VibeVoiceTokenizerEncoderOutput(mean=latents.permute(0, 2, 1)) @torch.no_grad() def sampling(self, encoder_output, dist_type=None): """Sample from the encoder output distribution""" return encoder_output.sample(dist_type='none') def forward(self, audio, cache=None, sample_indices=None, use_cache=False, debug=False): """Full forward pass: encode audio to latents, then decode back to audio""" encoder_output = self.encode(audio, cache=cache, sample_indices=sample_indices, use_cache=use_cache, debug=debug) sampled_latents, _ = self.sampling(encoder_output, dist_type='none') return None, sampled_latents AutoModel.register(VibeVoiceAcousticTokenizerConfig, VibeVoiceAcousticTokenizerModel) AutoModel.register(VibeVoiceSemanticTokenizerConfig, VibeVoiceSemanticTokenizerModel) __all__ = [ "VibeVoiceTokenizerStreamingCache", "VibeVoiceAcousticTokenizerModel", "VibeVoiceSemanticTokenizerModel", ]