omni entrypoint support tokenizer arg (#572)

Signed-off-by: Divyansh Singhvi <divyanshsinghvi@gmail.com>

tests/entrypoints/test_async_omni_diffusion_config.py

@@ -7,7 +7,7 @@ from vllm_omni.entrypoints.async_omni import AsyncOmni
 
 def test_default_stage_config_includes_cache_backend(monkeypatch):
     """Ensure cache_backend/cache_config are preserved in default diffusion stage."""
-    monkeypatch.setattr(omni_module, "load_stage_configs_from_model", lambda model: [])
+    monkeypatch.setattr(omni_module, "load_stage_configs_from_model", lambda model, base_engine_args=None: [])
     monkeypatch.setattr(omni_module, "resolve_model_config_path", lambda model: None)
     monkeypatch.setattr(AsyncOmni, "_start_stages", lambda self, model: None)
     monkeypatch.setattr(AsyncOmni, "_wait_for_stages_ready", lambda self, timeout=0: None)
@@ -37,7 +37,7 @@ def test_default_stage_config_includes_cache_backend(monkeypatch):
 
 def test_default_cache_config_used_when_missing(monkeypatch):
     """Ensure default cache_config is applied when cache_backend is set."""
-    monkeypatch.setattr(omni_module, "load_stage_configs_from_model", lambda model: [])
+    monkeypatch.setattr(omni_module, "load_stage_configs_from_model", lambda model, base_engine_args=None: [])
     monkeypatch.setattr(omni_module, "resolve_model_config_path", lambda model: None)
     monkeypatch.setattr(AsyncOmni, "_start_stages", lambda self, model: None)
     monkeypatch.setattr(AsyncOmni, "_wait_for_stages_ready", lambda self, timeout=0: None)
@@ -55,7 +55,7 @@ def test_default_cache_config_used_when_missing(monkeypatch):
 
 def test_default_stage_devices_from_sequence_parallel(monkeypatch):
     """Ensure devices list reflects sequence parallel size when no parallel_config is provided."""
-    monkeypatch.setattr(omni_module, "load_stage_configs_from_model", lambda model: [])
+    monkeypatch.setattr(omni_module, "load_stage_configs_from_model", lambda model, base_engine_args=None: [])
     monkeypatch.setattr(omni_module, "resolve_model_config_path", lambda model: None)
     monkeypatch.setattr(AsyncOmni, "_start_stages", lambda self, model: None)
     monkeypatch.setattr(AsyncOmni, "_wait_for_stages_ready", lambda self, timeout=0: None)

tests/entrypoints/test_omni_llm.py

@@ -445,7 +445,7 @@ def mock_get_config(monkeypatch):
 def test_initialize_stage_configs_called_when_none(monkeypatch, fake_stage_config):
     """Test that stage configs are auto-loaded when stage_configs_path is None."""
 
-    def _fake_loader(model: str):
+    def _fake_loader(model: str, base_engine_args=None):
         return [
             _FakeStageConfig(fake_stage_config),
             _FakeStageConfig(fake_stage_config),
@@ -510,7 +510,7 @@ def test_initialize_stage_configs_called_when_none(monkeypatch, fake_stage_confi
 def test_generate_raises_on_length_mismatch(monkeypatch, fake_stage_config):
     """Test that generate raises ValueError when sampling_params_list length doesn't match."""
 
-    def _fake_loader(model: str):
+    def _fake_loader(model: str, base_engine_args=None):
         return [_FakeStageConfig(fake_stage_config)]
 
     import sys
@@ -562,7 +562,7 @@ def test_generate_pipeline_and_final_outputs(monkeypatch, fake_stage_config):
     stage_cfg1 = dict(fake_stage_config)
     stage_cfg1["processed_input"] = ["processed-for-stage-1"]
 
-    def _fake_loader(model: str):
+    def _fake_loader(model: str, base_engine_args=None):
         return [_FakeStageConfig(stage_cfg0), _FakeStageConfig(stage_cfg1)]
 
     import sys
@@ -660,7 +660,7 @@ def test_generate_no_final_output_returns_empty(monkeypatch, fake_stage_config):
     stage_cfg0["final_output"] = False
     stage_cfg1["final_output"] = False
 
-    def _fake_loader(model: str):
+    def _fake_loader(model: str, base_engine_args=None):
         return [_FakeStageConfig(stage_cfg0), _FakeStageConfig(stage_cfg1)]
 
     import sys
@@ -734,7 +734,7 @@ def test_generate_sampling_params_none_use_default(monkeypatch, fake_stage_confi
     stage_cfg0["final_output"] = False
     stage_cfg1["final_output"] = False
 
-    def _fake_loader(model: str):
+    def _fake_loader(model: str, base_engine_args=None):
         return [_FakeStageConfig(stage_cfg0), _FakeStageConfig(stage_cfg1)]
 
     import sys
@@ -802,7 +802,7 @@ def test_generate_sampling_params_none_use_default(monkeypatch, fake_stage_confi
 def test_wait_for_stages_ready_timeout(monkeypatch, fake_stage_config):
     """Test that _wait_for_stages_ready handles timeout correctly."""
 
-    def _fake_loader(model: str):
+    def _fake_loader(model: str, base_engine_args=None):
         return [_FakeStageConfig(fake_stage_config)]
 
     import sys
@@ -858,7 +858,7 @@ def test_wait_for_stages_ready_timeout(monkeypatch, fake_stage_config):
 def test_generate_handles_error_messages(monkeypatch, fake_stage_config):
     """Test that generate handles error messages from stages correctly."""
 
-    def _fake_loader(model: str):
+    def _fake_loader(model: str, base_engine_args=None):
         return [_FakeStageConfig(fake_stage_config)]
 
     import sys
@@ -934,7 +934,7 @@ def test_generate_handles_error_messages(monkeypatch, fake_stage_config):
 def test_close_sends_shutdown_signal(monkeypatch, fake_stage_config):
     """Test that close() sends shutdown signal to all input queues."""
 
-    def _fake_loader(model: str):
+    def _fake_loader(model: str, base_engine_args=None):
         return [_FakeStageConfig(fake_stage_config)]
 
     import sys

vllm_omni/entrypoints/omni.py

@@ -201,16 +201,21 @@ class OmniBase:
         stage_configs_path = kwargs.get("stage_configs_path", None)
         log_stats = kwargs.get("log_stats", False)
 
+        ### base engine args
+        tokenizer = kwargs.get("tokenizer", None)
+
+        base_engine_args = {"tokenizer": tokenizer} if tokenizer is not None else None
+
         # Load stage configurations from YAML
         if stage_configs_path is None:
             self.config_path = resolve_model_config_path(model)
-            self.stage_configs = load_stage_configs_from_model(model)
+            self.stage_configs = load_stage_configs_from_model(model, base_engine_args=base_engine_args)
             if not self.stage_configs:
                 default_stage_cfg = self._create_default_diffusion_stage_cfg(kwargs)
                 self.stage_configs = OmegaConf.create(default_stage_cfg)
         else:
             self.config_path = stage_configs_path
-            self.stage_configs = load_stage_configs_from_yaml(stage_configs_path)
+            self.stage_configs = load_stage_configs_from_yaml(stage_configs_path, base_engine_args=base_engine_args)
 
         # Initialize connectors
         self.omni_transfer_config, self.connectors = initialize_orchestrator_connectors(