diff --git a/docs/user_guide/examples/online_serving/qwen3_tts.md b/docs/user_guide/examples/online_serving/qwen3_tts.md index f899e362e..83253491b 100644 --- a/docs/user_guide/examples/online_serving/qwen3_tts.md +++ b/docs/user_guide/examples/online_serving/qwen3_tts.md @@ -264,9 +264,30 @@ Returns binary audio data with appropriate `Content-Type` header (e.g., `audio/w | `ref_text` | string | No | Transcript of reference audio (for ICL mode) | | `x_vector_only_mode` | bool | No | Use speaker embedding only (no ICL) | +## Streaming + +Set `stream=true` with `response_format="pcm"` to receive raw PCM audio chunks as they are decoded +(one chunk per 25-frame Code2Wav window): + +```bash +curl -X POST http://localhost:8091/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "input": "Hello, how are you?", + "voice": "vivian", + "language": "English", + "stream": true, + "response_format": "pcm" + }' --no-buffer | play -t raw -r 24000 -e signed -b 16 -c 1 - +``` + +**Constraints:** +- `stream=true` requires `response_format="pcm"` (raw 16-bit signed PCM, 24 kHz mono). +- `speed` adjustment is not supported when streaming. +- Requires the server stage config to have `async_chunk: true` (default in `qwen3_tts.yaml`). + ## Limitations -- **No streaming**: Audio is generated completely before being returned. Streaming will be supported after the pipeline is disaggregated (see RFC #938). - **Single request**: Batch processing is not yet optimized for online serving. ## Troubleshooting diff --git a/examples/offline_inference/qwen3_tts/README.md b/examples/offline_inference/qwen3_tts/README.md index 2b04cad75..f4ce92f4d 100644 --- a/examples/offline_inference/qwen3_tts/README.md +++ b/examples/offline_inference/qwen3_tts/README.md @@ -87,6 +87,19 @@ Examples: python end2end.py --query-type Base --mode-tag icl ``` +## Streaming Mode + +Add `--streaming` to stream audio chunks progressively via `AsyncOmni` (requires `async_chunk: true` in the stage config): + +```bash +python end2end.py --query-type CustomVoice --streaming --output-dir /tmp/out_stream +``` + +Each 25-frame Code2Wav chunk is logged as it arrives. The final WAV file is written once generation +completes. This demonstrates that audio data is available progressively rather than only at the end. + +> **Note:** Streaming uses `AsyncOmni` internally. The non-streaming path (`Omni`) is unchanged. + ## Batched Decoding The Code2Wav stage (stage 1) supports batched decoding, where multiple requests are decoded in a single forward pass through the SpeechTokenizer. To use it, provide a stage config with `max_batch_size > 1` and pass multiple prompts via `--txt-prompts` with a matching `--batch-size`. diff --git a/examples/offline_inference/qwen3_tts/end2end.py b/examples/offline_inference/qwen3_tts/end2end.py index fec515d30..2839a6958 100644 --- a/examples/offline_inference/qwen3_tts/end2end.py +++ b/examples/offline_inference/qwen3_tts/end2end.py @@ -4,6 +4,7 @@ Provides single and batch sample inputs for CustomVoice, VoiceDesign, and Base tasks, then runs Omni generation and saves output wav files. """ +import asyncio import logging import os from typing import Any, NamedTuple @@ -15,7 +16,7 @@ os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" from vllm.utils.argparse_utils import FlexibleArgumentParser -from vllm_omni import Omni +from vllm_omni import AsyncOmni, Omni logger = logging.getLogger(__name__) @@ -242,12 +243,15 @@ def get_base_query(use_batch_sample: bool = False, mode_tag: str = "icl") -> Que ) -def main(args): - """Run offline inference with Omni using prepared sample inputs. +query_map = { + "CustomVoice": get_custom_voice_query, + "VoiceDesign": get_voice_design_query, + "Base": get_base_query, +} - Args: - args: Parsed CLI args from parse_args(). - """ + +def _build_inputs(args) -> tuple[str, list]: + """Resolve model name and inputs list from CLI args.""" if args.batch_size < 1 or (args.batch_size & (args.batch_size - 1)) != 0: raise ValueError( f"--batch-size must be a power of two (got {args.batch_size}); " @@ -259,40 +263,49 @@ def main(args): if args.query_type in {"CustomVoice", "VoiceDesign"}: query_result = query_func(use_batch_sample=args.use_batch_sample) elif args.query_type == "Base": - query_result = query_func( - use_batch_sample=args.use_batch_sample, - mode_tag=args.mode_tag, - ) + query_result = query_func(use_batch_sample=args.use_batch_sample, mode_tag=args.mode_tag) else: query_result = query_func() model_name = query_result.model_name - # Load prompts from text file if provided. - # Use the default query as a template so task-specific fields - # (e.g. ref_audio for Base) are preserved; only override text. if args.txt_prompts: with open(args.txt_prompts) as f: lines = [line.strip() for line in f if line.strip()] if not lines: raise ValueError(f"No valid prompts found in {args.txt_prompts}") - template = query_result.inputs - if isinstance(template, list): - template = template[0] + template = query_result.inputs if not isinstance(query_result.inputs, list) else query_result.inputs[0] template_info = template["additional_information"] - inputs = [] - for text in lines: - additional_information = {**template_info, "text": [text]} - inputs.append( - { - "prompt_token_ids": [0] * _estimate_prompt_len(additional_information, model_name), - "additional_information": additional_information, - } - ) + inputs = [ + { + "prompt_token_ids": [0] * _estimate_prompt_len({**template_info, "text": [t]}, model_name), + "additional_information": {**template_info, "text": [t]}, + } + for t in lines + ] else: - inputs = query_result.inputs - if not isinstance(inputs, list): - inputs = [inputs] + inputs = query_result.inputs if isinstance(query_result.inputs, list) else [query_result.inputs] + + return model_name, inputs + + +def _save_wav(output_dir: str, request_id: str, mm: dict) -> None: + """Concatenate audio chunks and write to a wav file.""" + audio_data = mm["audio"] + sr_raw = mm["sr"] + sr_val = sr_raw[-1] if isinstance(sr_raw, list) and sr_raw else sr_raw + sr = sr_val.item() if hasattr(sr_val, "item") else int(sr_val) + audio_tensor = torch.cat(audio_data, dim=-1) if isinstance(audio_data, list) else audio_data + out_wav = os.path.join(output_dir, f"output_{request_id}.wav") + sf.write(out_wav, audio_tensor.float().cpu().numpy().flatten(), samplerate=sr, format="WAV") + logger.info(f"Request ID: {request_id}, Saved audio to {out_wav}") + + +def main(args): + """Run offline inference with Omni.""" + model_name, inputs = _build_inputs(args) + output_dir = args.output_dir + os.makedirs(output_dir, exist_ok=True) omni = Omni( model=model_name, @@ -301,43 +314,40 @@ def main(args): stage_init_timeout=args.stage_init_timeout, ) - output_dir = args.output_dir - os.makedirs(output_dir, exist_ok=True) - batch_size = args.batch_size for batch_start in range(0, len(inputs), batch_size): batch = inputs[batch_start : batch_start + batch_size] - omni_generator = omni.generate(batch, sampling_params_list=None) - for stage_outputs in omni_generator: + for stage_outputs in omni.generate(batch): for output in stage_outputs.request_output: - request_id = output.request_id - audio_data = output.outputs[0].multimodal_output["audio"] - # async_chunk mode returns a list of chunks; concatenate them. - if isinstance(audio_data, list): - audio_tensor = torch.cat(audio_data, dim=-1) - else: - audio_tensor = audio_data - output_wav = os.path.join(output_dir, f"output_{request_id}.wav") - sr_val = output.outputs[0].multimodal_output["sr"] - audio_samplerate = sr_val.item() if hasattr(sr_val, "item") else int(sr_val[-1]) - # Convert to numpy array and ensure correct format - audio_numpy = audio_tensor.float().detach().cpu().numpy() + _save_wav(output_dir, output.request_id, output.outputs[0].multimodal_output) - # Ensure audio is 1D (flatten if needed) - if audio_numpy.ndim > 1: - audio_numpy = audio_numpy.flatten() - # Save audio file with explicit WAV format - sf.write(output_wav, audio_numpy, samplerate=audio_samplerate, format="WAV") - print(f"Request ID: {request_id}, Saved audio to {output_wav}") +async def main_streaming(args): + """Run offline inference with AsyncOmni, logging each audio chunk as it arrives.""" + model_name, inputs = _build_inputs(args) + output_dir = args.output_dir + os.makedirs(output_dir, exist_ok=True) + + omni = AsyncOmni( + model=model_name, + stage_configs_path=args.stage_configs_path, + log_stats=args.log_stats, + stage_init_timeout=args.stage_init_timeout, + ) + + for i, prompt in enumerate(inputs): + request_id = str(i) + async for stage_output in omni.generate(prompt, request_id=request_id): + mm = stage_output.request_output.outputs[0].multimodal_output + if not stage_output.finished: + audio = mm.get("audio") + n = len(audio) if isinstance(audio, list) else (0 if audio is None else 1) + logger.info(f"Request {request_id}: received chunk {n}") + else: + _save_wav(output_dir, request_id, mm) def parse_args(): - """Parse CLI arguments for offline TTS inference. - - Returns: - argparse.Namespace with CLI options. - """ parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models") parser.add_argument( "--query-type", @@ -438,6 +448,12 @@ def parse_args(): choices=["icl", "xvec_only"], help="Mode tag for Base query x_vector_only_mode (default: icl).", ) + parser.add_argument( + "--streaming", + action="store_true", + default=False, + help="Stream audio chunks as they arrive via AsyncOmni (async_chunk mode only).", + ) parser.add_argument( "--batch-size", type=int, @@ -448,13 +464,9 @@ def parse_args(): return parser.parse_args() -query_map = { - "CustomVoice": get_custom_voice_query, - "VoiceDesign": get_voice_design_query, - "Base": get_base_query, -} - - if __name__ == "__main__": args = parse_args() - main(args) + if args.streaming: + asyncio.run(main_streaming(args)) + else: + main(args) diff --git a/examples/online_serving/qwen3_tts/README.md b/examples/online_serving/qwen3_tts/README.md index dd4619cc4..71a3960f0 100644 --- a/examples/online_serving/qwen3_tts/README.md +++ b/examples/online_serving/qwen3_tts/README.md @@ -146,7 +146,8 @@ Lists available voices for the loaded model: "ref_audio": "URL or base64 for voice cloning", "ref_text": "Reference audio transcript", "x_vector_only_mode": false, - "max_new_tokens": 2048 + "max_new_tokens": 2048, + "stream": false } ``` @@ -166,7 +167,7 @@ Returns audio data in the requested format (default: WAV). | `model` | string | server's model | Model to use (optional, should match server if specified) | | `voice` | string | "vivian" | Speaker name (e.g., vivian, ryan, aiden) | | `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, aac, opus | -| `speed` | float | 1.0 | Playback speed (0.25-4.0) | +| `speed` | float | 1.0 | Playback speed (0.25-4.0, not supported with `stream=true`) | ### vLLM-Omni Extension Parameters @@ -176,6 +177,7 @@ Returns audio data in the requested format (default: WAV). | `language` | string | "Auto" | Language (see supported languages below) | | `instructions` | string | "" | Voice style/emotion instructions | | `max_new_tokens` | int | 2048 | Maximum tokens to generate | +| `stream` | bool | false | Stream raw PCM chunks as they are decoded (requires `response_format="pcm"`) | **Supported languages:** Auto, Chinese, English, Japanese, Korean, German, French, Russian, Portuguese, Spanish, Italian @@ -224,9 +226,37 @@ with open("output.wav", "wb") as f: f.write(response.content) ``` +## Streaming + +Set `stream=true` with `response_format="pcm"` to receive raw PCM audio chunks as they are decoded +(one chunk per 25-frame Code2Wav window, ~2 s latency to first byte): + +```bash +# Stream raw PCM and pipe directly to a player +curl -X POST http://localhost:8091/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "input": "Hello, how are you?", + "voice": "vivian", + "language": "English", + "stream": true, + "response_format": "pcm" + }' --no-buffer | play -t raw -r 24000 -e signed -b 16 -c 1 - + +# Or save the raw PCM for offline processing +curl -X POST http://localhost:8091/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{"input": "Hello!", "stream": true, "response_format": "pcm"}' \ + --output output.pcm +``` + +**Constraints:** +- `stream=true` requires `response_format="pcm"` (raw 16-bit signed PCM, 24 kHz mono). +- `speed` adjustment is not supported when streaming (`speed` must be 1.0 or omitted). +- Requires the server to be launched with a stage config that sets `async_chunk: true`. + ## Limitations -- **No streaming**: Audio is generated completely before being returned. Streaming will be supported after the pipeline is disaggregated (see RFC #938). - **Single request**: Batch processing is not yet optimized for online serving. ## Troubleshooting diff --git a/tests/entrypoints/openai_api/test_serving_speech.py b/tests/entrypoints/openai_api/test_serving_speech.py index 37ee72972..7af4e078f 100644 --- a/tests/entrypoints/openai_api/test_serving_speech.py +++ b/tests/entrypoints/openai_api/test_serving_speech.py @@ -7,6 +7,7 @@ import pytest import torch from fastapi import FastAPI from fastapi.testclient import TestClient +from pydantic import ValidationError from pytest_mock import MockerFixture from vllm_omni.entrypoints.openai.audio_utils_mixin import AudioMixin @@ -507,3 +508,110 @@ class TestTTSMethods: error = server._validate_tts_request(req) assert error is not None assert "max 10 characters" in error + + +class TestStreamingProtocolValidation: + """Unit tests for the stream field validators in OpenAICreateSpeechRequest.""" + + def test_stream_validation_errors(self): + """stream=True requires response_format='pcm' and speed=1.0.""" + with pytest.raises(ValidationError, match="response_format='pcm'"): + OpenAICreateSpeechRequest(input="Hello", stream=True, response_format="wav") + with pytest.raises(ValidationError, match="Speed adjustment is not supported"): + OpenAICreateSpeechRequest(input="Hello", stream=True, response_format="pcm", speed=2.0) + + def test_stream_valid(self): + """stream=True + response_format='pcm' + speed=1.0 is accepted.""" + req = OpenAICreateSpeechRequest(input="Hello", stream=True, response_format="pcm") + assert req.stream is True + + def test_sse_stream_format_is_blocked(self): + """stream_format='sse' is blocked.""" + with pytest.raises(ValidationError, match="sse"): + OpenAICreateSpeechRequest(input="Hello", stream_format="sse") + + +class TestStreamingResponse: + """Integration tests for the streaming audio response path.""" + + @pytest.fixture + def streaming_app(self, mocker: MockerFixture): + """Test app whose mock engine yields one intermediate chunk then a final chunk.""" + mock_engine_client = mocker.MagicMock() + mock_engine_client.errored = False + + def _make_output(finished: bool) -> OmniRequestOutput: + chunk = torch.sin(torch.linspace(0, 440 * 2 * torch.pi, 24000)) + + class MockCompletionOutput: + def __init__(self, index: int = 0): + self.index = index + self.text = "" + self.token_ids = [] + self.finish_reason = "stop" + self.stop_reason = None + self.logprobs = None + + class MockRequestOutput: + def __init__(self, audio_tensor: torch.Tensor): + self.request_id = "speech-stream-test" + self.outputs = [MockCompletionOutput(index=0)] + self.multimodal_output = {"audio": audio_tensor} + self.finished = finished + self.prompt_token_ids = None + self.encoder_prompt_token_ids = None + self.num_cached_tokens = None + self.prompt_logprobs = None + self.kv_transfer_params = None + + return OmniRequestOutput( + stage_id=0, + final_output_type="audio", + request_output=MockRequestOutput(audio_tensor=chunk), + finished=finished, + ) + + async def mock_generate_streaming(*args, **kwargs): + yield _make_output(finished=False) + yield _make_output(finished=True) + + mock_engine_client.generate = mocker.MagicMock(side_effect=mock_generate_streaming) + mock_engine_client.default_sampling_params_list = [{}] + mock_models = mocker.MagicMock() + mock_models.is_base_model.return_value = True + + speech_server = OmniOpenAIServingSpeech( + engine_client=mock_engine_client, + models=mock_models, + request_logger=mocker.MagicMock(), + ) + + original_create_speech = speech_server.create_speech + sig = signature(original_create_speech) + new_parameters = [p for name, p in sig.parameters.items() if name != "raw_request"] + new_sig = Signature(parameters=new_parameters, return_annotation=sig.return_annotation) + + async def awaitable_create_speech(*args, **kwargs): + return await original_create_speech(*args, **kwargs) + + awaitable_create_speech.__signature__ = new_sig + speech_server.create_speech = awaitable_create_speech + + app = FastAPI() + app.add_api_route("/v1/audio/speech", speech_server.create_speech, methods=["POST"], response_model=None) + return app + + def test_streaming(self, streaming_app): + """stream=True must return audio/pcm with non-empty body.""" + client = TestClient(streaming_app) + response = client.post("/v1/audio/speech", json={"input": "Hello", "stream": True, "response_format": "pcm"}) + assert response.status_code == 200 + assert "audio/pcm" in response.headers["content-type"] + assert len(response.content) > 0 + + def test_non_streaming_unchanged(self, streaming_app): + """Non-streaming path must still return audio/wav.""" + client = TestClient(streaming_app) + response = client.post("/v1/audio/speech", json={"input": "Hello", "response_format": "wav"}) + assert response.status_code == 200 + assert "audio/wav" in response.headers["content-type"] diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index 08c4ce903..79dc7a8b7 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/vllm_omni/entrypoints/async_omni.py @@ -528,12 +528,14 @@ class AsyncOmni(OmniBase): final_output_type=stage.final_output_type, request_output=engine_outputs, images=images, + finished=finished, ) else: output_to_yield = OmniRequestOutput( stage_id=stage_id, final_output_type=stage.final_output_type, request_output=engine_outputs, + finished=finished, ) # Mark last output time metrics.stage_last_ts[stage_id] = max(metrics.stage_last_ts[stage_id] or 0.0, time.time()) diff --git a/vllm_omni/entrypoints/openai/protocol/audio.py b/vllm_omni/entrypoints/openai/protocol/audio.py index d23460626..8915e544b 100644 --- a/vllm_omni/entrypoints/openai/protocol/audio.py +++ b/vllm_omni/entrypoints/openai/protocol/audio.py @@ -1,7 +1,7 @@ from typing import Literal import numpy as np -from pydantic import BaseModel, Field, field_validator +from pydantic import BaseModel, Field, field_validator, model_validator class OpenAICreateSpeechRequest(BaseModel): @@ -22,6 +22,13 @@ class OpenAICreateSpeechRequest(BaseModel): le=4.0, ) stream_format: Literal["sse", "audio"] | None = "audio" + stream: bool = Field( + default=False, + description=( + "If true, stream raw PCM audio chunks as they are decoded. " + "Requires response_format='pcm'. Speed adjustment is not supported when streaming." + ), + ) # Qwen3-TTS specific parameters task_type: Literal["CustomVoice", "VoiceDesign", "Base"] | None = Field( @@ -56,6 +63,22 @@ class OpenAICreateSpeechRequest(BaseModel): raise ValueError("'sse' is not a supported stream_format yet. Please use 'audio'.") return v + @model_validator(mode="after") + def validate_streaming_constraints(self) -> "OpenAICreateSpeechRequest": + if self.stream: + if self.response_format != "pcm": + raise ValueError( + "Streaming (stream=true) requires response_format='pcm'. " + f"Got response_format='{self.response_format}'." + ) + if self.speed is None: + self.speed = 1.0 + elif self.speed != 1.0: + raise ValueError( + "Speed adjustment is not supported when streaming (stream=true). Set speed=1.0 or omit it." + ) + return self + class CreateAudio(BaseModel): audio_tensor: np.ndarray diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index 98f79f87f..1275daae7 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -10,14 +10,13 @@ from urllib.request import urlopen import numpy as np import soundfile as sf from fastapi import Request -from fastapi.responses import Response +from fastapi.responses import Response, StreamingResponse from vllm.entrypoints.openai.engine.serving import OpenAIServing from vllm.logger import init_logger from vllm.utils import random_uuid from vllm_omni.entrypoints.openai.audio_utils_mixin import AudioMixin from vllm_omni.entrypoints.openai.protocol.audio import ( - AudioResponse, CreateAudio, OpenAICreateSpeechRequest, ) @@ -253,6 +252,85 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): wav_np, sr = await loop.run_in_executor(None, _fetch_sync) return wav_np.tolist(), sr + async def _generate_pcm_chunks(self, generator, request_id: str): + """Generate PCM audio chunks for streaming response. + + Handles two audio output modes from the engine: + - Cumulative mode (list): Engine returns growing list of chunks; + we emit only the new tail on each iteration. + - Per-step mode (tensor): Engine returns single tensor per iteration; + we emit it directly. + + Args: + generator: Async generator from the engine + request_id: Request identifier for logging + + Yields: + Raw PCM bytes for each audio chunk + """ + prev_count = 0 + sample_rate_val = 24000 + try: + async for res in generator: + audio_output, audio_key = self._extract_audio_output(res) + if audio_key is None: + continue + + sr_raw = audio_output.get("sr") + if sr_raw is not None: + sr_val = sr_raw[-1] if isinstance(sr_raw, list) and sr_raw else sr_raw + sample_rate_val = sr_val.item() if hasattr(sr_val, "item") else int(sr_val) + + audio_val = audio_output[audio_key] + if isinstance(audio_val, list): + # Cumulative mode: each update grows the list; emit only new tail. + new_chunks = audio_val[prev_count:] + prev_count = len(audio_val) + else: + # Per-step mode: each update is a single tensor; emit directly. + if audio_val is not None: + new_chunks = [audio_val] + prev_count += 1 + else: + new_chunks = [] + + for chunk_tensor in new_chunks: + chunk_np = ( + chunk_tensor.float().detach().cpu().numpy() if hasattr(chunk_tensor, "float") else chunk_tensor + ) + if chunk_np.ndim > 1: + chunk_np = chunk_np.squeeze() + audio_obj = CreateAudio( + audio_tensor=chunk_np, + sample_rate=sample_rate_val, + response_format="pcm", + speed=1.0, + stream_format="audio", + base64_encode=False, + ) + yield self.create_audio(audio_obj).audio_data + except asyncio.CancelledError: + logger.info("Streaming request %s cancelled by client", request_id) + raise + except Exception as e: + logger.exception("Streaming speech generation failed for %s: %s", request_id, e) + + @staticmethod + def _extract_audio_output(res) -> tuple[dict | None, str | None]: + """Return (audio_output dict, audio key) or (None, None). + + Returns the raw dict so callers can apply their own extraction strategy: + streaming needs per-chunk delta slicing; non-streaming needs full concatenation. + """ + mm = getattr(res, "multimodal_output", None) + if not mm: + ro = getattr(res, "request_output", None) + mm = getattr(ro, "multimodal_output", None) if ro else None + if not mm: + return None, None + key = "audio" if "audio" in mm else ("model_outputs" if "model_outputs" in mm else None) + return mm, key + def _build_tts_params(self, request: OpenAICreateSpeechRequest) -> dict[str, Any]: """Build TTS parameters from request. @@ -300,6 +378,11 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): else: params["max_new_tokens"] = [2048] + # VoiceDesign requires non_streaming_mode (match offline script behaviour). + # CustomVoice and Base rely on the model default (True and False respectively). + if params["task_type"][0] == "VoiceDesign": + params["non_streaming_mode"] = [True] + return params async def create_speech( @@ -323,9 +406,9 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): - ref_text: Transcript of reference audio (Base task) - x_vector_only_mode: Use speaker embedding only (Base task) - NOTE: Streaming audio generation is not currently supported. + Streaming is supported via stream=True with response_format='pcm'. + Each Code2Wav chunk is yielded as raw PCM bytes as soon as it is decoded. """ - error_check_ret = await self._check_model(request) if error_check_ret is not None: logger.error("Error with model %s", error_check_ret) @@ -343,23 +426,15 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): if validation_error: return self.create_error_response(validation_error) - # Must use prompt_token_ids (not text prompt): the AR Talker - # operates on codec tokens; text token IDs exceed codec vocab. - # model.preprocess replaces all embeddings, so placeholder value - # is irrelevant -- but length must match to avoid excess padding. tts_params = self._build_tts_params(request) - if request.ref_audio is not None: wav_list, sr = await self._resolve_ref_audio(request.ref_audio) tts_params["ref_audio"] = [[wav_list, sr]] + # Prompt length must match model-side embeddings; values are placeholders. ph_len = self._estimate_prompt_len(tts_params) - prompt = { - "prompt_token_ids": [1] * ph_len, - "additional_information": tts_params, - } + prompt = {"prompt_token_ids": [1] * ph_len, "additional_information": tts_params} else: - # Fallback for unsupported models tts_params = {} prompt = {"prompt": request.input} @@ -379,6 +454,13 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): output_modalities=["audio"], ) + if request.stream: + return StreamingResponse( + self._generate_pcm_chunks(generator, request_id), + media_type="audio/pcm", + ) + + # Non-streaming: collect final output final_output: OmniRequestOutput | None = None async for res in generator: final_output = res @@ -386,55 +468,34 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): if final_output is None: return self.create_error_response("No output generated from the model.") - # Extract audio from output - # Audio can be in final_output.multimodal_output or final_output.request_output.multimodal_output - # Support both "audio" and "model_outputs" keys for compatibility with different models - audio_output = None - if hasattr(final_output, "multimodal_output") and final_output.multimodal_output: - audio_output = final_output.multimodal_output - if not audio_output and hasattr(final_output, "request_output"): - if final_output.request_output and hasattr(final_output.request_output, "multimodal_output"): - audio_output = final_output.request_output.multimodal_output - - # Check for audio data using either "audio" or "model_outputs" key - audio_key = None - if audio_output: - if "audio" in audio_output: - audio_key = "audio" - elif "model_outputs" in audio_output: - audio_key = "model_outputs" - - if not audio_output or audio_key is None: + audio_output, audio_key = self._extract_audio_output(final_output) + if audio_key is None: return self.create_error_response("TTS model did not produce audio output.") audio_tensor = audio_output[audio_key] - sample_rate = audio_output.get("sr", 24000) - if hasattr(sample_rate, "item"): - sample_rate = sample_rate.item() + sr_raw = audio_output.get("sr", 24000) + sr_val = sr_raw[-1] if isinstance(sr_raw, list) and sr_raw else sr_raw + sample_rate = sr_val.item() if hasattr(sr_val, "item") else int(sr_val) - # Streaming accumulates chunks as a list; concat first. + # async_chunk mode accumulates chunks as a list; concat first. if isinstance(audio_tensor, list): import torch audio_tensor = torch.cat(audio_tensor, dim=-1) - # Convert tensor to numpy if hasattr(audio_tensor, "float"): audio_tensor = audio_tensor.float().detach().cpu().numpy() - - # Squeeze batch dimension if present, but preserve channel dimension for stereo if audio_tensor.ndim > 1: audio_tensor = audio_tensor.squeeze() audio_obj = CreateAudio( audio_tensor=audio_tensor, - sample_rate=int(sample_rate), + sample_rate=sample_rate, response_format=request.response_format or "wav", speed=request.speed or 1.0, stream_format=request.stream_format, base64_encode=False, ) - - audio_response: AudioResponse = self.create_audio(audio_obj) + audio_response = self.create_audio(audio_obj) return Response(content=audio_response.audio_data, media_type=audio_response.media_type) except asyncio.CancelledError: