[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>

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