508 lines
18 KiB
Python
508 lines
18 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""
|
|
OmniVoice model for vLLM-Omni two-stage TTS pipeline.
|
|
|
|
Stage 0 (Generator): Qwen3 backbone + iterative unmasking → 8-codebook tokens
|
|
Stage 1 (Decoder): HiggsAudioV2 decoder → 24kHz waveform
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
from collections.abc import Iterable, Mapping, Sequence
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn as nn
|
|
import torchaudio
|
|
from transformers.feature_extraction_utils import BatchFeature
|
|
from vllm.config import VllmConfig
|
|
from vllm.config.multimodal import BaseDummyOptions
|
|
from vllm.inputs.llm import MultiModalDataDict
|
|
from vllm.logger import init_logger
|
|
from vllm.multimodal.inputs import (
|
|
MultiModalFieldConfig,
|
|
MultiModalKwargsItems,
|
|
)
|
|
from vllm.multimodal.parse import MultiModalDataItems, MultiModalDataParser
|
|
from vllm.multimodal.processing import (
|
|
BaseDummyInputsBuilder,
|
|
BaseMultiModalProcessor,
|
|
BaseProcessingInfo,
|
|
ProcessorInputs,
|
|
PromptIndexTargets,
|
|
PromptInsertion,
|
|
PromptUpdate,
|
|
)
|
|
from vllm.sequence import IntermediateTensors
|
|
|
|
from vllm_omni.model_executor.models.output_templates import OmniOutput
|
|
from vllm_omni.platforms import current_omni_platform
|
|
from vllm_omni.transformers_utils.configs.omnivoice import OmniVoiceConfig
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Multimodal processing
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class OmniVoiceMultiModalProcessingInfo(BaseProcessingInfo):
|
|
def get_hf_config(self):
|
|
return self.ctx.get_hf_config(OmniVoiceConfig)
|
|
|
|
def get_supported_mm_limits(self) -> Mapping[str, int | None]:
|
|
return {"audio": None}
|
|
|
|
def get_data_parser(self):
|
|
return MultiModalDataParser(
|
|
target_sr=self.ctx.get_hf_config().sample_rate,
|
|
expected_hidden_size=self._get_expected_hidden_size(),
|
|
)
|
|
|
|
|
|
class OmniVoiceMultiModalProcessor(BaseMultiModalProcessor[OmniVoiceMultiModalProcessingInfo]):
|
|
"""Processes text + optional reference audio for OmniVoice.
|
|
|
|
For voice cloning: text + reference audio → tokenized reference
|
|
For auto voice: text only
|
|
"""
|
|
|
|
def _ensure_cached_runtime_components(self, model_dir: str, config: OmniVoiceConfig) -> None:
|
|
cached_model_dir = getattr(self, "_cached_model_dir", None)
|
|
if cached_model_dir == model_dir:
|
|
return
|
|
|
|
from transformers import AutoTokenizer
|
|
|
|
self.text_tokenizer = AutoTokenizer.from_pretrained(model_dir)
|
|
|
|
# Audio tokenizer for encoding reference audio (requires transformers>=5.3)
|
|
audio_tokenizer_path = os.path.join(model_dir, "audio_tokenizer")
|
|
try:
|
|
from transformers import (
|
|
AutoFeatureExtractor,
|
|
HiggsAudioV2TokenizerModel,
|
|
)
|
|
|
|
self.audio_tokenizer = HiggsAudioV2TokenizerModel.from_pretrained(audio_tokenizer_path, device_map="cpu")
|
|
self.feature_extractor = AutoFeatureExtractor.from_pretrained(audio_tokenizer_path)
|
|
self.audio_tokenizer.eval()
|
|
except ImportError:
|
|
self.audio_tokenizer = None
|
|
self.feature_extractor = None
|
|
logger.warning("Voice cloning disabled (requires transformers>=5.3.0).")
|
|
|
|
self._cached_model_dir = model_dir
|
|
|
|
def _call_hf_processor(
|
|
self,
|
|
prompt: str,
|
|
mm_data: Mapping[str, object],
|
|
mm_kwargs: Mapping[str, object],
|
|
tok_kwargs: Mapping[str, object],
|
|
) -> BatchFeature:
|
|
config = self.info.ctx.get_hf_config()
|
|
model_dir = self.info.ctx.model_config.model
|
|
self._ensure_cached_runtime_components(model_dir, config)
|
|
|
|
audio = mm_data.get("audio", None)
|
|
if audio is None:
|
|
audio = mm_data.get("audios")
|
|
if audio is not None:
|
|
audio = audio[0], config.sample_rate
|
|
|
|
# Build text prompt with control tokens
|
|
lang = mm_kwargs.get("lang", None)
|
|
instruct = mm_kwargs.get("instruct", None)
|
|
denoise = mm_kwargs.get("denoise", True)
|
|
ref_text = mm_kwargs.get("ref_text", None)
|
|
|
|
# Construct the style + text portion
|
|
style_text = ""
|
|
if denoise:
|
|
style_text += "<|denoise|>"
|
|
lang_str = lang if lang else "None"
|
|
instruct_str = instruct if instruct else "None"
|
|
style_text += f"<|lang_start|>{lang_str}<|lang_end|>"
|
|
style_text += f"<|instruct_start|>{instruct_str}<|instruct_end|>"
|
|
|
|
# Combine ref_text and main text
|
|
if ref_text:
|
|
full_text = f"{ref_text} {prompt}"
|
|
else:
|
|
full_text = prompt
|
|
|
|
text_prompt = f"{style_text}<|text_start|>{full_text}<|text_end|>"
|
|
text_tokens = self.text_tokenizer(text_prompt, return_tensors="pt").input_ids.squeeze(0) # [N_text]
|
|
|
|
if audio is None:
|
|
# Text-only path (auto voice mode)
|
|
return BatchFeature(
|
|
{
|
|
"input_ids": text_tokens,
|
|
"input_len": [len(text_tokens)],
|
|
}
|
|
)
|
|
|
|
# Voice cloning: encode reference audio to tokens
|
|
audio_signal, sr = audio
|
|
if isinstance(audio_signal, np.ndarray):
|
|
audio_signal = torch.from_numpy(audio_signal).float()
|
|
if audio_signal.dim() == 1:
|
|
audio_signal = audio_signal.unsqueeze(0)
|
|
|
|
# Resample to tokenizer sample rate if needed
|
|
if self.feature_extractor is not None:
|
|
target_sr = self.feature_extractor.sampling_rate
|
|
if sr != target_sr:
|
|
audio_signal = torchaudio.functional.resample(audio_signal, sr, target_sr)
|
|
|
|
# Encode reference audio to 8-codebook tokens
|
|
if self.audio_tokenizer is None:
|
|
raise RuntimeError("Voice cloning requires transformers>=5.3.0. Try: uv pip install 'transformers>=5.3.0'")
|
|
|
|
with torch.inference_mode():
|
|
ref_audio_tokens = self.audio_tokenizer.encode(audio_signal) # [8, T_ref]
|
|
if ref_audio_tokens.dim() == 3:
|
|
ref_audio_tokens = ref_audio_tokens.squeeze(0) # [8, T_ref]
|
|
|
|
ft = BatchFeature(
|
|
{
|
|
"input_ids": text_tokens,
|
|
"ref_audio_tokens": ref_audio_tokens, # [8, T_ref]
|
|
"ref_audio_len": [ref_audio_tokens.shape[1]],
|
|
}
|
|
)
|
|
return ft
|
|
|
|
def _get_mm_fields_config(
|
|
self,
|
|
hf_inputs: BatchFeature,
|
|
hf_processor_mm_kwargs: Mapping[str, object],
|
|
) -> Mapping[str, MultiModalFieldConfig]:
|
|
return {
|
|
"ref_audio_tokens": MultiModalFieldConfig.batched("audio"),
|
|
"ref_audio_len": MultiModalFieldConfig.batched("audio"),
|
|
}
|
|
|
|
def _hf_processor_applies_updates(
|
|
self,
|
|
prompt_text: str,
|
|
mm_items: MultiModalDataItems,
|
|
hf_processor_mm_kwargs: Mapping[str, object],
|
|
tokenization_kwargs: Mapping[str, object],
|
|
) -> bool:
|
|
return False
|
|
|
|
def _get_prompt_updates(
|
|
self,
|
|
mm_items: MultiModalDataItems,
|
|
hf_processor_mm_kwargs: Mapping[str, object],
|
|
out_mm_kwargs: MultiModalKwargsItems,
|
|
) -> Sequence[PromptUpdate]:
|
|
def insertion_end(item_idx):
|
|
if "audio" in out_mm_kwargs and out_mm_kwargs["audio"]:
|
|
ref_len = out_mm_kwargs["audio"][0]["ref_audio_len"].data[0].item()
|
|
return [1] * ref_len
|
|
return []
|
|
|
|
return [
|
|
PromptInsertion(
|
|
modality="audio",
|
|
target=PromptIndexTargets.start(),
|
|
insertion=insertion_end,
|
|
),
|
|
]
|
|
|
|
|
|
class OmniVoiceDummyInputsBuilder(BaseDummyInputsBuilder[OmniVoiceMultiModalProcessingInfo]):
|
|
def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
|
|
return "Hello, this is a test of the OmniVoice system."
|
|
|
|
def get_dummy_mm_data(
|
|
self,
|
|
seq_len: int,
|
|
mm_counts: Mapping[str, int],
|
|
mm_options: Mapping[str, BaseDummyOptions] | None = None,
|
|
) -> MultiModalDataDict:
|
|
num_audios = mm_counts.get("audio")
|
|
max_prompt_seconds = 10
|
|
prompt_sample_rate = 24000
|
|
target_audio_length = max_prompt_seconds * prompt_sample_rate
|
|
|
|
audio_overrides = mm_options.get("audio") if mm_options else None
|
|
mm_data = {
|
|
"audio": (
|
|
self._get_dummy_audios(
|
|
length=target_audio_length,
|
|
num_audios=num_audios,
|
|
overrides=audio_overrides,
|
|
)[0],
|
|
24000,
|
|
),
|
|
}
|
|
return mm_data
|
|
|
|
def get_dummy_processor_inputs(
|
|
self,
|
|
seq_len: int,
|
|
mm_counts: Mapping[str, int],
|
|
mm_options: Mapping[str, BaseDummyOptions] | None = None,
|
|
) -> ProcessorInputs:
|
|
inputs = super().get_dummy_processor_inputs(seq_len, mm_counts, mm_options)
|
|
inputs.hf_processor_mm_kwargs = {"ref_text": "Testing voice cloning."}
|
|
return inputs
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Main model class
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class OmniVoiceModel(
|
|
nn.Module,
|
|
):
|
|
"""OmniVoice model for vLLM-Omni two-stage pipeline.
|
|
|
|
Routes to generator (Stage 0) or decoder (Stage 1) based on model_stage.
|
|
"""
|
|
|
|
requires_raw_input_tokens = True
|
|
|
|
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
|
|
super().__init__()
|
|
self.config = vllm_config.model_config.hf_config
|
|
self.have_multimodal_outputs = True
|
|
self.model_stage = vllm_config.model_config.model_stage
|
|
self.model_dir = vllm_config.model_config.model
|
|
|
|
if self.model_stage == "omnivoice_generator":
|
|
from vllm_omni.model_executor.models.omnivoice.omnivoice_generator import (
|
|
OmniVoiceGenerator,
|
|
)
|
|
|
|
self.generator = OmniVoiceGenerator(self.config)
|
|
self.model = self.generator
|
|
elif self.model_stage == "omnivoice_decoder":
|
|
from vllm_omni.model_executor.models.omnivoice.omnivoice_decoder import (
|
|
OmniVoiceDecoder,
|
|
)
|
|
|
|
self.decoder = OmniVoiceDecoder(self.config)
|
|
self.model = self.decoder
|
|
else:
|
|
raise ValueError(f"Unsupported model_stage: {self.model_stage}")
|
|
|
|
def embed_input_ids(
|
|
self,
|
|
input_ids: torch.Tensor,
|
|
multimodal_embeddings=None,
|
|
is_multimodal=None,
|
|
) -> torch.Tensor:
|
|
if self.model_stage == "omnivoice_generator":
|
|
# Generator handles its own embedding in forward()
|
|
hidden = int(self.config.llm_hidden_size)
|
|
return torch.zeros((input_ids.shape[0], hidden), device=input_ids.device)
|
|
elif self.model_stage == "omnivoice_decoder":
|
|
hidden = int(self.config.llm_hidden_size)
|
|
return torch.zeros((input_ids.shape[0], hidden), device=input_ids.device)
|
|
else:
|
|
raise RuntimeError(f"embed_input_ids not valid for {self.model_stage}")
|
|
|
|
def forward(
|
|
self,
|
|
input_ids: torch.Tensor,
|
|
positions: torch.Tensor,
|
|
intermediate_tensors: IntermediateTensors | None = None,
|
|
inputs_embeds: torch.Tensor | None = None,
|
|
additional_information: dict[str, object] | None = None,
|
|
**kwargs: object,
|
|
) -> OmniOutput:
|
|
if self.model_stage == "omnivoice_generator":
|
|
return self._forward_generator(input_ids, kwargs)
|
|
elif self.model_stage == "omnivoice_decoder":
|
|
return self._forward_decoder(input_ids, kwargs)
|
|
else:
|
|
raise ValueError(f"Unsupported model_stage: {self.model_stage}")
|
|
|
|
def _forward_generator(self, input_ids: torch.Tensor, kwargs: dict) -> OmniOutput:
|
|
"""Run generator stage: text → 8-codebook audio tokens."""
|
|
runtime_info = kwargs.get("runtime_additional_information", [])
|
|
|
|
if not runtime_info:
|
|
# Profiling / dummy run — return a plain tensor (not OmniOutput)
|
|
# so the v1 model runner's _dummy_run can index into it.
|
|
return torch.zeros(
|
|
(input_ids.shape[0], self.config.llm_hidden_size),
|
|
device=input_ids.device,
|
|
dtype=torch.float32,
|
|
)
|
|
|
|
info = runtime_info[0]
|
|
device = input_ids.device
|
|
num_codebooks = self.config.num_audio_codebook
|
|
mask_id = self.config.audio_mask_id
|
|
|
|
# Extract text tokens from input_ids
|
|
text_tokens = input_ids # [N_text]
|
|
text_len = text_tokens.shape[0]
|
|
|
|
# Estimate target length using RuleDurationEstimator
|
|
# (same formula as reference OmniVoice: weight * 25 / 14.1)
|
|
from vllm_omni.model_executor.models.omnivoice.duration import (
|
|
RuleDurationEstimator,
|
|
)
|
|
|
|
if not hasattr(self, "_duration_estimator"):
|
|
self._duration_estimator = RuleDurationEstimator()
|
|
raw_text = info.get("raw_text", "")
|
|
if raw_text:
|
|
target_len = self._duration_estimator.estimate_duration(raw_text, "Nice to meet you.", 25)
|
|
target_len = max(1, int(target_len))
|
|
else:
|
|
# Fallback: use character weight formula on text tokens
|
|
# approximate ~1.77 frames per text token (25/14.1)
|
|
target_len = max(int(text_len * 1.77), 25)
|
|
|
|
# Get reference audio tokens if available
|
|
ref_audio_tokens = info.get("ref_audio_tokens", None)
|
|
|
|
# Build input_ids tensor: [2*B, 8, S]
|
|
# B=1, conditional + unconditional
|
|
|
|
# Replicate text tokens across 8 codebooks
|
|
text_ids = text_tokens.unsqueeze(0).repeat(num_codebooks, 1) # [8, N_text]
|
|
|
|
# Target: all MASK
|
|
target_ids = torch.full((num_codebooks, target_len), mask_id, dtype=torch.long, device=device)
|
|
|
|
# Conditional: [text] [ref_audio?] [target_mask]
|
|
cond_audio_start = text_ids.shape[1]
|
|
if ref_audio_tokens is not None:
|
|
ref_tokens = ref_audio_tokens.to(device) # [8, T_ref]
|
|
cond_ids = torch.cat([text_ids, ref_tokens, target_ids], dim=1)
|
|
else:
|
|
cond_ids = torch.cat([text_ids, target_ids], dim=1)
|
|
|
|
cond_len = cond_ids.shape[1]
|
|
|
|
# Unconditional: [target_mask only]
|
|
uncond_ids = target_ids.clone()
|
|
uncond_len = target_len
|
|
|
|
# Pad to same length
|
|
max_len = max(cond_len, uncond_len)
|
|
if cond_len < max_len:
|
|
pad = torch.full(
|
|
(num_codebooks, max_len - cond_len),
|
|
mask_id,
|
|
dtype=torch.long,
|
|
device=device,
|
|
)
|
|
cond_ids = torch.cat([cond_ids, pad], dim=1)
|
|
if uncond_len < max_len:
|
|
pad = torch.full(
|
|
(num_codebooks, max_len - uncond_len),
|
|
mask_id,
|
|
dtype=torch.long,
|
|
device=device,
|
|
)
|
|
uncond_ids = torch.cat([uncond_ids, pad], dim=1)
|
|
|
|
batch_input_ids = torch.stack([cond_ids, uncond_ids], dim=0) # [2, 8, max_len]
|
|
|
|
# Audio mask: True for audio positions
|
|
batch_audio_mask = torch.zeros((2, max_len), dtype=torch.bool, device=device)
|
|
batch_audio_mask[0, cond_audio_start:cond_len] = True
|
|
batch_audio_mask[1, :uncond_len] = True
|
|
|
|
# Attention mask: [2, 1, S, S]
|
|
batch_attention_mask = torch.zeros((2, 1, max_len, max_len), dtype=torch.bool, device=device)
|
|
batch_attention_mask[0, :, :cond_len, :cond_len] = True
|
|
batch_attention_mask[1, :, :uncond_len, :uncond_len] = True
|
|
|
|
# Run iterative generation
|
|
tokens = self.generator(
|
|
input_ids=batch_input_ids,
|
|
audio_mask=batch_audio_mask,
|
|
attention_mask=batch_attention_mask,
|
|
target_lens=[target_len],
|
|
num_step=self.config.num_step,
|
|
guidance_scale=self.config.guidance_scale,
|
|
t_shift=self.config.t_shift,
|
|
layer_penalty_factor=self.config.layer_penalty_factor,
|
|
position_temperature=self.config.position_temperature,
|
|
class_temperature=self.config.class_temperature,
|
|
) # [1, 8, target_len]
|
|
|
|
return OmniOutput(
|
|
text_hidden_states=None,
|
|
multimodal_outputs={"audio_tokens": tokens},
|
|
)
|
|
|
|
def _forward_decoder(self, input_ids: torch.Tensor, kwargs: dict) -> OmniOutput:
|
|
"""Run decoder stage: 8-codebook tokens → audio waveform."""
|
|
runtime_info = kwargs.get("runtime_additional_information", [])
|
|
|
|
if not runtime_info:
|
|
# Profiling / dummy run — return plain tensor for v1 runner compat
|
|
return torch.zeros(
|
|
(input_ids.shape[0], self.config.llm_hidden_size),
|
|
device=input_ids.device,
|
|
dtype=torch.float32,
|
|
)
|
|
|
|
info = runtime_info[0]
|
|
audio_tokens = info.get("audio_tokens", None)
|
|
|
|
if audio_tokens is None:
|
|
raise RuntimeError("No audio_tokens received from generator stage")
|
|
|
|
if isinstance(audio_tokens, np.ndarray):
|
|
audio_tokens = torch.from_numpy(audio_tokens)
|
|
|
|
# audio_tokens: [B, 8, T]; buffer may be CPU — move to decoder weights
|
|
if audio_tokens.dim() == 2:
|
|
audio_tokens = audio_tokens.unsqueeze(0) # Add batch dim
|
|
|
|
dec_device = next(self.decoder.parameters()).device
|
|
audio_tokens = audio_tokens.to(device=dec_device, dtype=torch.long)
|
|
|
|
tts_speech = self.decoder(audio_tokens)
|
|
|
|
return OmniOutput(
|
|
text_hidden_states=None,
|
|
multimodal_outputs={
|
|
"audio": tts_speech,
|
|
"sr": self.config.sample_rate,
|
|
},
|
|
)
|
|
|
|
def _resolve_model_dir(self) -> str:
|
|
"""Resolve model directory to local path (handles HF hub IDs)."""
|
|
model_dir = self.model_dir
|
|
if os.path.isdir(model_dir):
|
|
return model_dir
|
|
# HF hub model ID — resolve to local cache
|
|
from huggingface_hub import snapshot_download
|
|
|
|
return snapshot_download(model_dir)
|
|
|
|
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
|
try:
|
|
device = next(self.parameters()).device
|
|
except StopIteration:
|
|
device = current_omni_platform.get_torch_device()
|
|
|
|
model_dir = self._resolve_model_dir()
|
|
|
|
if self.model_stage == "omnivoice_generator":
|
|
self.generator.load_weights(model_dir, device)
|
|
elif self.model_stage == "omnivoice_decoder":
|
|
self.decoder.load_weights(model_dir, device)
|
|
else:
|
|
raise ValueError(f"{self.model_stage} not supported!")
|