"""VibeVoice-Embed: the VibeVoice acoustic encoder as a standalone voice-embedding model. Everything from the module docstring down to the VibeVoiceEmbedModel class is upstream VibeVoice code (MIT, https://github.com/vibevoice-community/VibeVoice), mechanically extracted from ``vibevoice/modular/modular_vibevoice_tokenizer.py`` with the decoder classes removed. Only VibeVoiceEmbedModel and its output type are new. Extraction is scripted, not hand-copied — see the source repo referenced in the model card. """ 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.configuration_utils import PretrainedConfig from transformers.utils import logging from transformers.modeling_utils import PreTrainedModel from transformers.activations import ACT2FN from transformers.utils import ModelOutput from transformers.models.auto import AutoModel # noqa: F401 (parity with upstream imports) from .configuration_vibevoice_embed import VibeVoiceEmbedConfig 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 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 # 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 @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 # --- VibeVoice-Embed model ------------------------------------------------- # This fragment is concatenated after the encoder classes extracted from upstream # VibeVoice by scripts/build_vibevoice_embed_repo.py. Names it uses (TokenizerEncoder, # VibeVoiceTokenizerEncoderOutput, torch, nn, copy, dataclass, ...) are defined above # in the assembled modeling_vibevoice_embed.py. @dataclass class VibeVoiceEmbedOutput(ModelOutput): """ Output of [`VibeVoiceEmbedModel`]. Args: pooler_output (`torch.FloatTensor` of shape `(batch, vae_dim)`): The voice embedding: latent frames mean-pooled over time. Always float32 — averaging tens of frames in bfloat16 loses precision the embedding is then judged on. Unnormalised by design; L2-normalise before cosine indexing. last_hidden_state (`torch.FloatTensor` of shape `(batch, frames, vae_dim)`): The per-frame latent means (7.5 Hz for the released checkpoint), in the model's compute dtype. """ pooler_output: torch.FloatTensor = None last_hidden_state: torch.FloatTensor = None class VibeVoiceEmbedModel(PreTrainedModel): """The acoustic ENCODER of VibeVoice, standalone, as a voice-embedding model. Upstream VibeVoice pairs this encoder with a decoder as a reconstruction VAE. For embedding only the encoder is needed, and ``encode()`` returns the latent distribution's MEAN — no sampling — so the same clip always embeds identically. """ config_class = VibeVoiceEmbedConfig base_model_prefix = "vibevoice_embed" main_input_name = "input_values" _supports_flash_attn_2 = True _supports_sdpa = True _no_split_modules = ["TokenizerEncoder"] 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") if isinstance(config.encoder_depths, str): encoder_depths = [int(d) for d in config.encoder_depths.split("-")] else: encoder_depths = config.encoder_depths # Identical to how upstream VibeVoiceAcousticTokenizerModel builds its encoder, # so the extracted weights load 1:1. 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 self.encoder = TokenizerEncoder(encoder_config) # post_init, not upstream's bare ``self.apply(self._init_weights)``: on # transformers 5.x it also builds the tied-weights bookkeeping # (all_tied_weights_keys) that from_pretrained requires, while on 4.x it # reduces to the same weight init. self.post_init() def _init_weights(self, module): 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) def encode(self, audio, cache=None, sample_indices=None, use_cache=False, debug=False): """Upstream-compatible: audio ``(batch, 1, samples)`` to latent distribution. Returns [`VibeVoiceTokenizerEncoderOutput`] whose ``mean`` is ``(batch, frames, vae_dim)`` — the deterministic latents ``forward`` pools. """ 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) def forward( self, input_values, padding_mask=None, return_dict: Optional[bool] = None, ): r""" Args: input_values (`torch.FloatTensor` of shape `(batch, samples)` or `(batch, 1, samples)`): Mono waveform at ``config.sampling_rate`` (24 kHz), roughly in [-1, 1]. padding_mask (`torch.Tensor` of shape `(batch, samples)`, *optional*): 1 for real samples, 0 for right-padding. Required for correct pooling of batched variable-length clips: the encoder is causal, so padding cannot corrupt the real frames, but the frames it emits FOR the padding would otherwise be averaged into the embedding — a clip's vector would depend on what it was batched with. """ return_dict = return_dict if return_dict is not None else self.config.use_return_dict if input_values.dim() == 2: input_values = input_values.unsqueeze(1) if input_values.dim() != 3 or input_values.shape[1] != 1: raise ValueError( f"input_values must be (batch, samples) or (batch, 1, samples), got " f"{tuple(input_values.shape)}" ) frames = self.encode(input_values.to(self.dtype)).mean # (batch, frames, vae_dim) # Pool in float32: the embedding is compared at tolerances bfloat16 cannot hold. frames32 = frames.float() if padding_mask is not None: hop = self.config.hop_length lengths = padding_mask.to(torch.long).sum(dim=-1) n_frames = torch.clamp((lengths + hop - 1) // hop, min=1) n_frames = torch.minimum( n_frames, torch.full_like(n_frames, frames32.shape[1]) ) frame_mask = ( torch.arange(frames32.shape[1], device=frames32.device)[None, :] < n_frames[:, None] ) pooled = (frames32 * frame_mask[..., None]).sum(dim=1) pooled = pooled / frame_mask.sum(dim=1).clamp(min=1)[..., None] else: pooled = frames32.mean(dim=1) if not return_dict: return (pooled, frames) return VibeVoiceEmbedOutput(pooler_output=pooled, last_hidden_state=frames) # No AutoModel.register() here: for a trust_remote_code repo the ``auto_map`` entry in # config.json is the registration, and an import-time register() call would raise on # the second import of the dynamic module. __all__ = ["VibeVoiceEmbedModel", "VibeVoiceEmbedOutput"]