refactor(inputs): build the Omni renderer as a subclass instead of a proxy

OmniRenderer wrapped the upstream renderer, forwarded unknown attributes via
__getattr__, and re-bound eight BaseRenderer methods (render_cmpl, render_chat,
process_for_engine, _process_enc_dec and their async twins) onto the wrapper
so that they would reach Omni's _process_singleton. That allow-list hid any
concrete-renderer override of those methods and would silently skip the
Omni extras for any future BaseRenderer entry point not on the list.

Make OmniRenderer a mixin and create the renderer as a real subclass of the
upstream class instead:

- omni_renderer_cls(base_cls) builds (and caches) Omni<Base> = (OmniRenderer,
  base_cls); only _process_singleton / _process_singleton_async are extended,
  everything else is the upstream renderer unchanged.
- build_omni_renderer(vllm_config) resolves the class the way upstream
  renderer_from_config does (cached_tokenizer_from_config,
  tokenizer_args_from_config, RENDERER_REGISTRY.load_renderer_cls) and
  instantiates the Omni subclass; renderer_cls/tokenizer can be passed to
  skip the resolution.
- build_stage0_input_processor passes the renderer to InputProcessor up
  front instead of replacing it afterwards; the skip-tokenizer branch keeps
  _build_token_only_renderer and never resolves a tokenizer.

Rebinding an existing instance's __class__ was rejected: CPython refuses it
for HfRenderer ("object layout differs"), so the subclass is constructed.

Tests: tests/inputs/test_preprocess.py builds the renderer with
omni_renderer_cls and adds checks that the subclass is created once, that a
concrete renderer's render_cmpl override is no longer suppressed, and that
build_omni_renderer resolves like upstream or honours an explicit class;
tests/engine/test_async_omni_engine_stage_init.py asserts the stage-0
processor receives the built subclass instance and that the skip-tokenizer
branch never calls build_omni_renderer. Two pre-existing mypy complaints in
test_preprocess.py are fixed with annotations.

Validation: ruff, pre-commit (incl. mypy) on the four files; the modules
import and omni_renderer_cls(HfRenderer) builds; tests/inputs (32),
tests/engine/test_async_omni_engine_stage_init.py (55) and tests/entrypoints
(548) pass on CPU.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01MpXe2LeRrUQGjZfKFaocMZ
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>

vllm_omni/inputs/preprocess.py

@@ -1,34 +1,46 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Omni renderer: the upstream renderer class with Omni engine-input routing.
 
-from typing import Any
+``OmniRenderer`` is a mixin. ``omni_renderer_cls`` puts it in front of a
+concrete ``BaseRenderer`` subclass, so the Omni instance *is* an upstream
+renderer (tokenizer, executors, multimodal cache and every public entry point
+come from upstream unchanged) and only ``_process_singleton`` /
+``_process_singleton_async`` are extended. ``build_omni_renderer`` resolves
+the concrete class the same way upstream ``renderer_from_config`` does.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, cast
 
 from vllm.renderers import BaseRenderer
+from vllm.renderers.registry import RENDERER_REGISTRY
+from vllm.tokenizers.registry import cached_tokenizer_from_config, tokenizer_args_from_config
+
+if TYPE_CHECKING:
+    from vllm.config import VllmConfig
+    from vllm.tokenizers import TokenizerLike
 
 
 class OmniRenderer:
-    """Delegate tokenization/chat rendering while extending engine-input routing.
+    """Mixin that routes Omni prompts through the upstream renderer.
 
-    Keep the original renderer's tokenizer, executors, and multimodal caches.
-    The upstream pipeline methods run on this adapter so both sync and async
-    entrypoints reach the Omni singleton processing below.
+    It has no base class of its own; ``omni_renderer_cls`` prepends it to a
+    concrete ``BaseRenderer`` subclass, so ``self`` and ``super()`` are that
+    renderer at runtime (hence the casts below).
+
+    Two extensions over upstream ``_process_singleton``:
+
+    * a token prompt that carries ``mm_processor_kwargs`` but no
+      ``multi_modal_data`` still goes through the multimodal processor (AR
+      image-generation models such as GLM-Image read their target size from
+      those kwargs);
+    * Omni-only prompt keys (``prompt``, ``cache_salt``,
+      ``additional_information``, ``model_intermediate_buffer``) are copied
+      onto the engine input.
     """
 
-    render_cmpl = BaseRenderer.render_cmpl
-    render_cmpl_async = BaseRenderer.render_cmpl_async
-    render_chat = BaseRenderer.render_chat
-    render_chat_async = BaseRenderer.render_chat_async
-    process_for_engine = BaseRenderer.process_for_engine
-    process_for_engine_async = BaseRenderer.process_for_engine_async
-    _process_enc_dec = BaseRenderer._process_enc_dec
-    _process_enc_dec_async = BaseRenderer._process_enc_dec_async
-
-    def __init__(self, renderer: BaseRenderer) -> None:
-        self._renderer = renderer
-
-    def __getattr__(self, name: str) -> Any:
-        return getattr(self._renderer, name)
-
     @staticmethod
     def _with_omni_extras(inputs, prompt):
         for key in ("prompt", "cache_salt", "additional_information", "model_intermediate_buffer"):
@@ -36,9 +48,14 @@ class OmniRenderer:
                 inputs[key] = prompt[key]
         return inputs
 
+    @staticmethod
+    def _routes_no_media_kwargs(prompt) -> bool:
+        return "prompt_embeds" not in prompt and "mm_processor_kwargs" in prompt and not prompt.get("multi_modal_data")
+
     def _process_singleton(self, prompt, *, skip_mm_cache: bool = False):
-        if "prompt_embeds" not in prompt and "mm_processor_kwargs" in prompt and not prompt.get("multi_modal_data"):
-            inputs = self._renderer._process_multimodal(
+        renderer = cast(BaseRenderer, self)
+        if self._routes_no_media_kwargs(prompt):
+            inputs = renderer._process_multimodal(
                 prompt["prompt_token_ids"],
                 {},
                 mm_processor_kwargs=prompt["mm_processor_kwargs"],
@@ -46,12 +63,13 @@ class OmniRenderer:
                 skip_mm_cache=skip_mm_cache,
             )
         else:
-            inputs = self._renderer._process_singleton(prompt, skip_mm_cache=skip_mm_cache)
+            inputs = cast(BaseRenderer, super())._process_singleton(prompt, skip_mm_cache=skip_mm_cache)
         return self._with_omni_extras(inputs, prompt)
 
     async def _process_singleton_async(self, prompt, *, skip_mm_cache: bool = False):
-        if "prompt_embeds" not in prompt and "mm_processor_kwargs" in prompt and not prompt.get("multi_modal_data"):
-            inputs = await self._renderer._process_multimodal_async(
+        renderer = cast(BaseRenderer, self)
+        if self._routes_no_media_kwargs(prompt):
+            inputs = await renderer._process_multimodal_async(
                 prompt["prompt_token_ids"],
                 {},
                 mm_processor_kwargs=prompt["mm_processor_kwargs"],
@@ -59,5 +77,42 @@ class OmniRenderer:
                 skip_mm_cache=skip_mm_cache,
             )
         else:
-            inputs = await self._renderer._process_singleton_async(prompt, skip_mm_cache=skip_mm_cache)
+            inputs = await cast(BaseRenderer, super())._process_singleton_async(prompt, skip_mm_cache=skip_mm_cache)
         return self._with_omni_extras(inputs, prompt)
+
+
+_OMNI_RENDERER_CLASSES: dict[type[BaseRenderer], type[BaseRenderer]] = {}
+
+
+def omni_renderer_cls(base_cls: type[BaseRenderer]) -> type[BaseRenderer]:
+    """Return the Omni subclass of ``base_cls`` (created once per base class)."""
+    if issubclass(base_cls, OmniRenderer):
+        return base_cls
+    cls = _OMNI_RENDERER_CLASSES.get(base_cls)
+    if cls is None:
+        name = f"Omni{base_cls.__name__}"
+        cls = type(name, (OmniRenderer, base_cls), {"__module__": __name__, "__qualname__": name})
+        _OMNI_RENDERER_CLASSES[base_cls] = cls
+    return cls
+
+
+def build_omni_renderer(
+    vllm_config: VllmConfig,
+    *,
+    renderer_cls: type[BaseRenderer] | None = None,
+    tokenizer: TokenizerLike | None = None,
+) -> BaseRenderer:
+    """Build the Omni renderer for ``vllm_config``.
+
+    With ``renderer_cls`` unset this mirrors upstream ``renderer_from_config``:
+    the tokenizer and renderer mode come from the model config and the class
+    from ``RENDERER_REGISTRY``. Pass ``renderer_cls`` (and optionally
+    ``tokenizer``) to skip that resolution, e.g. for tokenizer-less stages.
+    """
+    model_config = vllm_config.model_config
+    if renderer_cls is None:
+        if tokenizer is None:
+            tokenizer = cached_tokenizer_from_config(model_config)
+        renderer_mode, *_ = tokenizer_args_from_config(model_config)
+        renderer_cls = RENDERER_REGISTRY.load_renderer_cls(renderer_mode)
+    return omni_renderer_cls(renderer_cls)(vllm_config, tokenizer)