christian.schwarz/VibeVoice
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add VibeVoice-ASR

Browse Source
This commit is contained in:
Zhiliang Peng
2026-01-21 22:18:33 +08:00
committed by GitHub
parent 6c7369bb31
commit 56cb11e7b2
14 changed files with 4062 additions and 94 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vibevoice/modular/modular_vibevoice_tokenizer.py
+313 -10
View File
@@ -17,7 +17,7 @@ from transformers.utils import logging
from transformers.modeling_utils import PreTrainedModel
from transformers.activations import ACT2FN
from .configuration_vibevoice import VibeVoiceAcousticTokenizerConfig
from .configuration_vibevoice import VibeVoiceAcousticTokenizerConfig, VibeVoiceSemanticTokenizerConfig
logger = logging.get_logger(__name__)
@@ -26,14 +26,13 @@ import os
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")
# 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")
# 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
# logger.warning("APEX FusedRMSNorm not available, using native implementation")
# Normalization modules
class ConvLayerNorm(nn.LayerNorm):
@@ -297,7 +296,8 @@ class SConv1d(nn.Module):
cache: Optional[VibeVoiceTokenizerStreamingCache] = None,
sample_indices: Optional[torch.Tensor] = None,
use_cache: bool = False,
debug: bool = False) -> torch.Tensor:
debug: bool = False,
is_final_chunk: bool = False) -> torch.Tensor:
"""
Forward pass with optional streaming support via cache.
@@ -307,6 +307,7 @@ class SConv1d(nn.Module):
sample_indices: Indices identifying each sample for cache management
use_cache: Whether to use cached states for streaming
debug: Whether to print debug information
is_final_chunk: Whether this is the final chunk (adds extra padding for alignment)
Returns:
Output tensor
@@ -322,12 +323,13 @@ class SConv1d(nn.Module):
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)
return self._forward_streaming(x, cache, sample_indices, debug, is_final_chunk)
def _forward_streaming(self, x: torch.Tensor,
cache: VibeVoiceTokenizerStreamingCache,
sample_indices: torch.Tensor,
debug: bool = False) -> torch.Tensor:
debug: bool = False,
is_final_chunk: bool = False) -> torch.Tensor:
"""Streaming forward pass with cache operations kept separate from compiled code"""
B, C, T = x.shape
@@ -350,6 +352,16 @@ class SConv1d(nn.Module):
input_with_context = torch.cat([cached_states, x], dim=2)
else:
input_with_context = x
# For final chunk, add extra padding to ensure ceil behavior (same as non-streaming)
if is_final_chunk:
extra_padding = get_extra_padding_for_conv1d(
input_with_context, self.kernel_size, self.stride, self.padding_total
)
if extra_padding > 0:
input_with_context = pad1d(input_with_context, (0, extra_padding), mode=self.pad_mode)
if debug:
print(f"[DEBUG] Final chunk: added extra_padding={extra_padding}")
if debug:
print(f"[DEBUG] Input shape: {x.shape}, Cache shape: {cached_states.shape}, Combined: {input_with_context.shape}")
@@ -684,6 +696,135 @@ class Block1D(nn.Module):
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, is_final_chunk=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, is_final_chunk=is_final_chunk)
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, is_final_chunk=is_final_chunk)
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, is_final_chunk=False):
x = self.forward_features(x, cache=cache, sample_indices=sample_indices, use_cache=use_cache, debug=debug, is_final_chunk=is_final_chunk)
x = self.head(x, cache=cache, sample_indices=sample_indices, use_cache=use_cache, debug=debug, is_final_chunk=is_final_chunk)
return x
class TokenizerDecoder(nn.Module):
"""
Decoder component for the VibeVoice tokenizer that converts latent representations back to audio.
@@ -821,15 +962,63 @@ class TokenizerDecoder(nn.Module):
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 (only decoder) for acoustic tokens"""
"""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 = ["TokenizerDecoder"]
_no_split_modules = ["TokenizerEncoder", "TokenizerDecoder"]
def __init__(self, config):
super().__init__(config)
@@ -850,6 +1039,21 @@ class VibeVoiceAcousticTokenizerModel(PreTrainedModel):
# 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
@@ -865,6 +1069,8 @@ class VibeVoiceAcousticTokenizerModel(PreTrainedModel):
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
@@ -884,6 +1090,24 @@ class VibeVoiceAcousticTokenizerModel(PreTrainedModel):
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, is_final_chunk=False):
"""Convert audio to latent representations"""
latents = self.encoder(audio, cache=cache, sample_indices=sample_indices, use_cache=use_cache, debug=debug, is_final_chunk=is_final_chunk)
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"""
@@ -895,10 +1119,89 @@ class VibeVoiceAcousticTokenizerModel(PreTrainedModel):
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, is_final_chunk=False):
"""Convert audio to latent representations"""
latents = self.encoder(audio, cache=cache, sample_indices=sample_indices, use_cache=use_cache, debug=debug, is_final_chunk=is_final_chunk)
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",
]