Coverage for wrapper/vibevoice/modular_vibevoice_tokenizer.py: 36%
675 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-09 04:47 +0000
1# mypy: ignore-errors
2# Copy from: https://github.com/microsoft/VibeVoice/blob/main/vibevoice/modular/modular_vibevoice_tokenizer.py
4import math
6from functools import partial
7from dataclasses import dataclass
9from typing import Optional
10from typing import Union
11from typing import List
12from typing import Dict
13from typing import Any
14from typing import Tuple
16import copy
18import numpy as np
20import torch
21import torch.nn as nn
22import torch.nn.functional as F
24from transformers import AutoModel
25from transformers.utils import logging
26from transformers.modeling_utils import PreTrainedModel
27from transformers.activations import ACT2FN
29from configuration_vibevoice import VibeVoiceAcousticTokenizerConfig
30from configuration_vibevoice import VibeVoiceSemanticTokenizerConfig
32logger = logging.get_logger(__name__)
35# Normalization modules
36class ConvLayerNorm(nn.LayerNorm):
37 """
38 Convolution-friendly LayerNorm that moves channels to last dimensions
39 before running the normalization and moves them back to original position right after.
40 """
41 def __init__(
42 self,
43 normalized_shape: Union[int, List[int], torch.Size],
44 **kwargs: Any
45 ) -> None:
46 super().__init__(normalized_shape, **kwargs)
48 def forward(self, x: torch.Tensor) -> torch.Tensor:
49 x = x.transpose(1, 2) # b ... t -> b t ...
50 x = nn.functional.layer_norm(
51 x.float(),
52 self.normalized_shape,
53 self.weight.float(),
54 self.bias.float(),
55 self.eps
56 ).type_as(x)
57 x = x.transpose(1, 2) # b t ... -> b ... t
58 return x
61class RMSNorm(nn.Module):
62 def __init__(
63 self,
64 dim: int,
65 eps: float = 1e-5,
66 elementwise_affine: bool = True,
67 weight_shape: Optional[Tuple[int, ...]] = None
68 ) -> None:
69 super().__init__()
70 self.dim = dim
71 self.eps = eps
72 self.elementwise_affine = elementwise_affine
73 if self.elementwise_affine:
74 weight_shape = (dim,) if weight_shape is None else weight_shape
75 self.weight = nn.Parameter(torch.ones(weight_shape))
76 else:
77 self.register_parameter('weight', None)
79 def _norm(self, x: torch.Tensor) -> torch.Tensor:
80 return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
82 def forward(self, x: torch.Tensor) -> torch.Tensor:
83 output = self._norm(x.float()).type_as(x)
84 if self.weight is not None:
85 output = output * self.weight
86 return output
88 def extra_repr(self) -> str:
89 return f'dim={self.dim}, eps={self.eps}, elementwise_affine={self.elementwise_affine}'
92class ConvRMSNorm(RMSNorm):
93 def __init__(
94 self,
95 dim: int,
96 eps: float = 1e-5,
97 elementwise_affine: bool = True,
98 weight_shape: Optional[Tuple[int, ...]] = None
99 ) -> None:
100 super().__init__(dim, eps, elementwise_affine, weight_shape)
102 def forward(
103 self,
104 x: torch.Tensor
105 ) -> torch.Tensor:
106 x = x.transpose(1, 2) # b ... t -> b t ...
107 output = self._norm(x.float()).type_as(x)
108 if self.weight is not None:
109 output = output * self.weight
111 output = output.transpose(1, 2) # b t ... -> b ... t
112 return output
115# Convolutional layers and utilities
116CONV_NORMALIZATIONS = frozenset([
117 'none', 'weight_norm',
118 'spectral_norm',
119 'time_layer_norm',
120 'layer_norm',
121 'time_group_norm'])
124def apply_parametrization_norm(module: nn.Module, norm: str = 'none') -> nn.Module:
125 assert norm in CONV_NORMALIZATIONS
126 if norm == 'weight_norm':
127 return nn.utils.weight_norm(module)
128 if norm == 'spectral_norm':
129 return nn.utils.spectral_norm(module)
130 # We already check was in CONV_NORMALIZATION, so any other choice
131 # doesn't need reparametrization.
132 return module
135def get_norm_module(module: nn.Module, causal: bool = False, norm: str = 'none', **norm_kwargs: Any) -> nn.Module:
136 """Return the proper normalization module. If causal is True, this will ensure the returned
137 module is causal, or return an error if the normalization doesn't support causal evaluation.
138 """
139 assert norm in CONV_NORMALIZATIONS
140 if norm == 'layer_norm':
141 assert isinstance(module, nn.modules.conv._ConvNd)
142 return ConvLayerNorm(module.out_channels, **norm_kwargs)
143 if norm == 'time_group_norm':
144 if causal:
145 raise ValueError("GroupNorm doesn't support causal evaluation.")
146 assert isinstance(module, nn.modules.conv._ConvNd)
147 return nn.GroupNorm(1, module.out_channels, **norm_kwargs)
148 return nn.Identity()
151def get_extra_padding_for_conv1d(
152 x: torch.Tensor,
153 kernel_size: int,
154 stride: int,
155 padding_total: int = 0
156) -> int:
157 """Calculate extra padding needed for convolution to have the same output length"""
158 length = x.shape[-1]
159 n_frames = (length - kernel_size + padding_total) / stride + 1
160 ideal_length = (math.ceil(n_frames) - 1) * stride + (kernel_size - padding_total)
161 return ideal_length - length
164def pad1d(
165 x: torch.Tensor,
166 paddings: Tuple[int, int],
167 mode: str = 'zero',
168 value: float = 0.
169) -> torch.Tensor:
170 """Pad 1D input with handling for small inputs in reflect mode"""
171 length = x.shape[-1]
172 padding_left, padding_right = paddings
173 assert padding_left >= 0 and padding_right >= 0, (padding_left, padding_right)
174 if mode == 'reflect':
175 max_pad = max(padding_left, padding_right)
176 extra_pad = 0
177 if length <= max_pad:
178 extra_pad = max_pad - length + 1
179 x = F.pad(x, (0, extra_pad))
180 padded = F.pad(x, paddings, mode, value)
181 end = padded.shape[-1] - extra_pad
182 return padded[..., :end]
183 else:
184 return F.pad(x, paddings, mode, value)
187def unpad1d(
188 x: torch.Tensor,
189 paddings: Tuple[int, int]
190) -> torch.Tensor:
191 """Remove padding from x, handling properly zero padding. Only for 1d!"""
192 padding_left, padding_right = paddings
193 assert padding_left >= 0 and padding_right >= 0, (padding_left, padding_right)
194 assert (padding_left + padding_right) <= x.shape[-1]
195 end = x.shape[-1] - padding_right
196 return x[..., padding_left: end]
199class NormConv1d(nn.Module):
200 """Wrapper around Conv1d and normalization applied to this conv"""
201 def __init__(
202 self,
203 *args: Any,
204 causal: bool = False,
205 norm: str = 'none',
206 norm_kwargs: Dict[str, Any] = {},
207 **kwargs: Any
208 ) -> None:
209 super().__init__()
210 self.conv = apply_parametrization_norm(nn.Conv1d(*args, **kwargs), norm)
211 self.norm = get_norm_module(self.conv, causal, norm, **norm_kwargs)
212 self.norm_type = norm
214 def forward(
215 self,
216 x: torch.Tensor,
217 ) -> torch.Tensor:
218 x = self.conv(x)
219 x = self.norm(x)
220 return x
223class NormConvTranspose1d(nn.Module):
224 """Wrapper around ConvTranspose1d and normalization applied to this conv"""
225 def __init__(
226 self,
227 *args: Any,
228 causal: bool = False,
229 norm: str = 'none',
230 norm_kwargs: Dict[str, Any] = {},
231 **kwargs: Any
232 ) -> None:
233 super().__init__()
234 self.convtr = apply_parametrization_norm(nn.ConvTranspose1d(*args, **kwargs), norm)
235 self.norm = get_norm_module(self.convtr, causal, norm, **norm_kwargs)
236 self.norm_type = norm
238 def forward(self, x: torch.Tensor) -> torch.Tensor:
239 x = self.convtr(x)
240 x = self.norm(x)
241 return x
244class VibeVoiceTokenizerStreamingCache:
245 """Cache for streaming convolution, similar to KV cache in attention"""
246 def __init__(self) -> None:
247 self.cache: Dict[Tuple[str, int], torch.Tensor] = {} # Dict mapping (layer_id, sample_idx) to state tensor
249 def get(
250 self,
251 layer_id: str,
252 sample_indices: torch.Tensor
253 ) -> Optional[torch.Tensor]:
254 """Get cached states for given layer and sample indices"""
255 states = []
256 max_length = 0
258 # First pass: collect states and find max length
259 for idx in sample_indices.tolist():
260 key = (layer_id, idx)
261 if key not in self.cache:
262 return None # If any sample is missing, return None
263 state = self.cache[key]
264 states.append(state)
265 max_length = max(max_length, state.shape[-1])
267 # Second pass: pad states to max length if needed
268 if len(states) > 0 and states[0].dim() >= 2:
269 padded_states = []
270 for state in states:
271 if state.shape[-1] < max_length:
272 # Pad on the time dimension (last dimension)
273 pad_size = max_length - state.shape[-1]
274 # Pad with zeros on the LEFT to align the most recent samples
275 padded_state = F.pad(state, (pad_size, 0), mode='constant', value=0)
276 padded_states.append(padded_state)
277 else:
278 padded_states.append(state)
279 return torch.stack(padded_states, dim=0)
280 else:
281 return torch.stack(states, dim=0)
283 def set(self, layer_id: str, sample_indices: torch.Tensor, states: torch.Tensor) -> None:
284 """Set cached states for given layer and sample indices"""
285 for i, idx in enumerate(sample_indices.tolist()):
286 key = (layer_id, idx)
287 self.cache[key] = states[i].detach()
289 def set_to_zero(self, sample_indices: torch.Tensor) -> None:
290 """Set all cached states to zero for given sample indices"""
291 for key in list(self.cache.keys()):
292 layer_id, sample_idx = key
293 if sample_idx in sample_indices.tolist():
294 # Create zero tensor with same shape and dtype as cached tensor
295 cached_tensor = self.cache[key]
296 self.cache[key] = torch.zeros_like(cached_tensor)
298 def clear(
299 self,
300 layer_id: Optional[str] = None,
301 sample_indices: Optional[torch.Tensor] = None
302 ) -> None:
303 """Clear cache for specific layer/samples or everything"""
304 if layer_id is None and sample_indices is None:
305 self.cache.clear()
306 elif layer_id is not None and sample_indices is None:
307 # Clear all samples for a specific layer
308 keys_to_remove = [k for k in self.cache.keys() if k[0] == layer_id]
309 for k in keys_to_remove:
310 del self.cache[k]
311 elif layer_id is not None and sample_indices is not None:
312 # Clear specific samples for a specific layer
313 for idx in sample_indices.tolist():
314 key = (layer_id, idx)
315 self.cache.pop(key, None)
318class SConv1d(nn.Module):
319 """Conv1d with built-in handling of asymmetric or causal padding and normalization."""
320 def __init__(
321 self, in_channels: int,
322 out_channels: int,
323 kernel_size: int,
324 stride: int = 1,
325 dilation: int = 1,
326 groups: int = 1,
327 bias: bool = True,
328 causal: bool = False,
329 norm: str = 'none',
330 norm_kwargs: Dict[str, Any] = {},
331 pad_mode: str = 'reflect'
332 ) -> None:
333 super().__init__()
334 self.conv = NormConv1d(
335 in_channels,
336 out_channels,
337 kernel_size,
338 stride,
339 dilation=dilation,
340 groups=groups,
341 bias=bias,
342 causal=causal,
343 norm=norm,
344 norm_kwargs=norm_kwargs)
345 self.causal = causal
346 self.pad_mode = pad_mode
348 # Store configuration
349 self.kernel_size = kernel_size
350 self.dilation = dilation
351 self.stride = stride
352 self.in_channels = in_channels
353 self.out_channels = out_channels
355 # For causal convolution, we need to maintain kernel_size - 1 samples as context
356 # need to check use which context_size is more suitable
357 # self.context_size = (kernel_size - 1) * dilation
358 self.context_size = (kernel_size - 1) * dilation - (stride - 1)
360 # For non-streaming mode, calculate padding
361 self.padding_total = (kernel_size - 1) * dilation - (stride - 1)
363 # Create a unique layer ID for cache management
364 self._layer_id: Optional[str] = None
366 @property
367 def layer_id(self) -> str:
368 if self._layer_id is None:
369 self._layer_id = f"sconv1d_{id(self)}"
370 assert self._layer_id is not None
371 return self._layer_id
373 def forward(
374 self,
375 x: torch.Tensor,
376 cache: Optional[VibeVoiceTokenizerStreamingCache] = None,
377 sample_indices: Optional[torch.Tensor] = None,
378 use_cache: bool = False,
379 debug: bool = False
380 ) -> torch.Tensor:
381 """
382 Forward pass with optional streaming support via cache.
383 Args:
384 x: Input tensor [batch_size, channels, time]
385 cache: VibeVoiceTokenizerStreamingCache object for maintaining states
386 sample_indices: Indices identifying each sample for cache management
387 use_cache: Whether to use cached states for streaming
388 debug: Whether to print debug information
389 Returns:
390 Output tensor
391 """
392 B, C, T = x.shape
394 # Non-streaming mode
395 if not use_cache or cache is None:
396 return self._forward_non_streaming(x, debug=debug)
398 # Streaming mode
399 assert self.causal, "Streaming mode is only supported for causal convolutions"
400 assert sample_indices is not None, "sample_indices must be provided for streaming mode"
401 assert len(sample_indices) == B, "sample_indices must match batch size"
403 return self._forward_streaming(x, cache, sample_indices, debug)
405 def _forward_streaming(
406 self,
407 x: torch.Tensor,
408 cache: VibeVoiceTokenizerStreamingCache,
409 sample_indices: torch.Tensor,
410 debug: bool = False
411 ) -> torch.Tensor:
412 """Streaming forward pass with cache operations kept separate from compiled code"""
413 B, C, T = x.shape
415 # Cache operations (not compiled)
416 cached_states = cache.get(self.layer_id, sample_indices)
418 if cached_states is None:
419 # First chunk - initialize with zeros for context
420 if self.context_size > 0:
421 cached_states = torch.zeros(B, C, self.context_size, device=x.device, dtype=x.dtype)
422 logger.debug(f"Initialized cache with shape: {cached_states.shape}, context_size={self.context_size}")
423 else:
424 cached_states = torch.zeros(B, C, 0, device=x.device, dtype=x.dtype)
425 logger.debug("No context needed (kernel_size=stride)")
427 # Concatenate cached states with input
428 if cached_states.shape[2] > 0:
429 input_with_context = torch.cat([cached_states, x], dim=2)
430 else:
431 input_with_context = x
433 logger.debug(
434 f"Input shape: {x.shape}, Cache shape: {cached_states.shape}, Combined: {input_with_context.shape}")
436 # Apply convolution directly - no extra padding in streaming mode
437 # The conv layer will handle its own padding internally
438 output = self.conv(input_with_context)
440 logger.debug(f"Output shape: {output.shape}")
442 # Update cache for next chunk
443 if self.context_size > 0:
444 # Calculate how many samples to keep
445 total_input_length = input_with_context.shape[2]
447 # Keep the last context_size samples
448 if total_input_length >= self.context_size:
449 new_cache_start = total_input_length - self.context_size
450 new_cache = input_with_context[:, :, new_cache_start:]
451 else:
452 # If we have less than context_size samples, keep everything
453 new_cache = input_with_context
455 logger.debug(f"New cache shape: {new_cache.shape}")
457 cache.set(self.layer_id, sample_indices, new_cache)
459 return output
461 def _forward_non_streaming(self, x: torch.Tensor, debug: bool = False) -> torch.Tensor:
462 """Standard forward pass without streaming"""
463 B, C, T = x.shape
464 kernel_size = self.kernel_size
465 stride = self.stride
466 # dilation = self.dilation
467 padding_total = self.padding_total
469 # Compute extra padding for stride alignment
470 extra_padding = get_extra_padding_for_conv1d(x, kernel_size, stride, padding_total)
472 logger.debug(f"Input shape: {x.shape}, padding_total={padding_total}, extra_padding={extra_padding}")
474 if self.causal:
475 # Left padding for causal
476 if self.pad_mode == 'constant':
477 x = pad1d(x, (padding_total, extra_padding), mode=self.pad_mode, value=0)
478 else:
479 x = pad1d(x, (padding_total, extra_padding), mode=self.pad_mode)
480 else:
481 # Symmetric padding for non-causal
482 padding_right = padding_total // 2
483 padding_left = padding_total - padding_right
484 x = pad1d(x, (padding_left, padding_right + extra_padding), mode=self.pad_mode)
486 logger.debug(f"After padding: {x.shape}")
488 output = self.conv(x)
490 logger.debug(f"Output shape: {output.shape}")
492 return output
495class SConvTranspose1d(nn.Module):
496 """ConvTranspose1d with built-in handling of asymmetric or causal padding and normalization."""
497 def __init__(
498 self,
499 in_channels: int,
500 out_channels: int,
501 kernel_size: int,
502 stride: int = 1,
503 causal: bool = False,
504 norm: str = 'none',
505 trim_right_ratio: float = 1.,
506 norm_kwargs: Dict[str, Any] = {},
507 bias: bool = True
508 ):
509 super().__init__()
510 self.convtr = NormConvTranspose1d(
511 in_channels, out_channels,
512 kernel_size, stride,
513 causal=causal,
514 norm=norm, norm_kwargs=norm_kwargs,
515 bias=bias)
516 self.causal = causal
517 self.trim_right_ratio = trim_right_ratio
518 assert self.causal or self.trim_right_ratio == 1., \
519 "`trim_right_ratio` != 1.0 only makes sense for causal convolutions"
520 assert self.trim_right_ratio >= 0. and self.trim_right_ratio <= 1.
522 # Store configuration
523 self.kernel_size = kernel_size
524 self.stride = stride
525 self.in_channels = in_channels
526 self.out_channels = out_channels
528 # For transposed convolution, padding calculation is different
529 self.padding_total = kernel_size - stride
531 # For streaming, we need to keep track of input history
532 # Transposed conv needs to see multiple input samples to produce correct output
533 self.context_size = kernel_size - 1
535 # Create a unique layer ID for cache management
536 self._layer_id: Optional[str] = None
538 @property
539 def layer_id(self) -> str:
540 if self._layer_id is None:
541 self._layer_id = f"sconvtr1d_{id(self)}"
542 return self._layer_id
544 def forward(
545 self, x: torch.Tensor,
546 cache: Optional[VibeVoiceTokenizerStreamingCache] = None,
547 sample_indices: Optional[torch.Tensor] = None,
548 use_cache: bool = False,
549 debug: bool = False
550 ) -> torch.Tensor:
551 """
552 Forward pass with optional streaming support via cache.
553 """
554 B, C, T = x.shape
556 # Non-streaming mode
557 if not use_cache or cache is None:
558 return self._forward_non_streaming(x, debug=debug)
560 # Streaming mode
561 assert sample_indices is not None, "sample_indices must be provided for streaming mode"
562 assert len(sample_indices) == B, "sample_indices must match batch size"
564 return self._forward_streaming(x, cache, sample_indices, debug)
566 def _forward_streaming(
567 self,
568 x: torch.Tensor,
569 cache: VibeVoiceTokenizerStreamingCache,
570 sample_indices: torch.Tensor,
571 debug: bool = False
572 ) -> torch.Tensor:
573 """Streaming forward pass with cache operations kept separate from compiled code"""
574 B, C, T = x.shape
576 # Cache operations (not compiled)
577 cached_input = cache.get(self.layer_id, sample_indices)
579 if cached_input is None:
580 # First chunk - no history yet
581 cached_input = torch.zeros(B, C, 0, device=x.device, dtype=x.dtype)
582 logger.debug("Initialized empty cache for transposed conv")
584 # Concatenate cached input with new input
585 full_input = torch.cat([cached_input, x], dim=2)
587 logger.debug(f"Input shape: {x.shape}, Cache shape: {cached_input.shape}, Combined: {full_input.shape}")
589 # First chunk or debug mode - use uncompiled version
590 full_output = self.convtr(full_input)
592 logger.debug(f"Full transposed conv output shape: {full_output.shape}")
594 # Calculate padding to remove
595 if self.causal:
596 padding_right = math.ceil(self.padding_total * self.trim_right_ratio)
597 padding_left = self.padding_total - padding_right
598 else:
599 padding_right = self.padding_total // 2
600 padding_left = self.padding_total - padding_right
602 # Remove padding
603 if padding_left + padding_right > 0:
604 full_output = unpad1d(full_output, (padding_left, padding_right))
606 logger.debug(f"After unpadding: {full_output.shape}")
608 # Determine which part of the output corresponds to the new input
609 if cached_input.shape[2] == 0:
610 # First chunk - return all output
611 output = full_output
612 else:
613 # Subsequent chunks - return only the new output
614 expected_new_output = T * self.stride
616 # Take the last expected_new_output samples
617 if full_output.shape[2] >= expected_new_output:
618 output = full_output[:, :, -expected_new_output:]
619 else:
620 output = full_output
622 logger.debug(f"Final streaming output shape: {output.shape}")
624 # Update cache
625 if full_input.shape[2] > self.context_size:
626 new_cache = full_input[:, :, -self.context_size:]
627 else:
628 new_cache = full_input
630 logger.debug(f"New cache shape: {new_cache.shape}")
632 cache.set(self.layer_id, sample_indices, new_cache)
634 return output
636 def _forward_non_streaming(self, x: torch.Tensor, debug: bool = False) -> torch.Tensor:
637 """Standard forward pass without streaming"""
638 logger.debug(f"Input shape: {x.shape}")
640 # Apply transposed convolution
641 y = self.convtr(x)
643 logger.debug(f"After transposed conv: {y.shape}")
645 # Calculate and remove padding
646 if self.causal:
647 padding_right = math.ceil(self.padding_total * self.trim_right_ratio)
648 padding_left = self.padding_total - padding_right
649 else:
650 padding_right = self.padding_total // 2
651 padding_left = self.padding_total - padding_right
653 if padding_left + padding_right > 0:
654 y = unpad1d(y, (padding_left, padding_right))
656 logger.debug(f"Final output shape: {y.shape}")
658 return y
661# FFN
662class FFN(nn.Module):
663 def __init__(
664 self,
665 embed_dim: int,
666 ffn_dim: int,
667 bias: bool = False,
668 ) -> None:
669 super().__init__()
670 self.embed_dim = embed_dim
671 self.linear1 = nn.Linear(self.embed_dim, ffn_dim, bias=bias)
672 self.gelu = ACT2FN["gelu"]
673 self.linear2 = nn.Linear(ffn_dim, self.embed_dim, bias=bias)
675 def forward(self, x: torch.Tensor) -> torch.Tensor:
676 x = self.linear1(x)
677 x = self.gelu(x)
678 x = self.linear2(x)
679 return x
682class Convlayer(nn.Module):
683 def __init__(
684 self,
685 in_channels: int,
686 out_channels: int,
687 kernel_size: int,
688 stride: int = 1,
689 dilation: int = 1,
690 groups: int = 1,
691 bias: bool = True,
692 pad_mode: str = 'zeros',
693 norm: str = 'weight_norm',
694 causal: bool = True,
695 ) -> None:
696 super().__init__()
697 self.conv = SConv1d(
698 in_channels, out_channels, kernel_size, stride=stride, dilation=dilation,
699 groups=groups, bias=bias, pad_mode=pad_mode, norm=norm, causal=causal)
701 def forward(self, x: torch.Tensor) -> torch.Tensor:
702 return self.conv(x)
705class Block1D(nn.Module):
706 def __init__(
707 self,
708 dim: int,
709 kernel_size: int = 7,
710 drop_path: float = 0.,
711 mixer_layer: str = "conv",
712 layer_scale_init_value: float = 1e-6,
713 **kwargs: Any
714 ) -> None:
715 super().__init__()
717 if kwargs.get('layernorm', 'LN') == 'LN':
718 self.norm = ConvLayerNorm(dim, eps=kwargs.get('eps', 1e-6))
719 self.ffn_norm = ConvLayerNorm(dim, eps=kwargs.get('eps', 1e-6))
720 elif kwargs.get('layernorm', 'RMSNorm') == 'RMSNorm':
721 self.norm = ConvRMSNorm(dim, eps=kwargs.get('eps', 1e-6))
722 self.ffn_norm = ConvRMSNorm(dim, eps=kwargs.get('eps', 1e-6))
724 if mixer_layer == 'conv':
725 self.mixer = Convlayer(
726 dim, dim,
727 groups=kwargs.get('groups', 1),
728 kernel_size=kernel_size,
729 pad_mode=kwargs.get('pad_mode', 'reflect'),
730 norm=kwargs.get('norm', 'none'),
731 causal=kwargs.get('causal', True),
732 bias=kwargs.get('bias', True),
733 )
734 elif mixer_layer == 'depthwise_conv':
735 self.mixer = Convlayer(
736 dim, dim, groups=dim,
737 kernel_size=kernel_size,
738 pad_mode=kwargs.get('pad_mode', 'reflect'),
739 norm=kwargs.get('norm', 'none'),
740 causal=kwargs.get('causal', True),
741 bias=kwargs.get('bias', True),
742 )
743 else:
744 raise ValueError(f"Unsupported mixer layer: {mixer_layer}")
746 self.ffn = FFN(
747 dim,
748 kwargs.get('ffn_expansion', 4) * dim,
749 bias=kwargs.get('bias', False),
750 )
751 self.drop_path = nn.Identity() if drop_path <= 0. else nn.modules.DropPath(drop_path)
753 if layer_scale_init_value > 0:
754 self.gamma = nn.Parameter(layer_scale_init_value * torch.ones((dim)), requires_grad=True)
755 self.ffn_gamma = nn.Parameter(layer_scale_init_value * torch.ones((dim)), requires_grad=True)
756 else:
757 self.gamma = None
758 self.ffn_gamma = None
760 def forward(
761 self,
762 x: torch.Tensor,
763 ) -> torch.Tensor:
764 # mixer
765 residual = x
766 x = self.norm(x)
767 x = self.mixer(x)
768 if self.gamma is not None:
769 x = x * self.gamma.unsqueeze(-1)
770 x = residual + self.drop_path(x)
772 # ffn
773 residual = x
774 x = self.ffn_norm(x)
775 x = x.permute(0, 2, 1)
776 x = self.ffn(x)
777 x = x.permute(0, 2, 1)
778 if self.ffn_gamma is not None:
779 x = x * self.ffn_gamma.unsqueeze(-1)
780 x = residual + self.drop_path(x)
782 return x
785class TokenizerEncoder(nn.Module):
786 """
787 Encoder component for the VibeVoice tokenizer that converts audio to latent representations.
789 Args:
790 config: Configuration object with model parameters
791 """
792 def __init__(self, config: Any) -> None:
793 super().__init__()
795 # Extract parameters from config
796 self.channels = config.channels
797 self.dimension = config.dimension
798 self.n_filters = config.n_filters
799 self.ratios = list(reversed(config.ratios))
800 self.depths = config.depths
801 self.n_residual_layers = getattr(config, "n_residual_layers", 1)
802 self.hop_length = np.prod(self.ratios)
803 self.causal = config.causal
805 # Additional config parameters with defaults
806 kernel_size = getattr(config, "kernel_size", 7)
807 last_kernel_size = getattr(config, "last_kernel_size", 7)
808 norm = getattr(config, "norm", "none")
809 norm_params = getattr(config, "norm_params", {})
810 pad_mode = getattr(config, "pad_mode", "reflect")
811 bias = getattr(config, "bias", True)
812 layernorm = getattr(config, "layernorm", "LN")
813 layernorm_eps = getattr(config, "layernorm_eps", 1e-6)
814 layernorm_elementwise_affine = getattr(config, "layernorm_elementwise_affine", True)
815 drop_path_rate = getattr(config, "drop_path_rate", 0.0)
816 mixer_layer = getattr(config, "mixer_layer", "conv")
817 layer_scale_init_value = getattr(config, "layer_scale_init_value", 0)
818 disable_last_norm = getattr(config, "disable_last_norm", False)
820 # determine the norm type based on layernorm
821 # norm_type is Any because it can be either a class (ConvLayerNorm) or a partial callable (partial[ConvRMSNorm])
822 norm_type: Any
823 if layernorm == 'LN':
824 norm_type = ConvLayerNorm
825 elif layernorm == 'RMSNorm':
826 norm_type = partial(ConvRMSNorm, elementwise_affine=layernorm_elementwise_affine)
827 else:
828 raise ValueError(f"Unsupported norm type: {layernorm}")
830 # stem and intermediate downsampling conv layers
831 stem = nn.Sequential(
832 SConv1d(
833 self.channels,
834 self.n_filters,
835 kernel_size,
836 norm=norm,
837 norm_kwargs=norm_params,
838 causal=self.causal,
839 pad_mode=pad_mode, bias=bias))
841 self.downsample_layers = nn.ModuleList()
842 self.downsample_layers.append(stem)
843 for i in range(len(self.ratios)):
844 in_ch = self.n_filters * (2 ** i)
845 out_ch = self.n_filters * (2 ** (i + 1))
846 downsample_layer = nn.Sequential(
847 SConv1d(
848 in_ch, out_ch,
849 kernel_size=self.ratios[i] * 2,
850 stride=self.ratios[i], causal=self.causal,
851 pad_mode=pad_mode, norm=norm, bias=bias))
852 self.downsample_layers.append(downsample_layer)
854 # configure the transformer blocks
855 layer_type = partial(
856 Block1D,
857 mixer_layer=mixer_layer,
858 layernorm=layernorm,
859 eps=layernorm_eps,
860 causal=self.causal,
861 pad_mode=pad_mode,
862 norm=norm,
863 bias=bias,
864 layer_scale_init_value=layer_scale_init_value,
865 )
867 self.stages = nn.ModuleList()
868 dp_rates = [x.item() for x in torch.linspace(0, drop_path_rate, sum(self.depths))]
869 cur = 0
871 for i in range(len(self.depths)):
872 in_ch = self.n_filters * (2 ** i)
873 stage = nn.Sequential(
874 *[layer_type(dim=in_ch, drop_path=dp_rates[cur + j]) for j in range(self.depths[i])]
875 )
876 self.stages.append(stage)
877 cur += self.depths[i]
879 if not disable_last_norm:
880 self.norm = norm_type(in_ch, eps=layernorm_eps)
881 else:
882 self.norm = nn.Identity()
883 self.head = SConv1d(
884 in_ch, self.dimension,
885 kernel_size=last_kernel_size,
886 causal=self.causal,
887 pad_mode=pad_mode,
888 norm=norm,
889 bias=bias)
891 def forward_features(
892 self,
893 x: torch.Tensor,
894 cache: Optional[VibeVoiceTokenizerStreamingCache] = None,
895 sample_indices: Optional[torch.Tensor] = None,
896 use_cache: bool = False,
897 debug: bool = False
898 ) -> torch.Tensor:
899 for i in range(len(self.depths)):
900 # Apply downsampling
901 for layer in self.downsample_layers[i]:
902 if isinstance(layer, SConv1d):
903 x = layer(x, cache=cache, sample_indices=sample_indices, use_cache=use_cache, debug=debug)
904 else:
905 x = layer(x)
907 # Apply stage (Block1D contains Convlayer which contains SConv1d)
908 for block in self.stages[i]:
909 if hasattr(block, 'mixer') and hasattr(block.mixer, 'conv') and isinstance(block.mixer.conv, SConv1d):
910 # Block1D forward with cache support
911 residual = x
912 x = block.norm(x)
913 x = block.mixer.conv(
914 x,
915 cache=cache,
916 sample_indices=sample_indices,
917 use_cache=use_cache,
918 debug=debug)
919 if block.gamma is not None:
920 x = x * block.gamma.unsqueeze(-1)
921 x = residual + x
923 # FFN part
924 residual = x
925 x = block.ffn_norm(x)
926 x = x.permute(0, 2, 1)
927 x = block.ffn(x)
928 x = x.permute(0, 2, 1)
929 if block.ffn_gamma is not None:
930 x = x * block.ffn_gamma.unsqueeze(-1)
931 x = residual + x
932 else:
933 x = block(x)
935 return self.norm(x)
937 def forward(
938 self,
939 x: torch.Tensor,
940 cache: Optional[VibeVoiceTokenizerStreamingCache] = None,
941 sample_indices: Optional[torch.Tensor] = None,
942 use_cache: bool = False,
943 debug: bool = False
944 ) -> torch.Tensor:
945 x = self.forward_features(x, cache=cache, sample_indices=sample_indices, use_cache=use_cache, debug=debug)
946 x = self.head(x, cache=cache, sample_indices=sample_indices, use_cache=use_cache, debug=debug)
947 return x
950class TokenizerDecoder(nn.Module):
951 """
952 Decoder component for the VibeVoice tokenizer that converts latent representations back to audio.
954 Args:
955 config: Configuration object with model parameters
956 """
957 def __init__(self, config: Any) -> None:
958 super().__init__()
960 # Extract parameters from config
961 self.dimension = config.dimension
962 self.channels = config.channels
963 self.n_filters = config.n_filters
964 self.ratios = config.ratios
966 # IMPORTANT CHANGE: Don't reverse depths again since they're already reversed in VibeVoiceAcousticTokenizerModel
967 self.depths = config.depths # Changed from list(reversed(config.depths))
969 self.n_residual_layers = getattr(config, "n_residual_layers", 1)
970 self.hop_length = np.prod(self.ratios)
971 self.causal = config.causal
973 # Additional config parameters with defaults
974 kernel_size = getattr(config, "kernel_size", 7)
975 last_kernel_size = getattr(config, "last_kernel_size", 7)
976 norm = getattr(config, "norm", "none")
977 norm_params = getattr(config, "norm_params", {})
978 pad_mode = getattr(config, "pad_mode", "reflect")
979 bias = getattr(config, "bias", True)
980 layernorm = getattr(config, "layernorm", "LN")
981 layernorm_eps = getattr(config, "layernorm_eps", 1e-6)
982 trim_right_ratio = getattr(config, "trim_right_ratio", 1.0)
983 layernorm_elementwise_affine = getattr(config, "layernorm_elementwise_affine", True)
984 drop_path_rate = getattr(config, "drop_path_rate", 0.0)
985 mixer_layer = getattr(config, "mixer_layer", "conv")
986 layer_scale_init_value = getattr(config, "layer_scale_init_value", 0)
987 disable_last_norm = getattr(config, "disable_last_norm", False)
989 # determine the norm type based on layernorm
990 # norm_type is Any because it can be either a class (ConvLayerNorm) or a partial callable (partial[ConvRMSNorm])
991 norm_type: Any
992 if layernorm == 'LN':
993 norm_type = ConvLayerNorm
994 elif layernorm == 'RMSNorm':
995 norm_type = partial(ConvRMSNorm, elementwise_affine=layernorm_elementwise_affine)
996 else:
997 raise ValueError(f"Unsupported norm type: {layernorm}")
999 # stem and upsampling layers
1000 stem = nn.Sequential(
1001 SConv1d(
1002 self.dimension,
1003 self.n_filters * 2 ** (len(self.depths) - 1),
1004 kernel_size,
1005 norm=norm,
1006 norm_kwargs=norm_params,
1007 causal=self.causal,
1008 pad_mode=pad_mode,
1009 bias=bias))
1011 self.upsample_layers = nn.ModuleList()
1012 self.upsample_layers.append(stem)
1013 for i in range(len(self.ratios)):
1014 in_ch = self.n_filters * (2 ** (len(self.depths) - 1 - i))
1015 out_ch = self.n_filters * (2 ** (len(self.depths) - 1 - i - 1))
1016 upsample_layer = nn.Sequential(
1017 SConvTranspose1d(
1018 in_ch, out_ch,
1019 kernel_size=self.ratios[i] * 2, stride=self.ratios[i],
1020 norm=norm, norm_kwargs=norm_params, bias=bias,
1021 causal=self.causal, trim_right_ratio=trim_right_ratio))
1022 self.upsample_layers.append(upsample_layer)
1024 # configure transformer blocks
1025 layer_type = partial(
1026 Block1D,
1027 mixer_layer=mixer_layer,
1028 layernorm=layernorm,
1029 eps=layernorm_eps,
1030 causal=self.causal,
1031 pad_mode=pad_mode,
1032 norm=norm,
1033 bias=bias,
1034 layer_scale_init_value=layer_scale_init_value,
1035 )
1037 self.stages = nn.ModuleList()
1038 dp_rates = [x.item() for x in torch.linspace(0, drop_path_rate, sum(self.depths))]
1039 cur = 0
1041 # Create stages in the same order as the original model
1042 for i in range(len(self.depths)):
1043 in_ch = self.n_filters * (2 ** (len(self.depths) - 1 - i))
1044 stage = nn.Sequential(
1045 *[layer_type(dim=in_ch, drop_path=dp_rates[cur + j]) for j in range(self.depths[i])]
1046 )
1047 self.stages.append(stage)
1048 cur += self.depths[i]
1050 if not disable_last_norm:
1051 self.norm = norm_type(in_ch, eps=layernorm_eps)
1052 else:
1053 self.norm = nn.Identity()
1054 self.head = SConv1d(
1055 in_ch, self.channels,
1056 kernel_size=last_kernel_size,
1057 causal=self.causal,
1058 pad_mode=pad_mode,
1059 norm=norm,
1060 bias=bias)
1062 def forward_features(
1063 self,
1064 x: torch.Tensor,
1065 cache: Optional[VibeVoiceTokenizerStreamingCache] = None,
1066 sample_indices: Optional[torch.Tensor] = None,
1067 use_cache: bool = False,
1068 debug: bool = False
1069 ) -> torch.Tensor:
1070 for i in range(len(self.depths)):
1071 # Apply upsampling
1072 for layer in self.upsample_layers[i]:
1073 if isinstance(layer, (SConv1d, SConvTranspose1d)):
1074 x = layer(x, cache=cache, sample_indices=sample_indices, use_cache=use_cache, debug=debug)
1075 else:
1076 x = layer(x)
1078 # Apply stage (Block1D contains Convlayer which contains SConv1d)
1079 for block in self.stages[i]:
1080 if hasattr(block, 'mixer') and hasattr(block.mixer, 'conv') and isinstance(block.mixer.conv, SConv1d):
1081 # Block1D forward with cache support
1082 residual = x
1083 x = block.norm(x)
1084 x = block.mixer.conv(
1085 x, cache=cache,
1086 sample_indices=sample_indices,
1087 use_cache=use_cache,
1088 debug=debug)
1089 if block.gamma is not None:
1090 x = x * block.gamma.unsqueeze(-1)
1091 x = residual + x
1093 # FFN part
1094 residual = x
1095 x = block.ffn_norm(x)
1096 x = x.permute(0, 2, 1)
1097 x = block.ffn(x)
1098 x = x.permute(0, 2, 1)
1099 if block.ffn_gamma is not None:
1100 x = x * block.ffn_gamma.unsqueeze(-1)
1101 x = residual + x
1102 else:
1103 x = block(x)
1105 return self.norm(x)
1107 def forward(
1108 self,
1109 x: torch.Tensor,
1110 cache: Optional[VibeVoiceTokenizerStreamingCache] = None,
1111 sample_indices: Optional[torch.Tensor] = None,
1112 use_cache: bool = False,
1113 debug: bool = False
1114 ) -> torch.Tensor:
1115 x = self.forward_features(x, cache=cache, sample_indices=sample_indices, use_cache=use_cache, debug=debug)
1116 x = self.head(x, cache=cache, sample_indices=sample_indices, use_cache=use_cache, debug=debug)
1117 return x
1120@dataclass
1121class VibeVoiceTokenizerEncoderOutput:
1122 """
1123 Output of VibeVoice tokenizer encoder, representing a Gaussian distribution with fixed variance.
1125 Args:
1126 mean (`torch.FloatTensor`): The mean parameters of the distribution.
1127 std (`float` or `torch.FloatTensor`): Fixed standard deviation value.
1128 """
1129 mean: torch.Tensor
1130 std: Optional[Union[float, torch.Tensor]] = None
1132 def sample(
1133 self,
1134 dist_type: str = "fix"
1135 ) -> Tuple[torch.Tensor, Optional[Union[float, torch.Tensor]]]:
1136 """
1137 Sample from the distribution.
1139 Args:
1140 dist_type (`str`): Sampling method, either 'fix' or 'gaussian'.
1142 Returns:
1143 `torch.FloatTensor`: Sampled values.
1144 `torch.FloatTensor` (optional): Standard deviation used (only when dist_type='gaussian').
1145 """
1146 if dist_type == 'fix':
1147 x = self.mean + self.std * torch.randn_like(self.mean)
1148 return x, self.std
1149 elif dist_type == 'gaussian':
1150 assert self.std is not None
1151 batch_size = self.mean.size(0)
1152 value = self.std / 0.8
1153 std = torch.randn(batch_size, device=self.mean.device, dtype=self.mean.dtype) * value
1155 while std.dim() < self.mean.dim():
1156 std = std.unsqueeze(-1)
1158 x = self.mean + std * torch.randn_like(self.mean)
1159 return x, std
1160 else:
1161 return self.mean, self.std
1163 def kl(self) -> torch.Tensor:
1164 """Compute KL divergence between this distribution and a standard normal."""
1165 target = torch.zeros_like(self.mean)
1166 return F.mse_loss(self.mean, target, reduction='none')
1168 def mode(self) -> torch.Tensor:
1169 """Return the distribution mode (which is the mean for Gaussian)."""
1170 return self.mean
1173class VibeVoiceAcousticTokenizerModel(PreTrainedModel):
1174 """VibeVoice speech tokenizer model combining encoder and decoder for acoustic tokens"""
1176 config_class = VibeVoiceAcousticTokenizerConfig
1177 base_model_prefix = "vibevoice_acoustic_tokenizer"
1178 _supports_flash_attn_2 = True
1179 _supports_sdpa = True
1180 _no_split_modules = ["TokenizerEncoder", "TokenizerDecoder"]
1182 def __init__(self, config: Any) -> None:
1183 super().__init__(config)
1185 self.register_buffer('fix_std', torch.tensor(config.fix_std), persistent=False)
1186 self.std_dist_type = getattr(config, "std_dist_type", "fix")
1188 # Parse encoder depths
1189 if isinstance(config.encoder_depths, str):
1190 encoder_depths = [int(d) for d in config.encoder_depths.split('-')]
1191 else:
1192 encoder_depths = config.encoder_depths
1194 # Parse decoder depths if provided
1195 if config.decoder_depths is not None and isinstance(config.decoder_depths, str):
1196 decoder_depths = [int(d) for d in config.decoder_depths.split('-')]
1197 else:
1198 # Default: use reversed encoder depths if decoder_depths is None
1199 decoder_depths = list(reversed(encoder_depths))
1201 # Create encoder config
1202 encoder_config = copy.deepcopy(config)
1203 encoder_config.dimension = config.vae_dim
1204 encoder_config.n_filters = config.encoder_n_filters
1205 encoder_config.ratios = config.encoder_ratios
1206 encoder_config.depths = encoder_depths
1207 encoder_config.norm = config.conv_norm
1208 encoder_config.pad_mode = config.pad_mode
1209 encoder_config.bias = config.conv_bias
1210 encoder_config.layernorm_eps = config.layernorm_eps
1211 encoder_config.layernorm_elementwise_affine = config.layernorm_elementwise_affine
1212 encoder_config.mixer_layer = config.mixer_layer
1213 encoder_config.layer_scale_init_value = config.layer_scale_init_value
1214 encoder_config.disable_last_norm = config.disable_last_norm
1216 # Create decoder config
1217 decoder_config = copy.deepcopy(config)
1218 decoder_config.dimension = config.vae_dim
1219 decoder_config.n_filters = config.decoder_n_filters
1220 decoder_config.ratios = config.decoder_ratios
1221 decoder_config.depths = decoder_depths
1222 decoder_config.norm = config.conv_norm
1223 decoder_config.pad_mode = config.pad_mode
1224 decoder_config.bias = config.conv_bias
1225 decoder_config.layernorm_eps = config.layernorm_eps
1226 decoder_config.layernorm_elementwise_affine = config.layernorm_elementwise_affine
1227 decoder_config.mixer_layer = config.mixer_layer
1228 decoder_config.layer_scale_init_value = config.layer_scale_init_value
1229 decoder_config.disable_last_norm = config.disable_last_norm
1231 # Initialize encoder and decoder
1232 self.encoder = TokenizerEncoder(encoder_config)
1233 self.decoder = TokenizerDecoder(decoder_config)
1235 # Initialize weights
1236 self.apply(self._init_weights)
1238 def _init_weights(self, module: nn.Module) -> None:
1239 """Initialize weights for the model"""
1240 if isinstance(module, nn.Linear):
1241 nn.init.normal_(module.weight, std=self.config.weight_init_value)
1242 if module.bias is not None:
1243 nn.init.zeros_(module.bias)
1244 elif isinstance(module, nn.LayerNorm):
1245 nn.init.ones_(module.weight)
1246 nn.init.zeros_(module.bias)
1247 elif isinstance(module, nn.Conv1d):
1248 nn.init.normal_(module.weight, std=self.config.weight_init_value)
1249 if module.bias is not None:
1250 nn.init.zeros_(module.bias)
1252 @torch.no_grad()
1253 def encode(
1254 self,
1255 audio: torch.Tensor,
1256 cache: Optional[VibeVoiceTokenizerStreamingCache] = None,
1257 sample_indices: Optional[torch.Tensor] = None,
1258 use_cache: bool = False,
1259 debug: bool = False
1260 ) -> VibeVoiceTokenizerEncoderOutput:
1261 """Convert audio to latent representations"""
1262 latents = self.encoder(audio, cache=cache, sample_indices=sample_indices, use_cache=use_cache, debug=debug)
1263 return VibeVoiceTokenizerEncoderOutput(mean=latents.permute(0, 2, 1), std=self.fix_std)
1265 @torch.no_grad()
1266 def sampling(
1267 self,
1268 encoder_output: VibeVoiceTokenizerEncoderOutput,
1269 dist_type: Optional[str] = None
1270 ) -> Tuple[torch.Tensor, Optional[Union[float, torch.Tensor]]]:
1271 """Sample from the encoder output distribution"""
1272 dist_type = dist_type or self.std_dist_type
1274 if dist_type == 'fix':
1275 return encoder_output.sample(dist_type='fix')
1276 elif dist_type == 'gaussian':
1277 return encoder_output.sample(dist_type='gaussian')
1278 else:
1279 raise ValueError(f"Unsupported dist_type: {dist_type}, expected 'fix' or 'gaussian'")
1281 @torch.no_grad()
1282 def decode(
1283 self,
1284 latents: torch.Tensor,
1285 cache: Optional[VibeVoiceTokenizerStreamingCache] = None,
1286 sample_indices: Optional[torch.Tensor] = None,
1287 use_cache: bool = False,
1288 debug: bool = False
1289 ) -> torch.Tensor:
1290 """Convert latent representations back to audio"""
1291 if latents.shape[1] == self.config.vae_dim:
1292 pass
1293 else:
1294 latents = latents.permute(0, 2, 1)
1296 audio = self.decoder(latents, cache=cache, sample_indices=sample_indices, use_cache=use_cache, debug=debug)
1297 return audio
1299 def forward(
1300 self,
1301 audio: torch.Tensor,
1302 cache: Optional[VibeVoiceTokenizerStreamingCache] = None,
1303 sample_indices: Optional[torch.Tensor] = None,
1304 use_cache: bool = False,
1305 debug: bool = False
1306 ) -> Tuple[torch.Tensor, torch.Tensor]:
1307 """Full forward pass: encode audio to latents, then decode back to audio"""
1308 encoder_output = self.encode(
1309 audio, cache=cache,
1310 sample_indices=sample_indices,
1311 use_cache=use_cache, debug=debug)
1312 sampled_latents, _ = self.sampling(encoder_output)
1313 reconstructed = self.decode(
1314 sampled_latents, cache=cache,
1315 sample_indices=sample_indices,
1316 use_cache=use_cache, debug=debug)
1317 return reconstructed, sampled_latents
1320class VibeVoiceSemanticTokenizerModel(PreTrainedModel):
1321 """VibeVoice speech tokenizer model with only encoder for semantic tokens"""
1323 config_class = VibeVoiceSemanticTokenizerConfig
1324 base_model_prefix = "vibevoice_semantic_tokenizer"
1325 _supports_flash_attn_2 = True
1326 _supports_sdpa = True
1327 _no_split_modules = ["TokenizerEncoder"]
1329 def __init__(self, config: Any) -> None:
1330 super().__init__(config)
1332 # Parse encoder depths
1333 if isinstance(config.encoder_depths, str):
1334 encoder_depths = [int(d) for d in config.encoder_depths.split('-')]
1335 else:
1336 encoder_depths = config.encoder_depths
1338 # Create encoder config
1339 encoder_config = copy.deepcopy(config)
1340 encoder_config.dimension = config.vae_dim
1341 encoder_config.n_filters = config.encoder_n_filters
1342 encoder_config.ratios = config.encoder_ratios
1343 encoder_config.depths = encoder_depths
1344 encoder_config.norm = config.conv_norm
1345 encoder_config.pad_mode = config.pad_mode
1346 encoder_config.bias = config.conv_bias
1347 encoder_config.layernorm_eps = config.layernorm_eps
1348 encoder_config.layernorm_elementwise_affine = config.layernorm_elementwise_affine
1349 encoder_config.mixer_layer = config.mixer_layer
1350 encoder_config.layer_scale_init_value = config.layer_scale_init_value
1351 encoder_config.disable_last_norm = config.disable_last_norm
1353 # Initialize encoder and decoder
1354 self.encoder = TokenizerEncoder(encoder_config)
1356 # Initialize weights
1357 self.apply(self._init_weights)
1359 def _init_weights(self, module: nn.Module) -> None:
1360 """Initialize weights for the model"""
1361 if isinstance(module, nn.Linear):
1362 nn.init.normal_(module.weight, std=self.config.weight_init_value)
1363 if module.bias is not None:
1364 nn.init.zeros_(module.bias)
1365 elif isinstance(module, nn.LayerNorm):
1366 nn.init.ones_(module.weight)
1367 nn.init.zeros_(module.bias)
1368 elif isinstance(module, nn.Conv1d):
1369 nn.init.normal_(module.weight, std=self.config.weight_init_value)
1370 if module.bias is not None:
1371 nn.init.zeros_(module.bias)
1373 @torch.no_grad()
1374 def encode(
1375 self,
1376 audio: torch.Tensor,
1377 cache: Optional[VibeVoiceTokenizerStreamingCache] = None,
1378 sample_indices: Optional[torch.Tensor] = None,
1379 use_cache: bool = False,
1380 debug: bool = False
1381 ) -> VibeVoiceTokenizerEncoderOutput:
1382 """Convert audio to latent representations"""
1383 latents = self.encoder(
1384 audio, cache=cache, sample_indices=sample_indices,
1385 use_cache=use_cache, debug=debug)
1386 return VibeVoiceTokenizerEncoderOutput(mean=latents.permute(0, 2, 1))
1388 @torch.no_grad()
1389 def sampling(
1390 self,
1391 encoder_output: VibeVoiceTokenizerEncoderOutput,
1392 dist_type: Optional[str] = None
1393 ) -> Tuple[torch.Tensor, Optional[Union[float, torch.Tensor]]]:
1394 """Sample from the encoder output distribution"""
1395 return encoder_output.sample(dist_type='none')
1397 def forward(
1398 self,
1399 audio: torch.Tensor,
1400 cache: Optional[VibeVoiceTokenizerStreamingCache] = None,
1401 sample_indices: Optional[torch.Tensor] = None,
1402 use_cache: bool = False,
1403 debug: bool = False
1404 ) -> Tuple[None, torch.Tensor]:
1405 """Full forward pass: encode audio to latents, then decode back to audio"""
1406 encoder_output = self.encode(
1407 audio, cache=cache, sample_indices=sample_indices,
1408 use_cache=use_cache, debug=debug)
1409 sampled_latents, _ = self.sampling(encoder_output, dist_type='none')
1410 return None, sampled_latents
1413AutoModel.register(VibeVoiceAcousticTokenizerConfig, VibeVoiceAcousticTokenizerModel)
1414AutoModel.register(VibeVoiceSemanticTokenizerConfig, VibeVoiceSemanticTokenizerModel)
1416__all__ = [
1417 "VibeVoiceTokenizerStreamingCache",
1418 "VibeVoiceAcousticTokenizerModel",
1419 "VibeVoiceSemanticTokenizerModel",
1420]