microsoft--vibevoice
wehub-resource-sync
907 行
42 KiB
Python
907 行
42 KiB
Python
from dataclasses import dataclass
|
|
from typing import Any, Dict, List, Optional, Tuple, Union, Callable
|
|
from tqdm import tqdm
|
|
import inspect
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
from transformers.models.auto import AutoModel, AutoModelForCausalLM
|
|
from transformers.generation import GenerationMixin, GenerationConfig, LogitsProcessor, LogitsProcessorList, StoppingCriteriaList
|
|
from transformers.modeling_outputs import BaseModelOutputWithPast, ModelOutput
|
|
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
|
|
from .modular_vibevoice_diffusion_head import VibeVoiceDiffusionHead
|
|
from vibevoice.schedule.dpm_solver import DPMSolverMultistepScheduler
|
|
from .configuration_vibevoice_streaming import VibeVoiceStreamingConfig
|
|
from .modular_vibevoice_text_tokenizer import VibeVoiceTextTokenizer, VibeVoiceTextTokenizerFast
|
|
from .modeling_vibevoice_streaming import VibeVoiceStreamingPreTrainedModel, VibeVoiceStreamingModel, BinaryClassifier
|
|
from .streamer import AudioStreamer, AsyncAudioStreamer
|
|
|
|
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"]
|
|
|
|
TTS_TEXT_WINDOW_SIZE = 5
|
|
TTS_SPEECH_WINDOW_SIZE = 6
|
|
|
|
|
|
# ============================================================================
|
|
# Transformers >= 4.57 Compatibility Layer
|
|
# The cache system was refactored in transformers 4.57, requiring these helpers.
|
|
# ============================================================================
|
|
|
|
class MockCacheLayer:
|
|
"""
|
|
Mock cache layer for transformers >= 4.57 compatibility.
|
|
Provides the `layers` interface expected by DynamicCache in newer versions.
|
|
"""
|
|
|
|
def __init__(self, key_cache, value_cache, parent_cache=None, layer_idx=0):
|
|
self.key_cache = key_cache
|
|
self.value_cache = value_cache
|
|
self._parent_cache = parent_cache
|
|
self._layer_idx = layer_idx
|
|
|
|
def get_mask_sizes(self, cache_position):
|
|
"""Return KV length and offset for mask creation."""
|
|
seq_length = self.key_cache.shape[2] if self.key_cache is not None else 0
|
|
query_length = cache_position.shape[0]
|
|
return seq_length + query_length, 0
|
|
|
|
def update(self, key_states, value_states, cache_kwargs=None):
|
|
"""Update the cache with new key/value states."""
|
|
if self._parent_cache is None:
|
|
return self.key_cache, self.value_cache
|
|
|
|
parent = self._parent_cache
|
|
idx = self._layer_idx
|
|
|
|
# Extend cache lists if needed
|
|
while len(parent.key_cache) <= idx:
|
|
parent.key_cache.append(None)
|
|
parent.value_cache.append(None)
|
|
|
|
# Concatenate or initialize cache
|
|
if parent.key_cache[idx] is not None:
|
|
parent.key_cache[idx] = torch.cat([parent.key_cache[idx], key_states], dim=2)
|
|
parent.value_cache[idx] = torch.cat([parent.value_cache[idx], value_states], dim=2)
|
|
else:
|
|
parent.key_cache[idx] = key_states
|
|
parent.value_cache[idx] = value_states
|
|
|
|
# Update local references
|
|
self.key_cache = parent.key_cache[idx]
|
|
self.value_cache = parent.value_cache[idx]
|
|
return self.key_cache, self.value_cache
|
|
|
|
|
|
def _ensure_cache_has_layers(cache):
|
|
"""
|
|
Ensure the cache has all required attributes for transformers >= 4.57.
|
|
Creates MockCacheLayer wrappers to provide the expected `layers` interface.
|
|
"""
|
|
if cache is None:
|
|
return cache
|
|
|
|
# Add required attributes (skip if read-only)
|
|
for attr, default in [('layer_class_to_replicate', None), ('offloading', False), ('is_compileable', False)]:
|
|
if not hasattr(cache, attr):
|
|
try:
|
|
setattr(cache, attr, default)
|
|
except AttributeError:
|
|
pass
|
|
|
|
# Build layers list from key_cache/value_cache
|
|
if hasattr(cache, 'key_cache') and hasattr(cache, 'value_cache'):
|
|
try:
|
|
cache.layers = [
|
|
MockCacheLayer(cache.key_cache[i], cache.value_cache[i], parent_cache=cache, layer_idx=i)
|
|
for i in range(len(cache.key_cache))
|
|
]
|
|
except AttributeError:
|
|
pass
|
|
elif not hasattr(cache, 'layers'):
|
|
try:
|
|
cache.layers = []
|
|
except AttributeError:
|
|
pass
|
|
|
|
return cache
|
|
|
|
|
|
def _update_model_kwargs_for_generation(
|
|
outputs: ModelOutput,
|
|
model_kwargs: Dict[str, Any],
|
|
num_new_tokens: int = 1,
|
|
) -> Dict[str, Any]:
|
|
"""
|
|
Update model_kwargs after adding new tokens (supports multi-token windows).
|
|
|
|
Updates past_key_values, attention_mask, and cache_position for the next forward pass.
|
|
"""
|
|
model_kwargs["past_key_values"] = _ensure_cache_has_layers(outputs.past_key_values)
|
|
|
|
attention_mask = model_kwargs["attention_mask"]
|
|
model_kwargs["attention_mask"] = torch.cat(
|
|
[attention_mask, attention_mask.new_ones((attention_mask.shape[0], num_new_tokens))], dim=-1
|
|
)
|
|
|
|
cache_pos = model_kwargs["cache_position"]
|
|
model_kwargs["cache_position"] = torch.arange(
|
|
cache_pos[-1] + 1, cache_pos[-1] + num_new_tokens + 1, device=cache_pos.device
|
|
)
|
|
|
|
return model_kwargs
|
|
|
|
|
|
@dataclass
|
|
class VibeVoiceCausalLMOutputWithPast(BaseModelOutputWithPast):
|
|
logits: Optional[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
|
|
reach_max_step_sample: Optional[torch.BoolTensor] = None
|
|
|
|
|
|
class VibeVoiceStreamingForConditionalGenerationInference(VibeVoiceStreamingPreTrainedModel, GenerationMixin):
|
|
|
|
def __init__(self, config):
|
|
super().__init__(config)
|
|
|
|
# Initialize the base model
|
|
self.model = VibeVoiceStreamingModel(config)
|
|
|
|
# TTS generation EOS classifier
|
|
self.tts_eos_classifier = BinaryClassifier(config.decoder_config.hidden_size)
|
|
|
|
# inference configuration
|
|
self.ddpm_inference_steps = config.diffusion_head_config.ddpm_num_inference_steps
|
|
|
|
# Initialize weights and apply final processing
|
|
self.post_init()
|
|
|
|
@property
|
|
def noise_scheduler(self):
|
|
return self.model.noise_scheduler
|
|
|
|
@property
|
|
def prediction_head(self):
|
|
return self.model.prediction_head
|
|
|
|
@property
|
|
def speech_scaling_factor(self):
|
|
return self.model.speech_scaling_factor
|
|
|
|
@property
|
|
def speech_bias_factor(self):
|
|
return self.model.speech_bias_factor
|
|
|
|
@property
|
|
def acoustic_tokenizer(self):
|
|
return self.model.acoustic_tokenizer
|
|
|
|
@property
|
|
def acoustic_connector(self):
|
|
return self.model.acoustic_connector
|
|
|
|
def tie_weights(self):
|
|
"""
|
|
Tie the weights between the input embeddings and the output embeddings.
|
|
"""
|
|
# Tie lm_head.weight to language_model.embed_tokens.weight
|
|
if not getattr(self.config, 'tie_word_embeddings', False):
|
|
return
|
|
|
|
if hasattr(self, 'lm_head') and hasattr(self.model.language_model, 'embed_tokens'):
|
|
self.lm_head.weight = self.model.language_model.embed_tokens.weight
|
|
|
|
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):
|
|
"""
|
|
This model does not define an `lm_head` (vocabulary projection).
|
|
"""
|
|
return None
|
|
|
|
def set_output_embeddings(self, new_embeddings):
|
|
"""
|
|
No-op because there is no `lm_head`. Provided only to satisfy optional API calls.
|
|
To enable, first create `self.lm_head` then allow assignment.
|
|
"""
|
|
raise RuntimeError("Output embeddings (lm_head) are not defined for this model. "
|
|
"Create one before calling set_output_embeddings if needed.")
|
|
|
|
def set_speech_tokenizers(self, acoustic_tokenizer=None):
|
|
"""Set the speech tokenizers used for encoding and decoding speech."""
|
|
self.model.set_speech_tokenizers(acoustic_tokenizer)
|
|
|
|
def set_ddpm_inference_steps(self, num_steps=None):
|
|
self.ddpm_inference_steps = num_steps or self.config.diffusion_head_config.ddpm_num_inference_steps
|
|
|
|
def prepare_inputs_for_generation(
|
|
self,
|
|
input_ids: torch.LongTensor,
|
|
past_key_values=None,
|
|
attention_mask=None,
|
|
inputs_embeds=None,
|
|
cache_position=None,
|
|
**kwargs,
|
|
):
|
|
"""Prepare model inputs for generation (transformers >= 4.57 compatible)."""
|
|
model_inputs = {"cache_position": cache_position}
|
|
|
|
# Slice inputs when using cache
|
|
if past_key_values is not None:
|
|
model_inputs["past_key_values"] = past_key_values
|
|
if inputs_embeds is not None and input_ids.shape[1] == 0:
|
|
inputs_embeds = inputs_embeds[:, -cache_position.shape[0]:]
|
|
elif inputs_embeds is not None or (cache_position is not None and cache_position[-1] >= input_ids.shape[1]):
|
|
input_ids = input_ids[:, -cache_position.shape[0]:]
|
|
elif cache_position is not None and input_ids.shape[1] != cache_position.shape[0]:
|
|
input_ids = input_ids[:, cache_position]
|
|
|
|
# Set input_ids or inputs_embeds
|
|
use_embeds = inputs_embeds is not None and (
|
|
past_key_values is None or (cache_position is not None and len(cache_position) == inputs_embeds.shape[1])
|
|
)
|
|
if use_embeds:
|
|
model_inputs["input_ids"] = None
|
|
model_inputs["inputs_embeds"] = inputs_embeds
|
|
else:
|
|
model_inputs["input_ids"] = input_ids.clone(memory_format=torch.contiguous_format) if input_ids is not None else None
|
|
model_inputs["inputs_embeds"] = None
|
|
|
|
if attention_mask is not None:
|
|
model_inputs["attention_mask"] = attention_mask
|
|
|
|
# Create position_ids from attention_mask
|
|
if attention_mask is not None and kwargs.get("position_ids") is None:
|
|
position_ids = attention_mask.long().cumsum(-1) - 1
|
|
position_ids.masked_fill_(attention_mask == 0, 1)
|
|
kwargs["position_ids"] = position_ids
|
|
|
|
# Slice position_ids when using cache
|
|
if kwargs.get("position_ids") is not None:
|
|
if past_key_values is not None:
|
|
seq_len = model_inputs["inputs_embeds"].shape[1] if model_inputs.get("inputs_embeds") is not None else model_inputs["input_ids"].shape[1]
|
|
model_inputs["position_ids"] = kwargs["position_ids"][:, -seq_len:].clone(memory_format=torch.contiguous_format)
|
|
else:
|
|
model_inputs["position_ids"] = kwargs.pop("position_ids").clone(memory_format=torch.contiguous_format)
|
|
|
|
# Forward remaining kwargs
|
|
for key, value in kwargs.items():
|
|
if key not in model_inputs:
|
|
model_inputs[key] = value
|
|
|
|
model_inputs.pop("labels", None)
|
|
return model_inputs
|
|
|
|
def _update_model_kwargs_for_generation(
|
|
self,
|
|
outputs,
|
|
model_kwargs,
|
|
is_encoder_decoder=False,
|
|
num_new_tokens=1,
|
|
):
|
|
"""Override to ensure cache compatibility with transformers >= 4.57."""
|
|
model_kwargs = super()._update_model_kwargs_for_generation(
|
|
outputs, model_kwargs, is_encoder_decoder=is_encoder_decoder, num_new_tokens=num_new_tokens
|
|
)
|
|
if "past_key_values" in model_kwargs:
|
|
model_kwargs["past_key_values"] = _ensure_cache_has_layers(model_kwargs["past_key_values"])
|
|
return model_kwargs
|
|
|
|
def _init_cache_for_generation(self, generation_config, model_kwargs, batch_size, max_cache_length, device):
|
|
"""
|
|
Initialize cache for generation, handling different transformers versions.
|
|
For transformers >= 4.57, returns None to let the model create the cache dynamically.
|
|
"""
|
|
try:
|
|
from transformers.cache_utils import DynamicCache
|
|
sig = inspect.signature(DynamicCache.__init__)
|
|
if 'config' in sig.parameters:
|
|
# transformers >= 4.57: let model handle cache creation
|
|
return None
|
|
else:
|
|
# Older versions: use parent method
|
|
prep_sig = inspect.signature(self._prepare_cache_for_generation)
|
|
if 'device' in prep_sig.parameters:
|
|
self._prepare_cache_for_generation(generation_config, model_kwargs, None, batch_size, max_cache_length, device)
|
|
else:
|
|
self._prepare_cache_for_generation(generation_config, model_kwargs, None, batch_size, max_cache_length)
|
|
return model_kwargs.get("past_key_values")
|
|
except Exception:
|
|
return None
|
|
|
|
# @can_return_tuple
|
|
def forward_lm(
|
|
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,
|
|
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,
|
|
**kwargs,
|
|
) -> Union[Tuple, BaseModelOutputWithPast]:
|
|
"""
|
|
Single pass of the base text LM.
|
|
|
|
- Builds embeddings if `inputs_embeds` not provided.
|
|
- Uses (and returns) `past_key_values` when `use_cache=True`.
|
|
- No loss / no lm_head / no speech logic.
|
|
|
|
Args:
|
|
input_ids: (B, S) token ids.
|
|
attention_mask: (B, S) mask.
|
|
past_key_values: cache from previous steps.
|
|
cache_position: positions for cached tokens.
|
|
labels: unsupported (will raise).
|
|
|
|
Returns:
|
|
BaseModelOutputWithPast with `last_hidden_state` and `past_key_values`.
|
|
"""
|
|
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
|
|
|
# Get embeddings
|
|
if inputs_embeds is None:
|
|
inputs_embeds = self.model.get_input_embeddings()(input_ids)
|
|
|
|
outputs = self.model.language_model(
|
|
inputs_embeds=inputs_embeds,
|
|
attention_mask=attention_mask,
|
|
position_ids=position_ids,
|
|
past_key_values=past_key_values,
|
|
use_cache=use_cache,
|
|
output_attentions=output_attentions,
|
|
output_hidden_states=output_hidden_states,
|
|
return_dict=return_dict,
|
|
cache_position=cache_position,
|
|
**kwargs,
|
|
)
|
|
|
|
hidden_states = outputs[0] if not return_dict else outputs.last_hidden_state
|
|
|
|
if labels is not None:
|
|
raise NotImplementedError("Loss computation is not implemented in this version.")
|
|
|
|
return BaseModelOutputWithPast(
|
|
past_key_values=outputs.past_key_values,
|
|
last_hidden_state=hidden_states,
|
|
attentions=outputs.attentions,
|
|
)
|
|
|
|
# @can_return_tuple
|
|
def forward_tts_lm(
|
|
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,
|
|
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,
|
|
lm_last_hidden_state: Optional[torch.FloatTensor] = None,
|
|
tts_text_masks: Optional[torch.BoolTensor] = None,
|
|
**kwargs,
|
|
) -> Union[Tuple, VibeVoiceCausalLMOutputWithPast]:
|
|
"""
|
|
Single pass of the TTS LM.
|
|
|
|
- Overwrites tail embeddings with `lm_last_hidden_state`.
|
|
- Adds type embedding via `tts_text_masks` (1=text, 0=speech).
|
|
- Predicts EOS from last hidden state (binary classifier).
|
|
- No loss / no full acoustic decoding here.
|
|
|
|
Args:
|
|
input_ids: (B, S) token ids.
|
|
attention_mask: (B, S) mask.
|
|
lm_last_hidden_state: (B, K, H) hidden states to splice into the tail.
|
|
tts_text_masks: (B, 1) mask marking current position as text(1)/speech(0).
|
|
past_key_values: cache from previous TTS steps.
|
|
cache_position: positions for cached tokens.
|
|
labels: unsupported (will raise).
|
|
|
|
Returns:
|
|
VibeVoiceCausalLMOutputWithPast with `logits` (EOS), `last_hidden_state`, `past_key_values`.
|
|
"""
|
|
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
|
|
|
# Get embeddings
|
|
if inputs_embeds is None:
|
|
# Will be replaced with lm_last_hidden_state
|
|
inputs_embeds = self.model.get_input_embeddings()(input_ids)
|
|
|
|
# Replace the last part of inputs_embeds with lm_last_hidden_state
|
|
start_idx = inputs_embeds.shape[1] - lm_last_hidden_state.shape[1]
|
|
inputs_embeds[:, start_idx:, :] = lm_last_hidden_state
|
|
|
|
# Adds type embedding via `tts_text_masks`.
|
|
inputs_embeds = inputs_embeds + self.model.tts_input_types(tts_text_masks.long())
|
|
|
|
outputs = self.model.tts_language_model(
|
|
inputs_embeds=inputs_embeds,
|
|
attention_mask=attention_mask,
|
|
position_ids=position_ids,
|
|
past_key_values=past_key_values,
|
|
use_cache=use_cache,
|
|
output_attentions=output_attentions,
|
|
output_hidden_states=output_hidden_states,
|
|
return_dict=return_dict,
|
|
cache_position=cache_position,
|
|
**kwargs,
|
|
)
|
|
|
|
hidden_states = outputs[0] if not return_dict else outputs.last_hidden_state
|
|
logits = self.tts_eos_classifier(hidden_states[:, -1, :])
|
|
|
|
if labels is not None:
|
|
raise NotImplementedError("Loss computation is not implemented in this version.")
|
|
|
|
return VibeVoiceCausalLMOutputWithPast(
|
|
logits=logits,
|
|
past_key_values=outputs.past_key_values,
|
|
last_hidden_state=hidden_states,
|
|
attentions=outputs.attentions,
|
|
)
|
|
|
|
def forward(self, *args, **kwargs):
|
|
"""
|
|
Unified forward is intentionally disabled.
|
|
|
|
Reasons:
|
|
1. The inference pipeline is staged: base text LM, then TTS LM, plus streaming & diffusion handled in `generate`.
|
|
2. A monolithic call would hide required sequencing (prefill, window stepping, speech diffusion sampling).
|
|
|
|
Use instead:
|
|
- self.forward_lm(...) for a base text LM step (prefill or incremental).
|
|
- self.forward_tts_lm(...) for a single TTS LM step (needs LM hidden states).
|
|
- self.generate(...) for full streaming (text + speech + diffusion + audio assembly).
|
|
|
|
Raises:
|
|
RuntimeError: Always (by design).
|
|
"""
|
|
raise RuntimeError(
|
|
"Unified forward is disabled. Use `forward_lm`, `forward_tts_lm`, or `generate` instead."
|
|
)
|
|
|
|
def _build_generate_config_model_kwargs(self, generation_config, inputs, tokenizer, return_processors=False, **kwargs):
|
|
if generation_config is None:
|
|
generation_config = GenerationConfig(
|
|
bos_token_id=tokenizer.bos_token_id,
|
|
eos_token_id=tokenizer.eos_token_id,
|
|
pad_token_id = tokenizer.pad_token_id
|
|
)
|
|
else:
|
|
generation_config = GenerationConfig(
|
|
**generation_config,
|
|
bos_token_id=tokenizer.bos_token_id,
|
|
eos_token_id=tokenizer.eos_token_id,
|
|
pad_token_id = tokenizer.pad_token_id
|
|
)
|
|
|
|
generation_config, model_kwargs = self._prepare_generation_config(
|
|
generation_config,
|
|
True,
|
|
speech_start_id=tokenizer.speech_start_id,
|
|
speech_end_id=tokenizer.speech_end_id,
|
|
speech_diffusion_id=tokenizer.speech_diffusion_id,
|
|
**kwargs
|
|
)
|
|
generation_config.speech_start_id = tokenizer.speech_start_id
|
|
generation_config.speech_end_id = tokenizer.speech_end_id
|
|
generation_config.speech_diffusion_id = tokenizer.speech_diffusion_id
|
|
|
|
inputs_tensor, model_input_name, model_kwargs = self._prepare_model_inputs(inputs, generation_config.bos_token_id, model_kwargs)
|
|
batch_size = inputs_tensor.shape[0]
|
|
device = self.device
|
|
|
|
self._prepare_special_tokens(generation_config, True, device=device)
|
|
generation_config.use_cache = True
|
|
model_kwargs["use_cache"] = generation_config.use_cache
|
|
input_ids = inputs_tensor.to(self.device)
|
|
|
|
input_ids_length = input_ids.shape[1]
|
|
has_default_max_length = kwargs.get("max_length") is None and generation_config.max_length is not None
|
|
has_default_min_length = kwargs.get("min_length") is None and generation_config.min_length is not None
|
|
generation_config = self._prepare_generated_length(
|
|
generation_config=generation_config,
|
|
has_default_max_length=has_default_max_length,
|
|
has_default_min_length=has_default_min_length,
|
|
model_input_name=model_input_name,
|
|
inputs_tensor=inputs_tensor,
|
|
input_ids_length=input_ids_length,
|
|
)
|
|
|
|
max_cache_length = generation_config.max_length - 1
|
|
# Handle cache initialization for different transformers versions
|
|
model_kwargs["past_key_values"] = self._init_cache_for_generation(
|
|
generation_config, model_kwargs, batch_size, max_cache_length, device
|
|
)
|
|
model_kwargs['cache_position'] = torch.arange(input_ids_length, device=device, dtype=torch.long)
|
|
for k, v in model_kwargs.items():
|
|
if isinstance(v, torch.Tensor):
|
|
model_kwargs[k] = v.to(device=device)
|
|
|
|
if return_processors:
|
|
logits_processor = self._get_logits_processor(
|
|
generation_config=generation_config,
|
|
input_ids_seq_length=input_ids_length,
|
|
encoder_input_ids=inputs_tensor,
|
|
prefix_allowed_tokens_fn=None,
|
|
logits_processor=LogitsProcessorList(),
|
|
device=inputs_tensor.device,
|
|
model_kwargs=model_kwargs,
|
|
)
|
|
|
|
stopping_criteria = self._get_stopping_criteria(generation_config=generation_config, stopping_criteria=StoppingCriteriaList())
|
|
|
|
return generation_config, model_kwargs, input_ids, logits_processor, stopping_criteria
|
|
else:
|
|
return generation_config, model_kwargs, input_ids
|
|
|
|
@torch.no_grad()
|
|
def generate(
|
|
self,
|
|
inputs: Optional[torch.Tensor] = None,
|
|
generation_config: Optional[GenerationConfig] = None,
|
|
logits_processor: Optional[LogitsProcessorList] = None,
|
|
stopping_criteria: Optional[StoppingCriteriaList] = None,
|
|
prefix_allowed_tokens_fn: Optional[Callable[[int, torch.Tensor], List[int]]] = None,
|
|
synced_gpus: Optional[bool] = None,
|
|
assistant_model: Optional["PreTrainedModel"] = None,
|
|
audio_streamer: Optional[Union[AudioStreamer, AsyncAudioStreamer]] = None,
|
|
negative_prompt_ids: Optional[torch.Tensor] = None,
|
|
negative_prompt_attention_mask: Optional[torch.Tensor] = None,
|
|
speech_tensors: Optional[torch.FloatTensor] = None,
|
|
speech_masks: Optional[torch.BoolTensor] = None,
|
|
speech_input_mask: Optional[torch.BoolTensor] = None,
|
|
tts_text_ids: Optional[torch.LongTensor] = None,
|
|
return_speech: bool = True,
|
|
cfg_scale: float = 1.0,
|
|
stop_check_fn: Optional[Callable[[], bool]] = None,
|
|
**kwargs,
|
|
) -> Union[torch.LongTensor, VibeVoiceGenerationOutput]:
|
|
"""
|
|
Text is fed in small windows (dynamic slicing of `tts_text_ids`), which enables streaming text input: you don’t need the full text upfront. After each text window, a loop samples several speech latents (diffusion). The interleaved text encoding + speech generation enables streaming text input and realtime speech output.
|
|
The function only supports batch size = 1 currently.
|
|
|
|
- Windowed text prefill → incremental LM + TTS LM updates.
|
|
- Interleave speech token diffusion sampling (`sample_speech_tokens`).
|
|
- Stops on EOS (binary classifier) or max length / external `stop_check_fn`.
|
|
- Returns final token `sequences` and (optionally) concatenated speech audio.
|
|
|
|
Args (selected):
|
|
tts_text_ids: Full text tokens to stream in windows.
|
|
audio_streamer: If provided, emits audio chunks during generation.
|
|
cfg_scale: Classifier-free guidance scale for speech diffusion.
|
|
return_speech: If False, skips audio decode concatenation.
|
|
stop_check_fn: External early-stop hook (returns True to halt).
|
|
|
|
Returns:
|
|
VibeVoiceGenerationOutput with:
|
|
- sequences: final token ids
|
|
- speech_outputs: list of concatenated audio tensors (or None)
|
|
- reach_max_step_sample: flags for samples stopped by max length
|
|
"""
|
|
# 1. Handle `generation_config` and kwargs that might update it, and validate the `.generate()` call
|
|
tokenizer = kwargs.pop("tokenizer", None)
|
|
neg_text_input_id = tokenizer.convert_tokens_to_ids("<|image_pad|>")
|
|
|
|
tts_lm_input_ids = kwargs.pop("tts_lm_input_ids", None)
|
|
tts_lm_attention_mask = kwargs.pop("tts_lm_attention_mask", None)
|
|
# all_prefilled_outputs: cached prefilled prompt outputs for lm, tts_lm, neg_lm, neg_tts_lm
|
|
all_prefilled_outputs = kwargs.pop("all_prefilled_outputs", None)
|
|
tts_text_ids = tts_text_ids.to(self.device)
|
|
|
|
if kwargs.get('max_new_tokens', None) is None:
|
|
kwargs['max_new_tokens'] = self.config.decoder_config.max_position_embeddings - tts_lm_input_ids.shape[-1]
|
|
|
|
generation_config, model_kwargs, input_ids, logits_processor, stopping_criteria = self._build_generate_config_model_kwargs(
|
|
generation_config, inputs, tokenizer, return_processors=True, **kwargs
|
|
)
|
|
|
|
negative_kwargs = {
|
|
'input_ids': torch.full((kwargs['input_ids'].shape[0], 1), neg_text_input_id, dtype=torch.long, device=kwargs['input_ids'].device),
|
|
'attention_mask': torch.ones((kwargs['input_ids'].shape[0], 1), dtype=torch.long, device=kwargs['input_ids'].device),
|
|
'max_new_tokens': kwargs.get('max_new_tokens', 100)
|
|
}
|
|
negative_generation_config, negative_model_kwargs, negative_input_ids = self._build_generate_config_model_kwargs(
|
|
None, None, tokenizer, return_processors=False, **negative_kwargs
|
|
)
|
|
|
|
tts_lm_kwargs = {
|
|
'input_ids': tts_lm_input_ids,
|
|
'attention_mask': tts_lm_attention_mask,
|
|
'max_new_tokens': kwargs.get('max_new_tokens', 100)
|
|
}
|
|
tts_lm_generation_config, tts_lm_model_kwargs, tts_lm_input_ids = self._build_generate_config_model_kwargs(
|
|
None, None, tokenizer, return_processors=False, **tts_lm_kwargs
|
|
)
|
|
|
|
tts_lm_negative_kwargs = {
|
|
'input_ids': torch.full((kwargs['input_ids'].shape[0], 1), neg_text_input_id, dtype=torch.long, device=kwargs['input_ids'].device),
|
|
'attention_mask': torch.ones((kwargs['input_ids'].shape[0], 1), dtype=torch.long, device=kwargs['input_ids'].device),
|
|
'max_new_tokens': kwargs.get('max_new_tokens', 100)
|
|
}
|
|
tts_lm_negative_generation_config, tts_lm_negative_model_kwargs, tts_lm_negative_input_ids = self._build_generate_config_model_kwargs(
|
|
None, None, tokenizer, return_processors=False, **tts_lm_negative_kwargs
|
|
)
|
|
|
|
acoustic_cache = VibeVoiceTokenizerStreamingCache()
|
|
batch_size = input_ids.shape[0]
|
|
assert batch_size == 1, "Currently only supports batch size == 1"
|
|
device = input_ids.device
|
|
finished_tags = torch.zeros(batch_size, dtype=torch.bool, device=device)
|
|
verbose = kwargs.get("verbose", False)
|
|
|
|
# Initialize audio chunks storage for each sample
|
|
audio_chunks = [[] for _ in range(batch_size)]
|
|
tts_text_window_index = 0
|
|
reach_max_step_sample = torch.zeros(batch_size, dtype=torch.bool, device=device)
|
|
first_text_window_size = TTS_TEXT_WINDOW_SIZE if tts_text_ids.shape[1] >= TTS_TEXT_WINDOW_SIZE else tts_text_ids.shape[1]
|
|
|
|
outputs = all_prefilled_outputs["lm"]
|
|
tts_lm_outputs = all_prefilled_outputs["tts_lm"]
|
|
negative_outputs = all_prefilled_outputs["neg_lm"]
|
|
tts_lm_negative_outputs = all_prefilled_outputs["neg_tts_lm"]
|
|
|
|
model_kwargs = _update_model_kwargs_for_generation(
|
|
outputs, model_kwargs, num_new_tokens=first_text_window_size,
|
|
)
|
|
tts_lm_model_kwargs = _update_model_kwargs_for_generation(
|
|
tts_lm_outputs, tts_lm_model_kwargs, num_new_tokens=first_text_window_size,
|
|
)
|
|
negative_model_kwargs = self._update_model_kwargs_for_generation(
|
|
negative_outputs, negative_model_kwargs, is_encoder_decoder=False,
|
|
)
|
|
tts_lm_negative_model_kwargs = self._update_model_kwargs_for_generation(
|
|
tts_lm_negative_outputs, tts_lm_negative_model_kwargs, is_encoder_decoder=False,
|
|
)
|
|
|
|
step = tts_lm_input_ids.shape[1]
|
|
total_generated_speech_tokens = 0
|
|
total_prefilled_text_tokens = 0
|
|
if kwargs.get("show_progress_bar", True):
|
|
progress_bar = tqdm(
|
|
total=tts_lm_generation_config.max_length,
|
|
desc=f"Prefilled {step} tokens, current step ({step} / {tts_lm_generation_config.max_length})",
|
|
initial=step,
|
|
leave=False
|
|
)
|
|
else:
|
|
progress_bar = None
|
|
|
|
while True:
|
|
# Check for external stop signal
|
|
if stop_check_fn is not None and stop_check_fn():
|
|
if verbose:
|
|
print(f"Generation stopped externally at step {step + 1}")
|
|
# End the audio streamer if it exists
|
|
if audio_streamer is not None:
|
|
audio_streamer.end()
|
|
break
|
|
|
|
# # Check if audio_streamer has been ended (stopped externally)
|
|
# if audio_streamer is not None and hasattr(audio_streamer, 'finished_flags'):
|
|
# if any(audio_streamer.finished_flags):
|
|
# if verbose:
|
|
# print(f"Audio generation stopped externally at step {step + 1}")
|
|
# break
|
|
|
|
if finished_tags.all():
|
|
if hasattr(progress_bar, 'set_description'):
|
|
progress_bar.set_description("Generation complete")
|
|
break
|
|
|
|
cur_input_tts_text_ids = tts_text_ids[:, tts_text_window_index*TTS_TEXT_WINDOW_SIZE:(tts_text_window_index+1)*TTS_TEXT_WINDOW_SIZE]
|
|
next_text_window_size = tts_text_ids[:, (tts_text_window_index+1)*TTS_TEXT_WINDOW_SIZE:(tts_text_window_index+2)*TTS_TEXT_WINDOW_SIZE].shape[1]
|
|
tts_text_window_index += 1
|
|
|
|
if cur_input_tts_text_ids.shape[1] > 0:
|
|
input_ids = torch.cat([input_ids, cur_input_tts_text_ids], dim=-1)
|
|
tts_lm_input_ids = torch.cat([tts_lm_input_ids, cur_input_tts_text_ids], dim=-1)
|
|
|
|
if tts_lm_input_ids.shape[1] > tts_lm_generation_config.max_length:
|
|
if verbose:
|
|
print(f"Reached maximum generation length {generation_config.max_length}, stopped it.")
|
|
reached_samples = torch.arange(batch_size, device=device)[~finished_tags]
|
|
if reached_samples.numel() > 0:
|
|
reach_max_step_sample[reached_samples] = True
|
|
break
|
|
|
|
step += cur_input_tts_text_ids.shape[1]
|
|
total_prefilled_text_tokens += cur_input_tts_text_ids.shape[1]
|
|
if progress_bar is not None:
|
|
progress_bar.update(cur_input_tts_text_ids.shape[1])
|
|
progress_bar.set_description(f"Prefilled {total_prefilled_text_tokens} text tokens, generated {total_generated_speech_tokens} speech tokens, current step ({step} / {tts_lm_generation_config.max_length})")
|
|
|
|
model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)
|
|
# Forward pass through the model
|
|
outputs = self.forward_lm(
|
|
**model_inputs, return_dict=True, output_attentions=False, output_hidden_states=False,
|
|
)
|
|
model_kwargs = _update_model_kwargs_for_generation(
|
|
outputs, model_kwargs, num_new_tokens=next_text_window_size,
|
|
)
|
|
|
|
tts_lm_model_inputs = self.prepare_inputs_for_generation(tts_lm_input_ids, **tts_lm_model_kwargs)
|
|
tts_lm_additional_inputs = {
|
|
"tts_text_masks": torch.ones_like(tts_lm_input_ids[:, -1:]),
|
|
"lm_last_hidden_state": outputs.last_hidden_state,
|
|
}
|
|
# Forward pass through the model
|
|
tts_lm_outputs = self.forward_tts_lm(
|
|
**tts_lm_model_inputs, **tts_lm_additional_inputs, return_dict=True, output_attentions=False, output_hidden_states=False,
|
|
)
|
|
tts_lm_model_kwargs = self._update_model_kwargs_for_generation(
|
|
tts_lm_outputs, tts_lm_model_kwargs, is_encoder_decoder=False,
|
|
)
|
|
|
|
diffusion_indices = torch.LongTensor([0])
|
|
for cur_speech_index in range(TTS_SPEECH_WINDOW_SIZE):
|
|
positive_condition = tts_lm_outputs.last_hidden_state[diffusion_indices, -1, :]
|
|
negative_condition = tts_lm_negative_outputs.last_hidden_state[diffusion_indices, -1, :]
|
|
|
|
speech_latent = self.sample_speech_tokens(
|
|
positive_condition,
|
|
negative_condition,
|
|
cfg_scale=cfg_scale,
|
|
).unsqueeze(1)
|
|
|
|
# Decode acoustic latent to audio using acoustic streaming cache
|
|
scaled_latent = speech_latent / self.model.speech_scaling_factor.to(speech_latent.device) - self.model.speech_bias_factor.to(speech_latent.device)
|
|
audio_chunk = self.model.acoustic_tokenizer.decode(
|
|
scaled_latent.to(self.model.acoustic_tokenizer.device),
|
|
cache=acoustic_cache, # Use acoustic-specific cache
|
|
sample_indices=diffusion_indices.to(self.model.acoustic_tokenizer.device),
|
|
use_cache=True,
|
|
debug=False
|
|
)
|
|
|
|
# Store audio chunks for each sample
|
|
for i, sample_idx in enumerate(diffusion_indices):
|
|
idx = sample_idx.item()
|
|
# Only append audio chunk if the sample is not finished
|
|
if not finished_tags[idx]:
|
|
audio_chunks[idx].append(audio_chunk[i])
|
|
|
|
# Add streaming support here
|
|
if audio_streamer is not None:
|
|
# Stream the audio chunks immediately
|
|
audio_streamer.put(audio_chunk, diffusion_indices)
|
|
|
|
acoustic_embed = self.model.acoustic_connector(speech_latent)
|
|
tts_lm_input_ids = torch.cat([tts_lm_input_ids, torch.ones_like(tts_lm_input_ids[:, -1:])], dim=-1)
|
|
|
|
if tts_lm_input_ids.shape[1] > tts_lm_generation_config.max_length:
|
|
break
|
|
|
|
step += 1
|
|
total_generated_speech_tokens += 1
|
|
if progress_bar is not None:
|
|
progress_bar.update(1)
|
|
progress_bar.set_description(f"Prefilled {total_prefilled_text_tokens} text tokens, generated {total_generated_speech_tokens} speech tokens, current step ({step} / {tts_lm_generation_config.max_length})")
|
|
|
|
tts_lm_model_inputs = self.prepare_inputs_for_generation(tts_lm_input_ids, **tts_lm_model_kwargs)
|
|
tts_lm_additional_inputs = {
|
|
"tts_text_masks": torch.zeros_like(tts_lm_input_ids[:, -1:]),
|
|
"lm_last_hidden_state": acoustic_embed,
|
|
}
|
|
# Forward pass through the model
|
|
tts_lm_outputs = self.forward_tts_lm(
|
|
**tts_lm_model_inputs, **tts_lm_additional_inputs, return_dict=True, output_attentions=False, output_hidden_states=False,
|
|
)
|
|
if cur_speech_index == TTS_SPEECH_WINDOW_SIZE - 1 and next_text_window_size > 0:
|
|
tts_lm_model_kwargs = _update_model_kwargs_for_generation(
|
|
tts_lm_outputs, tts_lm_model_kwargs, num_new_tokens=next_text_window_size,
|
|
)
|
|
else:
|
|
tts_lm_model_kwargs = self._update_model_kwargs_for_generation(
|
|
tts_lm_outputs, tts_lm_model_kwargs, is_encoder_decoder=False,
|
|
)
|
|
|
|
tts_lm_negative_input_ids = torch.cat([tts_lm_negative_input_ids, torch.ones_like(tts_lm_input_ids[:, -1:])], dim=-1)
|
|
tts_lm_negative_model_inputs = self.prepare_inputs_for_generation(tts_lm_negative_input_ids, **tts_lm_negative_model_kwargs)
|
|
# Forward negative pass through the model
|
|
tts_lm_negative_additional_inputs = {
|
|
"tts_text_masks": torch.zeros_like(tts_lm_negative_input_ids[:, -1:]),
|
|
"lm_last_hidden_state": acoustic_embed,
|
|
}
|
|
tts_lm_negative_outputs = self.forward_tts_lm(
|
|
**tts_lm_negative_model_inputs, **tts_lm_negative_additional_inputs, return_dict=True, output_attentions=False, output_hidden_states=False,
|
|
)
|
|
tts_lm_negative_model_kwargs = self._update_model_kwargs_for_generation(
|
|
tts_lm_negative_outputs, tts_lm_negative_model_kwargs, is_encoder_decoder=False,
|
|
)
|
|
|
|
tts_eos_logits = torch.sigmoid(self.tts_eos_classifier(tts_lm_outputs.last_hidden_state[diffusion_indices, -1, :]))
|
|
if tts_eos_logits[0].item() > 0.5:
|
|
# If EOS token is predicted, we can stop generation for this sample
|
|
finished_tags[diffusion_indices] = True
|
|
if audio_streamer is not None:
|
|
audio_streamer.end(diffusion_indices)
|
|
|
|
if tts_lm_input_ids.shape[1] > tts_lm_generation_config.max_length:
|
|
if verbose:
|
|
print(f"Reached maximum generation length {tts_lm_generation_config.max_length}, stopped it.")
|
|
reached_samples = torch.arange(batch_size, device=device)[~finished_tags]
|
|
if reached_samples.numel() > 0:
|
|
reach_max_step_sample[reached_samples] = True
|
|
break
|
|
|
|
if audio_streamer is not None:
|
|
audio_streamer.end()
|
|
|
|
# Concatenate audio chunks for each sample
|
|
final_audio_outputs = []
|
|
for sample_chunks in audio_chunks:
|
|
if sample_chunks:
|
|
# Concatenate all chunks along the time dimension (assumed to be the last dimension)
|
|
concatenated_audio = torch.cat(sample_chunks, dim=-1)
|
|
final_audio_outputs.append(concatenated_audio)
|
|
else:
|
|
# If no audio was generated for this sample, append None
|
|
final_audio_outputs.append(None)
|
|
|
|
if reach_max_step_sample is not None and reach_max_step_sample.any():
|
|
print(f"Reached maximum generation length {tts_lm_generation_config.max_length}, stopped it.")
|
|
|
|
return VibeVoiceGenerationOutput(
|
|
sequences=tts_lm_input_ids,
|
|
speech_outputs=final_audio_outputs if return_speech else None,
|
|
reach_max_step_sample=reach_max_step_sample,
|
|
)
|
|
|
|
@torch.no_grad()
|
|
def sample_speech_tokens(self, condition, neg_condition, cfg_scale=3.0):
|
|
self.model.noise_scheduler.set_timesteps(self.ddpm_inference_steps)
|
|
condition = torch.cat([condition, neg_condition], dim=0).to(self.model.prediction_head.device)
|
|
speech = torch.randn(condition.shape[0], self.config.acoustic_vae_dim).to(condition)
|
|
for t in self.model.noise_scheduler.timesteps:
|
|
half = speech[: len(speech) // 2]
|
|
combined = torch.cat([half, half], dim=0)
|
|
eps = self.model.prediction_head(combined, t.repeat(combined.shape[0]).to(combined), condition=condition)
|
|
cond_eps, uncond_eps = torch.split(eps, len(eps) // 2, dim=0)
|
|
half_eps = uncond_eps + cfg_scale * (cond_eps - uncond_eps)
|
|
eps = torch.cat([half_eps, half_eps], dim=0)
|
|
speech = self.model.noise_scheduler.step(eps, t, speech).prev_sample
|
|
return speech[: len(speech) // 2]
|
|
|
|
|
|
AutoModelForCausalLM.register(VibeVoiceStreamingConfig, VibeVoiceStreamingForConditionalGenerationInference)
|
|
|
|
__all__ = [
|
|
"VibeVoiceStreamingForConditionalGenerationInference",
|
|
]
|