From 6ffdf1a8f1e47b01159b73b4cdd4c037e757dbba Mon Sep 17 00:00:00 2001 From: Sy03 <1370724210@qq.com> Date: Sat, 29 Aug 2026 05:35:27 +0800 Subject: [PATCH] [Bugfix][Qwen3-TTS] Handle codec generations that exhaust their token budget Signed-off-by: Sy03 <1370724210@qq.com> --- recipes/Qwen/Qwen3-TTS.md | 19 +- tests/engine/test_orchestrator.py | 8 +- .../openai_api/test_serving_speech.py | 167 +++++++++++++++++- .../openai_api/test_serving_speech_stream.py | 28 ++- .../openai_api/test_speech_usage.py | 9 + .../openai_api/test_tts_adapter.py | 56 +++++- vllm_omni/engine/stage_pool.py | 10 ++ .../entrypoints/openai/serving_speech.py | 113 ++++++++++-- .../openai/serving_speech_stream.py | 32 +++- vllm_omni/entrypoints/openai/speech_usage.py | 7 + .../openai/tts_adapters/__init__.py | 3 + .../entrypoints/openai/tts_adapters/base.py | 29 ++- .../openai/tts_adapters/qwen3_tts.py | 108 ++++++++++- vllm_omni/metrics/stats.py | 5 + 14 files changed, 551 insertions(+), 43 deletions(-) diff --git a/recipes/Qwen/Qwen3-TTS.md b/recipes/Qwen/Qwen3-TTS.md index 7365aafae..ede83d47d 100644 --- a/recipes/Qwen/Qwen3-TTS.md +++ b/recipes/Qwen/Qwen3-TTS.md @@ -18,11 +18,11 @@ examples in this repository. Qwen3-TTS supports three task types, each backed by a dedicated model checkpoint: -| Task Type | Model | Description | -| ------------- | ---------------------------------------- | ------------------------------------------------------------ | -| `CustomVoice` | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | Predefined speaker voices with optional style/emotion control | -| `VoiceDesign` | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | Generate speech from a natural language voice description | -| `Base` | `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | Voice cloning from reference audio + transcript | +| Task Type | Model | Description | +| ------------- | ---------------------------------------- | ------------------------------------------------------------- | +| `CustomVoice` | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | Predefined speaker voices with optional style/emotion control | +| `VoiceDesign` | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | Generate speech from a natural language voice description | +| `Base` | `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | Voice cloning from reference audio + transcript | Smaller 0.6B variants are also available for `CustomVoice` and `Base`. @@ -166,6 +166,15 @@ python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py --query-ty - Key flags: `--omni` is required. `--deploy-config` points to the bundled two-stage pipeline config. - Async chunking: Enabled by default in `qwen3_tts.yaml` for streaming-friendly first-audio latency. Raw audio streaming requires `stream=true`, `stream_format="audio"`, and `response_format="pcm"`. - Task/model matching: Each task type requires its matching model checkpoint. Using a CustomVoice model for a Base (voice clone) request will fail. +- Base codec termination: Base requests without an explicit `max_new_tokens` + use a text-scaled safety ceiling (at least 192 codec frames and no more than + the configured model limit). If the Talker reaches that ceiling without + codec EOS, non-streaming serving discards the incomplete audio and retries + once with a fresh seed; an explicit `seed` or `max_new_tokens` disables the + retry. SSE and WebSocket clients receive structured errors and must discard + previously emitted audio when the error contains `"action":"discard"`; + raw PCM streams terminate with a connection error after any partial bytes + already sent. - Known limitations: The server serves one model variant at a time. To switch task types (e.g., CustomVoice to Base), restart the server with the corresponding model. ## Hardware Support diff --git a/tests/engine/test_orchestrator.py b/tests/engine/test_orchestrator.py index 61a32fb86..5b257cb3c 100644 --- a/tests/engine/test_orchestrator.py +++ b/tests/engine/test_orchestrator.py @@ -2216,7 +2216,12 @@ def test_stage_pool_metrics_use_resumable_segment_token_count() -> None: ) output = SimpleNamespace( request_id="req-stream", - outputs=[SimpleNamespace(cumulative_token_ids=list(range(11)))], + outputs=[ + SimpleNamespace( + cumulative_token_ids=list(range(11)), + finish_reason=FinishReason.LENGTH, + ) + ], ) metrics = pool.build_stage_metrics( @@ -2228,6 +2233,7 @@ def test_stage_pool_metrics_use_resumable_segment_token_count() -> None: assert metrics.num_tokens_out == 3 assert metrics.output_unit_count == 3 + assert metrics.finish_reason == "length" def test_image_ttfo_preserves_request_time_and_tracks_stage_time() -> None: diff --git a/tests/entrypoints/openai_api/test_serving_speech.py b/tests/entrypoints/openai_api/test_serving_speech.py index b42de54f2..19d383504 100644 --- a/tests/entrypoints/openai_api/test_serving_speech.py +++ b/tests/entrypoints/openai_api/test_serving_speech.py @@ -45,9 +45,15 @@ from vllm_omni.entrypoints.openai.serving_speech import ( OmniOpenAIServingSpeech, _create_wav_header, ) -from vllm_omni.entrypoints.openai.tts_adapters.base import DEFAULT_TTS_LANGUAGES, PreparedRequest, SpeechServingContext +from vllm_omni.entrypoints.openai.tts_adapters.base import ( + DEFAULT_TTS_LANGUAGES, + PreparedRequest, + SpeechServingContext, + TTSGenerationError, +) from vllm_omni.entrypoints.openai.tts_adapters.capabilities import load_supported_speakers from vllm_omni.entrypoints.openai.tts_adapters.ming_tts import MingTTSAdapter +from vllm_omni.entrypoints.openai.tts_adapters.qwen3_tts import Qwen3TTSCodecLimitError from vllm_omni.entrypoints.openai.tts_adapters.voxtral import VoxtralTTSAdapter from vllm_omni.model_executor.models.fish_speech.prompt_utils import ( FISH_TEXT_ONLY_SYSTEM_PROMPT, @@ -2928,6 +2934,40 @@ class TestStreamingResponse: assert "text/event-stream" not in response.headers["content-type"] assert len(response.content) > 0 + def test_raw_stream_passes_tts_params_to_common_guard(self, streaming_app, monkeypatch): + """The raw route must not bypass the shared Qwen3-TTS EOS guard.""" + speech_server = streaming_app.state.speech_server + finalized_tts_params = {"_qwen3_tts_effective_max_tokens": [192]} + captured: dict = {} + + async def prepare(_request, request_id=None): + return request_id, object(), finalized_tts_params + + async def generate_chunks( + _generator, + _request_id, + _response_format="pcm", + raw_request=None, + request_start_s=None, + include_sample_rate=False, + usage_acc=None, + tts_params=None, + collect=None, + ): + captured["tts_params"] = tts_params + yield b"\x00\x00" + + monkeypatch.setattr(speech_server, "_prepare_speech_generation", prepare) + monkeypatch.setattr(speech_server, "_generate_audio_chunks", generate_chunks) + + response = TestClient(streaming_app).post( + "/v1/audio/speech", + json={"input": "Hello", "stream_format": "audio", "response_format": "pcm"}, + ) + + assert response.status_code == 200 + assert captured["tts_params"] is finalized_tts_params + def test_sse_streaming(self, streaming_app): """stream_format=sse without stream=True returns audio deltas as SSE.""" client = TestClient(streaming_app) @@ -4848,6 +4888,57 @@ class TestTTSAsyncOffloading: assert ("artifact-a", False) not in qwen3_tts_server._ref_audio_model_artifact_ready assert "artifact-b" in {entry[3] for entry in qwen3_tts_server._ref_audio_resolve_cache.values()} + @pytest.mark.asyncio + async def test_qwen3_tts_nonstream_retries_codec_limit_once(self, qwen3_tts_server, mocker: MockerFixture) -> None: + qwen3_tts_server._check_model = mocker.AsyncMock(return_value=None) + qwen3_tts_server._generate_audio_bytes = mocker.AsyncMock( + side_effect=[ + Qwen3TTSCodecLimitError("codec limit"), + (b"RIFF" + b"\x00" * 32, "audio/wav"), + ] + ) + request = OpenAICreateSpeechRequest(input="Hello", task_type="Base") + + response = await qwen3_tts_server.create_speech(request) + + assert response.status_code == 200 + assert qwen3_tts_server._generate_audio_bytes.await_count == 2 + first_call, retry_call = qwen3_tts_server._generate_audio_bytes.await_args_list + assert first_call.args[0] is request + assert first_call.kwargs["request_id"] != retry_call.kwargs["request_id"] + assert retry_call.args[0].seed is not None + assert request.seed is None + + @pytest.mark.asyncio + async def test_nonstream_does_not_retry_nonretryable_generation_error( + self, qwen3_tts_server, mocker: MockerFixture + ) -> None: + qwen3_tts_server._check_model = mocker.AsyncMock(return_value=None) + qwen3_tts_server._generate_audio_bytes = mocker.AsyncMock(side_effect=TTSGenerationError("invalid generation")) + + response = await qwen3_tts_server.create_speech(OpenAICreateSpeechRequest(input="Hello", task_type="Base")) + + assert response.status_code == 500 + qwen3_tts_server._generate_audio_bytes.assert_awaited_once() + + @pytest.mark.asyncio + @pytest.mark.parametrize("request_kwargs", [{"seed": 100423}, {"max_new_tokens": 192}]) + async def test_qwen3_tts_nonstream_does_not_retry_explicit_sampling_controls( + self, + qwen3_tts_server, + mocker: MockerFixture, + request_kwargs: dict[str, int], + ) -> None: + qwen3_tts_server._check_model = mocker.AsyncMock(return_value=None) + qwen3_tts_server._generate_audio_bytes = mocker.AsyncMock(side_effect=Qwen3TTSCodecLimitError("codec limit")) + + response = await qwen3_tts_server.create_speech( + OpenAICreateSpeechRequest(input="Hello", task_type="Base", **request_kwargs) + ) + + assert response.status_code == 500 + qwen3_tts_server._generate_audio_bytes.assert_awaited_once() + @pytest.mark.asyncio async def test_generate_audio_chunks_discards_ref_audio_artifact_warmup_on_error(self, qwen3_tts_server): async def failing_generator(): @@ -4862,6 +4953,80 @@ class TestTTSAsyncOffloading: assert "req-fail" not in qwen3_tts_server._request_ref_audio_artifact_keys assert ("artifact-fail", False) not in qwen3_tts_server._ref_audio_model_artifact_ready + @pytest.mark.asyncio + async def test_generate_audio_chunks_rejects_qwen3_codec_limit_before_success(self, qwen3_tts_server): + async def length_limited_generator(): + yield SimpleNamespace( + multimodal_output={"audio": torch.zeros(16, dtype=torch.float32), "sr": 24000}, + metrics={ + "stage_metrics": { + "0": { + "num_tokens_out": 191, + "finish_reason": "length", + } + } + }, + ) + + chunks = qwen3_tts_server._generate_audio_chunks( + length_limited_generator(), + "req-codec-limit", + tts_params={"task_type": ["Base"], "_qwen3_tts_effective_max_tokens": [192]}, + ) + with pytest.raises(Qwen3TTSCodecLimitError, match="191/192"): + async for _ in chunks: + pass + + @pytest.mark.asyncio + async def test_sse_codec_limit_marks_partial_audio_for_discard(self, qwen3_tts_server): + async def length_limited_generator(): + yield SimpleNamespace( + multimodal_output={"audio": torch.zeros(16, dtype=torch.float32), "sr": 24000}, + metrics={ + "stage_metrics": { + "0": { + "num_tokens_out": 191, + "finish_reason": "length", + } + } + }, + ) + + events = [ + event + async for event in qwen3_tts_server._generate_audio_sse_events( + length_limited_generator(), + "req-sse-codec-limit", + tts_params={"task_type": ["Base"], "_qwen3_tts_effective_max_tokens": [192]}, + ) + ] + + assert any("event: speech.audio.delta" in event for event in events) + assert not any("event: speech.audio.done" in event for event in events) + error_event = next(event for event in events if "event: speech.audio.error" in event) + payload = json.loads(next(line for line in error_event.splitlines() if line.startswith("data: "))[6:]) + assert payload["error"]["partial_audio"] is True + assert payload["error"]["action"] == "discard" + + @pytest.mark.asyncio + async def test_sse_error_before_audio_is_not_marked_partial(self, qwen3_tts_server): + async def failing_generator(): + raise RuntimeError("boom") + yield # pragma: no cover + + events = [ + event + async for event in qwen3_tts_server._generate_audio_sse_events( + failing_generator(), + "req-sse-no-audio", + ) + ] + + error_event = next(event for event in events if "event: speech.audio.error" in event) + payload = json.loads(next(line for line in error_event.splitlines() if line.startswith("data: "))[6:]) + assert "partial_audio" not in payload["error"] + assert "action" not in payload["error"] + @pytest.mark.asyncio async def test_generate_audio_chunks_discards_ref_audio_artifact_warmup_on_close(self, qwen3_tts_server): async def pcm_generator(): diff --git a/tests/entrypoints/openai_api/test_serving_speech_stream.py b/tests/entrypoints/openai_api/test_serving_speech_stream.py index 37e3509f5..7c553af8f 100644 --- a/tests/entrypoints/openai_api/test_serving_speech_stream.py +++ b/tests/entrypoints/openai_api/test_serving_speech_stream.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + import asyncio import base64 from types import SimpleNamespace @@ -44,7 +47,9 @@ def _build_test_app( speech_service._prepare_speech_generation = mocker.AsyncMock(return_value=("req-1", object(), {})) speech_service.forced_aligner_enabled = False - async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False, collect=None): + async def mock_generate_pcm_chunks( + _generator, _request_id, *, include_sample_rate=False, tts_params=None, collect=None + ): for chunk in (b"\x01\x02", b"\x03\x04\x05"): yield (chunk, 24000) if include_sample_rate else chunk @@ -244,6 +249,7 @@ class TestStreamingSpeechWebSocket: def test_streaming_multiple_binary_frames(self, mocker: MockerFixture): captured_requests = [] + captured_tts_params = [] speech_service = mocker.MagicMock(spec=OmniOpenAIServingSpeech) speech_service._generate_audio_bytes = mocker.AsyncMock(return_value=(b"", "audio/wav")) @@ -253,11 +259,12 @@ class TestStreamingSpeechWebSocket: async def mock_prepare_speech_generation(request): captured_requests.append(request) - return "req-stream", object(), {} + return "req-stream", object(), {"_qwen3_tts_effective_max_tokens": [192]} speech_service._prepare_speech_generation = mock_prepare_speech_generation - async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False): + async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False, tts_params=None): + captured_tts_params.append(tts_params) for chunk in (b"\x01\x02", b"\x03\x04\x05", b"\x06"): yield (chunk, 24000) if include_sample_rate else chunk @@ -302,6 +309,7 @@ class TestStreamingSpeechWebSocket: assert captured_requests[0].stream is True assert captured_requests[0].response_format == "pcm" assert captured_requests[0].initial_codec_chunk_frames == 12 + assert captured_tts_params == [{"_qwen3_tts_effective_max_tokens": [192]}] assert speech_service._generate_audio_bytes.await_count == 0 def test_word_timestamps_requires_configured_aligner(self, mocker: MockerFixture): @@ -344,7 +352,9 @@ class TestStreamingSpeechWebSocket: # The forced-aligner stage rides the same generator: its pooling output # is surfaced via the ``collect`` channel once the audio has streamed. - async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False, collect=None): + async def mock_generate_pcm_chunks( + _generator, _request_id, *, include_sample_rate=False, tts_params=None, collect=None + ): for chunk in (first_chunk, second_chunk): yield (chunk, 1000) if include_sample_rate else chunk if collect is not None: @@ -427,7 +437,9 @@ class TestStreamingSpeechWebSocket: speech_service.forced_aligner_enabled = True speech_service._prepare_speech_generation = mocker.AsyncMock(return_value=("req", object(), {})) - async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False, collect=None): + async def mock_generate_pcm_chunks( + _generator, _request_id, *, include_sample_rate=False, tts_params=None, collect=None + ): chunk = b"\x01" * 1000 yield (chunk, 1000) if include_sample_rate else chunk if collect is not None: @@ -608,7 +620,7 @@ class TestStreamingSpeechWebSocket: speech_service.engine_client.abort = mocker.AsyncMock() speech_service.forced_aligner_enabled = False - async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False): + async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False, tts_params=None): yield b"\x01\x02" raise RuntimeError("stream boom") @@ -633,6 +645,8 @@ class TestStreamingSpeechWebSocket: assert ws.receive_json() == { "type": "error", "message": "Generation failed for utterance 0, sentence 0: stream boom", + "partial_audio": True, + "action": "discard", } assert ws.receive_json() == { "type": "audio.done", @@ -702,7 +716,7 @@ class TestStreamingSpeechWebSocket: speech_service.engine_client.abort = mocker.AsyncMock() speech_service.forced_aligner_enabled = False - async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False): + async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False, tts_params=None): yield b"\x01\x02" speech_service._generate_pcm_chunks = mock_generate_pcm_chunks diff --git a/tests/entrypoints/openai_api/test_speech_usage.py b/tests/entrypoints/openai_api/test_speech_usage.py index 2fba10516..56b74c05b 100644 --- a/tests/entrypoints/openai_api/test_speech_usage.py +++ b/tests/entrypoints/openai_api/test_speech_usage.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Unit tests for speech token-usage accounting (issue #4646). These exercise the pure usage logic without a model/server: input-token @@ -205,6 +206,14 @@ def test_output_tokens_from_stage_metrics(): assert acc.total() == 77 +def test_output_tokens_preserves_stage0_finish_reason(): + acc = SpeechOutputTokenCounter() + res = _final_res(stage0_out=128) + res.metrics["stage_metrics"]["0"]["finish_reason"] = "length" + acc.observe(res) + assert acc.stage0_finish_reason == "length" + + def test_output_tokens_takes_max_across_stages(): acc = SpeechOutputTokenCounter() acc.observe(_final_res(stage0_out=42, stage1_out=0)) diff --git a/tests/entrypoints/openai_api/test_tts_adapter.py b/tests/entrypoints/openai_api/test_tts_adapter.py index e070bfebd..03faacfcc 100644 --- a/tests/entrypoints/openai_api/test_tts_adapter.py +++ b/tests/entrypoints/openai_api/test_tts_adapter.py @@ -32,7 +32,11 @@ from vllm_omni.entrypoints.openai.tts_adapters.moss_tts import ( MossTTSAdapter, MossTTSNanoAdapter, ) -from vllm_omni.entrypoints.openai.tts_adapters.qwen3_tts import Qwen3TTSAdapter +from vllm_omni.entrypoints.openai.tts_adapters.qwen3_tts import ( + QWEN3_TTS_EFFECTIVE_MAX_TOKENS_KEY, + Qwen3TTSAdapter, + Qwen3TTSCodecLimitError, +) from vllm_omni.model_executor.models.indextts2 import prompt_utils from vllm_omni.model_executor.models.indextts2.tokenizer_v2_5 import ( INDEXTTS25_TOKENIZER_FILE, @@ -162,6 +166,56 @@ def test_qwen3_tts_metadata(): assert issubclass(Qwen3TTSAdapter, ARTTSAdapter) +@pytest.mark.parametrize( + ("task_type", "text_tokens", "request_cap", "expected_cap"), + [ + ("Base", 0, None, 4096), + ("Base", 10, None, 192), + ("Base", 23, None, 276), + ("Base", 23, 128, 128), + ("Base", 23, 512, 512), + ("Base", 400, None, 4096), + ("CustomVoice", 10, None, 4096), + ("CustomVoice", 10, 256, 256), + ], +) +def test_qwen3_tts_applies_text_scaled_codec_safety_limit(task_type, text_tokens, request_cap, expected_cap, mocker): + server = mocker.Mock() + server._count_usage_text_tokens.return_value = text_tokens + adapter = Qwen3TTSAdapter(SpeechServingContext(server=server)) + stage_defaults = [SamplingParams(max_tokens=4096, min_tokens=2)] + request = OpenAICreateSpeechRequest( + input="test text", + task_type=task_type, + max_new_tokens=request_cap, + ) + prompt: dict[str, Any] = {"additional_information": {}} + + overridden = adapter.apply_sampling_overrides(stage_defaults, request, prompt) + + assert overridden[0].max_tokens == expected_cap + assert prompt["additional_information"][QWEN3_TTS_EFFECTIVE_MAX_TOKENS_KEY] == [expected_cap] + assert stage_defaults[0].max_tokens == 4096 + + +def test_qwen3_tts_rejects_only_length_finished_base_audio(mocker): + adapter = Qwen3TTSAdapter(SpeechServingContext(server=mocker.Mock())) + params = { + "task_type": ["Base"], + QWEN3_TTS_EFFECTIVE_MAX_TOKENS_KEY: [192], + } + + # EOS at the exact budget is valid; the token count alone is not a + # sufficient failure signal. + adapter.validate_generation(params, stage0_finish_reason="stop", output_tokens=192) + adapter.validate_generation(params, stage0_finish_reason=None, output_tokens=192) + + # Some frontends expose 191 decoded frames for max_new_tokens=192. The + # engine terminal reason still makes this an unambiguous limit failure. + with pytest.raises(Qwen3TTSCodecLimitError, match="191/192"): + adapter.validate_generation(params, stage0_finish_reason="length", output_tokens=191) + + def test_indextts_adapters_are_versioned(): assert resolve_adapter("indextts2") is IndexTTS2Adapter assert resolve_adapter("indextts2_5") is IndexTTS25Adapter diff --git a/vllm_omni/engine/stage_pool.py b/vllm_omni/engine/stage_pool.py index 5a8305803..e7c67cf9c 100644 --- a/vllm_omni/engine/stage_pool.py +++ b/vllm_omni/engine/stage_pool.py @@ -633,6 +633,15 @@ class StagePool: if callable(pop_native_text_metrics): native_text_metrics = pop_native_text_metrics(request_id) native_generation_tokens = native_text_metrics.get("num_generation_tokens") + finish_reason = next( + ( + str(reason) + for request_output in reversed(request_outputs) + for completion in reversed(getattr(request_output, "outputs", []) or []) + if (reason := getattr(completion, "finish_reason", None)) is not None + ), + None, + ) num_tokens_out = ( max(int(native_generation_tokens), 0) if isinstance(native_generation_tokens, int) and not isinstance(native_generation_tokens, bool) @@ -701,6 +710,7 @@ class StagePool: # happens inside the model runner and is not observable here. batch_size=1, replica_id=replica_id, + finish_reason=finish_reason, rx_decode_time_ms=0.0, rx_transfer_bytes=0, rx_in_flight_time_ms=0.0, diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index 61b2200cf..9ed81a862 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -56,6 +56,7 @@ from vllm_omni.entrypoints.openai.speech_usage import ( ) from vllm_omni.entrypoints.openai.tts_adapters import ( SpeechServingContext, + TTSGenerationError, all_tts_stage_keys, detect_tts_model_type, resolve_adapter, @@ -78,6 +79,7 @@ from vllm_omni.utils.speaker_cache import get_speaker_cache logger = init_logger(__name__) + _SPEECH_USAGE_INPUT_TOKENS_HEADER = "X-VLLM-OMNI-INPUT-TOKENS" _SPEECH_USAGE_OUTPUT_TOKENS_HEADER = "X-VLLM-OMNI-OUTPUT-TOKENS" _SPEECH_USAGE_TOTAL_TOKENS_HEADER = "X-VLLM-OMNI-TOTAL-TOKENS" @@ -750,6 +752,20 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): details = self._compute_speech_input_details(request, tts_params) return build_speech_usage(details, output_tokens) + def _validate_tts_generation( + self, + tts_params: dict[str, Any], + usage_acc: SpeechOutputTokenCounter, + ) -> None: + adapter = self._get_tts_adapter() + if adapter is None: + return + adapter.validate_generation( + tts_params, + stage0_finish_reason=usage_acc.stage0_finish_reason, + output_tokens=usage_acc.total(), + ) + @staticmethod def _build_speech_usage_headers(usage: SpeechTokenUsage | None) -> dict[str, str]: """Map speech usage into non-streaming response headers. @@ -2173,6 +2189,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): request_start_s: float | None = None, include_sample_rate: bool = False, usage_acc: SpeechOutputTokenCounter | None = None, + tts_params: dict[str, Any] | None = None, collect: dict | None = None, ): """Generate audio chunks for streaming response. @@ -2198,6 +2215,13 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): stream_start_s = request_start_s if request_start_s is not None else time.perf_counter() artifact_ready = False + # SSE supplies an accumulator for usage output. Raw-audio and WebSocket + # streams retain terminal metrics only when their model adapter needs + # generation validation. + adapter = self._get_tts_adapter() + if tts_params is not None and usage_acc is None and adapter is not None and adapter.validates_generation: + usage_acc = SpeechOutputTokenCounter() + try: async for res in generator: # Tally generated codec tokens for usage (reads per-stage metrics @@ -2275,6 +2299,11 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): # Audex contract: zero codec tokens must abort the stream, not # complete it cleanly with zero audio bytes. raise ValueError("Audex produced no audio (the thinker emitted zero or invalid codec tokens)") + # Check before committing the reference-audio artifact or logging + # success. Streaming protocols may already have emitted partial + # bytes, but they must terminate as an error rather than cleanly. + if tts_params is not None and usage_acc is not None: + self._validate_tts_generation(tts_params, usage_acc) self._mark_ref_audio_artifact_ready_for_request(request_id) artifact_ready = True total_ms = (time.perf_counter() - stream_start_s) * 1000.0 @@ -2359,6 +2388,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): ``input_tokens`` is computed from the request text + reference audio. """ usage_acc = SpeechOutputTokenCounter() + emitted_audio = False try: async for chunk in self._generate_audio_chunks( generator, @@ -2367,6 +2397,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): raw_request=raw_request, request_start_s=request_start_s, usage_acc=usage_acc, + tts_params=tts_params, ): payload = { "type": "speech.audio.delta", @@ -2374,6 +2405,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): "response_format": response_format, } data = json.dumps(payload, separators=(",", ":")) + emitted_audio = True yield f"event: speech.audio.delta\ndata: {data}\n\n" done_payload: dict[str, Any] = {"type": "speech.audio.done"} if request is not None: @@ -2385,16 +2417,19 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): except asyncio.CancelledError: raise except Exception as e: - payload = { - "type": "speech.audio.error", - "error": { - "message": str(e), - "type": "server_error", - "param": None, - "code": HTTPStatus.INTERNAL_SERVER_ERROR.value, - }, + error: dict[str, Any] = { + "message": str(e), + "type": "server_error", + "param": None, + "code": HTTPStatus.INTERNAL_SERVER_ERROR.value, } - data = json.dumps(payload, separators=(",", ":")) + if emitted_audio: + error.update(partial_audio=True, action="discard") + error_payload: dict[str, Any] = { + "type": "speech.audio.error", + "error": error, + } + data = json.dumps(error_payload, separators=(",", ":")) yield f"event: speech.audio.error\ndata: {data}\n\n" @staticmethod @@ -3109,7 +3144,13 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): return request_id, generator, tts_params async def _generate_pcm_chunks( - self, generator, request_id: str, *, include_sample_rate: bool = False, collect: dict | None = None + self, + generator, + request_id: str, + *, + include_sample_rate: bool = False, + tts_params: dict[str, Any] | None = None, + collect: dict | None = None, ): """Yield raw PCM byte chunks from the engine generator. @@ -3123,15 +3164,16 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): request_id, response_format="pcm", include_sample_rate=include_sample_rate, + tts_params=tts_params, collect=collect, ): yield chunk async def _iter_pcm_audio_bytes(self, request: OpenAICreateSpeechRequest): """Yield raw PCM bytes for a speech request as soon as chunks are decoded.""" - request_id, generator, _ = await self._prepare_speech_generation(request) + request_id, generator, tts_params = await self._prepare_speech_generation(request) try: - async for chunk in self._generate_pcm_chunks(generator, request_id): + async for chunk in self._generate_pcm_chunks(generator, request_id, tts_params=tts_params): yield chunk finally: self._discard_ref_audio_artifact_warmup(request_id) @@ -3203,6 +3245,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): if final_output is None: raise ValueError("No output generated from the model.") + self._validate_tts_generation(bytes_tts_params or {}, usage_acc) + # Extract audio from the audio-bearing res (not necessarily the last # yielded one, which may be the aligner's timestamps output). audio_source = audio_res if audio_res is not None else final_output @@ -3527,7 +3571,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): return error media_type = "audio/wav" if response_format == "wav" else "audio/pcm" - _, generator, _ = await self._prepare_speech_generation(request, request_id=request_id) + _, generator, raw_tts_params = await self._prepare_speech_generation(request, request_id=request_id) return StreamingResponse( self._generate_audio_chunks( generator, @@ -3535,6 +3579,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): response_format, raw_request=raw_request, request_start_s=request_start_s, + tts_params=raw_tts_params, ), media_type=media_type, ) @@ -3563,9 +3608,35 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): collect: dict = {} usage_box: list[SpeechTokenUsage] = [] - audio_bytes, media_type = await self._generate_audio_bytes( - request, request_id=request_id, usage_out=usage_box, collect=collect - ) + try: + audio_bytes, media_type = await self._generate_audio_bytes( + request, request_id=request_id, usage_out=usage_box, collect=collect + ) + except TTSGenerationError as error: + # An adapter can reject otherwise completed audio. Retry only + # retryable, stochastic, server-budgeted non-streaming + # requests: changing an explicit seed would violate + # reproducibility, and changing a caller-provided token limit + # would violate the requested budget. + if not error.retryable or request.seed is not None or request.max_new_tokens is not None: + raise + retry_seed = int(random_uuid()[:8], 16) + retry_request = request.model_copy(update={"seed": retry_seed}) + retry_request_id = f"speech-{random_uuid()}" + collect.clear() + usage_box.clear() + logger.warning( + "TTS request %s failed generation validation; retrying once as %s with seed=%d", + request_id, + retry_request_id, + retry_seed, + ) + audio_bytes, media_type = await self._generate_audio_bytes( + retry_request, + request_id=retry_request_id, + usage_out=usage_box, + collect=collect, + ) total_ms = (time.perf_counter() - request_start_s) * 1000.0 logger.info( "[SpeechE2E] request_id=%s stream=false status=ok total_ms=%.2f response_bytes=%d", @@ -3618,6 +3689,16 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): e, ) return self.create_error_response(e) + except TTSGenerationError as e: + total_ms = (time.perf_counter() - request_start_s) * 1000.0 + logger.error( + "[SpeechE2E] request_id=%s stream=%s status=invalid_generation total_ms=%.2f error=%s", + request_id, + bool(request.stream), + total_ms, + e, + ) + return self._diffusion_error_response(str(e), status_code=500) except Exception as e: total_ms = (time.perf_counter() - request_start_s) * 1000.0 logger.exception( diff --git a/vllm_omni/entrypoints/openai/serving_speech_stream.py b/vllm_omni/entrypoints/openai/serving_speech_stream.py index 2b71ad6c7..0f883f179 100644 --- a/vllm_omni/entrypoints/openai/serving_speech_stream.py +++ b/vllm_omni/entrypoints/openai/serving_speech_stream.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """WebSocket handler for streaming text input TTS. Accepts text incrementally via WebSocket, buffers it until input.done, and @@ -329,7 +332,7 @@ class OmniStreamingSpeechHandler: request_id = None try: if config.stream_audio: - request_id, generator, _ = await self._speech_service._prepare_speech_generation(request) + request_id, generator, tts_params = await self._speech_service._prepare_speech_generation(request) if config.word_timestamps: total_bytes = await self._stream_audio_with_alignments( websocket=websocket, @@ -339,9 +342,16 @@ class OmniStreamingSpeechHandler: utterance_index=utterance_index, sentence_index=sentence_index, language=config.language, + tts_params=tts_params, ) else: - async with aclosing(self._speech_service._generate_pcm_chunks(generator, request_id)) as stream: + async with aclosing( + self._speech_service._generate_pcm_chunks( + generator, + request_id, + tts_params=tts_params, + ) + ) as stream: async for chunk in stream: total_bytes += len(chunk) await websocket.send_bytes(chunk) @@ -367,6 +377,7 @@ class OmniStreamingSpeechHandler: await self._send_error( websocket, f"Generation failed for utterance {utterance_index}, sentence {sentence_index}: {e}", + partial_audio=total_bytes > 0, ) finally: try: @@ -392,6 +403,7 @@ class OmniStreamingSpeechHandler: utterance_index: int, sentence_index: int, language: str | None = None, + tts_params: dict | None = None, ) -> int: """Stream PCM as JSON ``audio.chunk`` frames, aligned per sentence. @@ -438,6 +450,7 @@ class OmniStreamingSpeechHandler: generator, request_id, include_sample_rate=True, + tts_params=tts_params, collect=collect, ) ) as stream: @@ -463,14 +476,15 @@ class OmniStreamingSpeechHandler: return total_bytes @staticmethod - async def _send_error(websocket: WebSocket, message: str) -> None: + async def _send_error(websocket: WebSocket, message: str, *, partial_audio: bool = False) -> None: """Send an error message to the client.""" try: - await websocket.send_json( - { - "type": "error", - "message": message, - } - ) + payload: dict[str, object] = { + "type": "error", + "message": message, + } + if partial_audio: + payload.update(partial_audio=True, action="discard") + await websocket.send_json(payload) except Exception: pass # Connection may already be closed; safe to ignore diff --git a/vllm_omni/entrypoints/openai/speech_usage.py b/vllm_omni/entrypoints/openai/speech_usage.py index 5e38405fd..cb7667b48 100644 --- a/vllm_omni/entrypoints/openai/speech_usage.py +++ b/vllm_omni/entrypoints/openai/speech_usage.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Token-usage accounting for the Speech (``/v1/audio/speech``) API. Why this module exists (issue #4646) @@ -160,6 +161,7 @@ class SpeechOutputTokenCounter: """ output_tokens: int = 0 + stage0_finish_reason: str | None = None def observe(self, res: Any) -> None: metrics = getattr(res, "metrics", None) @@ -168,6 +170,11 @@ class SpeechOutputTokenCounter: stage_metrics = metrics.get("stage_metrics") if not isinstance(stage_metrics, dict): return + stage0 = stage_metrics.get("0", stage_metrics.get(0)) + if isinstance(stage0, dict): + reason = stage0.get("finish_reason") + if reason is not None: + self.stage0_finish_reason = str(reason) for stage in stage_metrics.values(): if not isinstance(stage, dict): continue diff --git a/vllm_omni/entrypoints/openai/tts_adapters/__init__.py b/vllm_omni/entrypoints/openai/tts_adapters/__init__.py index e725e65cd..dcf8c4eaa 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/__init__.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/__init__.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Registry of TTS serving adapters, and the model detection built on it. Adapters register themselves by their ``name`` (the model-type discriminator) @@ -19,6 +20,7 @@ from vllm_omni.entrypoints.openai.tts_adapters.base import ( OutputPolicy, PreparedRequest, SpeechServingContext, + TTSGenerationError, TTSModelAdapter, ) @@ -126,6 +128,7 @@ __all__ = [ "OutputPolicy", "PreparedRequest", "SpeechServingContext", + "TTSGenerationError", "TTSModelAdapter", "TTS_ADAPTER_REGISTRY", "all_tts_model_types", diff --git a/vllm_omni/entrypoints/openai/tts_adapters/base.py b/vllm_omni/entrypoints/openai/tts_adapters/base.py index e5e186e18..4842acfd7 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/base.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/base.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Base contract for per-model TTS serving adapters. This package factors the per-model ``if self._tts_model_type == ...`` dispatch @@ -11,7 +12,7 @@ See the RFC for the full design (issue #4327). """ from abc import ABC, abstractmethod -from collections.abc import Callable +from collections.abc import Callable, Mapping from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, ClassVar @@ -69,6 +70,14 @@ def apply_max_new_tokens( return sampling_params_list +class TTSGenerationError(RuntimeError): + """A completed TTS generation that must not be returned as valid audio.""" + + def __init__(self, message: str, *, retryable: bool = False) -> None: + super().__init__(message) + self.retryable = retryable + + @dataclass class OutputPolicy: """How the orchestrator aggregates engine output for a model. @@ -153,6 +162,8 @@ class TTSModelAdapter(ABC): detect_priority: ClassVar[int] = 100 #: Serving backend: ``"ar"`` (engine_client) or ``"diffusion"``. backend: ClassVar[str] = "ar" + #: Whether streaming entrypoints must retain terminal metrics for validation. + validates_generation: ClassVar[bool] = False #: Whether the model consumes ``request.speed`` in its native parameters. native_speed_control: ClassVar[bool] = False @@ -233,6 +244,21 @@ class TTSModelAdapter(ABC): """ return sampling_params_list + def validate_generation( + self, + tts_params: Mapping[str, object], + *, + stage0_finish_reason: str | None, + output_tokens: int, + ) -> None: + """Reject a completed generation that is invalid for this model. + + Adapters that override this hook must also set + :attr:`validates_generation` so streaming entrypoints retain the + terminal metrics needed by the validation without charging that cost + to unrelated TTS models. + """ + async def warmup(self) -> None: return @@ -304,5 +330,6 @@ __all__ = [ "OutputPolicy", "PreparedRequest", "SpeechServingContext", + "TTSGenerationError", "TTSModelAdapter", ] diff --git a/vllm_omni/entrypoints/openai/tts_adapters/qwen3_tts.py b/vllm_omni/entrypoints/openai/tts_adapters/qwen3_tts.py index d29394ec7..c2fd4faf8 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/qwen3_tts.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/qwen3_tts.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Qwen3-TTS serving adapter.""" from collections.abc import Mapping @@ -12,7 +13,7 @@ from vllm_omni.entrypoints.openai.tts_adapters.base import ( DEFAULT_TTS_LANGUAGES, ARTTSAdapter, PreparedRequest, - apply_max_new_tokens, + TTSGenerationError, ) from vllm_omni.entrypoints.openai.tts_adapters.capabilities import load_precomputed_speakers from vllm_omni.utils.speaker_cache import validate_qwen3_tts_profile @@ -22,11 +23,23 @@ if TYPE_CHECKING: logger = init_logger(__name__) +QWEN3_TTS_EFFECTIVE_MAX_TOKENS_KEY = "_qwen3_tts_effective_max_tokens" +_MIN_CODEC_FRAMES = 192 +_MAX_CODEC_FRAMES_PER_TEXT_TOKEN = 12 + + +class Qwen3TTSCodecLimitError(TTSGenerationError): + """Qwen3-TTS Base exhausted its codec budget without emitting EOS.""" + + def __init__(self, message: str) -> None: + super().__init__(message, retryable=True) + @register_tts_adapter class Qwen3TTSAdapter(ARTTSAdapter): """Adapter for Qwen3-TTS (AR ``engine_client`` backend).""" + validates_generation = True stage_keys = frozenset({"qwen3_tts"}) name = "qwen3_tts" @@ -226,4 +239,95 @@ class Qwen3TTSAdapter(ARTTSAdapter): prompt: dict[str, Any] | None = None, request_id: str | None = None, ) -> list: - return apply_max_new_tokens(sampling_params_list, request) + """Apply a text-scaled safety ceiling to Base codec generation. + + Qwen3-TTS can rarely enter a repetitive state in which codec EOS is no + longer reachable through top-k sampling. A fixed 4096-frame ceiling + turns that into several minutes of unusable audio. Bound the default + Base-task budget by text length, while preserving an explicit caller + ``max_new_tokens`` override and the configured budget for other tasks. + """ + import copy + + del request_id + server = self.ctx.server + # Only scalar fields on stage 0 are changed below. Shallow-copy each + # stage so the shared defaults stay immutable without deep-copying the + # complete sampling configuration on every request. + sampling_params_list = [copy.copy(params) for params in sampling_params_list] + configured_cap = getattr(sampling_params_list[0], "max_tokens", None) + task_type = request.task_type or "CustomVoice" + text_tokens = None + dynamic_cap = None + effective_cap: int | None + + if request.max_new_tokens is not None: + # An explicit request budget is an opt-out from the automatic + # ceiling. It remains an upper bound, and a length finish is still + # surfaced as an incomplete generation rather than valid audio. + effective_cap = int(request.max_new_tokens) + elif task_type == "Base": + counted_text_tokens = server._count_usage_text_tokens(request.input) + if counted_text_tokens > 0: + text_tokens = counted_text_tokens + dynamic_cap = max(_MIN_CODEC_FRAMES, text_tokens * _MAX_CODEC_FRAMES_PER_TEXT_TOKEN) + effective_cap = min(dynamic_cap, int(configured_cap)) if configured_cap is not None else dynamic_cap + else: + # Token counting is best-effort. If the tokenizer is missing or + # rejects the input, preserve the configured budget instead of + # truncating an otherwise valid request at the minimum ceiling. + effective_cap = int(configured_cap) if configured_cap is not None else None + else: + effective_cap = int(configured_cap) if configured_cap is not None else None + + if effective_cap is not None: + effective_cap = max(1, effective_cap) + sampling_params_list[0].max_tokens = effective_cap + sampling_params_list[0].min_tokens = min( + int(getattr(sampling_params_list[0], "min_tokens", 0) or 0), + effective_cap, + ) + + if isinstance(prompt, dict): + additional_information = prompt.get("additional_information") + if isinstance(additional_information, dict) and effective_cap is not None: + additional_information[QWEN3_TTS_EFFECTIVE_MAX_TOKENS_KEY] = [effective_cap] + + logger.debug( + "Qwen3-TTS codec budget: task_type=%s text_tokens=%s dynamic_cap=%s " + "configured_cap=%s request_cap=%s effective_cap=%s", + task_type, + text_tokens, + dynamic_cap, + configured_cap, + request.max_new_tokens, + effective_cap, + ) + return sampling_params_list + + def validate_generation( + self, + tts_params: Mapping[str, object], + *, + stage0_finish_reason: str | None, + output_tokens: int, + ) -> None: + if QWEN3_TTS_EFFECTIVE_MAX_TOKENS_KEY not in tts_params: + return + task_type = tts_params.get("task_type") + if isinstance(task_type, (list, tuple)): + task_type = task_type[0] if task_type else None + if task_type != "Base" or stage0_finish_reason != "length": + return + + raw_limit = tts_params.get(QWEN3_TTS_EFFECTIVE_MAX_TOKENS_KEY) + if isinstance(raw_limit, (list, tuple)): + raw_limit = raw_limit[0] if raw_limit else None + try: + limit = int(raw_limit) if isinstance(raw_limit, (str, bytes, bytearray, int, float)) else 0 + except (TypeError, ValueError): + limit = 0 + raise Qwen3TTSCodecLimitError( + "Qwen3-TTS Base did not emit codec EOS before its token budget " + f"({output_tokens}/{limit} codec tokens); the generated audio is incomplete." + ) diff --git a/vllm_omni/metrics/stats.py b/vllm_omni/metrics/stats.py index 51d424508..9e9191eaa 100644 --- a/vllm_omni/metrics/stats.py +++ b/vllm_omni/metrics/stats.py @@ -45,6 +45,7 @@ class StageRequestStats: stage_id: int | None = None replica_id: int | None = None final_output_type: str | None = None + finish_reason: str | None = None request_id: str | None = None postprocess_time_ms: float = 0.0 diffusion_metrics: dict[str, float] = None @@ -126,6 +127,7 @@ STAGE_EXCLUDE = { "rx_decode_time_ms", "rx_in_flight_time_ms", "final_output_type", + "finish_reason", "pipeline_timings", } TRANSFER_EXCLUDE = {"from_stage", "to_stage", "request_id", "used_shm"} @@ -473,6 +475,7 @@ class OrchestratorAggregator: current = { "stage_id": sid, "final_output_type": evt.final_output_type, + "finish_reason": evt.finish_reason, defs.NUM_TOKENS_IN: int(evt.num_tokens_in), defs.NUM_TOKENS_OUT: int(evt.num_tokens_out), defs.STAGE_GEN_TIME_MS: float(evt.stage_gen_time_ms), @@ -498,6 +501,8 @@ class OrchestratorAggregator: current[defs.NUM_TOKENS_IN] = int(current.get(defs.NUM_TOKENS_IN, 0)) + int(evt.num_tokens_in) current[defs.NUM_TOKENS_OUT] = int(current.get(defs.NUM_TOKENS_OUT, 0)) + int(evt.num_tokens_out) + if evt.finish_reason is not None: + current["finish_reason"] = evt.finish_reason current[defs.STAGE_GEN_TIME_MS] = float(current.get(defs.STAGE_GEN_TIME_MS, 0.0)) + float(evt.stage_gen_time_ms) current[defs.POSTPROCESS_TIME_MS] = float(current.get(defs.POSTPROCESS_TIME_MS, 0.0)) + float( evt.postprocess_time_ms