fix for some reviews

vllm_omni/config/__init__.py

@@ -1,10 +1,11 @@
 """
 Configuration module for vLLM-omni.
 """
-from vllm.config import ModelConfig
 from typing import Optional
 from pydantic.dataclasses import dataclass
 from pydantic import ConfigDict
+
+from vllm.config import ModelConfig
 from vllm.config import config
 
 import vllm_omni.model_executor.models as me_models

vllm_omni/entrypoints/omni_llm.py

@@ -15,7 +15,7 @@ from vllm.logger import init_logger
 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.entrypoints.stage 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

vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml

@@ -4,7 +4,7 @@ stage_args:
     engine_args:
       model_stage: thinker
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_cls: vllm_omni.worker.AR_gpu_worker.ARGPUWorker
+      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
       scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler
       gpu_memory_utilization: 0.32
       enforce_eager: true  # need to discuss
@@ -16,7 +16,7 @@ stage_args:
     engine_args:
       model_stage: talker
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_cls: vllm_omni.worker.AR_gpu_worker.ARGPUWorker
+      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
       scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler
       gpu_memory_utilization: 0.32
       enforce_eager: true
@@ -24,12 +24,11 @@ stage_args:
       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
+      worker_cls: vllm_omni.worker.gpu_diffusion_worker.GPUDiffusionWorker
       scheduler_cls: vllm_omni.core.sched.diffusion_scheduler.DiffusionScheduler
       gpu_memory_utilization: 0.3
       enforce_eager: true

vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py

@@ -1,6 +1,8 @@
-from vllm.inputs import TextPrompt
 from typing import Union
 import torch
+
+from vllm.inputs import TextPrompt
+
 from vllm_omni.inputs.data import OmniTokensPrompt
 
 def thinker2talker(stage_list, engine_input_source, prompt: Union[OmniTokensPrompt, TextPrompt] = None):