feat(omni): add Nemotron Audex speech synthesis to /v1/audio/speech

components/src/dynamo/common/protocols/audio_protocol.py

@@ -80,6 +80,12 @@ class NvCreateAudioSpeechRequest(BaseModel):
     max_new_tokens: Optional[int] = None
     """Maximum tokens to generate (default: 2048)."""
 
+    # Audex-specific params
+    cfg_scale: Optional[float] = None
+    """Classifier-free guidance strength (Audex). 1.0 disables guidance; 1.5 is
+    the recommended TTS quality setting, and 3.0 the official TTA setting
+    (applied by default for text-to-audio)."""
+
     nvext: Optional[AudioNvExt] = None
     """NVIDIA extensions."""
 

components/src/dynamo/vllm/omni/audio_handler.py

@@ -25,6 +25,8 @@
 from dynamo.common.protocols import sanitize_media_passthrough
 from dynamo.common.protocols.audio_protocol import NvCreateAudioSpeechRequest
 from dynamo.common.utils.output_modalities import RequestType
+from dynamo.vllm.omni.audex import AudexRequestAdapter
+from dynamo.vllm.omni.utils import engine_model_stages, validate_audio_max_new_tokens
 
 logger = logging.getLogger(__name__)
 
@@ -54,9 +56,14 @@ class AudioGenerationHandler:
     """Handles audio/TTS request processing for the vLLM-Omni backend.
 
     Instantiated by OmniHandler during initialization and held as a
-    composition attribute (``self._audio_handler``).  This keeps
+    composition attribute (``self.audio``).  This keeps
     audio-specific logic (validation, prompt building, encoding) out
     of the orchestrator.
+
+    Model-specific request preparation lives outside this class: Audex in
+    ``omni.audex.AudexRequestAdapter`` (``self.audex``), Qwen3-TTS in the
+    ``_engine_inputs_tts`` path below. Everything else takes the generic
+    plain-text prompt.
     """
 
     def __init__(self, config, engine_client, media_output_fs, media_output_http_url):
@@ -65,6 +72,7 @@ def __init__(self, config, engine_client, media_output_fs, media_output_http_url
         self.media_output_fs = media_output_fs
         self.media_output_http_url = media_output_http_url
         self._tts_tokenizer: Any = None
+        self.audex = AudexRequestAdapter(config, engine_client)
 
         # Cache TTS capabilities from model config at init.
         self._tts_supported_speakers: set = self._load_supported_speakers()
@@ -128,32 +136,9 @@ def _is_tts_model(self) -> bool:
         Searches for a TTS model_stage in the engine's stage list,
         stage configs, or model config. Supports multiple vLLM-Omni versions.
         """
-        # Try stage_list
-        stage_list = getattr(self.engine_client, "stage_list", None)
-        if stage_list:
-            for stage in stage_list:
-                ms = getattr(stage, "model_stage", None)
-                logger.debug("_is_tts_model: stage=%s model_stage=%s", stage, ms)
-                if ms in _TTS_MODEL_STAGES:
-                    return True
-
-        # Try stage_configs
-        stage_configs = getattr(self.engine_client, "stage_configs", None)
-        if stage_configs:
-            for cfg in stage_configs:
-                engine_args = (
-                    cfg.get("engine_args", {})
-                    if isinstance(cfg, dict)
-                    else getattr(cfg, "engine_args", {})
-                )
-                ms = (
-                    engine_args.get("model_stage")
-                    if isinstance(engine_args, dict)
-                    else getattr(engine_args, "model_stage", None)
-                )
-                logger.debug("_is_tts_model: stage_config model_stage=%s", ms)
-                if ms in _TTS_MODEL_STAGES:
-                    return True
+        stages = engine_model_stages(self.engine_client)
+        if stages & _TTS_MODEL_STAGES:
+            return True
 
         # Try model_config.hf_config.model_type (universal fallback)
         try:
@@ -165,30 +150,40 @@ def _is_tts_model(self) -> bool:
             logger.debug("_is_tts_model: hf_config fallback failed: %s", e)
 
         logger.warning(
-            "_is_tts_model: could not detect TTS model. "
-            "stage_list=%s, stage_configs=%s",
-            stage_list is not None,
-            stage_configs is not None,
+            "_is_tts_model: could not detect TTS model. engine model stages=%s",
+            sorted(stages),
         )
         return False
 
     # -- Audio engine input construction --------------------------------------
 
-    async def build_engine_inputs(self, req: NvCreateAudioSpeechRequest):
+    async def build_engine_inputs(
+        self, req: NvCreateAudioSpeechRequest, request_id: str | None = None
+    ):
         """Build engine inputs for an audio/TTS request.
 
-        Two code paths (matching vLLM-Omni serving_speech.py):
+        Three code paths (matching vLLM-Omni serving_speech.py):
 
         * **TTS path** (Qwen3-TTS): ``prompt_token_ids`` +
           ``additional_information``.
+        * **Audex path** (delegated to ``self.audex``): literal ChatML prompt
+          priming codec generation, plus the CFG / RVQ-phase contracts on
+          stage-0 sampling params.
         * **Generic audio path** (MiMo-Audio, etc.): plain text prompt.
+
+        ``request_id`` is the final Dynamo request id; Audex uses it as the CFG
+        pair id that binds a guided request to its unconditional companion.
         """
         # Import here to avoid circular dependency
         from dynamo.vllm.omni.omni_handler import EngineInputs
 
         if not req.input or not req.input.strip():
             raise ValueError("Input text cannot be empty")
 
+        audex_model_type = self.audex.model_type()
+        if audex_model_type is not None:
+            return self._engine_inputs_audex(req, request_id, audex_model_type)
+
         output_format = (req.response_format or "wav").lower()
 
         # URL delivery, whole-file encoders, and speed adjustment require the
@@ -225,6 +220,40 @@ async def build_engine_inputs(self, req: NvCreateAudioSpeechRequest):
             stream_audio=stream_audio,
         )
 
+    # -- Audex ----------------------------------------------------------------
+
+    def _engine_inputs_audex(
+        self,
+        req: NvCreateAudioSpeechRequest,
+        request_id: str | None,
+        model_type: str,
+    ):
+        """Wrap the Audex adapter's prepared request in ``EngineInputs``.
+
+        Everything model-specific (validation, ChatML prompt, CFG/RVQ
+        contracts) belongs to ``omni.audex.AudexRequestAdapter``; this only adds
+        the response-delivery fields, which are the same for every audio model.
+        ``model_type`` is the resolution the caller already made to get here, so
+        the adapter does not repeat the stage lookup.
+
+        ``stream_audio`` is left off: per-payload chunk streaming has not been
+        validated for the Audex pipeline, so its requests take the aggregate
+        path and the response carries one complete waveform. Aggregation itself
+        follows the resolved output kind, not the model (see
+        ``utils.audio_output_is_cumulative``).
+        """
+        from dynamo.vllm.omni.omni_handler import EngineInputs
+
+        prepared = self.audex.prepare(req, request_id, model_type)
+        return EngineInputs(
+            prompt=OmniTextPrompt(prompt=prepared.prompt),
+            sampling_params_list=prepared.sampling_params_list,
+            request_type=RequestType.AUDIO_GENERATION,
+            response_format=req.data_source,
+            output_format=req.response_format,
+            speed=req.speed or 1.0,
+        )
+
     # -- Qwen3-TTS-specific helpers -------------------------------------------
 
     async def _engine_inputs_tts(
@@ -341,17 +370,7 @@ def _validate_tts_request(self, req: NvCreateAudioSpeechRequest) -> None:
                 f"(max {self.config.tts_max_instructions_length} characters)"
             )
 
-        if req.max_new_tokens is not None:
-            if req.max_new_tokens < self.config.tts_max_new_tokens_min:
-                raise ValueError(
-                    f"max_new_tokens must be at least "
-                    f"{self.config.tts_max_new_tokens_min}"
-                )
-            if req.max_new_tokens > self.config.tts_max_new_tokens_max:
-                raise ValueError(
-                    f"max_new_tokens cannot exceed "
-                    f"{self.config.tts_max_new_tokens_max}"
-                )
+        validate_audio_max_new_tokens(req.max_new_tokens, self.config)
 
     async def _resolve_ref_audio(self, ref_audio_str: str) -> tuple:
         """Download or decode reference audio for voice cloning (Base task)."""

components/src/dynamo/vllm/omni/omni_handler.py

@@ -54,6 +54,7 @@
     OutputFormatter,
 )
 from dynamo.vllm.omni.utils import (
+    audio_output_is_cumulative,
     build_image_generation_prompt,
     image_generation_negative_prompt_from_request,
     image_generation_sampling_overrides,
@@ -377,7 +378,7 @@ async def _generate_openai_mode(
 
         try:
             inputs = await self.build_engine_inputs(
-                parsed_request, request_type, image=image
+                parsed_request, request_type, image=image, request_id=request_id
             )
         except (ValueError, NotImplementedError, RuntimeError) as e:
             logger.error(f"Invalid request {request_id}: {e}")
@@ -417,8 +418,14 @@ async def _generate_openai_mode(
 
         previous_text = ""
         audio_stream_state = AudioStreamState() if inputs.stream_audio else None
+        # Read the coerced params, not the request's: the coercion above is what
+        # decides whether the engine emits disjoint deltas or whole-waveform
+        # snapshots, so aggregation has to follow its result rather than the
+        # model's identity.
         audio_aggregate_state = (
-            AudioAggregateState()
+            AudioAggregateState(
+                cumulative=audio_output_is_cumulative(inputs.sampling_params_list)
+            )
             if inputs.request_type == RequestType.AUDIO_GENERATION
             and not inputs.stream_audio
             else None
@@ -608,6 +615,7 @@ async def build_engine_inputs(
         ],
         request_type: RequestType,
         image: PIL.Image.Image | None = None,
+        request_id: str | None = None,
     ) -> EngineInputs:
         """Convert a parsed request into AsyncOmni engine inputs.
 
@@ -616,6 +624,8 @@ async def build_engine_inputs(
                 for image/video/audio requests, or a raw dict for chat completions.
             request_type: The RequestType determined by parse_request_type.
             image: Pre-loaded PIL Image for I2V requests (from input_reference).
+            request_id: Final request id, used by audio models (Audex) that bind
+                per-request state such as the CFG pair id to it.
 
         Returns:
             EngineInputs ready for engine_client.generate().
@@ -631,7 +641,9 @@ async def build_engine_inputs(
             return self._engine_inputs_from_video(parsed_request, image=image)
         elif request_type == RequestType.AUDIO_GENERATION:
             assert isinstance(parsed_request, NvCreateAudioSpeechRequest)
-            return await self.audio.build_engine_inputs(parsed_request)
+            return await self.audio.build_engine_inputs(
+                parsed_request, request_id=request_id
+            )
 
         raise ValueError(f"Unknown request type: {request_type}")
 

components/src/dynamo/vllm/omni/output_formatter.py

@@ -58,6 +58,16 @@ class AudioAggregateState:
     emitted_chunks: int = 0
     num_channels: int | None = None
     channel_axis: int | None = None
+    cumulative: bool = False
+    """Each payload is a snapshot of the whole waveform decoded so far.
+
+    Set from the output kind the engine was actually given, not from the model:
+    ``RequestOutputKind.CUMULATIVE`` consolidates the accumulated audio on every
+    step and drains nothing, so the snapshots must be de-duplicated to the
+    longest one rather than concatenated, while ``DELTA`` drains what it emits
+    and yields disjoint pieces that must all be kept. See
+    ``utils.audio_output_is_cumulative``, which the handler uses to fill this in.
+    """
 
 
 class TextFormatter:
@@ -325,11 +335,19 @@ async def format(
 
         try:
             start_time = time.time()
+            # A cumulative payload already carries the earlier frames, so it is
+            # taken whole and de-duplicated in _append_audio_chunk; tracking the
+            # newly appended entries would keep only the first snapshot.
+            chunk_state: AudioStreamState | AudioAggregateState | None
+            if stream_state is not None:
+                chunk_state = stream_state
+            elif aggregate_state is not None and not aggregate_state.cumulative:
+                chunk_state = aggregate_state
+            else:
+                chunk_state = None
+
             audio_np, sample_rate = self._extract_audio_tensor(
-                mm_output,
-                chunk_state=(
-                    stream_state if stream_state is not None else aggregate_state
-                ),
+                mm_output, chunk_state=chunk_state
             )
             if audio_np.size == 0:
                 return None
@@ -420,7 +438,15 @@ def _append_audio_chunk(
             expected_channel_axis=state.channel_axis,
         )
         self._validate_audio_metadata(state, sample_rate, num_channels, channel_axis)
-        state.chunks.append(audio_np)
+        if not state.cumulative:
+            state.chunks.append(audio_np)
+            return
+
+        # Snapshots only grow, but keep the longest rather than the latest so a
+        # truncated trailing payload cannot drop already-decoded audio.
+        if state.chunks and state.chunks[0].shape[-1] >= audio_np.shape[-1]:
+            return
+        state.chunks = [audio_np]
 
     @staticmethod
     def _validate_audio_metadata(
@@ -461,6 +487,13 @@ def _extract_audio_tensor(
 
         if isinstance(audio_val, list):
             if chunk_state is not None:
+                # Slicing by a running count assumes each list *extends* the
+                # previous one. That holds only while the engine does not drain
+                # what it emits: a draining (``DELTA``) stage restarts its list
+                # at index 0, and this would then keep just the tail. Audex
+                # takes the delta path but yields bare tensors, which skip this
+                # branch entirely. If a stage ever emits multi-entry lists under
+                # DELTA, the count has to be per-payload rather than running.
                 new_audio = audio_val[chunk_state.emitted_chunks :]
                 chunk_state.emitted_chunks = len(audio_val)
                 audio_val = new_audio