[1/N][Refactor] Clean up dead code in output processor (#1579)

Signed-off-by: gcanlin <canlinguosdu@gmail.com>

vllm_omni/engine/output_processor.py

@@ -1,4 +1,3 @@
-from collections.abc import Callable
 from typing import Any
 
 import numpy as np
@@ -60,9 +59,11 @@ class OmniRequestState(RequestState):
                 incoming: dict[str, Any] = {}
                 target_key = self.mm_type or "hidden"
 
-                # Iterate directly without unnecessary dict copy
                 for k, v in payload.items():
-                    # Optional remap: if producer used "model_outputs" or "hidden", rename to mm_type
+                    # Normalize producer keys to the modality name.
+                    # AR runners produce {"hidden": ...} and generation
+                    # runners produce {"model_outputs": ...}; remap both
+                    # to the semantic modality key (e.g. "audio", "latent").
                     if k == "model_outputs":
                         k = target_key
                     elif k == "hidden" and target_key != "hidden":
@@ -247,16 +248,17 @@ class OmniRequestState(RequestState):
 
 
 class MultimodalOutputProcessor(VLLMOutputProcessor):
-    """Handles multimodal output processing by normalizing EngineCoreOutput
-    before delegating to the base vLLM OutputProcessor.
+    """Handles multimodal output processing by capturing pooling_output
+    from EngineCoreOutput and accumulating it as multimodal tensors,
+    before delegating to the base vLLM OutputProcessor for text handling.
 
-    Strategy:
-    - Route by EngineCoreOutput.output_type when present
-      ("image", "text+image", "latents", "text").
-    - Fallback to pooling/text heuristics when output_type is absent.
-    - Mutate EngineCoreOutput in-place to ensure vLLM's base processor can
-      produce the correct RequestOutput/PoolingRequestOutput.
-    - Allow custom per-modality handlers via register_handler().
+    The actual data flow is:
+    1. For each EngineCoreOutput with pooling_output and a detokenizer:
+       - Capture pooling_output into OmniRequestState.add_multimodal_tensor()
+       - Clear eco.pooling_output to force text path in base processor
+    2. Base vLLM OutputProcessor handles text detokenization
+    3. On finish, _consolidate_multimodal_tensors() concatenates accumulated tensors
+    4. _new_completion_output() attaches mm_accumulated to CompletionOutput
     """
 
     def __init__(
@@ -273,27 +275,12 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
             log_stats: Whether to log statistics
             stream_interval: Stream interval for output generation
             engine_core_output_type: Optional output type specification
-                (e.g., "image", "audio", "latents"). Used to route outputs
-                to appropriate processors. If None, output type is inferred.
+                (e.g., "image", "audio", "latent"). Used to tag multimodal
+                outputs with the correct modality key.
         """
         super().__init__(tokenizer=tokenizer, log_stats=log_stats, stream_interval=stream_interval)
-        self.output_handlers: dict[str, Callable[[EngineCoreOutput], None]] = {}
-        self._reqid_to_mm_type: dict[str, str] = {}
         self.engine_core_output_type = engine_core_output_type
 
-    def register_handler(self, modality: str, handler: Callable[[EngineCoreOutput], None]) -> None:
-        """Register a custom handler for a specific modality.
-
-        Allows custom processing logic for specific output modalities.
-        The handler is called before default processing for outputs
-        matching the specified modality.
-
-        Args:
-            modality: Modality name (e.g., "image", "audio", "latents")
-            handler: Callable that takes an EngineCoreOutput and processes it
-        """
-        self.output_handlers[modality.lower()] = handler
-
     def add_request(
         self,
         request: EngineCoreRequest,
@@ -346,20 +333,13 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
         engine_core_timestamp: float | None = None,
         iteration_stats: IterationStats | None = None,
     ) -> OutputProcessorOutput:
-        self._reqid_to_mm_type.clear()
         for eco in engine_core_outputs:
-            mm_type = (self.engine_core_output_type or "").lower()
-            if mm_type:
-                self._reqid_to_mm_type[eco.request_id] = mm_type
-            self._route_and_normalize(eco)
             req_state = self.request_states.get(eco.request_id)
             if req_state is None or not isinstance(req_state, OmniRequestState):
                 continue
             if eco.pooling_output is not None and req_state.detokenizer is not None:
-                req_state.add_multimodal_tensor(
-                    eco.pooling_output,
-                    (getattr(eco, "output_type", self.engine_core_output_type) or "").lower(),
-                )
+                mm_type = (getattr(eco, "output_type", self.engine_core_output_type) or "").lower()
+                req_state.add_multimodal_tensor(eco.pooling_output, mm_type)
                 # Force text path in base processor for multimodal outputs.
                 eco.pooling_output = None
 
@@ -368,105 +348,3 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
             engine_core_timestamp=engine_core_timestamp,
             iteration_stats=iteration_stats,
         )
-
-    # ---- routing helpers ----
-    def _route_and_normalize(self, eco: EngineCoreOutput) -> None:
-        output_type = (getattr(eco, "output_type", self.engine_core_output_type) or "").lower()
-
-        # Custom handler first (if registered)
-        if output_type in self.output_handlers:
-            try:
-                self.output_handlers[output_type](eco)
-                # Fall through to default fixups in case the handler left gaps
-            except Exception:
-                logger.exception("Error in custom output handler for %s", output_type)
-
-        if output_type == "image":
-            self._process_image_output(eco)
-        elif output_type in ("text+image", "text,image", "image+text"):
-            self._process_text_image_output(eco)
-        elif output_type in ("latents", "latent"):
-            self._process_latents_output(eco)
-        elif output_type in ("audio", "speech"):
-            self._process_audio_output(eco)
-        elif output_type == "text":
-            self._process_text_output(eco)
-        else:
-            # Fallback heuristic
-            if eco.pooling_output is not None:
-                self._process_pooling_output(eco)
-            else:
-                self._process_text_output(eco)
-
-    # ---- modality processors ----
-    def _process_image_output(self, eco: EngineCoreOutput) -> None:
-        """Ensure image tensors are surfaced via pooling_output for vLLM."""
-        if eco.pooling_output is None:
-            tensor = self._extract_from_multimodal_outputs(eco, keys=("image", "images", "pixel_values", "pixels"))
-            if tensor is not None:
-                eco.pooling_output = tensor
-
-    def _process_text_image_output(self, eco: EngineCoreOutput) -> None:
-        """Allow text+image outputs. Text path stays as new_token_ids;
-        image/latents route via pooling_output."""
-        # Preserve text tokens as-is; ensure pooling_output carries image/latents
-        if eco.pooling_output is None:
-            tensor = self._extract_from_multimodal_outputs(
-                eco,
-                keys=(
-                    "image",
-                    "images",
-                    "pixel_values",
-                    "pixels",
-                    "latent",
-                    "latents",
-                    "z",
-                ),
-            )
-            if tensor is not None:
-                eco.pooling_output = tensor
-
-    def _process_latents_output(self, eco: EngineCoreOutput) -> None:
-        """Ensure latent tensors are surfaced via pooling_output."""
-        if eco.pooling_output is None:
-            tensor = self._extract_from_multimodal_outputs(eco, keys=("latent", "latents", "z", "posterior"))
-            if tensor is not None:
-                eco.pooling_output = tensor
-
-    def _process_audio_output(self, eco: EngineCoreOutput) -> None:
-        """Ensure audio tensors are surfaced via pooling_output."""
-        if eco.pooling_output is None:
-            tensor = self._extract_from_multimodal_outputs(
-                eco, keys=("audio", "audios", "wav", "waveform", "audio_pcm", "pcm")
-            )
-            if tensor is not None:
-                eco.pooling_output = tensor
-
-    def _process_text_output(self, eco: EngineCoreOutput) -> None:
-        """No-op; base processor will detokenize new_token_ids → text."""
-        return
-
-    def _process_pooling_output(self, eco: EngineCoreOutput) -> None:
-        """Optional sanity checks for pooling tensor."""
-        if eco.pooling_output is None:
-            return
-        if not isinstance(eco.pooling_output, torch.Tensor):
-            # Best-effort: convert to tensor if it's a list/ndarray-like
-            try:
-                eco.pooling_output = torch.as_tensor(eco.pooling_output)
-            except Exception:
-                pass
-
-    def _extract_from_multimodal_outputs(self, eco: EngineCoreOutput, keys: tuple[str, ...]) -> torch.Tensor | None:
-        mm = getattr(eco, "multimodal_outputs", None)
-        if not isinstance(mm, dict):
-            return None
-        for k in keys:
-            v = mm.get(k)
-            if isinstance(v, torch.Tensor):
-                return v
-        # Try the first tensor in the dict as a fallback
-        for v in mm.values():
-            if isinstance(v, torch.Tensor):
-                return v
-        return None

vllm_omni/entrypoints/openai/serving_speech.py

@@ -418,7 +418,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             mm = getattr(ro, "multimodal_output", None) if ro else None
         if not mm:
             return None, None
-        key = "audio" if "audio" in mm else ("model_outputs" if "model_outputs" in mm else None)
+        key = "audio" if "audio" in mm else None
         return mm, key
 
     def _build_tts_params(self, request: OpenAICreateSpeechRequest) -> dict[str, Any]: