omni entrypoint support tokenizer arg (#572)

Signed-off-by: Divyansh Singhvi <divyanshsinghvi@gmail.com>
This commit is contained in:
dsinghvi authored and GitHub committed 2026-01-08 15:44:14 +08:00
1 parent 4df794bdaa
commit 8c1259383c
3 files changed
+18 -13

No files matched your search

@@ -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)