[Diffusion] DreamZero world model integration with CFG parallel + OpenPI serving (#2162)

Signed-off-by: Meng <meng_chen99@163.com>
Signed-off-by: Yangshen Deng <yangshen.d@outlook.com>
Signed-off-by: Suli Wang <50828270+wsl2000@users.noreply.github.com>
Co-authored-by: Meng <meng_chen99@163.com>
Co-authored-by: Suli Wang <50828270+wsl2000@users.noreply.github.com>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>

tests/entrypoints/test_omni_entrypoints.py

@@ -413,6 +413,19 @@ def test_openai_serving_models_can_consume_async_omni_compat_attrs():
 
 
 def test_get_diffusion_od_config_returns_diffusion_stage_config():
+    diffusion_od_config = object()
+    omni = object.__new__(AsyncOmni)
+    omni.engine = SimpleNamespace(
+        stage_clients=[
+            SimpleNamespace(stage_type="llm"),
+            SimpleNamespace(stage_type="diffusion", od_config=diffusion_od_config),
+        ]
+    )
+
+    assert omni.get_diffusion_od_config() is diffusion_od_config
+
+
+def test_get_diffusion_od_config_falls_back_to_inner_engine():
     diffusion_od_config = object()
     omni = object.__new__(AsyncOmni)
     omni.engine = SimpleNamespace(

tests/entrypoints/test_resolve_dreamzero_config.py

@@ -0,0 +1,56 @@
+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.stage_config.StageConfigFactory._auto_detect_model_type",
+        classmethod(lambda _cls, _model, trust_remote_code=True: ("vla", None)),
+    )
+    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"