[Qwen3TTS][Feat] Streaming output (#1438)

Signed-off-by: pablo <pablo@agigo.ai>
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: pablo <pablo@agigo.ai>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

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

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`.

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)

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

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"]

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())

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

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: