diff --git a/tests/entrypoints/test_async_omni_diffusion_config.py b/tests/entrypoints/test_async_omni_diffusion_config.py index 1a4582681..6b49eba2c 100644 --- a/tests/entrypoints/test_async_omni_diffusion_config.py +++ b/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) diff --git a/tests/entrypoints/test_omni_llm.py b/tests/entrypoints/test_omni_llm.py index 2359a93f2..fd2ba2673 100644 --- a/tests/entrypoints/test_omni_llm.py +++ b/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 diff --git a/vllm_omni/entrypoints/omni.py b/vllm_omni/entrypoints/omni.py index 168999ac9..109d5f854 100644 --- a/vllm_omni/entrypoints/omni.py +++ b/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(