335 lines
13 KiB
Python
335 lines
13 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""IndexTTS2 serving adapter."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import math
|
|
from collections.abc import Mapping
|
|
from typing import TYPE_CHECKING, Any, ClassVar
|
|
|
|
import numpy as np
|
|
import torch
|
|
from vllm.inputs import tokens_input
|
|
from vllm.utils import random_uuid
|
|
|
|
from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter
|
|
from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest, apply_max_new_tokens
|
|
from vllm_omni.model_executor.models.indextts2.configuration_indextts2 import (
|
|
INDEXTTS25_MAX_DURATION_FACTOR,
|
|
INDEXTTS25_MIN_DURATION_FACTOR,
|
|
)
|
|
|
|
if TYPE_CHECKING:
|
|
from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest
|
|
|
|
_TTS_MAX_NEW_TOKENS_MIN = 1
|
|
_TTS_MAX_NEW_TOKENS_MAX = 4096
|
|
_INDEXTTS2_EMOTION_KEYS = ("emo_audio", "emo_vector", "emo_alpha", "emo_text", "use_emo_text", "use_random")
|
|
|
|
|
|
def _update_conditioning_hash(h: hashlib._Hash, value: Any) -> None:
|
|
"""Hash request conditioning values without serializing huge repr strings."""
|
|
if value is None:
|
|
h.update(b"none")
|
|
return
|
|
if isinstance(value, torch.Tensor):
|
|
arr = value.detach().cpu().numpy()
|
|
h.update(b"tensor")
|
|
h.update(repr((arr.shape, str(arr.dtype))).encode("utf-8"))
|
|
h.update(np.ascontiguousarray(arr).tobytes())
|
|
return
|
|
if isinstance(value, np.ndarray):
|
|
arr = np.ascontiguousarray(value)
|
|
h.update(b"ndarray")
|
|
h.update(repr((arr.shape, str(arr.dtype))).encode("utf-8"))
|
|
h.update(arr.tobytes())
|
|
return
|
|
if isinstance(value, (list, tuple)) and len(value) == 2 and isinstance(value[1], (int, np.integer)):
|
|
samples = value[0]
|
|
if isinstance(samples, (list, tuple, np.ndarray, torch.Tensor)):
|
|
if isinstance(samples, torch.Tensor):
|
|
samples_arr = samples.detach().cpu().numpy()
|
|
else:
|
|
samples_arr = np.asarray(samples, dtype=np.float32)
|
|
samples_arr = np.ascontiguousarray(samples_arr, dtype=np.float32)
|
|
h.update(b"audio")
|
|
h.update(int(value[1]).to_bytes(8, byteorder="little", signed=True))
|
|
h.update(int(samples_arr.size).to_bytes(8, byteorder="little", signed=False))
|
|
h.update(samples_arr.tobytes())
|
|
return
|
|
if isinstance(value, Mapping):
|
|
h.update(b"mapping")
|
|
for key in sorted(value):
|
|
h.update(b"\x01")
|
|
h.update(repr(key).encode("utf-8"))
|
|
h.update(b"\x02")
|
|
_update_conditioning_hash(h, value[key])
|
|
return
|
|
if isinstance(value, (list, tuple)):
|
|
h.update(f"sequence:{len(value)}".encode())
|
|
for item in value:
|
|
h.update(b"\x03")
|
|
_update_conditioning_hash(h, item)
|
|
return
|
|
h.update(repr(value).encode("utf-8"))
|
|
|
|
|
|
def _first_conditioning_value(value: Any) -> Any:
|
|
if isinstance(value, list) and len(value) == 1:
|
|
return value[0]
|
|
return value
|
|
|
|
|
|
def indextts2_conditioning_cache_salt(
|
|
request: OpenAICreateSpeechRequest,
|
|
tts_params: dict[str, Any],
|
|
*,
|
|
request_id: str | None = None,
|
|
) -> str:
|
|
"""Stable prefix-cache salt for IndexTTS2 placeholder prompts.
|
|
|
|
IndexTTS2 builds prefill embeddings from ``additional_information`` (speaker
|
|
audio plus emotion controls), while the scheduler only sees dummy token ids.
|
|
Include all supported control inputs so prefix-cache hits are only possible
|
|
when text, speaker, and emotion conditioning are identical.
|
|
"""
|
|
h = hashlib.sha256()
|
|
for key in (
|
|
"text",
|
|
"voice",
|
|
"voice_name",
|
|
"voice_created_at",
|
|
"emo_audio",
|
|
"emo_vector",
|
|
"emo_alpha",
|
|
"emo_text",
|
|
"use_emo_text",
|
|
"use_random",
|
|
"lang",
|
|
"text_normalization",
|
|
):
|
|
h.update(b"\x00")
|
|
h.update(key.encode("utf-8"))
|
|
h.update(b"\x00")
|
|
_update_conditioning_hash(h, tts_params.get(key))
|
|
|
|
for key, value in (
|
|
("input", request.input),
|
|
("ref_audio", request.ref_audio),
|
|
):
|
|
h.update(b"\x00")
|
|
h.update(key.encode("utf-8"))
|
|
h.update(b"\x00")
|
|
_update_conditioning_hash(h, value)
|
|
|
|
use_random = bool(_first_conditioning_value(tts_params.get("use_random")))
|
|
if use_random:
|
|
h.update(b"\x00indextts2-use-random-request\x00")
|
|
h.update((request_id or random_uuid()).encode("utf-8"))
|
|
return h.hexdigest()[:32]
|
|
|
|
|
|
@register_tts_adapter
|
|
class IndexTTS2Adapter(ARTTSAdapter):
|
|
stage_keys = frozenset({"indextts2_talker"})
|
|
name = "indextts2"
|
|
|
|
def validate(self, request: OpenAICreateSpeechRequest) -> str | None:
|
|
server = self.ctx.server
|
|
if not request.input or not request.input.strip():
|
|
return "Input text cannot be empty"
|
|
|
|
voice_lower = request.voice.lower() if isinstance(request.voice, str) else None
|
|
has_uploaded_voice = bool(
|
|
voice_lower
|
|
and voice_lower in server.uploaded_speakers
|
|
and server._get_uploaded_audio_data(voice_lower) is not None
|
|
)
|
|
if request.ref_audio is None and not has_uploaded_voice:
|
|
return "IndexTTS2 requires 'ref_audio' or an uploaded audio voice for voice cloning"
|
|
if request.ref_audio is not None:
|
|
fmt_err = server._validate_ref_audio_format(request.ref_audio)
|
|
if fmt_err:
|
|
return fmt_err
|
|
if request.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}"
|
|
if request.extra_params is not None:
|
|
if not isinstance(request.extra_params, dict):
|
|
return "extra_params must be a JSON object/dict."
|
|
extra_err = self._validate_extra_params(request.extra_params)
|
|
if extra_err:
|
|
return extra_err
|
|
return None
|
|
|
|
def _validate_extra_params(self, extras: Mapping[str, Any]) -> str | None:
|
|
server = self.ctx.server
|
|
if "emo_audio" in extras:
|
|
value = extras["emo_audio"]
|
|
if not isinstance(value, str):
|
|
return "extra_params.emo_audio must be a URL (http/https), base64 data URL (data:...), or file URI (file://...)"
|
|
fmt_err = server._validate_ref_audio_format(value)
|
|
if fmt_err:
|
|
return fmt_err.replace("ref_audio", "extra_params.emo_audio", 1)
|
|
|
|
if "emo_vector" in extras:
|
|
value = extras["emo_vector"]
|
|
if not isinstance(value, list) or len(value) != 8:
|
|
return "extra_params.emo_vector must be a list of 8 numbers"
|
|
for item in value:
|
|
if isinstance(item, bool) or not isinstance(item, (int, float)) or not math.isfinite(float(item)):
|
|
return "extra_params.emo_vector must be a list of 8 finite numbers"
|
|
if float(item) < 0.0 or float(item) > 1.2:
|
|
return "extra_params.emo_vector values must be in [0, 1.2]"
|
|
|
|
if "emo_alpha" in extras:
|
|
value = extras["emo_alpha"]
|
|
if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(float(value)):
|
|
return "extra_params.emo_alpha must be a finite number"
|
|
|
|
if "emo_text" in extras and extras["emo_text"] is not None and not isinstance(extras["emo_text"], str):
|
|
return "extra_params.emo_text must be a string"
|
|
|
|
for key in ("use_emo_text", "use_random"):
|
|
if key in extras and not isinstance(extras[key], bool):
|
|
return f"extra_params.{key} must be a boolean"
|
|
return None
|
|
|
|
async def build(
|
|
self,
|
|
request: OpenAICreateSpeechRequest,
|
|
sampling_params_list: list,
|
|
has_inline_ref_audio: bool,
|
|
) -> PreparedRequest:
|
|
del sampling_params_list, has_inline_ref_audio
|
|
|
|
server = self.ctx.server
|
|
tts_params = await self._build_params(request)
|
|
from vllm_omni.model_executor.models.indextts2.prompt_utils import (
|
|
estimate_indextts2_prefill_prompt_len,
|
|
)
|
|
|
|
prompt_kwargs: dict[str, Any] = {}
|
|
if self.name == "indextts2_5":
|
|
from vllm_omni.model_executor.models.indextts2.tokenizer_v2_5 import (
|
|
INDEXTTS25_TOKENIZER_FILE,
|
|
)
|
|
|
|
hf_config = getattr(server.engine_client.model_config, "hf_config", None)
|
|
prompt_kwargs["tokenizer_file"] = getattr(
|
|
hf_config,
|
|
"tokenizer_file",
|
|
INDEXTTS25_TOKENIZER_FILE,
|
|
)
|
|
ph_len = estimate_indextts2_prefill_prompt_len(
|
|
server.engine_client.model_config.model,
|
|
request.input,
|
|
model_type=self.name,
|
|
lang=str(_first_conditioning_value(tts_params.get("lang")) or "zh"),
|
|
text_normalization=bool(
|
|
_first_conditioning_value(tts_params.get("text_normalization"))
|
|
if "text_normalization" in tts_params
|
|
else True
|
|
),
|
|
**prompt_kwargs,
|
|
)
|
|
prompt = tokens_input(prompt_token_ids=[1] * ph_len)
|
|
prompt["additional_information"] = tts_params
|
|
prompt["cache_salt"] = indextts2_conditioning_cache_salt(request, tts_params)
|
|
return PreparedRequest(prompt=prompt, tts_params=tts_params, model_type=self.name)
|
|
|
|
def apply_sampling_overrides(
|
|
self,
|
|
sampling_params_list: list,
|
|
request: OpenAICreateSpeechRequest,
|
|
prompt: dict[str, Any] | None = None,
|
|
request_id: str | None = None,
|
|
) -> list:
|
|
return apply_max_new_tokens(sampling_params_list, request)
|
|
|
|
async def _build_params(self, request: OpenAICreateSpeechRequest) -> dict[str, Any]:
|
|
server = self.ctx.server
|
|
params: dict[str, Any] = {"text": [request.input]}
|
|
voice_lower = request.voice.lower() if isinstance(request.voice, str) else None
|
|
ref_audio_source = request.ref_audio
|
|
using_uploaded_voice = False
|
|
if ref_audio_source is None and voice_lower and voice_lower in server.uploaded_speakers:
|
|
ref_audio_source = server._get_uploaded_audio_data(voice_lower)
|
|
using_uploaded_voice = ref_audio_source is not None
|
|
if ref_audio_source is not None and isinstance(ref_audio_source, str):
|
|
wav_list, sr = await server._resolve_ref_audio(ref_audio_source)
|
|
params["voice"] = [[wav_list, sr]]
|
|
if using_uploaded_voice and voice_lower:
|
|
params["voice_name"] = [voice_lower]
|
|
params["voice_created_at"] = [server._voice_created_at(voice_lower)]
|
|
|
|
extras = request.extra_params if isinstance(request.extra_params, dict) else {}
|
|
for key in _INDEXTTS2_EMOTION_KEYS:
|
|
if key not in extras:
|
|
continue
|
|
if key == "emo_audio":
|
|
wav_list, sr = await server._resolve_ref_audio(extras[key])
|
|
params[key] = [[wav_list, sr]]
|
|
else:
|
|
params[key] = [extras[key]]
|
|
return params
|
|
|
|
|
|
@register_tts_adapter
|
|
class IndexTTS25Adapter(IndexTTS2Adapter):
|
|
stage_keys = frozenset({"indextts2_5_talker"})
|
|
name = "indextts2_5"
|
|
native_speed_control: ClassVar[bool] = True
|
|
|
|
def validate(self, request: OpenAICreateSpeechRequest) -> str | None:
|
|
error = super().validate(request)
|
|
if error is not None:
|
|
return error
|
|
speed = request.speed if request.speed is not None else 1.0
|
|
if speed < INDEXTTS25_MIN_DURATION_FACTOR or speed > INDEXTTS25_MAX_DURATION_FACTOR:
|
|
return "IndexTTS 2.5 speed must be between 0.5 and 2.0"
|
|
return None
|
|
|
|
def _validate_extra_params(self, extras: Mapping[str, Any]) -> str | None:
|
|
error = super()._validate_extra_params(extras)
|
|
if error is not None:
|
|
return error
|
|
|
|
if "lang" in extras:
|
|
value = extras["lang"]
|
|
if not isinstance(value, str):
|
|
return "extra_params.lang must be a string"
|
|
from vllm_omni.model_executor.models.indextts2.tokenizer_v2_5 import (
|
|
normalize_language_code,
|
|
)
|
|
|
|
try:
|
|
normalize_language_code(value)
|
|
except ValueError as exc:
|
|
return str(exc)
|
|
if "text_normalization" in extras and not isinstance(
|
|
extras["text_normalization"],
|
|
bool,
|
|
):
|
|
return "extra_params.text_normalization must be a boolean"
|
|
return None
|
|
|
|
async def _build_params(
|
|
self,
|
|
request: OpenAICreateSpeechRequest,
|
|
) -> dict[str, Any]:
|
|
params = await super()._build_params(request)
|
|
extras = request.extra_params if isinstance(request.extra_params, dict) else {}
|
|
from vllm_omni.model_executor.models.indextts2.tokenizer_v2_5 import (
|
|
normalize_language_code,
|
|
)
|
|
|
|
params["lang"] = [normalize_language_code(str(extras.get("lang", "zh")))]
|
|
params["text_normalization"] = [bool(extras.get("text_normalization", True))]
|
|
speed = request.speed if request.speed is not None else 1.0
|
|
params["duration_factor"] = [1.0 / speed]
|
|
return params
|