47 lines
1.7 KiB
Python
47 lines
1.7 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.tokenizer import IndexTTS2Tokenizer
|
||
|
|
|
||
|
|
_CONDITIONING_PREFIX_TOKENS = 34
|
||
|
|
_TEXT_WRAPPER_TOKENS = 2
|
||
|
|
_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) -> int:
|
||
|
|
"""Return the placeholder prompt length expected by the IndexTTS2 talker.
|
||
|
|
|
||
|
|
The model-side prefill replaces placeholder embeddings with:
|
||
|
|
34 conditioning tokens, text tokens wrapped with start/stop markers, and
|
||
|
|
one start-mel token.
|
||
|
|
"""
|
||
|
|
tokenizer = _get_text_tokenizer(model_id_or_path)
|
||
|
|
text_token_count = len(tokenizer.encode(text, add_special_tokens=False)) + _TEXT_WRAPPER_TOKENS
|
||
|
|
return _CONDITIONING_PREFIX_TOKENS + text_token_count + _START_MEL_TOKENS
|
||
|
|
|
||
|
|
|
||
|
|
def build_indextts2_prefill_prompt_ids(
|
||
|
|
model_id_or_path: str,
|
||
|
|
text: str,
|
||
|
|
*,
|
||
|
|
placeholder_token_id: int = 1,
|
||
|
|
) -> list[int]:
|
||
|
|
return [placeholder_token_id] * estimate_indextts2_prefill_prompt_len(model_id_or_path, text)
|