From b25c563f9f3061386e17e424f8f00347695e8449 Mon Sep 17 00:00:00 2001
From: yinpeiqi <yinpeiqi809@gmail.com>
Date: Thu, 27 Nov 2025 16:49:04 +0700
Subject: [PATCH] add preprocess mixin

---
 .../models/qwen2_5_omni/qwen2_5_omni.py       | 11 +++++---
 vllm_omni/model_executor/preprocess_mixin.py  | 27 +++++++++++++++++++
 vllm_omni/worker/gpu_ar_model_runner.py       |  4 +--
 3 files changed, 36 insertions(+), 6 deletions(-)
 create mode 100644 vllm_omni/model_executor/preprocess_mixin.py

diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py
index d3acaebf3..b102b8ffb 100644
--- a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py
+++ b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py
@@ -34,6 +34,7 @@ from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker import (
 )
 from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights, split_list_into_ranges
 from vllm_omni.model_executor.models.vision import get_llm_pos_ids_for_vision
+from vllm_omni.model_executor.preprocess_mixin import PreprocessMixin
 
 TALKER_CODEC_EOS_TOKEN_ID = 8294
 TALKER_CODEC_BOS_TOKEN_ID = 8293
@@ -56,11 +57,12 @@ logger = init_logger(__name__)
     dummy_inputs=Qwen2_5OmniThinkerDummyInputsBuilder,
 )
 class Qwen2_5OmniForConditionalGeneration(
-    nn.Module, SupportsMultiModal, SupportsPP, Qwen2_5OmniConditionalGenerationMixin, SupportsMRoPE
+    nn.Module, SupportsMultiModal, SupportsPP, SupportsMRoPE,
+    Qwen2_5OmniConditionalGenerationMixin, PreprocessMixin
 ):
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__()
-        self.has_process = False
+        self.has_preprocess = False
         self.have_multimodal_outputs = True
         config: Qwen2_5OmniConfig = vllm_config.model_config.hf_config
         multimodal_config = vllm_config.model_config.multimodal_config
@@ -91,8 +93,9 @@ class Qwen2_5OmniForConditionalGeneration(
             self.token2wav = None
 
         elif self.model_stage == "talker":
-            # temporal use this prepresation
-            self.has_process = True
+            # register the process function for the talker stage
+            self.has_preprocess = True
+            self.set_preprocess(self.process_fn)
             self.thinker = None
             # Initialize talker model wrapper (handles projection + LM)
             self.talker = init_vllm_registered_model(
diff --git a/vllm_omni/model_executor/preprocess_mixin.py b/vllm_omni/model_executor/preprocess_mixin.py
new file mode 100644
index 000000000..15027c153
--- /dev/null
+++ b/vllm_omni/model_executor/preprocess_mixin.py
@@ -0,0 +1,27 @@
+from collections.abc import Callable
+import torch
+
+class PreprocessMixin:
+    """
+    Mixin class for all stages in the Omni model.
+    """
+    def set_preprocess(self, preprocess_fn: Callable) -> None:
+        """
+        Set a preprocess function for the stage.
+        Args:
+            preprocess_fn: The preprocess function to register.
+        """
+        self.preprocess = preprocess_fn
+
+    def preprocess(self,
+                input_ids: torch.Tensor,
+                inputs_embeds: torch.Tensor,
+                **input_dict: object
+    ) -> tuple[torch.Tensor, torch.Tensor, dict]:
+        """
+        Process the input_ids and inputs_embeds for the given input_dict.
+        Returns the processed input_ids, inputs_embeds, and the input_dict.
+
+        If the stage don't applicable, return the original input_ids, inputs_embeds, and an empty dict.
+        """
+        return input_ids, inputs_embeds, {}
diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py
index 84d312cf0..a0f6780ca 100644
--- a/vllm_omni/worker/gpu_ar_model_runner.py
+++ b/vllm_omni/worker/gpu_ar_model_runner.py
@@ -172,7 +172,7 @@ class GPUARModelRunner(OmniGPUModelRunner):
             encoder_inputs = self._extract_encoder_inputs(scheduler_output)
             model_kwargs.update(encoder_inputs)
 
-        if self.model.has_process:
+        if self.model.has_preprocess:
             # Overlay custom prompt_embeds per request for the prompt portion;
             # collect additional_information (tensor/list) for prefill portion only
             for req_index, req_id in enumerate(self.input_batch.req_ids):
@@ -186,7 +186,7 @@ class GPUARModelRunner(OmniGPUModelRunner):
                 span_len = int(e) - int(s)
 
                 # call the custom process function
-                req_input_ids, req_embeds, update_dict = self.model.process_fn(
+                req_input_ids, req_embeds, update_dict = self.model.preprocess(
                     input_ids=input_ids[s:e],
                     inputs_embeds=inputs_embeds[s:e],
                     **req_infos