Sy03andHongsheng Liu
efbe4116ff
Signed-off-by: Sy03 <1370724210@qq.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
410 lines
17 KiB
Python
410 lines
17 KiB
Python
import asyncio
|
|
import base64
|
|
import io
|
|
import ipaddress
|
|
import socket
|
|
from typing import Any
|
|
from urllib.parse import urlparse
|
|
from urllib.request import urlopen
|
|
|
|
import numpy as np
|
|
import soundfile as sf
|
|
from fastapi import Request
|
|
from fastapi.responses import Response
|
|
from vllm.entrypoints.openai.engine.serving import OpenAIServing
|
|
from vllm.logger import init_logger
|
|
from vllm.utils import random_uuid
|
|
|
|
from vllm_omni.entrypoints.openai.audio_utils_mixin import AudioMixin
|
|
from vllm_omni.entrypoints.openai.protocol.audio import (
|
|
AudioResponse,
|
|
CreateAudio,
|
|
OpenAICreateSpeechRequest,
|
|
)
|
|
from vllm_omni.outputs import OmniRequestOutput
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
_REF_AUDIO_TIMEOUT_S = 15
|
|
_REF_AUDIO_MAX_BYTES = 50 * 1024 * 1024 # 50 MB
|
|
_REF_AUDIO_BLOCKED_NETWORKS = [
|
|
ipaddress.ip_network("127.0.0.0/8"),
|
|
ipaddress.ip_network("10.0.0.0/8"),
|
|
ipaddress.ip_network("172.16.0.0/12"),
|
|
ipaddress.ip_network("192.168.0.0/16"),
|
|
ipaddress.ip_network("169.254.0.0/16"),
|
|
ipaddress.ip_network("::1/128"),
|
|
ipaddress.ip_network("fc00::/7"),
|
|
ipaddress.ip_network("fe80::/10"),
|
|
]
|
|
|
|
# TTS Configuration (currently supports Qwen3-TTS)
|
|
_TTS_MODEL_STAGES: set[str] = {"qwen3_tts"}
|
|
_TTS_LANGUAGES: set[str] = {
|
|
"Auto",
|
|
"Chinese",
|
|
"English",
|
|
"Japanese",
|
|
"Korean",
|
|
"German",
|
|
"French",
|
|
"Russian",
|
|
"Portuguese",
|
|
"Spanish",
|
|
"Italian",
|
|
}
|
|
_TTS_MAX_INSTRUCTIONS_LENGTH = 500
|
|
_TTS_MAX_NEW_TOKENS_MIN = 1
|
|
_TTS_MAX_NEW_TOKENS_MAX = 4096
|
|
|
|
|
|
class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
# Load supported speakers
|
|
self.supported_speakers = self._load_supported_speakers()
|
|
logger.info(f"Loaded {len(self.supported_speakers)} supported speakers: {sorted(self.supported_speakers)}")
|
|
self._tts_tokenizer = None
|
|
|
|
def _load_supported_speakers(self) -> set[str]:
|
|
"""Load supported speakers (case-insensitive) from the model configuration."""
|
|
try:
|
|
talker_config = self.engine_client.model_config.hf_config.talker_config
|
|
|
|
# Check for speakers in either spk_id or speaker_id
|
|
for attr_name in ["spk_id", "speaker_id"]:
|
|
speakers_dict = getattr(talker_config, attr_name, None)
|
|
if speakers_dict and isinstance(speakers_dict, dict):
|
|
# Normalize to lowercase for case-insensitive matching
|
|
return {speaker.lower() for speaker in speakers_dict.keys()}
|
|
|
|
logger.warning("No speakers found in talker_config (checked spk_id and speaker_id)")
|
|
except Exception as e:
|
|
logger.warning(f"Could not load speakers from model config: {e}")
|
|
|
|
return set()
|
|
|
|
def _estimate_prompt_len(self, tts_params: dict[str, Any]) -> int:
|
|
"""Estimate prompt length so the placeholder matches model-side embeddings."""
|
|
try:
|
|
from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import (
|
|
Qwen3TTSTalkerForConditionalGeneration,
|
|
)
|
|
|
|
if self._tts_tokenizer is None:
|
|
from transformers import AutoTokenizer
|
|
|
|
model_name = self.engine_client.model_config.model
|
|
self._tts_tokenizer = AutoTokenizer.from_pretrained(
|
|
model_name,
|
|
trust_remote_code=True,
|
|
padding_side="left",
|
|
)
|
|
hf_config = self.engine_client.model_config.hf_config
|
|
talker_config = hf_config.talker_config
|
|
task_type = (tts_params.get("task_type") or ["CustomVoice"])[0]
|
|
return Qwen3TTSTalkerForConditionalGeneration.estimate_prompt_len_from_additional_information(
|
|
additional_information=tts_params,
|
|
task_type=task_type,
|
|
tokenize_prompt=lambda t: self._tts_tokenizer(t, padding=False)["input_ids"],
|
|
codec_language_id=getattr(talker_config, "codec_language_id", None),
|
|
spk_is_dialect=getattr(talker_config, "spk_is_dialect", None),
|
|
)
|
|
except Exception as e:
|
|
logger.warning("Failed to estimate TTS prompt length, using fallback 2048: %s", e)
|
|
return 2048
|
|
|
|
def _is_tts_model(self) -> bool:
|
|
"""Check if the current model is a supported TTS model."""
|
|
stage_list = getattr(self.engine_client, "stage_list", None)
|
|
if stage_list:
|
|
for stage in stage_list:
|
|
model_stage = getattr(stage, "model_stage", None)
|
|
if model_stage in _TTS_MODEL_STAGES:
|
|
return True
|
|
return False
|
|
|
|
def _validate_tts_request(self, request: OpenAICreateSpeechRequest) -> str | None:
|
|
"""Validate TTS request parameters. Returns error message or None."""
|
|
task_type = request.task_type or "CustomVoice"
|
|
|
|
# Normalize voice to lowercase for case-insensitive matching
|
|
if request.voice is not None:
|
|
request.voice = request.voice.lower()
|
|
|
|
# Validate input is not empty
|
|
if not request.input or not request.input.strip():
|
|
return "Input text cannot be empty"
|
|
|
|
# Validate language
|
|
if request.language is not None and request.language not in _TTS_LANGUAGES:
|
|
return f"Invalid language '{request.language}'. Supported: {', '.join(sorted(_TTS_LANGUAGES))}"
|
|
|
|
# Validate speaker for CustomVoice task
|
|
if task_type == "CustomVoice" and request.voice is not None:
|
|
if self.supported_speakers and request.voice not in self.supported_speakers:
|
|
return f"Invalid speaker '{request.voice}'. Supported: {', '.join(sorted(self.supported_speakers))}"
|
|
|
|
# Validate Base task requirements
|
|
if task_type == "Base":
|
|
if request.ref_audio is None:
|
|
return "Base task requires 'ref_audio' for voice cloning"
|
|
# Validate ref_audio format
|
|
if not (request.ref_audio.startswith(("http://", "https://")) or request.ref_audio.startswith("data:")):
|
|
return "ref_audio must be a URL (http/https) or base64 data URL (data:...)"
|
|
|
|
# Validate cross-parameter dependencies
|
|
if task_type != "Base":
|
|
if request.ref_text is not None:
|
|
return "'ref_text' is only valid for Base task"
|
|
if request.x_vector_only_mode is not None:
|
|
return "'x_vector_only_mode' is only valid for Base task"
|
|
|
|
# Validate VoiceDesign task requirements
|
|
if task_type == "VoiceDesign" and not request.instructions:
|
|
return "VoiceDesign task requires 'instructions' to describe the voice"
|
|
|
|
# Validate instructions length
|
|
if request.instructions and len(request.instructions) > _TTS_MAX_INSTRUCTIONS_LENGTH:
|
|
return f"Instructions too long (max {_TTS_MAX_INSTRUCTIONS_LENGTH} characters)"
|
|
|
|
# Validate max_new_tokens range
|
|
if request.max_new_tokens is not None:
|
|
if request.max_new_tokens < _TTS_MAX_NEW_TOKENS_MIN:
|
|
return f"max_new_tokens must be at least {_TTS_MAX_NEW_TOKENS_MIN}"
|
|
if request.max_new_tokens > _TTS_MAX_NEW_TOKENS_MAX:
|
|
return f"max_new_tokens cannot exceed {_TTS_MAX_NEW_TOKENS_MAX}"
|
|
|
|
return None
|
|
|
|
@staticmethod
|
|
async def _resolve_ref_audio(ref_audio_str: str) -> tuple[list[float], int]:
|
|
"""Resolve ref_audio URL/base64 to (wav_samples, sample_rate)."""
|
|
parsed = urlparse(ref_audio_str)
|
|
|
|
def _check_ssrf(url: str) -> None:
|
|
host = urlparse(url).hostname
|
|
if not host:
|
|
raise ValueError("ref_audio URL must include a hostname")
|
|
for info in socket.getaddrinfo(host, None):
|
|
ip_str = str(info[4][0]).split("%", 1)[0]
|
|
addr = ipaddress.ip_address(ip_str)
|
|
if any(addr in net for net in _REF_AUDIO_BLOCKED_NETWORKS):
|
|
raise ValueError(f"ref_audio URL resolves to blocked address: {addr}")
|
|
|
|
def _fetch_sync() -> tuple[np.ndarray, int]:
|
|
if parsed.scheme in ("http", "https"):
|
|
_check_ssrf(ref_audio_str)
|
|
with urlopen(ref_audio_str, timeout=_REF_AUDIO_TIMEOUT_S) as resp:
|
|
data = resp.read(_REF_AUDIO_MAX_BYTES + 1)
|
|
if len(data) > _REF_AUDIO_MAX_BYTES:
|
|
raise ValueError(f"ref_audio URL exceeds {_REF_AUDIO_MAX_BYTES} bytes")
|
|
buf = io.BytesIO(data)
|
|
elif ref_audio_str.startswith("data:"):
|
|
b64 = ref_audio_str
|
|
if "," in b64:
|
|
b64 = b64.split(",", 1)[1]
|
|
buf = io.BytesIO(base64.b64decode(b64))
|
|
else:
|
|
raise ValueError("ref_audio must be an http(s) URL or data: base64 URI")
|
|
audio, sr = sf.read(buf, dtype="float32", always_2d=False)
|
|
if isinstance(audio, np.ndarray) and audio.ndim > 1:
|
|
audio = np.mean(audio, axis=-1)
|
|
return np.asarray(audio, dtype=np.float32), int(sr)
|
|
|
|
loop = asyncio.get_running_loop()
|
|
wav_np, sr = await loop.run_in_executor(None, _fetch_sync)
|
|
return wav_np.tolist(), sr
|
|
|
|
def _build_tts_params(self, request: OpenAICreateSpeechRequest) -> dict[str, Any]:
|
|
"""Build TTS parameters from request.
|
|
|
|
Processes each parameter if present, skips if not.
|
|
Values are wrapped in lists as required by the model.
|
|
"""
|
|
params: dict[str, Any] = {}
|
|
|
|
# Text content (always required)
|
|
params["text"] = [request.input]
|
|
|
|
# Task type
|
|
if request.task_type is not None:
|
|
params["task_type"] = [request.task_type]
|
|
else:
|
|
params["task_type"] = ["CustomVoice"]
|
|
|
|
# Language
|
|
if request.language is not None:
|
|
params["language"] = [request.language]
|
|
else:
|
|
params["language"] = ["Auto"]
|
|
|
|
# Speaker (voice)
|
|
if request.voice is not None:
|
|
params["speaker"] = [request.voice]
|
|
elif params["task_type"][0] == "CustomVoice":
|
|
params["speaker"] = ["Vivian"] # Default for CustomVoice
|
|
|
|
# Instructions for style/emotion control
|
|
if request.instructions is not None:
|
|
params["instruct"] = [request.instructions]
|
|
else:
|
|
params["instruct"] = [""]
|
|
|
|
# Voice clone: ref_audio resolved in create_speech(), not here.
|
|
if request.ref_text is not None:
|
|
params["ref_text"] = [request.ref_text]
|
|
if request.x_vector_only_mode is not None:
|
|
params["x_vector_only_mode"] = [request.x_vector_only_mode]
|
|
|
|
# Generation parameters
|
|
if request.max_new_tokens is not None:
|
|
params["max_new_tokens"] = [request.max_new_tokens]
|
|
else:
|
|
params["max_new_tokens"] = [2048]
|
|
|
|
return params
|
|
|
|
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)
|
|
|
|
NOTE: Streaming audio generation is not currently supported.
|
|
"""
|
|
|
|
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
|
|
|
|
if self.engine_client.errored:
|
|
raise self.engine_client.dead_error
|
|
|
|
request_id = f"speech-{random_uuid()}"
|
|
|
|
try:
|
|
if self._is_tts_model():
|
|
# Validate TTS parameters
|
|
validation_error = self._validate_tts_request(request)
|
|
if validation_error:
|
|
return self.create_error_response(validation_error)
|
|
|
|
# Must use prompt_token_ids (not text prompt): the AR Talker
|
|
# operates on codec tokens; text token IDs exceed codec vocab.
|
|
# model.preprocess replaces all embeddings, so placeholder value
|
|
# is irrelevant -- but length must match to avoid excess padding.
|
|
tts_params = self._build_tts_params(request)
|
|
|
|
if request.ref_audio is not None:
|
|
wav_list, sr = await self._resolve_ref_audio(request.ref_audio)
|
|
tts_params["ref_audio"] = [[wav_list, sr]]
|
|
|
|
ph_len = self._estimate_prompt_len(tts_params)
|
|
prompt = {
|
|
"prompt_token_ids": [1] * ph_len,
|
|
"additional_information": tts_params,
|
|
}
|
|
else:
|
|
# Fallback for unsupported models
|
|
tts_params = {}
|
|
prompt = {"prompt": request.input}
|
|
|
|
logger.info(
|
|
"TTS speech request %s: text=%r, task_type=%s",
|
|
request_id,
|
|
request.input[:50] + "..." if len(request.input) > 50 else request.input,
|
|
tts_params.get("task_type", ["unknown"])[0],
|
|
)
|
|
|
|
sampling_params_list = self.engine_client.default_sampling_params_list
|
|
|
|
generator = self.engine_client.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:
|
|
return self.create_error_response("No output generated from the model.")
|
|
|
|
# Extract audio from output
|
|
# Audio can be in final_output.multimodal_output or final_output.request_output.multimodal_output
|
|
# Support both "audio" and "model_outputs" keys for compatibility with different models
|
|
audio_output = None
|
|
if hasattr(final_output, "multimodal_output") and final_output.multimodal_output:
|
|
audio_output = final_output.multimodal_output
|
|
if not audio_output and hasattr(final_output, "request_output"):
|
|
if final_output.request_output and hasattr(final_output.request_output, "multimodal_output"):
|
|
audio_output = final_output.request_output.multimodal_output
|
|
|
|
# Check for audio data using either "audio" or "model_outputs" key
|
|
audio_key = None
|
|
if audio_output:
|
|
if "audio" in audio_output:
|
|
audio_key = "audio"
|
|
elif "model_outputs" in audio_output:
|
|
audio_key = "model_outputs"
|
|
|
|
if not audio_output or audio_key is None:
|
|
return self.create_error_response("TTS model did not produce audio output.")
|
|
|
|
audio_tensor = audio_output[audio_key]
|
|
sample_rate = audio_output.get("sr", 24000)
|
|
if hasattr(sample_rate, "item"):
|
|
sample_rate = sample_rate.item()
|
|
|
|
# Streaming accumulates chunks as a list; concat first.
|
|
if isinstance(audio_tensor, list):
|
|
import torch
|
|
|
|
audio_tensor = torch.cat(audio_tensor, dim=-1)
|
|
# Convert tensor to numpy
|
|
if hasattr(audio_tensor, "float"):
|
|
audio_tensor = audio_tensor.float().detach().cpu().numpy()
|
|
|
|
# Squeeze batch dimension if present, but preserve channel dimension for stereo
|
|
if audio_tensor.ndim > 1:
|
|
audio_tensor = audio_tensor.squeeze()
|
|
|
|
audio_obj = CreateAudio(
|
|
audio_tensor=audio_tensor,
|
|
sample_rate=int(sample_rate),
|
|
response_format=request.response_format or "wav",
|
|
speed=request.speed or 1.0,
|
|
stream_format=request.stream_format,
|
|
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.create_error_response("Client disconnected")
|
|
except ValueError as e:
|
|
return self.create_error_response(e)
|
|
except Exception as e:
|
|
logger.exception("Speech generation failed: %s", e)
|
|
return self.create_error_response(f"Speech generation failed: {e}")
|