Sy03andYueqian Lin
58007c0084
Signed-off-by: Sysy <1370724210@qq.com> Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>
90 lines
3.0 KiB
Python
90 lines
3.0 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Prompt helpers for IndexTTS2 talker prefill."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from functools import lru_cache
|
|
|
|
from vllm_omni.model_executor.models.indextts2.preprocess_utils import resolve_model_file
|
|
from vllm_omni.model_executor.models.indextts2.text_processing_v2_5 import (
|
|
prepare_indextts25_text,
|
|
)
|
|
from vllm_omni.model_executor.models.indextts2.tokenizer import IndexTTS2Tokenizer
|
|
from vllm_omni.model_executor.models.indextts2.tokenizer_v2_5 import (
|
|
INDEXTTS25_TOKENIZER_FILE,
|
|
)
|
|
|
|
_V2_CONDITIONING_PREFIX_TOKENS = 34
|
|
_V25_CONDITIONING_PREFIX_TOKENS = 3
|
|
_V2_TEXT_WRAPPER_TOKENS = 2
|
|
_V25_TEXT_WRAPPER_TOKENS = 2
|
|
_V25_TEXT_WRAPPER_TOKEN_IDS = frozenset({0, 1})
|
|
_START_MEL_TOKENS = 1
|
|
|
|
|
|
def _resolve_bpe_model_path(model_id_or_path: str) -> str:
|
|
path = resolve_model_file(model_id_or_path, "bpe.model")
|
|
if path is None:
|
|
raise FileNotFoundError(f"Could not resolve bpe.model for {model_id_or_path!r}")
|
|
return path
|
|
|
|
|
|
@lru_cache(maxsize=16)
|
|
def _get_text_tokenizer(model_id_or_path: str) -> IndexTTS2Tokenizer:
|
|
return IndexTTS2Tokenizer(_resolve_bpe_model_path(model_id_or_path), model_dir=model_id_or_path)
|
|
|
|
|
|
def estimate_indextts2_prefill_prompt_len(
|
|
model_id_or_path: str,
|
|
text: str,
|
|
*,
|
|
model_type: str = "indextts2",
|
|
lang: str = "zh",
|
|
text_normalization: bool = True,
|
|
tokenizer_file: str = INDEXTTS25_TOKENIZER_FILE,
|
|
) -> int:
|
|
"""Return the placeholder prompt length expected by the IndexTTS2 talker.
|
|
|
|
IndexTTS 2 uses 34 conditioning tokens. IndexTTS 2.5 uses one projected
|
|
CAMPPlus speaker token followed by two zero tokens.
|
|
"""
|
|
if model_type == "indextts2_5":
|
|
text_ids, _ = prepare_indextts25_text(
|
|
text,
|
|
lang=lang,
|
|
model_dir=model_id_or_path,
|
|
text_normalization=text_normalization,
|
|
tokenizer_file=tokenizer_file,
|
|
)
|
|
conditioning_prefix_tokens = _V25_CONDITIONING_PREFIX_TOKENS
|
|
text_token_count = (
|
|
sum(token_id not in _V25_TEXT_WRAPPER_TOKEN_IDS for token_id in text_ids) + _V25_TEXT_WRAPPER_TOKENS
|
|
)
|
|
else:
|
|
tokenizer = _get_text_tokenizer(model_id_or_path)
|
|
text_token_count = len(tokenizer.encode(text, add_special_tokens=False)) + _V2_TEXT_WRAPPER_TOKENS
|
|
conditioning_prefix_tokens = _V2_CONDITIONING_PREFIX_TOKENS
|
|
return conditioning_prefix_tokens + text_token_count + _START_MEL_TOKENS
|
|
|
|
|
|
def build_indextts2_prefill_prompt_ids(
|
|
model_id_or_path: str,
|
|
text: str,
|
|
*,
|
|
model_type: str = "indextts2",
|
|
lang: str = "zh",
|
|
text_normalization: bool = True,
|
|
tokenizer_file: str = INDEXTTS25_TOKENIZER_FILE,
|
|
placeholder_token_id: int = 1,
|
|
) -> list[int]:
|
|
prompt_len = estimate_indextts2_prefill_prompt_len(
|
|
model_id_or_path,
|
|
text,
|
|
model_type=model_type,
|
|
lang=lang,
|
|
text_normalization=text_normalization,
|
|
tokenizer_file=tokenizer_file,
|
|
)
|
|
return [placeholder_token_id] * prompt_len
|