diff --git a/tests/engine/test_async_omni_engine_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py index c2bbc2ed6..859e03289 100644 --- a/tests/engine/test_async_omni_engine_stage_init.py +++ b/tests/engine/test_async_omni_engine_stage_init.py @@ -399,12 +399,20 @@ def test_initialize_local_diffusion_replica_restores_device_visibility_after_loc os.environ[env_var] = old_env -def test_initialize_local_diffusion_replica_passes_stage_init_timeout_and_inline_flag(monkeypatch): +@pytest.mark.parametrize( + ("num_stages", "expected_inline"), + [(1, True), (2, False)], +) +def test_initialize_local_diffusion_replica_passes_stage_init_timeout_and_inline_flag( + monkeypatch, + num_stages, + expected_inline, +): import vllm_omni.engine.stage_runtime as runtime_mod from vllm_omni.engine.stage_engine_startup import StageReplicaResources runtime = StageRuntime( - stage_configs=[types.SimpleNamespace()], + stage_configs=[types.SimpleNamespace()] * num_stages, model="dummy-model", config_path="dummy-config", stage_init_timeout=1, @@ -434,7 +442,7 @@ def test_initialize_local_diffusion_replica_passes_stage_init_timeout_and_inline "stage_id": 0, "stage_init_timeout": 302, "batch_size": 4, - "use_inline": False, + "use_inline": expected_inline, "omni_master_server": None, } diff --git a/vllm_omni/engine/stage_runtime.py b/vllm_omni/engine/stage_runtime.py index add223ea5..6088814ce 100644 --- a/vllm_omni/engine/stage_runtime.py +++ b/vllm_omni/engine/stage_runtime.py @@ -673,7 +673,8 @@ class StageRuntime: metadata=plan.metadata, stage_init_timeout=stage_init_timeout, batch_size=self._diffusion_batch_size, - use_inline=plan.num_replicas == 1 and bool(inline_diffusion or custom_pipeline_args), + use_inline=plan.num_replicas == 1 + and bool(self._num_stages == 1 or inline_diffusion or custom_pipeline_args), replica_id=plan.replica_id, omni_master_server=self._get_omni_master_server(), omni_coordinator_address=self._get_coordinator_address(),