[Bugfix] Skip generic text-only warmup for DreamZero (#7548)

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>

tests/diffusion/test_diffusion_engine_dummy_run.py

@@ -7,6 +7,7 @@ from unittest.mock import Mock
 import pytest
 
 from vllm_omni.diffusion import io_support
+from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.diffusion_engine import DiffusionEngine, DiffusionExecutionMode
 from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
 from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVRequest
@@ -16,6 +17,22 @@ from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 
 
+@pytest.mark.parametrize("model_class_name", ["DreamZeroPipeline", "LingBotWorldCausalDMDPipeline"])
+def test_observation_conditioned_startup_skips_generic_warmup(model_class_name: str) -> None:
+    """Observation-conditioned pipelines must not receive generic text requests."""
+    engine = DiffusionEngine.__new__(DiffusionEngine)
+    engine.od_config = OmniDiffusionConfig.__new__(OmniDiffusionConfig)
+    engine.od_config.model_class_name = model_class_name
+    engine.od_config.diffusion_load_format = "default"
+    engine.add_req_and_wait_for_response = Mock(side_effect=AssertionError("generic text warmup submitted"))
+    engine.close = Mock()
+
+    engine.run_startup_warmup()
+
+    engine.add_req_and_wait_for_response.assert_not_called()
+    engine.close.assert_not_called()
+
+
 @pytest.mark.parametrize("step_execution", [False, True])
 def test_dummy_run_uses_enough_steps_for_execution_mode(
     monkeypatch: pytest.MonkeyPatch,

vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py

@@ -17,6 +17,7 @@ import re as re_module
 from collections import OrderedDict
 from collections.abc import Iterable
 from contextlib import contextmanager
+from typing import ClassVar
 
 import numpy as np
 import torch
@@ -121,6 +122,10 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
     methods below. The runner binds one session state only for ``forward()``.
     """
 
+    # Generic warmup cannot synthesize robot observations. AR-Diffusion uses
+    # ar_diffusion_warmup_requests() for model-specific warmup instead.
+    dummy_run_num_frames: ClassVar[int] = 0
+
     _POSITIVE_BRANCH = "positive"
     _NEGATIVE_BRANCH = "negative"
     _ar_diffusion_kv_state = None