2988 lines
137 KiB
Python
2988 lines
137 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
import asyncio
|
|
import base64
|
|
import hashlib
|
|
import io
|
|
import json
|
|
import math
|
|
import os
|
|
import re
|
|
import struct
|
|
import time
|
|
from collections import OrderedDict
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from contextlib import aclosing
|
|
from http import HTTPStatus
|
|
from pathlib import Path
|
|
from typing import Any, cast
|
|
from urllib.parse import urlparse
|
|
from urllib.request import url2pathname
|
|
|
|
import anyio
|
|
import numpy as np
|
|
import soundfile as sf
|
|
import torch
|
|
from fastapi import HTTPException, Request, UploadFile
|
|
from fastapi.responses import Response, StreamingResponse
|
|
from vllm.entrypoints.generate.base.protocol import RequestResponseMetadata
|
|
from vllm.entrypoints.generate.base.serving import GenerateBaseServing as OpenAIServing
|
|
from vllm.entrypoints.launchers.launcher import terminate_if_errored
|
|
from vllm.entrypoints.serve.engine.protocol import ErrorResponse
|
|
from vllm.logger import init_logger
|
|
from vllm.multimodal.media import MediaConnector
|
|
from vllm.sampling_params import RequestOutputKind, SamplingParams
|
|
from vllm.utils import random_uuid
|
|
from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError
|
|
|
|
from vllm_omni.config.speech_cache import SpeechCacheConfig
|
|
from vllm_omni.config.stage_config import StagePipelineConfig
|
|
from vllm_omni.entrypoints.openai.audio_utils_mixin import AudioMixin, StreamingAudioResampler
|
|
from vllm_omni.entrypoints.openai.errors import InvalidPresetVoiceReferenceError, InvalidVoiceReferenceError
|
|
from vllm_omni.entrypoints.openai.protocol.audio import (
|
|
AudioResponse,
|
|
BatchSpeechRequest,
|
|
BatchSpeechResponse,
|
|
CreateAudio,
|
|
OpenAICreateSpeechRequest,
|
|
RegisteredVoiceReference,
|
|
SpeechBatchItem,
|
|
SpeechBatchItemResult,
|
|
SpeechInputTokenDetails,
|
|
SpeechTokenUsage,
|
|
)
|
|
from vllm_omni.entrypoints.openai.speech_usage import (
|
|
SpeechOutputTokenCounter,
|
|
build_speech_usage,
|
|
qwen3_tts_input_token_details,
|
|
)
|
|
from vllm_omni.entrypoints.openai.tts_adapters import (
|
|
OutputPolicy,
|
|
SpeechServingContext,
|
|
TTSGenerationError,
|
|
all_tts_stage_keys,
|
|
detect_tts_model_type,
|
|
resolve_adapter,
|
|
tts_entry_stage_archs,
|
|
)
|
|
from vllm_omni.entrypoints.utils import coerce_param_message_types
|
|
from vllm_omni.metrics.modality import observe_audio_first_packet, observe_audio_streaming_finalize
|
|
from vllm_omni.outputs import OmniRequestOutput
|
|
from vllm_omni.utils.speaker_cache import get_speaker_cache
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
_SPEECH_USAGE_INPUT_TOKENS_HEADER = "X-VLLM-OMNI-INPUT-TOKENS"
|
|
_SPEECH_USAGE_OUTPUT_TOKENS_HEADER = "X-VLLM-OMNI-OUTPUT-TOKENS"
|
|
_SPEECH_USAGE_TOTAL_TOKENS_HEADER = "X-VLLM-OMNI-TOTAL-TOKENS"
|
|
_SPEECH_USAGE_INPUT_TEXT_TOKENS_HEADER = "X-VLLM-OMNI-INPUT-TEXT-TOKENS"
|
|
_SPEECH_USAGE_INPUT_AUDIO_TOKENS_HEADER = "X-VLLM-OMNI-INPUT-AUDIO-TOKENS"
|
|
# Raw ``pcm`` bodies carry no format header and the rate/channel count is
|
|
# model-native, so responses state the s16le layout explicitly.
|
|
SPEECH_AUDIO_SAMPLE_RATE_HEADER = "X-Audio-Sample-Rate"
|
|
SPEECH_AUDIO_CHANNELS_HEADER = "X-Audio-Channels"
|
|
|
|
|
|
def _stage_speech_metadata(stage: Any) -> tuple[str | None, str | None, str | None]:
|
|
"""Return speech routing metadata from typed or legacy stage configs."""
|
|
topology = getattr(stage, "stage_pipeline_config", None)
|
|
if isinstance(topology, StagePipelineConfig):
|
|
return (
|
|
getattr(topology, "model_stage", None),
|
|
getattr(getattr(stage, "model_config", None), "model_arch", None),
|
|
getattr(stage, "worker_type", None),
|
|
)
|
|
engine_args = getattr(stage, "engine_args", None)
|
|
return (
|
|
getattr(engine_args, "model_stage", None),
|
|
getattr(engine_args, "model_arch", None),
|
|
getattr(engine_args, "worker_type", None),
|
|
)
|
|
|
|
|
|
# TTS Configuration
|
|
#
|
|
# The stage-key -> model-type mapping is NOT declared here: it is derived from
|
|
# the ``stage_keys`` / ``model_archs`` each adapter declares, via
|
|
# ``tts_adapters.detect_tts_model_type``. Adding a TTS model must not require an
|
|
# edit to this module.
|
|
#
|
|
# Audex contract: zero-codec / invalid generations arrive as empty terminal
|
|
# payloads and must fail the request, never serialize as a successful empty
|
|
# WAV. Covers both the TTS ("audex") and TTA ("audex_tta") pipelines.
|
|
_AUDEX_NO_AUDIO_GUARD_MODEL_TYPES = frozenset({"audex", "audex_tta"})
|
|
_REF_AUDIO_MIN_DURATION = 1.0 # seconds
|
|
_REF_AUDIO_MAX_DURATION = 30.0 # seconds
|
|
_REF_AUDIO_METADATA_FETCH_ATTEMPTS = 3
|
|
_REMOTE_REF_AUDIO_SCHEMES = frozenset({"http", "https", "data"})
|
|
_TTS_MAX_INSTRUCTIONS_LENGTH = 500
|
|
_DEFAULT_VOICE_NAME = "default"
|
|
|
|
|
|
def _is_default_voice(voice, supported_speakers):
|
|
"""Check if a lowercased voice name is the placeholder default and not
|
|
an actual registered/built-in speaker."""
|
|
return voice == _DEFAULT_VOICE_NAME and voice not in supported_speakers
|
|
|
|
|
|
def _create_wav_header(sample_rate: int, num_channels: int = 1, bits_per_sample: int = 16) -> bytes:
|
|
"""Create a WAV header with placeholder size values for streaming.
|
|
|
|
Uses 0xFFFFFFFF as placeholder for data size fields, which is accepted
|
|
by most audio clients and matches OpenAI's streaming WAV implementation.
|
|
|
|
Args:
|
|
sample_rate: Audio sample rate in Hz
|
|
num_channels: Number of audio channels (1 for mono, 2 for stereo)
|
|
bits_per_sample: Bits per sample (typically 16)
|
|
|
|
Returns:
|
|
44-byte WAV header as bytes
|
|
"""
|
|
byte_rate = sample_rate * num_channels * bits_per_sample // 8
|
|
block_align = num_channels * bits_per_sample // 8
|
|
|
|
# Use 0xFFFFFFFF as placeholder for unknown size (streaming)
|
|
placeholder_size = 0xFFFFFFFF
|
|
|
|
# ref https://docs.fileformat.com/audio/wav/
|
|
header = struct.pack(
|
|
"<4sI4s4sIHHIIHH4sI",
|
|
b"RIFF", # ChunkID
|
|
placeholder_size, # ChunkSize (placeholder)
|
|
b"WAVE", # Format
|
|
b"fmt ", # Subchunk1ID
|
|
16, # Subchunk1Size (16 for PCM)
|
|
1, # AudioFormat (1 for PCM)
|
|
num_channels, # NumChannels
|
|
sample_rate, # SampleRate
|
|
byte_rate, # ByteRate
|
|
block_align, # BlockAlign
|
|
bits_per_sample, # BitsPerSample
|
|
b"data", # Subchunk2ID
|
|
placeholder_size, # Subchunk2Size (placeholder)
|
|
)
|
|
|
|
return header
|
|
|
|
|
|
def _infer_audio_num_channels(audio: np.ndarray) -> int:
|
|
"""Infer channel count before streaming PCM bytes are wrapped as WAV."""
|
|
if audio.ndim == 3 and audio.shape[0] == 1:
|
|
audio = audio[0]
|
|
if audio.ndim == 2:
|
|
if audio.shape[0] in (1, 2):
|
|
return int(audio.shape[0])
|
|
if audio.shape[1] in (1, 2):
|
|
return int(audio.shape[1])
|
|
return 1
|
|
|
|
|
|
def _audio_format_headers(audio_format: dict[str, int] | None) -> dict[str, str]:
|
|
"""Response headers describing the emitted audio's rate and channel count."""
|
|
if not audio_format:
|
|
return {}
|
|
return {
|
|
SPEECH_AUDIO_SAMPLE_RATE_HEADER: str(int(audio_format["sample_rate"])),
|
|
SPEECH_AUDIO_CHANNELS_HEADER: str(int(audio_format["channels"])),
|
|
# Without this, CORS hides both headers from cross-origin browser clients.
|
|
"Access-Control-Expose-Headers": f"{SPEECH_AUDIO_SAMPLE_RATE_HEADER}, {SPEECH_AUDIO_CHANNELS_HEADER}",
|
|
}
|
|
|
|
|
|
class _PrimedAudioStream:
|
|
"""Async iterator over an audio stream whose first chunk was already pulled.
|
|
|
|
Raw audio streams pull the first chunk before the HTTP response starts so
|
|
its headers can state the real sample rate and channel count.
|
|
"""
|
|
|
|
def __init__(self, chunks: Any, first: Any, has_first: bool) -> None:
|
|
self._chunks = chunks
|
|
self._first = first
|
|
self._has_first = has_first
|
|
|
|
@classmethod
|
|
async def start(cls, chunks: Any) -> "_PrimedAudioStream":
|
|
try:
|
|
first = await anext(chunks)
|
|
except StopAsyncIteration:
|
|
return cls(chunks, None, False)
|
|
except BaseException:
|
|
with anyio.CancelScope(shield=True):
|
|
await chunks.aclose()
|
|
raise
|
|
return cls(chunks, first, True)
|
|
|
|
def __aiter__(self) -> "_PrimedAudioStream":
|
|
return self
|
|
|
|
async def __anext__(self) -> Any:
|
|
if self._has_first:
|
|
first, self._first, self._has_first = self._first, None, False
|
|
return first
|
|
return await anext(self._chunks)
|
|
|
|
async def aclose(self) -> None:
|
|
self._first, self._has_first = None, False
|
|
await self._chunks.aclose()
|
|
|
|
|
|
def _encoded_audio_format(audio_response: Any, sample_rate: int, audio: np.ndarray) -> dict[str, int]:
|
|
"""Rate and channel count of an encoded chunk, preferring encoder metadata."""
|
|
metadata = getattr(audio_response, "audio_metadata", None)
|
|
if metadata is not None:
|
|
return {"sample_rate": int(metadata.sample_rate_hz), "channels": int(metadata.channels)}
|
|
return {"sample_rate": int(sample_rate), "channels": _infer_audio_num_channels(np.asarray(audio))}
|
|
|
|
|
|
def _sanitize_filename(filename: str) -> str:
|
|
"""Sanitize filename to prevent path traversal attacks.
|
|
|
|
Only allows alphanumeric characters, underscores, hyphens, and dots.
|
|
Replaces any other characters with underscores.
|
|
"""
|
|
# Remove any path components
|
|
filename = os.path.basename(filename)
|
|
# Replace any non-alphanumeric, underscore, hyphen, or dot with underscore
|
|
sanitized = re.sub(r"[^a-zA-Z0-9_.\-]", "_", filename)
|
|
# Ensure filename is not empty
|
|
if not sanitized:
|
|
sanitized = "file"
|
|
# Limit length to prevent potential issues
|
|
if len(sanitized) > 255:
|
|
sanitized = sanitized[:255]
|
|
return sanitized
|
|
|
|
|
|
def _validate_speaker_name(name: str) -> str:
|
|
"""Trim and reject empty / path-separator / NUL / reserved voice names."""
|
|
trimmed = (name or "").strip()
|
|
if not trimmed or trimmed in (".", "..") or any(c in trimmed for c in "/\\\x00"):
|
|
raise ValueError(f"Invalid voice name {name!r}: must be non-empty, no path separators or NUL")
|
|
return trimmed
|
|
|
|
|
|
def _validate_path_within_directory(file_path: Path, directory: Path) -> bool:
|
|
"""Validate that file_path is within the specified directory.
|
|
|
|
Prevents path traversal attacks by ensuring the resolved path
|
|
is within the target directory.
|
|
"""
|
|
try:
|
|
# Resolve both paths to absolute paths
|
|
file_path_resolved = file_path.resolve()
|
|
directory_resolved = directory.resolve()
|
|
# Check if file_path is within directory
|
|
return directory_resolved in file_path_resolved.parents or directory_resolved == file_path_resolved
|
|
except Exception:
|
|
return False
|
|
|
|
|
|
class _SpeechStreamingResponse(StreamingResponse):
|
|
async def stream_response(self, send) -> None:
|
|
try:
|
|
await super().stream_response(send)
|
|
finally:
|
|
# A disconnect during send leaves the iterator suspended at yield.
|
|
# Close it explicitly, even inside Starlette's cancelled scope.
|
|
with anyio.CancelScope(shield=True):
|
|
await cast(Any, self.body_iterator).aclose()
|
|
|
|
|
|
class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
|
|
_diffusion_mode: bool = False
|
|
_media_connector: MediaConnector | None = None
|
|
_allowed_local_media_path: str = ""
|
|
_tts_executor: ThreadPoolExecutor | None = None
|
|
|
|
def _init_speaker_storage(self) -> None:
|
|
"""Initialize speaker storage + cache, restoring any persisted uploads."""
|
|
speaker_samples_dir = os.environ.get("SPEAKER_SAMPLES_DIR", os.path.expanduser("~/.cache/vllm-omni/speakers"))
|
|
self.uploaded_speakers_dir = Path(speaker_samples_dir).expanduser()
|
|
self.uploaded_speakers_dir.mkdir(parents=True, exist_ok=True)
|
|
_raw_cap = os.environ.get("SPEAKER_MAX_UPLOADED", "")
|
|
try:
|
|
self._max_uploaded_speakers = int(_raw_cap) if _raw_cap else 1000
|
|
except ValueError:
|
|
logger.warning("Invalid SPEAKER_MAX_UPLOADED=%r; using default 1000", _raw_cap)
|
|
self._max_uploaded_speakers = 1000
|
|
_policy = os.environ.get("VLLM_OMNI_SPEAKER_REGISTRATION_POLICY", "overwrite").lower()
|
|
if _policy not in ("overwrite", "immutable"):
|
|
raise ValueError(
|
|
f"Invalid VLLM_OMNI_SPEAKER_REGISTRATION_POLICY={_policy!r}; expected 'overwrite' or 'immutable'."
|
|
)
|
|
self._registration_policy = _policy
|
|
self.uploaded_speakers: dict[str, dict[str, Any]] = {}
|
|
self._ref_audio_data_url_cache: dict[str, str] = {}
|
|
self._ref_audio_resolve_cache: OrderedDict[str, tuple[np.ndarray, int, int, str]] = OrderedDict()
|
|
self._ref_audio_resolve_cache_bytes = 0
|
|
config = self.speech_cache_config
|
|
self._ref_audio_resolve_cache_max_entries = config.resolve_max_entries
|
|
self._ref_audio_resolve_cache_max_bytes = config.resolve_max_bytes
|
|
logger.info("Speech cache configuration: %s", config)
|
|
# Readiness is keyed by (artifact_key, x_vector_only). An x-vector-only
|
|
# request caches a speaker embedding but no ref_code, so its artifact
|
|
# must not satisfy a later ICL request that needs ref_code (#5049).
|
|
self._ref_audio_model_artifact_ready: set[tuple[str, bool]] = set()
|
|
self._request_ref_audio_artifact_keys: dict[str, tuple[str, bool]] = {}
|
|
self._speaker_cache = get_speaker_cache(max_bytes=config.speaker_max_bytes)
|
|
self._last_upload_ts = 0
|
|
self._upload_lock = asyncio.Lock()
|
|
self._restore_uploaded_speakers()
|
|
logger.info(
|
|
"Speaker storage: dir=%s, max_speakers=%d, restored=%d",
|
|
self.uploaded_speakers_dir,
|
|
self._max_uploaded_speakers,
|
|
len(self.uploaded_speakers),
|
|
)
|
|
|
|
def _next_upload_timestamp(self) -> int:
|
|
ts = max(int(time.time()), self._last_upload_ts + 1)
|
|
self._last_upload_ts = ts
|
|
return ts
|
|
|
|
_META_SCALAR_INT_KEYS: tuple[str, ...] = (
|
|
"created_at",
|
|
"file_size",
|
|
"sample_rate",
|
|
"embedding_dim",
|
|
)
|
|
|
|
@classmethod
|
|
def _speaker_metadata_to_header(cls, speaker_data: dict[str, Any]) -> dict[str, str]:
|
|
"""Serialize a speaker_data dict into safetensors' ``dict[str, str]`` header."""
|
|
header: dict[str, str] = {}
|
|
for k, v in speaker_data.items():
|
|
if v is None:
|
|
continue
|
|
# file_path is re-derived from the path on load; don't persist it.
|
|
if k == "file_path":
|
|
continue
|
|
header[k] = str(v)
|
|
return header
|
|
|
|
@classmethod
|
|
def _speaker_metadata_from_header(cls, header: dict[str, str], file_path: str) -> dict[str, Any]:
|
|
"""Reverse of :meth:`_speaker_metadata_to_header`: coerce ints back and re-inject file_path."""
|
|
data: dict[str, Any] = dict(header)
|
|
for k in cls._META_SCALAR_INT_KEYS:
|
|
if k in data:
|
|
try:
|
|
data[k] = int(data[k])
|
|
except ValueError:
|
|
logger.warning(
|
|
"Speaker metadata %r in %s is not a valid int (got %r); leaving as string",
|
|
k,
|
|
file_path,
|
|
data[k],
|
|
)
|
|
data["file_path"] = file_path
|
|
return data
|
|
|
|
def _restore_uploaded_speakers(self) -> None:
|
|
"""Scan ``uploaded_speakers_dir`` for safetensors files and rebuild state."""
|
|
try:
|
|
from safetensors import safe_open
|
|
except ImportError:
|
|
logger.warning("safetensors unavailable; uploaded voices will not persist across restarts")
|
|
return
|
|
|
|
restored = 0
|
|
for path in sorted(self.uploaded_speakers_dir.glob("*.safetensors")):
|
|
try:
|
|
with safe_open(str(path), framework="pt") as f:
|
|
header = dict(f.metadata() or {})
|
|
except Exception as e:
|
|
logger.warning("Could not read voice file %s: %s", path, e)
|
|
continue
|
|
voice_name_lower = header.get("voice_name_lower") or header.get("name", "").lower()
|
|
if not voice_name_lower:
|
|
logger.warning("Voice file %s has no voice name in metadata; skipping", path)
|
|
continue
|
|
speaker_data = self._speaker_metadata_from_header(header, str(path))
|
|
speaker_data.setdefault("name", voice_name_lower)
|
|
speaker_data.setdefault("file_size", int(path.stat().st_size))
|
|
self.uploaded_speakers[voice_name_lower] = speaker_data
|
|
self._last_upload_ts = max(self._last_upload_ts, int(speaker_data.get("created_at", 0)))
|
|
restored += 1
|
|
if restored:
|
|
logger.info("Restored %d uploaded voice(s) from %s", restored, self.uploaded_speakers_dir)
|
|
|
|
@classmethod
|
|
def for_diffusion(
|
|
cls,
|
|
diffusion_engine: "Any",
|
|
model_name: str,
|
|
stage_configs: "list[Any] | None" = None,
|
|
allowed_local_media_path: str = "",
|
|
allowed_media_domains: list[str] | None = None,
|
|
speech_cache_config: SpeechCacheConfig | None = None,
|
|
) -> "OmniOpenAIServingSpeech":
|
|
"""Create a speech serving instance for pure diffusion TTS models.
|
|
|
|
Bypasses OpenAIServing.__init__ which requires a fully configured
|
|
engine client that pure diffusion engines don't provide.
|
|
"""
|
|
instance = cls.__new__(cls)
|
|
instance.speech_cache_config = speech_cache_config or SpeechCacheConfig()
|
|
instance._diffusion_mode = True
|
|
instance._diffusion_engine = diffusion_engine
|
|
instance._diffusion_model_name = model_name
|
|
instance._diffusion_stage_configs = stage_configs
|
|
instance._allowed_local_media_path = allowed_local_media_path
|
|
instance._media_connector = MediaConnector(
|
|
allowed_local_media_path=allowed_local_media_path,
|
|
allowed_media_domains=allowed_media_domains,
|
|
)
|
|
instance._tts_model_type = "omnivoice"
|
|
instance._is_tts = False
|
|
# Diffusion-only instances don't have a TTS stage; set None so any
|
|
# ``_is_tts_model()`` / ``_tts_stage`` access doesn't raise AttributeError.
|
|
instance._tts_stage = None
|
|
instance._adapter = None
|
|
instance._init_speaker_storage()
|
|
return instance
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
self._media_connector = None
|
|
self._allowed_local_media_path = ""
|
|
self.model_name = kwargs.pop("model_name", None)
|
|
self.speech_cache_config = kwargs.pop("speech_cache_config", None) or SpeechCacheConfig()
|
|
# True when the server was launched with --forced-aligner (a pooling
|
|
# aligner stage is appended to the pipeline). Gates word_timestamps.
|
|
self.forced_aligner_enabled: bool = bool(kwargs.pop("forced_aligner_enabled", False))
|
|
super().__init__(*args, **kwargs)
|
|
self._init_speaker_storage()
|
|
|
|
# Find and cache the TTS stage (if any) during initialization
|
|
self._tts_stage = self._find_tts_stage()
|
|
self._is_tts = self._tts_stage is not None
|
|
|
|
# Determine TTS model type or None
|
|
self._tts_model_type = self._detect_tts_model_type()
|
|
|
|
# Shared executor for blocking adapter preprocessing. It must exist
|
|
# before adapter construction so adapters can create their make_async
|
|
# wrappers during their own lifecycle initialization.
|
|
self._tts_executor = ThreadPoolExecutor(max_workers=1)
|
|
# Resolve the per-model serving adapter (RFC #4327), keyed on the
|
|
# detected model-type. Every dedicated TTS model has an adapter; the
|
|
# adapter owns request validation, prompt/param building, capability
|
|
# metadata, and sampling overrides. The model-type label remains in the
|
|
# orchestrator for compatibility during this incremental migration.
|
|
self._adapter = None
|
|
if self._tts_stage is not None:
|
|
adapter_cls = resolve_adapter(self._tts_model_type)
|
|
if adapter_cls is not None:
|
|
ctx = SpeechServingContext(server=self, engine_client=self.engine_client)
|
|
self._adapter = adapter_cls(ctx)
|
|
logger.info("Resolved TTS serving adapter: %s", adapter_cls.__name__)
|
|
|
|
adapter = self._adapter
|
|
if adapter is not None:
|
|
adapter.load_capabilities()
|
|
self._drop_shadowing_uploads()
|
|
available_speakers = self._get_available_speakers()
|
|
logger.info("Loaded %d supported speakers: %s", len(available_speakers), sorted(available_speakers))
|
|
|
|
# Cache TTS configuration values (computed once, reused per request)
|
|
self._max_instructions_length = self._compute_max_instructions_length()
|
|
|
|
self._tts_tokenizer = None
|
|
# Per-request output policy from the adapter's PreparedRequest. Keyed
|
|
# by request_id so concurrent speech requests cannot overwrite each
|
|
# other. Stored only for non-streaming requests and consumed (popped)
|
|
# by the non-streaming accumulator in `_generate_audio_bytes`.
|
|
self._speech_output_policies: dict[str, OutputPolicy] = {}
|
|
|
|
# Batch configuration
|
|
self._batch_max_items: int = getattr(self.engine_client, "tts_batch_max_items", 32)
|
|
|
|
def _get_tts_adapter(self):
|
|
"""Return the per-model serving adapter for the current ``_tts_model_type``.
|
|
|
|
Pure-diffusion speech uses its dedicated request path and does not
|
|
resolve adapters for AR-stage TTS models.
|
|
|
|
Resolved lazily (rebuilt if ``_tts_model_type`` changed since the cached
|
|
instance was built) so callers that set ``_tts_model_type`` after
|
|
construction still dispatch to the matching adapter. In production
|
|
``_tts_model_type`` is fixed at init, so the cached instance is reused.
|
|
"""
|
|
if self._diffusion_mode:
|
|
return None
|
|
|
|
adapter_cls = resolve_adapter(self._tts_model_type)
|
|
if adapter_cls is None:
|
|
self._adapter = None
|
|
return None
|
|
if self._adapter is None or type(self._adapter) is not adapter_cls:
|
|
ctx = SpeechServingContext(server=self, engine_client=self.engine_client)
|
|
self._adapter = adapter_cls(ctx)
|
|
self._adapter.load_capabilities()
|
|
self._drop_shadowing_uploads()
|
|
return self._adapter
|
|
|
|
def _uses_native_speed_control(self) -> bool:
|
|
adapter = self._get_tts_adapter()
|
|
return bool(adapter is not None and adapter.native_speed_control)
|
|
|
|
def _get_available_speakers(self) -> set[str]:
|
|
"""Return all built-in, precomputed, and runtime-uploaded speakers."""
|
|
available_speakers = set(self.uploaded_speakers)
|
|
if self._adapter is not None:
|
|
available_speakers.update(self._adapter.capabilities.supported_speakers)
|
|
available_speakers.update(self._adapter.capabilities.precomputed_speakers)
|
|
return available_speakers
|
|
|
|
def _audio_encode_speed(self, request: OpenAICreateSpeechRequest) -> float:
|
|
if self._uses_native_speed_control():
|
|
return 1.0
|
|
return float(request.speed or 1.0)
|
|
|
|
async def warmup(self) -> None:
|
|
"""Run model-specific startup warmup through the resolved adapter.
|
|
|
|
Warmup requirements differ by model. For example, Qwen3-TTS can warm up
|
|
its standalone tokenizer decoder during model initialization, while
|
|
VoxCPM2 requires a real inference request with an active vLLM
|
|
``ForwardContext``. The adapter owns that model-specific lifecycle logic.
|
|
"""
|
|
if self._adapter is None:
|
|
return
|
|
await self._adapter.warmup()
|
|
|
|
def shutdown(self) -> None:
|
|
"""Shut down the TTS thread pool executor."""
|
|
if self._tts_executor is not None:
|
|
self._tts_executor.shutdown(wait=False, cancel_futures=True)
|
|
self._tts_executor = None
|
|
for name in list(self.uploaded_speakers):
|
|
self._speaker_cache.clear(name)
|
|
|
|
def _find_tts_stage(self):
|
|
"""Find and return the TTS stage config, or None if not found."""
|
|
tts_stage_keys = all_tts_stage_keys()
|
|
entry_stage_archs = tts_entry_stage_archs()
|
|
all_stages = frozenset(
|
|
model_stage
|
|
for stage in self.engine_client.stage_configs
|
|
if (model_stage := _stage_speech_metadata(stage)[0]) is not None
|
|
)
|
|
for stage in self.engine_client.stage_configs:
|
|
model_stage, model_arch, worker_type = _stage_speech_metadata(stage)
|
|
if model_stage in tts_stage_keys:
|
|
# Owning the stage key is not always enough: a model may be
|
|
# speech-capable only in some deployment topologies. Ask the
|
|
# adapter. Stages that resolve to no adapter keep the legacy
|
|
# behaviour of being accepted here (and detected as ``None``).
|
|
adapter_cls = resolve_adapter(detect_tts_model_type(model_stage, model_arch))
|
|
if adapter_cls is not None and not adapter_cls.stage_serves_speech(model_stage, all_stages):
|
|
continue
|
|
return stage
|
|
# Models with no dedicated TTS model_stage value identify their AR
|
|
# entry stage by architecture (Ming dense).
|
|
if model_arch in entry_stage_archs and worker_type == "ar":
|
|
return stage
|
|
return None
|
|
|
|
def _detect_tts_model_type(self) -> str | None:
|
|
"""Detect TTS model type from the resolved stage's deployment metadata.
|
|
|
|
The mapping lives on the adapters (``stage_keys`` / ``model_archs``);
|
|
this only supplies the stage under inspection.
|
|
"""
|
|
if self._tts_stage is None:
|
|
return None
|
|
return detect_tts_model_type(*_stage_speech_metadata(self._tts_stage)[:2])
|
|
|
|
def _compute_max_instructions_length(self) -> int:
|
|
"""Compute max instructions length with precedence: CLI > stage config > default.
|
|
|
|
Called once during initialization; result is cached in self._max_instructions_length.
|
|
"""
|
|
# 1. CLI override takes highest priority (stored in engine_client)
|
|
cli_override = getattr(self.engine_client, "tts_max_instructions_length", None)
|
|
if cli_override is not None:
|
|
return cli_override
|
|
|
|
# 2. Try to get from TTS stage config
|
|
if self._tts_stage is not None:
|
|
topology = getattr(self._tts_stage, "stage_pipeline_config", None)
|
|
tts_args = (
|
|
getattr(topology, "extras", {}).get("tts_args", {})
|
|
if isinstance(topology, StagePipelineConfig)
|
|
else getattr(self._tts_stage, "tts_args", {})
|
|
)
|
|
if "max_instructions_length" in tts_args:
|
|
return tts_args["max_instructions_length"]
|
|
|
|
# 3. Default fallback
|
|
return _TTS_MAX_INSTRUCTIONS_LENGTH
|
|
|
|
def _get_available_voices(self) -> set[str]:
|
|
"""Get all voice names accepted by the API, including the placeholder default."""
|
|
return self._get_available_speakers() | {_DEFAULT_VOICE_NAME}
|
|
|
|
def _get_usage_text_tokenizer(self):
|
|
"""Return a text tokenizer for counting input-text usage tokens.
|
|
|
|
Prefer the per-model tokenizer already loaded for prompt-length
|
|
estimation (`_tts_tokenizer`, which is the *correct* tokenizer for the
|
|
active model). Fall back to a lazily-loaded, cached generic tokenizer
|
|
for models that never populate `_tts_tokenizer`. Returns None if no
|
|
tokenizer can be obtained (usage then reports text_tokens=0).
|
|
"""
|
|
if self._tts_tokenizer is not None:
|
|
return self._tts_tokenizer
|
|
if getattr(self, "_usage_text_tokenizer", None) is None:
|
|
try:
|
|
from transformers import AutoTokenizer
|
|
|
|
self._usage_text_tokenizer = AutoTokenizer.from_pretrained(
|
|
self.engine_client.model_config.model, trust_remote_code=True
|
|
)
|
|
except Exception as e: # pragma: no cover - environment dependent
|
|
logger.warning("Usage: could not load a text tokenizer (%s); text_tokens will be 0", e)
|
|
self._usage_text_tokenizer = None
|
|
return self._usage_text_tokenizer
|
|
|
|
def _count_usage_text_tokens(self, text: str) -> int:
|
|
"""Token count of `text` using the model's text tokenizer (0 on failure)."""
|
|
if not text:
|
|
return 0
|
|
tok = self._get_usage_text_tokenizer()
|
|
if tok is None:
|
|
return 0
|
|
try:
|
|
return len(tok(text, padding=False)["input_ids"])
|
|
except Exception:
|
|
return 0
|
|
|
|
def _compute_speech_input_details(
|
|
self, request: OpenAICreateSpeechRequest, tts_params: dict[str, Any]
|
|
) -> SpeechInputTokenDetails:
|
|
"""Input-token breakdown (text + reference-audio) for a speech request.
|
|
|
|
Counts `input` (+ `instructions`) as text tokens, and reference-audio
|
|
codec frames as audio tokens *only* when in-context voice cloning is
|
|
active (see `qwen3_tts_input_token_details` / `gate_audio_tokens`). The
|
|
audio gating reads Qwen3-TTS `tts_params` conventions; other TTS models
|
|
do not set those keys, so they degrade cleanly to text-only counts.
|
|
"""
|
|
return qwen3_tts_input_token_details(
|
|
input_text=request.input,
|
|
instructions=request.instructions,
|
|
tts_params=tts_params or {},
|
|
count_text_tokens=self._count_usage_text_tokens,
|
|
)
|
|
|
|
def _build_speech_usage(
|
|
self,
|
|
request: OpenAICreateSpeechRequest,
|
|
tts_params: dict[str, Any],
|
|
output_tokens: int,
|
|
) -> SpeechTokenUsage:
|
|
"""Assemble the full usage object (input breakdown + generated tokens)."""
|
|
details = self._compute_speech_input_details(request, tts_params)
|
|
return build_speech_usage(details, output_tokens)
|
|
|
|
def _validate_tts_generation(
|
|
self,
|
|
tts_params: dict[str, Any],
|
|
usage_acc: SpeechOutputTokenCounter,
|
|
) -> None:
|
|
adapter = self._get_tts_adapter()
|
|
if adapter is None:
|
|
return
|
|
adapter.validate_generation(
|
|
tts_params,
|
|
stage0_finish_reason=usage_acc.stage0_finish_reason,
|
|
output_tokens=usage_acc.total(),
|
|
)
|
|
|
|
@staticmethod
|
|
def _build_speech_usage_headers(usage: SpeechTokenUsage | None) -> dict[str, str]:
|
|
"""Map speech usage into non-streaming response headers.
|
|
|
|
Returns an empty dict when usage is unavailable, allowing callers to
|
|
merge these headers with other optional response headers.
|
|
"""
|
|
if usage is None:
|
|
return {}
|
|
|
|
return {
|
|
_SPEECH_USAGE_INPUT_TOKENS_HEADER: str(usage.input_tokens),
|
|
_SPEECH_USAGE_OUTPUT_TOKENS_HEADER: str(usage.output_tokens),
|
|
_SPEECH_USAGE_TOTAL_TOKENS_HEADER: str(usage.total_tokens),
|
|
_SPEECH_USAGE_INPUT_TEXT_TOKENS_HEADER: str(usage.input_token_details.text_tokens),
|
|
_SPEECH_USAGE_INPUT_AUDIO_TOKENS_HEADER: str(usage.input_token_details.audio_tokens),
|
|
}
|
|
|
|
def _voice_created_at(self, voice_lower: str) -> int:
|
|
"""Return the upload timestamp of an uploaded voice, or 0 for built-ins.
|
|
|
|
Plumbed through to the model-side cache key so that delete + re-upload
|
|
of the same name yields a fresh cache slot.
|
|
"""
|
|
info = self.uploaded_speakers.get(voice_lower)
|
|
return int(info.get("created_at", 0)) if info else 0
|
|
|
|
def _load_uploaded_audio(self, voice_name: str) -> tuple[np.ndarray, int] | None:
|
|
"""Load decoded audio samples + sample rate from an uploaded voice's safetensors."""
|
|
voice_name_lower = voice_name.lower()
|
|
info = self.uploaded_speakers.get(voice_name_lower)
|
|
if info is None or info.get("embedding_source") != "audio":
|
|
return None
|
|
file_path = Path(info["file_path"])
|
|
if not file_path.exists():
|
|
logger.warning("Voice file not found for %s: %s", voice_name, file_path)
|
|
return None
|
|
try:
|
|
from safetensors import safe_open
|
|
except ImportError:
|
|
logger.error("The 'safetensors' package is required to load uploaded voices")
|
|
return None
|
|
try:
|
|
with safe_open(str(file_path), framework="pt") as f:
|
|
if "audio" not in f.keys():
|
|
return None
|
|
samples = f.get_tensor("audio").numpy()
|
|
sr = int((f.metadata() or {}).get("sample_rate", info.get("sample_rate", 0)))
|
|
except Exception as e:
|
|
logger.error("Could not load audio for voice %s: %s", voice_name, e)
|
|
return None
|
|
if sr <= 0:
|
|
return None
|
|
return samples, sr
|
|
|
|
def _load_registered_reference(self, reference: RegisteredVoiceReference) -> tuple[np.ndarray, int]:
|
|
"""Load one captured generation without retaining a data URI or waveform.
|
|
|
|
Keep the legacy WAV conversion on cache misses to preserve PCM16
|
|
quantization. A missing/deleted generation fails instead of substituting
|
|
a newer upload with the same name. Called off the event loop.
|
|
"""
|
|
from safetensors import safe_open
|
|
|
|
path = Path(reference.file_path)
|
|
if not _validate_path_within_directory(path, self.uploaded_speakers_dir):
|
|
raise ValueError("Invalid registered reference path")
|
|
with safe_open(str(path), framework="pt") as f:
|
|
metadata = f.metadata() or {}
|
|
if int(metadata.get("created_at", 0)) != reference.created_at:
|
|
raise ValueError("Registered reference generation changed")
|
|
samples = f.get_tensor("audio").numpy()
|
|
sr = int(metadata["sample_rate"])
|
|
# Match _get_uploaded_audio_data's WAV subtype and the resolver's float32
|
|
# channel mixing, without base64 serialization or any persistent copy.
|
|
buf = io.BytesIO()
|
|
sf.write(buf, samples, sr, format="WAV")
|
|
buf.seek(0)
|
|
waveform, sr = sf.read(buf, dtype="float32")
|
|
waveform, sr, _, _ = self._finalize_fetched_ref_audio(waveform, sr)
|
|
return waveform, sr
|
|
|
|
def _get_uploaded_audio_data(self, voice_name: str) -> str | None:
|
|
"""Return a base64-encoded WAV data URL for an uploaded voice.
|
|
|
|
Memoized so the WAV re-encode runs once per voice per process.
|
|
"""
|
|
voice_name_lower = voice_name.lower()
|
|
cached = self._ref_audio_data_url_cache.get(voice_name_lower)
|
|
if cached is not None:
|
|
return cached
|
|
|
|
data = self._load_uploaded_audio(voice_name)
|
|
if data is None:
|
|
return None
|
|
samples, sr = data
|
|
try:
|
|
buf = io.BytesIO()
|
|
sf.write(buf, samples, sr, format="WAV")
|
|
audio_b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
|
|
data_url = f"data:audio/wav;base64,{audio_b64}"
|
|
except Exception as e:
|
|
logger.error("Could not encode voice %s as WAV: %s", voice_name, e)
|
|
return None
|
|
self._ref_audio_data_url_cache[voice_name_lower] = data_url
|
|
return data_url
|
|
|
|
def _get_uploaded_speaker_embedding(self, voice_name: str) -> list[float] | None:
|
|
"""Load a pre-computed speaker embedding from an uploaded voice's safetensors.
|
|
|
|
Returns ``None`` if the voice has audio (not a direct embedding)."""
|
|
voice_name_lower = voice_name.lower()
|
|
info = self.uploaded_speakers.get(voice_name_lower)
|
|
if info is None or info.get("embedding_source") != "direct":
|
|
return None
|
|
file_path = Path(info["file_path"])
|
|
if not file_path.exists():
|
|
logger.warning("Embedding file not found for voice %s: %s", voice_name, file_path)
|
|
return None
|
|
if not _validate_path_within_directory(file_path, self.uploaded_speakers_dir):
|
|
logger.error("File path traversal detected for voice %s: %s", voice_name, file_path)
|
|
return None
|
|
try:
|
|
from safetensors.torch import load_file
|
|
except ImportError:
|
|
logger.error("The 'safetensors' package is required to load speaker embeddings")
|
|
return None
|
|
try:
|
|
tensors = load_file(str(file_path))
|
|
if "speaker_embedding" not in tensors:
|
|
logger.warning("Key 'speaker_embedding' missing in %s", file_path)
|
|
return None
|
|
return tensors["speaker_embedding"].squeeze().tolist()
|
|
except Exception as e:
|
|
logger.error("Could not load embedding for voice %s: %s", voice_name, e)
|
|
return None
|
|
|
|
def _apply_uploaded_speaker(self, request: OpenAICreateSpeechRequest) -> str | None:
|
|
"""Resolve ``request.voice`` against uploaded speakers, mutating
|
|
``request.ref_audio`` / ``request.ref_text`` in place. Returns an
|
|
error string if the voice is invalid, else ``None``.
|
|
"""
|
|
if request.voice is None or request.ref_audio is not None:
|
|
return None
|
|
|
|
voice_lower = request.voice.lower()
|
|
if voice_lower not in self.uploaded_speakers:
|
|
if self._tts_model_type in (
|
|
"cosyvoice3",
|
|
"fish_tts",
|
|
"audio8_tts",
|
|
"omnivoice",
|
|
"moss_tts_nano",
|
|
"glm_tts",
|
|
"higgs_audio_v2",
|
|
"higgs_audio_v3",
|
|
):
|
|
label = {
|
|
"cosyvoice3": "CosyVoice3",
|
|
"fish_tts": "Fish Speech",
|
|
"audio8_tts": "Audio8 TTS",
|
|
"omnivoice": "OmniVoice",
|
|
"moss_tts_nano": "MOSS-TTS-Nano",
|
|
"higgs_audio_v2": "Higgs-Audio V2",
|
|
"higgs_audio_v3": "Higgs-Audio V3",
|
|
"glm_tts": "GLM-TTS",
|
|
}.get(self._tts_model_type, self._tts_model_type)
|
|
return (
|
|
f"Unknown voice '{request.voice}'. {label} has no "
|
|
f"built-in speakers. Upload a voice first via "
|
|
f"POST /v1/audio/voices, or use ref_audio + ref_text."
|
|
)
|
|
return None
|
|
|
|
speaker_info = self.uploaded_speakers[voice_lower]
|
|
if speaker_info.get("embedding_source") == "direct":
|
|
return (
|
|
f"Uploaded voice '{request.voice}' uses a speaker embedding "
|
|
f"(Qwen3-only). Re-upload with an audio file for this model."
|
|
)
|
|
|
|
audio_data = self._get_uploaded_audio_data(request.voice)
|
|
if not audio_data:
|
|
return f"Audio file for uploaded voice '{request.voice}' is missing"
|
|
|
|
request.ref_audio = audio_data
|
|
if not request.ref_text or not request.ref_text.strip():
|
|
stored_ref_text = speaker_info.get("ref_text")
|
|
if stored_ref_text:
|
|
request.ref_text = stored_ref_text
|
|
|
|
logger.info("Resolved uploaded voice '%s' for %s", voice_lower, self._tts_model_type)
|
|
return None
|
|
|
|
def _check_upload_cap(self) -> None:
|
|
if len(self.uploaded_speakers) >= self._max_uploaded_speakers:
|
|
raise ValueError(
|
|
f"Uploaded voice limit reached ({self._max_uploaded_speakers}). "
|
|
f"Delete an existing voice before registering a new one, or raise "
|
|
f"the cap via SPEAKER_MAX_UPLOADED."
|
|
)
|
|
|
|
def _drop_shadowing_uploads(self) -> None:
|
|
"""Uploads are restored from disk before the adapter exists, so a name
|
|
registered under a built-in speaker before the collision guard would
|
|
keep shadowing it across restarts. Drop such entries from the registry;
|
|
the file stays on disk for the operator to remove."""
|
|
caps = self._adapter.capabilities if self._adapter is not None else None
|
|
if caps is None:
|
|
return
|
|
for voice_name_lower in list(self.uploaded_speakers):
|
|
if voice_name_lower in caps.supported_speakers or voice_name_lower in caps.precomputed_speakers:
|
|
info = self.uploaded_speakers.pop(voice_name_lower)
|
|
logger.warning(
|
|
"Uploaded voice %r shadows a built-in speaker and is ignored; remove %s to clear this warning.",
|
|
voice_name_lower,
|
|
info.get("file_path"),
|
|
)
|
|
|
|
def _check_registration_allowed(self, voice_name_lower: str, name: str) -> None:
|
|
"""Reject names that would shadow a built-in voice or, under the
|
|
immutable policy, silently overwrite an existing upload."""
|
|
caps = self._adapter.capabilities if self._adapter is not None else None
|
|
if caps is not None and (
|
|
voice_name_lower in caps.supported_speakers or voice_name_lower in caps.precomputed_speakers
|
|
):
|
|
raise ValueError(
|
|
f"Voice name '{name}' is reserved by a built-in speaker of this model; choose a different name."
|
|
)
|
|
if self._registration_policy == "immutable" and voice_name_lower in self.uploaded_speakers:
|
|
raise ValueError(
|
|
f"Voice '{name}' already exists and VLLM_OMNI_SPEAKER_REGISTRATION_POLICY is 'immutable'; "
|
|
f"delete it first via DELETE /v1/audio/voices/{name}."
|
|
)
|
|
|
|
def _evict_existing_upload(self, voice_name_lower: str, name: str) -> None:
|
|
"""Drop an existing upload with this name so the caller can re-register it."""
|
|
if voice_name_lower not in self.uploaded_speakers:
|
|
return
|
|
old = self.uploaded_speakers.pop(voice_name_lower)
|
|
self._ref_audio_data_url_cache.pop(voice_name_lower, None)
|
|
old_path = old.get("file_path")
|
|
if old_path:
|
|
try:
|
|
Path(old_path).unlink(missing_ok=True)
|
|
except Exception as e:
|
|
logger.warning("Failed to remove previous file for '%s': %s", name, e)
|
|
self._speaker_cache.clear(voice_name_lower)
|
|
logger.info("Speaker '%s' re-uploaded; previous cache and file overwritten", name)
|
|
|
|
async def upload_voice(
|
|
self,
|
|
audio_file: UploadFile,
|
|
consent: str,
|
|
name: str,
|
|
*,
|
|
ref_text: str | None = None,
|
|
speaker_description: str | None = None,
|
|
) -> dict:
|
|
"""Upload a new voice sample."""
|
|
name = _validate_speaker_name(name)
|
|
# Normalize optional strings: treat whitespace-only as absent
|
|
if ref_text is not None:
|
|
ref_text = ref_text.strip() or None
|
|
if speaker_description is not None:
|
|
speaker_description = speaker_description.strip() or None
|
|
# Validate file size (max 10MB)
|
|
MAX_FILE_SIZE = 10 * 1024 * 1024 # 10MB
|
|
audio_file.file.seek(0, 2) # Seek to end
|
|
file_size = audio_file.file.tell()
|
|
audio_file.file.seek(0) # Reset to beginning
|
|
|
|
if file_size > MAX_FILE_SIZE:
|
|
raise ValueError(f"File size exceeds maximum limit of 10MB. Got {file_size} bytes.")
|
|
|
|
# Detect MIME type from filename if content_type is generic
|
|
mime_type = audio_file.content_type
|
|
if mime_type == "application/octet-stream":
|
|
# Simple MIME type detection based on file extension
|
|
filename_lower = audio_file.filename.lower()
|
|
if filename_lower.endswith(".wav"):
|
|
mime_type = "audio/wav"
|
|
elif filename_lower.endswith((".mp3", ".mpeg")):
|
|
mime_type = "audio/mpeg"
|
|
elif filename_lower.endswith(".flac"):
|
|
mime_type = "audio/flac"
|
|
elif filename_lower.endswith(".ogg"):
|
|
mime_type = "audio/ogg"
|
|
elif filename_lower.endswith(".aac"):
|
|
mime_type = "audio/aac"
|
|
elif filename_lower.endswith(".webm"):
|
|
mime_type = "audio/webm"
|
|
elif filename_lower.endswith(".mp4"):
|
|
mime_type = "audio/mp4"
|
|
else:
|
|
mime_type = "audio/wav" # Default
|
|
|
|
# Validate MIME type
|
|
allowed_mime_types = {
|
|
"audio/mpeg",
|
|
"audio/wav",
|
|
"audio/x-wav",
|
|
"audio/ogg",
|
|
"audio/aac",
|
|
"audio/flac",
|
|
"audio/webm",
|
|
"audio/mp4",
|
|
}
|
|
|
|
if mime_type not in allowed_mime_types:
|
|
raise ValueError(f"Unsupported MIME type: {mime_type}. Allowed: {allowed_mime_types}")
|
|
|
|
# Read content before acquiring the lock; decode happens inside.
|
|
content = await audio_file.read()
|
|
|
|
async with self._upload_lock:
|
|
voice_name_lower = name.lower()
|
|
self._check_registration_allowed(voice_name_lower, name)
|
|
self._evict_existing_upload(voice_name_lower, name)
|
|
self._check_upload_cap()
|
|
|
|
sanitized_name = _sanitize_filename(name)
|
|
sanitized_consent = _sanitize_filename(consent)
|
|
timestamp = self._next_upload_timestamp()
|
|
file_suffix = Path(audio_file.filename).suffix
|
|
file_ext = file_suffix[1:] if file_suffix and len(file_suffix) > 1 else "wav"
|
|
sanitized_ext = _sanitize_filename(file_ext)
|
|
if not sanitized_ext or sanitized_ext == "file":
|
|
sanitized_ext = "wav"
|
|
|
|
filename = f"{sanitized_name}_{sanitized_consent}_{timestamp}.safetensors"
|
|
file_path = self.uploaded_speakers_dir / filename
|
|
if not _validate_path_within_directory(file_path, self.uploaded_speakers_dir):
|
|
raise ValueError("Invalid file path: potential path traversal attack detected")
|
|
|
|
try:
|
|
wav_np, sr = sf.read(io.BytesIO(content))
|
|
except Exception as e:
|
|
raise ValueError(f"Could not decode audio file: {e}")
|
|
duration = len(wav_np) / sr if sr > 0 else 0.0
|
|
if duration < _REF_AUDIO_MIN_DURATION:
|
|
raise ValueError(
|
|
f"Reference audio too short ({duration:.1f}s). "
|
|
f"At least {_REF_AUDIO_MIN_DURATION:.0f}s of clear speech is required."
|
|
)
|
|
if duration > _REF_AUDIO_MAX_DURATION:
|
|
raise ValueError(
|
|
f"Reference audio too long ({duration:.1f}s). "
|
|
f"Maximum {_REF_AUDIO_MAX_DURATION:.0f}s supported — use a shorter clip."
|
|
)
|
|
|
|
speaker_data: dict[str, Any] = {
|
|
"name": name,
|
|
"voice_name_lower": voice_name_lower,
|
|
"consent": consent,
|
|
"file_path": str(file_path),
|
|
"created_at": timestamp,
|
|
"mime_type": mime_type,
|
|
"original_filename": audio_file.filename,
|
|
"file_size": file_size,
|
|
"sample_rate": int(sr),
|
|
"ref_text": ref_text,
|
|
"embedding_source": "audio",
|
|
}
|
|
if speaker_description:
|
|
speaker_data["speaker_description"] = speaker_description
|
|
|
|
try:
|
|
from safetensors.torch import save_file
|
|
except ImportError as exc:
|
|
raise ValueError("safetensors is required for voice upload") from exc
|
|
try:
|
|
audio_tensor = torch.from_numpy(np.asarray(wav_np, dtype=np.float32)).contiguous()
|
|
save_file(
|
|
{"audio": audio_tensor},
|
|
str(file_path),
|
|
metadata=self._speaker_metadata_to_header(speaker_data),
|
|
)
|
|
except Exception as e:
|
|
raise ValueError(f"Failed to save voice file: {e}")
|
|
|
|
self.uploaded_speakers[voice_name_lower] = speaker_data
|
|
|
|
logger.info("Uploaded new voice '%s' with consent ID '%s'", name, consent)
|
|
|
|
# Return voice information without exposing the server file path
|
|
result = {
|
|
"name": name,
|
|
"consent": consent,
|
|
"created_at": timestamp,
|
|
"mime_type": mime_type,
|
|
"file_size": file_size,
|
|
}
|
|
if speaker_data.get("ref_text"):
|
|
result["ref_text"] = speaker_data["ref_text"]
|
|
if speaker_data.get("speaker_description"):
|
|
result["speaker_description"] = speaker_data["speaker_description"]
|
|
return result
|
|
|
|
async def upload_voice_embedding(self, embedding_json: str, consent: str, name: str) -> dict:
|
|
"""Upload a voice from a pre-computed speaker embedding.
|
|
|
|
Stores the embedding as a safetensors file and marks it immediately
|
|
ready (no audio processing needed).
|
|
|
|
Args:
|
|
embedding_json: JSON-encoded list of floats (1024 or 2048 dim).
|
|
consent: Consent recording ID.
|
|
name: Name for the new voice.
|
|
|
|
Returns:
|
|
dict with voice information.
|
|
"""
|
|
name = _validate_speaker_name(name)
|
|
try:
|
|
embedding = json.loads(embedding_json)
|
|
except (json.JSONDecodeError, TypeError) as exc:
|
|
raise ValueError(f"'speaker_embedding' must be valid JSON: {exc}") from exc
|
|
|
|
if not isinstance(embedding, list) or not embedding:
|
|
raise ValueError("'speaker_embedding' must be a non-empty list of numbers")
|
|
|
|
if len(embedding) > 4096:
|
|
raise ValueError("'speaker_embedding' exceeds maximum length (4096 elements)")
|
|
|
|
if not all(isinstance(x, (int, float)) for x in embedding):
|
|
raise ValueError("'speaker_embedding' must contain only numeric values")
|
|
|
|
if not all(math.isfinite(x) for x in embedding):
|
|
raise ValueError("'speaker_embedding' values must be finite (no NaN or Inf)")
|
|
|
|
emb_dim = len(embedding)
|
|
dim_err = self._adapter.validate_tts_embedding_dim(emb_dim) if self._adapter is not None else None
|
|
if dim_err is not None:
|
|
raise ValueError(dim_err)
|
|
|
|
async with self._upload_lock:
|
|
voice_name_lower = name.lower()
|
|
self._check_registration_allowed(voice_name_lower, name)
|
|
self._evict_existing_upload(voice_name_lower, name)
|
|
self._check_upload_cap()
|
|
|
|
sanitized_name = _sanitize_filename(name)
|
|
sanitized_consent = _sanitize_filename(consent)
|
|
timestamp = self._next_upload_timestamp()
|
|
|
|
tensor = torch.tensor(embedding, dtype=torch.float32)
|
|
filename = f"{sanitized_name}_{sanitized_consent}_{timestamp}.safetensors"
|
|
file_path = self.uploaded_speakers_dir / filename
|
|
if not _validate_path_within_directory(file_path, self.uploaded_speakers_dir):
|
|
raise ValueError("Invalid file path: potential path traversal attack detected")
|
|
|
|
speaker_data: dict[str, Any] = {
|
|
"name": name,
|
|
"voice_name_lower": voice_name_lower,
|
|
"consent": consent,
|
|
"file_path": str(file_path),
|
|
"created_at": timestamp,
|
|
"mime_type": "application/x-safetensors",
|
|
"original_filename": filename,
|
|
"embedding_source": "direct",
|
|
"embedding_dim": emb_dim,
|
|
}
|
|
try:
|
|
from safetensors.torch import save_file
|
|
except ImportError as exc:
|
|
raise ValueError("safetensors is required for embedding upload") from exc
|
|
save_file(
|
|
{"speaker_embedding": tensor},
|
|
str(file_path),
|
|
metadata=self._speaker_metadata_to_header(speaker_data),
|
|
)
|
|
speaker_data["file_size"] = file_path.stat().st_size
|
|
|
|
self.uploaded_speakers[voice_name_lower] = speaker_data
|
|
|
|
logger.info("Uploaded voice '%s' from speaker embedding (%d-dim)", name, emb_dim)
|
|
|
|
return {
|
|
"name": name,
|
|
"consent": consent,
|
|
"created_at": timestamp,
|
|
"embedding_source": "direct",
|
|
"embedding_dim": emb_dim,
|
|
}
|
|
|
|
async def delete_voice(self, name: str):
|
|
"""
|
|
Delete an uploaded voice.
|
|
|
|
Args:
|
|
name: Voice name to delete
|
|
"""
|
|
async with self._upload_lock:
|
|
voice_name_lower = name.lower()
|
|
built_in_speakers = self._get_available_voices() - set(self.uploaded_speakers)
|
|
|
|
if voice_name_lower in built_in_speakers:
|
|
err = f"Cannot delete built-in voice '{name}'"
|
|
raise InvalidPresetVoiceReferenceError(err)
|
|
|
|
if voice_name_lower not in self.uploaded_speakers:
|
|
err = f"Voice '{name}' not found"
|
|
raise InvalidVoiceReferenceError(err)
|
|
|
|
speaker_info = self.uploaded_speakers.pop(voice_name_lower)
|
|
self._ref_audio_data_url_cache.pop(voice_name_lower, None)
|
|
|
|
file_path = speaker_info.get("file_path")
|
|
if file_path:
|
|
try:
|
|
Path(file_path).unlink(missing_ok=True)
|
|
except Exception as e:
|
|
logger.warning("Failed to delete audio file for '%s': %s", name, e)
|
|
|
|
self._speaker_cache.clear(voice_name_lower)
|
|
|
|
logger.info("Deleted voice '%s'", name)
|
|
|
|
def _is_tts_model(self) -> bool:
|
|
"""Check if the current model is a supported TTS model."""
|
|
return self._find_tts_stage() is not None
|
|
|
|
def _validate_tts_request(self, request: OpenAICreateSpeechRequest) -> str | None:
|
|
"""Validate TTS request parameters. Returns error message or None."""
|
|
sample_rate_error = self._validate_speech_sample_rate(request)
|
|
if sample_rate_error is not None:
|
|
return sample_rate_error
|
|
|
|
adapter = self._get_tts_adapter()
|
|
if adapter is not None:
|
|
adapter.normalize(request)
|
|
return adapter.validate(request)
|
|
|
|
adapter_cls = resolve_adapter("qwen3_tts")
|
|
if adapter_cls is None:
|
|
raise ValueError("Qwen3-TTS adapter is not registered")
|
|
|
|
ctx = SpeechServingContext(
|
|
server=self,
|
|
engine_client=self.engine_client,
|
|
)
|
|
adapter = adapter_cls(ctx)
|
|
adapter.normalize(request)
|
|
return adapter.validate(request)
|
|
|
|
def _validate_speech_sample_rate(self, request: OpenAICreateSpeechRequest) -> str | None:
|
|
if request.sample_rate is None:
|
|
return None
|
|
|
|
adapter_cls = resolve_adapter(self._tts_model_type)
|
|
supported_rates = adapter_cls.supported_output_sample_rates if adapter_cls is not None else frozenset()
|
|
if request.sample_rate not in supported_rates:
|
|
if supported_rates:
|
|
rates = ", ".join(str(rate) for rate in sorted(supported_rates))
|
|
return (
|
|
f"sample_rate={request.sample_rate} is not supported by the current TTS model; "
|
|
f"supported rates: {rates}"
|
|
)
|
|
return "sample_rate is not supported by the current TTS model"
|
|
return None
|
|
|
|
def _validate_ref_audio_format(self, ref_audio: str | list[str] | None) -> str | None:
|
|
"""Validate ref_audio is a supported URI format. Returns error or None."""
|
|
if not isinstance(ref_audio, str):
|
|
return "ref_audio must be a URL (http/https), base64 data URL (data:...), or file URI (file://...)"
|
|
scheme = (urlparse(ref_audio).scheme or "").lower()
|
|
if scheme not in {"http", "https", "data", "file"}:
|
|
return "ref_audio must be a URL (http/https), base64 data URL (data:...), or file URI (file://...)"
|
|
return None
|
|
|
|
@staticmethod
|
|
def _local_ref_audio_stat_path(ref_audio_str: str) -> str | None:
|
|
"""Filesystem path to stat for a local locator, else ``None``.
|
|
|
|
Scheme comparison is case-insensitive (``urlparse`` lowercases it), so
|
|
``FILE:///...`` is treated as a local file the same way ``file://`` is.
|
|
Remote ``http`` / ``https`` / ``data`` locators, including ``HTTP://``,
|
|
return ``None`` and stay string-keyed.
|
|
"""
|
|
parsed = urlparse(ref_audio_str)
|
|
scheme = (parsed.scheme or "").lower()
|
|
if scheme in _REMOTE_REF_AUDIO_SCHEMES:
|
|
return None
|
|
if scheme == "file":
|
|
netloc = parsed.netloc or ""
|
|
if netloc.lower() not in ("", "localhost"):
|
|
raise OSError(f"file:// URI with non-local authority {netloc!r} cannot be stat'd")
|
|
return url2pathname(parsed.path or "")
|
|
if scheme:
|
|
return None
|
|
return ref_audio_str
|
|
|
|
@staticmethod
|
|
def _get_ref_audio_cache_key(
|
|
ref_audio_str: str,
|
|
allowed_local_media_path: str | None = None,
|
|
) -> str:
|
|
"""Compute a cache key hash for *ref_audio_str*.
|
|
|
|
For local files (bare paths and ``file://`` URIs) the key folds in
|
|
``st_mtime_ns`` and ``st_size`` so that an on-disk edit automatically
|
|
invalidates the cached waveform without a server restart. Remote URLs
|
|
and ``data:`` URIs are keyed on the raw string alone — the server does
|
|
not re-fetch them to check for changes.
|
|
|
|
``os.stat`` runs only when *allowed_local_media_path* is a non-empty
|
|
path. vLLM's default is ``""`` (not ``None``); both mean local media
|
|
is refused, so a stat would be exposure for a request that cannot
|
|
load the file.
|
|
|
|
Note: when the key changes the *previous* entry stays in
|
|
``_ref_audio_resolve_cache`` until LRU eviction, so
|
|
``_discard_ref_audio_artifact_ready_if_unreferenced`` will not fire for
|
|
the replaced file and its stale artifact key remains in
|
|
``_ref_audio_model_artifact_ready``. This is functionally correct
|
|
(the stale entry is never *used*) but doubles memory until eviction.
|
|
"""
|
|
cache_key_source = ref_audio_str
|
|
try:
|
|
path = OmniOpenAIServingSpeech._local_ref_audio_stat_path(ref_audio_str)
|
|
except OSError as exc:
|
|
path = None
|
|
logger.debug("Skipping local-file cache metadata for %s: %s", ref_audio_str[:80], exc)
|
|
if path is not None:
|
|
try:
|
|
# Only stat paths the server operator has explicitly permitted
|
|
# via --allowed-local-media-path. The default is "" (and
|
|
# diffusion mode passes None); MediaConnector refuses local
|
|
# file loads in both cases, so a stat is exposure for zero
|
|
# benefit. ``if not`` covers "" and None; ``is None`` would
|
|
# treat "" as an allowlist of the process cwd.
|
|
if not allowed_local_media_path:
|
|
raise OSError("no allowed_local_media_path configured; skipping stat")
|
|
resolved = os.path.realpath(path)
|
|
allowed_base = os.path.realpath(allowed_local_media_path)
|
|
if os.path.commonpath([resolved, allowed_base]) != allowed_base:
|
|
raise OSError("path outside allowed_local_media_path; skipping stat")
|
|
st = os.stat(path)
|
|
cache_key_source = f"{ref_audio_str}:{st.st_mtime_ns}:{st.st_size}"
|
|
except (OSError, ValueError):
|
|
# Truncate the value to avoid dumping huge base64 blobs into
|
|
# the server log when a client omits the ``data:`` prefix.
|
|
display = ref_audio_str[:80]
|
|
if len(ref_audio_str) > 80:
|
|
display += f"... ({len(ref_audio_str)} chars)"
|
|
logger.debug(
|
|
"Failed to stat ref_audio path %s; falling back to string-only cache key (stale cache possible)",
|
|
display,
|
|
)
|
|
return hashlib.sha1(cache_key_source.encode("utf-8")).hexdigest()
|
|
|
|
async def _ref_audio_cache_key(self, ref_audio_str: str, allowed_local_media_path: str | None) -> str:
|
|
"""Stat local files off the event loop; remote locators stay a cheap hash."""
|
|
return await asyncio.to_thread(
|
|
self._get_ref_audio_cache_key,
|
|
ref_audio_str,
|
|
allowed_local_media_path,
|
|
)
|
|
|
|
def _finalize_fetched_ref_audio(self, wav_np: np.ndarray, sr: int) -> tuple[np.ndarray, int, str, float]:
|
|
wav_np = np.asarray(wav_np, dtype=np.float32)
|
|
if wav_np.ndim > 1:
|
|
wav_np = np.mean(wav_np, axis=-1)
|
|
sr = int(sr)
|
|
duration = len(wav_np) / sr if sr > 0 else 0.0
|
|
if duration < _REF_AUDIO_MIN_DURATION:
|
|
raise ValueError(
|
|
f"Reference audio too short ({duration:.1f}s). "
|
|
f"At least {_REF_AUDIO_MIN_DURATION:.0f}s of clear speech is required."
|
|
)
|
|
if duration > _REF_AUDIO_MAX_DURATION:
|
|
raise ValueError(
|
|
f"Reference audio too long ({duration:.1f}s). "
|
|
f"Maximum {_REF_AUDIO_MAX_DURATION:.0f}s supported — use a shorter clip."
|
|
)
|
|
# Own the buffer: a view could retain a much larger decoded allocation.
|
|
waveform = np.array(wav_np, dtype=np.float32, order="C", copy=True)
|
|
artifact_key = self._make_ref_audio_artifact_cache_key(waveform, sr)
|
|
return waveform, sr, artifact_key, duration
|
|
|
|
async def _resolve_ref_audio(self, ref_audio_str: str) -> tuple[list[float], int, str]:
|
|
"""Keep list-based consumers and their engine transport compatible.
|
|
|
|
Only the numeric array is cached; this temporary list belongs to the
|
|
caller. Server-side reference encoders should use the array resolver.
|
|
"""
|
|
waveform, sr, cache_key = await self._resolve_ref_audio_array(ref_audio_str)
|
|
return waveform.tolist(), sr, cache_key
|
|
|
|
async def _resolve_ref_audio_array(self, ref_audio_str: str) -> tuple[np.ndarray, int, str]:
|
|
"""Resolve ref_audio to (wav_samples, sample_rate, cache_key).
|
|
|
|
The float32 array can be shared with the cache. Callers must not mutate
|
|
it; copy before any in-place processing.
|
|
|
|
Delegates to upstream vLLM's MediaConnector which handles http(s)
|
|
URLs, ``data:`` base64 URIs, and ``file:`` local paths (the latter
|
|
gated by ``--allowed-local-media-path``).
|
|
|
|
Local file references incorporate mtime and size into the cache key
|
|
so that modified files are automatically reloaded without a server
|
|
restart. Remote URLs remain cached by their original string locator.
|
|
|
|
The returned *cache_key* should be passed to
|
|
``_get_resolved_ref_audio_artifact_key`` when the caller needs the
|
|
artifact key, avoiding a redundant ``os.stat`` and the TOCTOU window.
|
|
After a cache miss the file is re-stat'd; a metadata change retries
|
|
the fetch so the stored waveform matches the key.
|
|
"""
|
|
# Pass the allowed-local-media-path so the stat is restricted to
|
|
# paths the server operator has explicitly permitted.
|
|
allowed_path: str | None
|
|
if self._diffusion_mode:
|
|
allowed_path = self._allowed_local_media_path
|
|
else:
|
|
allowed_path = getattr(self.model_config, "allowed_local_media_path", None)
|
|
|
|
waveform: np.ndarray | None = None
|
|
sr = 0
|
|
artifact_key = ""
|
|
post_key = ""
|
|
for attempt in range(_REF_AUDIO_METADATA_FETCH_ATTEMPTS):
|
|
cache_key = await self._ref_audio_cache_key(ref_audio_str, allowed_path)
|
|
cached = self._ref_audio_resolve_cache.get(cache_key)
|
|
if cached is not None:
|
|
self._ref_audio_resolve_cache.move_to_end(cache_key)
|
|
waveform, sr, _, _ = cached
|
|
logger.debug(
|
|
"Resolved ref_audio from cache: samples=%d sr=%d duration_s=%.3f",
|
|
len(waveform),
|
|
sr,
|
|
len(waveform) / sr if sr > 0 else 0.0,
|
|
)
|
|
return waveform, sr, cache_key
|
|
|
|
if self._media_connector is None:
|
|
model_config = self.model_config
|
|
self._media_connector = MediaConnector(
|
|
allowed_local_media_path=model_config.allowed_local_media_path,
|
|
allowed_media_domains=model_config.allowed_media_domains,
|
|
)
|
|
|
|
fetch_start_s = time.perf_counter()
|
|
wav_np, fetched_sr = await self._media_connector.fetch_audio_async(ref_audio_str)
|
|
fetch_decode_ms = (time.perf_counter() - fetch_start_s) * 1000.0
|
|
finalize_start_s = time.perf_counter()
|
|
waveform, sr, artifact_key, duration = self._finalize_fetched_ref_audio(wav_np, fetched_sr)
|
|
finalize_ms = (time.perf_counter() - finalize_start_s) * 1000.0
|
|
logger.debug(
|
|
"Resolved ref_audio: fetch_decode_ms=%.3f finalize_ms=%.3f samples=%d sr=%d duration_s=%.3f",
|
|
fetch_decode_ms,
|
|
finalize_ms,
|
|
len(waveform),
|
|
sr,
|
|
duration,
|
|
)
|
|
post_key = await self._ref_audio_cache_key(ref_audio_str, allowed_path)
|
|
if post_key == cache_key:
|
|
self._put_resolved_ref_audio(cache_key, waveform, sr, artifact_key)
|
|
return waveform, sr, cache_key
|
|
logger.debug(
|
|
"ref_audio metadata changed during fetch (attempt %d/%d); retrying",
|
|
attempt + 1,
|
|
_REF_AUDIO_METADATA_FETCH_ATTEMPTS,
|
|
)
|
|
|
|
logger.warning(
|
|
"ref_audio file changed during fetch after %d attempts; skipping resolve cache",
|
|
_REF_AUDIO_METADATA_FETCH_ATTEMPTS,
|
|
)
|
|
assert waveform is not None and post_key
|
|
return waveform, sr, post_key
|
|
|
|
@staticmethod
|
|
def _make_ref_audio_artifact_cache_key(wav: np.ndarray, sr: int) -> str:
|
|
wav_f32 = wav.astype(np.float32, copy=False).reshape(-1)
|
|
h = hashlib.sha1()
|
|
h.update(int(sr).to_bytes(4, byteorder="little", signed=False))
|
|
h.update(int(wav_f32.size).to_bytes(8, byteorder="little", signed=False))
|
|
h.update(wav_f32.tobytes(order="C"))
|
|
return h.hexdigest()
|
|
|
|
def _get_resolved_ref_audio_artifact_key(self, cache_key: str) -> str | None:
|
|
"""Look up the artifact key for a previously resolved ref_audio.
|
|
|
|
*cache_key* must be the exact key returned by the preceding
|
|
``_resolve_ref_audio`` call. Passing the key explicitly avoids a
|
|
second ``os.stat`` syscall and eliminates the TOCTOU window that
|
|
would arise from recomputing the key independently.
|
|
"""
|
|
cached = self._ref_audio_resolve_cache.get(cache_key)
|
|
if cached is None:
|
|
return None
|
|
self._ref_audio_resolve_cache.move_to_end(cache_key)
|
|
return cached[3]
|
|
|
|
def _put_resolved_ref_audio(
|
|
self, cache_key: str, waveform: np.ndarray | list[float], sr: int, artifact_key: str
|
|
) -> None:
|
|
if self._ref_audio_resolve_cache_max_entries <= 0 or self._ref_audio_resolve_cache_max_bytes <= 0:
|
|
return
|
|
# Keep the finalized ndarray's ownership and identity for the array
|
|
# resolver. Legacy list callers are materialized into a compact array;
|
|
# their Python list remains independent from cached storage.
|
|
waveform = np.asarray(waveform, dtype=np.float32)
|
|
if not waveform.flags.c_contiguous:
|
|
waveform = np.ascontiguousarray(waveform)
|
|
size = int(waveform.nbytes)
|
|
if size > self._ref_audio_resolve_cache_max_bytes:
|
|
return
|
|
previous = self._ref_audio_resolve_cache.pop(cache_key, None)
|
|
if previous is not None:
|
|
self._ref_audio_resolve_cache_bytes -= previous[2]
|
|
if previous[3] != artifact_key:
|
|
self._discard_ref_audio_artifact_ready_if_unreferenced(previous[3])
|
|
self._ref_audio_resolve_cache[cache_key] = (waveform, int(sr), size, artifact_key)
|
|
self._ref_audio_resolve_cache_bytes += size
|
|
while len(self._ref_audio_resolve_cache) > self._ref_audio_resolve_cache_max_entries:
|
|
_, (_, _, old_size, old_artifact_key) = self._ref_audio_resolve_cache.popitem(last=False)
|
|
self._ref_audio_resolve_cache_bytes -= old_size
|
|
self._discard_ref_audio_artifact_ready_if_unreferenced(old_artifact_key)
|
|
while self._ref_audio_resolve_cache_bytes > self._ref_audio_resolve_cache_max_bytes:
|
|
_, (_, _, old_size, old_artifact_key) = self._ref_audio_resolve_cache.popitem(last=False)
|
|
self._ref_audio_resolve_cache_bytes -= old_size
|
|
self._discard_ref_audio_artifact_ready_if_unreferenced(old_artifact_key)
|
|
|
|
def _discard_ref_audio_artifact_ready_if_unreferenced(self, artifact_key: str) -> None:
|
|
if artifact_key and all(entry[3] != artifact_key for entry in self._ref_audio_resolve_cache.values()):
|
|
self._ref_audio_model_artifact_ready = {
|
|
(key, mode) for (key, mode) in self._ref_audio_model_artifact_ready if key != artifact_key
|
|
}
|
|
|
|
@staticmethod
|
|
def _tts_x_vector_only(tts_params: dict[str, Any]) -> bool:
|
|
return bool((tts_params.get("x_vector_only_mode") or [False])[0])
|
|
|
|
def _track_ref_audio_artifact_warmup(
|
|
self, request_id: str, artifact_key: str | None, x_vector_only: bool = False
|
|
) -> None:
|
|
if artifact_key:
|
|
self._request_ref_audio_artifact_keys[request_id] = (artifact_key, bool(x_vector_only))
|
|
|
|
def _mark_ref_audio_artifact_ready_for_request(self, request_id: str) -> None:
|
|
tracked = self._request_ref_audio_artifact_keys.pop(request_id, None)
|
|
if tracked is None:
|
|
return
|
|
artifact_key, x_vector_only = tracked
|
|
if artifact_key and any(entry[3] == artifact_key for entry in self._ref_audio_resolve_cache.values()):
|
|
self._ref_audio_model_artifact_ready.add((artifact_key, x_vector_only))
|
|
|
|
def _discard_ref_audio_artifact_warmup(self, request_id: str) -> None:
|
|
self._request_ref_audio_artifact_keys.pop(request_id, None)
|
|
|
|
async def _resolve_ref_audio_many(self, ref_audio_list: list[str]) -> list[tuple[list[float], int]]:
|
|
resolved = []
|
|
for ref_audio in ref_audio_list:
|
|
wav_list, sr, _ = await self._resolve_ref_audio(ref_audio)
|
|
resolved.append((wav_list, sr))
|
|
return resolved
|
|
|
|
def _observe_speech_audio_ttfp(
|
|
self,
|
|
*,
|
|
request_id: str,
|
|
result: Any,
|
|
request_arrival_ts: float | None = None,
|
|
first_packet_ts: float | None = None,
|
|
) -> tuple[int | None, int | None, bool]:
|
|
"""Emit the Speech API TTFP sample once the first PCM payload exists."""
|
|
engine_client = getattr(self, "engine_client", None)
|
|
mod_metrics = getattr(engine_client, "mod_metrics", None)
|
|
if mod_metrics is None:
|
|
return None, None, False
|
|
|
|
req_state = next(
|
|
(
|
|
state
|
|
for state in getattr(engine_client, "request_states", {}).values()
|
|
if state.external_request_id == request_id
|
|
),
|
|
None,
|
|
)
|
|
if req_state is not None and req_state.first_audio_ts is not None:
|
|
return req_state.audio_emit_stage_id, req_state.audio_emit_replica_id, True
|
|
|
|
arrival_ts = request_arrival_ts
|
|
if arrival_ts is None and req_state is not None:
|
|
arrival_ts = req_state.request_arrival_ts
|
|
stage_id = getattr(result, "stage_id", None)
|
|
replica_id = getattr(result, "replica_id", None)
|
|
if replica_id is None and req_state is not None and stage_id is not None:
|
|
stage_pools = getattr(getattr(engine_client, "engine", None), "stage_pools", None)
|
|
if stage_pools is not None and 0 <= stage_id < len(stage_pools):
|
|
replica_id = stage_pools[stage_id].get_bound_replica_id(req_state.request_id)
|
|
if arrival_ts is None or arrival_ts <= 0 or stage_id is None or replica_id is None:
|
|
return stage_id, replica_id, False
|
|
|
|
now_ts = first_packet_ts if first_packet_ts is not None else time.time()
|
|
observe_audio_first_packet(
|
|
mod_metrics,
|
|
stage_id=stage_id,
|
|
replica_id=replica_id,
|
|
arrival_ts=arrival_ts,
|
|
now_ts=now_ts,
|
|
)
|
|
if req_state is not None:
|
|
req_state.first_audio_ts = now_ts
|
|
req_state.audio_emit_stage_id = stage_id
|
|
req_state.audio_emit_replica_id = replica_id
|
|
return stage_id, replica_id, True
|
|
|
|
async def _generate_audio_chunks(
|
|
self,
|
|
generator,
|
|
request_id: str,
|
|
response_format: str = "pcm",
|
|
raw_request: Request | None = None,
|
|
request_start_s: float | None = None,
|
|
request_arrival_ts: float | None = None,
|
|
include_sample_rate: bool = False,
|
|
usage_acc: SpeechOutputTokenCounter | None = None,
|
|
tts_params: dict[str, Any] | None = None,
|
|
collect: dict | None = None,
|
|
target_sample_rate: int | None = None,
|
|
cumulative_audio: bool = False,
|
|
audio_format: dict[str, int] | None = None,
|
|
):
|
|
"""Generate audio chunks for streaming response.
|
|
|
|
Handles two audio output modes from the engine:
|
|
- Cumulative mode (list): Engine returns growing list of chunks;
|
|
we emit only the new tail on each iteration.
|
|
- Tensor mode: Emit each tensor directly, or only its new samples when
|
|
cumulative_audio is set for requests that retain audio for alignment.
|
|
|
|
Args:
|
|
generator: Async generator from the engine
|
|
request_id: Request identifier for logging
|
|
response_format: Audio format (pcm or wav)
|
|
audio_format: Optional out-parameter. Before the first audio
|
|
chunk is yielded it receives that chunk's ``sample_rate`` and
|
|
``channels``, so callers can describe a raw PCM stream.
|
|
|
|
Yields:
|
|
Raw audio bytes for each chunk (with WAV header for first chunk if wav format)
|
|
"""
|
|
prev_count = 0
|
|
emitted_samples = 0
|
|
sample_rate_val = 24000
|
|
first_chunk = True
|
|
first_audio_chunk_s: float | None = None
|
|
first_audio_packet_ts: float | None = None
|
|
ttfp_observed = False
|
|
stream_start_s = request_start_s if request_start_s is not None else time.perf_counter()
|
|
artifact_ready = False
|
|
audio_chunk_arrivals_s: list[float] = []
|
|
audio_chunk_bytes: list[int] = []
|
|
audio_stage_id: int | None = None
|
|
audio_replica_id: int | None = None
|
|
audio_channels = 1
|
|
source_sample_rate: int | None = None
|
|
resampler: StreamingAudioResampler | None = None
|
|
last_audio_result: Any = None
|
|
|
|
def record_stream_abort(reason: str) -> None:
|
|
mod_metrics = getattr(getattr(self, "engine_client", None), "mod_metrics", None)
|
|
if mod_metrics is not None:
|
|
mod_metrics.inc_speech_stream_aborted(reason)
|
|
|
|
def record_audio_chunk(audio_bytes: bytes, chunk_np: np.ndarray, result: Any) -> None:
|
|
nonlocal first_audio_chunk_s, first_audio_packet_ts, ttfp_observed
|
|
nonlocal audio_stage_id, audio_replica_id, audio_channels
|
|
if not audio_bytes:
|
|
return
|
|
if first_audio_chunk_s is None:
|
|
first_audio_chunk_s = time.perf_counter()
|
|
first_audio_packet_ts = time.time()
|
|
audio_channels = _infer_audio_num_channels(np.asarray(chunk_np))
|
|
if not ttfp_observed:
|
|
audio_stage_id, audio_replica_id, ttfp_observed = self._observe_speech_audio_ttfp(
|
|
request_id=request_id,
|
|
result=result,
|
|
request_arrival_ts=request_arrival_ts,
|
|
first_packet_ts=first_audio_packet_ts,
|
|
)
|
|
audio_chunk_arrivals_s.append(max(time.perf_counter() - stream_start_s, 0.0))
|
|
audio_chunk_bytes.append(len(audio_bytes))
|
|
|
|
# SSE supplies an accumulator for usage output. Raw-audio and WebSocket
|
|
# streams retain terminal metrics only when their model adapter needs
|
|
# generation validation.
|
|
adapter = self._get_tts_adapter()
|
|
if tts_params is not None and usage_acc is None and adapter is not None and adapter.validates_generation:
|
|
usage_acc = SpeechOutputTokenCounter()
|
|
|
|
try:
|
|
async for res in generator:
|
|
# Tally generated codec tokens for usage (reads per-stage metrics
|
|
# off the final output; a cheap early-return on every other res).
|
|
if usage_acc is not None:
|
|
usage_acc.observe(res)
|
|
audio_output, audio_key = self._extract_audio_output(res)
|
|
if audio_key is None:
|
|
# Stash the aligner's timestamps output for streaming callers.
|
|
if collect is not None and self._is_timestamps_output(res):
|
|
collect["aligner_res"] = res
|
|
continue
|
|
audio_output = cast(dict, audio_output)
|
|
|
|
sr_raw = audio_output.get("sr")
|
|
if sr_raw is not None:
|
|
sr_val = sr_raw[-1] if isinstance(sr_raw, list) and sr_raw else sr_raw
|
|
sample_rate_val = sr_val.item() if hasattr(sr_val, "item") else int(sr_val)
|
|
if source_sample_rate is not None and sample_rate_val != source_sample_rate:
|
|
raise ValueError(
|
|
"Audio sample rate changed during streaming: "
|
|
f"{source_sample_rate} Hz to {sample_rate_val} Hz"
|
|
)
|
|
audio_val = audio_output[audio_key]
|
|
if isinstance(audio_val, list):
|
|
# Cumulative mode: each update grows the list; emit only new tail.
|
|
new_chunks = audio_val[prev_count:]
|
|
prev_count = len(audio_val)
|
|
else:
|
|
# Per-step mode: each update is a single tensor; emit directly.
|
|
if audio_val is not None:
|
|
new_chunks = [audio_val]
|
|
prev_count += 1
|
|
else:
|
|
new_chunks = []
|
|
|
|
if new_chunks and source_sample_rate is None:
|
|
if response_format == "wav" and sr_raw is None:
|
|
raise ValueError("First audio output must include sample rate metadata for WAV streaming")
|
|
source_sample_rate = sample_rate_val
|
|
if target_sample_rate is not None and target_sample_rate != source_sample_rate:
|
|
resampler = StreamingAudioResampler(source_sample_rate, target_sample_rate)
|
|
|
|
output_sample_rate = target_sample_rate or sample_rate_val
|
|
|
|
for chunk_tensor in new_chunks:
|
|
# Flush may run after a non-audio (e.g. aligner) result.
|
|
# Retain the producer of the audio buffered by the resampler.
|
|
last_audio_result = res
|
|
chunk_np = (
|
|
chunk_tensor.float().detach().cpu().numpy() if hasattr(chunk_tensor, "float") else chunk_tensor
|
|
)
|
|
if chunk_np.ndim > 1:
|
|
chunk_np = chunk_np.squeeze()
|
|
if cumulative_audio and not isinstance(audio_val, list):
|
|
total_samples = chunk_np.shape[-1]
|
|
chunk_np = chunk_np[..., emitted_samples:]
|
|
emitted_samples = total_samples
|
|
if chunk_np.size == 0:
|
|
continue
|
|
if resampler is not None:
|
|
chunk_np = resampler.process(chunk_np)
|
|
if chunk_np.size == 0:
|
|
continue
|
|
if int(np.size(chunk_np)) == 0:
|
|
# Zero-size chunks must not emit a WAV header or count
|
|
# as first audio; the post-loop guard below needs to
|
|
# see an audio-less stream to fail the request.
|
|
continue
|
|
wav_header: bytes | None = None
|
|
# For WAV format, emit header before first audio chunk
|
|
if response_format == "wav" and first_chunk:
|
|
num_channels = _infer_audio_num_channels(np.asarray(chunk_np))
|
|
wav_header = _create_wav_header(
|
|
sample_rate=output_sample_rate,
|
|
num_channels=num_channels,
|
|
bits_per_sample=16,
|
|
)
|
|
first_chunk = False
|
|
|
|
# Convert audio to PCM bytes
|
|
audio_obj = CreateAudio(
|
|
audio_tensor=chunk_np,
|
|
sample_rate=output_sample_rate,
|
|
response_format="pcm",
|
|
speed=1.0,
|
|
base64_encode=False,
|
|
)
|
|
audio_response = self.create_audio(audio_obj)
|
|
audio_bytes = cast(bytes, audio_response.audio_data)
|
|
record_audio_chunk(audio_bytes, chunk_np, res)
|
|
if audio_format is not None and not audio_format and audio_bytes:
|
|
audio_format.update(_encoded_audio_format(audio_response, output_sample_rate, chunk_np))
|
|
if wav_header is not None:
|
|
yield wav_header
|
|
if include_sample_rate:
|
|
yield audio_bytes, output_sample_rate
|
|
else:
|
|
yield audio_bytes
|
|
|
|
if resampler is not None:
|
|
final_chunk = resampler.process(np.empty((0,), dtype=np.float32), final=True)
|
|
if final_chunk.size:
|
|
output_sample_rate = target_sample_rate or sample_rate_val
|
|
wav_header = None
|
|
if response_format == "wav" and first_chunk:
|
|
wav_header = _create_wav_header(
|
|
sample_rate=output_sample_rate,
|
|
num_channels=1,
|
|
bits_per_sample=16,
|
|
)
|
|
first_chunk = False
|
|
audio_obj = CreateAudio(
|
|
audio_tensor=final_chunk,
|
|
sample_rate=output_sample_rate,
|
|
response_format="pcm",
|
|
speed=1.0,
|
|
base64_encode=False,
|
|
)
|
|
audio_response = self.create_audio(audio_obj)
|
|
audio_bytes = cast(bytes, audio_response.audio_data)
|
|
record_audio_chunk(audio_bytes, final_chunk, last_audio_result)
|
|
if audio_format is not None and not audio_format and audio_bytes:
|
|
audio_format.update(_encoded_audio_format(audio_response, output_sample_rate, final_chunk))
|
|
if wav_header is not None:
|
|
yield wav_header
|
|
if include_sample_rate:
|
|
yield audio_bytes, output_sample_rate
|
|
else:
|
|
yield audio_bytes
|
|
if self._tts_model_type in _AUDEX_NO_AUDIO_GUARD_MODEL_TYPES and first_audio_chunk_s is None:
|
|
# Audex contract: zero codec tokens must abort the stream, not
|
|
# complete it cleanly with zero audio bytes.
|
|
raise ValueError("Audex produced no audio (the thinker emitted zero or invalid codec tokens)")
|
|
if adapter is not None:
|
|
adapter.validate_stream_audio(has_audio=first_audio_chunk_s is not None)
|
|
# Check before committing the reference-audio artifact or logging
|
|
# success. Streaming protocols may already have emitted partial
|
|
# bytes, but they must terminate as an error rather than cleanly.
|
|
if tts_params is not None and usage_acc is not None:
|
|
self._validate_tts_generation(tts_params, usage_acc)
|
|
mod_metrics = getattr(getattr(self, "engine_client", None), "mod_metrics", None)
|
|
if mod_metrics is not None and audio_stage_id is not None and audio_replica_id is not None:
|
|
observe_audio_streaming_finalize(
|
|
mod_metrics,
|
|
stage_id=audio_stage_id,
|
|
replica_id=audio_replica_id,
|
|
chunk_arrival_times_s=audio_chunk_arrivals_s,
|
|
chunk_bytes=audio_chunk_bytes,
|
|
sample_rate=target_sample_rate or sample_rate_val,
|
|
channels=audio_channels,
|
|
)
|
|
self._mark_ref_audio_artifact_ready_for_request(request_id)
|
|
artifact_ready = True
|
|
if mod_metrics is not None:
|
|
mod_metrics.inc_speech_stream_completed()
|
|
total_ms = (time.perf_counter() - stream_start_s) * 1000.0
|
|
if first_audio_chunk_s is not None:
|
|
first_chunk_ms = (first_audio_chunk_s - stream_start_s) * 1000.0
|
|
logger.info(
|
|
"[SpeechE2E] request_id=%s stream=true status=ok total_ms=%.2f first_chunk_ms=%.2f",
|
|
request_id,
|
|
total_ms,
|
|
first_chunk_ms,
|
|
)
|
|
else:
|
|
logger.info(
|
|
"[SpeechE2E] request_id=%s stream=true status=ok total_ms=%.2f first_chunk_ms=NA",
|
|
request_id,
|
|
total_ms,
|
|
)
|
|
except GeneratorExit:
|
|
record_stream_abort("closed")
|
|
raise
|
|
except asyncio.CancelledError:
|
|
record_stream_abort("cancelled")
|
|
total_ms = (time.perf_counter() - stream_start_s) * 1000.0
|
|
logger.info(
|
|
"[SpeechE2E] request_id=%s stream=true status=cancelled total_ms=%.2f",
|
|
request_id,
|
|
total_ms,
|
|
)
|
|
logger.info("Streaming request %s cancelled by client", request_id)
|
|
raise
|
|
except EngineDeadError as e:
|
|
record_stream_abort("engine_dead")
|
|
total_ms = (time.perf_counter() - stream_start_s) * 1000.0
|
|
logger.error(
|
|
"[SpeechE2E] request_id=%s stream=true status=engine_dead total_ms=%.2f",
|
|
request_id,
|
|
total_ms,
|
|
)
|
|
logger.error(
|
|
"EngineDeadError during streaming speech for %s: %s",
|
|
request_id,
|
|
e,
|
|
)
|
|
# Actively signal shutdown rather than relying on the watchdog.
|
|
if raw_request is not None:
|
|
terminate_if_errored(
|
|
server=raw_request.app.state.server,
|
|
engine=self.engine_client,
|
|
)
|
|
raise
|
|
except Exception as e:
|
|
record_stream_abort("error")
|
|
total_ms = (time.perf_counter() - stream_start_s) * 1000.0
|
|
logger.exception(
|
|
"[SpeechE2E] request_id=%s stream=true status=error total_ms=%.2f error=%s",
|
|
request_id,
|
|
total_ms,
|
|
e,
|
|
)
|
|
logger.exception("Streaming speech generation failed for %s: %s", request_id, e)
|
|
raise
|
|
finally:
|
|
if not artifact_ready:
|
|
self._discard_ref_audio_artifact_warmup(request_id)
|
|
# Disconnects can arrive while suspended at yield. Closing the
|
|
# engine iterator must survive the cancelled ASGI scope.
|
|
close = getattr(generator, "aclose", None)
|
|
if close is not None:
|
|
with anyio.CancelScope(shield=True):
|
|
await close()
|
|
|
|
async def _generate_audio_sse_events(
|
|
self,
|
|
generator,
|
|
request_id: str,
|
|
response_format: str = "pcm",
|
|
raw_request: Request | None = None,
|
|
request_start_s: float | None = None,
|
|
request_arrival_ts: float | None = None,
|
|
request: OpenAICreateSpeechRequest | None = None,
|
|
tts_params: dict[str, Any] | None = None,
|
|
):
|
|
"""Generate OpenAI-style SSE events with base64 audio deltas.
|
|
|
|
Field naming follows the OpenAI ``speech.audio.delta`` schema, which
|
|
carries the base64 chunk in ``audio`` (not ``delta`` — that is the
|
|
Realtime API ``response.output_audio.delta`` convention, a different event).
|
|
See https://platform.openai.com/docs/api-reference/audio-streaming.
|
|
|
|
The terminal ``speech.audio.done`` event carries a ``usage`` object
|
|
(``input_tokens``/``output_tokens``/``total_tokens`` + a per-modality
|
|
``input_token_details`` breakdown), matching OpenAI's documented
|
|
``speech.audio.done`` schema. ``output_tokens`` is accumulated from the
|
|
stage-0 deltas as they stream (see ``SpeechOutputTokenCounter``);
|
|
``input_tokens`` is computed from the request text + reference audio.
|
|
"""
|
|
usage_acc = SpeechOutputTokenCounter()
|
|
emitted_audio = False
|
|
audio_format: dict[str, int] = {}
|
|
try:
|
|
async with aclosing(
|
|
self._generate_audio_chunks(
|
|
generator,
|
|
request_id,
|
|
response_format,
|
|
raw_request=raw_request,
|
|
request_start_s=request_start_s,
|
|
request_arrival_ts=request_arrival_ts,
|
|
usage_acc=usage_acc,
|
|
tts_params=tts_params,
|
|
target_sample_rate=request.sample_rate if request is not None else None,
|
|
audio_format=audio_format,
|
|
)
|
|
) as chunks:
|
|
async for chunk in chunks:
|
|
payload = {
|
|
"type": "speech.audio.delta",
|
|
"audio": base64.b64encode(chunk).decode("ascii"),
|
|
"response_format": response_format,
|
|
}
|
|
if audio_format:
|
|
# Additive to the OpenAI schema: pcm deltas are
|
|
# model-native s16le, not always 24 kHz mono.
|
|
payload["sample_rate"] = audio_format["sample_rate"]
|
|
payload["channels"] = audio_format["channels"]
|
|
data = json.dumps(payload, separators=(",", ":"))
|
|
emitted_audio = True
|
|
yield f"event: speech.audio.delta\ndata: {data}\n\n"
|
|
done_payload: dict[str, Any] = {"type": "speech.audio.done"}
|
|
if request is not None:
|
|
# Streaming path: output_tokens = sum of stage-0 deltas.
|
|
usage = self._build_speech_usage(request, tts_params or {}, usage_acc.total())
|
|
done_payload["usage"] = usage.model_dump()
|
|
done = json.dumps(done_payload, separators=(",", ":"))
|
|
yield f"event: speech.audio.done\ndata: {done}\n\n"
|
|
except asyncio.CancelledError:
|
|
raise
|
|
except Exception as e:
|
|
error: dict[str, Any] = {
|
|
"message": str(e),
|
|
"type": "server_error",
|
|
"param": None,
|
|
"code": HTTPStatus.INTERNAL_SERVER_ERROR.value,
|
|
}
|
|
if emitted_audio:
|
|
error.update(partial_audio=True, action="discard")
|
|
error_payload: dict[str, Any] = {
|
|
"type": "speech.audio.error",
|
|
"error": error,
|
|
}
|
|
data = json.dumps(error_payload, separators=(",", ":"))
|
|
yield f"event: speech.audio.error\ndata: {data}\n\n"
|
|
|
|
@staticmethod
|
|
def _is_timestamps_output(res) -> bool:
|
|
"""True when ``res`` is the forced-aligner stage's terminal timestamps output."""
|
|
from vllm_omni.model_executor.stage_input_processors.forced_aligner import TIMESTAMPS_MODALITY
|
|
|
|
return getattr(res, "final_output_type", None) == TIMESTAMPS_MODALITY
|
|
|
|
@staticmethod
|
|
def _extract_audio_output(res) -> tuple[dict | None, str | None]:
|
|
"""Return (audio_output dict, audio key) or (None, None).
|
|
|
|
Returns the raw dict so callers can apply their own extraction strategy:
|
|
streaming needs per-chunk delta slicing; non-streaming needs full concatenation.
|
|
"""
|
|
mm = getattr(res, "multimodal_output", None)
|
|
ro = None
|
|
if not mm:
|
|
ro = res
|
|
mm = getattr(ro, "multimodal_output", None) if ro else None
|
|
if not mm:
|
|
# MultimodalOutputProcessor attaches mm_accumulated on per-completion outputs.
|
|
container = res if hasattr(res, "outputs") else ro
|
|
outputs = getattr(container, "outputs", None) if container is not None else None
|
|
if outputs:
|
|
for completion_output in outputs:
|
|
completion_mm = getattr(completion_output, "multimodal_output", None)
|
|
if completion_mm:
|
|
mm = completion_mm
|
|
break
|
|
if not mm:
|
|
return None, None
|
|
key = "audio" if "audio" in mm else ("model_outputs" if "model_outputs" in mm else None)
|
|
return mm, key
|
|
|
|
async def _prepare_speech_generation(
|
|
self,
|
|
request: OpenAICreateSpeechRequest,
|
|
request_id: str | None = None,
|
|
has_inline_ref_audio: bool | None = None,
|
|
arrival_time: float | None = None,
|
|
) -> tuple[str, Any, dict[str, Any]]:
|
|
if self.engine_client.errored:
|
|
raise self.engine_client.dead_error
|
|
|
|
sample_rate_error = self._validate_speech_sample_rate(request)
|
|
if sample_rate_error is not None:
|
|
raise ValueError(sample_rate_error)
|
|
|
|
request_id = request_id or f"speech-{random_uuid()}"
|
|
qwen3_ref_audio_warmup_artifact_key: str | None = None
|
|
|
|
# If this is a streaming request with real async chunks, we need to
|
|
# coerce cumulative outputs to delta outputs; this ensures we don't
|
|
# emit redundant MM data & drain after emitting. Qwen3-TTS full-payload
|
|
# (async_chunk=False) has no incremental audio chunks, so keep
|
|
# FINAL_ONLY semantics and let the streaming response send the final
|
|
# waveform once. Scoped to qwen3_tts: other async_chunk=False models
|
|
# keep the DELTA coercion they stream with today.
|
|
# list() makes a copy to avoid mutating the params.
|
|
sampling_params_list = list(self.engine_client.default_sampling_params_list)
|
|
async_chunk = getattr(self.model_config, "async_chunk", True)
|
|
qwen3_full_payload = self._tts_model_type == "qwen3_tts" and not bool(async_chunk)
|
|
is_streaming_request = request.is_streaming() and not qwen3_full_payload
|
|
sampling_params_list = coerce_param_message_types(sampling_params_list, is_streaming_request)
|
|
|
|
# Build prompt + tts_params via the per-model adapter (RFC #4327). Every
|
|
# dedicated TTS model resolves to an adapter that owns its validation,
|
|
# uploaded-speaker handling, prompt/param building, and sampling
|
|
# overrides. The model-type label remains available for compatibility.
|
|
# Non-TTS deployments (no adapter) fall through to the rejection below.
|
|
# Capture inline-ref-audio status BEFORE validate(): several adapters
|
|
# apply uploaded speakers inside validate(), which sets request.ref_audio
|
|
# in place. The builders need to know whether the caller supplied audio
|
|
# inline vs. via an uploaded voice.
|
|
model_type: str | None = None
|
|
has_inline_ref_audio = (request.ref_audio is not None) if has_inline_ref_audio is None else has_inline_ref_audio
|
|
if (adapter := self._get_tts_adapter()) is not None:
|
|
adapter.normalize(request)
|
|
validation_error = adapter.validate(request)
|
|
if validation_error:
|
|
raise ValueError(validation_error)
|
|
prepared = await adapter.build(request, sampling_params_list, has_inline_ref_audio)
|
|
prompt = prepared.prompt
|
|
tts_params = prepared.tts_params
|
|
model_type = prepared.model_type
|
|
qwen3_ref_audio_warmup_artifact_key = prepared.warmup_artifact_key
|
|
output_policy = prepared.output_policy
|
|
else:
|
|
# Qwen omni models (Qwen3-Omni, Qwen2.5-Omni) use a "talker"
|
|
# stage whose preprocess requires chat-templated tokens. The
|
|
# async-chunk orchestrator prewarms the talker via
|
|
# compute_talker_prompt_ids_length(), which scans for Qwen
|
|
# chat-template markers (im_start_token_id 151644). A raw-text
|
|
# prompt produces a 1-token placeholder that crashes the talker's
|
|
# prefill/decode handoff. Reject early with an actionable message.
|
|
stage_names = {_stage_speech_metadata(s)[0] for s in self.engine_client.stage_configs}
|
|
if "talker" in stage_names:
|
|
raise ValueError(
|
|
"The /v1/audio/speech endpoint is only supported for "
|
|
"dedicated TTS models (e.g., Qwen3-TTS, Voxtral, Fish "
|
|
"Speech, CosyVoice3, OmniVoice, VoxCPM2). For omni "
|
|
"models like Qwen3-Omni, use /v1/chat/completions with "
|
|
'\'"modalities": ["audio"]\' instead.'
|
|
)
|
|
tts_params = {}
|
|
prompt = {"prompt": request.input}
|
|
output_policy = OutputPolicy()
|
|
|
|
if model_type is None:
|
|
if self._is_tts:
|
|
model_type = tts_params.get("task_type", ["unknown"])[0]
|
|
else:
|
|
model_type = "generic"
|
|
logger.info("TTS speech request %s: model=%s", request_id, model_type)
|
|
_rl = getattr(self, "request_logger", None)
|
|
if _rl:
|
|
base_len = len(f"TTS speech request {request_id}: text=")
|
|
raw_max = getattr(_rl, "max_log_len", None)
|
|
cap = raw_max if isinstance(raw_max, int) else 200
|
|
text = request.input[: max(cap - base_len, 0)]
|
|
logger.debug("TTS speech request %s: text=%r", request_id, text)
|
|
|
|
# Apply model-specific extra parameters
|
|
if request.extra_params is not None and sampling_params_list:
|
|
if not isinstance(request.extra_params, dict):
|
|
raise HTTPException(
|
|
status_code=HTTPStatus.BAD_REQUEST.value,
|
|
detail="extra_params must be a JSON object/dict.",
|
|
)
|
|
import copy
|
|
|
|
sampling_params_list = copy.deepcopy(sampling_params_list)
|
|
for name in ("temperature", "top_p", "top_k"):
|
|
if (value := request.extra_params.get(name)) is not None:
|
|
setattr(sampling_params_list[0], name, value)
|
|
if sampling_params_list[0].extra_args is None:
|
|
sampling_params_list[0].extra_args = {}
|
|
sampling_params_list[0].extra_args.update(request.extra_params)
|
|
logger.info("Applied extra_params: %s", request.extra_params)
|
|
|
|
# Apply adapter-owned sampling overrides, including request-level token
|
|
# limits and model-specific dynamic token or stop-token configuration.
|
|
if sampling_params_list and (adapter := self._get_tts_adapter()) is not None:
|
|
sampling_params_list = adapter.apply_sampling_overrides(sampling_params_list, request, prompt, request_id)
|
|
|
|
if request.seed is not None and sampling_params_list:
|
|
import copy
|
|
|
|
sampling_params_list = copy.deepcopy(sampling_params_list)
|
|
stage0_params = sampling_params_list[0]
|
|
stage0_params.seed = request.seed
|
|
if stage0_params.extra_args is None:
|
|
stage0_params.extra_args = {}
|
|
stage0_params.extra_args["tts_local_seed"] = request.seed
|
|
|
|
# When word_timestamps is requested, also ask for the aligner stage's
|
|
# output so the orchestrator drives the request through the forced-aligner
|
|
# stage (final_stage_id extends to it). Harmless if no aligner stage exists.
|
|
output_modalities = ["audio"]
|
|
if getattr(request, "word_timestamps", False):
|
|
from vllm_omni.model_executor.stage_input_processors.forced_aligner import TIMESTAMPS_MODALITY
|
|
|
|
output_modalities.append(TIMESTAMPS_MODALITY)
|
|
if is_streaming_request:
|
|
# The aligner consumes the terminal waveform, so retain earlier
|
|
# audio chunks instead of draining them after each DELTA output.
|
|
sampling_params_list = [
|
|
sp.clone() if isinstance(sp, SamplingParams) else sp for sp in sampling_params_list
|
|
]
|
|
for sp in sampling_params_list:
|
|
if isinstance(sp, SamplingParams):
|
|
sp.output_kind = RequestOutputKind.CUMULATIVE
|
|
|
|
generator = self.engine_client.generate(
|
|
prompt=prompt,
|
|
request_id=request_id,
|
|
sampling_params_list=sampling_params_list,
|
|
output_modalities=output_modalities,
|
|
arrival_time=arrival_time,
|
|
)
|
|
self._track_ref_audio_artifact_warmup(
|
|
request_id,
|
|
qwen3_ref_audio_warmup_artifact_key,
|
|
x_vector_only=self._tts_x_vector_only(tts_params),
|
|
)
|
|
# Only the non-streaming path consumes output policies (in
|
|
# `_generate_audio_bytes`); streaming never reads them. Store late —
|
|
# after every fallible step above — and skip streaming requests
|
|
# entirely, so no caller has to discard entries (request ids are
|
|
# unique per request, so a stranded entry would never be reclaimed).
|
|
if not request.is_streaming():
|
|
self._speech_output_policies[request_id] = output_policy
|
|
return request_id, generator, tts_params
|
|
|
|
async def _generate_pcm_chunks(
|
|
self,
|
|
generator,
|
|
request_id: str,
|
|
*,
|
|
request_start_s: float | None = None,
|
|
request_arrival_ts: float | None = None,
|
|
include_sample_rate: bool = False,
|
|
tts_params: dict[str, Any] | None = None,
|
|
collect: dict | None = None,
|
|
target_sample_rate: int | None = None,
|
|
cumulative_audio: bool = False,
|
|
audio_format: dict[str, int] | None = None,
|
|
):
|
|
"""Yield raw PCM byte chunks from the engine generator.
|
|
|
|
Delegates to ``_generate_audio_chunks`` with ``response_format="pcm"``.
|
|
Used by the WebSocket streaming handler and ``_iter_pcm_audio_bytes``.
|
|
``collect`` (when given) receives the forced-aligner stage's pooling
|
|
output under ``"aligner_res"`` for downstream word-timestamp extraction.
|
|
"""
|
|
async with aclosing(
|
|
self._generate_audio_chunks(
|
|
generator,
|
|
request_id,
|
|
response_format="pcm",
|
|
request_start_s=request_start_s,
|
|
request_arrival_ts=request_arrival_ts,
|
|
include_sample_rate=include_sample_rate,
|
|
tts_params=tts_params,
|
|
collect=collect,
|
|
target_sample_rate=target_sample_rate,
|
|
cumulative_audio=cumulative_audio,
|
|
audio_format=audio_format,
|
|
)
|
|
) as chunks:
|
|
async for chunk in chunks:
|
|
yield chunk
|
|
|
|
async def _iter_pcm_audio_bytes(self, request: OpenAICreateSpeechRequest):
|
|
"""Yield raw PCM bytes for a speech request as soon as chunks are decoded."""
|
|
request_id, generator, tts_params = await self._prepare_speech_generation(request)
|
|
try:
|
|
async with aclosing(
|
|
self._generate_pcm_chunks(
|
|
generator,
|
|
request_id,
|
|
tts_params=tts_params,
|
|
target_sample_rate=request.sample_rate,
|
|
cumulative_audio=request.word_timestamps,
|
|
)
|
|
) as chunks:
|
|
async for chunk in chunks:
|
|
yield chunk
|
|
finally:
|
|
self._discard_ref_audio_artifact_warmup(request_id)
|
|
|
|
async def _generate_audio_bytes(
|
|
self,
|
|
request: OpenAICreateSpeechRequest,
|
|
base64_encode: bool = False,
|
|
request_id: str | None = None,
|
|
usage_out: list[SpeechTokenUsage] | None = None,
|
|
has_inline_ref_audio: bool | None = None,
|
|
collect: dict | None = None,
|
|
request_arrival_ts: float | None = None,
|
|
) -> tuple[bytes | str, str]:
|
|
# ``usage_out`` is an opt-in output channel: when a list is passed, the
|
|
# computed SpeechTokenUsage is appended to it. The return stays a
|
|
# 2-tuple so existing callers (and their test mocks) are unaffected;
|
|
# batch and non-streaming response-header paths opt in when surfacing
|
|
# usage outside the raw audio body.
|
|
request_arrival_ts = request_arrival_ts if request_arrival_ts is not None else time.time()
|
|
request_id, generator, bytes_tts_params = await self._prepare_speech_generation(
|
|
request,
|
|
request_id=request_id,
|
|
has_inline_ref_audio=has_inline_ref_audio,
|
|
arrival_time=request_arrival_ts,
|
|
)
|
|
artifact_ready = False
|
|
|
|
try:
|
|
# Sparse-audio models (MOSS-TTS-Nano, Gepard) emit delta chunks per
|
|
# yield with async_chunk=false. The engine surfaces each yield as its
|
|
# own RequestOutput, so we accumulate across the async-for loop when
|
|
# the adapter asks for it — final_output alone may only carry the
|
|
# last (often empty) sentinel. Prefer the engine's own concatenated
|
|
# waveform when FINAL_ONLY already produced one.
|
|
accumulate_nonstreaming = self._speech_output_policies.pop(
|
|
request_id, OutputPolicy()
|
|
).accumulate_nonstreaming
|
|
delta_chunks: list[Any] = []
|
|
delta_sample_rate: int | None = None
|
|
|
|
final_output: OmniRequestOutput | None = None
|
|
# Non-streaming is FINAL_ONLY, so the stage-0 output carries the full
|
|
# token sequence; the counter records its length for output_tokens.
|
|
usage_acc = SpeechOutputTokenCounter()
|
|
audio_res: OmniRequestOutput | None = None
|
|
aligner_res: OmniRequestOutput | None = None
|
|
async for res in generator:
|
|
final_output = res
|
|
usage_acc.observe(res)
|
|
# The generator yields both the audio output (Code2Wav) and, with
|
|
# a forced-aligner stage, a timestamps output. Keep the audio res
|
|
# for the WAV and the aligner res for word timestamps.
|
|
if self._is_timestamps_output(res):
|
|
aligner_res = res
|
|
else:
|
|
_, audio_key = self._extract_audio_output(res)
|
|
if audio_key is not None:
|
|
audio_res = res
|
|
if not accumulate_nonstreaming:
|
|
continue
|
|
try:
|
|
step_audio, step_key = self._extract_audio_output(res)
|
|
except Exception:
|
|
continue
|
|
if step_key is None:
|
|
continue
|
|
step_audio = cast(dict, step_audio)
|
|
chunk = step_audio[step_key]
|
|
candidates = chunk if isinstance(chunk, list) else [chunk]
|
|
for cand in candidates:
|
|
if hasattr(cand, "numel") and cand.numel() > 0:
|
|
delta_chunks.append(cand)
|
|
sr_step = step_audio.get("sr")
|
|
if sr_step is not None:
|
|
sr_val_step = sr_step[-1] if isinstance(sr_step, list) and sr_step else sr_step
|
|
delta_sample_rate = int(sr_val_step.item()) if hasattr(sr_val_step, "item") else int(sr_val_step)
|
|
|
|
if final_output is None:
|
|
raise ValueError("No output generated from the model.")
|
|
|
|
self._validate_tts_generation(bytes_tts_params or {}, usage_acc)
|
|
|
|
# Extract audio from the audio-bearing res (not necessarily the last
|
|
# yielded one, which may be the aligner's timestamps output).
|
|
audio_source = audio_res if audio_res is not None else final_output
|
|
audio_output, audio_key = self._extract_audio_output(audio_source)
|
|
if audio_key is None:
|
|
raise ValueError("TTS model did not produce audio output.")
|
|
audio_output = cast(dict, audio_output)
|
|
|
|
# Surface forced-aligner word timestamps to the caller (set as a
|
|
# response header) when requested and an aligner stage produced them.
|
|
if collect is not None and getattr(request, "word_timestamps", False):
|
|
from vllm_omni.utils.forced_aligner import extract_word_timestamps
|
|
|
|
ts = (
|
|
extract_word_timestamps(aligner_res, request.input, getattr(request, "language", None))
|
|
if aligner_res is not None
|
|
else None
|
|
)
|
|
if ts is not None:
|
|
collect["word_timestamps"] = ts
|
|
|
|
# Let the adapter fold engine-side metadata (e.g. YuE2's
|
|
# meta.truncated) into ``collect`` for response headers.
|
|
if collect is not None and (adapter := self._get_tts_adapter()) is not None:
|
|
adapter.collect_response_metadata(audio_output, collect)
|
|
|
|
# A model can flag a per-request synthesis failure through the
|
|
# adapter (e.g. YuE2's terminal NAR/VAE pass OOMing on one
|
|
# request); answer 500 instead of shipping a zero-length WAV.
|
|
# Raising (not returning a Response) keeps this function's
|
|
# tuple contract; create_speech maps TTSGenerationError to 500.
|
|
if collect is not None and collect.get("audio_synthesis_error"):
|
|
raise TTSGenerationError(
|
|
"The model failed to synthesize audio for this request",
|
|
retryable=False,
|
|
)
|
|
|
|
audio_tensor = audio_output[audio_key]
|
|
sr_raw = audio_output.get("sr", 24000)
|
|
sr_val = sr_raw[-1] if isinstance(sr_raw, list) and sr_raw else sr_raw
|
|
sample_rate = sr_val.item() if hasattr(sr_val, "item") else int(sr_val)
|
|
|
|
if accumulate_nonstreaming:
|
|
# Prefer the engine's own consolidated audio when present. After the
|
|
# vllm 0.20 rebase non-stream requests resolve to FINAL_ONLY, so
|
|
# final_output already carries the full concatenated waveform; the
|
|
# delta-accumulator below is kept as a fallback for DELTA-style
|
|
# engines that surface chunks one yield at a time.
|
|
if isinstance(audio_tensor, list):
|
|
non_empty_final = [c for c in audio_tensor if hasattr(c, "numel") and c.numel() > 0]
|
|
final_audio = torch.cat(non_empty_final, dim=-1) if non_empty_final else None
|
|
elif hasattr(audio_tensor, "numel") and audio_tensor.numel() > 0:
|
|
final_audio = audio_tensor
|
|
else:
|
|
final_audio = None
|
|
|
|
if final_audio is not None:
|
|
audio_tensor = final_audio
|
|
elif delta_chunks:
|
|
audio_tensor = torch.cat(delta_chunks, dim=-1)
|
|
else:
|
|
audio_tensor = np.zeros((0,), dtype=np.float32)
|
|
if delta_sample_rate is not None:
|
|
sample_rate = delta_sample_rate
|
|
elif isinstance(audio_tensor, list):
|
|
async_chunk = bool(getattr(self.engine_client.model_config, "async_chunk", False))
|
|
if async_chunk:
|
|
non_empty_chunks = [candidate for candidate in audio_tensor if candidate.numel() > 0]
|
|
audio_tensor = (
|
|
torch.cat(non_empty_chunks, dim=-1) if non_empty_chunks else np.zeros((0,), dtype=np.float32)
|
|
)
|
|
else:
|
|
audio_history = audio_tensor
|
|
audio_tensor = np.zeros((0,), dtype=np.float32)
|
|
# Non-async Qwen3-TTS returns cumulative history snapshots, so keep the latest non-empty tensor.
|
|
for candidate in reversed(audio_history):
|
|
if candidate.numel() > 0:
|
|
audio_tensor = candidate
|
|
break
|
|
if hasattr(audio_tensor, "float"):
|
|
audio_tensor = audio_tensor.float().detach().cpu().numpy()
|
|
|
|
if audio_tensor.ndim > 1:
|
|
audio_tensor = audio_tensor.squeeze()
|
|
|
|
if self._tts_model_type in _AUDEX_NO_AUDIO_GUARD_MODEL_TYPES and int(np.size(audio_tensor)) == 0:
|
|
# Audex contract: zero codec tokens must fail the request, not
|
|
# serialize as an empty-but-successful WAV.
|
|
raise ValueError("Audex produced no audio (the thinker emitted zero or invalid codec tokens)")
|
|
|
|
audio_obj = CreateAudio(
|
|
audio_tensor=audio_tensor,
|
|
sample_rate=sample_rate,
|
|
output_sample_rate=request.sample_rate,
|
|
response_format=request.response_format or "wav",
|
|
speed=self._audio_encode_speed(request),
|
|
base64_encode=base64_encode,
|
|
)
|
|
audio_response: AudioResponse = self.create_audio(audio_obj)
|
|
if collect is not None:
|
|
collect["audio_format"] = _encoded_audio_format(
|
|
audio_response, request.sample_rate or sample_rate, audio_tensor
|
|
)
|
|
self._mark_ref_audio_artifact_ready_for_request(request_id)
|
|
artifact_ready = True
|
|
if usage_out is not None:
|
|
usage_out.append(self._build_speech_usage(request, bytes_tts_params or {}, usage_acc.total()))
|
|
return audio_response.audio_data, audio_response.media_type
|
|
finally:
|
|
if not artifact_ready:
|
|
self._discard_ref_audio_artifact_warmup(request_id)
|
|
|
|
def _get_normalized_voice(self, voice: str | None) -> str | None:
|
|
"""Get the normalized voice to be used; currently this means that
|
|
the voice is a:
|
|
- lowercase str if it's a valid supported/uploaded speaker
|
|
- None if the voice is the placeholder default or not provided
|
|
"""
|
|
if voice is not None:
|
|
voice = voice.lower()
|
|
available_speakers = self._get_available_speakers()
|
|
if voice not in available_speakers:
|
|
raise ValueError(
|
|
f"Invalid voice '{voice}'. Supported: {', '.join(sorted(self._get_available_voices()))}"
|
|
)
|
|
return voice
|
|
|
|
async def _create_diffusion_speech(
|
|
self,
|
|
request: OpenAICreateSpeechRequest,
|
|
) -> Response:
|
|
"""Handle speech generation for pure diffusion TTS models (e.g. OmniVoice)."""
|
|
from vllm_omni.outputs import OmniRequestOutput
|
|
|
|
try:
|
|
if not request.input or not request.input.strip():
|
|
raise ValueError("Input text cannot be empty")
|
|
|
|
if request.ref_audio is not None:
|
|
fmt_err = self._validate_ref_audio_format(request.ref_audio)
|
|
if fmt_err:
|
|
return self._diffusion_error_response(fmt_err, status_code=400)
|
|
|
|
request.voice = self._get_normalized_voice(request.voice)
|
|
|
|
has_inline_ref_audio = request.ref_audio is not None
|
|
err = self._apply_uploaded_speaker(request)
|
|
if err:
|
|
raise ValueError(err)
|
|
|
|
request_id = f"speech-{random_uuid()}"
|
|
prompt: dict[str, Any] = {"input": request.input}
|
|
if request.ref_audio:
|
|
wav, sr, _ = await self._resolve_ref_audio(cast(str, request.ref_audio))
|
|
prompt["ref_audio"] = (np.asarray(wav, dtype=np.float32), sr)
|
|
if request.ref_text:
|
|
prompt["ref_text"] = request.ref_text
|
|
if request.voice:
|
|
if request.voice in self.uploaded_speakers and not has_inline_ref_audio:
|
|
prompt["voice_name"] = request.voice
|
|
prompt["voice_created_at"] = self._voice_created_at(request.voice)
|
|
if request.language:
|
|
prompt["lang"] = request.language
|
|
if request.instructions:
|
|
prompt["instruct"] = request.instructions
|
|
|
|
logger.info(
|
|
"Diffusion TTS speech request %s: voice_clone=%s",
|
|
request_id,
|
|
"ref_audio" in prompt,
|
|
)
|
|
_rl = getattr(self, "request_logger", None)
|
|
if _rl:
|
|
base_len = len(f"Diffusion TTS speech request {request_id}: text=")
|
|
raw_max = getattr(_rl, "max_log_len", None)
|
|
cap = raw_max if isinstance(raw_max, int) else 200
|
|
text = request.input[: max(cap - base_len, 0)]
|
|
logger.debug("Diffusion TTS speech request %s: text=%r", request_id, text)
|
|
if request.extra_params is not None and not isinstance(request.extra_params, dict):
|
|
raise ValueError("extra_params must be a JSON object/dict.")
|
|
extra = dict(request.extra_params or {})
|
|
if request.seed is not None:
|
|
extra["seed"] = request.seed
|
|
# Apply extra_params from the request to sampling params
|
|
sampling_params_list = self._diffusion_engine.default_sampling_params_list
|
|
if extra:
|
|
import copy
|
|
|
|
sampling_params_list = copy.deepcopy(sampling_params_list)
|
|
if sampling_params_list[0].extra_args is None:
|
|
sampling_params_list[0].extra_args = {}
|
|
sampling_params_list[0].extra_args.update(extra)
|
|
|
|
sampling = sampling_params_list[0]
|
|
|
|
# This change allows StepScheduler read total_steps from upper
|
|
# sampling.num_inference_steps, check diffusion/sched/step_scheduler:_get_total_steps
|
|
if "num_inference_steps" in extra:
|
|
value = extra["num_inference_steps"]
|
|
try:
|
|
sampling.num_inference_steps = int(value)
|
|
except (TypeError, ValueError) as exc:
|
|
raise ValueError("num_inference_steps must be an integer") from exc
|
|
|
|
if "guidance_scale" in extra:
|
|
value = extra["guidance_scale"]
|
|
try:
|
|
sampling.guidance_scale = float(value)
|
|
except (TypeError, ValueError) as exc:
|
|
raise ValueError("guidance_scale must be a number") from exc
|
|
|
|
logger.info("Applied extra_params to diffusion: %s", extra)
|
|
|
|
generator = self._diffusion_engine.generate(
|
|
prompt=prompt,
|
|
request_id=request_id,
|
|
sampling_params_list=sampling_params_list,
|
|
output_modalities=["audio"],
|
|
)
|
|
|
|
final_output: OmniRequestOutput | None = None
|
|
async for res in generator:
|
|
final_output = res
|
|
|
|
if final_output is None:
|
|
raise ValueError("No output generated from the model.")
|
|
|
|
audio_output, audio_key = self._extract_audio_output(final_output)
|
|
if audio_key is None:
|
|
raise ValueError("TTS model did not produce audio output.")
|
|
audio_output = cast(dict, audio_output)
|
|
|
|
audio_tensor = audio_output[audio_key]
|
|
sr_raw = audio_output.get("sr", 24000)
|
|
sr_val = sr_raw[-1] if isinstance(sr_raw, list) and sr_raw else sr_raw
|
|
sample_rate = sr_val.item() if hasattr(sr_val, "item") else int(sr_val)
|
|
|
|
if isinstance(audio_tensor, list):
|
|
non_empty = [c for c in audio_tensor if c.numel() > 0]
|
|
audio_tensor = torch.cat(non_empty, dim=-1) if non_empty else np.zeros((0,), dtype=np.float32)
|
|
if hasattr(audio_tensor, "float"):
|
|
audio_tensor = audio_tensor.float().detach().cpu().numpy()
|
|
if audio_tensor.ndim > 1:
|
|
audio_tensor = audio_tensor.squeeze()
|
|
|
|
audio_obj = CreateAudio(
|
|
audio_tensor=audio_tensor,
|
|
sample_rate=sample_rate,
|
|
output_sample_rate=request.sample_rate,
|
|
response_format=request.response_format or "wav",
|
|
speed=self._audio_encode_speed(request),
|
|
base64_encode=False,
|
|
)
|
|
audio_response: AudioResponse = self.create_audio(audio_obj)
|
|
return Response(content=audio_response.audio_data, media_type=audio_response.media_type)
|
|
|
|
except asyncio.CancelledError:
|
|
return self._diffusion_error_response("Client disconnected")
|
|
except (EngineGenerateError, EngineDeadError):
|
|
raise # Propagate to the global Omni exception handler
|
|
except ValueError as e:
|
|
return self._diffusion_error_response(str(e), status_code=400)
|
|
except Exception as e:
|
|
logger.exception("Diffusion speech generation failed: %s", e)
|
|
return self._diffusion_error_response(f"Speech generation failed: {e}")
|
|
|
|
@staticmethod
|
|
def _diffusion_error_response(message: str, status_code: int = 500) -> Response:
|
|
"""Create a JSON error response without depending on OpenAIServing.
|
|
|
|
Args:
|
|
message: Error message to surface to the client.
|
|
status_code: HTTP status code; defaults to 500. Pass a 4xx code for
|
|
client-input validation failures so the response semantics match
|
|
the OpenAI-compatible behavior used by ``create_speech``.
|
|
"""
|
|
err_type = "BadRequestError" if 400 <= status_code < 500 else "server_error"
|
|
error_body = json.dumps({"error": {"message": message, "type": err_type, "param": None, "code": status_code}})
|
|
return Response(content=error_body, media_type="application/json", status_code=status_code)
|
|
|
|
def _validate_speech_streaming_request(
|
|
self,
|
|
request: OpenAICreateSpeechRequest,
|
|
*,
|
|
mode_label: str,
|
|
) -> tuple[str, Response | None]:
|
|
"""Validate pcm/wav + speed constraints for streaming speech responses."""
|
|
response_format = (request.response_format or "wav").lower()
|
|
if response_format not in ("pcm", "wav"):
|
|
return response_format, self.create_error_response(
|
|
f"{mode_label} is only supported for 'pcm' and 'wav' formats. Got '{response_format}'."
|
|
)
|
|
if request.speed is not None and request.speed != 1.0 and not self._uses_native_speed_control():
|
|
return response_format, self.create_error_response(
|
|
f"{mode_label} is not supported with speed adjustment. "
|
|
"Use a non-streaming request or remove the speed parameter."
|
|
)
|
|
return response_format, None
|
|
|
|
async def create_speech(
|
|
self,
|
|
request: OpenAICreateSpeechRequest,
|
|
raw_request: Request | None = None,
|
|
):
|
|
"""
|
|
Create Speech API similar to OpenAI's API.
|
|
|
|
See https://platform.openai.com/docs/api-reference/audio/createSpeech
|
|
for the API specification. This API mimics the OpenAI
|
|
Create Speech API.
|
|
|
|
For Qwen3-TTS models, additional parameters are supported:
|
|
- task_type: "CustomVoice", "VoiceDesign", or "Base"
|
|
- language: Language code (e.g., "Chinese", "English", "Auto")
|
|
- voice: Speaker name (e.g., "Vivian", "Ryan") for CustomVoice
|
|
- instructions: Voice style/emotion instructions
|
|
- ref_audio: Reference audio for voice cloning (Base task)
|
|
- ref_text: Transcript of reference audio (Base task)
|
|
- x_vector_only_mode: Use speaker embedding only (Base task)
|
|
- sample_rate: Target output sample rate (8000 or 24000 Hz)
|
|
|
|
Streaming is supported via the ``stream=True`` switch or ``stream_format='sse'``,
|
|
which return OpenAI ``speech.audio.*`` SSE events. ``stream_format='audio'``
|
|
opts into raw audio streaming with ``response_format='pcm'`` or ``'wav'``.
|
|
Raw audio streaming yields each Code2Wav chunk as raw bytes as soon as it is
|
|
decoded. Raw WAV streaming emits a header with placeholder size values first.
|
|
"""
|
|
if request.voice is not None:
|
|
if _is_default_voice(request.voice.lower(), self._get_available_speakers()):
|
|
request.voice = None
|
|
|
|
sample_rate_error = self._validate_speech_sample_rate(request)
|
|
if sample_rate_error is not None:
|
|
return self.create_error_response(sample_rate_error)
|
|
|
|
if self._diffusion_mode:
|
|
return await self._create_diffusion_speech(request)
|
|
|
|
error_check_ret = await self._check_model(request)
|
|
if error_check_ret is not None:
|
|
logger.error("Error with model %s", error_check_ret)
|
|
return error_check_ret
|
|
|
|
request_id = f"speech-{random_uuid()}"
|
|
request_start_s = time.perf_counter()
|
|
request_arrival_ts = (
|
|
float(getattr(raw_request.state, "request_timestamp", time.time()))
|
|
if raw_request is not None
|
|
else time.time()
|
|
)
|
|
if raw_request:
|
|
raw_request.state.request_metadata = RequestResponseMetadata(
|
|
request_id=request_id,
|
|
)
|
|
|
|
try:
|
|
if request.is_streaming() and request.word_timestamps:
|
|
return self.create_error_response(
|
|
"word_timestamps=true is currently supported by the WebSocket "
|
|
"/v1/audio/speech/stream path. Use session.config with "
|
|
"stream_audio=true and response_format='pcm'."
|
|
)
|
|
if request.word_timestamps and not self.forced_aligner_enabled:
|
|
# Fail loud instead of silently returning 200 with no timestamps header.
|
|
return self.create_error_response(
|
|
"word_timestamps=true requires the server to be launched with --forced-aligner."
|
|
)
|
|
|
|
if request.is_raw_audio_stream():
|
|
response_format, error = self._validate_speech_streaming_request(
|
|
request,
|
|
mode_label="Streaming",
|
|
)
|
|
if error is not None:
|
|
return error
|
|
|
|
media_type = "audio/wav" if response_format == "wav" else "audio/pcm"
|
|
_, generator, raw_tts_params = await self._prepare_speech_generation(
|
|
request,
|
|
request_id=request_id,
|
|
arrival_time=request_arrival_ts,
|
|
)
|
|
audio_format: dict[str, int] = {}
|
|
# Pull the first chunk before sending headers: the body has no
|
|
# self-describing format for pcm and the rate/channel count is
|
|
# model-native (not always 24 kHz mono), so the headers state it.
|
|
body = await _PrimedAudioStream.start(
|
|
self._generate_audio_chunks(
|
|
generator,
|
|
request_id,
|
|
response_format,
|
|
raw_request=raw_request,
|
|
request_start_s=request_start_s,
|
|
request_arrival_ts=request_arrival_ts,
|
|
tts_params=raw_tts_params,
|
|
target_sample_rate=request.sample_rate,
|
|
audio_format=audio_format,
|
|
)
|
|
)
|
|
return _SpeechStreamingResponse(
|
|
body,
|
|
media_type=media_type,
|
|
headers=_audio_format_headers(audio_format),
|
|
)
|
|
|
|
if request.is_sse_stream():
|
|
response_format, error = self._validate_speech_streaming_request(
|
|
request,
|
|
mode_label="SSE streaming",
|
|
)
|
|
if error is not None:
|
|
return error
|
|
|
|
_, generator, sse_tts_params = await self._prepare_speech_generation(
|
|
request,
|
|
request_id=request_id,
|
|
arrival_time=request_arrival_ts,
|
|
)
|
|
return _SpeechStreamingResponse(
|
|
self._generate_audio_sse_events(
|
|
generator,
|
|
request_id,
|
|
response_format,
|
|
raw_request=raw_request,
|
|
request_start_s=request_start_s,
|
|
request_arrival_ts=request_arrival_ts,
|
|
request=request,
|
|
tts_params=sse_tts_params,
|
|
),
|
|
media_type="text/event-stream",
|
|
)
|
|
|
|
collect: dict = {}
|
|
usage_box: list[SpeechTokenUsage] = []
|
|
try:
|
|
audio_bytes, media_type = await self._generate_audio_bytes(
|
|
request,
|
|
request_id=request_id,
|
|
usage_out=usage_box,
|
|
collect=collect,
|
|
request_arrival_ts=request_arrival_ts,
|
|
)
|
|
except TTSGenerationError as error:
|
|
# An adapter can reject otherwise completed audio. Retry only
|
|
# retryable, stochastic, server-budgeted non-streaming
|
|
# requests: changing an explicit seed would violate
|
|
# reproducibility, and changing a caller-provided token limit
|
|
# would violate the requested budget.
|
|
if not error.retryable or request.seed is not None or request.max_new_tokens is not None:
|
|
raise
|
|
retry_seed = int(random_uuid()[:8], 16)
|
|
retry_request = request.model_copy(update={"seed": retry_seed})
|
|
retry_request_id = f"speech-{random_uuid()}"
|
|
collect.clear()
|
|
usage_box.clear()
|
|
logger.warning(
|
|
"TTS request %s failed generation validation; retrying once as %s with seed=%d",
|
|
request_id,
|
|
retry_request_id,
|
|
retry_seed,
|
|
)
|
|
audio_bytes, media_type = await self._generate_audio_bytes(
|
|
retry_request,
|
|
request_id=retry_request_id,
|
|
usage_out=usage_box,
|
|
collect=collect,
|
|
request_arrival_ts=request_arrival_ts,
|
|
)
|
|
total_ms = (time.perf_counter() - request_start_s) * 1000.0
|
|
logger.info(
|
|
"[SpeechE2E] request_id=%s stream=false status=ok total_ms=%.2f response_bytes=%d",
|
|
request_id,
|
|
total_ms,
|
|
len(audio_bytes) if isinstance(audio_bytes, (bytes, bytearray)) else len(str(audio_bytes)),
|
|
)
|
|
headers = self._build_speech_usage_headers(usage_box[0] if usage_box else None)
|
|
if collect.get("word_timestamps") is not None:
|
|
# Default ensure_ascii keeps the header latin-1 encodable (non-ASCII words \uXXXX-escaped).
|
|
ts_json = json.dumps(collect["word_timestamps"])
|
|
# Cap at 4 KB: oversized headers turn into opaque 502s at common reverse-proxy defaults.
|
|
if len(ts_json) <= 4096:
|
|
headers["X-Word-Timestamps"] = ts_json
|
|
else:
|
|
# Marker header so clients can tell an oversized alignment from no alignment.
|
|
headers["X-Word-Timestamps-Omitted"] = f"oversize; bytes={len(ts_json)}; limit=4096"
|
|
logger.warning(
|
|
"X-Word-Timestamps header omitted: %d bytes exceeds the 4 KB budget "
|
|
"(use the WebSocket streaming path for long transcripts)",
|
|
len(ts_json),
|
|
)
|
|
if collect.get("audio_truncated") is not None:
|
|
headers["X-Audio-Truncated"] = "true" if collect["audio_truncated"] else "false"
|
|
headers.update(_audio_format_headers(collect.get("audio_format")))
|
|
return Response(content=audio_bytes, media_type=media_type, headers=headers)
|
|
|
|
except asyncio.CancelledError:
|
|
total_ms = (time.perf_counter() - request_start_s) * 1000.0
|
|
logger.info(
|
|
"[SpeechE2E] request_id=%s stream=%s status=cancelled total_ms=%.2f",
|
|
request_id,
|
|
bool(request.stream),
|
|
total_ms,
|
|
)
|
|
return self.create_error_response("Client disconnected")
|
|
except (EngineGenerateError, EngineDeadError):
|
|
total_ms = (time.perf_counter() - request_start_s) * 1000.0
|
|
logger.error(
|
|
"[SpeechE2E] request_id=%s stream=%s status=engine_error total_ms=%.2f",
|
|
request_id,
|
|
bool(request.stream),
|
|
total_ms,
|
|
)
|
|
raise # Propagate to the global Omni exception handler
|
|
except ValueError as e:
|
|
total_ms = (time.perf_counter() - request_start_s) * 1000.0
|
|
logger.warning(
|
|
"[SpeechE2E] request_id=%s stream=%s status=bad_request total_ms=%.2f error=%s",
|
|
request_id,
|
|
bool(request.stream),
|
|
total_ms,
|
|
e,
|
|
)
|
|
return self.create_error_response(e)
|
|
except TTSGenerationError as e:
|
|
total_ms = (time.perf_counter() - request_start_s) * 1000.0
|
|
logger.error(
|
|
"[SpeechE2E] request_id=%s stream=%s status=invalid_generation total_ms=%.2f error=%s",
|
|
request_id,
|
|
bool(request.stream),
|
|
total_ms,
|
|
e,
|
|
)
|
|
return self._diffusion_error_response(str(e), status_code=500)
|
|
except Exception as e:
|
|
total_ms = (time.perf_counter() - request_start_s) * 1000.0
|
|
logger.exception(
|
|
"[SpeechE2E] request_id=%s stream=%s status=error total_ms=%.2f error=%s",
|
|
request_id,
|
|
bool(request.stream),
|
|
total_ms,
|
|
e,
|
|
)
|
|
return self.create_error_response(
|
|
f"Speech generation failed: {e}",
|
|
err_type="InternalServerError",
|
|
status_code=HTTPStatus.INTERNAL_SERVER_ERROR,
|
|
)
|
|
|
|
@staticmethod
|
|
def _merge_batch_item(
|
|
batch: BatchSpeechRequest,
|
|
item: SpeechBatchItem,
|
|
) -> OpenAICreateSpeechRequest:
|
|
"""Merge batch-level defaults with per-item overrides into a full request."""
|
|
|
|
def _pick(field: str):
|
|
"""Return item-level value if set, else batch-level value."""
|
|
item_val = getattr(item, field, None)
|
|
return item_val if item_val is not None else getattr(batch, field, None)
|
|
|
|
picked_speed = _pick("speed")
|
|
return OpenAICreateSpeechRequest(
|
|
input=item.input,
|
|
model=batch.model,
|
|
voice=_pick("voice"),
|
|
instructions=_pick("instructions"),
|
|
response_format=_pick("response_format") or "wav",
|
|
sample_rate=_pick("sample_rate"),
|
|
speed=picked_speed if picked_speed is not None else 1.0,
|
|
stream=False,
|
|
task_type=_pick("task_type"),
|
|
language=_pick("language"),
|
|
ref_audio=_pick("ref_audio"),
|
|
ref_text=_pick("ref_text"),
|
|
x_vector_only_mode=_pick("x_vector_only_mode"),
|
|
max_new_tokens=_pick("max_new_tokens"),
|
|
initial_codec_chunk_frames=_pick("initial_codec_chunk_frames"),
|
|
non_streaming_mode=_pick("non_streaming_mode"),
|
|
)
|
|
|
|
async def create_speech_batch(
|
|
self,
|
|
batch_request: BatchSpeechRequest,
|
|
) -> BatchSpeechResponse | ErrorResponse:
|
|
"""Generate speech for multiple items concurrently."""
|
|
if self._diffusion_mode:
|
|
raise ValueError("Batch speech is not supported in diffusion mode")
|
|
if len(batch_request.items) > self._batch_max_items:
|
|
raise ValueError(
|
|
f"Batch contains {len(batch_request.items)} items, exceeding the maximum of {self._batch_max_items}."
|
|
)
|
|
|
|
error_check_ret = await self._check_model(batch_request)
|
|
if error_check_ret is not None:
|
|
return error_check_ret
|
|
|
|
if self.engine_client.errored:
|
|
raise self.engine_client.dead_error
|
|
|
|
batch_id = f"speech-batch-{random_uuid()}"
|
|
|
|
merged_requests = [self._merge_batch_item(batch_request, item) for item in batch_request.items]
|
|
|
|
async def _run_item(idx: int, req: OpenAICreateSpeechRequest) -> SpeechBatchItemResult:
|
|
has_inline_ref_audio = req.ref_audio is not None
|
|
validation_error = self._validate_tts_request(req)
|
|
if validation_error is not None:
|
|
return SpeechBatchItemResult(index=idx, status="error", error=validation_error)
|
|
usage_box: list[SpeechTokenUsage] = []
|
|
try:
|
|
audio_data, media_type = await self._generate_audio_bytes(
|
|
req, base64_encode=True, usage_out=usage_box, has_inline_ref_audio=has_inline_ref_audio
|
|
)
|
|
except Exception as e:
|
|
logger.exception("Batch item %d failed: %s", idx, e)
|
|
return SpeechBatchItemResult(index=idx, status="error", error=str(e))
|
|
return SpeechBatchItemResult(
|
|
index=idx,
|
|
status="success",
|
|
audio_data=audio_data,
|
|
media_type=media_type,
|
|
usage=usage_box[0] if usage_box else None,
|
|
)
|
|
|
|
results = await asyncio.gather(
|
|
*[_run_item(i, req) for i, req in enumerate(merged_requests)],
|
|
return_exceptions=True,
|
|
)
|
|
|
|
final_results: list[SpeechBatchItemResult] = []
|
|
for i, r in enumerate(results):
|
|
if isinstance(r, BaseException):
|
|
logger.exception("Batch item %d raised unexpected exception: %s", i, r)
|
|
final_results.append(SpeechBatchItemResult(index=i, status="error", error=str(r)))
|
|
else:
|
|
final_results.append(r)
|
|
|
|
succeeded = sum(1 for r in final_results if r.status == "success")
|
|
return BatchSpeechResponse(
|
|
id=batch_id,
|
|
results=final_results,
|
|
total=len(final_results),
|
|
succeeded=succeeded,
|
|
failed=len(final_results) - succeeded,
|
|
)
|
|
|
|
|
|
ServingSpeech = OmniOpenAIServingSpeech
|