diff --git a/tests/diffusion/test_omni_diffusion_config_offline_model_path.py b/tests/diffusion/test_omni_diffusion_config_offline_model_path.py new file mode 100644 index 000000000..674f8d93f --- /dev/null +++ b/tests/diffusion/test_omni_diffusion_config_offline_model_path.py @@ -0,0 +1,60 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for OmniDiffusionConfig HF-offline model-path resolution.""" + +import os + +import huggingface_hub +import pytest + +import vllm_omni.diffusion.data as diffusion_data +from vllm_omni.diffusion.data import OmniDiffusionConfig + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_LOCAL_SNAPSHOT = "/root/.cache/huggingface/hub/models--foo--bar/snapshots/deadbeef" + + +@pytest.fixture +def record_get_model_path(monkeypatch: pytest.MonkeyPatch): + calls: list[tuple[str, str | None]] = [] + + def fake_get_model_path(model, revision=None): + calls.append((model, revision)) + if os.path.exists(model): + return model + return _LOCAL_SNAPSHOT + + monkeypatch.setattr(diffusion_data, "get_model_path", fake_get_model_path) + return calls + + +class TestOmniDiffusionConfigOfflineModelPath: + def test_offline_repo_id_replaced_with_local_path( + self, monkeypatch: pytest.MonkeyPatch, record_get_model_path + ) -> None: + monkeypatch.setattr(huggingface_hub.constants, "HF_HUB_OFFLINE", True) + config = OmniDiffusionConfig(model="foo/bar") + assert config.model == _LOCAL_SNAPSHOT + assert record_get_model_path == [("foo/bar", None)] + + def test_online_leaves_model_untouched(self, monkeypatch: pytest.MonkeyPatch, record_get_model_path) -> None: + monkeypatch.setattr(huggingface_hub.constants, "HF_HUB_OFFLINE", False) + config = OmniDiffusionConfig(model="foo/bar") + assert config.model == "foo/bar" + assert record_get_model_path == [] + + def test_offline_local_path_unchanged(self, monkeypatch: pytest.MonkeyPatch, tmp_path) -> None: + monkeypatch.setattr(huggingface_hub.constants, "HF_HUB_OFFLINE", True) + local = str(tmp_path) + config = OmniDiffusionConfig(model=local) + assert config.model == local + + @pytest.mark.parametrize("model", [None, ""]) + def test_offline_empty_model_unchanged( + self, monkeypatch: pytest.MonkeyPatch, record_get_model_path, model: str | None + ) -> None: + monkeypatch.setattr(huggingface_hub.constants, "HF_HUB_OFFLINE", True) + config = OmniDiffusionConfig(model=model) + assert config.model == model + assert record_get_model_path == [] diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 31617ccdb..86f8d8980 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -11,6 +11,7 @@ from enum import Enum from typing import TYPE_CHECKING, Any import diffusers +import huggingface_hub import torch from PIL import Image from pydantic import Field, model_validator @@ -20,6 +21,7 @@ from vllm.logger import init_logger from vllm.model_executor.layers.quantization.base_config import ( QuantizationConfig, ) +from vllm.transformers_utils.repo_utils import get_model_path from vllm_omni.diffusion.model_metadata import get_diffusion_model_metadata from vllm_omni.diffusion.utils.network_utils import is_port_available @@ -1018,6 +1020,18 @@ class OmniDiffusionConfig: "valid together with diffusion_load_format=diffusers" ) + # when use hf offline, replace model to local model path + # align with ar stage behavior, see vllm/engine/arg_utils.py + if huggingface_hub.constants.HF_HUB_OFFLINE and self.model: + model_id = self.model + self.model = get_model_path(self.model, self.revision) + if model_id != self.model: + logger.info( + "HF_HUB_OFFLINE is True, replace model_id [%s] to model_path [%s]", + model_id, + self.model, + ) + def _propagate_quantization_from_tf_config(self, tf_config: "TransformerConfig") -> None: if tf_config.quant_config is None: return