121 lines
5.4 KiB
Python
121 lines
5.4 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""Voxtral TTS serving adapter."""
|
|
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
from vllm.utils.async_utils import make_async
|
|
|
|
from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter
|
|
from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest, apply_max_new_tokens
|
|
from vllm_omni.entrypoints.openai.tts_adapters.capabilities import load_supported_speakers
|
|
|
|
if TYPE_CHECKING:
|
|
from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest
|
|
|
|
|
|
@register_tts_adapter
|
|
class VoxtralTTSAdapter(ARTTSAdapter):
|
|
stage_keys = frozenset({"audio_generation"})
|
|
name = "voxtral_tts"
|
|
|
|
def __init__(self, ctx) -> None:
|
|
super().__init__(ctx)
|
|
self._build_prompt_async = make_async(self._build_prompt, executor=getattr(ctx.server, "_tts_executor", None))
|
|
self._encoder_loaded = None
|
|
|
|
def _build_prompt(self, request: "OpenAICreateSpeechRequest") -> dict[str, Any]:
|
|
"""Build a Voxtral prompt for either a preset voice or inline reference audio."""
|
|
from mistral_common.protocol.speech.request import SpeechRequest
|
|
from vllm.inputs import tokens_input
|
|
|
|
ref_audio = request.ref_audio
|
|
if not request.voice and not ref_audio:
|
|
raise ValueError("Voxtral requires either a voice name or ref_audio.")
|
|
|
|
if ref_audio is not None and not self._encoder_loaded:
|
|
validation_error = (
|
|
"Voice cloning with 'ref_audio' requires the full Voxtral checkpoint "
|
|
"with encoder weights. The open-source variant only supports "
|
|
"preset voices via the 'voice' parameter."
|
|
)
|
|
raise ValueError(validation_error)
|
|
|
|
server = self.ctx.server
|
|
if server._tts_tokenizer is None:
|
|
from vllm.tokenizers import cached_tokenizer_from_config
|
|
|
|
server._tts_tokenizer = cached_tokenizer_from_config(self.ctx.engine_client.model_config).instruct
|
|
|
|
if isinstance(ref_audio, str) and ref_audio.startswith("data:"):
|
|
_, _, ref_audio = ref_audio.partition(",")
|
|
|
|
if request.voice is not None:
|
|
tokenized = server._tts_tokenizer.encode_speech_request(
|
|
SpeechRequest(input=request.input, voice=request.voice)
|
|
)
|
|
prompt = tokens_input(prompt_token_ids=tokenized.tokens)
|
|
prompt["additional_information"] = {"voice": [request.voice]}
|
|
return prompt
|
|
|
|
tokenized = server._tts_tokenizer.encode_speech_request(SpeechRequest(input=request.input, ref_audio=ref_audio))
|
|
audio = tokenized.audios[0]
|
|
return {
|
|
"prompt_token_ids": tokenized.tokens,
|
|
"multi_modal_data": {"audio": [(audio.audio_array, audio.sampling_rate)]},
|
|
}
|
|
|
|
def validate(self, request: "OpenAICreateSpeechRequest") -> str | None:
|
|
"""Validate Voxtral TTS request parameters. Returns error message or None."""
|
|
server = self.ctx.server
|
|
if not request.input or not request.input.strip():
|
|
return "Input text cannot be empty"
|
|
|
|
# Voxtral TTS requires either a preset voice or ref_audio for voice cloning.
|
|
if request.voice is None and request.ref_audio is None:
|
|
return "Either 'voice' (preset speaker) or 'ref_audio' (voice cloning) must be provided"
|
|
|
|
if request.ref_audio is not None:
|
|
fmt_err = server._validate_ref_audio_format(request.ref_audio)
|
|
if fmt_err:
|
|
return fmt_err
|
|
|
|
if request.voice is not None:
|
|
request.voice = request.voice.lower()
|
|
available_speakers = server._get_available_speakers()
|
|
if available_speakers and request.voice not in available_speakers:
|
|
return f"Invalid speaker '{request.voice}'. Supported: {', '.join(sorted(available_speakers))}"
|
|
|
|
if request.max_new_tokens is not None:
|
|
if request.max_new_tokens < self.max_new_tokens_min:
|
|
return f"max_new_tokens must be at least {self.max_new_tokens_min}"
|
|
if request.max_new_tokens > self.max_new_tokens_max:
|
|
return f"max_new_tokens cannot exceed {self.max_new_tokens_max}"
|
|
|
|
return None
|
|
|
|
async def build(
|
|
self, request: "OpenAICreateSpeechRequest", sampling_params_list: list, has_inline_ref_audio: bool
|
|
) -> PreparedRequest:
|
|
# Cache encoder_loaded property to avoid RPC call on every request with ref_audio
|
|
if self._encoder_loaded is None:
|
|
loaded = await self.ctx.server.engine_client.collective_rpc("encoder_loaded")
|
|
# Flatten results before checking if encoder loaded property is set
|
|
self._encoder_loaded = any(sum(loaded, []))
|
|
|
|
prompt = await self._build_prompt_async(request)
|
|
return PreparedRequest(prompt=prompt, tts_params={}, model_type="voxtral_tts")
|
|
|
|
def _load_supported_speakers(self) -> set[str]:
|
|
config = self.ctx.engine_client.model_config.hf_config.audio_config
|
|
return load_supported_speakers(self.ctx.engine_client, config)
|
|
|
|
def apply_sampling_overrides(
|
|
self,
|
|
sampling_params_list: list,
|
|
request: "OpenAICreateSpeechRequest",
|
|
prompt: dict[str, Any] | None = None,
|
|
request_id: str | None = None,
|
|
) -> list:
|
|
return apply_max_new_tokens(sampling_params_list, request)
|