From 87d6414a26fdce8e5c11e32a15f54e3400313156 Mon Sep 17 00:00:00 2001 From: Gaohan123 Date: Wed, 22 Oct 2025 18:14:11 +0800 Subject: [PATCH] supplement some imports --- vllm_omni/entrypoints/omni_llm.py | 6 ++++ .../stage_configs/qwen2_5_omni.yaml | 7 ++++ .../stage_input_processors/__init__.py | 0 .../stage_input_processors/qwen2_5_omni.py | 32 +++++++++++++++++++ 4 files changed, 45 insertions(+) create mode 100644 vllm_omni/model_executor/stage_input_processors/__init__.py create mode 100644 vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py diff --git a/vllm_omni/entrypoints/omni_llm.py b/vllm_omni/entrypoints/omni_llm.py index 063d37570..214060a40 100644 --- a/vllm_omni/entrypoints/omni_llm.py +++ b/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() diff --git a/vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml b/vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml index 05e45d3b2..aa38bae56 100644 --- a/vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml +++ b/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 diff --git a/vllm_omni/model_executor/stage_input_processors/__init__.py b/vllm_omni/model_executor/stage_input_processors/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py b/vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py new file mode 100644 index 000000000..a854fc562 --- /dev/null +++ b/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 \ No newline at end of file