diff --git a/tests/entrypoints/test_utils.py b/tests/entrypoints/test_utils.py index 11b26d8c3..352ed2aad 100644 --- a/tests/entrypoints/test_utils.py +++ b/tests/entrypoints/test_utils.py @@ -1,5 +1,6 @@ """Unit tests for vllm_omni.entrypoints.utils module.""" +import os from collections import Counter from dataclasses import dataclass @@ -12,6 +13,7 @@ from vllm_omni.entrypoints.utils import ( _convert_dataclasses_to_dict, _filter_dict_like_object, filter_dataclass_kwargs, + resolve_model_config_path, ) pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -269,3 +271,36 @@ class TestFilterDataclassKwargs: assert "unknown_top" not in result assert result["cache_config"]["rel_l1_thresh"] == 0.3 assert "extra_param" not in result["cache_config"] + + +class TestResolveModelConfigPath: + """Test suite for resolve_model_config_path function with diffusers format models.""" + + def test_glm_image_diffusers_format_resolution(self, mocker: MockerFixture): + """Test GlmImagePipeline diffusers class resolves to glm_image config.""" + mocker.patch( + "vllm_omni.entrypoints.utils.file_or_path_exists", + return_value=True, + ) + mocker.patch( + "vllm_omni.entrypoints.utils._try_get_class_name_from_diffusers_config", + return_value="GlmImagePipeline", + ) + mocker.patch( + "vllm_omni.entrypoints.utils.current_omni_platform.get_default_stage_config_path", + return_value="vllm_omni/model_executor/stage_configs", + ) + + original_exists = os.path.exists + + def mock_exists(path): + if "glm_image.yaml" in str(path): + return True + return original_exists(path) + + mocker.patch("os.path.exists", side_effect=mock_exists) + + result = resolve_model_config_path("zai-org/GLM-Image") + + assert result is not None + assert "glm_image.yaml" in result diff --git a/vllm_omni/entrypoints/utils.py b/vllm_omni/entrypoints/utils.py index 40d5ae6e6..c664fe80a 100644 --- a/vllm_omni/entrypoints/utils.py +++ b/vllm_omni/entrypoints/utils.py @@ -18,6 +18,10 @@ PROJECT_ROOT = Path(__file__).parent.parent.parent logger = init_logger(__name__) +_DIFFUSERS_CLASS_TO_CONFIG: dict[str, str] = { + "GlmImagePipeline": "glm_image", +} + def inject_omni_kv_config(stage: Any, omni_conn_cfg: dict[str, Any], omni_from: str, omni_to: str) -> None: """Inject connector configuration into stage engine arguments.""" @@ -218,13 +222,16 @@ def resolve_model_config_path(model: str) -> str: ) default_config_path = current_omni_platform.get_default_stage_config_path() - model_type_str = f"{model_type}.yaml" + if model_type in _DIFFUSERS_CLASS_TO_CONFIG: + normalized_model_type = _DIFFUSERS_CLASS_TO_CONFIG[model_type] + else: + normalized_model_type = model_type.replace("-", "_") + model_type_str = f"{normalized_model_type}.yaml" complete_config_path = PROJECT_ROOT / default_config_path / model_type_str if os.path.exists(complete_config_path): return str(complete_config_path) - # Fall back to default config - stage_config_file = f"vllm_omni/model_executor/stage_configs/{model_type}.yaml" + stage_config_file = f"vllm_omni/model_executor/stage_configs/{normalized_model_type}.yaml" stage_config_path = PROJECT_ROOT / stage_config_file if not os.path.exists(stage_config_path): return None