diff --git a/docs/serving/speech_api.md b/docs/serving/speech_api.md index 804439e40..c11bbd1e9 100644 --- a/docs/serving/speech_api.md +++ b/docs/serving/speech_api.md @@ -122,7 +122,8 @@ Content-Type: application/json | `max_new_tokens` | integer | 2048 | Maximum tokens to generate | | `initial_codec_chunk_frames` | integer | null | Per-request initial chunk size override for TTFA tuning. When null, IC is computed dynamically based on server load. | | `non_streaming_mode` | bool | null | Qwen3-TTS prompt construction mode override. Does not affect HTTP response streaming or async-chunk pipelining. When null, Qwen3-TTS uses model defaults: Base=false, CustomVoice/VoiceDesign=true. | -| `stream` | bool | false | Stream raw PCM chunks as they are decoded (requires `response_format="pcm"`) | +| `stream_format` | string | null | Streaming output format. `"audio"` streams raw audio bytes as they are decoded; `"sse"` streams OpenAI `speech.audio.*` Server-Sent Events. Omit (null) for a single non-streaming binary response. See [Response Format](#response-format). | +| `stream` | bool | false | Legacy streaming switch, equivalent to `stream_format="audio"`. Requires `response_format="pcm"` or `"wav"`. Speed adjustment is not supported when streaming. | **Supported languages:** Only applicable to Qwen3-TTS. Derived from the model configuration (`talker_config.codec_language_id` in the checkpoint's `config.json`), plus `Auto`, which is always accepted. Official Qwen3-TTS checkpoints support: Auto, Chinese, English, Japanese, Korean, German, French, Russian, Portuguese, Spanish, Italian. @@ -136,7 +137,60 @@ Content-Type: application/json ### Response Format -Returns binary audio data with appropriate `Content-Type` header (e.g., `audio/wav`). +The response shape depends on the streaming parameters: + +**Non-streaming (default).** With `stream=false` and no `stream_format`, returns the +complete clip as binary audio data with an appropriate `Content-Type` header (e.g. +`audio/wav`). The raw-bytes body has no JSON carrier, so no `usage` is reported. + +**Raw audio stream** (`stream=true` or `stream_format="audio"`). Streams raw +audio bytes (PCM or WAV) as they are decoded. + +Both streaming modes (`stream_format="audio"` and `"sse"`) require +`response_format="pcm"` or `"wav"`, and `speed` must be `1.0` (or omitted). + +**SSE stream** (`stream_format="sse"`). Streams [OpenAI `speech.audio.*` Server-Sent +Events](https://platform.openai.com/docs/api-reference/audio-streaming). Set +`stream_format="sse"` and leave the legacy `stream` unset: `stream=true` takes +precedence and forces the raw-audio path, so the response would be audio bytes with +no `speech.audio.*` events (and no terminal `usage`). Each event has an `event:` +line and a JSON `data:` line: + +- `speech.audio.delta` — a base64 audio chunk: + + ```json + { "type": "speech.audio.delta", "audio": "", "response_format": "pcm" } + ``` + +- `speech.audio.done` — terminal event, carrying token `usage`: + + ```json + { + "type": "speech.audio.done", + "usage": { + "input_tokens": 119, + "output_tokens": 77, + "total_tokens": 196, + "input_token_details": { "text_tokens": 18, "audio_tokens": 101 } + } + } + ``` + +- `speech.audio.error` — emitted instead of `speech.audio.done` if generation fails: + + ```json + { "type": "speech.audio.error", "error": { "message": "...", "type": "server_error", "param": null, "code": 500 } } + ``` + +The `usage` object on `speech.audio.done` is the same shape returned per item by the +[batch endpoint](#batch-speech-generation): + +- `input_tokens` = `text_tokens` + `audio_tokens` + - `text_tokens`: tokens of the synthesized text (`input` plus `instructions`) + - `audio_tokens`: reference-audio codec frames, non-zero only for in-context + voice cloning (Base task); `0` for CustomVoice/VoiceDesign or x-vector-only +- `output_tokens`: generated codec tokens +- `total_tokens` = `input_tokens` + `output_tokens` ### Voices Endpoint @@ -490,7 +544,13 @@ Each item in the `items` array requires only `input` (the text). All other field "index": 0, "status": "success", "audio_data": "", - "media_type": "audio/wav" + "media_type": "audio/wav", + "usage": { + "input_tokens": 119, + "output_tokens": 77, + "total_tokens": 196, + "input_token_details": { "text_tokens": 18, "audio_tokens": 101 } + } }, { "index": 1, @@ -504,6 +564,18 @@ Each item in the `items` array requires only `input` (the text). All other field } ``` +Each successful item carries a `usage` object (errored items omit it): + +- `input_tokens` = `text_tokens` + `audio_tokens` + - `text_tokens`: tokens of the synthesized text (`input` plus `instructions`) + - `audio_tokens`: reference-audio codec frames, non-zero only for in-context + voice cloning (Base task); `0` for CustomVoice/VoiceDesign or x-vector-only +- `output_tokens`: generated codec tokens +- `total_tokens` = `input_tokens` + `output_tokens` + +This is the same `usage` object emitted on the terminal `speech.audio.done` event +of the single endpoint's [SSE stream](#response-format) (`stream_format="sse"`). + ### Examples **Basic batch with shared defaults:** diff --git a/tests/entrypoints/openai_api/test_serving_speech.py b/tests/entrypoints/openai_api/test_serving_speech.py index 1c85d7fa0..a2e31dcd7 100644 --- a/tests/entrypoints/openai_api/test_serving_speech.py +++ b/tests/entrypoints/openai_api/test_serving_speech.py @@ -2535,6 +2535,61 @@ def test_api_server_create_speech_batch_without_handler_returns_404(mocker: Mock ) +def test_api_server_create_speech_batch_omits_null_fields(mocker: MockerFixture): + # The batch response must omit optional null fields rather than serialize them + # as null (issue #4646 follow-up): errored items drop usage/audio_data/media_type, + # successful items drop error. This is the shape documented in speech_api.md. + from vllm_omni.entrypoints.openai.protocol.audio import ( + BatchSpeechResponse, + SpeechBatchItemResult, + SpeechInputTokenDetails, + SpeechTokenUsage, + ) + + handler = mocker.MagicMock() + handler.create_speech_batch = mocker.AsyncMock( + return_value=BatchSpeechResponse( + id="speech-batch-test", + results=[ + SpeechBatchItemResult( + index=0, + status="success", + audio_data="YWJj", + media_type="audio/wav", + usage=SpeechTokenUsage( + input_tokens=119, + output_tokens=77, + total_tokens=196, + input_token_details=SpeechInputTokenDetails(text_tokens=18, audio_tokens=101), + ), + ), + SpeechBatchItemResult(index=1, status="error", error="Input text cannot be empty"), + ], + total=2, + succeeded=1, + failed=1, + ) + ) + raw_request = _make_api_server_request(handler, path="/v1/audio/speech/batch") + request = BatchSpeechRequest(items=[SpeechBatchItem(input="hi"), SpeechBatchItem(input="")]) + + response = asyncio.run(api_server_module.create_speech_batch(request, raw_request)) + + assert isinstance(response, JSONResponse) + assert response.status_code == 200 + body = json.loads(response.body) + success, errored = body["results"][0], body["results"][1] + # Successful item carries usage and drops the null `error`. + assert success["usage"]["total_tokens"] == 196 + assert success["usage"]["input_token_details"] == {"text_tokens": 18, "audio_tokens": 101} + assert "error" not in success + # Errored item drops usage/audio_data/media_type instead of serializing null. + assert "usage" not in errored + assert "audio_data" not in errored + assert "media_type" not in errored + assert errored["error"] == "Input text cannot be empty" + + def test_api_server_create_audio_generate_without_handler_returns_404(mocker: MockerFixture): fake_base = _patch_api_server_base(mocker) raw_request = _make_api_server_request(None, path="/v1/audio/generate") diff --git a/tests/entrypoints/openai_api/test_speech_usage.py b/tests/entrypoints/openai_api/test_speech_usage.py new file mode 100644 index 000000000..8211dab88 --- /dev/null +++ b/tests/entrypoints/openai_api/test_speech_usage.py @@ -0,0 +1,220 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Unit tests for speech token-usage accounting (issue #4646). + +These exercise the pure usage logic without a model/server: input-token +breakdown (text + reference-audio gating) and output-token accumulation. +""" + +from dataclasses import dataclass + +import pytest + +from vllm_omni.entrypoints.openai.speech_usage import ( + SpeechOutputTokenCounter, + build_speech_usage, + gate_audio_tokens, + qwen3_tts_input_token_details, +) + + +# A deterministic stand-in for a real tokenizer: 1 token per whitespace word. +def word_count(text: str) -> int: + return len(text.split()) + + +# --- input breakdown: gating ------------------------------------------------ + + +def test_gate_audio_tokens_base_in_context_counts_ref_frames(): + assert ( + gate_audio_tokens(task_type="Base", x_vector_only_mode=False, icl_mode_override=None, ref_code_length=120) + == 120 + ) + + +def test_gate_audio_tokens_customvoice_is_zero(): + assert ( + gate_audio_tokens( + task_type="CustomVoice", x_vector_only_mode=False, icl_mode_override=None, ref_code_length=120 + ) + == 0 + ) + + +def test_gate_audio_tokens_x_vector_only_is_zero_even_with_ref(): + # x-vector cloning inserts NO codec frames, so audio_tokens must be 0 + # even though ref_code_length is populated. + assert ( + gate_audio_tokens(task_type="Base", x_vector_only_mode=True, icl_mode_override=None, ref_code_length=120) == 0 + ) + + +def test_gate_audio_tokens_icl_override_wins(): + assert ( + gate_audio_tokens(task_type="Base", x_vector_only_mode=False, icl_mode_override=False, ref_code_length=120) == 0 + ) + assert ( + gate_audio_tokens(task_type="Base", x_vector_only_mode=True, icl_mode_override=True, ref_code_length=80) == 80 + ) + + +def test_gate_audio_tokens_icl_override_ignored_for_non_base(): + # icl_mode_override is honored only for the Base task; CustomVoice stays 0 + # even if such a flag leaks into its params (mirrors the estimator). + assert ( + gate_audio_tokens( + task_type="CustomVoice", x_vector_only_mode=False, icl_mode_override=True, ref_code_length=120 + ) + == 0 + ) + + +def test_gate_audio_tokens_bad_ref_length_is_zero(): + assert ( + gate_audio_tokens(task_type="Base", x_vector_only_mode=False, icl_mode_override=None, ref_code_length="oops") + == 0 + ) + + +# --- input breakdown: Qwen3-TTS reproduces issue #4646 observation ---------- + + +def test_customvoice_text_scales_with_input_no_audio(): + short = qwen3_tts_input_token_details( + input_text="hi there", + instructions=None, + tts_params={"task_type": ["CustomVoice"]}, + count_text_tokens=word_count, + ) + long = qwen3_tts_input_token_details( + input_text="hi there this is a much longer sentence", + instructions=None, + tts_params={"task_type": ["CustomVoice"]}, + count_text_tokens=word_count, + ) + assert short.audio_tokens == 0 and long.audio_tokens == 0 + assert long.text_tokens > short.text_tokens # scales with input + + +def test_base_icl_text_tracks_input_audio_tracks_ref(): + # The core #4646 fix: text_tokens reflect the input (not dropped), and + # audio_tokens track the reference audio independently. + params = {"task_type": ["Base"], "x_vector_only_mode": [False], "ref_code_length": [100]} + short = qwen3_tts_input_token_details( + input_text="hi", instructions=None, tts_params=params, count_text_tokens=word_count + ) + long = qwen3_tts_input_token_details( + input_text="hi there friend", instructions=None, tts_params=params, count_text_tokens=word_count + ) + # text now varies with input (the bug was that it was dropped for Base) + assert long.text_tokens > short.text_tokens + # audio fixed by the reference, independent of input text + assert short.audio_tokens == long.audio_tokens == 100 + + +def test_instructions_count_toward_text(): + without = qwen3_tts_input_token_details( + input_text="hello world", + instructions=None, + tts_params={"task_type": ["CustomVoice"]}, + count_text_tokens=word_count, + ) + with_instr = qwen3_tts_input_token_details( + input_text="hello world", + instructions="speak slowly and warmly", + tts_params={"task_type": ["CustomVoice"]}, + count_text_tokens=word_count, + ) + assert with_instr.text_tokens == without.text_tokens + word_count("speak slowly and warmly") + + +# --- build usage ------------------------------------------------------------ + + +def test_build_speech_usage_aggregates_and_totals(): + details = qwen3_tts_input_token_details( + input_text="hi there", + instructions=None, + tts_params={"task_type": ["Base"], "x_vector_only_mode": [False], "ref_code_length": [100]}, + count_text_tokens=word_count, + ) + usage = build_speech_usage(details, output_tokens=250) + assert usage.input_tokens == details.text_tokens + 100 + assert usage.output_tokens == 250 + assert usage.total_tokens == usage.input_tokens + 250 + assert usage.input_token_details.text_tokens == 2 + assert usage.input_token_details.audio_tokens == 100 + + +def test_batch_item_result_serializes_usage(): + # Guards the batch per-item usage carrier (and that it is optional/None on error). + from vllm_omni.entrypoints.openai.protocol.audio import SpeechBatchItemResult + + details = qwen3_tts_input_token_details( + input_text="hello world", + instructions=None, + tts_params={"task_type": ["Base"], "x_vector_only_mode": [False], "ref_code_length": [50]}, + count_text_tokens=word_count, + ) + item = SpeechBatchItemResult( + index=0, + status="success", + audio_data="x", + media_type="audio/wav", + usage=build_speech_usage(details, output_tokens=30), + ) + dumped = item.model_dump() + assert dumped["usage"]["input_tokens"] == 2 + 50 + assert dumped["usage"]["output_tokens"] == 30 + assert dumped["usage"]["input_token_details"] == {"text_tokens": 2, "audio_tokens": 50} + # Errored items carry no usage. + assert SpeechBatchItemResult(index=1, status="error", error="boom").model_dump()["usage"] is None + + +# --- output token accumulation ---------------------------------------------- +# The generated codec-token count is read from per-stage pipeline metrics on the +# final output: res.metrics["stage_metrics"][]["num_tokens_out"]. Only the +# AR stage (0) is non-zero; code2wav (stage 1) reports 0. + + +@dataclass +class _FakeRes: + metrics: dict + + +def _final_res(stage0_out, stage1_out=0): + return _FakeRes( + metrics={ + "stage_metrics": { + "0": {"num_tokens_in": 110, "num_tokens_out": stage0_out}, + "1": {"num_tokens_in": 0, "num_tokens_out": stage1_out}, + } + } + ) + + +def test_output_tokens_from_stage_metrics(): + acc = SpeechOutputTokenCounter() + # Intermediate outputs carry empty metrics -> contribute nothing. + acc.observe(_FakeRes(metrics={})) + acc.observe(_FakeRes(metrics={})) + # Final output carries stage metrics; stage-0 num_tokens_out is the count. + acc.observe(_final_res(stage0_out=77)) + assert acc.total() == 77 + + +def test_output_tokens_takes_max_across_stages(): + acc = SpeechOutputTokenCounter() + acc.observe(_final_res(stage0_out=42, stage1_out=0)) + assert acc.total() == 42 + + +def test_output_tokens_zero_without_metrics(): + acc = SpeechOutputTokenCounter() + acc.observe(_FakeRes(metrics={})) + acc.observe(_FakeRes(metrics={"stage_metrics": {}})) + assert acc.total() == 0 + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-v"])) diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 3560f26b8..07e5d2bfa 100755 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -1317,7 +1317,10 @@ async def create_speech_batch(request: BatchSpeechRequest, raw_request: Request) result = await handler.create_speech_batch(request) if isinstance(result, ErrorResponse): return _error_response_to_json_response(result) - return JSONResponse(content=result.model_dump()) + # exclude_none so optional per-item fields are omitted rather than + # serialized as null: errored items drop `usage`/`audio_data`/`media_type`, + # successful items drop `error`. Matches the documented batch response shape. + return JSONResponse(content=result.model_dump(exclude_none=True)) except (EngineGenerateError, EngineDeadError) as exc: return _create_engine_error_json_response(raw_request, exc) except ValueError as e: diff --git a/vllm_omni/entrypoints/openai/protocol/audio.py b/vllm_omni/entrypoints/openai/protocol/audio.py index cb2836b2f..48a646edd 100644 --- a/vllm_omni/entrypoints/openai/protocol/audio.py +++ b/vllm_omni/entrypoints/openai/protocol/audio.py @@ -421,12 +421,65 @@ class BatchSpeechRequest(BaseModel): non_streaming_mode: bool | None = None +class SpeechInputTokenDetails(BaseModel): + """Per-modality breakdown of the speech request's *input* tokens. + + The aggregate ``input_tokens`` on :class:`SpeechTokenUsage` is the sum of + these. We surface the split (rather than one opaque number) for the same + reason OpenAI's realtime/chat usage does: text and audio inputs are billed + and reasoned about differently, and folding them together is misleading. + + Fields: + text_tokens: Tokens of the text to synthesize. This is ``input`` plus + ``instructions`` (style/emotion prompt), because both are tokenized + into the model prefill. This is the number that should scale with + how much text the caller asked to speak. + audio_tokens: Reference-audio codec frames used as voice-cloning + conditioning. NON-ZERO only when in-context voice cloning is + actually active (Qwen3-TTS ``task_type='Base'`` ICL). It is 0 for + CustomVoice/VoiceDesign and for x-vector-only cloning, because those + paths put no reference codec frames into the prefill. See issue + #4646: this is the value that previously leaked into ``prompt_tokens`` + and made Base usage look independent of the input text. + """ + + text_tokens: int = 0 + audio_tokens: int = 0 + + +class SpeechTokenUsage(BaseModel): + """Token usage for a speech (TTS) request. + + Field naming follows OpenAI's documented ``speech.audio.done`` event + (``input_tokens``/``output_tokens``/``total_tokens``), NOT chat's + ``prompt_tokens``/``completion_tokens``. + + input_tokens = text_tokens + audio_tokens (see SpeechInputTokenDetails) + output_tokens = generated codec/audio tokens (stage-0 decode steps) + total_tokens = input_tokens + output_tokens + + IMPORTANT: ``input_tokens`` is computed from the *semantic* request inputs + (tokenized text + reference-audio frames), NOT from ``len(prompt_token_ids)``. + For staged TTS models that engine prompt is a ``[1] * prefill_len`` + placeholder whose length mirrors the model prefill, so it is not a faithful + input-token count (issue #4646). + """ + + input_tokens: int = 0 + output_tokens: int = 0 + total_tokens: int = 0 + input_token_details: SpeechInputTokenDetails = Field(default_factory=SpeechInputTokenDetails) + + class SpeechBatchItemResult(BaseModel): index: int status: Literal["success", "error"] audio_data: str | None = None media_type: str | None = None error: str | None = None + # Per-item token usage (input text + reference-audio conditioning, and + # generated audio tokens). None when the item errored before generation. + usage: SpeechTokenUsage | None = None class BatchSpeechResponse(BaseModel): diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index 1a3f13391..1b31d9b9d 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -43,6 +43,13 @@ from vllm_omni.entrypoints.openai.protocol.audio import ( OpenAICreateSpeechRequest, SpeechBatchItem, SpeechBatchItemResult, + SpeechInputTokenDetails, + SpeechTokenUsage, +) +from vllm_omni.entrypoints.openai.speech_usage import ( + SpeechOutputTokenCounter, + build_speech_usage, + qwen3_tts_input_token_details, ) from vllm_omni.entrypoints.openai.tts_adapters import ( SpeechServingContext, @@ -895,6 +902,69 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): logger.warning("Failed to estimate TTS prompt length, using fallback 2048: %s", e) return 2048 + def _get_usage_text_tokenizer(self): + """Return a text tokenizer for counting input-text usage tokens. + + Prefer the per-model tokenizer already loaded for prompt-length + estimation (`_tts_tokenizer`, which is the *correct* tokenizer for the + active model). Fall back to a lazily-loaded, cached generic tokenizer + for models that never populate `_tts_tokenizer`. Returns None if no + tokenizer can be obtained (usage then reports text_tokens=0). + """ + if self._tts_tokenizer is not None: + return self._tts_tokenizer + if getattr(self, "_usage_text_tokenizer", None) is None: + try: + from transformers import AutoTokenizer + + self._usage_text_tokenizer = AutoTokenizer.from_pretrained( + self.engine_client.model_config.model, trust_remote_code=True + ) + except Exception as e: # pragma: no cover - environment dependent + logger.warning("Usage: could not load a text tokenizer (%s); text_tokens will be 0", e) + self._usage_text_tokenizer = None + return self._usage_text_tokenizer + + def _count_usage_text_tokens(self, text: str) -> int: + """Token count of `text` using the model's text tokenizer (0 on failure).""" + if not text: + return 0 + tok = self._get_usage_text_tokenizer() + if tok is None: + return 0 + try: + return len(tok(text, padding=False)["input_ids"]) + except Exception: + return 0 + + def _compute_speech_input_details( + self, request: OpenAICreateSpeechRequest, tts_params: dict[str, Any] + ) -> SpeechInputTokenDetails: + """Input-token breakdown (text + reference-audio) for a speech request. + + Counts `input` (+ `instructions`) as text tokens, and reference-audio + codec frames as audio tokens *only* when in-context voice cloning is + active (see `qwen3_tts_input_token_details` / `gate_audio_tokens`). The + audio gating reads Qwen3-TTS `tts_params` conventions; other TTS models + do not set those keys, so they degrade cleanly to text-only counts. + """ + return qwen3_tts_input_token_details( + input_text=request.input, + instructions=request.instructions, + tts_params=tts_params or {}, + count_text_tokens=self._count_usage_text_tokens, + ) + + def _build_speech_usage( + self, + request: OpenAICreateSpeechRequest, + tts_params: dict[str, Any], + output_tokens: int, + ) -> SpeechTokenUsage: + """Assemble the full usage object (input breakdown + generated tokens).""" + details = self._compute_speech_input_details(request, tts_params) + return build_speech_usage(details, output_tokens) + def _estimate_fish_ref_code_len(self, ref_audio: object) -> int | None: """Estimate Fish Speech semantic token length from raw reference audio.""" from vllm_omni.model_executor.models.fish_speech.dac_utils import ( @@ -2635,6 +2705,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): raw_request: Request | None = None, request_start_s: float | None = None, include_sample_rate: bool = False, + usage_acc: SpeechOutputTokenCounter | None = None, ): """Generate audio chunks for streaming response. @@ -2661,6 +2732,10 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): try: async for res in generator: + # Tally generated codec tokens for usage (reads per-stage metrics + # off the final output; a cheap early-return on every other res). + if usage_acc is not None: + usage_acc.observe(res) audio_output, audio_key = self._extract_audio_output(res) if audio_key is None: continue @@ -2786,6 +2861,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): response_format: str = "pcm", raw_request: Request | None = None, request_start_s: float | None = None, + request: OpenAICreateSpeechRequest | None = None, + tts_params: dict[str, Any] | None = None, ): """Generate OpenAI-style SSE events with base64 audio deltas. @@ -2793,7 +2870,15 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): carries the base64 chunk in ``audio`` (not ``delta`` — that is the Realtime API ``response.audio.delta`` convention, a different event). See https://platform.openai.com/docs/api-reference/audio-streaming. + + The terminal ``speech.audio.done`` event carries a ``usage`` object + (``input_tokens``/``output_tokens``/``total_tokens`` + a per-modality + ``input_token_details`` breakdown), matching OpenAI's documented + ``speech.audio.done`` schema. ``output_tokens`` is accumulated from the + stage-0 deltas as they stream (see ``SpeechOutputTokenCounter``); + ``input_tokens`` is computed from the request text + reference audio. """ + usage_acc = SpeechOutputTokenCounter() try: async for chunk in self._generate_audio_chunks( generator, @@ -2801,6 +2886,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): response_format, raw_request=raw_request, request_start_s=request_start_s, + usage_acc=usage_acc, ): payload = { "type": "speech.audio.delta", @@ -2809,7 +2895,12 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): } data = json.dumps(payload, separators=(",", ":")) yield f"event: speech.audio.delta\ndata: {data}\n\n" - done = json.dumps({"type": "speech.audio.done"}, separators=(",", ":")) + done_payload: dict[str, Any] = {"type": "speech.audio.done"} + if request is not None: + # Streaming path: output_tokens = sum of stage-0 deltas. + usage = self._build_speech_usage(request, tts_params or {}, usage_acc.total()) + done_payload["usage"] = usage.model_dump() + done = json.dumps(done_payload, separators=(",", ":")) yield f"event: speech.audio.done\ndata: {done}\n\n" except asyncio.CancelledError: raise @@ -3650,8 +3741,13 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): request: OpenAICreateSpeechRequest, base64_encode: bool = False, request_id: str | None = None, + usage_out: list[SpeechTokenUsage] | None = None, ) -> tuple[bytes | str, str]: - request_id, generator, _ = await self._prepare_speech_generation(request, request_id=request_id) + # ``usage_out`` is an opt-in output channel: when a list is passed, the + # computed SpeechTokenUsage is appended to it. The return stays a + # 2-tuple so existing callers (and their test mocks) are unaffected; + # only the batch path, which surfaces per-item usage, opts in. + request_id, generator, bytes_tts_params = await self._prepare_speech_generation(request, request_id=request_id) artifact_ready = False try: @@ -3664,8 +3760,12 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): moss_sample_rate: int | None = None final_output: OmniRequestOutput | None = None + # Non-streaming is FINAL_ONLY, so the stage-0 output carries the full + # token sequence; the counter records its length for output_tokens. + usage_acc = SpeechOutputTokenCounter() async for res in generator: final_output = res + usage_acc.observe(res) if not is_moss: continue try: @@ -3749,6 +3849,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): audio_response: AudioResponse = self.create_audio(audio_obj) self._mark_ref_audio_artifact_ready_for_request(request_id) artifact_ready = True + if usage_out is not None: + usage_out.append(self._build_speech_usage(request, bytes_tts_params or {}, usage_acc.total())) return audio_response.audio_data, audio_response.media_type finally: if not artifact_ready: @@ -3983,7 +4085,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): if error is not None: return error - _, generator, _ = await self._prepare_speech_generation(request, request_id=request_id) + _, generator, sse_tts_params = await self._prepare_speech_generation(request, request_id=request_id) return StreamingResponse( self._generate_audio_sse_events( generator, @@ -3991,6 +4093,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): response_format, raw_request=raw_request, request_start_s=request_start_s, + request=request, + tts_params=sse_tts_params, ), media_type="text/event-stream", ) @@ -4108,8 +4212,9 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): validation_error = self._validate_tts_request(req) if validation_error is not None: return SpeechBatchItemResult(index=idx, status="error", error=validation_error) + usage_box: list[SpeechTokenUsage] = [] try: - audio_data, media_type = await self._generate_audio_bytes(req, base64_encode=True) + audio_data, media_type = await self._generate_audio_bytes(req, base64_encode=True, usage_out=usage_box) except Exception as e: logger.exception("Batch item %d failed: %s", idx, e) return SpeechBatchItemResult(index=idx, status="error", error=str(e)) @@ -4118,6 +4223,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): status="success", audio_data=audio_data, media_type=media_type, + usage=usage_box[0] if usage_box else None, ) results = await asyncio.gather( diff --git a/vllm_omni/entrypoints/openai/speech_usage.py b/vllm_omni/entrypoints/openai/speech_usage.py new file mode 100644 index 000000000..5e38405fd --- /dev/null +++ b/vllm_omni/entrypoints/openai/speech_usage.py @@ -0,0 +1,180 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Token-usage accounting for the Speech (``/v1/audio/speech``) API. + +Why this module exists (issue #4646) +------------------------------------ +For staged TTS models the engine prompt is a PLACEHOLDER: the serving layer +builds ``prompt_token_ids = [1] * prefill_len`` and lets the model rebuild the +real conditioning (text / ref_audio / ref_text) into ``inputs_embeds`` later. +So ``len(prompt_token_ids)`` is NOT a faithful count of what the caller sent. + +For Qwen3-TTS specifically the placeholder length mirrors the model prefill: + * CustomVoice/VoiceDesign: the full input text is embedded in the prefill, so + the placeholder length scales with the input text. + * Base in-context voice cloning: the prefill embeds only ``codec_bos`` + the + reference-audio codec frames; the input text is consumed incrementally + during DECODE, not prefill. So the placeholder length tracks the *reference + audio* and is independent of the input text. + +That is why ``usage.prompt_tokens`` (== ``len(prompt_token_ids)``) looked wrong +for Base: it counted reference-audio frames instead of the synthesized text. + +This module computes usage from the *semantic* inputs instead: + + input_tokens = text_tokens + audio_tokens + text_tokens -> tokens of ``input`` (+ ``instructions``); the text to speak + audio_tokens -> reference-audio codec frames used as voice-clone + conditioning, counted ONLY when in-context cloning is + actually active (see ``gate_audio_tokens``) + output_tokens = generated codec/audio tokens (stage-0 decode steps) + total_tokens = input_tokens + output_tokens + +Naming (``input_tokens``/``output_tokens``) follows OpenAI's ``speech.audio.done`` +event; the ``input_token_details`` breakdown follows OpenAI's realtime/chat +convention of never folding audio into an opaque text count. +""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any + +from vllm_omni.entrypoints.openai.protocol.audio import ( + SpeechInputTokenDetails, + SpeechTokenUsage, +) + + +def _first(value: Any, default: Any = None) -> Any: + """Unwrap the singleton-list convention used by ``tts_params``. + + ``tts_params`` wraps scalars in 1-element lists (e.g. ``task_type=["Base"]``) + because the model side batches per request. This returns the inner scalar. + """ + if isinstance(value, (list, tuple)): + return value[0] if value else default + return value if value is not None else default + + +def gate_audio_tokens( + *, + task_type: str | None, + x_vector_only_mode: bool, + icl_mode_override: bool | None, + ref_code_length: Any, +) -> int: + """Reference-audio codec frames that actually enter the prefill. + + Mirrors the ``in_context_mode`` decision in + ``Qwen3TTSPromptEmbedsBuilder.estimate_prompt_len_from_additional_information``: + reference codec frames are prefilled ONLY for Base in-context voice cloning. + They are NOT prefilled for: + * CustomVoice / VoiceDesign (no reference audio at all), or + * Base ``x_vector_only_mode`` (the reference audio is reduced to a single + speaker embedding vector; no per-frame codec context is inserted). + + Counting ``ref_code_length`` outside the in-context path would re-introduce + the issue #4646 inconsistency in the other direction, so we gate it here. + """ + in_context = (task_type == "Base") and not x_vector_only_mode + # An explicit per-request ``voice_clone_prompt.icl_mode`` wins if present, + # but ONLY for the Base task -- the estimator consults ``icl_mode`` inside + # its ``task_type == "Base"`` branch only, so CustomVoice / VoiceDesign must + # never have audio counted even if such a flag leaks into their params. + if task_type == "Base" and icl_mode_override is not None: + in_context = bool(icl_mode_override) + if not in_context: + return 0 + try: + return max(0, int(ref_code_length)) if ref_code_length is not None else 0 + except (TypeError, ValueError): + return 0 + + +def qwen3_tts_input_token_details( + *, + input_text: str, + instructions: str | None, + tts_params: dict[str, Any], + count_text_tokens: Callable[[str], int], +) -> SpeechInputTokenDetails: + """Compute the input-token breakdown for a Qwen3-TTS request. + + ``count_text_tokens`` tokenizes a string with the model's *text* tokenizer + and returns the token count. ``tts_params`` is the finalized param dict from + the adapter build (it already carries ``ref_code_length`` for ICL clones and + the resolved ``task_type`` / ``x_vector_only_mode``). + """ + # Text tokens = the text to synthesize plus the style/emotion instructions, + # because Qwen3-TTS tokenizes and prepends the instruction block too. + text = input_text or "" + instr = instructions or "" + text_tokens = count_text_tokens(text) if text else 0 + if instr.strip(): + text_tokens += count_text_tokens(instr) + + # Audio tokens = reference codec frames, gated to the in-context clone path. + voice_clone_prompt = _first(tts_params.get("voice_clone_prompt"), None) + icl_override = None + if isinstance(voice_clone_prompt, dict): + icl_flag = voice_clone_prompt.get("icl_mode") + if isinstance(icl_flag, bool): + icl_override = icl_flag + audio_tokens = gate_audio_tokens( + task_type=_first(tts_params.get("task_type"), "CustomVoice"), + x_vector_only_mode=bool(_first(tts_params.get("x_vector_only_mode"), False)), + icl_mode_override=icl_override, + ref_code_length=_first(tts_params.get("ref_code_length"), None), + ) + return SpeechInputTokenDetails(text_tokens=int(text_tokens), audio_tokens=int(audio_tokens)) + + +def build_speech_usage(details: SpeechInputTokenDetails, output_tokens: int) -> SpeechTokenUsage: + """Assemble the final usage object from the input breakdown + output count.""" + input_tokens = int(details.text_tokens) + int(details.audio_tokens) + out = max(0, int(output_tokens)) + return SpeechTokenUsage( + input_tokens=input_tokens, + output_tokens=out, + total_tokens=input_tokens + out, + input_token_details=details, + ) + + +@dataclass +class SpeechOutputTokenCounter: + """Accumulates the count of generated output (codec) tokens off the stream. + + Why not count ``res.outputs[0].token_ids``? In the staged TTS pipeline the + token-generating stage (stage 0, the AR talker) is NOT surfaced to the audio + consumer loop — only the downstream code2wav (audio) outputs reach it, and + those carry an empty ``token_ids`` (their payload is a waveform, not tokens). + + The generated codec-token count is instead reported by the pipeline itself + in per-stage metrics on the final (finished) output: + ``res.metrics['stage_metrics'][]['num_tokens_out']``. Only the AR + stage has a non-zero ``num_tokens_out`` (code2wav reports 0), so we take the + max across stages. Tracking the max is safe because the metrics appear on + the terminal output (intermediate outputs carry empty metrics). + """ + + output_tokens: int = 0 + + def observe(self, res: Any) -> None: + metrics = getattr(res, "metrics", None) + if not isinstance(metrics, dict): + return + stage_metrics = metrics.get("stage_metrics") + if not isinstance(stage_metrics, dict): + return + for stage in stage_metrics.values(): + if not isinstance(stage, dict): + continue + n = stage.get("num_tokens_out") + if isinstance(n, int) and n > self.output_tokens: + self.output_tokens = n + + def total(self) -> int: + """Generated codec/audio token count (max stage ``num_tokens_out``).""" + return self.output_tokens