[Bugfix][Diffusion] Restore inline single-stage execution (#6813)

Signed-off-by: Gao Han <hgaoaf@connect.ust.hk>

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,
     }
 

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(),