[Voxtral] Raise BadRequest for ref_audio (#7012)

Signed-off-by: Clodagh Walsh <clodaghwalsh17@gmail.com>

recipes/mistralai/Voxtral-TTS.md

@@ -178,7 +178,7 @@ mistralai/Voxtral-4B-TTS-2603 checkpoint. running the
 voice-cloning e2e test ([`Voice cloning (capability gated upstream)`](../../examples/offline_inference/text_to_speech/README.md#voice-cloning-capability-gated-upstream)
 ) or adding a `ref_audio` request parameter will fail
 because the public checkpoint does not provide the encoder weights needed to
-convert reference audio into conditioning features. The expected failure is: `RuntimeError: encode_waveforms requires encoder weights which are not available in the open-source checkpoint.`
+convert reference audio into conditioning features. The expected failure is a 400 BadRequest `Voice cloning with 'ref_audio' requires the full Voxtral checkpoint with encoder weights. The open-source variant only supports preset voices via the 'voice' parameter.`.
 
 ### 1x NVIDIA RTX A6000 48GB
 

tests/e2e/online_serving/test_voxtral_tts_expansion.py

@@ -12,11 +12,14 @@ import os
 os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
 
 import pytest
+from openai import BadRequestError as OpenAIBadRequestError
 
 from tests.helpers.mark import hardware_test
+from tests.helpers.media import generate_synthetic_audio
 from tests.helpers.runtime import OmniServerParams
 from tests.helpers.stage_config import get_deploy_config_path
 
+# Note this model is the public checkpoint without the encoder weights required for voice cloning
 MODEL = "mistralai/Voxtral-4B-TTS-2603"
 STAGE_CONFIG = get_deploy_config_path("voxtral_tts.yaml")
 EXTRA_ARGS = ["--trust-remote-code", "--enforce-eager", "--disable-log-stats"]
@@ -173,3 +176,21 @@ def test_speech_batches(omni_server, online_client) -> None:
             "timeout": 120.0,
         }
     )
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_speech_ref_audio(omni_server, online_client) -> None:
+    """Test ref_audio failure is handled as a BadRequest instead of a RuntimeError"""
+    audio_data = f"data:audio/wav;base64,{generate_synthetic_audio(5, 1)['base64']}"
+    with pytest.raises(OpenAIBadRequestError) as exc_info:
+        online_client.send_audio_speech_request(
+            {
+                "model": omni_server.model,
+                "input": "how are you",
+                "ref_audio": audio_data,
+                "language": "English",
+                "response_format": "wav",
+            }
+        )
+
+    assert exc_info.value.status_code == 400

tests/entrypoints/openai_api/test_serving_speech.py

@@ -5133,6 +5133,7 @@ class TestTTSAsyncOffloading:
 
     def test_prepare_speech_generation_awaits_voxtral_async(self, voxtral_server, mocker: MockerFixture):
         """Voxtral path in _prepare_speech_generation should call the async wrapper."""
+        voxtral_server._adapter._encoder_loaded = mocker.AsyncMock(return_value=False)
         voxtral_server._adapter._build_prompt_async = mocker.AsyncMock(
             return_value={
                 "prompt_token_ids": [1, 2, 3],
@@ -5288,6 +5289,7 @@ class TestTTSAsyncOffloading:
         """FINAL_ONLY streaming for async_chunk=False is scoped to qwen3_tts only."""
         voxtral_server.engine_client.model_config.async_chunk = False
         mocker.patch.object(voxtral_server._get_tts_adapter(), "validate", return_value=None)
+        voxtral_server._adapter._encoder_loaded = mocker.AsyncMock(return_value=False)
         voxtral_server._adapter._build_prompt_async = mocker.AsyncMock(
             return_value={
                 "prompt_token_ids": [1, 2, 3],

vllm_omni/entrypoints/openai/tts_adapters/voxtral.py

@@ -22,6 +22,7 @@ class VoxtralTTSAdapter(ARTTSAdapter):
     def __init__(self, ctx) -> None:
         super().__init__(ctx)
         self._build_prompt_async = make_async(self._build_prompt, executor=getattr(ctx.server, "_tts_executor", None))
+        self._encoder_loaded = None
 
     def _build_prompt(self, request: "OpenAICreateSpeechRequest") -> dict[str, Any]:
         """Build a Voxtral prompt for either a preset voice or inline reference audio."""
@@ -31,13 +32,24 @@ class VoxtralTTSAdapter(ARTTSAdapter):
         ref_audio = request.ref_audio
         if not request.voice and not ref_audio:
             raise ValueError("Voxtral requires either a voice name or ref_audio.")
-        if isinstance(ref_audio, str) and ref_audio.startswith("data:"):
-            _, _, ref_audio = ref_audio.partition(",")
+
+        if ref_audio is not None and not self._encoder_loaded:
+            validation_error = (
+                "Voice cloning with 'ref_audio' requires the full Voxtral checkpoint "
+                "with encoder weights. The open-source variant only supports "
+                "preset voices via the 'voice' parameter."
+            )
+            raise ValueError(validation_error)
+
         server = self.ctx.server
         if server._tts_tokenizer is None:
             from vllm.tokenizers import cached_tokenizer_from_config
 
             server._tts_tokenizer = cached_tokenizer_from_config(self.ctx.engine_client.model_config).instruct
+
+        if isinstance(ref_audio, str) and ref_audio.startswith("data:"):
+            _, _, ref_audio = ref_audio.partition(",")
+
         if request.voice is not None:
             tokenized = server._tts_tokenizer.encode_speech_request(
                 SpeechRequest(input=request.input, voice=request.voice)
@@ -45,6 +57,7 @@ class VoxtralTTSAdapter(ARTTSAdapter):
             prompt = tokens_input(prompt_token_ids=tokenized.tokens)
             prompt["additional_information"] = {"voice": [request.voice]}
             return prompt
+
         tokenized = server._tts_tokenizer.encode_speech_request(SpeechRequest(input=request.input, ref_audio=ref_audio))
         audio = tokenized.audios[0]
         return {
@@ -84,6 +97,12 @@ class VoxtralTTSAdapter(ARTTSAdapter):
     async def build(
         self, request: "OpenAICreateSpeechRequest", sampling_params_list: list, has_inline_ref_audio: bool
     ) -> PreparedRequest:
+        # Cache encoder_loaded property to avoid RPC call on every request with ref_audio
+        if self._encoder_loaded is None:
+            loaded = await self.ctx.server.engine_client.collective_rpc("encoder_loaded")
+            #  Flatten results before checking if encoder loaded property is set
+            self._encoder_loaded = any(sum(loaded, []))
+
         prompt = await self._build_prompt_async(request)
         return PreparedRequest(prompt=prompt, tts_params={}, model_type="voxtral_tts")
 

vllm_omni/model_executor/models/voxtral_tts/voxtral_tts.py

@@ -174,6 +174,11 @@ class VoxtralTTSForConditionalGeneration(
             return self.model.get_language_model()
         return self.model
 
+    def encoder_loaded(self):
+        if self.audio_tokenizer is not None:
+            return self.audio_tokenizer.encoder_loaded
+        return False
+
     def _enable_acoustic_transformer_cudagraph(self):
         """Initialize and capture CUDA graphs for compute_mm_logits."""
         if self.model_stage != "audio_generation" or not hasattr(self, "_vllm_config"):

vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_tokenizer.py

@@ -900,6 +900,10 @@ class VoxtralTTSAudioTokenizer(nn.Module):
         """List of size of each codebook"""
         return self.quantizer.codebook_sizes
 
+    @property
+    def encoder_loaded(self) -> bool:
+        return self._encoder_loaded
+
     def load_weight(self, weight: tuple[str, torch.Tensor]) -> str:
         params_dict = dict(self.named_parameters())
         name, loaded_weight = weight
@@ -985,7 +989,7 @@ class VoxtralTTSAudioTokenizer(nn.Module):
 
     def encode_waveforms(self, x: list[torch.Tensor]) -> list[torch.Tensor]:
         if not self._encoder_loaded:
-            raise RuntimeError(
+            raise ValueError(
                 "encode_waveforms requires encoder weights which are not available in the open-source checkpoint."
             )
         audio_codes = []

vllm_omni/worker/base.py

@@ -308,3 +308,16 @@ class OmniGPUWorkerBase(GPUWorker):
                     pass
             tid = task.task_id if hasattr(task, "task_id") else "unknown"
             return OmniACK(task_id=tid, status="ERROR", error_msg=str(e))
+
+    def encoder_loaded(self) -> bool:
+        """Check if encoder weights are loaded in the model.
+
+        This method is exposed via collective_rpc to check encoder availability
+        for models that support voice cloning with reference audio.
+
+        Returns:
+            bool: True if encoder weights are available, False otherwise.
+        """
+        if hasattr(self.model_runner, "model") and hasattr(self.model_runner.model, "encoder_loaded"):
+            return self.model_runner.model.encoder_loaded()
+        return False