supplement some imports

vllm_omni/entrypoints/omni_llm.py

@@ -17,6 +17,8 @@ import vllm.envs as envs
 from vllm_omni.entrypoints.utils import load_stage_configs_from_model
 from vllm_omni.entrypoints.stage_manager import Stage
 from vllm_omni.engine.arg_utils import OmniEngineArgs
+from vllm_omni.engine.output_processor import MultimodalOutputProcessor
+from vllm_omni.engine.processor import OmniProcessor
 from vllm_omni.outputs import OmniRequestOutput
 
 
@@ -134,6 +136,10 @@ class StageLLM(LLM):
         # Create the Engine (autoselects V0 vs V1)
         self.llm_engine = LLMEngine.from_engine_args(
             engine_args=engine_args, usage_context=UsageContext.LLM_CLASS)
+        self.llm_engine.output_processor = MultimodalOutputProcessor(tokenizer=self.llm_engine.tokenizer, 
+                                                                    log_stats=self.llm_engine.log_stats)
+        self.llm_engine.processor = OmniProcessor(vllm_config=self.llm_engine.vllm_config,
+                                                  tokenizer=self.llm_engine.tokenizer)
         self.engine_class = type(self.llm_engine)
 
         self.request_counter = Counter()

vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml

@@ -4,6 +4,8 @@ stage_args:
     engine_args:
       model_stage: thinker
       model_arch: Qwen2_5OmniForConditionalGeneration
+      worker_cls: vllm_omni.worker.AR_gpu_worker.ARGPUWorker
+      scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler
       gpu_memory_utilization: 0.32
       enforce_eager: true  # need to discuss
       trust_remote_code: true
@@ -14,16 +16,21 @@ stage_args:
     engine_args:
       model_stage: talker
       model_arch: Qwen2_5OmniForConditionalGeneration
+      worker_cls: vllm_omni.worker.AR_gpu_worker.ARGPUWorker
+      scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler
       gpu_memory_utilization: 0.32
       enforce_eager: true
       trust_remote_code: true
       engine_output_type: latent
     engine_input_source: [0]
+    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen2_5_omni.thinker2talker
 
   - stage_id: 2
     engine_args:
       model_stage: code2wav
       model_arch: Qwen2_5OmniForConditionalGeneration
+      worker_cls: vllm_omni.worker.diffusion_gpu_worker.DiffusionGPUWorker
+      scheduler_cls: vllm_omni.core.sched.diffusion_scheduler.DiffusionScheduler
       gpu_memory_utilization: 0.3
       enforce_eager: true
       trust_remote_code: true

vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py

@@ -0,0 +1,32 @@
+from vllm.inputs import TextPrompt
+from typing import Union
+import torch
+from vllm_omni.inputs.data import OmniTokensPrompt
+
+def thinker2talker(stage_list, engine_input_source, prompt: Union[OmniTokensPrompt, TextPrompt] = None):
+    source_stage_id = engine_input_source[0]
+    thinker_outputs = stage_list[source_stage_id].engine_outputs
+    talker_inputs = []
+    multi_modal_data = {thinker_output.request_id: 
+            prompt.get('multi_modal_data', None) for thinker_output, prompt in zip(thinker_outputs, prompt)}
+    
+    for i, thinker_output in enumerate(thinker_outputs):
+        output = thinker_output.outputs[0]
+        prompt_token_ids = thinker_output.prompt_token_ids
+        thinker_output_ids = output.token_ids
+        prompt_token_ids_len = len(prompt_token_ids)
+        thinker_hidden_states = output.multimodal_output["latent"].clone().detach().cuda()
+        talker_inputs.append(
+            OmniTokensPrompt(
+                prompt_token_ids=[0] * (len(prompt_token_ids) + 2), # add 2 for codec pad and start token
+                additional_information={
+                    "thinker_result": thinker_hidden_states[prompt_token_ids_len:].to(torch.float32),
+                    "prompt_embeds": thinker_hidden_states[:prompt_token_ids_len].to(torch.float32),
+                    "prompt_token_ids": prompt_token_ids,
+                    "thinker_output_token_ids": thinker_output_ids,
+                },
+                multi_modal_data=multi_modal_data[thinker_output.request_id] if multi_modal_data is not None else None,
+                mm_processor_kwargs=None,
+            )
+        )
+    return talker_inputs