add openai create speech endpoint (#305)

Signed-off-by: Bhanu068 <voutharoja.bhanu06@gmail.com>
Signed-off-by: Bhanu Prakash Voutharoja <59905694+Bhanu068@users.noreply.github.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>

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

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
 
 

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

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

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

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__}")