From 2f69748ee48d69c51dc924c9f84028ac858fbc87 Mon Sep 17 00:00:00 2001 From: Bhanu Prakash Voutharoja <59905694+Bhanu068@users.noreply.github.com> Date: Sat, 27 Dec 2025 16:56:51 +1100 Subject: [PATCH] add openai create speech endpoint (#305) Signed-off-by: Bhanu068 Signed-off-by: Bhanu Prakash Voutharoja <59905694+Bhanu068@users.noreply.github.com> Co-authored-by: Hongsheng Liu Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com> --- .../openai_api/test_serving_speech.py | 267 ++++++++++++++++++ vllm_omni/entrypoints/openai/api_server.py | 32 +++ .../entrypoints/openai/audio_utils_mixin.py | 93 ++++++ .../entrypoints/openai/protocol/audio.py | 42 +++ vllm_omni/entrypoints/openai/serving_chat.py | 29 +- .../entrypoints/openai/serving_speech.py | 80 ++++++ 6 files changed, 527 insertions(+), 16 deletions(-) create mode 100644 tests/entrypoints/openai_api/test_serving_speech.py create mode 100644 vllm_omni/entrypoints/openai/audio_utils_mixin.py create mode 100644 vllm_omni/entrypoints/openai/protocol/audio.py create mode 100644 vllm_omni/entrypoints/openai/serving_speech.py diff --git a/tests/entrypoints/openai_api/test_serving_speech.py b/tests/entrypoints/openai_api/test_serving_speech.py new file mode 100644 index 000000000..e2d33f317 --- /dev/null +++ b/tests/entrypoints/openai_api/test_serving_speech.py @@ -0,0 +1,267 @@ +# tests/entrypoints/openai/test_serving_speech.py +import logging +from inspect import Signature, signature +from unittest.mock import MagicMock, patch + +import numpy as np +import pytest +import torch +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from vllm_omni.entrypoints.openai.audio_utils_mixin import AudioMixin +from vllm_omni.entrypoints.openai.protocol.audio import CreateAudio +from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech +from vllm_omni.outputs import OmniRequestOutput + +logger = logging.getLogger(__name__) + + +class TestAudioMixin: + @pytest.fixture + def audio_mixin(self): + return AudioMixin() + + def test_stereo_to_mono_conversion(self, audio_mixin): + stereo_tensor = np.random.rand(24000, 2).astype(np.float32) + audio_obj = CreateAudio(audio_tensor=stereo_tensor) + + with ( + patch.object( + audio_mixin, "_apply_speed_adjustment", side_effect=lambda tensor, speed, sr: (tensor, sr) + ) as mock_speed, + patch("soundfile.write") as _, + ): + audio_mixin.create_audio(audio_obj) + + # Check that the tensor passed to speed adjustment is mono + mock_speed.assert_called_once() + adjusted_tensor = mock_speed.call_args[0][0] + assert len(adjusted_tensor) == 24000 + + @patch("librosa.effects.time_stretch") + def test_speed_adjustment(self, mock_time_stretch, audio_mixin): + mock_time_stretch.return_value = np.zeros(12000) + audio_tensor = np.random.rand(24000).astype(np.float32) + + adjusted_audio, _ = audio_mixin._apply_speed_adjustment(audio_tensor, speed=2.0, sample_rate=24000) + + mock_time_stretch.assert_called_with(y=audio_tensor, rate=2.0) + assert adjusted_audio.shape == (12000,) + + @patch("soundfile.write") + def test_unsupported_format_fallback(self, mock_write, audio_mixin, caplog): + audio_tensor = np.random.rand(24000).astype(np.float32) + # Use a format that is not in the list of supported formats + audio_obj = CreateAudio(audio_tensor=audio_tensor, response_format="vorbis") + + audio_mixin.create_audio(audio_obj) + + # Should fall back to 'wav' + mock_write.assert_called_once() + write_kwargs = mock_write.call_args.kwargs + assert write_kwargs["format"] == "WAV" + + def test_mono_audio_preservation(self, audio_mixin): + """Test that mono (1D) audio tensors are processed correctly and passed to writer.""" + mono_tensor = np.random.rand(24000).astype(np.float32) + audio_obj = CreateAudio(audio_tensor=mono_tensor) + + with patch("soundfile.write") as mock_write: + audio_mixin.create_audio(audio_obj) + + mock_write.assert_called_once() + # Verify the tensor passed to soundfile.write is the exact 1D tensor + output_tensor = mock_write.call_args[0][1] + assert output_tensor.ndim == 1 + assert output_tensor.shape == (24000,) + assert np.array_equal(output_tensor, mono_tensor) + + def test_stereo_audio_preservation(self, audio_mixin): + """Test that stereo (2D) audio tensors are processed correctly and preserved.""" + stereo_tensor = np.random.rand(24000, 2).astype(np.float32) + audio_obj = CreateAudio(audio_tensor=stereo_tensor) + + with patch("soundfile.write") as mock_write: + audio_mixin.create_audio(audio_obj) + + mock_write.assert_called_once() + # Verify the tensor passed to soundfile.write is the exact 2D tensor + output_tensor = mock_write.call_args[0][1] + assert output_tensor.ndim == 2 + assert output_tensor.shape == (24000, 2) + assert np.array_equal(output_tensor, stereo_tensor) + + def test_speed_adjustment_bypass(self, audio_mixin): + """Test that speed=1.0 bypasses the expensive librosa time stretching.""" + audio_tensor = np.random.rand(24000).astype(np.float32) + + with patch("librosa.effects.time_stretch") as mock_time_stretch: + # speed=1.0 should return immediately without calling librosa + result, _ = audio_mixin._apply_speed_adjustment(audio_tensor, speed=1.0, sample_rate=24000) + + mock_time_stretch.assert_not_called() + assert np.array_equal(result, audio_tensor) + + @patch("librosa.effects.time_stretch") + def test_speed_adjustment_stereo_handling(self, mock_time_stretch, audio_mixin): + """Test that speed adjustment is attempted on stereo inputs.""" + stereo_tensor = np.random.rand(24000, 2).astype(np.float32) + # Mock return value representing a sped-up version (half length) + mock_time_stretch.return_value = np.zeros((12000, 2), dtype=np.float32) + + result, _ = audio_mixin._apply_speed_adjustment(stereo_tensor, speed=2.0, sample_rate=24000) + + mock_time_stretch.assert_called_once() + # Ensure the stereo tensor was passed to librosa + call_args = mock_time_stretch.call_args + assert np.array_equal(call_args.kwargs["y"], stereo_tensor) + assert call_args.kwargs["rate"] == 2.0 + assert result.shape == (12000, 2) + + +# Helper to create mock model output for endpoint tests +def create_mock_audio_output_for_test( + request_id: str = "speech-mock-123", +) -> OmniRequestOutput: + 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, request_id: str, audio_tensor: torch.Tensor): + self.request_id = request_id + self.outputs = [MockCompletionOutput(index=0)] + self.multimodal_output = {"audio": audio_tensor} + self.finished = True + 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 + + num_samples = 24000 + audio_tensor = torch.sin(torch.linspace(0, 440 * 2 * torch.pi, num_samples)) + mock_request_output = MockRequestOutput(request_id=request_id, audio_tensor=audio_tensor) + + return OmniRequestOutput( + stage_id=0, + final_output_type="audio", + request_output=mock_request_output, + ) + + +@pytest.fixture +def test_app(): + # Mock the engine client + mock_engine_client = MagicMock() + mock_engine_client.errored = False + + async def mock_generate_fn(*args, **kwargs): + yield create_mock_audio_output_for_test(request_id=kwargs.get("request_id")) + + mock_engine_client.generate = MagicMock(side_effect=mock_generate_fn) + mock_engine_client.default_sampling_params_list = [{}] + + # Mock models to have an is_base_model method + mock_models = MagicMock() + mock_models.is_base_model.return_value = True + + mock_request_logger = MagicMock() + + speech_server = OmniOpenAIServingSpeech( + engine_client=mock_engine_client, + models=mock_models, + request_logger=mock_request_logger, + ) + + # Patch the signature of create_speech to remove 'raw_request' for FastAPI route introspection + original_create_speech = speech_server.create_speech + _ = MagicMock(side_effect=original_create_speech) + + sig = signature(original_create_speech) + + new_parameters = [param for name, param in sig.parameters.items() if name != "raw_request"] + + new_sig = Signature(parameters=new_parameters, return_annotation=sig.return_annotation) + + async def awaitable_patched_create_speech(*args, **kwargs): + return await original_create_speech(*args, **kwargs) + + awaitable_patched_create_speech.__signature__ = new_sig + speech_server.create_speech = awaitable_patched_create_speech + + app = FastAPI() + app.add_api_route("/v1/audio/speech", speech_server.create_speech, methods=["POST"], response_model=None) + + return app + + +@pytest.fixture +def client(test_app): + return TestClient(test_app) + + +class TestSpeechAPI: + def test_create_speech_success(self, client): + payload = { + "input": "Hello world", + "model": "tts-model", + "voice": "alloy", + "response_format": "wav", + } + response = client.post("/v1/audio/speech", json=payload) + assert response.status_code == 200 + assert response.headers["content-type"] == "audio/wav" + assert len(response.content) > 0 + + def test_create_speech_mp3_format(self, client): + payload = { + "input": "Hello world", + "model": "tts-model", + "voice": "alloy", + "response_format": "mp3", + } + response = client.post("/v1/audio/speech", json=payload) + assert response.status_code == 200 + assert response.headers["content-type"] == "audio/mpeg" + assert len(response.content) > 0 + + def test_create_speech_invalid_format(self, client): + payload = { + "input": "Hello world", + "model": "tts-model", + "voice": "alloy", + "response_format": "invalid_format", + } + response = client.post("/v1/audio/speech", json=payload) + assert response.status_code == 422 # Unprocessable Entity + + @patch("vllm_omni.entrypoints.openai.serving_speech.OmniOpenAIServingSpeech.create_audio") + def test_speed_parameter_is_used(self, mock_create_audio, test_app): + client = TestClient(test_app) + + mock_audio_response = MagicMock() + mock_audio_response.audio_data = b"dummy_audio" + mock_audio_response.media_type = "audio/wav" + mock_create_audio.return_value = mock_audio_response + + payload = { + "input": "This should be fast.", + "model": "tts-model", + "voice": "alloy", + "response_format": "wav", + "speed": 2.5, + } + client.post("/v1/audio/speech", json=payload) + + mock_create_audio.assert_called_once() + call_args = mock_create_audio.call_args[0] + audio_obj = call_args[0] + assert isinstance(audio_obj, CreateAudio) + assert audio_obj.speed == 2.5 diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 04f41c9ea..3e0042549 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -46,12 +46,14 @@ from vllm_omni.entrypoints.openai.image_api_utils import ( encode_image_base64, parse_size, ) +from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest from vllm_omni.entrypoints.openai.protocol.images import ( ImageData, ImageGenerationRequest, ImageGenerationResponse, ) from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat +from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech logger = init_logger(__name__) @@ -432,6 +434,10 @@ async def omni_init_app_state( log_error_stack=args.log_error_stack, ) + state.openai_serving_speech = OmniOpenAIServingSpeech( + engine_client, state.openai_serving_models, request_logger=request_logger + ) + state.enable_server_load_tracking = args.enable_server_load_tracking state.server_load_metrics = 0 @@ -440,6 +446,10 @@ def Omnichat(request: Request) -> OmniOpenAIServingChat | None: return request.app.state.openai_serving_chat +def Omnispeech(request: Request) -> OmniOpenAIServingSpeech | None: + return request.app.state.openai_serving_speech + + @router.post( "/v1/chat/completions", dependencies=[Depends(validate_json_request)], @@ -496,6 +506,28 @@ async def create_chat_completion(request: ChatCompletionRequest, raw_request: Re return StreamingResponse(content=generator, media_type="text/event-stream") +@router.post( + "/v1/audio/speech", + dependencies=[Depends(validate_json_request)], + responses={ + HTTPStatus.OK.value: {"content": {"audio/*": {}}}, + HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse}, + HTTPStatus.NOT_FOUND.value: {"model": ErrorResponse}, + HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse}, + }, +) +@with_cancellation +@load_aware_call +async def create_speech(request: OpenAICreateSpeechRequest, raw_request: Request): + handler = Omnispeech(raw_request) + if handler is None: + return base(raw_request).create_error_response(message="The model does not support Speech API") + try: + return await handler.create_speech(request, raw_request) + except Exception as e: + raise HTTPException(status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, detail=str(e)) from e + + # Image generation API endpoints diff --git a/vllm_omni/entrypoints/openai/audio_utils_mixin.py b/vllm_omni/entrypoints/openai/audio_utils_mixin.py new file mode 100644 index 000000000..13df32ebe --- /dev/null +++ b/vllm_omni/entrypoints/openai/audio_utils_mixin.py @@ -0,0 +1,93 @@ +from io import BytesIO + +import numpy as np +from vllm.logger import init_logger + +from vllm_omni.entrypoints.openai.protocol.audio import AudioResponse, CreateAudio + +try: + import soundfile +except ImportError: + soundfile = None + +try: + import librosa +except ImportError: + librosa = None + +logger = init_logger(__name__) + + +class AudioMixin: + """Mixin class to add audio-related utilities.""" + + def create_audio(self, audio_obj: CreateAudio) -> AudioResponse: + """Convert audio tensor to bytes in the specified format.""" + + audio_tensor = audio_obj.audio_tensor + sample_rate = audio_obj.sample_rate + response_format = audio_obj.response_format.lower() + stream_format = audio_obj.stream_format + base64_encode = audio_obj.base64_encode + speed = audio_obj.speed + + if stream_format != "audio": + raise ValueError(f"Unsupported stream format: {stream_format}") + + if soundfile is None: + raise ImportError( + "soundfile is required for audio generation. Please install it with: pip install soundfile" + ) + + if audio_tensor.ndim > 2: + raise ValueError( + f"Unsupported audio tensor dimension: {audio_tensor.ndim}. " + "Only mono (1D) and stereo (2D) are supported." + ) + + audio_tensor, sample_rate = self._apply_speed_adjustment(audio_tensor, speed, sample_rate) + + supported_formats = { + "wav": ("WAV", "audio/wav", {}), + "pcm": ("RAW", "audio/pcm", {"subtype": "PCM_16"}), + "flac": ("FLAC", "audio/flac", {}), + "mp3": ("MP3", "audio/mpeg", {}), + "aac": ("AAC", "audio/aac", {}), + "opus": ("OGG", "audio/ogg", {"subtype": "OPUS"}), + } + + if response_format not in supported_formats: + logger.warning(f"Unsupported response format '{response_format}', defaulting to 'wav'.") + response_format = "wav" + + soundfile_format, media_type, kwargs = supported_formats[response_format] + + with BytesIO() as buffer: + soundfile.write(buffer, audio_tensor, sample_rate, format=soundfile_format, **kwargs) + audio_data = buffer.getvalue() + + if base64_encode: + import base64 + + audio_data = base64.b64encode(audio_data).decode("utf-8") + + return AudioResponse(audio_data=audio_data, media_type=media_type) + + def _apply_speed_adjustment(self, audio_tensor: np.ndarray, speed: float, sample_rate: int): + """Apply speed adjustment to the audio tensor while preserving pitch.""" + if speed == 1.0: + return audio_tensor, sample_rate + + if librosa is None: + raise ImportError("librosa is required for speed adjustment. Please install it with: pip install librosa") + + try: + # librosa.effects.time_stretch requires a float audio tensor. + if not np.issubdtype(audio_tensor.dtype, np.floating): + audio_tensor = audio_tensor.astype(np.float32) + + stretched_audio = librosa.effects.time_stretch(y=audio_tensor, rate=speed) + return stretched_audio, sample_rate + except Exception as e: + logger.error(f"An error occurred during speed adjustment: {e}") + raise ValueError("Failed to apply speed adjustment.") from e diff --git a/vllm_omni/entrypoints/openai/protocol/audio.py b/vllm_omni/entrypoints/openai/protocol/audio.py new file mode 100644 index 000000000..267d06291 --- /dev/null +++ b/vllm_omni/entrypoints/openai/protocol/audio.py @@ -0,0 +1,42 @@ +from typing import Literal + +import numpy as np +from pydantic import BaseModel, Field, field_validator + + +class OpenAICreateSpeechRequest(BaseModel): + input: str + model: str | None = None + voice: Literal["alloy", "ash", "ballad", "coral", "echo", "fable", "onyx", "nova", "sage", "shimmer", "verse"] + instructions: str | None = None + response_format: Literal["wav", "pcm", "flac", "mp3", "aac", "opus"] + speed: float | None = Field( + default=1.0, + ge=0.25, + le=4.0, + ) + stream_format: Literal["sse", "audio"] | None = "audio" + + @field_validator("stream_format") + @classmethod + def validate_stream_format(cls, v: str) -> str: + if v == "sse": + raise ValueError("'sse' is not a supported stream_format yet. Please use 'audio'.") + return v + + +class CreateAudio(BaseModel): + audio_tensor: np.ndarray + sample_rate: int = 24000 + response_format: str = "wav" + speed: float = 1.0 + stream_format: Literal["sse", "audio"] | None = "audio" + base64_encode: bool = True + + class Config: + arbitrary_types_allowed = True + + +class AudioResponse(BaseModel): + audio_data: bytes | str + media_type: str diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index 9a8407d84..72012d926 100644 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -71,6 +71,8 @@ from vllm.tokenizers.mistral import ( from vllm.utils.collection_utils import as_list from vllm_omni.entrypoints.chat_utils import parse_chat_messages_futures +from vllm_omni.entrypoints.openai.audio_utils_mixin import AudioMixin +from vllm_omni.entrypoints.openai.protocol.audio import AudioResponse, CreateAudio from vllm_omni.outputs import OmniRequestOutput if TYPE_CHECKING: @@ -79,7 +81,7 @@ if TYPE_CHECKING: logger = init_logger(__name__) -class OmniOpenAIServingChat(OpenAIServingChat): +class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): """OpenAI-compatible chat serving for both LLM and Diffusion models. This class extends OpenAIServingChat to support: @@ -860,26 +862,21 @@ class OmniOpenAIServingChat(OpenAIServingChat): final_res = omni_outputs.request_output audio_tensor = final_res.multimodal_output["audio"].float().detach().cpu().numpy() - # Convert numpy array to WAV bytes and encode as base64 - if soundfile is None: - raise ImportError( - "soundfile is required for audio generation. Please install it with: pip install soundfile" - ) - - # Default sample rate for TTS models (typically 24000 Hz) - # You may need to adjust this based on your model's configuration - sample_rate = 24000 - # Ensure audio is 1D (flatten if needed) if audio_tensor.ndim > 1: audio_tensor = audio_tensor.flatten() - # Convert to WAV format and encode as base64 - with BytesIO() as buffer: - soundfile.write(buffer, audio_tensor, sample_rate, format="WAV") - wav_bytes = buffer.getvalue() + audio_obj = CreateAudio( + audio_tensor=audio_tensor, + sample_rate=24000, + response_format="wav", + speed=1.0, + stream_format="audio", + base64_encode=True, + ) - audio_base64 = base64.b64encode(wav_bytes).decode("utf-8") + audio_response: AudioResponse = self.create_audio(audio_obj) + audio_base64 = audio_response.audio_data # Generate unique ID for the audio audio_id = f"audio-{uuid.uuid4().hex[:16]}" diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py new file mode 100644 index 000000000..c6b87810e --- /dev/null +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -0,0 +1,80 @@ +import asyncio + +from fastapi import Request +from fastapi.responses import Response +from vllm.entrypoints.openai.serving_engine 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 +from vllm_omni.outputs import OmniRequestOutput + +logger = init_logger(__name__) + + +class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): + async def create_speech( + self, + request: OpenAICreateSpeechRequest, + raw_request: Request | None = None, + ): + """ + Create Speech API similar to OpenAI's API. + + See https://platform.openai.com/docs/api-reference/audio/createSpeech + for the API specification. This API mimics the OpenAI + Create Speech API. + + NOTE: This implementation does not currently support streaming audio + generation. The full audio is generated before the response is sent. + """ + + error_check_ret = await self._check_model(request) + if error_check_ret is not None: + logger.error("Error with model %s", error_check_ret) + return error_check_ret + + if self.engine_client.errored: + raise self.engine_client.dead_error + + request_id = f"speech-{random_uuid()}" + + try: + prompt = {"prompt": request.input} + sampling_params_list = self.engine_client.default_sampling_params_list + + generator = self.engine_client.generate( + prompt=prompt, request_id=request_id, sampling_params_list=sampling_params_list + ) + + final_output: OmniRequestOutput | None = None + async for res in generator: + final_output = res + + if final_output is None: + return self.create_error_response("No output generated from the model.") + + if final_output.final_output_type != "audio": + return self.create_error_response(f"Unexpected final output type: {final_output.final_output_type}") + + audio_tensor = final_output.request_output.multimodal_output["audio"].float().detach().cpu().numpy() + audio_obj = CreateAudio( + audio_tensor=audio_tensor, + sample_rate=24000, + response_format=request.response_format, + speed=request.speed, + stream_format=request.stream_format, + base64_encode=False, + ) + + audio_response: AudioResponse = self.create_audio(audio_obj) + return Response(content=audio_response.audio_data, media_type=audio_response.media_type) + + except asyncio.CancelledError: + return self.create_error_response("Client disconnected") + except ValueError as e: + # TODO: Use a vllm-specific Validation Error + return self.create_error_response(str(e)) + except Exception as e: + return self.create_error_response(f"{e} {e.__cause__}")