From ebe93f5d1e794202598e652a45fa74a29a7f698b Mon Sep 17 00:00:00 2001 From: Yueqian Lin <70319226+linyueqian@users.noreply.github.com> Date: Sun, 9 Aug 2026 06:30:05 -0700 Subject: [PATCH] [Refactor][TTS] Derive TTS model detection from adapter metadata (#5682) Signed-off-by: Yueqian Lin --- .claude/skills/add-tts-model/SKILL.md | 115 +++--- .pre-commit-config.yaml | 7 + docs/contributing/model/adding_tts_model.md | 173 +++++---- .../openai_api/test_audex_serving_guards.py | 4 + .../openai_api/test_serving_speech.py | 13 +- .../openai_api/test_tts_detection.py | 345 ++++++++++++++++++ tests/tools/test_check_tts_adapter.py | 79 ++++ tools/pre_commit/check_tts_adapter.py | 181 +++++++++ .../entrypoints/openai/serving_speech.py | 139 ++----- .../openai/tts_adapters/__init__.py | 88 ++++- .../entrypoints/openai/tts_adapters/audex.py | 9 + .../entrypoints/openai/tts_adapters/base.py | 66 +++- .../openai/tts_adapters/covo_audio.py | 7 + .../openai/tts_adapters/ming_tts.py | 10 +- .../openai/tts_adapters/voxcpm2.py | 5 + 15 files changed, 1002 insertions(+), 239 deletions(-) create mode 100644 tests/entrypoints/openai_api/test_tts_detection.py create mode 100644 tests/tools/test_check_tts_adapter.py create mode 100644 tools/pre_commit/check_tts_adapter.py diff --git a/.claude/skills/add-tts-model/SKILL.md b/.claude/skills/add-tts-model/SKILL.md index dc423825a..ff86ad457 100644 --- a/.claude/skills/add-tts-model/SKILL.md +++ b/.claude/skills/add-tts-model/SKILL.md @@ -220,65 +220,70 @@ See `plan/voxcpm2_native_ar_design.md`. ### Steps -1. **Register in `serving_speech.py`** — add all 5 points in a **single commit**; - partial integration causes hard-to-debug failures. This file is modified by every - model PR and is the most common source of rebase conflicts — see conflict note below. +1. **Write one adapter** under `vllm_omni/entrypoints/openai/tts_adapters/`. + `serving_speech.py` should not need an edit — detection, stage discovery and + dispatch are all derived from what the adapter declares. - **Point 1** — stage constant (near the top, alongside the other `_*_TTS_MODEL_STAGES` sets): + Create `vllm_omni/entrypoints/openai/tts_adapters/your_model.py`: ```python - _YOUR_MODEL_TTS_MODEL_STAGES = {"your_stage_key"} + @register_tts_adapter + class YourModelAdapter(ARTTSAdapter): + name = "your_model" # registry key + log label + stage_keys = frozenset({"your_stage_key"}) # the deploy yaml's model_stage + + def validate(self, request) -> str | None: + if not request.input or not request.input.strip(): + return "Input text cannot be empty" + return None + + async def build(self, request, sampling_params_list, has_inline_ref_audio): + params = {"text": [request.input]} + if request.voice is not None: + params["voice"] = [request.voice] + return PreparedRequest( + prompt={"prompt": request.input}, + tts_params=params, + model_type=self.name, + ) ``` - **Point 2** — union into `_TTS_MODEL_STAGES`: - ```python - _TTS_MODEL_STAGES: set[str] = ( - ... - | _YOUR_MODEL_TTS_MODEL_STAGES - ) - ``` + Then add the module to the import block at the bottom of + `tts_adapters/__init__.py` so it registers. That import line is the only shared + file a new model touches — which is also why the old rebase-conflict hotspot is + gone. - **Point 3** — model type detection in `_detect_tts_model_type()`: - ```python - if model_stage in _YOUR_MODEL_TTS_MODEL_STAGES: - return "your_model" - ``` + > **Pure-diffusion TTS does not go through adapters yet.** Under + > `for_diffusion()`, `create_speech()` routes straight to + > `_create_diffusion_speech()` and never calls `validate()`/`build()`. + > `DiffusionTTSAdapter` is scaffolding with no production subclass, so logic + > placed in one would silently never run. Diffusion-engine models follow the + > existing diffusion path; wiring it through adapters is open work (#4855). - **Point 4** — validation dispatch in `_validate_tts_request()`: - ```python - if self._tts_model_type == "your_model": - return self._validate_your_model_request(request) - ``` + **If a stage key alone cannot identify the model**, declare + `model_archs = frozenset({"YourModelForConditionalGeneration"})`; add + `arch_identifies_entry_stage = True` when the model owns no stage key at all + (Ming dense). For a rule that is not set membership, override `matches()` (see + `covo_audio.py`). For a genuine overlap with another adapter, give one an + explicit `detect_priority` — `test_tts_detection.py` fails on an unordered + overlap. Models that only serve speech in some topologies override + `stage_serves_speech()` (see `audex.py`). - **Point 5** — validation + parameter-builder methods: - ```python - def _validate_your_model_request(self, request) -> str | None: - if not request.input or not request.input.strip(): - return "Input text cannot be empty" - return None + **Reuse shared helpers via `self.ctx.server`** rather than reimplementing them: + `_resolve_ref_audio`, `_apply_uploaded_speaker`, `_validate_ref_audio_format`, + `_max_instructions_length`. Read a comparable adapter first — `fish_speech.py` + (voice cloning), `higgs_audio_v3.py` (parameter-heavy), `moss_tts.py` (family + sharing a base class). - def _build_your_model_params(self, request) -> dict: - params = {"text": [request.input]} - if request.voice is not None: - params["voice"] = [request.voice] - return params - ``` - Wire `_build_your_model_params` into `_create_tts_request()` alongside the other - model-specific param builders. + > **Do not add `self._tts_model_type == ...` branches to `serving_speech.py`.** + > Older models predate the adapter framework and still have them; they are being + > migrated out (RFC #4327, #4855). `tools/pre_commit/check_tts_adapter.py` is a + > ratchet on the remaining count and fails the commit if it grows. Behaviour that + > no adapter hook can express is a missing hook — propose it on the RFC. - > **Two dispatch patterns coexist**: Fish Speech uses a `self._is_fish_speech` boolean - > instance attribute checked before `elif self._is_tts`, while all newer models - > (CosyVoice3, MOSS-TTS-Nano) use the `_tts_model_type` string returned by - > `_detect_tts_model_type()`. For new models, always use the `_tts_model_type` string - > pattern — do not add new `_is_*` flags. - - > **Unused variable rule**: only extract fields in `_build_your_model_params` that - > are actually forwarded to the model. Unused extractions fail `ruff F841`. - > For voice-cloning fields (`ref_audio` → `prompt_audio_path`, `ref_text` → - > `prompt_text`), add them to the param builder and verify they reach the model call. - - **Rebase conflict note**: when rebasing onto `main` after another model was merged, - `serving_speech.py` will conflict. Resolution: always keep *both* the upstream - model's additions and your own — never discard either side. + > **Unused variable rule**: only extract fields in `build()` that are actually + > forwarded to the model. Unused extractions fail `ruff F841`. For voice-cloning + > fields (`ref_audio` -> `prompt_audio_path`, `ref_text` -> `prompt_text`), add + > them to the params and verify they reach the model call. 2. **Handle model-specific parameters**: - Voice cloning: `ref_audio` encoding and prompt injection @@ -295,7 +300,7 @@ import base64 from pathlib import Path def build_voice_clone_prompt(ref_audio_path: str, text: str, codec) -> list: - """Build prompt with reference audio for voice cloning in serving_speech.py.""" + """Build prompt with reference audio for voice cloning, called from the adapter.""" audio_bytes = Path(ref_audio_path).read_bytes() codes = codec.encode(audio_bytes) # Encode on CPU using model's codec (e.g., DAC) token_ids = [code + codec.vocab_offset for code in codes.flatten().tolist()] @@ -321,7 +326,7 @@ Classify the model's **CI priority** first (high / medium / low). High-priority | Level | Location | Marker | CI pipeline | Notes | |-------|----------|--------|-------------|-------| -| **L1** | `tests/model_executor/…`, `tests/entrypoints/openai_api/…`, stage-processor tests | `core_model` + `cpu` | `test-ready.yml` | Prompt assembly, async_chunk helpers, `serving_speech` validation — no GPU | +| **L1** | `tests/model_executor/…`, `tests/entrypoints/openai_api/…`, stage-processor tests | `core_model` + `cpu` | `test-ready.yml` | Prompt assembly, async_chunk helpers, adapter validation — no GPU | | **L2** | `tests/e2e/online_serving/test_{slug}.py` | **`core_model` + `advanced_model`** (both on baseline smoke) + `tts` + `@hardware_test(...)` | `test-ready.yml` (`ready` label) | Default deploy smoke: single `/v1/audio/speech` or offline `OmniRunner` path | | **L3** | `tests/e2e/online_serving/test_{slug}.py` **and** `tests/e2e/offline_inference/test_{slug}.py` | Baseline smoke: **`core_model` + `advanced_model`**; heavier cases: `advanced_model` only (+ `tts`) | `test-merge.yml` **or** merged into nightly TTS function job | Streaming, voice clone, batch/queue, async_chunk | | **L4** | `tests/e2e/online_serving/test_{slug}_expansion.py`, optional offline expansion | `full_model` + `tts` | `test-nightly.yml` (`:full_moon: TTS · Function Test with L4`) | Feature matrix; perf → `tests/dfx/perf/tests/test_tts.json` | @@ -366,7 +371,7 @@ See [vllm-omni-test skill](../vllm-omni-test/SKILL.md) § **Runtime send helpers ### Deliverables -- Updated `serving_speech.py` with all 5 integration points (single commit) +- One adapter file under `tts_adapters/` plus its line in the package import block - Client scripts and server launcher under `examples/online_serving/text_to_speech//` - Gradio demo with streaming and voice cloning UI in the same dir - E2E tests per **Test Case Writing (CI Levels)** above (priority tier determines L1–L4 scope) @@ -514,8 +519,8 @@ Use this checklist when integrating a new TTS model: - [ ] README.md written ### Phase 3: Online Serving -- [ ] All 5 `serving_speech.py` integration points added in one commit -- [ ] Only extract params in `_build_*_params` that are forwarded to the model call (ruff F841) +- [ ] Adapter written under `tts_adapters/` and registered in the import block +- [ ] Only extract params in `build()` that are forwarded to the model call (ruff F841) - [ ] Prompt builder handles text input correctly - [ ] Voice cloning works (if supported) - [ ] All response formats work (wav, mp3, flac, pcm) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index be649785d..a2ccc9a38 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -61,6 +61,13 @@ repos: types: [python] files: ^tests/ + - id: check-tts-adapter-migration + name: Ratchet TTS per-model branches out of serving_speech.py + entry: python tools/pre_commit/check_tts_adapter.py + language: python + pass_filenames: false + files: ^vllm_omni/entrypoints/openai/(serving_speech\.py|tts_adapters/__init__\.py)$ + # Keep `suggestion` last - id: suggestion name: Suggestion diff --git a/docs/contributing/model/adding_tts_model.md b/docs/contributing/model/adding_tts_model.md index 6b70fdb1e..45680ee5d 100644 --- a/docs/contributing/model/adding_tts_model.md +++ b/docs/contributing/model/adding_tts_model.md @@ -672,100 +672,147 @@ down and spawn a new one mid-module. A few rules that save real CI debugging tim ## Online Serving Integration -To expose your model through the `/v1/audio/speech` OpenAI-compatible endpoint, add -**all five** of the following integration points to -`vllm_omni/entrypoints/openai/serving_speech.py` in a **single commit**. Adding them -piecemeal causes partial-integration failures that are hard to debug. +To expose your model through the `/v1/audio/speech` OpenAI-compatible endpoint, write +one adapter under `vllm_omni/entrypoints/openai/tts_adapters/`. The adapter declares +how your model is recognized and owns its request handling; `serving_speech.py` should +not need an edit. -### 1. Stage constant +### 1. Write the adapter -Near the top of the file, alongside the other `_*_TTS_MODEL_STAGES` constants: +Create `vllm_omni/entrypoints/openai/tts_adapters/your_model.py`: ```python -_YOUR_MODEL_TTS_MODEL_STAGES = {"your_model_stage_key"} +# SPDX-License-Identifier: Apache-2.0 +"""YourModel serving adapter.""" + +from typing import TYPE_CHECKING + +from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter +from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest + +if TYPE_CHECKING: + from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest + + +@register_tts_adapter +class YourModelAdapter(ARTTSAdapter): + # The registry key, and the model-type string used in logs. + name = "your_model" + # The engine ``model_stage`` value(s) your deployment yaml sets. Detection, + # stage discovery and `/v1/audio/speech` routing all derive from this. + stage_keys = frozenset({"your_model_stage_key"}) + + def validate(self, request: "OpenAICreateSpeechRequest") -> str | None: + """Return an error string, or None if the request is valid.""" + if not request.input or not request.input.strip(): + return "Input text cannot be empty" + return None + + async def build( + self, + request: "OpenAICreateSpeechRequest", + sampling_params_list: list, + has_inline_ref_audio: bool, + ) -> PreparedRequest: + """Build the engine prompt and per-model parameters.""" + params: dict[str, list] = {"text": [request.input]} + if request.voice is not None: + params["voice"] = [request.voice] + return PreparedRequest( + prompt={"prompt": request.input}, + tts_params=params, + model_type=self.name, + ) ``` -### 2. Union into `_TTS_MODEL_STAGES` +Then add it to the import block at the bottom of +`vllm_omni/entrypoints/openai/tts_adapters/__init__.py` so it registers on import. +That import line is the only shared file a new model has to touch. -Add to the `_TTS_MODEL_STAGES` set union: +> **Pure-diffusion TTS does not go through adapters yet.** When the server is built +> by `OmniOpenAIServingSpeech.for_diffusion()`, `create_speech()` routes straight to +> `_create_diffusion_speech()` and never calls `validate()` or `build()`. +> `DiffusionTTSAdapter` is scaffolding for that migration and has no production +> subclass today, so validation or prompt construction placed in one would silently +> never run. If your model is served by the diffusion engine rather than +> `engine_client`, follow the existing diffusion path and say so on #4855 — wiring +> that path through the adapter layer is open work, not something to do inside a +> new-model PR. + +### 2. If a stage key alone cannot identify your model + +Two escape hatches, in order of preference: ```python -_TTS_MODEL_STAGES: set[str] = ( - ... - | _YOUR_MODEL_TTS_MODEL_STAGES -) +# The architecture is authoritative (e.g. your model has no dedicated +# model_stage value, or shares one with another model). +model_archs = frozenset({"YourModelForConditionalGeneration"}) + +# Set this too if your model owns NO stage key, so stage discovery has to find +# your AR entry stage by architecture. +arch_identifies_entry_stage = True ``` -### 3. Model type detection - -In `_detect_tts_model_type()`, add before the final `return None`: +If your rule is not a set-membership test, override `matches()` — see +`covo_audio.py`, which shares the generic `fused_thinker_talker` stage key with +non-CoVo deployments and confirms with the architecture: ```python -if model_stage in _YOUR_MODEL_TTS_MODEL_STAGES: - return "your_model" +@classmethod +def matches(cls, model_stage: str | None, model_arch: str | None) -> bool: + return super().matches(model_stage, model_arch) and bool(model_arch) and "CovoAudio" in model_arch ``` -### 4. Request validation dispatch +If your model can match the same deployment as another adapter, give one of them an +explicit `detect_priority` (lower runs first). `tests/entrypoints/openai_api/ +test_tts_detection.py` fails if two same-priority detectors can both match one input, +so an unordered overlap cannot be merged. -In `_validate_tts_request()`, add before the fallback `return`: +Models that are speech-capable only in some deployment topologies override +`stage_serves_speech()` — see `audex.py`, whose omni thinker is text-final unless the +speech decoder is deployed alongside it. -```python -if self._tts_model_type == "your_model": - return self._validate_your_model_request(request) -``` +### 3. Reuse the shared helpers -### 5. Validation and parameter-builder methods +Adapters reach shared serving helpers through `self.ctx.server`: reference-audio +resolution (`_resolve_ref_audio`), uploaded-speaker handling +(`_apply_uploaded_speaker`), format validation (`_validate_ref_audio_format`), and +prompt-length limits (`_max_instructions_length`). Read a comparable adapter before +writing your own — `fish_speech.py` for voice cloning, `higgs_audio_v3.py` for +parameter-heavy models, `moss_tts.py` for a family sharing one base class. -Add two new methods: +### 4. Test it -```python -def _validate_your_model_request( - self, request: OpenAICreateSpeechRequest -) -> str | None: - """Validate YourModel request. Returns an error string or None.""" - if not request.input or not request.input.strip(): - return "Input text cannot be empty" - return None +`tests/entrypoints/openai_api/test_tts_adapter.py` covers registry invariants and +`test_tts_detection.py` covers detection. Add your model to +`EXPECTED_MODEL_TYPES` in the former, and a case to the latter if you added a +`model_archs` or `matches()` rule. -def _build_your_model_params( - self, request: OpenAICreateSpeechRequest -) -> dict[str, Any]: - """Build additional_information dict for YourModel.""" - params: dict[str, Any] = {"text": [request.input]} - if request.voice is not None: - params["voice"] = [request.voice] - # Add any other model-specific fields here - return params -``` +### Do not add branches to `serving_speech.py` -Then wire `_build_your_model_params` into the request-dispatch block in -`_create_tts_request()` (search for the equivalent `_build_*_params` call for an -existing model to find the right location). If the model supports voice cloning -(`ref_audio` → `prompt_audio_path`, `ref_text` → `prompt_text`), add those mappings -here too — follow any existing `_build__params` in `serving_speech.py` (e.g. -`_build_moss_tts_params` for the voice-cloning variant) for the pattern. +Older models predate the adapter framework and still have per-model branches in +`serving_speech.py`; they are being migrated out (RFC #4327, #4855). Do not copy +that pattern. `tools/pre_commit/check_tts_adapter.py` is a ratchet on the remaining +`self._tts_model_type == ...` comparisons and fails the commit if the count grows. -> **Two dispatch patterns coexist:** Fish Speech uses a `self._is_fish_speech` boolean -> checked *before* `elif self._is_tts`. All newer models use the `_tts_model_type` -> string pattern shown above. For new models, always use the string pattern — do not -> add new `_is_*` boolean flags. - -> **Note on unused variables:** Only extract parameters in `_build_your_model_params` -> that you actually pass to the model's generate / `inference_stream` call. Extracting -> a variable without forwarding it will trigger a `ruff F841` pre-commit failure. +If your model genuinely needs behaviour that no adapter hook can express, that is a +missing hook — propose it on the RFC rather than adding a branch. ### Merge conflicts -`serving_speech.py` is modified by every new model PR and is the most common source of -rebase conflicts. When rebasing onto `main` and a conflict appears here, the resolution -is always to **keep both** the upstream model's additions and your own — never discard -either side. After resolving: +Your adapter is a new file, so it cannot conflict. The one shared line is your entry +in the import block at the bottom of `tts_adapters/__init__.py`, which conflicts only +when another model lands at the same time. The resolution is always to **keep both** +imports — never discard either side. After resolving: ```bash -git add vllm_omni/entrypoints/openai/serving_speech.py +git add vllm_omni/entrypoints/openai/tts_adapters/__init__.py git rebase --continue ``` +If you find yourself resolving a conflict inside `serving_speech.py`, something in +your model is still going through the legacy path; see the note above. + ## Single-Stage Models Some TTS models (e.g. MOSS-TTS-Nano) do not use a two-stage pipeline. Instead the diff --git a/tests/entrypoints/openai_api/test_audex_serving_guards.py b/tests/entrypoints/openai_api/test_audex_serving_guards.py index 343eb9816..68661865f 100644 --- a/tests/entrypoints/openai_api/test_audex_serving_guards.py +++ b/tests/entrypoints/openai_api/test_audex_serving_guards.py @@ -20,6 +20,10 @@ from vllm_omni.entrypoints.openai.serving_speech import ( OmniOpenAIServingSpeech, ) +# Serving-layer only: the engine is stubbed, so these run on the CPU lane +# alongside the other tests/entrypoints/openai_api suites. +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + def _serving(model_type: str = "audex") -> OmniOpenAIServingSpeech: serving = OmniOpenAIServingSpeech.__new__(OmniOpenAIServingSpeech) diff --git a/tests/entrypoints/openai_api/test_serving_speech.py b/tests/entrypoints/openai_api/test_serving_speech.py index 3739f8860..d36bc84be 100644 --- a/tests/entrypoints/openai_api/test_serving_speech.py +++ b/tests/entrypoints/openai_api/test_serving_speech.py @@ -3445,16 +3445,17 @@ class TestCosyVoice3Serving: def test_cosyvoice3_model_type_detection(self, cosyvoice3_server): assert cosyvoice3_server._tts_model_type == "cosyvoice3" assert cosyvoice3_server._is_tts is True - assert cosyvoice3_server._is_cosyvoice3 is True def test_cosyvoice3_stage_registered(self): - from vllm_omni.entrypoints.openai.serving_speech import ( - _COSYVOICE3_TTS_MODEL_STAGES, - _TTS_MODEL_STAGES, + from vllm_omni.entrypoints.openai.tts_adapters import ( + all_tts_stage_keys, + detect_tts_model_type, + resolve_adapter, ) - assert "cosyvoice3_talker" in _COSYVOICE3_TTS_MODEL_STAGES - assert "cosyvoice3_talker" in _TTS_MODEL_STAGES + assert "cosyvoice3_talker" in resolve_adapter("cosyvoice3").stage_keys + assert "cosyvoice3_talker" in all_tts_stage_keys() + assert detect_tts_model_type("cosyvoice3_talker", None) == "cosyvoice3" def test_validate_cosyvoice3_empty_input(self, cosyvoice3_server): request = OpenAICreateSpeechRequest(input="", ref_audio="data:audio/wav;base64,abc", ref_text="hello") diff --git a/tests/entrypoints/openai_api/test_tts_detection.py b/tests/entrypoints/openai_api/test_tts_detection.py new file mode 100644 index 000000000..27c6082c1 --- /dev/null +++ b/tests/entrypoints/openai_api/test_tts_detection.py @@ -0,0 +1,345 @@ +# SPDX-License-Identifier: Apache-2.0 +"""TTS model detection: registry-driven stage -> model-type resolution. + +``serving_speech.py`` used to carry a hand-written 20-branch ladder mapping +``model_stage``/``model_arch`` to a model-type string, duplicating the +``stage_keys`` every adapter already declared. Detection is now derived from the +adapters. ``test_detection_matches_legacy_ladder`` pins the new implementation +against a verbatim copy of the ladder it replaced, over every input the ladder +could distinguish, so the migration is checkable rather than reviewable-by-eye. +""" + +import itertools + +import pytest + +from vllm_omni.entrypoints.openai.tts_adapters import ( + LEGACY_TTS_DETECTORS, + TTS_ADAPTER_REGISTRY, + all_tts_stage_keys, + detect_tts_model_type, + iter_tts_detectors, + resolve_adapter, + tts_entry_stage_archs, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +# -------------------------------------------------------------------------- +# The legacy ladder, copied verbatim from serving_speech.py at 67c54777b +# (the commit this refactor is based on), with the module constants inlined. +# Do not "clean up" this function: its value is being a frozen oracle. +# -------------------------------------------------------------------------- +_LEGACY_MING_TTS_MODEL_ARCHS = {"MingTTSForConditionalGeneration"} +_LEGACY_VOXTRAL_TTS_MODEL_STAGES = {"audio_generation"} +_LEGACY_QWEN3_TTS_MODEL_STAGES = {"qwen3_tts"} +_LEGACY_FISH_TTS_MODEL_STAGES = {"fish_speech_slow_ar"} +_LEGACY_COSYVOICE3_TTS_MODEL_STAGES = {"cosyvoice3_talker"} +_LEGACY_OMNIVOICE_TTS_MODEL_STAGES = {"omnivoice_generator"} +_LEGACY_COVO_AUDIO_MODEL_STAGES = {"fused_thinker_talker"} +_LEGACY_VOXCPM2_TTS_MODEL_STAGES = {"latent_generator"} +_LEGACY_MING_TTS_MODEL_STAGES = {"ming_tts"} +_LEGACY_MOSS_TTS_MODEL_STAGES = {"moss_tts_nano"} +_LEGACY_MOSS_TTS_FULL_MODEL_STAGES = {"moss_tts", "moss_tts_codec"} +_LEGACY_MOSS_TTS_LOCAL_MODEL_STAGES = {"moss_tts_local", "moss_tts_local_codec"} +_LEGACY_HIGGS_AUDIO_V2_TTS_MODEL_STAGES = {"higgs_audio_v2"} +_LEGACY_HIGGS_V3_TTS_MODEL_STAGES = {"higgs_audio_v3"} +_LEGACY_GLM_TTS_MODEL_STAGES = {"glm_tts"} +_LEGACY_STEP_AUDIO2_TTS_MODEL_STAGES = {"step_audio2_thinker"} +_LEGACY_INDEXTTS2_TTS_MODEL_STAGES = {"indextts2_talker"} +_LEGACY_AUDEX_TTS_MODEL_STAGES = {"audex_thinker", "audex_omni"} +_LEGACY_AUDEX_TTA_MODEL_STAGES = {"audex_tta_thinker"} + +_LEGACY_TTS_MODEL_STAGES = ( + _LEGACY_VOXTRAL_TTS_MODEL_STAGES + | _LEGACY_QWEN3_TTS_MODEL_STAGES + | _LEGACY_FISH_TTS_MODEL_STAGES + | _LEGACY_COSYVOICE3_TTS_MODEL_STAGES + | _LEGACY_OMNIVOICE_TTS_MODEL_STAGES + | _LEGACY_HIGGS_AUDIO_V2_TTS_MODEL_STAGES + | _LEGACY_HIGGS_V3_TTS_MODEL_STAGES + | _LEGACY_COVO_AUDIO_MODEL_STAGES + | _LEGACY_VOXCPM2_TTS_MODEL_STAGES + | _LEGACY_MING_TTS_MODEL_STAGES + | _LEGACY_MOSS_TTS_MODEL_STAGES + | _LEGACY_MOSS_TTS_FULL_MODEL_STAGES + | _LEGACY_MOSS_TTS_LOCAL_MODEL_STAGES + | _LEGACY_GLM_TTS_MODEL_STAGES + | _LEGACY_STEP_AUDIO2_TTS_MODEL_STAGES + | _LEGACY_INDEXTTS2_TTS_MODEL_STAGES + | _LEGACY_AUDEX_TTS_MODEL_STAGES + | _LEGACY_AUDEX_TTA_MODEL_STAGES +) + + +def _legacy_detect(model_stage, model_arch): + if model_arch == "VoxCPM2TalkerForConditionalGeneration": + return "voxcpm2" + if model_stage in _LEGACY_QWEN3_TTS_MODEL_STAGES: + return "qwen3_tts" + if model_stage in _LEGACY_VOXTRAL_TTS_MODEL_STAGES: + return "voxtral_tts" + if model_stage in _LEGACY_FISH_TTS_MODEL_STAGES: + return "fish_tts" + if model_stage in _LEGACY_COSYVOICE3_TTS_MODEL_STAGES: + return "cosyvoice3" + if model_stage in _LEGACY_OMNIVOICE_TTS_MODEL_STAGES: + return "omnivoice" + if model_stage in _LEGACY_COVO_AUDIO_MODEL_STAGES: + if model_arch and "CovoAudio" in model_arch: + return "covo_audio" + if model_stage in _LEGACY_VOXCPM2_TTS_MODEL_STAGES: + return "voxcpm2" + if model_stage in _LEGACY_MING_TTS_MODEL_STAGES: + return "ming_flash_omni_tts" + if model_arch in _LEGACY_MING_TTS_MODEL_ARCHS: + return "ming_tts" + if model_stage in _LEGACY_MOSS_TTS_MODEL_STAGES: + return "moss_tts_nano" + if model_stage in _LEGACY_MOSS_TTS_FULL_MODEL_STAGES: + return "moss_tts" + if model_stage in _LEGACY_MOSS_TTS_LOCAL_MODEL_STAGES: + return "moss_tts" + if model_stage in _LEGACY_HIGGS_AUDIO_V2_TTS_MODEL_STAGES: + return "higgs_audio_v2" + if model_stage in _LEGACY_HIGGS_V3_TTS_MODEL_STAGES: + return "higgs_audio_v3" + if model_stage in _LEGACY_GLM_TTS_MODEL_STAGES: + return "glm_tts" + if model_stage in _LEGACY_STEP_AUDIO2_TTS_MODEL_STAGES: + return "step_audio2" + if model_stage in _LEGACY_INDEXTTS2_TTS_MODEL_STAGES: + return "indextts2" + if model_stage in _LEGACY_AUDEX_TTS_MODEL_STAGES: + return "audex" + if model_stage in _LEGACY_AUDEX_TTA_MODEL_STAGES: + return "audex_tta" + return None + + +# Every ``model_stage`` value any in-tree pipeline emits, plus any stage key an +# adapter claims, so the parametrization below covers the real deployment space +# rather than a hand-picked subset. Refresh the pipeline half with: +# grep -rhoE 'model_stage="[^"]+"' --include='pipeline*.py' vllm_omni/model_executor/ \ +# | sed 's/model_stage="//;s/"$//' | sort -u +# ``test_pipeline_stage_list_is_complete`` enforces the adapter half. +_PIPELINE_STAGES = [ + "AR", + "ar", + "asr", + "audex_code2wav", + "audex_omni", + "audex_thinker", + "audex_tta_thinker", + "audex_xcodec", + "audio_generation", + "audio_tokenizer", + "audio_vae", + "aura", + "code2wav", + "cosyvoice3_code2wav", + "cosyvoice3_talker", + "dac_decoder", + "diffusion", + "dit", + "fish_speech_slow_ar", + "fused_thinker_talker", + "glm_tts", + "glm_tts_dit", + "higgs_audio_v2", + "higgs_audio_v3", + "indextts2_s2mel_decoder", + "indextts2_talker", + "latent_generator", + "llm", + "ming_tts", + "moss_tts", + "moss_tts_codec", + "moss_tts_local", + "moss_tts_local_codec", + "moss_tts_nano", + # Claimed by the OmniVoice adapter but emitted by no pipeline: OmniVoice is + # served through the diffusion path (``OmniOpenAIServingSpeech.for_diffusion`` + # sets ``_tts_model_type`` directly and leaves ``_tts_stage`` None), so + # detection is never consulted for it. The key is vestigial in the adapter + # exactly as it was in the deleted ladder; kept here so detection is still + # exercised on it. + "omnivoice_generator", + "qwen3_tts", + "step_audio2_thinker", + "talker", + "thinker", + "token2audio", + "token2image", + "token2text", + "token2wav", + "tts", +] + +_STAGES = [*_PIPELINE_STAGES, None, "vae", "not_a_real_stage"] +_ARCHS = [ + None, + "VoxCPM2TalkerForConditionalGeneration", + "MingTTSForConditionalGeneration", + "CovoAudioForConditionalGeneration", + "MyCovoAudioThing", + "Qwen3TTSForConditionalGeneration", +] + +#: Architectures that no adapter claims. Paired with any stage key these are +#: inert, so the full cross product is meaningful for them. +_UNCLAIMED_ARCHS = [None, "Qwen3TTSForConditionalGeneration"] + +#: Architectures a detector matches on. Pairing one of these with *another* +#: model's stage key describes a deployment no pipeline can produce (the +#: architecture comes from the checkpoint, the stage key from the pipeline +#: definition), and the two implementations deliberately disagree there — +#: see ``test_arch_matching_is_a_fallback_not_an_override``. +_CLAIMED_ARCH_CASES = [ + # Ming dense: its AR stage is the generic ``llm``. + ("llm", "MingTTSForConditionalGeneration"), + ("audio_vae", "MingTTSForConditionalGeneration"), + (None, "MingTTSForConditionalGeneration"), + # Ming flash-omni owns the ``ming_tts`` stage key and must win it. + ("ming_tts", "MingTTSForConditionalGeneration"), + # VoxCPM2: talker architecture plus its own stage key. + ("latent_generator", "VoxCPM2TalkerForConditionalGeneration"), + (None, "VoxCPM2TalkerForConditionalGeneration"), + # CoVo shares ``fused_thinker_talker`` with non-CoVo fused deployments. + ("fused_thinker_talker", "CovoAudioForConditionalGeneration"), + ("fused_thinker_talker", "MyCovoAudioThing"), + ("fused_thinker_talker", None), + ("fused_thinker_talker", "Qwen3TTSForConditionalGeneration"), +] + +_REACHABLE = list(itertools.product(_STAGES, _UNCLAIMED_ARCHS)) + _CLAIMED_ARCH_CASES + + +@pytest.mark.parametrize(("model_stage", "model_arch"), _REACHABLE) +def test_detection_matches_legacy_ladder(model_stage, model_arch): + """Registry detection is identical to the ladder it replaced. + + Covers every ``model_stage`` any in-tree pipeline emits (``_PIPELINE_STAGES``, + plus negatives) against every architecture no adapter claims, and each + claimed architecture against the stage keys it really ships with. Pairing a + claimed architecture with an unrelated model's stage key is excluded because + no pipeline can produce it, and is covered separately by + ``test_arch_matching_is_a_fallback_not_an_override``. + """ + assert detect_tts_model_type(model_stage, model_arch) == _legacy_detect(model_stage, model_arch) + + +def test_arch_matching_is_a_fallback_not_an_override(): + """The one deliberate divergence from the ladder, pinned. + + The ladder tested Ming's architecture at a fixed position partway down, so a + Ming-architecture model deployed under *another* model's stage key resolved + to ``ming_tts`` if that key was checked below Ming, and to the other model if + above. Detection now treats an architecture-only match as a fallback that + runs after every adapter owning a real stage key, so the stage key always + wins. That is both self-consistent and unreachable in practice: a + checkpoint's architecture and a pipeline's stage key cannot be mixed across + models. + """ + assert detect_tts_model_type("audex_omni", "MingTTSForConditionalGeneration") == "audex" + assert _legacy_detect("audex_omni", "MingTTSForConditionalGeneration") == "ming_tts" + # Where the two real Ming deployments live, they agree. + assert detect_tts_model_type("llm", "MingTTSForConditionalGeneration") == "ming_tts" + assert detect_tts_model_type("ming_tts", "MingTTSForConditionalGeneration") == "ming_flash_omni_tts" + + +def test_stage_keys_cover_legacy_stage_set(): + """No stage key was dropped when the module constants were deleted.""" + assert all_tts_stage_keys() == frozenset(_LEGACY_TTS_MODEL_STAGES) + + +def test_pipeline_stage_list_is_complete(): + """``_PIPELINE_STAGES`` must still list every stage key adapters claim. + + Guards the parametrization above from silently narrowing: a new adapter + whose stage key is missing here would be detected but never exercised. + Regenerate the list from the pipelines (see the comment beside it) when a + new model lands. + """ + missing = all_tts_stage_keys() - set(_PIPELINE_STAGES) + assert not missing, f"stage keys claimed by adapters but absent from _PIPELINE_STAGES: {sorted(missing)}" + + +def test_entry_stage_archs_is_ming_only(): + """Only Ming dense identifies its entry stage by architecture. + + Widening this set would change which stage ``_find_tts_stage`` selects in + mixed deployments — notably VoxCPM2, which declares ``model_archs`` but is + found by its ``latent_generator`` stage key. + """ + assert tts_entry_stage_archs() == frozenset({"MingTTSForConditionalGeneration"}) + + +def test_every_detected_type_is_adapter_backed_or_declared_legacy(): + """Detection may only name a model that has an adapter or an explicit + :data:`LEGACY_TTS_DETECTORS` entry — never an undeclared string.""" + legacy_names = {d.name for d in LEGACY_TTS_DETECTORS} + for stage, arch in itertools.product(_STAGES, _ARCHS): # full cross product on purpose + detected = detect_tts_model_type(stage, arch) + if detected is None: + continue + assert detected in TTS_ADAPTER_REGISTRY or detected in legacy_names, ( + f"detection produced undeclared model type {detected!r} for stage={stage!r} arch={arch!r}" + ) + + +def test_legacy_detectors_have_no_adapter(): + """A legacy entry that gained an adapter must be deleted from the list.""" + for detector in LEGACY_TTS_DETECTORS: + assert resolve_adapter(detector.name) is None, ( + f"{detector.name!r} now has an adapter; remove it from LEGACY_TTS_DETECTORS" + ) + + +def test_no_ambiguous_detectors_at_equal_priority(): + """Two detectors that can match the same deployment must be ordered. + + This is what makes ``detect_tts_model_type`` independent of adapter import + order: any real overlap has to be resolved by an explicit + ``detect_priority``, not by whichever module happened to register first. + """ + detectors = iter_tts_detectors() + for left, right in itertools.combinations(detectors, 2): + if left.detect_priority != right.detect_priority: + continue + for stage, arch in itertools.product(_STAGES, _ARCHS): + assert not (left.matches(stage, arch) and right.matches(stage, arch)), ( + f"{left.name!r} and {right.name!r} both match stage={stage!r} arch={arch!r} " + f"at priority {left.detect_priority}; give one an explicit detect_priority" + ) + + +def test_detection_order_is_total_and_stable(): + """Resolution order must not depend on registry insertion order.""" + keys = [(d.detect_priority, d.name) for d in iter_tts_detectors()] + assert keys == sorted(keys) + assert len(set(keys)) == len(keys) + + +def test_adapters_declare_at_least_one_match_rule(): + """An adapter with neither stage keys nor architectures is unreachable.""" + for name, cls in TTS_ADAPTER_REGISTRY.items(): + assert cls.stage_keys or cls.model_archs, f"adapter {name!r} declares no stage_keys and no model_archs" + + +def test_audex_omni_requires_speech_decoder(): + """The Audex omni thinker only serves speech alongside its decoder.""" + audex = resolve_adapter("audex") + assert audex.stage_serves_speech("audex_omni", frozenset({"audex_omni", "audex_code2wav"})) is True + assert audex.stage_serves_speech("audex_omni", frozenset({"audex_omni"})) is False + # The dedicated TTS thinker is unconditional. + assert audex.stage_serves_speech("audex_thinker", frozenset({"audex_thinker"})) is True + + +def test_default_stage_serves_speech_is_unconditional(): + for name, cls in TTS_ADAPTER_REGISTRY.items(): + if name == "audex": + continue + for key in cls.stage_keys: + assert cls.stage_serves_speech(key, frozenset({key})) is True diff --git a/tests/tools/test_check_tts_adapter.py b/tests/tools/test_check_tts_adapter.py new file mode 100644 index 000000000..486a629d5 --- /dev/null +++ b/tests/tools/test_check_tts_adapter.py @@ -0,0 +1,79 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Tests for the TTS adapter migration ratchet. + +The gate only earns its place if it actually catches the ways a per-model branch +gets reintroduced. These pin what it detects — and, just as importantly, the one +shape it is known to miss, so nobody mistakes it for an adversarial sandbox. +""" + +from __future__ import annotations + +import ast +import sys +from pathlib import Path + +import pytest + +sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "tools" / "pre_commit")) + +from check_tts_adapter import ( # noqa: E402 + MAX_LEGACY_DETECTORS, + MAX_MODEL_TYPE_BRANCHES, + _legacy_detectors, + _model_type_branches, + main, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _count(src: str) -> int: + return len(_model_type_branches(ast.parse(src))) + + +@pytest.mark.parametrize( + ("label", "src", "expected"), + [ + ("direct comparison", 'if self._tts_model_type == "foo": pass', 1), + ("negated comparison", 'if self._tts_model_type != "foo": pass', 1), + ("membership test", 'if self._tts_model_type in ("a", "b"): pass', 1), + ("local alias", 'kind = self._tts_model_type\nif kind == "foo": pass', 1), + ("getattr read", 'if getattr(self, "_tts_model_type") == "foo": pass', 1), + ("getattr with default", 'if getattr(self, "_tts_model_type", None) == "foo": pass', 1), + ("match statement counts each case", 'match self._tts_model_type:\n case "a": pass\n case _: pass', 2), + ("alias then match", 'k = getattr(self, "_tts_model_type")\nmatch k:\n case "a": pass', 1), + ("unrelated attribute is ignored", 'if self._other == "foo": pass', 0), + ("unrelated local is ignored", 'kind = self._backend\nif kind == "foo": pass', 0), + ], +) +def test_branch_detection(label, src, expected): + assert _count(src) == expected, label + + +def test_dict_dispatch_is_a_known_gap(): + """Documented limitation, asserted so it is a decision and not a surprise. + + Table dispatch keyed on the model type reads as zero branches. The ratchet + guards against drift in the shape people actually write; catching every + indirection is a review question, not a lint one. + """ + assert _count('TABLE = {"foo": 1}\nvalue = TABLE[self._tts_model_type]') == 0 + + +def test_legacy_detector_counting(): + src = 'X = (LegacyDetector(name="a", stage_keys=frozenset({"s"})), LegacyDetector(name="b"))' + found = _legacy_detectors(ast.parse(src)) + assert [name for _, name in found] == ["a", "b"] + + +def test_gate_passes_on_current_tree(): + """The recorded budgets must match reality, or the ratchet is decorative.""" + assert main([]) == 0 + + +def test_budgets_are_the_real_counts(): + root = Path(__file__).resolve().parents[2] + serving = root / "vllm_omni" / "entrypoints" / "openai" / "serving_speech.py" + adapters = root / "vllm_omni" / "entrypoints" / "openai" / "tts_adapters" / "__init__.py" + assert len(_model_type_branches(ast.parse(serving.read_text()))) == MAX_MODEL_TYPE_BRANCHES + assert len(_legacy_detectors(ast.parse(adapters.read_text()))) == MAX_LEGACY_DETECTORS diff --git a/tools/pre_commit/check_tts_adapter.py b/tools/pre_commit/check_tts_adapter.py new file mode 100644 index 000000000..233e1a0e6 --- /dev/null +++ b/tools/pre_commit/check_tts_adapter.py @@ -0,0 +1,181 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Ratchet on the TTS adapter migration (RFC #4327, #4855). + +Per-model behaviour is being moved out of ``serving_speech.py`` and into +``entrypoints/openai/tts_adapters/``. The migration has stalled before, because +nothing stopped a new model from adding one more ``self._tts_model_type == ...`` +branch faster than the migration removed them. + +This gate makes the two debt counters monotonically non-increasing: + +1. ``self._tts_model_type`` comparisons in ``serving_speech.py`` +2. entries in ``tts_adapters.LEGACY_TTS_DETECTORS`` (models with no adapter) + +Lowering a budget after removing branches is expected and encouraged — the check +insists on it, so the recorded numbers cannot silently drift upward. Raising one +requires editing this file, which makes it a reviewable decision instead of an +invisible one. + +Source-only (``ast``), so it runs in pre-commit without importing vllm. +""" + +from __future__ import annotations + +import argparse +import ast +import sys +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[2] +SERVING_SPEECH = Path("vllm_omni/entrypoints/openai/serving_speech.py") +ADAPTERS_INIT = Path("vllm_omni/entrypoints/openai/tts_adapters/__init__.py") + +# Budgets. These may only go DOWN. See the module docstring. +MAX_MODEL_TYPE_BRANCHES = 29 +MAX_LEGACY_DETECTORS = 1 + +_GUIDANCE = """ +Move the per-model behaviour into its adapter under +``vllm_omni/entrypoints/openai/tts_adapters/`` instead of branching in the +shared serving module. ``docs/contributing/model/adding_tts_model.md`` walks +through it. If a branch genuinely cannot be expressed as an adapter hook, raise +the budget in this file and say why in the PR description. +""".strip() + + +_ATTR = "_tts_model_type" + + +def _reads_model_type(node: ast.AST) -> bool: + """Whether ``node`` evaluates to the model-type discriminator.""" + if isinstance(node, ast.Attribute) and node.attr == _ATTR: + return True + # getattr(self, "_tts_model_type") / getattr(self, "_tts_model_type", None) + return ( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id == "getattr" + and len(node.args) >= 2 + and isinstance(node.args[1], ast.Constant) + and node.args[1].value == _ATTR + ) + + +def _local_aliases(tree: ast.AST) -> set[str]: + """Local names bound directly to the model type, e.g. ``kind = self._tts_model_type``. + + Shallow on purpose: one hop is what an author reaches for when a branch gets + long, which is the drift this ratchet exists to catch. + """ + aliases: set[str] = set() + for node in ast.walk(tree): + value = None + if isinstance(node, ast.Assign): + value = node.value + elif isinstance(node, ast.AnnAssign | ast.AugAssign): + value = node.value + if value is None or not _reads_model_type(value): + continue + targets = node.targets if isinstance(node, ast.Assign) else [node.target] + for target in targets: + if isinstance(target, ast.Name): + aliases.add(target.id) + return aliases + + +def _model_type_branches(tree: ast.AST) -> list[tuple[int, str]]: + """Every branch keyed on the model-type discriminator. + + Counts direct comparisons, ``getattr`` reads, ``match`` subjects (one per + ``case``), and comparisons against a local alias of the attribute. This is a + ratchet against drift, not an adversarial sandbox: dispatch tables and other + indirections still slip past, and are a review question. + """ + aliases = _local_aliases(tree) + + def _is_subject(node: ast.AST) -> bool: + return _reads_model_type(node) or (isinstance(node, ast.Name) and node.id in aliases) + + found: list[tuple[int, str]] = [] + for node in ast.walk(tree): + if isinstance(node, ast.Compare): + if any(_is_subject(op) for op in [node.left, *node.comparators]): + found.append((node.lineno, ast.unparse(node))) + elif isinstance(node, ast.Match) and _is_subject(node.subject): + subject = ast.unparse(node.subject) + for case in node.cases: + found.append((case.pattern.lineno, f"match {subject}: case {ast.unparse(case.pattern)}")) + found.sort() + return found + + +def _legacy_detectors(tree: ast.AST) -> list[tuple[int, str]]: + """Every ``LegacyDetector(...)`` construction.""" + found: list[tuple[int, str]] = [] + for node in ast.walk(tree): + if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "LegacyDetector": + name = "?" + for kw in node.keywords: + if kw.arg == "name" and isinstance(kw.value, ast.Constant): + name = str(kw.value.value) + found.append((node.lineno, name)) + found.sort() + return found + + +def _check(label: str, path: Path, actual: int, budget: int, sample: list[str], budget_name: str) -> list[str]: + if actual > budget: + listing = "\n".join(f" {line}" for line in sample) + return [ + f"{path}: {actual} {label} exceeds the budget of {budget}.\n{listing}\n\n{_GUIDANCE}\n", + ] + if actual < budget: + return [ + f"{path}: {actual} {label}, down from a budget of {budget}. " + f"Lower {budget_name} to {actual} in {Path(__file__).relative_to(REPO_ROOT)} " + f"so the ratchet holds the new ground.", + ] + return [] + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + # pre-commit passes the staged filenames; we always audit the same two files, + # so they are accepted and ignored. + parser.add_argument("filenames", nargs="*") + parser.parse_args(argv) + + serving_path = REPO_ROOT / SERVING_SPEECH + adapters_path = REPO_ROOT / ADAPTERS_INIT + for path in (serving_path, adapters_path): + if not path.is_file(): + print(f"check_tts_adapter: expected file not found: {path}", file=sys.stderr) + return 1 + + branches = _model_type_branches(ast.parse(serving_path.read_text())) + detectors = _legacy_detectors(ast.parse(adapters_path.read_text())) + + errors = _check( + "`_tts_model_type` comparisons", + SERVING_SPEECH, + len(branches), + MAX_MODEL_TYPE_BRANCHES, + [f"line {lineno}: {src}" for lineno, src in branches], + "MAX_MODEL_TYPE_BRANCHES", + ) + errors += _check( + "unmigrated model types in LEGACY_TTS_DETECTORS", + ADAPTERS_INIT, + len(detectors), + MAX_LEGACY_DETECTORS, + [f"line {lineno}: {name}" for lineno, name in detectors], + "MAX_LEGACY_DETECTORS", + ) + + for error in errors: + print(f"check_tts_adapter: {error}", file=sys.stderr) + return 1 if errors else 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index 6cec05db8..abb7d939e 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -53,7 +53,10 @@ from vllm_omni.entrypoints.openai.speech_usage import ( ) from vllm_omni.entrypoints.openai.tts_adapters import ( SpeechServingContext, + 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.model_executor.models.fish_speech.prompt_utils import ( @@ -78,52 +81,17 @@ from vllm_omni.utils.speaker_cache import ( logger = init_logger(__name__) # TTS Configuration -_MING_TTS_MODEL_ARCHS = {"MingTTSForConditionalGeneration"} -_VOXTRAL_TTS_MODEL_STAGES = {"audio_generation"} -_QWEN3_TTS_MODEL_STAGES = {"qwen3_tts"} -_FISH_TTS_MODEL_STAGES = {"fish_speech_slow_ar"} -_COSYVOICE3_TTS_MODEL_STAGES = {"cosyvoice3_talker"} +# +# 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. +# # CosyVoice3 talker expects its reference transcript wrapped in the model # instruction template; without the delimiter the talker re-speaks the # reference (issue #4644). Matches the offline example/test and upstream demo. _COSYVOICE3_PROMPT_DELIMITER = "<|endofprompt|>" _COSYVOICE3_PROMPT_PREFIX = f"You are a helpful assistant.{_COSYVOICE3_PROMPT_DELIMITER}" -_OMNIVOICE_TTS_MODEL_STAGES = {"omnivoice_generator"} -_COVO_AUDIO_MODEL_STAGES = {"fused_thinker_talker"} -_VOXCPM2_TTS_MODEL_STAGES = {"latent_generator"} -_MING_TTS_MODEL_STAGES = {"ming_tts"} -_MOSS_TTS_MODEL_STAGES = {"moss_tts_nano"} -_MOSS_TTS_FULL_MODEL_STAGES = {"moss_tts", "moss_tts_codec"} -_MOSS_TTS_LOCAL_MODEL_STAGES = {"moss_tts_local", "moss_tts_local_codec"} -_HIGGS_AUDIO_V2_TTS_MODEL_STAGES = {"higgs_audio_v2"} -_HIGGS_V3_TTS_MODEL_STAGES = {"higgs_audio_v3"} -_GLM_TTS_MODEL_STAGES = {"glm_tts"} -_STEP_AUDIO2_TTS_MODEL_STAGES = {"step_audio2_thinker"} -_INDEXTTS2_TTS_MODEL_STAGES = {"indextts2_talker"} -# audex_omni covers the audex_s2s S2S deployment, whose -# TTS pass uses the same /v1/audio/speech surface. -_AUDEX_TTS_MODEL_STAGES = {"audex_thinker", "audex_omni"} -_AUDEX_TTA_MODEL_STAGES = {"audex_tta_thinker"} -_TTS_MODEL_STAGES: set[str] = ( - _VOXTRAL_TTS_MODEL_STAGES - | _QWEN3_TTS_MODEL_STAGES - | _FISH_TTS_MODEL_STAGES - | _COSYVOICE3_TTS_MODEL_STAGES - | _OMNIVOICE_TTS_MODEL_STAGES - | _HIGGS_AUDIO_V2_TTS_MODEL_STAGES - | _HIGGS_V3_TTS_MODEL_STAGES - | _COVO_AUDIO_MODEL_STAGES - | _VOXCPM2_TTS_MODEL_STAGES - | _MING_TTS_MODEL_STAGES - | _MOSS_TTS_MODEL_STAGES - | _MOSS_TTS_FULL_MODEL_STAGES - | _MOSS_TTS_LOCAL_MODEL_STAGES - | _GLM_TTS_MODEL_STAGES - | _STEP_AUDIO2_TTS_MODEL_STAGES - | _INDEXTTS2_TTS_MODEL_STAGES - | _AUDEX_TTS_MODEL_STAGES - | _AUDEX_TTA_MODEL_STAGES -) _SAMPLING_MAX_TOKENS_TTS_MODEL_TYPES = { "fish_tts", "qwen3_tts", @@ -446,7 +414,6 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): instance._diffusion_stage_configs = stage_configs instance._tts_model_type = "omnivoice" instance._is_tts = False - instance._is_fish_speech = 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 @@ -462,10 +429,6 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): # 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 - self._is_fish_speech = ( - self._tts_stage is not None - and getattr(getattr(self._tts_stage, "engine_args", None), "model_stage", None) == "fish_speech_slow_ar" - ) self._fish_speech_tokenizer = None self._covo_audio_tokenizer = None # Cached per process: the CosyVoice3 Qwen tokenizer + resolved model @@ -474,11 +437,6 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): # TTFP critical path) in _apply_cosyvoice3_dynamic_tokens. self._cosyvoice3_tokenizer = None - self._is_cosyvoice3 = ( - self._tts_stage is not None - and getattr(getattr(self._tts_stage, "engine_args", None), "model_stage", None) - in _COSYVOICE3_TTS_MODEL_STAGES - ) # Determine TTS model type or None self._tts_model_type = self._detect_tts_model_type() self.precomputed_speakers = self._load_precomputed_speakers() @@ -687,74 +645,43 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): def _find_tts_stage(self): """Find and return the TTS stage config, or None if not found.""" - all_stages = {getattr(stage.engine_args, "model_stage", None) for stage in self.engine_client.stage_configs} + tts_stage_keys = all_tts_stage_keys() + entry_stage_archs = tts_entry_stage_archs() + all_stages = frozenset( + getattr(stage.engine_args, "model_stage", None) for stage in self.engine_client.stage_configs + ) for stage in self.engine_client.stage_configs: engine_args = stage.engine_args model_stage = engine_args.model_stage model_arch = getattr(engine_args, "model_arch", None) worker_type = getattr(engine_args, "worker_type", None) - if model_stage in _TTS_MODEL_STAGES: - # The audio-capable Audex thinker is only speech-capable when - # deployed WITH the speech decoder (audex_s2s); - # the thinker-only deployment is text-final and must not - # accept /v1/audio/speech requests. - if model_stage == "audex_omni" and "audex_code2wav" not in all_stages: + 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 - # Ming dense identifies its AR entry stage by architecture because - # it does not use a dedicated TTS model_stage value. - if model_arch in _MING_TTS_MODEL_ARCHS and worker_type == "ar": + # 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 stage's model_stage attribute.""" + """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 - model_stage = getattr(self._tts_stage.engine_args, "model_stage", None) - model_arch = getattr(self._tts_stage.engine_args, "model_arch", None) - if model_arch == "VoxCPM2TalkerForConditionalGeneration": - return "voxcpm2" - if model_stage in _QWEN3_TTS_MODEL_STAGES: - return "qwen3_tts" - if model_stage in _VOXTRAL_TTS_MODEL_STAGES: - return "voxtral_tts" - if model_stage in _FISH_TTS_MODEL_STAGES: - return "fish_tts" - if model_stage in _COSYVOICE3_TTS_MODEL_STAGES: - return "cosyvoice3" - if model_stage in _OMNIVOICE_TTS_MODEL_STAGES: - return "omnivoice" - if model_stage in _COVO_AUDIO_MODEL_STAGES: - if model_arch and "CovoAudio" in model_arch: - return "covo_audio" - if model_stage in _VOXCPM2_TTS_MODEL_STAGES: - return "voxcpm2" - if model_stage in _MING_TTS_MODEL_STAGES: - return "ming_flash_omni_tts" - if model_arch in _MING_TTS_MODEL_ARCHS: - return "ming_tts" - if model_stage in _MOSS_TTS_MODEL_STAGES: - return "moss_tts_nano" - if model_stage in _MOSS_TTS_FULL_MODEL_STAGES: - return "moss_tts" - if model_stage in _MOSS_TTS_LOCAL_MODEL_STAGES: - return "moss_tts" - if model_stage in _HIGGS_AUDIO_V2_TTS_MODEL_STAGES: - return "higgs_audio_v2" - if model_stage in _HIGGS_V3_TTS_MODEL_STAGES: - return "higgs_audio_v3" - if model_stage in _GLM_TTS_MODEL_STAGES: - return "glm_tts" - if model_stage in _STEP_AUDIO2_TTS_MODEL_STAGES: - return "step_audio2" - if model_stage in _INDEXTTS2_TTS_MODEL_STAGES: - return "indextts2" - if model_stage in _AUDEX_TTS_MODEL_STAGES: - return "audex" - if model_stage in _AUDEX_TTA_MODEL_STAGES: - return "audex_tta" - return None + return detect_tts_model_type( + getattr(self._tts_stage.engine_args, "model_stage", None), + getattr(self._tts_stage.engine_args, "model_arch", None), + ) def _get_custom_voice_dir(self) -> str | None: try: diff --git a/vllm_omni/entrypoints/openai/tts_adapters/__init__.py b/vllm_omni/entrypoints/openai/tts_adapters/__init__.py index c1317c3af..f51fa4140 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/__init__.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/__init__.py @@ -1,12 +1,16 @@ # SPDX-License-Identifier: Apache-2.0 -"""Registry of TTS serving adapters. +"""Registry of TTS serving adapters, and the model detection built on it. -Adapters register themselves by their ``name`` (the model-type discriminator -returned by ``OmniOpenAIServingSpeech._detect_tts_model_type``) via -``@register_tts_adapter``. Resolution is by that name: detection already encodes -the (order-sensitive, partly ``model_arch``-based) stage logic, so keying the -registry on the resolved type is both simpler and correct by construction. -``stage_keys`` is retained on each adapter as documentation metadata. +Adapters register themselves by their ``name`` (the model-type discriminator) +via ``@register_tts_adapter``, and declare the deployment metadata that +identifies them: ``stage_keys``, plus ``model_archs`` for the few models a stage +key alone cannot distinguish. + +That metadata is the single source of truth. :func:`detect_tts_model_type` and +:func:`all_tts_stage_keys` derive the whole stage -> model-type mapping from it, +so a new TTS model is one adapter file and needs no edit to ``serving_speech.py``. +Models still awaiting an adapter are listed explicitly in +:data:`LEGACY_TTS_DETECTORS`. """ from vllm.logger import init_logger @@ -14,6 +18,7 @@ from vllm.logger import init_logger from vllm_omni.entrypoints.openai.tts_adapters.base import ( ARTTSAdapter, DiffusionTTSAdapter, + LegacyDetector, OutputPolicy, PreparedRequest, SpeechServingContext, @@ -24,6 +29,19 @@ logger = init_logger(__name__) TTS_ADAPTER_REGISTRY: dict[str, type[TTSModelAdapter]] = {} +#: Model types the serving layer detects that have no adapter yet. See +#: :class:`LegacyDetector`. This list must only ever shrink. +LEGACY_TTS_DETECTORS: tuple[LegacyDetector, ...] = ( + # Ming *flash-omni* claims the ``ming_tts`` stage key, while Ming *dense* + # (which does have an adapter) is identified by architecture. The stage key + # must therefore be tested first, hence the explicit priority. + LegacyDetector( + name="ming_flash_omni_tts", + stage_keys=frozenset({"ming_tts"}), + detect_priority=50, + ), +) + def register_tts_adapter(cls: type[TTSModelAdapter]) -> type[TTSModelAdapter]: """Class decorator: index ``cls`` under its ``name`` (model-type).""" @@ -48,6 +66,56 @@ def resolve_adapter(model_type: str | None) -> type[TTSModelAdapter] | None: return TTS_ADAPTER_REGISTRY.get(model_type) +def iter_tts_detectors() -> list[type[TTSModelAdapter] | LegacyDetector]: + """Every detector, in resolution order. + + Sorted by ``detect_priority`` then ``name``, so the order is total and does + not depend on adapter import order. + """ + detectors: list[type[TTSModelAdapter] | LegacyDetector] = [ + *TTS_ADAPTER_REGISTRY.values(), + *LEGACY_TTS_DETECTORS, + ] + detectors.sort(key=lambda d: (d.detect_priority, d.name)) + return detectors + + +def detect_tts_model_type(model_stage: str | None, model_arch: str | None) -> str | None: + """Resolve a deployed stage to its TTS model-type, or ``None``. + + The first detector that matches wins; ties are impossible for distinct model + types because overlapping detectors must declare an explicit + ``detect_priority`` (enforced by ``test_tts_detection.py``). + """ + for detector in iter_tts_detectors(): + if detector.matches(model_stage, model_arch): + return detector.name + return None + + +def all_tts_stage_keys() -> frozenset[str]: + """Every ``model_stage`` value served by ``/v1/audio/speech``.""" + keys: set[str] = set() + for detector in iter_tts_detectors(): + keys |= detector.stage_keys + return frozenset(keys) + + +def tts_entry_stage_archs() -> frozenset[str]: + """``model_arch`` values that identify an AR entry stage on their own. + + Only for models that own no ``model_stage`` key; stage discovery falls back + to these. Deliberately *not* every declared ``model_arch``: a model that also + owns a stage key is found by that key, and widening this would change which + stage is selected in mixed deployments. + """ + archs: set[str] = set() + for detector in iter_tts_detectors(): + if detector.arch_identifies_entry_stage: + archs |= detector.model_archs + return frozenset(archs) + + # Import adapter modules for their registration side effects. Keep at the bottom # so the registry helpers above are defined first. from vllm_omni.entrypoints.openai.tts_adapters import ( # noqa: E402,F401 @@ -72,12 +140,18 @@ from vllm_omni.entrypoints.openai.tts_adapters import ( # noqa: E402,F401 __all__ = [ "ARTTSAdapter", "DiffusionTTSAdapter", + "LegacyDetector", "OutputPolicy", "PreparedRequest", "SpeechServingContext", "TTSModelAdapter", + "LEGACY_TTS_DETECTORS", "TTS_ADAPTER_REGISTRY", "all_tts_model_types", + "all_tts_stage_keys", + "detect_tts_model_type", + "iter_tts_detectors", "register_tts_adapter", "resolve_adapter", + "tts_entry_stage_archs", ] diff --git a/vllm_omni/entrypoints/openai/tts_adapters/audex.py b/vllm_omni/entrypoints/openai/tts_adapters/audex.py index 890e132de..4533a420b 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/audex.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/audex.py @@ -29,6 +29,15 @@ class AudexAdapter(ARTTSAdapter): stage_keys = frozenset({"audex_thinker", "audex_omni"}) name = "audex" + @classmethod + def stage_serves_speech(cls, model_stage: str | None, all_stage_keys: frozenset[str]) -> bool: + """``audex_omni`` covers the audex_s2s deployment, whose TTS pass uses + the same ``/v1/audio/speech`` surface. Without the speech decoder the + same thinker is text-final and must not accept speech requests.""" + if model_stage == "audex_omni": + return "audex_code2wav" in all_stage_keys + return True + def validate(self, request: "OpenAICreateSpeechRequest") -> str | None: if not request.input or not request.input.strip(): return "Audex TTS requires non-empty input text" diff --git a/vllm_omni/entrypoints/openai/tts_adapters/base.py b/vllm_omni/entrypoints/openai/tts_adapters/base.py index 58c135044..69d6c5c6e 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/base.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/base.py @@ -99,8 +99,20 @@ class TTSModelAdapter(ABC): #: Stable discriminator string (the model-type from detection); registry key. name: ClassVar[str] - #: Engine ``model_stage`` key(s) this model uses, for documentation only. + #: Engine ``model_stage`` key(s) this model uses. Load-bearing: the serving + #: layer discovers the TTS stage and resolves the model type from these. stage_keys: ClassVar[frozenset[str]] = frozenset() + #: Engine ``model_arch`` value(s) that identify this model type. Checked + #: before ``stage_keys`` in :meth:`matches`. Only needed by models whose + #: ``model_stage`` alone is ambiguous or absent. + model_archs: ClassVar[frozenset[str]] = frozenset() + #: Set when the model has no dedicated ``model_stage`` value and its AR + #: entry stage must be discovered by ``model_archs`` instead (Ming dense). + arch_identifies_entry_stage: ClassVar[bool] = False + #: Detection order; lower runs first. Only set this to break a genuine + #: overlap with another adapter — see ``tests/.../test_tts_detection.py``, + #: which fails if two same-priority detectors can match the same input. + detect_priority: ClassVar[int] = 100 #: Serving backend: ``"ar"`` (engine_client) or ``"diffusion"``. backend: ClassVar[str] = "ar" @@ -111,6 +123,31 @@ class TTSModelAdapter(ABC): def __init__(self, ctx: SpeechServingContext) -> None: self.ctx = ctx + @classmethod + def matches(cls, model_stage: str | None, model_arch: str | None) -> bool: + """Whether a deployed stage with this ``model_stage``/``model_arch`` + is served by this adapter. + + Architecture wins over stage key, so a model that declares both is + recognized even when deployed under a stage key it shares with another + model. Override only for match rules that are not a set membership + test (see ``covo_audio``). + """ + if model_arch is not None and model_arch in cls.model_archs: + return True + return model_stage is not None and model_stage in cls.stage_keys + + @classmethod + def stage_serves_speech(cls, model_stage: str | None, all_stage_keys: frozenset[str]) -> bool: + """Whether a matched stage really accepts ``/v1/audio/speech`` in *this* + deployment. + + Lets a model that is speech-capable only in some topologies say so + (Audex: its omni thinker is text-final unless the speech decoder is + deployed alongside). Default: always. + """ + return True + def normalize(self, request: "OpenAICreateSpeechRequest") -> None: """In-place request normalization/mutation (e.g. infer task type, lowercase voice). Default: no-op.""" @@ -189,10 +226,37 @@ class DiffusionTTSAdapter(TTSModelAdapter): return frozenset(params) if params is not None else frozenset() +@dataclass(frozen=True) +class LegacyDetector: + """A model type the serving layer detects but has no adapter for yet. + + Detection has to keep naming these models so ``serving_speech.py`` can route + them down its legacy path, but they own none of the adapter contract. Each + entry is a migration debt: delete it when the model gets an adapter. The + pre-commit gate (``tools/pre_commit/check_tts_adapter.py``) fails if the list + grows. + + Exposes the subset of the :class:`TTSModelAdapter` class surface that + detection reads, so both kinds live in one sorted sequence. + """ + + name: str + stage_keys: frozenset[str] = frozenset() + model_archs: frozenset[str] = frozenset() + arch_identifies_entry_stage: bool = False + detect_priority: int = 100 + + def matches(self, model_stage: str | None, model_arch: str | None) -> bool: + if model_arch is not None and model_arch in self.model_archs: + return True + return model_stage is not None and model_stage in self.stage_keys + + # Re-exported here to avoid import cycles at call sites. __all__ = [ "ARTTSAdapter", "DiffusionTTSAdapter", + "LegacyDetector", "OutputPolicy", "PreparedRequest", "SpeechServingContext", diff --git a/vllm_omni/entrypoints/openai/tts_adapters/covo_audio.py b/vllm_omni/entrypoints/openai/tts_adapters/covo_audio.py index 70673ac54..28e64bc2a 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/covo_audio.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/covo_audio.py @@ -15,6 +15,13 @@ class CovoAudioAdapter(ARTTSAdapter): stage_keys = frozenset({"fused_thinker_talker"}) name = "covo_audio" + @classmethod + def matches(cls, model_stage: str | None, model_arch: str | None) -> bool: + """``fused_thinker_talker`` is a generic stage key that non-CoVo fused + deployments also use, so the architecture has to confirm it. A fused + stage from another model falls through to "not a TTS model".""" + return super().matches(model_stage, model_arch) and bool(model_arch) and "CovoAudio" in model_arch + def validate(self, request: "OpenAICreateSpeechRequest") -> str | None: if not request.input or not request.input.strip(): return "Input text cannot be empty" diff --git a/vllm_omni/entrypoints/openai/tts_adapters/ming_tts.py b/vllm_omni/entrypoints/openai/tts_adapters/ming_tts.py index adce42867..dd8b8b621 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/ming_tts.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/ming_tts.py @@ -17,8 +17,16 @@ logger = init_logger(__name__) @register_tts_adapter class MingTTSAdapter(ARTTSAdapter): - # Detected by model_arch (MingTTSForConditionalGeneration), not stage key. + # Ming dense has no dedicated model_stage value, so both stage discovery and + # model-type detection go through the architecture. ``ming_flash_omni_tts`` + # is the model that owns the ``ming_tts`` *stage key*, and is matched first. name = "ming_tts" + model_archs = frozenset({"MingTTSForConditionalGeneration"}) + arch_identifies_entry_stage = True + # Ming dense deploys its AR stage as the generic ``model_stage="llm"``, so + # this is an architecture *fallback*: it must run after every adapter that + # owns a real stage key, or it would claim stages those adapters own. + detect_priority = 200 def validate(self, request: "OpenAICreateSpeechRequest") -> str | None: """Validate Ming TTS request parameters. Returns error message or None.""" diff --git a/vllm_omni/entrypoints/openai/tts_adapters/voxcpm2.py b/vllm_omni/entrypoints/openai/tts_adapters/voxcpm2.py index 361aa2522..b0f188fbb 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/voxcpm2.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/voxcpm2.py @@ -16,7 +16,12 @@ class VoxCPM2Adapter(ARTTSAdapter): stage is present (and/or via ``model_arch``).""" stage_keys = frozenset({"latent_generator"}) + model_archs = frozenset({"VoxCPM2TalkerForConditionalGeneration"}) name = "voxcpm2" + # The talker architecture is authoritative and is tested ahead of every + # stage key, so a VoxCPM2 talker deployed under a stage key another model + # claims still resolves to VoxCPM2. + detect_priority = 10 def validate(self, request: "OpenAICreateSpeechRequest") -> str | None: """Validate VoxCPM2 request parameters. Returns error message or None."""