christian.schwarz/VibeVoice
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
142a00112e4594f02da4b3bb2458f07b77883278
VibeVoice/vllm_plugin/model.py
T

1330 lines
54 KiB
Python
Raw Blame History

"""
VibeVoice vLLM Plugin Model - Native Multimodal Integration
This module implements the VibeVoice ASR model with full vLLM multimodal registry
integration for speech-to-text inference.
"""
from typing import List, Optional, Tuple, Union, Dict, Any, Iterable, Mapping, Sequence, ClassVar, Literal
import json
import math
import os
import sys
from pathlib import Path
import torch
import torch.nn as nn
import numpy as np
from io import BytesIO
import tempfile
import base64
# ============================================================================
# Audio Loading: FFmpeg-based AudioMediaIO
# ============================================================================
# VibeVoice uses FFmpeg for audio decoding to ensure consistent behavior
# across different audio formats (MP3, WAV, FLAC, etc.).
from vibevoice.processor.audio_utils import load_audio_use_ffmpeg, load_audio_bytes_use_ffmpeg, AudioNormalizer
def _suffix_from_media_type(media_type: str | None) -> str:
if not media_type:
return ".bin"
mt = media_type.lower().strip()
if mt in ("audio/wav", "audio/x-wav", "audio/wave"):
return ".wav"
if mt in ("audio/mpeg", "audio/mp3", "audio/x-mp3"):
return ".mp3"
if mt in ("audio/flac",):
return ".flac"
if mt in ("audio/ogg", "audio/opus"):
return ".ogg"
if mt in ("audio/mp4", "audio/m4a"):
return ".m4a"
if mt in ("video/mp4",):
return ".mp4"
return ".bin"
def _ffmpeg_load_bytes(data: bytes, *, media_type: str | None = None) -> tuple[np.ndarray, int]:
"""Load audio bytes using FFmpeg.
Returns:
Tuple of (audio_waveform, sample_rate). Sample rate is always 24000.
"""
# Prefer stdin-pipe decoding to avoid temp-file IO under high concurrency.
audio, sr = load_audio_bytes_use_ffmpeg(data, resample=True, target_sr=24000)
normalizer = AudioNormalizer()
audio = normalizer(audio)
return audio, sr
def _ffmpeg_load_file(filepath) -> tuple[np.ndarray, int]:
"""Load audio file using FFmpeg.
Returns:
Tuple of (audio_waveform, sample_rate). Sample rate is always 24000.
"""
audio, sr = load_audio_use_ffmpeg(str(filepath), resample=True, target_sr=24000)
normalizer = AudioNormalizer()
audio = normalizer(audio)
return audio, sr
# Register FFmpeg-based audio loader
import vllm.multimodal.audio as _vllm_audio_module
_OriginalAudioMediaIO = _vllm_audio_module.AudioMediaIO
class _PatchedAudioMediaIO(_OriginalAudioMediaIO):
"""AudioMediaIO implementation using FFmpeg for audio decoding."""
def load_bytes(self, data: bytes) -> tuple[np.ndarray, int]:
return _ffmpeg_load_bytes(data, media_type=None)
def load_base64(self, media_type: str, data: str) -> tuple[np.ndarray, int]:
return _ffmpeg_load_bytes(base64.b64decode(data), media_type=media_type)
def load_file(self, filepath) -> tuple[np.ndarray, int]:
return _ffmpeg_load_file(filepath)
# Replace globally
_vllm_audio_module.AudioMediaIO = _PatchedAudioMediaIO
# Also patch in utils module where it's imported
import vllm.multimodal.utils as _vllm_utils_module
_vllm_utils_module.AudioMediaIO = _PatchedAudioMediaIO
# ============================================================================
from transformers import Qwen2Config, BatchFeature
from transformers.models.whisper import WhisperFeatureExtractor
from vllm.model_executor.models.qwen2 import Qwen2ForCausalLM
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
from vllm.config import VllmConfig, ModelConfig
from vllm.config.speech_to_text import SpeechToTextConfig
from vllm.multimodal import MULTIMODAL_REGISTRY
from vllm.multimodal.parse import MultiModalDataParser
from vllm.sequence import IntermediateTensors
from vllm.model_executor.models.interfaces import SupportsMultiModal, SupportsPP, MultiModalEmbeddings
from vllm.inputs import PromptType
from vllm.model_executor.models.utils import (
init_vllm_registered_model,
maybe_prefix,
AutoWeightsLoader,
WeightsMapper,
)
from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargsItems
from vllm.multimodal.processing import (
BaseMultiModalProcessor,
BaseProcessingInfo,
PromptReplacement,
PromptUpdate,
PromptUpdateDetails,
)
from vllm.multimodal.profiling import BaseDummyInputsBuilder, ProcessorInputs
# Import VibeVoice components
from vibevoice.modular.modular_vibevoice_tokenizer import (
VibeVoiceAcousticTokenizerModel,
VibeVoiceSemanticTokenizerModel,
VibeVoiceTokenizerStreamingCache,
VibeVoiceTokenizerEncoderOutput,
)
from vibevoice.modular.configuration_vibevoice import (
VibeVoiceAcousticTokenizerConfig,
VibeVoiceSemanticTokenizerConfig,
)
class SpeechConnector(nn.Module):
"""Projects speech features to language model hidden dimension.
Architecture: fc1 -> RMSNorm -> fc2 (no activation function)
"""
def __init__(self, input_dim: int, output_dim: int):
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, x: torch.Tensor) -> torch.Tensor:
x = self.fc1(x)
x = self.norm(x)
x = self.fc2(x)
return x
class LlamaRMSNorm(nn.Module):
"""RMSNorm layer used in SpeechConnector."""
def __init__(self, hidden_size: int, eps: float = 1e-6):
super().__init__()
self.weight = nn.Parameter(torch.ones(hidden_size))
self.variance_epsilon = eps
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
input_dtype = hidden_states.dtype
hidden_states = hidden_states.to(torch.float32)
variance = hidden_states.pow(2).mean(-1, keepdim=True)
hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
return self.weight * hidden_states.to(input_dtype)
class VibeVoiceAudioEncoder(nn.Module):
"""
VibeVoice Audio Encoder module.
Encapsulates Acoustic/Semantic VAE Tokenizers and projection Connectors.
Converts raw audio waveforms into embeddings compatible with the language model.
Features:
- Streaming support for long audio (>60s by default)
- Configurable dtype for numerical precision
- Supports both sampling and deterministic (mean) modes
"""
def __init__(self, config):
super().__init__()
self.config = config
import sys
def get_cfg(obj, key, default=None):
if isinstance(obj, dict):
return obj.get(key, default)
return getattr(obj, key, default)
self.acoustic_vae_dim = get_cfg(config, "acoustic_vae_dim", 64)
self.semantic_vae_dim = get_cfg(config, "semantic_vae_dim", 128)
decoder_config = get_cfg(config, "decoder_config")
text_config = get_cfg(config, "text_config")
target_hidden_size = None
if decoder_config is not None:
target_hidden_size = get_cfg(decoder_config, "hidden_size")
if target_hidden_size is None and text_config is not None:
target_hidden_size = get_cfg(text_config, "hidden_size")
if target_hidden_size is None:
target_hidden_size = get_cfg(config, "hidden_size")
if target_hidden_size is None:
print("[VibeVoice] WARN: Could not find hidden_size in config! Defaulting to 3584 (7B).", file=sys.stderr)
self.hidden_size = 3584
else:
self.hidden_size = target_hidden_size
ac_cfg = get_cfg(config, "acoustic_tokenizer_config")
sc_cfg = get_cfg(config, "semantic_tokenizer_config")
if ac_cfg is None or sc_cfg is None:
raise ValueError("Missing acoustic/semantic tokenizer config in model config")
# Handle both dict and already-constructed config objects
if isinstance(ac_cfg, VibeVoiceAcousticTokenizerConfig):
acoustic_config = ac_cfg
elif isinstance(ac_cfg, dict):
acoustic_config = VibeVoiceAcousticTokenizerConfig(**ac_cfg)
else:
raise TypeError(f"acoustic_tokenizer_config has unexpected type: {type(ac_cfg)}")
if isinstance(sc_cfg, VibeVoiceSemanticTokenizerConfig):
semantic_config = sc_cfg
elif isinstance(sc_cfg, dict):
semantic_config = VibeVoiceSemanticTokenizerConfig(**sc_cfg)
else:
raise TypeError(f"semantic_tokenizer_config has unexpected type: {type(sc_cfg)}")
# Tokenizers use float32 for numerical precision
self.acoustic_tokenizer = VibeVoiceAcousticTokenizerModel(acoustic_config)
self.semantic_tokenizer = VibeVoiceSemanticTokenizerModel(semantic_config)
# Get audio encoder dtype from config (defaults to float32 for precision)
root_torch_dtype = get_cfg(config, "torch_dtype", None)
if root_torch_dtype is not None:
if isinstance(root_torch_dtype, str):
self._audio_encoder_dtype = getattr(torch, root_torch_dtype)
else:
self._audio_encoder_dtype = root_torch_dtype
else:
self._audio_encoder_dtype = torch.float32
self.acoustic_connector = SpeechConnector(self.acoustic_vae_dim, self.hidden_size)
self.semantic_connector = SpeechConnector(self.semantic_vae_dim, self.hidden_size)
self.compress_ratio = get_cfg(config, "speech_tok_compress_ratio", 3200)
# Streaming controls
self.sample_rate = get_cfg(config, "target_sample_rate", 24000)
# Default to True (per requirement): segment + cache inside one forward call.
self.enable_streaming = get_cfg(config, "enable_streaming", True)
self.streaming_segment_duration = get_cfg(config, "streaming_segment_duration", 60.0)
# Control whether to use sample() or .mean for acoustic tokens
# Default: use sample() for training-consistent behavior
# Set VIBEVOICE_USE_MEAN=1 for deterministic output
use_mean_env = os.getenv("VIBEVOICE_USE_MEAN", "").strip().lower()
self.use_sample = use_mean_env not in ("1", "true", "yes")
# Language model dtype (set by VibeVoiceForCausalLM.__init__)
# This is the dtype that audio embeddings will be converted to before
# being passed to the language model. Defaults to bfloat16.
self._lm_dtype: torch.dtype = torch.bfloat16
def _ensure_audio_encoder_dtype(self):
"""Ensure all audio encoder components use the correct dtype from config.
vLLM may convert weights to a different dtype (e.g., bfloat16) during loading.
This method converts audio encoder components back to the config-specified dtype
(typically float32) for numerical precision during audio encoding.
"""
import sys
target_dtype = self._audio_encoder_dtype
# Check and convert acoustic_tokenizer
try:
acoustic_dtype = next(self.acoustic_tokenizer.parameters()).dtype
if acoustic_dtype != target_dtype:
self.acoustic_tokenizer = self.acoustic_tokenizer.to(dtype=target_dtype)
print(f"[VibeVoice] Converted acoustic_tokenizer to {target_dtype} (was {acoustic_dtype})", file=sys.stderr)
except StopIteration:
pass
# Check and convert semantic_tokenizer
try:
semantic_dtype = next(self.semantic_tokenizer.parameters()).dtype
if semantic_dtype != target_dtype:
self.semantic_tokenizer = self.semantic_tokenizer.to(dtype=target_dtype)
print(f"[VibeVoice] Converted semantic_tokenizer to {target_dtype} (was {semantic_dtype})", file=sys.stderr)
except StopIteration:
pass
# Check and convert acoustic_connector
try:
ac_conn_dtype = next(self.acoustic_connector.parameters()).dtype
if ac_conn_dtype != target_dtype:
self.acoustic_connector = self.acoustic_connector.to(dtype=target_dtype)
print(f"[VibeVoice] Converted acoustic_connector to {target_dtype} (was {ac_conn_dtype})", file=sys.stderr)
except StopIteration:
pass
# Check and convert semantic_connector
try:
sc_conn_dtype = next(self.semantic_connector.parameters()).dtype
if sc_conn_dtype != target_dtype:
self.semantic_connector = self.semantic_connector.to(dtype=target_dtype)
print(f"[VibeVoice] Converted semantic_connector to {target_dtype} (was {sc_conn_dtype})", file=sys.stderr)
except StopIteration:
pass
def forward(
self,
audio: torch.Tensor,
*,
use_streaming: bool = True,
segment_duration_s: Optional[float] = None,
use_sample: Optional[bool] = None,
) -> torch.Tensor:
"""Encode audio with optional streaming for long clips.
Args:
audio: Input audio tensor [B, T] or [T]
use_streaming: Whether to enable segmented encoding for long audio
segment_duration_s: Segment length in seconds (defaults to 60s)
use_sample: If True, use sampling for acoustic tokens; if False, use mean
Defaults to self.use_sample (controlled by VIBEVOICE_USE_MEAN env var)
Returns:
Audio embeddings tensor compatible with the language model
"""
# Ensure audio encoder components use correct dtype
self._ensure_audio_encoder_dtype()
# Audio input should match the audio encoder dtype
audio = audio.to(dtype=self._audio_encoder_dtype)
if audio.ndim == 1:
audio = audio.unsqueeze(0)
# Resolve streaming options
segment_duration = segment_duration_s or self.streaming_segment_duration
sample_rate = self.sample_rate
total_samples = audio.shape[-1]
segment_samples = int(segment_duration * sample_rate)
use_streaming = use_streaming and self.enable_streaming and total_samples > segment_samples
# Resolve use_sample flag
if use_sample is None:
use_sample = self.use_sample
# Keep encoding in inference mode to avoid autograd build-up
with torch.no_grad():
if not use_streaming:
acoustic_input = audio.unsqueeze(1)
acoustic_out = self.acoustic_tokenizer.encode(acoustic_input)
# Use sample() or .mean based on use_sample flag
if use_sample:
acoustic_tokens = acoustic_out.sample(
dist_type=self.acoustic_tokenizer.std_dist_type
)[0]
else:
acoustic_tokens = acoustic_out.mean
# Connector is now float32, no conversion needed
acoustic_embeds = self.acoustic_connector(acoustic_tokens)
semantic_out = self.semantic_tokenizer.encode(acoustic_input)
# Semantic always uses .mean for consistency
semantic_tokens = semantic_out.mean
# Connector is now float32, no conversion needed
semantic_embeds = self.semantic_connector(semantic_tokens)
else:
# ==========================================
# Streaming path (Retained for future use)
# ==========================================
acoustic_cache = VibeVoiceTokenizerStreamingCache()
semantic_cache = VibeVoiceTokenizerStreamingCache()
acoustic_mean_segments = []
semantic_mean_segments = []
batch_size = audio.shape[0]
sample_indices = torch.arange(batch_size, device=audio.device)
def _iter_segments(total_length: int, segment_length: int):
for start in range(0, total_length, segment_length):
end = min(start + segment_length, total_length)
if end > start:
yield start, end
segments = list(_iter_segments(total_samples, segment_samples))
num_segments = len(segments)
for seg_idx, (start, end) in enumerate(segments):
chunk = audio[:, start:end].contiguous()
if chunk.numel() == 0:
continue
# Check if this is the final segment
is_final = (seg_idx == num_segments - 1)
# --- Acoustic Encode ---
acoustic_enc_out = self.acoustic_tokenizer.encode(
chunk.unsqueeze(1),
cache=acoustic_cache,
sample_indices=sample_indices,
use_cache=True,
is_final_chunk=is_final,
)
acoustic_mean_segments.append(acoustic_enc_out.mean)
# --- Semantic Encode ---
semantic_enc_out = self.semantic_tokenizer.encode(
chunk.unsqueeze(1),
cache=semantic_cache,
sample_indices=sample_indices,
use_cache=True,
is_final_chunk=is_final,
)
semantic_mean_segments.append(semantic_enc_out.mean)
# Concatenate sequence outputs (Acoustic)
if len(acoustic_mean_segments) == 0:
acoustic_mean_full = torch.zeros(
(batch_size, 0, self.acoustic_vae_dim),
device=audio.device,
dtype=self._audio_encoder_dtype # Use config dtype
)
else:
acoustic_mean_full = torch.cat(acoustic_mean_segments, dim=1).contiguous()
# Get acoustic tokens based on use_sample flag
acoustic_enc_full = VibeVoiceTokenizerEncoderOutput(
mean=acoustic_mean_full,
std=self.acoustic_tokenizer.fix_std,
)
if use_sample:
acoustic_tokens = acoustic_enc_full.sample(
dist_type=self.acoustic_tokenizer.std_dist_type
)[0]
else:
acoustic_tokens = acoustic_enc_full.mean
# Connector uses same dtype as tokenizer
acoustic_embeds = self.acoustic_connector(acoustic_tokens)
# Concatenate sequence outputs (Semantic)
if len(semantic_mean_segments) == 0:
semantic_tokens = torch.zeros(
(batch_size, 0, self.semantic_vae_dim),
device=audio.device,
dtype=self._audio_encoder_dtype # Use config dtype
)
else:
semantic_tokens = torch.cat(semantic_mean_segments, dim=1).contiguous()
# Connector uses same dtype as tokenizer
semantic_embeds = self.semantic_connector(semantic_tokens)
# Combine acoustic and semantic embeddings
combined_embeds = acoustic_embeds + semantic_embeds
# Convert to language model dtype for compatibility
# Audio encoder uses config.torch_dtype (typically float32) for numerical precision,
# but LM expects the dtype specified by vLLM's --dtype flag (e.g., bfloat16, float16)
combined_embeds = combined_embeds.to(dtype=self._lm_dtype)
return combined_embeds
# ============================================================================
# vLLM Multimodal Processing Infrastructure
# ============================================================================
class VibeVoiceProcessingInfo(BaseProcessingInfo):
"""Processing info for VibeVoice multimodal model."""
def get_hf_config(self):
return self.ctx.get_hf_config()
def get_feature_extractor(self, **kwargs) -> WhisperFeatureExtractor:
"""
Get a WhisperFeatureExtractor for vLLM profiling compatibility.
IMPORTANT: This is NOT used in actual inference!
VibeVoice uses its own acoustic/semantic VAE tokenizers operating on
raw 24kHz waveforms, NOT Whisper mel spectrograms.
This feature extractor exists only to satisfy vLLM's multimodal
profiling infrastructure which may query audio parameters like
sampling_rate and chunk_length for memory estimation.
"""
# Read config from preprocessor_config.json if available
import json
import os
model_path = self.ctx.model_config.model
preprocessor_path = os.path.join(model_path, "preprocessor_config.json")
# Default values: keep a coherent (sr, hop) pair.
# VibeVoice runs at 24kHz in this repo (see demo/asr_transcribe_file.py).
config = {
"sampling_rate": 24000,
"feature_size": 128,
# 10ms hop at 24kHz
"hop_length": 240,
"chunk_length": 30,
"n_fft": 400,
"padding_value": 0.0,
}
# Try to load from config file
if os.path.exists(preprocessor_path):
try:
with open(preprocessor_path, "r") as f:
file_config = json.load(f)
config.update({k: file_config[k] for k in config.keys() if k in file_config})
except Exception:
pass # Use defaults
return WhisperFeatureExtractor(
feature_size=config["feature_size"],
sampling_rate=config["sampling_rate"],
hop_length=config["hop_length"],
chunk_length=config["chunk_length"],
n_fft=config["n_fft"],
padding_value=config["padding_value"],
)
def get_audio_token_info(self) -> dict:
"""
Get audio special tokens and their IDs.
Returns dict with:
audio_token, audio_bos_token, audio_eos_token,
audio_token_id, audio_bos_id, audio_eos_id
"""
tokenizer = self.get_tokenizer()
vocab = tokenizer.get_vocab()
# VibeVoice special tokens
tokens = {
"audio_token": "<|AUDIO|>",
"audio_bos_token": "<|audio_bos|>",
"audio_eos_token": "<|audio_eos|>",
}
# Get IDs
tokens["audio_token_id"] = vocab.get(tokens["audio_token"])
tokens["audio_bos_id"] = vocab.get(tokens["audio_bos_token"])
tokens["audio_eos_id"] = vocab.get(tokens["audio_eos_token"])
return tokens
def get_supported_mm_limits(self) -> Mapping[str, int | None]:
return {"audio": None}
class VibeVoiceDummyInputsBuilder(BaseDummyInputsBuilder[VibeVoiceProcessingInfo]):
"""
Build dummy inputs for multimodal profiling.
Dummy text uses the raw <|AUDIO|> token(s). vLLM's processing pipeline will
expand each <|AUDIO|> via `VibeVoiceMultiModalProcessor._get_prompt_updates`
into the full ASR format:
[speech_start_id] + [speech_pad_id] * N + [speech_end_id] + [newline_id]
where N is derived from audio length / compress_ratio.
"""
def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
num_audios = mm_counts.get("audio", 0)
if num_audios <= 0:
return ""
# Get the audio token from our token info helper
token_info = self.info.get_audio_token_info()
audio_token = token_info["audio_token"]
# Return ONLY the audio tokens - the HF processor adds bos/eos
return audio_token * num_audios
def get_dummy_mm_data(
self,
seq_len: int,
mm_counts: Mapping[str, int],
mm_options: Mapping[str, Any] | None = None,
) -> Dict[str, Any]:
"""Generate dummy audio data for profiling."""
feature_extractor = self.info.get_feature_extractor()
sampling_rate = feature_extractor.sampling_rate
audio_len = feature_extractor.chunk_length * sampling_rate
num_audios = mm_counts.get("audio", 0)
# Generate dummy audio as numpy arrays (what the HF processor expects)
return {
"audio": [np.zeros(audio_len, dtype=np.float32) for _ in range(num_audios)]
}
def get_dummy_processor_inputs(
self,
seq_len: int,
mm_counts: Mapping[str, int],
mm_options: Mapping[str, Any] | None = None,
) -> ProcessorInputs:
"""Build ProcessorInputs for dummy profiling."""
return ProcessorInputs(
prompt=self.get_dummy_text(mm_counts),
mm_data=self.get_dummy_mm_data(seq_len, mm_counts, mm_options),
)
def _vibevoice_field_config(hf_inputs: Mapping[str, torch.Tensor]):
"""Map HF processor output keys to audio modality.
Returns a config dict that tells vLLM how to batch multimodal data.
"""
# Always define the config for all fields we use
# Even if the field isn't in hf_inputs, vLLM needs to know how to batch it
config = {
# These are our custom fields for VibeVoice
"raw_audio": MultiModalFieldConfig.batched("audio"),
"raw_audio_lengths": MultiModalFieldConfig.batched("audio"),
"salt": MultiModalFieldConfig.batched("audio"),
}
# Add optional Whisper features if present
if "input_features" in hf_inputs:
config["input_features"] = MultiModalFieldConfig.batched("audio")
if "feature_attention_mask" in hf_inputs:
config["feature_attention_mask"] = MultiModalFieldConfig.batched("audio")
return config
class VibeVoiceMultiModalProcessor(BaseMultiModalProcessor[VibeVoiceProcessingInfo]):
"""
Multimodal processor for VibeVoice.
Handles the conversion of raw audio inputs to model-ready features,
and manages the prompt token replacement for audio placeholders.
"""
def _get_data_parser(self) -> MultiModalDataParser:
"""Create a data parser with the correct target sample rate (24kHz)."""
# VibeVoice requires 24kHz, not 16kHz (Whisper default)
target_sr = 24000
return MultiModalDataParser(target_sr=target_sr)
def _call_hf_processor(
self,
prompt: str,
mm_data: Mapping[str, object],
mm_kwargs: Mapping[str, object],
tok_kwargs: Mapping[str, object],
) -> BatchFeature:
"""
Process prompt and audio for vLLM multimodal pipeline.
We intentionally do NOT run a HF processor that would pre-expand
`<|AUDIO|>` inside this method. Instead we:
1) Tokenize the prompt as-is (so `<|AUDIO|>` stays a single token)
2) Store raw audio tensors for `embed_multimodal` to encode later
3) Let vLLM call `_get_prompt_updates` to expand the single `<|AUDIO|>`
into the full ASR format: [speech_start] + N*[speech_pad] + [speech_end] + [\\n]
"""
# Handle the case where 'audios' key is used (transformers deprecation)
mm_data = dict(mm_data) # Make a mutable copy
audios = mm_data.pop("audios", None)
if audios is not None and "audio" not in mm_data:
mm_data["audio"] = audios
# Text-only input handling
if not mm_data.get("audio"):
prompt_ids = self.info.get_tokenizer().encode(prompt)
prompt_ids = self._apply_hf_processor_tokens_only(prompt_ids)
return BatchFeature(dict(input_ids=[prompt_ids]), tensor_type="pt")
# Get raw audio data
raw_audio_list = mm_data.get("audio")
if isinstance(raw_audio_list, np.ndarray):
raw_audio_list = [raw_audio_list]
elif not isinstance(raw_audio_list, list):
raw_audio_list = list(raw_audio_list)
# Tokenize prompt directly to preserve <|AUDIO|> as a single token
# vLLM will expand it via _get_prompt_updates
tokenizer = self.info.get_tokenizer()
prompt_ids = tokenizer.encode(prompt, add_special_tokens=False)
prompt_ids = self._apply_hf_processor_tokens_only(prompt_ids)
# Create result with input_ids
result = BatchFeature(dict(input_ids=[prompt_ids]), tensor_type="pt")
# Add raw audio tensors for VibeVoice encoder
# Stack into a single tensor for vLLM's batched field config
max_len = max(len(a) for a in raw_audio_list)
raw_audio_tensors = []
audio_lengths = []
for audio in raw_audio_list:
audio_len = len(audio)
audio_lengths.append(audio_len)
if audio_len < max_len:
audio = np.pad(audio, (0, max_len - audio_len), mode='constant')
raw_audio_tensors.append(torch.from_numpy(audio).float())
# Stack into [num_audios, max_len] tensor
stacked_audio = torch.stack(raw_audio_tensors, dim=0) # Shape: [num_audios, max_len]
result["raw_audio"] = stacked_audio
# Convert lengths to tensor as well
result["raw_audio_lengths"] = torch.tensor(audio_lengths, dtype=torch.long)
# Add a random salt to ensure unique hash and bypass cache
import uuid
# Use a random integer for salt
salt_val = hash(str(uuid.uuid4())) % 100000
result["salt"] = torch.tensor([salt_val], dtype=torch.long).expand(len(raw_audio_list))
return result
def _hf_processor_applies_updates(
self,
prompt_text: str,
mm_items,
hf_processor_mm_kwargs: Mapping[str, object],
tokenization_kwargs: Mapping[str, object],
) -> bool:
"""Return whether the HF processor applies prompt updates.
Returns False because we handle token expansion via _get_prompt_updates.
"""
return False
def _get_mm_fields_config(
self,
hf_inputs: BatchFeature,
hf_processor_mm_kwargs: Mapping[str, object],
) -> Mapping[str, MultiModalFieldConfig]:
"""Configure which HF output fields map to which modality."""
return _vibevoice_field_config(hf_inputs)
def _get_prompt_updates(
self,
mm_items,
hf_processor_mm_kwargs: Mapping[str, object],
out_mm_kwargs: MultiModalKwargsItems,
) -> Sequence[PromptUpdate]:
"""
Define how to replace the audio placeholder in the prompt.
vLLM's OpenAI multimodal parsing inserts the model placeholder string
from `get_placeholder_str` (here: `<|AUDIO|>`) into the conversation.
We expand that single token into N repeated `<|AUDIO|>` tokens, where N
is derived from waveform length and `speech_tok_compress_ratio`.
"""
token_info = self.info.get_audio_token_info()
audio_token = token_info["audio_token"]
audio_token_id = token_info["audio_token_id"]
audio_bos_id = token_info.get("audio_bos_id")
audio_eos_id = token_info.get("audio_eos_id")
tokenizer = self.info.get_tokenizer()
vocab = tokenizer.get_vocab()
def _tok_id(name: str) -> int | None:
return vocab.get(name)
# Look up speech token IDs from vocabulary
# These tokens mark the start/end of audio embeddings in the prompt
speech_start_id = (
_tok_id("<|object_ref_start|>")
or getattr(tokenizer, "speech_start_id", None)
or _tok_id("<|speech_start|>")
)
speech_end_id = (
_tok_id("<|object_ref_end|>")
or getattr(tokenizer, "speech_end_id", None)
or _tok_id("<|speech_end|>")
)
speech_pad_id = (
_tok_id("<|box_start|>")
or getattr(tokenizer, "speech_pad_id", None)
or _tok_id("<|speech_pad|>")
)
if audio_token_id is None:
return []
# Get raw audio lengths (in samples, after resampling to 24kHz) from our stored data
out_mm_data = out_mm_kwargs.get_data()
raw_audio_lengths = out_mm_data.get("raw_audio_lengths", [])
# Fetch defaults from model config when available (falls back to 3200)
hf_config = self.info.get_hf_config()
if isinstance(hf_config, dict):
compress_ratio = int(hf_config.get("speech_tok_compress_ratio", 3200))
else:
compress_ratio = int(getattr(hf_config, "speech_tok_compress_ratio", 3200))
def _to_int_len(x) -> int:
if x is None:
return 0
if isinstance(x, torch.Tensor):
# Accept 0-dim or 1-dim scalar-like tensors
if x.numel() == 1:
return int(x.item())
# If a full tensor is passed accidentally, fall back to its length
return int(x.shape[0])
return int(x)
def get_replacement(item_idx: int):
if raw_audio_lengths and item_idx < len(raw_audio_lengths):
audio_len = _to_int_len(raw_audio_lengths[item_idx])
num_features = max(1, int(np.ceil(audio_len / compress_ratio)))
else:
# Fallback: estimate for 30 second audio at 24kHz
num_features = int(np.ceil(30 * 24000 / compress_ratio))
if num_features == 0:
raise ValueError(
f"Audio at index {item_idx} is too short to be represented"
)
# Build replacement token sequence:
# <|speech_start|> + N * <|speech_pad|> + <|speech_end|> + \n
# The newline is important for correct prompt structure.
newline_id = 198 # '\n' token
if speech_start_id is not None and speech_pad_id is not None and speech_end_id is not None:
embed_id = int(speech_pad_id)
replacement_ids = [int(speech_start_id)] + [embed_id] * num_features + [int(speech_end_id), newline_id]
# Fallback: add audio BOS/EOS boundaries around repeated <|AUDIO|>.
elif audio_bos_id is not None and audio_eos_id is not None:
embed_id = int(audio_token_id)
replacement_ids = [int(audio_bos_id)] + [embed_id] * num_features + [int(audio_eos_id)]
else:
embed_id = int(audio_token_id)
replacement_ids = [embed_id] * num_features
return PromptUpdateDetails.select_token_id(
replacement_ids,
embed_token_id=int(embed_id),
)
return [
PromptReplacement(
modality="audio",
# Keep string placeholder matching for maximum vLLM compatibility.
target=audio_token,
replacement=get_replacement,
)
]
# ============================================================================
# Main Model Class
# ============================================================================
@MULTIMODAL_REGISTRY.register_processor(
VibeVoiceMultiModalProcessor,
info=VibeVoiceProcessingInfo,
dummy_inputs=VibeVoiceDummyInputsBuilder,
)
class VibeVoiceForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):
"""
VibeVoice ASR model with native vLLM multimodal integration.
This model combines VibeVoice acoustic/semantic tokenizers for audio encoding
with a causal language model for text generation.
"""
# SupportsTranscription interface
supports_transcription: ClassVar[Literal[True]] = True
supports_transcription_only: ClassVar[bool] = False
supports_segment_timestamp: ClassVar[bool] = False
# Supported languages (Chinese as primary target)
supported_languages: ClassVar[Mapping[str, str]] = {
"zh": "Chinese",
"en": "English",
}
@classmethod
def get_placeholder_str(cls, modality: str, i: int) -> str | None:
"""Return the placeholder string format for a given modality.
Returns "<|AUDIO|>" which vLLM inserts into the conversation prompt.
This single placeholder is later expanded by `_get_prompt_updates` into:
[speech_start_id] + [speech_pad_id] * N + [speech_end_id] + [newline_id]
where N = ceil(audio_samples / compress_ratio).
"""
if modality.startswith("audio"):
return "<|AUDIO|>"
raise ValueError(f"Unsupported modality: {modality}")
@classmethod
def get_generation_prompt(
cls,
audio: np.ndarray,
stt_config: SpeechToTextConfig,
model_config: ModelConfig,
language: str | None,
task_type: Literal["transcribe", "translate"],
request_prompt: str,
to_language: str | None,
) -> PromptType:
"""Get the prompt for the ASR model.
Generates a chat-formatted prompt for speech-to-text transcription
with JSON output format.
"""
# If user provides custom prompt, use it
if request_prompt:
return request_prompt
# Calculate audio duration for the prompt
# Audio should be at 24kHz, so duration = len(audio) / 24000
duration = len(audio) / 24000 if audio is not None else 10.0
system_prompt = "You are a helpful assistant that transcribes audio input into text output in JSON format."
show_keys = ["Start time", "End time", "Speaker ID", "Content"]
user_suffix = (
f"This is a {duration:.2f} seconds audio, please transcribe it with these keys: "
+ ", ".join(show_keys)
)
# IMPORTANT: keep <|AUDIO|> as the only placeholder token here.
# `_get_prompt_updates` expands it into repeated `<|AUDIO|>` placeholders.
user_content = "<|AUDIO|>\n" + user_suffix
prompt = (
f"<|im_start|>system\n{system_prompt}<|im_end|>\n"
f"<|im_start|>user\n{user_content}<|im_end|>\n"
"<|im_start|>assistant\n"
)
return prompt
@classmethod
def get_speech_to_text_config(
cls, model_config: ModelConfig, task_type: Literal["transcribe", "translate"]
) -> SpeechToTextConfig:
"""Get the speech to text config for the ASR model."""
return SpeechToTextConfig(
language=None, # Auto-detect or use request language
task_type=task_type,
)
@classmethod
def get_num_audio_tokens(
cls,
audio_duration_s: float,
stt_config: SpeechToTextConfig,
model_config: ModelConfig,
) -> int | None:
"""Estimate number of audio tokens from duration.
Returns the number of audio EMBEDDING positions (speech_pad_id tokens).
Note: _get_prompt_updates actually generates:
[speech_start_id] + [speech_pad_id] * N + [speech_end_id] + [newline_id]
So total prompt tokens = N + 3, but this returns N (the embedding count).
"""
sampling_rate = 24000
compress_ratio = 3200
samples = int(audio_duration_s * sampling_rate)
num_tokens = int(np.ceil(samples / compress_ratio))
return num_tokens
@classmethod
def get_other_languages(cls) -> Mapping[str, str]:
"""Get languages from Whisper map not natively supported."""
# Import LANGUAGES from vllm
try:
from vllm.transformers_utils.tokenizer import LANGUAGES
except ImportError:
# Fallback to empty dict if import fails
return {}
return {k: v for k, v in LANGUAGES.items() if k not in cls.supported_languages}
@classmethod
def validate_language(cls, language: str | None) -> str | None:
"""Validate the language code."""
if language is None or language in cls.supported_languages:
return language
elif language in cls.get_other_languages():
print(f"Warning: Language {language!r} is not natively supported")
return language
else:
raise ValueError(
f"Unsupported language: {language!r}. "
f"Supported: {list(cls.supported_languages.keys())}"
)
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
super().__init__()
config = vllm_config.model_config.hf_config
self.config = config
# Keep a copy of the resolved model path for any custom weight-loading
# logic (e.g., loading audio encoder weights in fp32 directly from
# safetensors shards).
self._model_path = vllm_config.model_config.model
self.audio_encoder = VibeVoiceAudioEncoder(config)
# Pass decoder_config to the language model initialization
decoder_config = getattr(config, "decoder_config", config)
self.language_model = init_vllm_registered_model(
vllm_config=vllm_config,
hf_config=decoder_config,
prefix=maybe_prefix(prefix, "language_model"),
architectures=["Qwen2ForCausalLM"],
)
self.make_empty_intermediate_tensors = (
self.language_model.make_empty_intermediate_tensors
)
# Set the language model dtype for audio encoder output conversion
# This should match vLLM's --dtype flag (e.g., bfloat16, float16, float32)
# Audio encoder internal computation stays in fp32 for numerical precision,
# but output is converted to LM dtype for compatibility
lm_dtype = vllm_config.model_config.dtype
if lm_dtype is not None:
self.audio_encoder._lm_dtype = lm_dtype
# Ensure audio encoder uses correct dtype (typically fp32 for precision)
try:
self.audio_encoder._ensure_audio_encoder_dtype()
except Exception:
pass
def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:
return self.language_model.compute_logits(hidden_states)
def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
"""
Extract audio embeddings using VibeVoice's acoustic/semantic tokenizers.
Called by vLLM to get audio embeddings that replace audio placeholder tokens.
Returns:
Tuple of embedding tensors, one per audio input.
"""
# Get raw audio data (stored by our processor)
raw_audio = kwargs.get("raw_audio")
raw_audio_lengths = kwargs.get("raw_audio_lengths")
# Handle no audio input - this happens during memory profiling
if raw_audio is None:
return []
# Handle empty audio list
if isinstance(raw_audio, (list, tuple)) and len(raw_audio) == 0:
return []
# Flatten raw_audio_lengths if it's nested
def flatten_lengths(lengths):
"""Flatten nested lists/tensors of lengths to a single list."""
if lengths is None:
return []
result = []
if isinstance(lengths, torch.Tensor):
lengths = lengths.tolist()
if isinstance(lengths, (list, tuple)):
for item in lengths:
if isinstance(item, (list, tuple)):
result.extend(flatten_lengths(item))
elif isinstance(item, torch.Tensor):
if item.dim() == 0:
result.append(item.item())
else:
result.extend(item.tolist())
else:
result.append(item)
else:
result.append(lengths)
return result
raw_audio_lengths = flatten_lengths(raw_audio_lengths)
# Streaming controls. Enabled by default; can be overridden per-call.
use_streaming_flag = bool(
kwargs.get(
"use_streaming",
getattr(self.audio_encoder, "enable_streaming", True),
)
)
streaming_segment_duration = kwargs.get(
"streaming_segment_duration",
getattr(self.audio_encoder, "streaming_segment_duration", 60.0),
)
# Process each audio through the VibeVoice encoder
embeddings = []
# Get model device and dtype for alignment
try:
device = next(self.audio_encoder.parameters()).device
dtype = next(self.audio_encoder.parameters()).dtype
except StopIteration:
# Fallback if no parameters (shouldn't happen)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
dtype = torch.bfloat16
# Handle both stacked tensor and list of tensors
# vLLM batches as: [batch_size, 1, seq_len] or [batch_size, seq_len]
if isinstance(raw_audio, torch.Tensor):
if raw_audio.dim() == 3:
# Shape: [batch_size, 1, seq_len] - squeeze the middle dimension
num_audios = raw_audio.shape[0]
audio_list = [raw_audio[i].squeeze(0) for i in range(num_audios)]
elif raw_audio.dim() == 2:
# Shape: [batch_size, seq_len]
num_audios = raw_audio.shape[0]
audio_list = [raw_audio[i] for i in range(num_audios)]
else:
# Single 1D tensor
audio_list = [raw_audio]
elif isinstance(raw_audio, (list, tuple)):
audio_list = list(raw_audio)
else:
# Single tensor
audio_list = [raw_audio]
for i, audio_tensor in enumerate(audio_list):
try:
if isinstance(audio_tensor, list):
audio_tensor = torch.stack(audio_tensor)
# Ensure tensor
if not isinstance(audio_tensor, torch.Tensor):
audio_tensor = torch.tensor(audio_tensor)
# Let vLLM handle dtype (bfloat16 by default)
audio_tensor = audio_tensor.to(device=device)
# Get actual length if available, otherwise use full length
if raw_audio_lengths and i < len(raw_audio_lengths):
actual_len = int(raw_audio_lengths[i])
if actual_len > 0 and actual_len <= audio_tensor.shape[-1]:
# Truncate from the last dimension (sequence length)
audio_tensor = audio_tensor[..., :actual_len]
# Skip if audio is too short (< 1 frame)
if audio_tensor.numel() < 160: # Minimum ~1ms at 24kHz
continue
# Encode audio through VibeVoice encoder
audio_embeds = self.audio_encoder(
audio_tensor,
use_streaming=use_streaming_flag,
segment_duration_s=streaming_segment_duration,
)
# audio_embeds shape: [1, seq_len, hidden_size]
# We need to return it as a single embedding tensor per audio
final_embed = audio_embeds.squeeze(0)
embeddings.append(final_embed)
except Exception as e:
# Log error but don't crash - this helps debug profiling issues
print(f"[VibeVoice] Error encoding audio {i}: {e}")
import traceback
traceback.print_exc()
# Return empty embedding to avoid crash
continue
return tuple(embeddings)
def get_input_embeddings(self) -> torch.nn.Module:
"""Return the text embedding layer (embed_tokens).
vLLM uses this to get the embedding module for converting token IDs
to embeddings during decode phase.
Returns:
The embed_tokens module from the language model
"""
# Get embed_tokens from the language model
if hasattr(self.language_model, 'model') and hasattr(self.language_model.model, 'embed_tokens'):
return self.language_model.model.embed_tokens
elif hasattr(self.language_model, 'embed_tokens'):
return self.language_model.embed_tokens
else:
# Try to get from inner model
inner = self.language_model
if hasattr(inner, 'language_model'):
inner = inner.language_model
if hasattr(inner, 'model') and hasattr(inner.model, 'embed_tokens'):
return inner.model.embed_tokens
raise AttributeError("Cannot find embed_tokens layer")
def embed_input_ids(
self,
input_ids: torch.Tensor,
multimodal_embeddings: Optional[Union[torch.Tensor, List[torch.Tensor]]] = None,
is_multimodal: Optional[torch.Tensor] = None,
**kwargs, # Accept any additional kwargs for compatibility
) -> torch.Tensor:
"""Apply token embeddings to input_ids and merge with multimodal embeddings.
This is the preferred method in vLLM V1 for converting token IDs
to embeddings and merging multimodal (audio) embeddings.
Args:
input_ids: Tensor of token IDs to embed
multimodal_embeddings: Pre-computed multimodal embeddings (audio).
Can be a Tensor or a List of Tensors (vLLM standard).
is_multimodal: Boolean mask indicating which positions are multimodal
**kwargs: Additional arguments for compatibility
Returns:
Tensor of embeddings with multimodal content merged in
"""
from vllm.model_executor.models.utils import _merge_multimodal_embeddings
# Get text embeddings
embed_tokens = self.get_input_embeddings()
inputs_embeds = embed_tokens(input_ids)
# Merge multimodal embeddings if provided
if multimodal_embeddings is not None and is_multimodal is not None:
# Use vLLM's standard merge function which handles List[Tensor] correctly
inputs_embeds = _merge_multimodal_embeddings(
inputs_embeds,
multimodal_embeddings,
is_multimodal,
)
return inputs_embeds
def get_language_model(self) -> torch.nn.Module:
"""Return the language model backbone."""
return self.language_model
def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> set[str]:
"""Load model weights from checkpoint.
The checkpoint has weights named like:
- lm_head.weight -> language_model.lm_head.weight
- model.language_model.layers.X... -> language_model.model.layers.X...
- model.acoustic_tokenizer... -> audio_encoder.acoustic_tokenizer...
- model.semantic_tokenizer... -> audio_encoder.semantic_tokenizer...
- model.acoustic_connector... -> audio_encoder.acoustic_connector...
- model.semantic_connector... -> audio_encoder.semantic_connector...
Let vLLM handle all dtype conversions according to --dtype flag.
"""
# Map weight prefixes for VibeVoice
# The checkpoint uses "model." prefix, we need to remap it
mapper = WeightsMapper(
orig_to_new_prefix={
# Audio encoder components: model.X -> audio_encoder.X
"model.acoustic_tokenizer.": "audio_encoder.acoustic_tokenizer.",
"model.semantic_tokenizer.": "audio_encoder.semantic_tokenizer.",
"model.acoustic_connector.": "audio_encoder.acoustic_connector.",
"model.semantic_connector.": "audio_encoder.semantic_connector.",
# Language model: model.language_model.X -> language_model.model.X
"model.language_model.": "language_model.model.",
# LM head: lm_head.X -> language_model.lm_head.X
"lm_head.": "language_model.lm_head.",
}
)
loader = AutoWeightsLoader(self)
return loader.load_weights(weights, mapper=mapper)
def forward(
self,
input_ids: Optional[torch.Tensor],
positions: torch.Tensor,
intermediate_tensors: Optional[IntermediateTensors] = None,
inputs_embeds: Optional[torch.Tensor] = None,
**kwargs: object,
) -> Union[torch.Tensor, IntermediateTensors]:
"""
Forward pass for VibeVoice ASR model.
Handles embedding computation and language model forward pass.
Uses inputs_embeds if provided (from vLLM multimodal merge),
otherwise computes embeddings from input_ids.
Args:
input_ids: Token IDs. May be None when inputs_embeds is provided.
positions: Position indices for the input tokens.
intermediate_tensors: Intermediate tensors for pipeline parallelism.
inputs_embeds: Pre-computed embeddings (from multimodal merge or decode).
"""
try:
# PRIORITY: Use inputs_embeds if provided (from vLLM multimodal merge or decode)
# Only compute from input_ids if inputs_embeds is not available
if inputs_embeds is None and input_ids is not None:
# Compute embeddings from input_ids
inputs_embeds = self.get_input_embeddings()(input_ids)
# If we have intermediate tensors (pipeline parallelism), don't use inputs_embeds
if intermediate_tensors is not None:
inputs_embeds = None
# Get the inner model - handle both wrapped and direct language models
language_model = self.language_model
if hasattr(language_model, "language_model"):
language_model = language_model.language_model
# Call the language model's model (Qwen2Model)
# vLLM V1 passes kv_caches and attn_metadata via context, not arguments
# IMPORTANT: Pass input_ids=None when using inputs_embeds to avoid double embedding
hidden_states = language_model.model(
input_ids=None, # Always None when we have inputs_embeds
positions=positions,
intermediate_tensors=intermediate_tensors,
inputs_embeds=inputs_embeds
)
return hidden_states
except Exception as e:
raise
# Alias for training checkpoint compatibility
VibeVoiceForASRTraining = VibeVoiceForCausalLM
Reference in New Issue View Git Blame Copy Permalink