diff --git a/tests/entrypoints/openai_api/test_serving_speech.py b/tests/entrypoints/openai_api/test_serving_speech.py index 7e35fec0d..e90992317 100644 --- a/tests/entrypoints/openai_api/test_serving_speech.py +++ b/tests/entrypoints/openai_api/test_serving_speech.py @@ -2,6 +2,7 @@ import asyncio import logging import os +import struct from inspect import Signature, signature from pathlib import Path from unittest.mock import AsyncMock, MagicMock, patch @@ -22,6 +23,7 @@ from vllm_omni.entrypoints.openai.audio_utils_mixin import AudioMixin from vllm_omni.entrypoints.openai.protocol.audio import CreateAudio, OpenAICreateSpeechRequest from vllm_omni.entrypoints.openai.serving_speech import ( OmniOpenAIServingSpeech, + _create_wav_header, ) from vllm_omni.outputs import OmniRequestOutput @@ -920,17 +922,20 @@ 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") + """stream=True requires response_format not in ('pcm', 'wav') and speed=1.0.""" + with pytest.raises(ValidationError, match="requires response_format not in \\('pcm', 'wav'\\)"): + OpenAICreateSpeechRequest(input="Hello", stream=True, response_format="mp3") 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.""" + """stream=True + response_format in ('pcm', 'wav') + speed=1.0 is accepted.""" req = OpenAICreateSpeechRequest(input="Hello", stream=True, response_format="pcm") assert req.stream is True + req = OpenAICreateSpeechRequest(input="Hello", stream=True, response_format="wav") + assert req.stream is True + def test_sse_stream_format_is_blocked(self): """stream_format='sse' is blocked.""" with pytest.raises(ValidationError, match="sse"): @@ -1086,3 +1091,191 @@ def test_api_server_create_speech_wraps_error_response_status(): assert isinstance(response, JSONResponse) assert response.status_code == 400 + + +class TestWAVHeaderGeneration: + """Unit tests for WAV header generation with placeholder values.""" + + def test_wav_header_basic_structure(self): + """Test basic WAV header structure with default parameters.""" + header = _create_wav_header(sample_rate=24000, num_channels=1, bits_per_sample=16) + + # Verify header length (should be 44 bytes) + assert len(header) == 44, f"Expected 44 bytes, got {len(header)}" + + # Parse and verify header structure + ( + chunk_id, + chunk_size, + format_type, + subchunk1_id, + subchunk1_size, + audio_format, + num_channels, + sample_rate, + byte_rate, + block_align, + bits_per_sample, + subchunk2_id, + subchunk2_size, + ) = struct.unpack("<4sI4s4sIHHIIHH4sI", header) + + # Verify RIFF header + assert chunk_id == b"RIFF", f"Expected RIFF, got {chunk_id}" + assert chunk_size == 0xFFFFFFFF, f"Expected placeholder 0xFFFFFFFF, got {chunk_size:#x}" + assert format_type == b"WAVE", f"Expected WAVE, got {format_type}" + + # Verify fmt chunk + assert subchunk1_id == b"fmt ", f"Expected 'fmt ', got {subchunk1_id}" + assert subchunk1_size == 16, f"Expected 16, got {subchunk1_size}" + assert audio_format == 1, f"Expected PCM (1), got {audio_format}" + assert num_channels == 1, f"Expected 1 channel, got {num_channels}" + assert sample_rate == 24000, f"Expected 24000 Hz, got {sample_rate}" + assert byte_rate == 48000, f"Expected 48000 byte/s, got {byte_rate}" + assert block_align == 2, f"Expected 2 bytes block align, got {block_align}" + assert bits_per_sample == 16, f"Expected 16 bits, got {bits_per_sample}" + + # Verify data chunk + assert subchunk2_id == b"data", f"Expected 'data', got {subchunk2_id}" + assert subchunk2_size == 0xFFFFFFFF, f"Expected placeholder 0xFFFFFFFF, got {subchunk2_size:#x}" + + def test_wav_header_different_sample_rates(self): + """Test WAV header with different sample rates.""" + test_cases = [ + (16000, 1, 16), + (22050, 1, 16), + (24000, 1, 16), + (44100, 1, 16), + (48000, 1, 16), + ] + + for sample_rate, num_channels, bits_per_sample in test_cases: + header = _create_wav_header(sample_rate, num_channels, bits_per_sample) + assert len(header) == 44, f"Header length mismatch for {sample_rate} Hz" + + # Parse sample rate from header + parsed_sample_rate = struct.unpack(" 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-wav-stream-test" + self.outputs = [MockCompletionOutput(index=0)] + self.multimodal_output = {"audio": audio_tensor, "sr": 24000} + 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_wav_streaming_success(self, wav_streaming_app): + """Test WAV format streaming returns correct content type and includes WAV header.""" + client = TestClient(wav_streaming_app) + response = client.post("/v1/audio/speech", json={"input": "Hello", "stream": True, "response_format": "wav"}) + + assert response.status_code == 200 + assert "audio/wav" in response.headers["content-type"] + assert len(response.content) > 44 # Should have WAV header + audio data + + # Verify WAV header is present + header = response.content[:44] + chunk_id = header[0:4] + format_type = header[8:12] + assert chunk_id == b"RIFF", "Should start with RIFF" + assert format_type == b"WAVE", "Should contain WAVE format" + + def test_streaming_unsupported_format_rejected(self, wav_streaming_app): + """Test that unsupported formats are rejected for streaming.""" + client = TestClient(wav_streaming_app) + + unsupported_formats = ["mp3"] + for fmt in unsupported_formats: + response = client.post("/v1/audio/speech", json={"input": "Hello", "stream": True, "response_format": fmt}) + assert response.status_code == 422 diff --git a/vllm_omni/entrypoints/openai/protocol/audio.py b/vllm_omni/entrypoints/openai/protocol/audio.py index eb32d7f05..e57743226 100644 --- a/vllm_omni/entrypoints/openai/protocol/audio.py +++ b/vllm_omni/entrypoints/openai/protocol/audio.py @@ -71,9 +71,9 @@ class OpenAICreateSpeechRequest(BaseModel): @model_validator(mode="after") def validate_streaming_constraints(self) -> "OpenAICreateSpeechRequest": if self.stream: - if self.response_format != "pcm": + if self.response_format not in ("pcm", "wav"): raise ValueError( - "Streaming (stream=true) requires response_format='pcm'. " + "Streaming (stream=true) requires response_format not in ('pcm', 'wav'). " f"Got response_format='{self.response_format}'." ) if self.speed is None: diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index 6b5656474..331adda8e 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -4,6 +4,7 @@ import json import math import os import re +import struct import time from pathlib import Path from typing import Any @@ -48,6 +49,47 @@ _TTS_MAX_NEW_TOKENS_MIN = 1 _TTS_MAX_NEW_TOKENS_MAX = 4096 +def _create_wav_header(sample_rate: int, num_channels: int = 1, bits_per_sample: int = 16) -> bytes: + """Create a WAV header with placeholder size values for streaming. + + Uses 0xFFFFFFFF as placeholder for data size fields, which is accepted + by most audio clients and matches OpenAI's streaming WAV implementation. + + Args: + sample_rate: Audio sample rate in Hz + num_channels: Number of audio channels (1 for mono, 2 for stereo) + bits_per_sample: Bits per sample (typically 16) + + Returns: + 44-byte WAV header as bytes + """ + byte_rate = sample_rate * num_channels * bits_per_sample // 8 + block_align = num_channels * bits_per_sample // 8 + + # Use 0xFFFFFFFF as placeholder for unknown size (streaming) + placeholder_size = 0xFFFFFFFF + + # ref https://docs.fileformat.com/audio/wav/ + header = struct.pack( + "<4sI4s4sIHHIIHH4sI", + b"RIFF", # ChunkID + placeholder_size, # ChunkSize (placeholder) + b"WAVE", # Format + b"fmt ", # Subchunk1ID + 16, # Subchunk1Size (16 for PCM) + 1, # AudioFormat (1 for PCM) + num_channels, # NumChannels + sample_rate, # SampleRate + byte_rate, # ByteRate + block_align, # BlockAlign + bits_per_sample, # BitsPerSample + b"data", # Subchunk2ID + placeholder_size, # Subchunk2Size (placeholder) + ) + + return header + + def _sanitize_filename(filename: str) -> str: """Sanitize filename to prevent path traversal attacks. @@ -575,8 +617,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): wav_np = np.mean(wav_np, axis=-1) return wav_np.tolist(), int(sr) - async def _generate_pcm_chunks(self, generator, request_id: str): - """Generate PCM audio chunks for streaming response. + async def _generate_audio_chunks(self, generator, request_id: str, response_format: str = "pcm"): + """Generate audio chunks for streaming response. Handles two audio output modes from the engine: - Cumulative mode (list): Engine returns growing list of chunks; @@ -587,12 +629,15 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): Args: generator: Async generator from the engine request_id: Request identifier for logging + response_format: Audio format (pcm or wav) Yields: - Raw PCM bytes for each audio chunk + Raw audio bytes for each chunk (with WAV header for first chunk if wav format) """ prev_count = 0 sample_rate_val = 24000 + first_chunk = True + try: async for res in generator: audio_output, audio_key = self._extract_audio_output(res) @@ -623,6 +668,18 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): ) if chunk_np.ndim > 1: chunk_np = chunk_np.squeeze() + # For WAV format, emit header before first audio chunk + if response_format == "wav" and first_chunk: + # Assert that sample rate has been set from chunk metadata (not just default) + # This ensures the WAV header contains the correct sample rate + assert sr_raw is not None, ( + "First audio chunk must include sample rate metadata for WAV streaming" + ) + wav_header = _create_wav_header(sample_rate=sample_rate_val, num_channels=1, bits_per_sample=16) + yield wav_header + first_chunk = False + + # Convert audio to PCM bytes audio_obj = CreateAudio( audio_tensor=chunk_np, sample_rate=sample_rate_val, @@ -833,8 +890,9 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): - ref_text: Transcript of reference audio (Base task) - x_vector_only_mode: Use speaker embedding only (Base task) - 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. + Streaming is supported via stream=True with response_format='pcm' or 'wav'. + Each Code2Wav chunk is yielded as raw audio bytes as soon as it is decoded. + For WAV format, a header with placeholder size values is emitted first. """ error_check_ret = await self._check_model(request) if error_check_ret is not None: @@ -843,10 +901,28 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): try: if request.stream: + # Determine response format and media type for streaming + response_format = (request.response_format or "wav").lower() + + # Only pcm and wav support streaming without post-processing + if response_format not in ["pcm", "wav"]: + return self.create_error_response( + f"Streaming is only supported for 'pcm' and 'wav' formats. " + f"Got '{response_format}'. For other formats, use stream=False." + ) + + # Check if speed adjustment is requested (not compatible with streaming) + if request.speed is not None and request.speed != 1.0: + return self.create_error_response( + "Streaming is not supported with speed adjustment. " + "Use stream=False or remove the speed parameter." + ) + + media_type = "audio/wav" if response_format == "wav" else "audio/pcm" request_id, generator, _ = await self._prepare_speech_generation(request) return StreamingResponse( - self._generate_pcm_chunks(generator, request_id), - media_type="audio/pcm", + self._generate_audio_chunks(generator, request_id, response_format), + media_type=media_type, ) audio_bytes, media_type = await self._generate_audio_bytes(request)