Add VibeVoice-ASR
This commit is contained in:
Zhiliang Peng
@@ -240,9 +240,112 @@ class VibeVoiceConfig(PretrainedConfig):
|
||||
|
||||
super().__init__(**kwargs)
|
||||
|
||||
class VibeVoiceASRConfig(PretrainedConfig):
|
||||
model_type = "vibevoice"
|
||||
is_composition = True
|
||||
sub_configs = {
|
||||
"acoustic_tokenizer_config": VibeVoiceAcousticTokenizerConfig,
|
||||
"semantic_tokenizer_config": VibeVoiceSemanticTokenizerConfig,
|
||||
"decoder_config": Qwen2Config,
|
||||
}
|
||||
# keys_to_ignore_at_inference = ["past_key_values"]
|
||||
# Default tensor parallel plan for base model `Qwen2`
|
||||
base_model_tp_plan = {
|
||||
"layers.*.self_attn.q_proj": "colwise",
|
||||
"layers.*.self_attn.k_proj": "colwise",
|
||||
"layers.*.self_attn.v_proj": "colwise",
|
||||
"layers.*.self_attn.o_proj": "rowwise",
|
||||
"layers.*.mlp.gate_proj": "colwise",
|
||||
"layers.*.mlp.up_proj": "colwise",
|
||||
"layers.*.mlp.down_proj": "rowwise",
|
||||
}
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
acoustic_tokenizer_config=None,
|
||||
semantic_tokenizer_config=None,
|
||||
decoder_config=None,
|
||||
**kwargs
|
||||
):
|
||||
|
||||
# kwargs["_attn_implementation"] = "flash_attention_2"
|
||||
kwargs["_attn_implementation_autoset"] = False
|
||||
|
||||
if acoustic_tokenizer_config is None:
|
||||
self.acoustic_tokenizer_config = self.sub_configs["acoustic_tokenizer_config"]()
|
||||
elif isinstance(acoustic_tokenizer_config, dict):
|
||||
acoustic_tokenizer_config["model_type"] = "vibevoice_acoustic_tokenizer"
|
||||
self.acoustic_tokenizer_config = self.sub_configs["acoustic_tokenizer_config"](**acoustic_tokenizer_config)
|
||||
elif isinstance(acoustic_tokenizer_config, VibeVoiceAcousticTokenizerConfig):
|
||||
# If an instance of the config class is provided
|
||||
self.acoustic_tokenizer_config = acoustic_tokenizer_config
|
||||
|
||||
if semantic_tokenizer_config is None:
|
||||
self.semantic_tokenizer_config = self.sub_configs["semantic_tokenizer_config"]()
|
||||
elif isinstance(semantic_tokenizer_config, dict):
|
||||
semantic_tokenizer_config["model_type"] = "vibevoice_semantic_tokenizer"
|
||||
self.semantic_tokenizer_config = self.sub_configs["semantic_tokenizer_config"](**semantic_tokenizer_config)
|
||||
elif isinstance(semantic_tokenizer_config, VibeVoiceSemanticTokenizerConfig):
|
||||
# If an instance of the config class is provided
|
||||
self.semantic_tokenizer_config = semantic_tokenizer_config
|
||||
|
||||
if decoder_config is None:
|
||||
self.decoder_config = self.sub_configs["decoder_config"]()
|
||||
elif isinstance(decoder_config, dict):
|
||||
# If a dictionary is provided, instantiate the config class with it
|
||||
# self.decoder_config = self.sub_configs["decoder_config"](**decoder_config)
|
||||
if decoder_config.get("model_type", '') == "qwen2":
|
||||
self.decoder_config = Qwen2Config(**decoder_config)
|
||||
else:
|
||||
raise ValueError(f"Unsupported decoder model type: {decoder_config.get('model_type', '')}")
|
||||
elif isinstance(decoder_config, Qwen2Config):
|
||||
# If an instance of the config class is provided
|
||||
self.decoder_config = decoder_config
|
||||
|
||||
# other parameters
|
||||
self.acoustic_vae_dim = getattr(self.acoustic_tokenizer_config, 'vae_dim', 64)
|
||||
self.semantic_vae_dim = getattr(self.semantic_tokenizer_config, 'vae_dim', 128)
|
||||
|
||||
super().__init__(**kwargs)
|
||||
|
||||
def get_text_config(self, decoder: bool = False):
|
||||
"""Return the text (decoder) config for generation."""
|
||||
return self.decoder_config
|
||||
|
||||
@property
|
||||
def vocab_size(self):
|
||||
"""Return vocab_size from decoder config for generation compatibility."""
|
||||
return self.decoder_config.vocab_size
|
||||
|
||||
@property
|
||||
def num_attention_heads(self):
|
||||
"""Return num_attention_heads from decoder config for Ulysses SP compatibility."""
|
||||
return self.decoder_config.num_attention_heads
|
||||
|
||||
@property
|
||||
def num_key_value_heads(self):
|
||||
"""Return num_key_value_heads from decoder config for Ulysses SP compatibility."""
|
||||
return self.decoder_config.num_key_value_heads
|
||||
|
||||
@property
|
||||
def hidden_size(self):
|
||||
"""Return hidden_size from decoder config for model compatibility."""
|
||||
return self.decoder_config.hidden_size
|
||||
|
||||
@property
|
||||
def num_hidden_layers(self):
|
||||
"""Return num_hidden_layers from decoder config for Ulysses SP compatibility."""
|
||||
return self.decoder_config.num_hidden_layers
|
||||
|
||||
@property
|
||||
def head_dim(self):
|
||||
"""Return head_dim from decoder config for Ulysses SP compatibility."""
|
||||
return getattr(self.decoder_config, 'head_dim', self.hidden_size // self.num_attention_heads)
|
||||
|
||||
__all__ = [
|
||||
"VibeVoiceAcousticTokenizerConfig",
|
||||
"VibeVoiceSemanticTokenizerConfig",
|
||||
"VibeVoiceDiffusionHeadConfig",
|
||||
"VibeVoiceConfig"
|
||||
"VibeVoiceConfig",
|
||||
"VibeVoiceASRConfig"
|
||||
]
|
||||
@@ -0,0 +1,496 @@
|
||||
# copied from https://github.com/vibevoice-community/VibeVoice/blob/main/vibevoice/modular/modeling_vibevoice.py
|
||||
from dataclasses import dataclass
|
||||
from typing import Dict, List, Optional, Tuple, Union, Callable
|
||||
from tqdm import tqdm
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import torch.distributed as dist
|
||||
|
||||
from transformers.models.auto import AutoModel, AutoModelForCausalLM
|
||||
|
||||
from transformers.activations import ACT2FN
|
||||
from transformers.modeling_outputs import CausalLMOutput, BaseModelOutputWithPast, ModelOutput
|
||||
from transformers.models.llama.modeling_llama import LlamaRMSNorm
|
||||
from transformers import modeling_utils
|
||||
from transformers.modeling_utils import PreTrainedModel
|
||||
from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
|
||||
from transformers.utils import logging
|
||||
|
||||
|
||||
from .modular_vibevoice_tokenizer import VibeVoiceTokenizerStreamingCache, VibeVoiceAcousticTokenizerModel, VibeVoiceSemanticTokenizerModel
|
||||
from .modular_vibevoice_diffusion_head import VibeVoiceDiffusionHead
|
||||
from vibevoice.schedule.dpm_solver import DPMSolverMultistepScheduler
|
||||
|
||||
from .configuration_vibevoice import VibeVoiceConfig
|
||||
|
||||
|
||||
logger = logging.get_logger(__name__)
|
||||
|
||||
if not hasattr(modeling_utils, "ALL_PARALLEL_STYLES") or modeling_utils.ALL_PARALLEL_STYLES is None:
|
||||
modeling_utils.ALL_PARALLEL_STYLES = ["tp", "none", "colwise", "rowwise"]
|
||||
|
||||
@dataclass
|
||||
class VibeVoiceCausalLMOutputWithPast(ModelOutput):
|
||||
loss: Optional[torch.FloatTensor] = None
|
||||
diffusion_loss: Optional[torch.FloatTensor] = None
|
||||
speech_token_num: Optional[int] = None
|
||||
logits: torch.FloatTensor = None
|
||||
past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
|
||||
hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
|
||||
attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class VibeVoiceGenerationOutput(ModelOutput):
|
||||
"""
|
||||
Output type for VibeVoice generation.
|
||||
|
||||
Args:
|
||||
sequences (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
|
||||
The generated sequences.
|
||||
speech_outputs (`List[torch.FloatTensor]`, *optional*):
|
||||
List of generated speech waveforms or latents for each speech segment.
|
||||
"""
|
||||
sequences: torch.LongTensor = None
|
||||
speech_outputs: Optional[List[torch.FloatTensor]] = None
|
||||
|
||||
|
||||
class SpeechConnector(nn.Module):
|
||||
def __init__(self, input_dim, output_dim):
|
||||
super().__init__()
|
||||
self.fc1 = nn.Linear(input_dim, output_dim)
|
||||
self.norm = LlamaRMSNorm(output_dim, eps=1e-6)
|
||||
self.fc2 = nn.Linear(output_dim, output_dim)
|
||||
|
||||
def forward(self, features, **kwargs):
|
||||
x = self.fc1(features)
|
||||
x = self.norm(x)
|
||||
x = self.fc2(x)
|
||||
return x
|
||||
|
||||
|
||||
# @auto_docstring
|
||||
class VibeVoicePreTrainedModel(PreTrainedModel):
|
||||
config_class = VibeVoiceConfig
|
||||
base_model_prefix = "model"
|
||||
supports_gradient_checkpointing = True
|
||||
_skip_keys_device_placement = "past_key_values"
|
||||
_supports_cache_class = True
|
||||
_supports_flash_attn_2 = True
|
||||
_supports_sdpa = True
|
||||
_supports_quantized_cache = True
|
||||
_supports_static_cache = True
|
||||
_supports_attention_backend = True
|
||||
|
||||
def _init_weights(self, module):
|
||||
if isinstance(module, VibeVoiceDiffusionHead):
|
||||
module.initialize_weights()
|
||||
return
|
||||
|
||||
# Use the language model's initializer_range if available
|
||||
if hasattr(self.config, 'language_model_config') and hasattr(self.config.language_model_config, 'initializer_range'):
|
||||
std = self.config.language_model_config.initializer_range
|
||||
elif hasattr(self.config, 'decoder_config') and hasattr(self.config.decoder_config, 'initializer_range'):
|
||||
std = self.config.decoder_config.initializer_range
|
||||
else:
|
||||
std = 0.02 # Default value
|
||||
|
||||
if isinstance(module, nn.Linear):
|
||||
module.weight.data.normal_(mean=0.0, std=std)
|
||||
if module.bias is not None:
|
||||
module.bias.data.zero_()
|
||||
elif isinstance(module, nn.LayerNorm):
|
||||
module.weight.data.fill_(1.0)
|
||||
module.bias.data.zero_()
|
||||
|
||||
# @auto_docstring
|
||||
class VibeVoiceModel(VibeVoicePreTrainedModel):
|
||||
def __init__(self, config):
|
||||
super().__init__(config)
|
||||
|
||||
if hasattr(config, 'torch_dtype') and config.torch_dtype is not None:
|
||||
if isinstance(config.torch_dtype, str):
|
||||
dtype = getattr(torch, config.torch_dtype)
|
||||
else:
|
||||
dtype = config.torch_dtype
|
||||
else:
|
||||
dtype = torch.float32
|
||||
|
||||
# Initialize Qwen2 model for language modeling
|
||||
lm_config = config.decoder_config
|
||||
self.language_model = AutoModel.from_config(lm_config)
|
||||
|
||||
# Initialize speech components if needed
|
||||
self.acoustic_tokenizer = AutoModel.from_config(config.acoustic_tokenizer_config).to(dtype)
|
||||
self.semantic_tokenizer = AutoModel.from_config(config.semantic_tokenizer_config).to(dtype)
|
||||
|
||||
self.acoustic_connector = SpeechConnector(config.acoustic_vae_dim, lm_config.hidden_size).to(dtype)
|
||||
self.semantic_connector = SpeechConnector(config.semantic_vae_dim, lm_config.hidden_size).to(dtype)
|
||||
|
||||
# Register scaling factors as buffers - use 1D tensors for FSDP compatibility
|
||||
self.register_buffer('speech_scaling_factor', torch.tensor(float('nan')))
|
||||
self.register_buffer('speech_bias_factor', torch.tensor(float('nan')))
|
||||
|
||||
# Initialize prediction head for speech generation
|
||||
self.prediction_head = AutoModel.from_config(config.diffusion_head_config).to(dtype)
|
||||
|
||||
# Initialize noise scheduler
|
||||
self.noise_scheduler = DPMSolverMultistepScheduler(
|
||||
num_train_timesteps=config.diffusion_head_config.ddpm_num_steps,
|
||||
beta_schedule=config.diffusion_head_config.ddpm_beta_schedule,
|
||||
prediction_type=config.diffusion_head_config.prediction_type
|
||||
)
|
||||
|
||||
def get_input_embeddings(self):
|
||||
if hasattr(self.language_model, 'embed_tokens'):
|
||||
# If the language model has an embed_tokens attribute, return it
|
||||
return self.language_model.embed_tokens
|
||||
|
||||
for name, attr in self.language_model.fullmap.items(): # parallel by nnscaler, the name is changed
|
||||
if attr.orig_name == 'embed_tokens.weight':
|
||||
return getattr(self.language_model, name)
|
||||
assert False, 'should not arrive here'
|
||||
|
||||
def set_input_embeddings(self, value):
|
||||
self.language_model.embed_tokens = value
|
||||
|
||||
def set_speech_tokenizers(self, acoustic_tokenizer=None, semantic_tokenizer=None):
|
||||
"""Set the speech tokenizers used for encoding and decoding speech."""
|
||||
self.acoustic_tokenizer = acoustic_tokenizer
|
||||
self.semantic_tokenizer = semantic_tokenizer
|
||||
|
||||
# Reset the encoder to evaluation mode
|
||||
if self.acoustic_tokenizer is not None:
|
||||
self.acoustic_tokenizer.eval()
|
||||
|
||||
if self.semantic_tokenizer is not None:
|
||||
self.semantic_tokenizer.eval()
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.LongTensor = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
position_ids: Optional[torch.LongTensor] = None,
|
||||
past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
|
||||
inputs_embeds: Optional[torch.FloatTensor] = None,
|
||||
use_cache: Optional[bool] = None,
|
||||
output_attentions: Optional[bool] = None,
|
||||
output_hidden_states: Optional[bool] = None,
|
||||
return_dict: Optional[bool] = None,
|
||||
cache_position: Optional[torch.LongTensor] = None,
|
||||
**kwargs,
|
||||
) -> Union[Tuple, BaseModelOutputWithPast]:
|
||||
|
||||
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
||||
|
||||
# Forward through language model
|
||||
outputs = self.language_model(
|
||||
input_ids=input_ids,
|
||||
attention_mask=attention_mask,
|
||||
position_ids=position_ids,
|
||||
past_key_values=past_key_values,
|
||||
inputs_embeds=inputs_embeds,
|
||||
use_cache=use_cache,
|
||||
output_attentions=output_attentions,
|
||||
output_hidden_states=output_hidden_states,
|
||||
return_dict=return_dict,
|
||||
cache_position=cache_position,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
if not return_dict:
|
||||
return outputs
|
||||
|
||||
return BaseModelOutputWithPast(
|
||||
last_hidden_state=outputs.last_hidden_state,
|
||||
past_key_values=outputs.past_key_values,
|
||||
hidden_states=outputs.hidden_states,
|
||||
attentions=outputs.attentions,
|
||||
)
|
||||
|
||||
|
||||
class VibeVoiceForConditionalGeneration(VibeVoicePreTrainedModel):
|
||||
_tied_weights_keys = ["lm_head.weight"]
|
||||
_tp_plan = {"lm_head": "colwise_rep"}
|
||||
|
||||
def __init__(self, config):
|
||||
super().__init__(config)
|
||||
self.model = VibeVoiceModel(config)
|
||||
self.vocab_size = config.decoder_config.vocab_size
|
||||
self.lm_head = nn.Linear(config.decoder_config.hidden_size, self.vocab_size, bias=False)
|
||||
|
||||
self.post_init()
|
||||
|
||||
def get_input_embeddings(self):
|
||||
return self.model.get_input_embeddings()
|
||||
|
||||
def set_input_embeddings(self, value):
|
||||
self.model.set_input_embeddings(value)
|
||||
|
||||
def get_output_embeddings(self):
|
||||
return self.lm_head
|
||||
|
||||
def set_decoder(self, decoder):
|
||||
self.model.language_model = decoder
|
||||
|
||||
def get_decoder(self):
|
||||
return self.model.language_model
|
||||
|
||||
def tie_weights(self):
|
||||
"""
|
||||
Tie the weights between the input embeddings and the output embeddings.
|
||||
"""
|
||||
if getattr(self.config.decoder_config, 'tie_word_embeddings', False):
|
||||
# The standard PreTrainedModel method will handle the tying.
|
||||
# It typically does a simple parameter object assignment, which is
|
||||
# CORRECT to do BEFORE FSDP wraps the model.
|
||||
output_embeddings = self.get_output_embeddings()
|
||||
input_embeddings = self.get_input_embeddings()
|
||||
if hasattr(input_embeddings, 'weight'):
|
||||
output_embeddings.weight = input_embeddings.weight
|
||||
else:
|
||||
# maybe returned input_embeddings a tensor directly
|
||||
output_embeddings.weight = input_embeddings
|
||||
|
||||
if getattr(output_embeddings, "bias", None) is not None:
|
||||
output_embeddings.bias.data = nn.functional.pad(
|
||||
output_embeddings.bias.data,
|
||||
(0, output_embeddings.weight.shape[0] - output_embeddings.bias.shape[0]),
|
||||
"constant",
|
||||
0,
|
||||
)
|
||||
print("Tied input and output embeddings using standard assignment.")
|
||||
else:
|
||||
print("tie_word_embeddings is False, not tying weights.")
|
||||
|
||||
# Also, ensure set_output_embeddings is safe, though your implementation looks okay.
|
||||
# The key is to avoid calling it after accelerator.prepare().
|
||||
def set_output_embeddings(self, new_embeddings):
|
||||
# Your current implementation using data.copy_ is good practice,
|
||||
# but the best way is to not call this after prepare().
|
||||
self.lm_head = new_embeddings
|
||||
|
||||
def forward_speech_features(
|
||||
self,
|
||||
speech_tensors=None,
|
||||
speech_masks=None,
|
||||
speech_type="audio",
|
||||
return_unmask=False
|
||||
):
|
||||
if speech_tensors is None:
|
||||
# Use config to get vae_dim instead of non-existent self.args
|
||||
vae_dim = self.config.acoustic_tokenizer_config.vae_dim
|
||||
audio_features = torch.zeros(1, 1, vae_dim).to(self.get_input_embeddings().weight)
|
||||
connect_features = self.model.acoustic_connector(audio_features)
|
||||
return audio_features, connect_features
|
||||
else:
|
||||
with torch.no_grad():
|
||||
if speech_type == "audio":
|
||||
with torch.no_grad():
|
||||
frames = self.model.acoustic_tokenizer.encode(speech_tensors.unsqueeze(1))[0][0]
|
||||
audio_tokens = frames.sample(self.model.acoustic_tokenizer.std_dist_type)[0]
|
||||
|
||||
elif speech_type == "vae":
|
||||
# Use config to get vae_dim instead of non-existent self.args
|
||||
vae_dim = self.config.acoustic_tokenizer_config.vae_dim
|
||||
speech_mode = speech_tensors.reshape(speech_tensors.size(0), -1, vae_dim)
|
||||
|
||||
# gaussian sample from the speech_mode
|
||||
batch_size = speech_mode.size(0)
|
||||
value = self.model.acoustic_tokenizer.fix_std / 0.8
|
||||
std = torch.randn(batch_size, dtype=speech_mode.dtype, device=speech_mode.device) * value
|
||||
std = std.view(-1, *[1] * (speech_mode.dim() - 1))
|
||||
audio_tokens = speech_mode + std * torch.randn(speech_mode.shape).to(speech_mode)
|
||||
else:
|
||||
raise NotImplementedError(f"Speech type {speech_type} not implemented")
|
||||
|
||||
if torch.isnan(self.model.speech_scaling_factor) or torch.isnan(self.model.speech_bias_factor):
|
||||
scaling_factor = 1. / audio_tokens[speech_masks].flatten().std()
|
||||
bias_factor = -audio_tokens[speech_masks].flatten().mean()
|
||||
|
||||
# Only use distributed operations if the process group is initialized
|
||||
if dist.is_available() and dist.is_initialized():
|
||||
dist.all_reduce(scaling_factor, op=dist.ReduceOp.SUM)
|
||||
dist.all_reduce(bias_factor, op=dist.ReduceOp.SUM)
|
||||
world_size = dist.get_world_size()
|
||||
self.model.speech_scaling_factor.copy_(scaling_factor / world_size)
|
||||
self.model.speech_bias_factor.copy_(bias_factor / world_size)
|
||||
print(f"Speech scaling factor (distributed): {self.model.speech_scaling_factor}, bias factor: {self.model.speech_bias_factor}", flush=True)
|
||||
else:
|
||||
# Single process case
|
||||
self.model.speech_scaling_factor.copy_(scaling_factor)
|
||||
self.model.speech_bias_factor.copy_(bias_factor)
|
||||
print(f"Speech scaling factor (single process): {self.model.speech_scaling_factor}, bias factor: {self.model.speech_bias_factor}", flush=True)
|
||||
|
||||
audio_features = (audio_tokens + self.model.speech_bias_factor) * self.model.speech_scaling_factor
|
||||
|
||||
connect_features = self.model.acoustic_connector(audio_features)
|
||||
if return_unmask:
|
||||
return audio_features, connect_features
|
||||
return audio_features[speech_masks], connect_features[speech_masks]
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.LongTensor = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
position_ids: Optional[torch.LongTensor] = None,
|
||||
past_key_values: Optional[List[torch.FloatTensor]] = None,
|
||||
inputs_embeds: Optional[torch.FloatTensor] = None,
|
||||
labels: Optional[torch.LongTensor] = None,
|
||||
use_cache: Optional[bool] = False,
|
||||
output_attentions: Optional[bool] = None,
|
||||
output_hidden_states: Optional[bool] = None,
|
||||
return_dict: Optional[bool] = None,
|
||||
cache_position: Optional[torch.LongTensor] = None,
|
||||
# New arguments for speech processing and loss calculation
|
||||
speech_tensors: Optional[torch.FloatTensor] = None,
|
||||
speech_masks: Optional[torch.BoolTensor] = None,
|
||||
speeches_loss_input: Optional[torch.FloatTensor] = None,
|
||||
speech_semantic_tensors: Optional[torch.FloatTensor] = None,
|
||||
acoustic_input_mask: Optional[torch.BoolTensor] = None,
|
||||
acoustic_loss_mask: Optional[torch.BoolTensor] = None,
|
||||
ddpm_batch_mul: int = 1,
|
||||
**kwargs: Optional[Dict[str, Union[torch.Tensor, str]]],
|
||||
) -> Union[Tuple, VibeVoiceCausalLMOutputWithPast]:
|
||||
|
||||
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
||||
|
||||
x = self.get_input_embeddings()(input_ids)
|
||||
|
||||
semantic_speech_all_connect_features = self.model.semantic_connector(speech_semantic_tensors)
|
||||
if speeches_loss_input is not None:
|
||||
# only part audio need diffuse
|
||||
speech_all_features, speech_all_connect_features = self.forward_speech_features(
|
||||
speech_tensors=speech_tensors.type_as(x) if speech_tensors is not None else None,
|
||||
speech_masks=speech_masks,
|
||||
speech_type=kwargs.get("speech_type", "audio"),
|
||||
return_unmask=True
|
||||
)
|
||||
if speech_tensors is not None:
|
||||
if semantic_speech_all_connect_features is not None:
|
||||
x[acoustic_input_mask] = (
|
||||
speech_all_connect_features[speech_masks]
|
||||
+ semantic_speech_all_connect_features[speech_masks]
|
||||
)
|
||||
else:
|
||||
x[acoustic_input_mask] = speech_all_connect_features[speech_masks]
|
||||
|
||||
# Select only the target segments' latents for diffusion loss.
|
||||
# Both masks are [num_segments, max_latent_len]; using 2D mask on [B,T,D] selects [N_true, D].
|
||||
target_latent_mask = speeches_loss_input & speech_masks
|
||||
speech_features = speech_all_features[target_latent_mask]
|
||||
speech_connect_features = speech_all_connect_features[target_latent_mask]
|
||||
else:
|
||||
speech_features, speech_connect_features = self.forward_speech_features(
|
||||
speech_tensors=speech_tensors.type_as(x) if speech_tensors is not None else None,
|
||||
speech_masks=speech_masks,
|
||||
speech_type=kwargs.get("speech_type", "audio"),
|
||||
)
|
||||
if speech_tensors is not None:
|
||||
x[acoustic_input_mask] = speech_connect_features
|
||||
|
||||
outputs = self.model(
|
||||
input_ids=None,
|
||||
attention_mask=attention_mask,
|
||||
position_ids=position_ids,
|
||||
past_key_values=past_key_values,
|
||||
inputs_embeds=x,
|
||||
use_cache=use_cache,
|
||||
output_attentions=output_attentions,
|
||||
output_hidden_states=False,
|
||||
return_dict=return_dict,
|
||||
cache_position=cache_position,
|
||||
)
|
||||
|
||||
hidden_states = outputs.last_hidden_state
|
||||
logits = self.lm_head(hidden_states)
|
||||
# logits = logits.float()
|
||||
|
||||
loss = None
|
||||
if labels is not None:
|
||||
# The custom CE loss with masking is calculated in the training script.
|
||||
# We leave the standard loss calculation here as None.
|
||||
pass
|
||||
|
||||
# --- Diffusion Loss Calculation ---
|
||||
diffusion_loss = None
|
||||
# This block is executed only if we are in a context that involves speech.
|
||||
if speech_tensors is not None and acoustic_loss_mask.sum().item() > 0:
|
||||
condition_features = hidden_states[acoustic_loss_mask]
|
||||
|
||||
speech_len, latent_size = speech_features.shape
|
||||
|
||||
noise = torch.randn(
|
||||
(speech_len * ddpm_batch_mul, latent_size),
|
||||
device=hidden_states.device,
|
||||
dtype=hidden_states.dtype
|
||||
)
|
||||
|
||||
timesteps = torch.multinomial(
|
||||
torch.ones(self.config.diffusion_head_config.ddpm_num_steps),
|
||||
speech_len * ddpm_batch_mul,
|
||||
replacement=True,
|
||||
).to(hidden_states.device)
|
||||
|
||||
speech_features_repeated = speech_features.repeat_interleave(ddpm_batch_mul, dim=0)
|
||||
condition_features_repeated = condition_features.repeat_interleave(ddpm_batch_mul, dim=0)
|
||||
|
||||
noisy_speech_features = self.model.noise_scheduler.add_noise(
|
||||
speech_features_repeated, noise, timesteps
|
||||
)
|
||||
|
||||
model_output = self.model.prediction_head(
|
||||
noisy_speech_features,
|
||||
timesteps.type_as(x),
|
||||
condition_features_repeated
|
||||
)
|
||||
|
||||
prediction_type = self.config.diffusion_head_config.prediction_type
|
||||
if prediction_type == "epsilon":
|
||||
target_for_loss = noise
|
||||
elif prediction_type == "v_prediction":
|
||||
target_for_loss = self.model.noise_scheduler.get_velocity(
|
||||
speech_features_repeated, noise, timesteps
|
||||
)
|
||||
else:
|
||||
raise NotImplementedError(f"Prediction type {prediction_type} not implemented")
|
||||
|
||||
diffusion_loss = F.mse_loss(model_output.float(), target_for_loss.float(), reduction='sum')
|
||||
if latent_size > 0 and ddpm_batch_mul > 0:
|
||||
diffusion_loss = diffusion_loss / latent_size / ddpm_batch_mul
|
||||
else:
|
||||
diffusion_loss = torch.tensor(0.0, device=diffusion_loss.device)
|
||||
|
||||
else:
|
||||
# Dummy loss for DDP to work when there are no speech samples in a batch,
|
||||
# but we are in a speech context.
|
||||
diffusion_loss = sum(p.sum() for p in self.model.prediction_head.parameters()) * 0.0
|
||||
diffusion_loss += sum(p.sum() for p in self.model.acoustic_connector.parameters()) * 0.0
|
||||
diffusion_loss += sum(p.sum() for p in self.model.semantic_connector.parameters()) * 0.0
|
||||
# --- End Diffusion Loss Calculation ---
|
||||
|
||||
if not return_dict:
|
||||
output = (logits, speech_len) + outputs.to_tuple()[1:]
|
||||
return (loss, diffusion_loss) + output
|
||||
|
||||
return VibeVoiceCausalLMOutputWithPast(
|
||||
loss=loss,
|
||||
diffusion_loss=diffusion_loss,
|
||||
speech_token_num=speech_len if speech_tensors is not None else 0,
|
||||
logits=logits,
|
||||
past_key_values=outputs.past_key_values,
|
||||
hidden_states=outputs.hidden_states,
|
||||
attentions=outputs.attentions,
|
||||
)
|
||||
|
||||
AutoModel.register(VibeVoiceConfig, VibeVoiceModel)
|
||||
AutoModelForCausalLM.register(VibeVoiceConfig, VibeVoiceForConditionalGeneration)
|
||||
|
||||
__all__ = [
|
||||
"VibeVoiceModel",
|
||||
"VibeVoicePreTrainedModel",
|
||||
"VibeVoiceForConditionalGeneration",
|
||||
"VibeVoiceCausalLMOutputWithPast",
|
||||
"VibeVoiceGenerationOutput",
|
||||
]
|
||||
@@ -0,0 +1,520 @@
|
||||
from typing import List, Optional, Tuple, Union
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from transformers.models.auto import AutoModel, AutoModelForCausalLM
|
||||
|
||||
from transformers.modeling_outputs import CausalLMOutput, BaseModelOutputWithPast
|
||||
from transformers import modeling_utils
|
||||
from transformers.modeling_utils import PreTrainedModel
|
||||
from transformers.utils import logging
|
||||
from transformers.generation import GenerationMixin
|
||||
|
||||
from .modular_vibevoice_tokenizer import (
|
||||
VibeVoiceTokenizerStreamingCache,
|
||||
VibeVoiceTokenizerEncoderOutput
|
||||
)
|
||||
|
||||
from .configuration_vibevoice import VibeVoiceASRConfig
|
||||
from .modeling_vibevoice import (
|
||||
VibeVoiceCausalLMOutputWithPast,
|
||||
SpeechConnector
|
||||
)
|
||||
|
||||
logger = logging.get_logger(__name__)
|
||||
|
||||
if not hasattr(modeling_utils, "ALL_PARALLEL_STYLES") or modeling_utils.ALL_PARALLEL_STYLES is None:
|
||||
modeling_utils.ALL_PARALLEL_STYLES = ["tp", "none", "colwise", "rowwise"]
|
||||
|
||||
# @auto_docstring
|
||||
class VibeVoiceASRPreTrainedModel(PreTrainedModel):
|
||||
config_class = VibeVoiceASRConfig
|
||||
base_model_prefix = "model"
|
||||
supports_gradient_checkpointing = True
|
||||
_skip_keys_device_placement = "past_key_values"
|
||||
_supports_cache_class = True
|
||||
_supports_flash_attn = True
|
||||
_supports_flash_attn_2 = True
|
||||
_supports_sdpa = True
|
||||
_supports_quantized_cache = True
|
||||
_supports_static_cache = True
|
||||
_supports_attention_backend = True
|
||||
|
||||
def _init_weights(self, module):
|
||||
|
||||
# Use the language model's initializer_range if available
|
||||
if hasattr(self.config, 'language_model_config') and hasattr(self.config.language_model_config, 'initializer_range'):
|
||||
std = self.config.language_model_config.initializer_range
|
||||
elif hasattr(self.config, 'decoder_config') and hasattr(self.config.decoder_config, 'initializer_range'):
|
||||
std = self.config.decoder_config.initializer_range
|
||||
else:
|
||||
std = 0.02 # Default value
|
||||
|
||||
if isinstance(module, nn.Linear):
|
||||
module.weight.data.normal_(mean=0.0, std=std)
|
||||
if module.bias is not None:
|
||||
module.bias.data.zero_()
|
||||
elif isinstance(module, nn.LayerNorm):
|
||||
module.weight.data.fill_(1.0)
|
||||
module.bias.data.zero_()
|
||||
|
||||
# @auto_docstring
|
||||
class VibeVoiceASRModel(VibeVoiceASRPreTrainedModel):
|
||||
def __init__(self, config):
|
||||
super().__init__(config)
|
||||
|
||||
if hasattr(config, 'torch_dtype') and config.torch_dtype is not None:
|
||||
if isinstance(config.torch_dtype, str):
|
||||
dtype = getattr(torch, config.torch_dtype)
|
||||
else:
|
||||
dtype = config.torch_dtype
|
||||
else:
|
||||
dtype = torch.float32
|
||||
|
||||
# Initialize Qwen2 model for language modeling
|
||||
lm_config = config.decoder_config
|
||||
self.language_model = AutoModel.from_config(lm_config)
|
||||
|
||||
# Initialize speech components if needed
|
||||
self.acoustic_tokenizer = AutoModel.from_config(config.acoustic_tokenizer_config).to(dtype)
|
||||
self.semantic_tokenizer = AutoModel.from_config(config.semantic_tokenizer_config).to(dtype)
|
||||
|
||||
self.acoustic_connector = SpeechConnector(config.acoustic_vae_dim, lm_config.hidden_size).to(dtype)
|
||||
self.semantic_connector = SpeechConnector(config.semantic_vae_dim, lm_config.hidden_size).to(dtype)
|
||||
|
||||
def get_input_embeddings(self):
|
||||
if hasattr(self.language_model, 'embed_tokens'):
|
||||
# If the language model has an embed_tokens attribute, return it
|
||||
return self.language_model.embed_tokens
|
||||
|
||||
for name, attr in self.language_model.fullmap.items(): # parallel by nnscaler, the name is changed
|
||||
if attr.orig_name == 'embed_tokens.weight':
|
||||
return getattr(self.language_model, name)
|
||||
assert False, 'should not arrive here'
|
||||
|
||||
def set_input_embeddings(self, value):
|
||||
self.language_model.embed_tokens = value
|
||||
|
||||
def set_speech_tokenizers(self, acoustic_tokenizer=None, semantic_tokenizer=None):
|
||||
"""Set the speech tokenizers used for encoding and decoding speech."""
|
||||
self.acoustic_tokenizer = acoustic_tokenizer
|
||||
self.semantic_tokenizer = semantic_tokenizer
|
||||
|
||||
# Reset the encoder to evaluation mode
|
||||
if self.acoustic_tokenizer is not None:
|
||||
self.acoustic_tokenizer.eval()
|
||||
|
||||
if self.semantic_tokenizer is not None:
|
||||
self.semantic_tokenizer.eval()
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.LongTensor = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
position_ids: Optional[torch.LongTensor] = None,
|
||||
past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
|
||||
inputs_embeds: Optional[torch.FloatTensor] = None,
|
||||
use_cache: Optional[bool] = None,
|
||||
output_attentions: Optional[bool] = None,
|
||||
output_hidden_states: Optional[bool] = None,
|
||||
return_dict: Optional[bool] = None,
|
||||
cache_position: Optional[torch.LongTensor] = None,
|
||||
**kwargs,
|
||||
) -> Union[Tuple, BaseModelOutputWithPast]:
|
||||
|
||||
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
||||
|
||||
# Forward through language model
|
||||
outputs = self.language_model(
|
||||
input_ids=input_ids,
|
||||
attention_mask=attention_mask,
|
||||
position_ids=position_ids,
|
||||
past_key_values=past_key_values,
|
||||
inputs_embeds=inputs_embeds,
|
||||
use_cache=use_cache,
|
||||
output_attentions=output_attentions,
|
||||
output_hidden_states=output_hidden_states,
|
||||
return_dict=return_dict,
|
||||
cache_position=cache_position,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
if not return_dict:
|
||||
return outputs
|
||||
|
||||
return BaseModelOutputWithPast(
|
||||
last_hidden_state=outputs.last_hidden_state,
|
||||
past_key_values=outputs.past_key_values,
|
||||
hidden_states=outputs.hidden_states,
|
||||
attentions=outputs.attentions,
|
||||
)
|
||||
|
||||
class VibeVoiceASRForConditionalGeneration(VibeVoiceASRPreTrainedModel, GenerationMixin):
|
||||
"""
|
||||
VibeVoice model for Automatic Speech Recognition (ASR) with language modeling head for conditional generation.
|
||||
This class is designed for inference and generation tasks.
|
||||
"""
|
||||
_tied_weights_keys = ["lm_head.weight"]
|
||||
_tp_plan = {"lm_head": "colwise_rep"}
|
||||
|
||||
def __init__(self, config):
|
||||
super().__init__(config)
|
||||
self.model = VibeVoiceASRModel(config)
|
||||
self.vocab_size = config.decoder_config.vocab_size
|
||||
|
||||
# Determine the dtype to use
|
||||
if hasattr(config, 'torch_dtype') and config.torch_dtype is not None:
|
||||
if isinstance(config.torch_dtype, str):
|
||||
dtype = getattr(torch, config.torch_dtype)
|
||||
else:
|
||||
dtype = config.torch_dtype
|
||||
else:
|
||||
dtype = torch.float32
|
||||
|
||||
# Initialize lm_head with the correct dtype
|
||||
self.lm_head = nn.Linear(config.decoder_config.hidden_size, self.vocab_size, bias=False).to(dtype)
|
||||
|
||||
# Initialize weights and apply final processing
|
||||
self.post_init()
|
||||
|
||||
def get_input_embeddings(self):
|
||||
return self.model.get_input_embeddings()
|
||||
|
||||
def set_input_embeddings(self, value):
|
||||
self.model.set_input_embeddings(value)
|
||||
|
||||
def get_output_embeddings(self):
|
||||
return self.lm_head
|
||||
|
||||
def set_output_embeddings(self, new_embeddings):
|
||||
self.lm_head = new_embeddings
|
||||
|
||||
def set_decoder(self, decoder):
|
||||
self.model.language_model = decoder
|
||||
|
||||
def get_decoder(self):
|
||||
return self.model.language_model
|
||||
|
||||
def tie_weights(self):
|
||||
"""Tie the weights between the input embeddings and the output embeddings."""
|
||||
if getattr(self.config.decoder_config, 'tie_word_embeddings', False):
|
||||
output_embeddings = self.get_output_embeddings()
|
||||
input_embeddings = self.get_input_embeddings()
|
||||
if hasattr(input_embeddings, 'weight'):
|
||||
output_embeddings.weight = input_embeddings.weight
|
||||
else:
|
||||
output_embeddings.weight = input_embeddings
|
||||
|
||||
def encode_speech(
|
||||
self,
|
||||
speech_tensors: torch.FloatTensor,
|
||||
speech_masks: Optional[torch.BoolTensor] = None,
|
||||
speech_semantic_tensors: Optional[torch.FloatTensor] = None,
|
||||
streaming_segment_duration: float = 60.0, # seconds
|
||||
):
|
||||
"""
|
||||
Encode speech input into features that can be used by the language model.
|
||||
This method is called once before generation to process the speech input.
|
||||
|
||||
For long audio (>600s by default), uses streaming processing to avoid conv overflow (>2^32).
|
||||
Segments are processed independently, then concatenated before final sampling.
|
||||
|
||||
Args:
|
||||
speech_tensors: Input audio tensor [batch_size, samples]
|
||||
speech_masks: Optional mask for speech features
|
||||
speech_semantic_tensors: Optional pre-computed semantic tokens
|
||||
streaming_segment_duration: Segment duration in seconds for streaming processing (default: 60s)
|
||||
"""
|
||||
if hasattr(self.config, 'torch_dtype') and self.config.torch_dtype is not None:
|
||||
if isinstance(self.config.torch_dtype, str):
|
||||
dtype = getattr(torch, self.config.torch_dtype)
|
||||
else:
|
||||
dtype = self.config.torch_dtype
|
||||
else:
|
||||
dtype = torch.float32
|
||||
|
||||
speech_tensors = speech_tensors.to(dtype)
|
||||
|
||||
# Ensure proper shape: (batch, samples)
|
||||
if speech_tensors.ndim == 1:
|
||||
speech_tensors = speech_tensors.unsqueeze(0)
|
||||
|
||||
batch_size, total_samples = speech_tensors.shape
|
||||
sample_rate = 24000 # fix 24kHz sample rate
|
||||
|
||||
# Calculate segment size in samples
|
||||
segment_samples = int(streaming_segment_duration * sample_rate)
|
||||
|
||||
# Decide whether to use streaming based on audio length
|
||||
use_streaming = total_samples > segment_samples
|
||||
|
||||
with torch.no_grad():
|
||||
if not use_streaming:
|
||||
# Short audio: direct processing (original behavior)
|
||||
encoder_output = self.model.acoustic_tokenizer.encode(speech_tensors.unsqueeze(1))
|
||||
audio_tokens = encoder_output.sample(dist_type=self.model.acoustic_tokenizer.std_dist_type)[0]
|
||||
acoustic_features = self.model.acoustic_connector(audio_tokens)
|
||||
|
||||
# Encode semantic features
|
||||
if speech_semantic_tensors is not None:
|
||||
semantic_features = self.model.semantic_connector(speech_semantic_tensors)
|
||||
else:
|
||||
semantic_tokens = self.model.semantic_tokenizer.encode(speech_tensors.unsqueeze(1)).mean
|
||||
semantic_features = self.model.semantic_connector(semantic_tokens)
|
||||
else:
|
||||
# Long audio: streaming processing
|
||||
# print(f"Using streaming processing for long audio: {total_samples/sample_rate:.1f}s "
|
||||
# f"(segment size: {streaming_segment_duration}s)")
|
||||
|
||||
# Initialize caches for both tokenizers
|
||||
acoustic_encoder_cache = VibeVoiceTokenizerStreamingCache()
|
||||
semantic_encoder_cache = VibeVoiceTokenizerStreamingCache()
|
||||
acoustic_mean_segments = []
|
||||
semantic_mean_segments = []
|
||||
sample_indices = torch.arange(batch_size, device=speech_tensors.device)
|
||||
|
||||
# Helper function from batch_asr_sft_cache.py
|
||||
def _iter_segments(total_length: int, segment_length: int):
|
||||
"""Iterate over audio segments with a given segment length."""
|
||||
if segment_length <= 0:
|
||||
raise ValueError("segment_length must be positive")
|
||||
for start in range(0, total_length, segment_length):
|
||||
end = min(start + segment_length, total_length)
|
||||
if end > start:
|
||||
yield start, end
|
||||
|
||||
# Process each segment for both acoustic and semantic tokenizers
|
||||
segments = list(_iter_segments(total_samples, segment_samples))
|
||||
num_segments = len(segments)
|
||||
for seg_idx, (start, end) in enumerate(segments):
|
||||
chunk = speech_tensors[:, start:end].contiguous()
|
||||
if chunk.numel() == 0:
|
||||
continue
|
||||
|
||||
# Check if this is the final segment
|
||||
is_final = (seg_idx == num_segments - 1)
|
||||
|
||||
# Encode chunk for acoustic tokenizer (don't sample yet)
|
||||
acoustic_encoder_output = self.model.acoustic_tokenizer.encode(
|
||||
chunk.unsqueeze(1),
|
||||
cache=acoustic_encoder_cache,
|
||||
sample_indices=sample_indices,
|
||||
use_cache=True,
|
||||
is_final_chunk=is_final,
|
||||
)
|
||||
acoustic_mean_segments.append(acoustic_encoder_output.mean)
|
||||
|
||||
# Encode chunk for semantic tokenizer (take mean directly)
|
||||
semantic_encoder_output = self.model.semantic_tokenizer.encode(
|
||||
chunk.unsqueeze(1),
|
||||
cache=semantic_encoder_cache,
|
||||
sample_indices=sample_indices,
|
||||
use_cache=True,
|
||||
is_final_chunk=is_final,
|
||||
)
|
||||
semantic_mean_segments.append(semantic_encoder_output.mean)
|
||||
|
||||
# print(f"Processed {len(acoustic_mean_segments)} segments.")
|
||||
# Concatenate all acoustic means and sample once
|
||||
acoustic_mean_full = torch.cat(acoustic_mean_segments, dim=1).contiguous()
|
||||
acoustic_encoder_output = VibeVoiceTokenizerEncoderOutput(
|
||||
mean=acoustic_mean_full,
|
||||
std=self.model.acoustic_tokenizer.fix_std
|
||||
)
|
||||
audio_tokens = acoustic_encoder_output.sample(
|
||||
dist_type=self.model.acoustic_tokenizer.std_dist_type
|
||||
)[0]
|
||||
acoustic_features = self.model.acoustic_connector(audio_tokens)
|
||||
|
||||
# Concatenate all semantic means
|
||||
semantic_tokens = torch.cat(semantic_mean_segments, dim=1).contiguous()
|
||||
semantic_features = self.model.semantic_connector(semantic_tokens)
|
||||
|
||||
# Combine acoustic and semantic features
|
||||
if speech_masks is not None:
|
||||
combined_features = acoustic_features[speech_masks] + semantic_features[speech_masks]
|
||||
else:
|
||||
combined_features = acoustic_features + semantic_features
|
||||
|
||||
return combined_features
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: Optional[torch.LongTensor] = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
position_ids: Optional[torch.LongTensor] = None,
|
||||
past_key_values: Optional[List[torch.FloatTensor]] = None,
|
||||
inputs_embeds: Optional[torch.FloatTensor] = None,
|
||||
labels: Optional[torch.LongTensor] = None,
|
||||
use_cache: Optional[bool] = None,
|
||||
output_attentions: Optional[bool] = None,
|
||||
output_hidden_states: Optional[bool] = None,
|
||||
return_dict: Optional[bool] = None,
|
||||
cache_position: Optional[torch.LongTensor] = None,
|
||||
# Speech-specific arguments
|
||||
speech_tensors: Optional[torch.FloatTensor] = None,
|
||||
speech_masks: Optional[torch.BoolTensor] = None,
|
||||
speech_semantic_tensors: Optional[torch.FloatTensor] = None,
|
||||
acoustic_input_mask: Optional[torch.BoolTensor] = None,
|
||||
**kwargs,
|
||||
) -> Union[Tuple, CausalLMOutput]:
|
||||
"""
|
||||
Forward pass for the model. Handles both training and generation scenarios.
|
||||
"""
|
||||
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
||||
output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
|
||||
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
||||
use_cache = use_cache if use_cache is not None else self.config.use_cache
|
||||
|
||||
# Process inputs
|
||||
if inputs_embeds is None and input_ids is not None:
|
||||
inputs_embeds = self.get_input_embeddings()(input_ids)
|
||||
|
||||
# If we have speech input and acoustic_input_mask, encode and insert speech features
|
||||
if speech_tensors is not None and acoustic_input_mask is not None:
|
||||
speech_features = self.encode_speech(
|
||||
speech_tensors=speech_tensors,
|
||||
speech_masks=speech_masks,
|
||||
speech_semantic_tensors=speech_semantic_tensors,
|
||||
)
|
||||
inputs_embeds[acoustic_input_mask] = speech_features
|
||||
|
||||
# Forward through the model
|
||||
outputs = self.model(
|
||||
input_ids=None,
|
||||
attention_mask=attention_mask,
|
||||
position_ids=position_ids,
|
||||
past_key_values=past_key_values,
|
||||
inputs_embeds=inputs_embeds,
|
||||
use_cache=use_cache,
|
||||
output_attentions=output_attentions,
|
||||
output_hidden_states=output_hidden_states,
|
||||
return_dict=return_dict,
|
||||
cache_position=cache_position,
|
||||
)
|
||||
|
||||
hidden_states = outputs[0] if not return_dict else outputs.last_hidden_state
|
||||
logits = self.lm_head(hidden_states)
|
||||
|
||||
loss = None
|
||||
if labels is not None:
|
||||
# Shift so that tokens < n predict n
|
||||
shift_logits = logits[..., :-1, :].contiguous()
|
||||
shift_labels = labels[..., 1:].contiguous()
|
||||
# Flatten the tokens
|
||||
loss_fct = nn.CrossEntropyLoss()
|
||||
shift_logits = shift_logits.view(-1, self.vocab_size)
|
||||
shift_labels = shift_labels.view(-1)
|
||||
# Enable model parallelism
|
||||
shift_labels = shift_labels.to(shift_logits.device)
|
||||
loss = loss_fct(shift_logits, shift_labels)
|
||||
|
||||
if not return_dict:
|
||||
output = (logits,) + outputs[1:]
|
||||
return (loss,) + output if loss is not None else output
|
||||
|
||||
return VibeVoiceCausalLMOutputWithPast(
|
||||
loss=loss,
|
||||
logits=logits,
|
||||
past_key_values=outputs.past_key_values,
|
||||
hidden_states=outputs.hidden_states,
|
||||
attentions=outputs.attentions,
|
||||
)
|
||||
|
||||
def prepare_inputs_for_generation(
|
||||
self,
|
||||
input_ids,
|
||||
past_key_values=None,
|
||||
attention_mask=None,
|
||||
inputs_embeds=None,
|
||||
cache_position=None,
|
||||
position_ids=None,
|
||||
use_cache=True,
|
||||
speech_tensors=None,
|
||||
speech_masks=None,
|
||||
speech_semantic_tensors=None,
|
||||
acoustic_input_mask=None,
|
||||
**kwargs,
|
||||
):
|
||||
"""
|
||||
Prepare inputs for generation step. This method is called by generate()
|
||||
for each token generation step.
|
||||
|
||||
Following Qwen2-VL's approach: speech inputs are only forwarded on the first pass
|
||||
(when cache_position[0] == 0), and are excluded in subsequent generation steps.
|
||||
"""
|
||||
# If we have past key values, we only need to process the new tokens
|
||||
if past_key_values is not None:
|
||||
if isinstance(past_key_values, tuple):
|
||||
past_length = past_key_values[0][0].shape[2]
|
||||
else:
|
||||
past_length = past_key_values.get_seq_length()
|
||||
|
||||
# Keep only the new tokens
|
||||
if input_ids is not None and input_ids.shape[1] > past_length:
|
||||
input_ids = input_ids[:, past_length:]
|
||||
|
||||
# Prepare position ids
|
||||
if position_ids is None and attention_mask is not None:
|
||||
position_ids = attention_mask.long().cumsum(-1) - 1
|
||||
position_ids.masked_fill_(attention_mask == 0, 1)
|
||||
if past_key_values is not None and input_ids is not None:
|
||||
position_ids = position_ids[:, -input_ids.shape[1]:]
|
||||
|
||||
# Prepare cache position
|
||||
if cache_position is None:
|
||||
past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
|
||||
cache_position = torch.arange(
|
||||
past_seen_tokens,
|
||||
past_seen_tokens + (input_ids.shape[1] if input_ids is not None else inputs_embeds.shape[1]),
|
||||
device=input_ids.device if input_ids is not None else inputs_embeds.device
|
||||
)
|
||||
|
||||
# Prepare model inputs
|
||||
if inputs_embeds is not None and past_key_values is None:
|
||||
model_inputs = {"inputs_embeds": inputs_embeds}
|
||||
else:
|
||||
model_inputs = {"input_ids": input_ids}
|
||||
|
||||
model_inputs.update(
|
||||
{
|
||||
"position_ids": position_ids,
|
||||
"cache_position": cache_position,
|
||||
"past_key_values": past_key_values,
|
||||
"use_cache": use_cache,
|
||||
"attention_mask": attention_mask,
|
||||
}
|
||||
)
|
||||
|
||||
# Following Qwen2-VL pattern: only include speech inputs on the first forward pass
|
||||
# (when cache_position[0] == 0), exclude them in subsequent generation steps
|
||||
if cache_position is not None and len(cache_position) > 0 and cache_position[0] == 0:
|
||||
# First forward pass - include speech inputs if provided
|
||||
model_inputs.update({
|
||||
"speech_tensors": speech_tensors,
|
||||
"speech_masks": speech_masks,
|
||||
"speech_semantic_tensors": speech_semantic_tensors,
|
||||
"acoustic_input_mask": acoustic_input_mask,
|
||||
})
|
||||
else:
|
||||
# Subsequent generation steps - exclude speech inputs
|
||||
model_inputs.update({
|
||||
"speech_tensors": None,
|
||||
"speech_masks": None,
|
||||
"speech_semantic_tensors": None,
|
||||
"acoustic_input_mask": None,
|
||||
})
|
||||
|
||||
# Include any remaining kwargs that might be needed
|
||||
model_inputs.update(kwargs)
|
||||
|
||||
return model_inputs
|
||||
|
||||
AutoModel.register(VibeVoiceASRConfig, VibeVoiceASRModel)
|
||||
AutoModelForCausalLM.register(VibeVoiceASRConfig, VibeVoiceASRForConditionalGeneration)
|
||||
|
||||
__all__ = [
|
||||
"VibeVoiceASRPreTrainedModel",
|
||||
"VibeVoiceASRModel",
|
||||
"VibeVoiceASRForConditionalGeneration",
|
||||
]
|
||||
@@ -207,8 +207,107 @@ class VibeVoiceTextTokenizerFast(Qwen2TokenizerFast):
|
||||
"""Id used for padding (returns -100 for loss masking)."""
|
||||
return self._pad_id
|
||||
|
||||
class VibeVoiceASRTextTokenizerFast(Qwen2TokenizerFast):
|
||||
"""
|
||||
Construct a "fast" VibeVoice tokenizer (backed by HuggingFace's *tokenizers* library).
|
||||
Based on the Qwen2 tokenizer with additional special tokens for speech.
|
||||
|
||||
Args:
|
||||
vocab_file (`str`, *optional*):
|
||||
Path to the vocabulary file.
|
||||
merges_file (`str`, *optional*):
|
||||
Path to the merges file.
|
||||
tokenizer_file (`str`, *optional*):
|
||||
Path to [tokenizers](https://github.com/huggingface/tokenizers) file.
|
||||
unk_token (`str`, *optional*, defaults to `"<|endoftext|>"`):
|
||||
The unknown token.
|
||||
bos_token (`str`, *optional*):
|
||||
The beginning of sequence token. Not used for vibevoice.
|
||||
eos_token (`str`, *optional*, defaults to `"<|endoftext|>"`):
|
||||
The end of sequence token.
|
||||
pad_token (`str`, *optional*, defaults to `"<|endoftext|>"`):
|
||||
The token used for padding.
|
||||
"""
|
||||
|
||||
model_input_names = ["input_ids", "attention_mask"]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
vocab_file=None,
|
||||
merges_file=None,
|
||||
tokenizer_file=None,
|
||||
unk_token="<|endoftext|>",
|
||||
bos_token=None,
|
||||
eos_token="<|endoftext|>",
|
||||
pad_token="<|endoftext|>",
|
||||
add_prefix_space=False,
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__(
|
||||
vocab_file=vocab_file,
|
||||
merges_file=merges_file,
|
||||
tokenizer_file=tokenizer_file,
|
||||
unk_token=unk_token,
|
||||
bos_token=bos_token,
|
||||
eos_token=eos_token,
|
||||
pad_token=pad_token,
|
||||
add_prefix_space=add_prefix_space,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
# Add VibeVoice-specific special tokens
|
||||
self._add_vibevoice_special_tokens()
|
||||
|
||||
# https://github.com/QwenLM/Qwen2.5-VL/blob/d2240f11656bfe404b9ba56db4e51cd09f522ff1/qwen-vl-finetune/qwenvl/data/data_qwen_packed.py#L57C5-L57C222
|
||||
self.chat_template = "{% for message in messages %}{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}"
|
||||
|
||||
def _add_vibevoice_special_tokens(self):
|
||||
"""Add VibeVoice-specific special tokens."""
|
||||
special_tokens = {
|
||||
"additional_special_tokens": [
|
||||
"<|object_ref_start|>", # Speech start (reusing vision tokens)
|
||||
"<|object_ref_end|>", # Speech end
|
||||
"<|box_start|>", # Speech diffusion pad
|
||||
]
|
||||
}
|
||||
num_added = self.add_special_tokens(special_tokens)
|
||||
|
||||
# Cache special token IDs
|
||||
self._speech_start_id = self.convert_tokens_to_ids("<|object_ref_start|>")
|
||||
self._speech_end_id = self.convert_tokens_to_ids("<|object_ref_end|>")
|
||||
self._speech_pad_id = self.convert_tokens_to_ids("<|box_start|>")
|
||||
|
||||
self._eos_id = self.eos_token_id # qwen2 / qwen3
|
||||
self._pad_id = self.convert_tokens_to_ids('<|image_pad|>')
|
||||
|
||||
return num_added
|
||||
|
||||
@property
|
||||
def eos_id(self) -> int:
|
||||
"""Id of the end of sequence token."""
|
||||
return self._eos_id
|
||||
|
||||
@property
|
||||
def speech_start_id(self) -> int:
|
||||
"""Id of the speech start token."""
|
||||
return self._speech_start_id
|
||||
|
||||
@property
|
||||
def speech_end_id(self) -> int:
|
||||
"""Id of the speech end token."""
|
||||
return self._speech_end_id
|
||||
|
||||
@property
|
||||
def speech_pad_id(self) -> int:
|
||||
"""Id of the speech diffusion token."""
|
||||
return self._speech_pad_id
|
||||
|
||||
@property
|
||||
def pad_id(self) -> int:
|
||||
return self._pad_id
|
||||
|
||||
__all__ = [
|
||||
"VibeVoiceTextTokenizer",
|
||||
"VibeVoiceTextTokenizerFast",
|
||||
"VibeVoiceASRTextTokenizerFast",
|
||||
]
|
||||
@@ -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",
|
||||
]
|
||||
Reference in New Issue
Block a user