diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py
index 2f744480c..41a74ca47 100644
--- a/vllm_omni/engine/arg_utils.py
+++ b/vllm_omni/engine/arg_utils.py
@@ -26,7 +26,7 @@ class OmniEngineArgs(EngineArgs):
                 "model_config.engine_output_type for schedulers to use."
             ),
         )
-        parser.add_argument("--model-stage", type=str, default=EngineArgs.model_stage, 
+        parser.add_argument("--model-stage", type=str, default=OmniEngineArgs.model_stage, 
         help="Declare model stage (e.g., 'thinker', 'talker', 'token2wav'). This will be written into model_config.model_stage for schedulers to use.")
         return parser
 
diff --git a/vllm_omni/entrypoints/omni_llm.py b/vllm_omni/entrypoints/omni_llm.py
index aee4fcf15..063d37570 100644
--- a/vllm_omni/entrypoints/omni_llm.py
+++ b/vllm_omni/entrypoints/omni_llm.py
@@ -39,8 +39,8 @@ class OmniLLM:
     def initialize_stages(self, model: str):
         for stage_config in self.stage_configs:
             stage = Stage(stage_config)
-            omni_llm = OmniLLM(model=model, **stage_config.engine_args)
-            stage.set_engine(omni_llm)
+            stage_llm = StageLLM(model=model, **stage_config.engine_args)
+            stage.set_engine(stage_llm)
             self.stage_list.append(stage)
     
     def generate(
@@ -51,6 +51,7 @@ class OmniLLM:
     ) -> list[OmniRequestOutput]:
         """Generate text outputs for the given prompts."""
         final_outputs: list[OmniRequestOutput] = []
+        assert len(sampling_params_list) == len(self.stage_list), "sampling_params_list must be a list of length equal to the number of stages"
         for stage_id, stage in enumerate(self.stage_list):
             if stage_id > 0:
                 engine_inputs = stage.process_engine_inputs(self.stage_list, prompts)