wehub-resource-sync
57 lines
2.0 KiB
Python
57 lines
2.0 KiB
Python
import pytest
|
|
|
|
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
|
from vllm_omni.diffusion.stage_diffusion_proc import StageDiffusionProc
|
|
from vllm_omni.entrypoints.utils import load_stage_configs_from_model, resolve_model_config_path
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
|
|
def test_dreamzero_vla_resolves_to_dreamzero_config(monkeypatch):
|
|
monkeypatch.setattr(
|
|
"vllm_omni.entrypoints.utils.get_config",
|
|
lambda _model, trust_remote_code=True: type("Cfg", (), {"model_type": "vla"})(),
|
|
)
|
|
monkeypatch.setattr(
|
|
"vllm_omni.entrypoints.utils._looks_like_dreamzero",
|
|
lambda _model: True,
|
|
)
|
|
result = resolve_model_config_path("GEAR-Dreams/DreamZero-DROID")
|
|
|
|
assert result is not None
|
|
assert result.endswith("vllm_omni/deploy/dreamzero.yaml")
|
|
|
|
|
|
def test_dreamzero_config_sets_model_class_and_policy_config(monkeypatch):
|
|
monkeypatch.setattr(
|
|
"vllm_omni.config.config_factory.StageConfigFactory._try_infer_model_type",
|
|
classmethod(lambda _cls, model, trust_remote_code=True: "vla"),
|
|
)
|
|
monkeypatch.setattr(
|
|
"vllm_omni.diffusion.utils.hf_utils._looks_like_dreamzero",
|
|
lambda _model: True,
|
|
)
|
|
|
|
stage_configs, _ = load_stage_configs_from_model("GEAR-Dreams/DreamZero-DROID")
|
|
engine_args = stage_configs[0].engine_args
|
|
|
|
assert engine_args.model_class_name == "DreamZeroPipeline"
|
|
assert engine_args.model_config.policy_server_config.action_space == "joint_position"
|
|
|
|
|
|
def test_dreamzero_enrich_config_preserves_explicit_model_class_name(monkeypatch):
|
|
monkeypatch.setattr(
|
|
"vllm.transformers_utils.config.get_hf_file_to_dict",
|
|
lambda path, _model: None if path == "model_index.json" else {"model_type": "vla", "architectures": ["VLA"]},
|
|
)
|
|
|
|
od_config = OmniDiffusionConfig(
|
|
model="GEAR-Dreams/DreamZero-DROID",
|
|
model_class_name="DreamZeroPipeline",
|
|
)
|
|
proc = StageDiffusionProc(od_config.model, od_config)
|
|
|
|
proc._enrich_config()
|
|
|
|
assert od_config.model_class_name == "DreamZeroPipeline"
|