Refactor StageDiffusionClient and StageEngineCoreClient (#2006)

docs/api/README.md

@@ -5,7 +5,6 @@
 Main entry points for vLLM-Omni inference and serving.
 
 - [vllm_omni.entrypoints.async_omni.AsyncOmni][]
-- [vllm_omni.entrypoints.async_omni_diffusion.AsyncOmniDiffusion][]
 - [vllm_omni.entrypoints.cfg_companion_tracker.CfgCompanionTracker][]
 - [vllm_omni.entrypoints.cli.benchmark.base.OmniBenchmarkSubcommandBase][]
 - [vllm_omni.entrypoints.cli.benchmark.main.OmniBenchmarkSubcommand][]

docs/contributing/ci/CI_5levels.md

@@ -168,7 +168,6 @@ vllm_omni/                                    tests/
 │   └── arg_utils.py                            │   └── test_arg_utils.py               ⬜
 │
 ├── entrypoints/                        →     ├── entrypoints/
-│   ├── async_omni_diffusion.py                 │   ├── test_async_omni_diffusion_config.py  ✅
 │   ├── stage_utils.py                          │   ├── test_stage_utils.py            ✅
 │   ├── cli/                                     │   ├── cli/                           (benchmarks/test_serve_cli.py covers CLI serve)
 │   │   └── ...                                  │   │   └── test_*.py                  ⬜

docs/contributing/ci/tests_style.md

@@ -73,7 +73,6 @@ vllm_omni/                                    tests/
 │   └── arg_utils.py                            │   └── test_arg_utils.py               ⬜
 │
 ├── entrypoints/                        →     ├── entrypoints/
-│   ├── async_omni_diffusion.py                 │   ├── test_async_omni_diffusion_config.py  ✅
 │   ├── stage_utils.py                          │   ├── test_stage_utils.py            ✅
 │   ├── cli/                                     │   ├── cli/                           (benchmarks/test_serve_cli.py covers CLI serve)
 │   │   └── ...                                  │   │   └── test_*.py                  ⬜

docs/design/module/async_omni_architecture.md

@@ -69,7 +69,7 @@
 [5] Orchestrator._orchestration_loop (loop)
     -> poll stage output
        - llm stage: await get_output_async()
-       - diffusion stage: get_diffusion_output_async()
+       - diffusion stage: get_diffusion_output_nowait()
     -> (llm stage) output_processors[i].process_outputs(...)
     -> _route_output(...)
     -> if finished and not final_stage and non-async-chunk:
@@ -112,7 +112,7 @@ sequenceDiagram
     ORCH->>S0: add_request_async
 
     loop poll route forward
-        ORCH->>S0: get_output_async / get_diffusion_output_async
+        ORCH->>S0: get_output_async / get_diffusion_output_nowait
         ORCH->>ORCH: _route_output
         alt need forward to next stage
             ORCH->>SN: add_request_async

tests/e2e/offline_inference/test_qwen_image_diffusion_batching.py

@@ -509,7 +509,7 @@ def test_diffusion_batching_async_explicit_batch(model_name: str):
     all prompts in a single engine call and returns a single combined result.
 
     The list-prompt path routes through the orchestrator's
-    ``add_batch_request_async`` → ``AsyncOmniDiffusion.generate_batch``
+    ``add_batch_request_async`` → ``AsyncOmni.generate_batch``
     and yields ONE ``OmniRequestOutput`` with ALL images combined.
     """
 

tests/engine/test_async_omni_engine_stage_init.py

@@ -113,6 +113,7 @@ def test_attach_llm_stage_uses_omni_input_preprocessor(monkeypatch):
         executor_class=object,
         engine_manager=object(),
         coordinator=object(),
+        proc=None,
         addresses=types.SimpleNamespace(
             inputs=["inproc://input"],
             outputs=["inproc://output"],

tests/entrypoints/openai_api/test_image_server.py

@@ -106,7 +106,7 @@ def test_encode_image_base64():
 
 
 class MockGenerationResult:
-    """Mock result object from AsyncOmniDiffusion.generate()"""
+    """Mock result object from AsyncOmni.generate()"""
 
     def __init__(self, images):
         self.images = images

tests/entrypoints/test_async_omni_diffusion.py

@@ -1,113 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import asyncio
-import threading
-from concurrent.futures import ThreadPoolExecutor
-from types import SimpleNamespace
-from unittest.mock import Mock
-
-import pytest
-
-import vllm_omni.diffusion.stage_diffusion_client as stage_diffusion_client_module
-from vllm_omni.diffusion.data import DiffusionRequestAbortedError
-from vllm_omni.diffusion.stage_diffusion_client import StageDiffusionClient
-from vllm_omni.entrypoints.async_omni_diffusion import AsyncOmniDiffusion
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-def test_get_diffusion_od_config_returns_direct_config():
-    diffusion = object.__new__(AsyncOmniDiffusion)
-    diffusion.od_config = object()
-
-    assert diffusion.get_diffusion_od_config() is diffusion.od_config
-
-
-def test_async_omni_diffusion_generate_aborts_engine_on_cancel():
-    async def run_test():
-        started = threading.Event()
-        release = threading.Event()
-        abort = Mock()
-
-        def step(request):
-            del request
-            started.set()
-            release.wait(timeout=5)
-            return [SimpleNamespace(request_id="req-1")]
-
-        diffusion = object.__new__(AsyncOmniDiffusion)
-        diffusion.engine = SimpleNamespace(step=step, abort=abort)
-        diffusion._executor = ThreadPoolExecutor(max_workers=1)
-
-        task = asyncio.create_task(
-            diffusion.generate(
-                prompt="hello",
-                sampling_params=OmniDiffusionSamplingParams(),
-                request_id="req-1",
-            )
-        )
-        try:
-            assert await asyncio.to_thread(started.wait, 1)
-            task.cancel()
-            with pytest.raises(asyncio.CancelledError):
-                await task
-        finally:
-            release.set()
-            diffusion._executor.shutdown(wait=True)
-
-        abort.assert_called_once_with("req-1")
-
-    asyncio.run(run_test())
-
-
-def test_stage_diffusion_client_abort_requests_forwards_to_engine():
-    async def run_test():
-        aborted_request_ids: list[list[str]] = []
-
-        async def abort(request_ids):
-            aborted_request_ids.append(request_ids)
-
-        client = object.__new__(StageDiffusionClient)
-        client._engine = SimpleNamespace(abort=abort)
-        client._tasks = {}
-
-        task = asyncio.create_task(asyncio.sleep(60))
-        client._tasks["req-1"] = task
-
-        await client.abort_requests_async(["req-1", "req-2"])
-
-        with pytest.raises(asyncio.CancelledError):
-            await task
-        assert client._tasks == {}
-        assert aborted_request_ids == [["req-1", "req-2"]]
-
-    asyncio.run(run_test())
-
-
-def test_stage_diffusion_client_run_treats_abort_as_normal_path(monkeypatch):
-    async def run_test():
-        async def generate(prompt, sampling_params, request_id):
-            del prompt, sampling_params
-            raise DiffusionRequestAbortedError(f"Request {request_id} aborted.")
-
-        info = Mock()
-        exception = Mock()
-        monkeypatch.setattr(stage_diffusion_client_module.logger, "info", info)
-        monkeypatch.setattr(stage_diffusion_client_module.logger, "exception", exception)
-
-        client = object.__new__(StageDiffusionClient)
-        client.stage_id = 3
-        client._engine = SimpleNamespace(generate=generate)
-        client._output_queue = asyncio.Queue()
-        client._tasks = {"req-1": object()}
-
-        await client._run("req-1", "prompt", OmniDiffusionSamplingParams())
-
-        assert client._output_queue.empty()
-        assert client._tasks == {}
-        info.assert_called_once()
-        exception.assert_not_called()
-
-    asyncio.run(run_test())

vllm_omni/diffusion/ipc.py

@@ -31,6 +31,12 @@ def _tensor_to_shm(tensor: torch.Tensor) -> dict[str, Any]:
     import numpy as np
 
     tensor = tensor.detach().cpu().contiguous()
+    original_dtype = tensor.dtype
+    # NumPy does not support bfloat16; promote to float32 for the SHM
+    # transfer and record the original dtype so _tensor_from_shm can
+    # convert back.  The round-trip is lossless for bfloat16 values.
+    if original_dtype == torch.bfloat16:
+        tensor = tensor.to(torch.float32)
     arr = tensor.numpy()
     nbytes = arr.nbytes
     shm = shared_memory.SharedMemory(create=True, size=nbytes)
@@ -40,7 +46,7 @@ def _tensor_to_shm(tensor: torch.Tensor) -> dict[str, Any]:
         "__tensor_shm__": True,
         "name": shm.name,
         "shape": list(tensor.shape),
-        "torch_dtype": str(tensor.dtype),
+        "torch_dtype": str(original_dtype),
         "numpy_dtype": str(arr.dtype),
         "nbytes": nbytes,
     }
@@ -59,6 +65,13 @@ def _tensor_from_shm(handle: dict[str, Any]) -> torch.Tensor:
         np_dtype = np.dtype(handle["numpy_dtype"])
         arr = np.ndarray(handle["shape"], dtype=np_dtype, buffer=shm.buf[: handle["nbytes"]])
         tensor = torch.from_numpy(arr.copy())
+        # Restore the original dtype if it differs from the numpy-compatible
+        # dtype used for the SHM transfer (e.g. bfloat16 → float32 → bfloat16).
+        torch_dtype_str = handle.get("torch_dtype", "")
+        if torch_dtype_str:
+            original_dtype = getattr(torch, torch_dtype_str.replace("torch.", ""), None)
+            if original_dtype is not None and tensor.dtype != original_dtype:
+                tensor = tensor.to(original_dtype)
     finally:
         shm.close()
         shm.unlink()

vllm_omni/diffusion/stage_diffusion_client.py

@@ -1,20 +1,30 @@
 """Stage Diffusion Client for vLLM-Omni multi-stage runtime.
 
-Wraps AsyncOmniDiffusion to expose the same interface the Orchestrator
-expects from any stage client.
+Spawns StageDiffusionProc in a subprocess and communicates via ZMQ
+(PUSH/PULL) to expose the same interface the Orchestrator expects
+from any stage client.
 """
 
 from __future__ import annotations
 
 import asyncio
 import time
+import uuid
+from dataclasses import fields, is_dataclass
 from typing import TYPE_CHECKING, Any
 
+import zmq
 from vllm.logger import init_logger
 
-from vllm_omni.diffusion.data import DiffusionRequestAbortedError
-from vllm_omni.engine.stage_init_utils import StageMetadata
-from vllm_omni.entrypoints.async_omni_diffusion import AsyncOmniDiffusion
+from vllm_omni.diffusion.stage_diffusion_proc import (
+    complete_diffusion_handshake,
+    spawn_diffusion_proc,
+)
+from vllm_omni.distributed.omni_connectors.utils.serialization import (
+    OmniMsgpackDecoder,
+    OmniMsgpackEncoder,
+)
+from vllm_omni.engine.stage_init_utils import StageMetadata, terminate_alive_proc
 from vllm_omni.outputs import OmniRequestOutput
 
 if TYPE_CHECKING:
@@ -25,11 +35,12 @@ logger = init_logger(__name__)
 
 
 class StageDiffusionClient:
-    """Wraps AsyncOmniDiffusion for use inside the Orchestrator.
+    """Communicates with StageDiffusionProc via ZMQ for use inside the Orchestrator.
 
     Exposes the same attributes and async methods the Orchestrator
     uses on StageEngineCoreClient, but routes execution through
-    DiffusionEngine instead of vLLM EngineCore.
+    a StageDiffusionProc subprocess instead of running the diffusion
+    engine in-process.
     """
 
     stage_type: str = "diffusion"
@@ -48,56 +59,137 @@ class StageDiffusionClient:
         self.custom_process_input_func = metadata.custom_process_input_func
         self.engine_input_source = metadata.engine_input_source
 
-        self._engine = AsyncOmniDiffusion(model=model, od_config=od_config, batch_size=batch_size)
+        # Spawn StageDiffusionProc subprocess and wait for READY.
+        proc, handshake_address, request_address, response_address = spawn_diffusion_proc(model, od_config)
+        complete_diffusion_handshake(proc, handshake_address)
+        self._proc = proc
+
+        # ZMQ sockets (sync) for communicating with the subprocess.
+        self._zmq_ctx = zmq.Context()
+        self._request_socket = self._zmq_ctx.socket(zmq.PUSH)
+        self._request_socket.connect(request_address)
+        self._response_socket = self._zmq_ctx.socket(zmq.PULL)
+        self._response_socket.connect(response_address)
+
+        self._encoder = OmniMsgpackEncoder()
+        self._decoder = OmniMsgpackDecoder()
+
+        # Buffers for demultiplexing response messages.
         self._output_queue: asyncio.Queue[OmniRequestOutput] = asyncio.Queue()
+        self._rpc_results: dict[str, Any] = {}
+        self._pending_rpcs: set[str] = set()
         self._tasks: dict[str, asyncio.Task] = {}
+        self._shutting_down = False
 
         logger.info("[StageDiffusionClient] Stage-%s initialized (batch_size=%d)", self.stage_id, batch_size)
 
+    # ------------------------------------------------------------------
+    # Internal helpers
+    # ------------------------------------------------------------------
+
+    def _drain_responses(self) -> None:
+        """Non-blocking drain of all available responses from the subprocess."""
+        while True:
+            try:
+                raw = self._response_socket.recv(zmq.NOBLOCK)
+            except zmq.Again:
+                break
+
+            msg = self._decoder.decode(raw)
+            msg_type = msg.get("type")
+
+            if msg_type == "result":
+                self._output_queue.put_nowait(msg["output"])
+            elif msg_type == "rpc_result":
+                self._rpc_results[msg["rpc_id"]] = msg["result"]
+            elif msg_type == "error":
+                req_id = msg.get("request_id")
+                rpc_id = msg.get("rpc_id")
+                error_msg = msg.get("error")
+                logger.error(
+                    "[StageDiffusionClient] Stage-%s subprocess error for %s: %s",
+                    self.stage_id,
+                    rpc_id or req_id,
+                    error_msg,
+                )
+                # Route RPC errors so collective_rpc_async can unblock.
+                if rpc_id is not None and rpc_id in self._pending_rpcs:
+                    self._rpc_results[rpc_id] = {
+                        "error": True,
+                        "reason": error_msg,
+                    }
+
+    # Fields that are subprocess-local and cannot be serialized across
+    # process boundaries.  They are recreated in the subprocess with
+    # their default values.
+    _NON_SERIALIZABLE_FIELDS = frozenset(
+        {
+            "generator",  # torch.Generator — recreated from seed
+            "modules",  # model components — loaded in subprocess
+        }
+    )
+
+    @staticmethod
+    def _sampling_params_to_dict(sampling_params: Any) -> dict[str, Any]:
+        """Convert sampling params to a plain dict for serialization.
+
+        Uses ``dataclasses.fields`` + ``getattr`` instead of ``asdict``
+        to avoid deep-copying large tensors, and skips fields that
+        cannot cross process boundaries.
+
+        When a ``torch.Generator`` is present but ``seed`` is not set,
+        the generator's initial seed is extracted so the subprocess can
+        recreate an equivalent generator via ``diffusion_model_runner``.
+        """
+        if is_dataclass(sampling_params) and not isinstance(sampling_params, type):
+            result = {
+                f.name: getattr(sampling_params, f.name)
+                for f in fields(sampling_params)
+                if f.name not in StageDiffusionClient._NON_SERIALIZABLE_FIELDS
+            }
+        elif not isinstance(sampling_params, dict):
+            raise TypeError(f"sampling_params is not a dict but {sampling_params.__class__.__name__}")
+        else:
+            result = {
+                k: v for k, v in sampling_params.items() if k not in StageDiffusionClient._NON_SERIALIZABLE_FIELDS
+            }
+
+        # Preserve the generator's seed across the process boundary so
+        # the subprocess can recreate deterministic random state.
+        if result.get("seed") is None:
+            generator = (
+                getattr(sampling_params, "generator", None)
+                if not isinstance(sampling_params, dict)
+                else sampling_params.get("generator")
+            )
+            if generator is not None:
+                if isinstance(generator, list) and generator:
+                    generator = generator[0]
+                if hasattr(generator, "initial_seed"):
+                    result["seed"] = generator.initial_seed()
+
+        return result
+
+    # ------------------------------------------------------------------
+    # Public API (matches the interface the Orchestrator expects)
+    # ------------------------------------------------------------------
+
     async def add_request_async(
         self,
         request_id: str,
         prompt: OmniPromptType,
         sampling_params: OmniDiffusionSamplingParams,
     ) -> None:
-        task = asyncio.create_task(
-            self._run(request_id, prompt, sampling_params),
-            name=f"diffusion-{request_id}",
+        self._request_socket.send(
+            self._encoder.encode(
+                {
+                    "type": "add_request",
+                    "request_id": request_id,
+                    "prompt": prompt,
+                    "sampling_params": self._sampling_params_to_dict(sampling_params),
+                }
+            )
         )
-        self._tasks[request_id] = task
-
-    async def _run(
-        self,
-        request_id: str,
-        prompt: OmniPromptType,
-        sampling_params: OmniDiffusionSamplingParams,
-    ) -> None:
-        try:
-            result = await self._engine.generate(prompt, sampling_params, request_id)
-            await self._output_queue.put(result)
-        except asyncio.CancelledError:
-            logger.info(
-                "[StageDiffusionClient] Stage-%s req=%s cancelled",
-                self.stage_id,
-                request_id,
-            )
-            raise
-        except DiffusionRequestAbortedError as e:
-            logger.info(
-                "[StageDiffusionClient] Stage-%s req=%s aborted: %s",
-                self.stage_id,
-                request_id,
-                e,
-            )
-        except Exception as e:
-            logger.exception(
-                "[StageDiffusionClient] Stage-%s req=%s failed: %s",
-                self.stage_id,
-                request_id,
-                e,
-            )
-        finally:
-            self._tasks.pop(request_id, None)
 
     # TODO(Long): Temporary solution to boost performance of diffusion stages.
     # Remove this after scheduling algorithm is implemented
@@ -126,12 +218,16 @@ class StageDiffusionClient:
         sampling_params: OmniDiffusionSamplingParams,
     ) -> None:
         try:
-            result = await self._engine.generate_batch(
-                prompts,
-                sampling_params,
-                request_id,
+            self._request_socket.send(
+                self._encoder.encode(
+                    {
+                        "type": "add_batch_request",
+                        "request_id": request_id,
+                        "prompts": prompts,
+                        "sampling_params": self._sampling_params_to_dict(sampling_params),
+                    }
+                )
             )
-            await self._output_queue.put(result)
         except Exception as e:
             logger.exception(
                 "[StageDiffusionClient] Stage-%s batch req=%s failed: %s",
@@ -142,18 +238,24 @@ class StageDiffusionClient:
         finally:
             self._tasks.pop(request_id, None)
 
-    def get_diffusion_output_async(self) -> OmniRequestOutput | None:
+    def get_diffusion_output_nowait(self) -> OmniRequestOutput | None:
+        self._drain_responses()
         try:
             return self._output_queue.get_nowait()
         except asyncio.QueueEmpty:
+            if not self._shutting_down and self._proc is not None and not self._proc.is_alive():
+                raise RuntimeError(f"StageDiffusionProc died unexpectedly (exit code {self._proc.exitcode})")
             return None
 
     async def abort_requests_async(self, request_ids: list[str]) -> None:
-        for rid in request_ids:
-            task = self._tasks.pop(rid, None)
-            if task:
-                task.cancel()
-        await self._engine.abort(request_ids)
+        self._request_socket.send(
+            self._encoder.encode(
+                {
+                    "type": "abort",
+                    "request_ids": list(request_ids),
+                }
+            )
+        )
 
     async def collective_rpc_async(
         self,
@@ -162,60 +264,66 @@ class StageDiffusionClient:
         args: tuple[Any, ...] = (),
         kwargs: dict[str, Any] | None = None,
     ) -> Any:
-        """Best-effort control RPC shim for diffusion stages.
-
-        TODO(AsyncOmni): add dedicated wrappers on AsyncOmniDiffusion for the
-        remaining control APIs instead of reaching into its underlying engine.
-        """
-        kwargs = kwargs or {}
-
-        # Handle profile method: inject stage_id into profile_prefix for diffusion stages
+        """Forward control RPCs to the diffusion subprocess."""
+        # Inject a default profile_prefix that includes stage_id when profiling.
         if method == "profile":
-            target = getattr(self._engine, method, None)
-            if target is None:
-                return {
-                    "supported": False,
-                    "todo": True,
-                    "reason": f"AsyncOmniDiffusion.{method} is not implemented",
-                }
-            # Extract is_start and profile_prefix from args
-            is_start = args[0] if args else True
-            profile_prefix = args[1] if len(args) > 1 else None
-            # Generate profile_prefix with stage_id if starting and no prefix provided
+            args_list = list(args)
+            is_start = args_list[0] if args_list else True
+            profile_prefix = args_list[1] if len(args_list) > 1 else None
             if is_start and profile_prefix is None:
                 profile_prefix = f"stage_{self.stage_id}_diffusion_{int(time.time())}"
-            result = target(is_start, profile_prefix)
-            if timeout is not None:
-                return await asyncio.wait_for(result, timeout=timeout)
-            return await result
+                if len(args_list) > 1:
+                    args_list[1] = profile_prefix
+                else:
+                    args_list.append(profile_prefix)
+                args = tuple(args_list)
 
-        if method in {"add_lora", "remove_lora", "list_loras", "pin_lora"}:
-            target = getattr(self._engine, method, None)
-            if target is None:
-                return {
-                    "supported": False,
-                    "todo": True,
-                    "reason": f"AsyncOmniDiffusion.{method} is not implemented",
+        kwargs = kwargs or {}
+        rpc_id = uuid.uuid4().hex
+        self._pending_rpcs.add(rpc_id)
+
+        self._request_socket.send(
+            self._encoder.encode(
+                {
+                    "type": "collective_rpc",
+                    "rpc_id": rpc_id,
+                    "method": method,
+                    "timeout": timeout,
+                    "args": list(args),
+                    "kwargs": kwargs,
                 }
-            result = target(*args, **kwargs)
-            if timeout is not None:
-                return await asyncio.wait_for(result, timeout=timeout)
-            return await result
-
-        # Fall back to collective RPC for other methods
-        loop = asyncio.get_running_loop()
-        return await loop.run_in_executor(
-            self._engine._executor,
-            self._engine.engine.collective_rpc,
-            method,
-            timeout,
-            args,
-            kwargs,
-            None,
+            )
         )
 
+        deadline = time.monotonic() + timeout if timeout else None
+        # Wait for the matching RPC response, buffering result messages.
+        try:
+            while True:
+                self._drain_responses()
+                if rpc_id in self._rpc_results:
+                    return self._rpc_results.pop(rpc_id)
+                if self._proc is not None and not self._proc.is_alive():
+                    raise RuntimeError(
+                        f"StageDiffusionProc died while waiting for "
+                        f"collective_rpc '{method}' (exit code {self._proc.exitcode})"
+                    )
+                if deadline and time.monotonic() > deadline:
+                    raise TimeoutError(f"collective_rpc_async '{method}' timed out after {timeout}s")
+                await asyncio.sleep(0.01)
+        finally:
+            self._pending_rpcs.discard(rpc_id)
+
     def shutdown(self) -> None:
-        for task in self._tasks.values():
-            task.cancel()
-        self._tasks.clear()
-        self._engine.close()
+        self._shutting_down = True
+        try:
+            self._request_socket.send(self._encoder.encode({"type": "shutdown"}))
+        except Exception:
+            pass
+
+        if self._proc is not None and self._proc.is_alive():
+            self._proc.join(timeout=10)
+            terminate_alive_proc(self._proc)
+
+        self._request_socket.close(linger=0)
+        self._response_socket.close(linger=0)
+        self._zmq_ctx.term()

vllm_omni/diffusion/stage_diffusion_proc.py

@@ -0,0 +1,604 @@
+"""Subprocess entry point for the diffusion engine.
+
+StageDiffusionProc runs DiffusionEngine in a child process,
+communicating with StageDiffusionClient via ZMQ (PUSH/PULL).
+"""
+
+from __future__ import annotations
+
+import asyncio
+import signal
+import time
+from concurrent.futures import ThreadPoolExecutor
+from multiprocessing.process import BaseProcess
+from typing import TYPE_CHECKING, Any
+
+import msgspec
+import zmq
+import zmq.asyncio
+from vllm.logger import init_logger
+from vllm.transformers_utils.config import get_hf_file_to_dict
+from vllm.utils.network_utils import get_open_zmq_ipc_path, zmq_socket_ctx
+from vllm.utils.system_utils import get_mp_context
+from vllm.v1.utils import shutdown
+
+from vllm_omni.diffusion.data import DiffusionRequestAbortedError, TransformerConfig
+from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.distributed.omni_connectors.utils.serialization import (
+    OmniMsgpackDecoder,
+    OmniMsgpackEncoder,
+)
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.outputs import OmniRequestOutput
+
+if TYPE_CHECKING:
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+logger = init_logger(__name__)
+
+_HANDSHAKE_POLL_TIMEOUT_S = 600
+
+
+class StageDiffusionProc:
+    """Subprocess entry point for diffusion inference.
+
+    Manages DiffusionEngine lifecycle, async request processing,
+    and ZMQ-based communication with StageDiffusionClient.
+    """
+
+    def __init__(self, model: str, od_config: OmniDiffusionConfig) -> None:
+        self._model = model
+        self._od_config = od_config
+        self._engine: DiffusionEngine | None = None
+        self._executor: ThreadPoolExecutor | None = None
+        self._closed = False
+
+    # ------------------------------------------------------------------
+    # Initialization
+    # ------------------------------------------------------------------
+
+    def initialize(self) -> None:
+        """Enrich config, create DiffusionEngine and thread pool."""
+        self._enrich_config()
+        self._engine = DiffusionEngine.make_engine(self._od_config)
+        self._executor = ThreadPoolExecutor(max_workers=1)
+        logger.info("StageDiffusionProc initialized with model: %s", self._model)
+
+    def _enrich_config(self) -> None:
+        """Load model metadata from HuggingFace and populate od_config fields.
+
+        Diffusers-style models expose ``model_index.json`` with ``_class_name``.
+        Non-diffusers models (e.g. Bagel, NextStep) only have ``config.json``,
+        so we fall back to reading that and mapping model_type manually.
+        """
+        od_config = self._od_config
+
+        try:
+            config_dict = get_hf_file_to_dict("model_index.json", od_config.model)
+            if config_dict is not None:
+                if od_config.model_class_name is None:
+                    od_config.model_class_name = config_dict.get("_class_name", None)
+                od_config.update_multimodal_support()
+
+                tf_config_dict = get_hf_file_to_dict("transformer/config.json", od_config.model)
+                od_config.tf_model_config = TransformerConfig.from_dict(tf_config_dict)
+            else:
+                raise FileNotFoundError("model_index.json not found")
+        except (AttributeError, OSError, ValueError, FileNotFoundError):
+            cfg = get_hf_file_to_dict("config.json", od_config.model)
+            if cfg is None:
+                raise ValueError(f"Could not find config.json or model_index.json for model {od_config.model}")
+
+            od_config.tf_model_config = TransformerConfig.from_dict(cfg)
+            model_type = cfg.get("model_type")
+            architectures = cfg.get("architectures") or []
+
+            if model_type == "bagel" or "BagelForConditionalGeneration" in architectures:
+                od_config.model_class_name = "BagelPipeline"
+                od_config.tf_model_config = TransformerConfig()
+                od_config.update_multimodal_support()
+            elif model_type == "nextstep":
+                if od_config.model_class_name is None:
+                    od_config.model_class_name = "NextStep11Pipeline"
+                od_config.tf_model_config = TransformerConfig()
+                od_config.update_multimodal_support()
+            elif architectures and len(architectures) == 1:
+                od_config.model_class_name = architectures[0]
+            else:
+                raise
+
+    # ------------------------------------------------------------------
+    # Request processing
+    # ------------------------------------------------------------------
+
+    def _reconstruct_sampling_params(self, sampling_params_dict: dict) -> OmniDiffusionSamplingParams:
+        """Reconstruct OmniDiffusionSamplingParams from a dict, handling LoRA."""
+        lora_req = sampling_params_dict.get("lora_request")
+        if lora_req is not None:
+            from vllm.lora.request import LoRARequest
+
+            if not isinstance(lora_req, LoRARequest):
+                sampling_params_dict["lora_request"] = msgspec.convert(lora_req, LoRARequest)
+
+        return OmniDiffusionSamplingParams(**sampling_params_dict)
+
+    async def _process_request(
+        self,
+        request_id: str,
+        prompt: Any,
+        sampling_params_dict: dict,
+    ) -> OmniRequestOutput:
+        """Build a diffusion request and run DiffusionEngine.step()."""
+        sampling_params = self._reconstruct_sampling_params(sampling_params_dict)
+
+        request = OmniDiffusionRequest(
+            prompts=[prompt],
+            sampling_params=sampling_params,
+            request_ids=[request_id],
+        )
+
+        loop = asyncio.get_running_loop()
+        results = await loop.run_in_executor(self._executor, self._engine.step, request)
+        result = results[0]
+        if not result.request_id:
+            result.request_id = request_id
+        return result
+
+    async def _process_batch_request(
+        self,
+        request_id: str,
+        prompts: list[Any],
+        sampling_params_dict: dict,
+    ) -> OmniRequestOutput:
+        """Build a batched diffusion request and run DiffusionEngine.step().
+
+        All prompts are processed in a single step() call.  The per-prompt
+        results are merged into one :class:`OmniRequestOutput` whose
+        ``images`` list contains every generated image, matching the
+        contract expected by the orchestrator and tests.
+        """
+        sampling_params = self._reconstruct_sampling_params(sampling_params_dict)
+
+        request = OmniDiffusionRequest(
+            prompts=prompts,
+            sampling_params=sampling_params,
+            request_ids=[request_id] * len(prompts),
+        )
+
+        loop = asyncio.get_running_loop()
+        results = await loop.run_in_executor(self._executor, self._engine.step, request)
+
+        # Merge per-prompt results into a single combined output.
+        all_images: list = []
+        merged_mm: dict[str, Any] = {}
+        merged_metrics: dict[str, Any] = {}
+        merged_durations: dict[str, float] = {}
+        peak_mem = 0.0
+        latents = None
+        final_output_type = "image"
+
+        for r in results:
+            all_images.extend(r.images)
+            merged_mm.update(r._multimodal_output)
+            merged_metrics.update(r.metrics)
+            merged_durations.update(r.stage_durations)
+            peak_mem = max(peak_mem, r.peak_memory_mb)
+            if latents is None and r.latents is not None:
+                latents = r.latents
+            if r.final_output_type != "image":
+                final_output_type = r.final_output_type
+
+        return OmniRequestOutput.from_diffusion(
+            request_id=request_id,
+            images=all_images,
+            prompt=prompts[0] if len(prompts) == 1 else None,
+            metrics=merged_metrics,
+            latents=latents,
+            multimodal_output=merged_mm or None,
+            final_output_type=final_output_type,
+            stage_durations=merged_durations,
+            peak_memory_mb=peak_mem,
+        )
+
+    # ------------------------------------------------------------------
+    # Collective RPC dispatch
+    # ------------------------------------------------------------------
+
+    async def _handle_collective_rpc(
+        self,
+        method: str,
+        timeout: float | None,
+        args: tuple,
+        kwargs: dict,
+    ) -> Any:
+        """Dispatch collective RPC calls to DiffusionEngine.
+
+        LoRA methods remap arguments and post-process results to match
+        the contract that ``AsyncOmni`` provides.
+        """
+        loop = asyncio.get_running_loop()
+
+        if method == "profile":
+            is_start = args[0] if args else True
+            profile_prefix = args[1] if len(args) > 1 else None
+            return await loop.run_in_executor(
+                self._executor,
+                self._engine.profile,
+                is_start,
+                profile_prefix,
+            )
+
+        if method == "add_lora":
+            # Reconstruct LoRARequest after IPC if needed.
+            lora_request = args[0] if args else kwargs.get("lora_request")
+            if lora_request is not None:
+                from vllm.lora.request import LoRARequest
+
+                if not isinstance(lora_request, LoRARequest):
+                    lora_request = msgspec.convert(lora_request, LoRARequest)
+            results = await loop.run_in_executor(
+                self._executor,
+                self._engine.collective_rpc,
+                "add_lora",
+                timeout,
+                (),
+                {"lora_request": lora_request},
+                None,
+            )
+            return all(results) if isinstance(results, list) else results
+
+        if method == "remove_lora":
+            results = await loop.run_in_executor(
+                self._executor,
+                self._engine.collective_rpc,
+                "remove_lora",
+                timeout,
+                args,
+                kwargs or {},
+                None,
+            )
+            return all(results) if isinstance(results, list) else results
+
+        if method == "list_loras":
+            results = await loop.run_in_executor(
+                self._executor,
+                self._engine.collective_rpc,
+                "list_loras",
+                timeout,
+                (),
+                {},
+                None,
+            )
+            if not isinstance(results, list):
+                return results or []
+            merged: set[int] = set()
+            for part in results:
+                merged.update(part or [])
+            return sorted(merged)
+
+        if method == "pin_lora":
+            lora_id = args[0] if args else kwargs.get("adapter_id")
+            results = await loop.run_in_executor(
+                self._executor,
+                self._engine.collective_rpc,
+                "pin_lora",
+                timeout,
+                (),
+                {"adapter_id": lora_id},
+                None,
+            )
+            return all(results) if isinstance(results, list) else results
+
+        # Fall back to DiffusionEngine.collective_rpc for all other methods
+        # (e.g. worker extension RPCs like "test_extension_name").
+        return await loop.run_in_executor(
+            self._executor,
+            self._engine.collective_rpc,
+            method,
+            timeout,
+            args,
+            kwargs or {},
+            None,
+        )
+
+    # ------------------------------------------------------------------
+    # ZMQ event loop
+    # ------------------------------------------------------------------
+
+    async def run_loop(
+        self,
+        request_address: str,
+        response_address: str,
+    ) -> None:
+        """Async event loop handling ZMQ messages from StageDiffusionClient."""
+        ctx = zmq.asyncio.Context()
+
+        request_socket = ctx.socket(zmq.PULL)
+        request_socket.bind(request_address)
+
+        response_socket = ctx.socket(zmq.PUSH)
+        response_socket.bind(response_address)
+
+        encoder = OmniMsgpackEncoder()
+        decoder = OmniMsgpackDecoder()
+
+        tasks: dict[str, asyncio.Task] = {}
+
+        async def _dispatch_request(request_id: str, prompt: Any, sampling_params_dict: dict) -> None:
+            """Process a single diffusion request and send the response."""
+            try:
+                result = await self._process_request(request_id, prompt, sampling_params_dict)
+                await response_socket.send(encoder.encode({"type": "result", "output": result}))
+            except DiffusionRequestAbortedError as e:
+                logger.info(
+                    "request_id: %s aborted: %s",
+                    request_id,
+                    str(e),
+                )
+            except Exception as e:
+                logger.exception("Diffusion request %s failed: %s", request_id, e)
+                await response_socket.send(
+                    encoder.encode(
+                        {
+                            "type": "error",
+                            "request_id": request_id,
+                            "error": str(e),
+                        }
+                    )
+                )
+            finally:
+                tasks.pop(request_id, None)
+
+        try:
+            while True:
+                raw = await request_socket.recv()
+                msg = decoder.decode(raw)
+                msg_type = msg.get("type")
+
+                if msg_type == "add_request":
+                    request_id = msg["request_id"]
+                    task = asyncio.create_task(
+                        _dispatch_request(
+                            request_id,
+                            msg["prompt"],
+                            msg["sampling_params"],
+                        )
+                    )
+                    tasks[request_id] = task
+
+                elif msg_type == "add_batch_request":
+                    request_id = msg["request_id"]
+
+                    async def _dispatch_batch(rid: str, prompts: list, sp_dict: dict) -> None:
+                        try:
+                            result = await self._process_batch_request(rid, prompts, sp_dict)
+                            await response_socket.send(encoder.encode({"type": "result", "output": result}))
+                        except DiffusionRequestAbortedError as e:
+                            logger.info(
+                                "request_id: %s aborted: %s",
+                                rid,
+                                str(e),
+                            )
+                        except Exception as e:
+                            logger.exception("Batch diffusion request %s failed: %s", rid, e)
+                            await response_socket.send(
+                                encoder.encode(
+                                    {
+                                        "type": "error",
+                                        "request_id": rid,
+                                        "error": str(e),
+                                    }
+                                )
+                            )
+                        finally:
+                            tasks.pop(rid, None)
+
+                    task = asyncio.create_task(
+                        _dispatch_batch(
+                            request_id,
+                            msg["prompts"],
+                            msg["sampling_params"],
+                        )
+                    )
+                    tasks[request_id] = task
+
+                elif msg_type == "abort":
+                    for rid in msg.get("request_ids", []):
+                        task = tasks.pop(rid, None)
+                        if task:
+                            task.cancel()
+                        self._engine.abort(rid)
+
+                elif msg_type == "collective_rpc":
+                    rpc_id = msg["rpc_id"]
+                    try:
+                        result = await self._handle_collective_rpc(
+                            msg["method"],
+                            msg.get("timeout"),
+                            tuple(msg.get("args", ())),
+                            msg.get("kwargs", {}),
+                        )
+                        await response_socket.send(
+                            encoder.encode(
+                                {
+                                    "type": "rpc_result",
+                                    "rpc_id": rpc_id,
+                                    "result": result,
+                                }
+                            )
+                        )
+                    except Exception as e:
+                        logger.exception("Collective RPC %s failed: %s", msg["method"], e)
+                        await response_socket.send(
+                            encoder.encode(
+                                {
+                                    "type": "error",
+                                    "rpc_id": rpc_id,
+                                    "error": str(e),
+                                }
+                            )
+                        )
+
+                elif msg_type == "shutdown":
+                    break
+
+        finally:
+            for task in tasks.values():
+                task.cancel()
+            if tasks:
+                await asyncio.gather(*tasks.values(), return_exceptions=True)
+
+            request_socket.close()
+            response_socket.close()
+            ctx.term()
+
+    # ------------------------------------------------------------------
+    # Lifecycle
+    # ------------------------------------------------------------------
+
+    def close(self) -> None:
+        """Release engine and thread pool resources."""
+        if self._closed:
+            return
+        self._closed = True
+
+        if self._engine is not None:
+            try:
+                self._engine.close()
+            except Exception as e:
+                logger.warning("Error closing diffusion engine: %s", e)
+
+        if self._executor is not None:
+            try:
+                self._executor.shutdown(wait=False)
+            except Exception as e:
+                logger.warning("Error shutting down executor: %s", e)
+
+    # ------------------------------------------------------------------
+    # Subprocess entry point
+    # ------------------------------------------------------------------
+
+    @classmethod
+    def run_diffusion_proc(
+        cls,
+        model: str,
+        od_config: OmniDiffusionConfig,
+        handshake_address: str,
+        request_address: str,
+        response_address: str,
+    ) -> None:
+        """Entry point for the diffusion subprocess."""
+        shutdown_requested = False
+
+        def signal_handler(signum: int, frame: Any) -> None:
+            nonlocal shutdown_requested
+            if not shutdown_requested:
+                shutdown_requested = True
+                raise SystemExit()
+
+        signal.signal(signal.SIGTERM, signal_handler)
+        signal.signal(signal.SIGINT, signal_handler)
+
+        proc = cls(model, od_config)
+        try:
+            proc.initialize()
+
+            # Send READY via handshake socket
+            handshake_ctx = zmq.Context()
+            handshake_socket = handshake_ctx.socket(zmq.DEALER)
+            handshake_socket.connect(handshake_address)
+            handshake_socket.send(msgspec.msgpack.encode({"status": "READY"}))
+            handshake_socket.close()
+            handshake_ctx.term()
+
+            # Run async event loop
+            asyncio.run(proc.run_loop(request_address, response_address))
+
+        except SystemExit:
+            logger.debug("StageDiffusionProc exiting.")
+            raise
+        except Exception:
+            logger.exception("StageDiffusionProc encountered a fatal error.")
+            raise
+        finally:
+            proc.close()
+
+
+# -- Free functions for backward compatibility with StageDiffusionClient ------
+
+
+def spawn_diffusion_proc(
+    model: str,
+    od_config: OmniDiffusionConfig,
+) -> tuple[BaseProcess, str, str, str]:
+    """Spawn a StageDiffusionProc subprocess.
+
+    Returns ``(proc, handshake_address, request_address, response_address)``.
+    """
+    handshake_address = get_open_zmq_ipc_path()
+    request_address = get_open_zmq_ipc_path()
+    response_address = get_open_zmq_ipc_path()
+
+    ctx = get_mp_context()
+    proc = ctx.Process(
+        target=StageDiffusionProc.run_diffusion_proc,
+        name="StageDiffusionProc",
+        kwargs={
+            "model": model,
+            "od_config": od_config,
+            "handshake_address": handshake_address,
+            "request_address": request_address,
+            "response_address": response_address,
+        },
+    )
+    proc.start()
+    # Wait for the process to become alive before returning.
+    deadline = time.monotonic() + 10
+    while not proc.is_alive():
+        if proc.exitcode is not None:
+            raise RuntimeError(f"StageDiffusionProc failed to start (exit code {proc.exitcode})")
+        if time.monotonic() > deadline:
+            raise TimeoutError("StageDiffusionProc did not become alive within 10s")
+        time.sleep(0.01)
+    return proc, handshake_address, request_address, response_address
+
+
+def complete_diffusion_handshake(
+    proc: BaseProcess,
+    handshake_address: str,
+) -> None:
+    """Wait for the diffusion subprocess to signal READY.
+
+    On failure the process is terminated before re-raising.
+    """
+    try:
+        _perform_diffusion_handshake(proc, handshake_address)
+    except Exception:
+        shutdown([proc])
+        raise
+
+
+def _perform_diffusion_handshake(
+    proc: BaseProcess,
+    handshake_address: str,
+) -> None:
+    """Run the handshake with the diffusion subprocess."""
+    with zmq_socket_ctx(handshake_address, zmq.ROUTER, bind=True) as handshake_socket:
+        poller = zmq.Poller()
+        poller.register(handshake_socket, zmq.POLLIN)
+        poller.register(proc.sentinel, zmq.POLLIN)
+
+        timeout_ms = _HANDSHAKE_POLL_TIMEOUT_S * 1000
+        while True:
+            events = dict(poller.poll(timeout=timeout_ms))
+            if not events:
+                raise TimeoutError("Timed out waiting for READY from StageDiffusionProc")
+            if handshake_socket in events:
+                identity, raw = handshake_socket.recv_multipart()
+                msg = msgspec.msgpack.decode(raw)
+                if msg.get("status") == "READY":
+                    return
+                raise RuntimeError(f"Expected READY, got: {msg}")
+            if proc.exitcode is not None:
+                raise RuntimeError(f"StageDiffusionProc died during handshake (exit code {proc.exitcode})")

vllm_omni/engine/async_omni_engine.py

@@ -32,7 +32,6 @@ from vllm.logger import init_logger
 from vllm.tokenizers import cached_tokenizer_from_config
 from vllm.v1.engine import EngineCoreRequest
 from vllm.v1.engine.input_processor import InputProcessor
-from vllm.v1.engine.utils import get_engine_zmq_addresses, launch_core_engines
 
 from vllm_omni.diffusion.data import DiffusionParallelConfig
 from vllm_omni.distributed.omni_connectors.utils.initialization import (
@@ -45,6 +44,10 @@ from vllm_omni.engine.orchestrator import Orchestrator
 from vllm_omni.engine.output_processor import MultimodalOutputProcessor
 from vllm_omni.engine.serialization import serialize_additional_information
 from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClient
+from vllm_omni.engine.stage_engine_core_proc import (
+    complete_stage_handshake,
+    spawn_stage_core,
+)
 from vllm_omni.engine.stage_init_utils import (
     StartedLlmStage,
     acquire_device_locks,
@@ -334,21 +337,17 @@ class AsyncOmniEngine:
                         engine_args_dict,
                         stage_init_timeout,
                     )
-                    addresses = get_engine_zmq_addresses(vllm_config)
-                    launch_cm = launch_core_engines(
+                    addresses, proc, handshake_address = spawn_stage_core(
                         vllm_config=vllm_config,
                         executor_class=executor_class,
                         log_stats=False,
-                        addresses=addresses,
                     )
-                    engine_manager, coordinator, addresses = launch_cm.__enter__()
                     started_stage = StartedLlmStage(
                         stage_id=metadata.stage_id,
                         metadata=metadata,
                         vllm_config=vllm_config,
                         executor_class=executor_class,
-                        engine_manager=engine_manager,
-                        coordinator=coordinator,
+                        proc=proc,
                         addresses=addresses,
                     )
                 finally:
@@ -358,7 +357,7 @@ class AsyncOmniEngine:
                         current_omni_platform.set_device_control_env_var(previous_visible_devices)
 
             logger.info("[AsyncOmniEngine] Stage %s engine launch started", metadata.stage_id)
-            launch_cm.__exit__(None, None, None)
+            complete_stage_handshake(proc, handshake_address, addresses, vllm_config)
             logger.info("[AsyncOmniEngine] Stage %s engine startup completed", metadata.stage_id)
             assert started_stage is not None
             return started_stage
@@ -389,11 +388,9 @@ class AsyncOmniEngine:
                 executor_class=started.executor_class,
                 metadata=started.metadata,
                 client_addresses=client_addresses,
-                engine_manager=started.engine_manager,
-                coordinator=started.coordinator,
+                proc=started.proc,
             )
-            started.engine_manager = None
-            started.coordinator = None
+            started.proc = None
         except Exception:
             close_started_llm_stage(started)
             raise

vllm_omni/engine/orchestrator.py

@@ -241,7 +241,7 @@ class Orchestrator:
                 # the output format in the future to simplify the processing logic in Orchestrator.
                 stage_client = self.stage_clients[stage_id]
                 if stage_client.stage_type == "diffusion":
-                    output = stage_client.get_diffusion_output_async()
+                    output = stage_client.get_diffusion_output_nowait()
                     if output is not None:
                         idle = False
                         req_state = self.request_states.get(output.request_id)

vllm_omni/engine/stage_engine_core_client.py

@@ -25,13 +25,13 @@ logger = init_logger(__name__)
 class StageEngineCoreClient(AsyncMPClient):
     """Stage async client that inherits from vLLM's AsyncMPClient.
 
-    Fully reuses AsyncMPClient.__init__ for:
+    Fully reuses AsyncMPClient for:
     - ZMQ setup, sockets
-    - launch_core_engines() -> EngineCoreProc
     - outputs_queue, output_queue_task
-    - All utility methods (shutdown, get_output_async, abort_requests_async, etc.)
+    - All utility methods (get_output_async, abort_requests_async, etc.)
 
-    This is the async version of StageMPClient, designed for use with AsyncOmniEngine.
+    The subprocess is spawned externally via ``spawn_stage_core`` /
+    ``complete_stage_handshake`` from *stage_engine_core_proc.py*.
     """
 
     def __init__(
@@ -40,6 +40,7 @@ class StageEngineCoreClient(AsyncMPClient):
         executor_class: type,
         log_stats: bool = False,
         client_addresses: dict[str, str] | None = None,
+        proc: Any = None,
         client_count: int = 1,
         client_index: int = 0,
         *,
@@ -53,6 +54,11 @@ class StageEngineCoreClient(AsyncMPClient):
         engine args building, device locking) is done by the Orchestrator
         via helpers in stage_init_utils.py. This constructor just stores metadata
         and calls super().__init__().
+
+        The subprocess is spawned externally via ``spawn_stage_core`` /
+        ``complete_stage_handshake`` (see *stage_engine_core_proc.py*).
+        The resulting ``proc`` handle is passed in so this client can
+        manage the process lifecycle on shutdown.
         """
         # -------- Stage metadata (public fields used at runtime) --------
         if metadata is not None:
@@ -69,6 +75,7 @@ class StageEngineCoreClient(AsyncMPClient):
             self.model_stage = metadata.model_stage
 
         self.engine_outputs: Any = None
+        self._proc = proc
 
         logger.info(
             "[StageEngineCoreClient] Stage-%s initializing EngineCore",
@@ -83,10 +90,6 @@ class StageEngineCoreClient(AsyncMPClient):
                 client_count=client_count,
                 client_index=client_index,
             )
-            if engine_manager is not None:
-                self.resources.engine_manager = engine_manager
-            if coordinator is not None:
-                self.resources.coordinator = coordinator
         except Exception:
             logger.exception(
                 "[StageEngineCoreClient] Stage-%s EngineCore init failed",
@@ -173,3 +176,12 @@ class StageEngineCoreClient(AsyncMPClient):
             args=args,
             kwargs=kwargs,
         )
+
+    def shutdown(self) -> None:
+        """Shutdown ZMQ connections and the subprocess."""
+        super().shutdown()
+        if self._proc is not None and self._proc.is_alive():
+            self._proc.terminate()
+            self._proc.join(timeout=5)
+            if self._proc.is_alive():
+                self._proc.kill()

vllm_omni/engine/stage_engine_core_proc.py

@@ -0,0 +1,206 @@
+"""
+Stage Core Process for vLLM-Omni V1 architecture.
+
+StageEngineCoreProc inherits from vLLM's EngineCoreProc and runs the engine core
+busy loop in a subprocess, communicating with StageEngineCoreClient via ZMQ.
+"""
+
+from __future__ import annotations
+
+import signal
+from multiprocessing.process import BaseProcess
+from typing import TYPE_CHECKING, Any
+
+import msgspec
+import zmq
+from vllm.logger import init_logger
+from vllm.transformers_utils.config import (
+    maybe_register_config_serialize_by_value,
+)
+from vllm.utils.network_utils import get_open_zmq_ipc_path, zmq_socket_ctx
+from vllm.utils.system_utils import (
+    decorate_logs,
+    get_mp_context,
+    set_process_title,
+)
+from vllm.v1.engine.core import EngineCoreProc
+from vllm.v1.engine.utils import (
+    EngineHandshakeMetadata,
+    EngineZmqAddresses,
+    get_engine_zmq_addresses,
+)
+from vllm.v1.utils import shutdown
+
+if TYPE_CHECKING:
+    from vllm.config import VllmConfig
+    from vllm.v1.executor import Executor
+
+logger = init_logger(__name__)
+
+_HANDSHAKE_POLL_TIMEOUT_S = 600
+
+
+class StageEngineCoreProc(EngineCoreProc):
+    """Stage-specific engine core process for vLLM-Omni.
+
+    Inherits from EngineCoreProc and provides its own ``run_stage_core``
+    entry point for launching in a subprocess.  Does **not** delegate to
+    ``EngineCoreProc.run_engine_core()``.
+    """
+
+    @staticmethod
+    def run_stage_core(
+        *args: Any,
+        dp_rank: int = 0,
+        local_dp_rank: int = 0,
+        **kwargs: Any,
+    ) -> None:
+        """Launch StageEngineCoreProc busy loop in background process."""
+        shutdown_requested = False
+        maybe_register_config_serialize_by_value()
+
+        def signal_handler(signum: int, frame: Any) -> None:
+            nonlocal shutdown_requested
+            if not shutdown_requested:
+                shutdown_requested = True
+                raise SystemExit()
+
+        signal.signal(signal.SIGTERM, signal_handler)
+        signal.signal(signal.SIGINT, signal_handler)
+
+        engine_core: StageEngineCoreProc | None = None
+        try:
+            vllm_config: VllmConfig = kwargs["vllm_config"]
+            parallel_config = vllm_config.parallel_config
+
+            set_process_title(f"StageEngineCoreProc_DP{dp_rank}")
+            decorate_logs()
+
+            # the current vllm-omni does not support data parallelism,
+            # so we set the data parallel size to 1.
+            # [TODO] support data parallelism in the future.
+            # https://github.com/vllm-project/vllm-omni/issues/984
+            parallel_config.data_parallel_size = 1
+            parallel_config.data_parallel_size_local = 1
+            parallel_config.data_parallel_rank = 0
+            parallel_config.data_parallel_index = dp_rank
+
+            engine_core = StageEngineCoreProc(
+                *args,
+                engine_index=dp_rank,
+                **kwargs,
+            )
+            engine_core.run_busy_loop()
+
+        except SystemExit:
+            logger.debug("StageEngineCoreProc exiting.")
+            raise
+        except Exception:
+            if engine_core is None:
+                logger.exception("StageEngineCoreProc failed to start.")
+            else:
+                logger.exception("StageEngineCoreProc encountered a fatal error.")
+                engine_core._send_engine_dead()
+            raise
+        finally:
+            if engine_core is not None:
+                engine_core.shutdown()
+
+
+def spawn_stage_core(
+    vllm_config: VllmConfig,
+    executor_class: type[Executor],
+    log_stats: bool = False,
+) -> tuple[EngineZmqAddresses, BaseProcess, str]:
+    """Spawn a *StageEngineCoreProc* subprocess without performing the handshake.
+
+    Must be called while the correct device env vars are set (e.g. under
+    the stage-launch lock).  Call ``complete_stage_handshake`` afterwards.
+
+    Returns ``(addresses, process, handshake_address)``.
+    """
+    addresses = get_engine_zmq_addresses(vllm_config)
+    handshake_address = get_open_zmq_ipc_path()
+
+    ctx = get_mp_context()
+    proc = ctx.Process(
+        target=StageEngineCoreProc.run_stage_core,
+        name="StageEngineCoreProc",
+        kwargs={
+            "vllm_config": vllm_config,
+            "local_client": True,
+            "handshake_address": handshake_address,
+            "executor_class": executor_class,
+            "log_stats": log_stats,
+            "dp_rank": 0,
+            "local_dp_rank": 0,
+        },
+    )
+    proc.start()
+    return addresses, proc, handshake_address
+
+
+def complete_stage_handshake(
+    proc: BaseProcess,
+    handshake_address: str,
+    addresses: EngineZmqAddresses,
+    vllm_config: VllmConfig,
+) -> None:
+    """Perform the HELLO/INIT/READY handshake with an already-spawned proc.
+
+    On failure the process is terminated before re-raising.
+    """
+    try:
+        _perform_handshake(proc, handshake_address, addresses, vllm_config)
+    except Exception:
+        shutdown([proc])
+        raise
+
+
+def _perform_handshake(
+    proc: BaseProcess,
+    handshake_address: str,
+    addresses: EngineZmqAddresses,
+    vllm_config: VllmConfig,
+) -> None:
+    """Run the HELLO / INIT / READY handshake with the subprocess."""
+    with zmq_socket_ctx(handshake_address, zmq.ROUTER, bind=True) as handshake_socket:
+        poller = zmq.Poller()
+        poller.register(handshake_socket, zmq.POLLIN)
+        poller.register(proc.sentinel, zmq.POLLIN)
+
+        identity, msg = _recv(poller, handshake_socket, proc, "HELLO")
+        if msg.get("status") != "HELLO":
+            raise RuntimeError(f"Expected HELLO, got: {msg}")
+
+        init_payload = EngineHandshakeMetadata(
+            addresses=addresses,
+            parallel_config={},
+        )
+        handshake_socket.send_multipart([identity, msgspec.msgpack.encode(init_payload)])
+
+        identity, msg = _recv(poller, handshake_socket, proc, "READY")
+        if msg.get("status") != "READY":
+            raise RuntimeError(f"Expected READY, got: {msg}")
+        num_gpu_blocks = msg.get("num_gpu_blocks")
+        if num_gpu_blocks is not None:
+            vllm_config.cache_config.num_gpu_blocks = num_gpu_blocks
+
+
+def _recv(
+    poller: zmq.Poller,
+    handshake_socket: zmq.Socket,
+    proc: BaseProcess,
+    expected: str,
+) -> tuple[bytes, dict]:
+    """Wait for one handshake message; raise if the process dies first."""
+    timeout_ms = _HANDSHAKE_POLL_TIMEOUT_S * 1000
+    while True:
+        events = dict(poller.poll(timeout=timeout_ms))
+        if not events:
+            raise TimeoutError(f"Timed out waiting for {expected} from StageEngineCoreProc")
+        if handshake_socket in events:
+            identity, raw = handshake_socket.recv_multipart()
+            return identity, msgspec.msgpack.decode(raw)
+        if proc.exitcode is not None:
+            raise RuntimeError(f"StageEngineCoreProc died during {expected} (exit code {proc.exitcode})")

vllm_omni/engine/stage_init_utils.py

@@ -75,6 +75,14 @@ def _resolve_model_tokenizer_paths(model: str, engine_args: dict[str, Any]) -> s
     return model
 
 
+def terminate_alive_proc(proc, timeout=5):
+    if proc.is_alive():
+        proc.terminate()
+        proc.join(timeout=timeout)
+        if proc.is_alive():
+            proc.kill()
+
+
 def resolve_worker_cls(engine_args: dict[str, Any]) -> None:
     """Resolve worker_cls from worker_type for non-diffusion stages."""
     worker_type = engine_args.get("worker_type", None)
@@ -121,8 +129,7 @@ class StartedLlmStage:
     metadata: Any
     vllm_config: Any
     executor_class: type
-    engine_manager: Any
-    coordinator: Any
+    proc: Any
     addresses: Any
 
 
@@ -446,7 +453,7 @@ def initialize_diffusion_stage(
         metadata: Extracted stage metadata.
         batch_size: Maximum number of requests to batch together in the
             diffusion engine.  Passed through to ``StageDiffusionClient``
-            and ultimately to ``AsyncOmniDiffusion``.
+            and ultimately to ``AsyncOmni``.
     """
     from vllm_omni.diffusion.data import OmniDiffusionConfig
     from vllm_omni.diffusion.stage_diffusion_client import StageDiffusionClient
@@ -461,23 +468,17 @@ def initialize_diffusion_stage(
 
 
 def close_started_llm_stage(started: StartedLlmStage) -> None:
-    """Close managers owned by a launched stage that never attached."""
-    resources = (
-        ("engine manager", started.engine_manager),
-        ("coordinator", started.coordinator),
-    )
-    for resource_name, resource in resources:
-        if resource is None:
-            continue
-        try:
-            resource.close()
-        except Exception as cleanup_error:
-            logger.warning(
-                "[stage_init] Failed to close launched %s for stage %s: %s",
-                resource_name,
-                started.stage_id,
-                cleanup_error,
-            )
+    """Terminate the subprocess owned by a launched stage that never attached."""
+    if started.proc is None:
+        return
+    try:
+        terminate_alive_proc(started.proc)
+    except Exception as cleanup_error:
+        logger.warning(
+            "[stage_init] Failed to terminate process for stage %s: %s",
+            started.stage_id,
+            cleanup_error,
+        )
 
 
 def finalize_initialized_stages(

vllm_omni/entrypoints/__init__.py

@@ -6,11 +6,9 @@ vLLM-Omni entrypoints module.
 """
 
 from vllm_omni.entrypoints.async_omni import AsyncOmni
-from vllm_omni.entrypoints.async_omni_diffusion import AsyncOmniDiffusion
 from vllm_omni.entrypoints.omni import Omni
 
 __all__ = [
     "AsyncOmni",
-    "AsyncOmniDiffusion",
     "Omni",
 ]

vllm_omni/entrypoints/async_omni_diffusion.py

@@ -1,473 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-"""
-Async entrypoint for vLLM-Omni diffusion model inference.
-
-Provides an asynchronous interface for running diffusion models,
-enabling concurrent request handling and streaming generation.
-"""
-
-import asyncio
-import uuid
-import weakref
-from collections.abc import AsyncGenerator, Iterable
-from concurrent.futures import ThreadPoolExecutor
-from typing import Any
-
-from vllm.logger import init_logger
-from vllm.transformers_utils.config import get_hf_file_to_dict
-
-from vllm_omni.diffusion.data import (
-    DiffusionRequestAbortedError,
-    OmniDiffusionConfig,
-    TransformerConfig,
-)
-from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
-from vllm_omni.diffusion.request import OmniDiffusionRequest
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType
-from vllm_omni.lora.request import LoRARequest
-from vllm_omni.outputs import OmniRequestOutput
-
-logger = init_logger(__name__)
-
-
-def _weak_close_async_omni_diffusion(engine: DiffusionEngine, executor: ThreadPoolExecutor) -> None:
-    """Best-effort diffusion cleanup for GC finalization."""
-    try:
-        engine.close()
-    except Exception:
-        pass
-    try:
-        executor.shutdown(wait=False)
-    except Exception:
-        pass
-
-
-class AsyncOmniDiffusion:
-    """Async entry point for vLLM-Omni diffusion model inference.
-
-    This class provides an asynchronous interface for running diffusion models,
-    enabling concurrent request handling. It wraps the DiffusionEngine and
-    provides async methods for image generation.
-
-    Args:
-        model: Model name or path to load
-        od_config: Optional OmniDiffusionConfig. If not provided, it will be
-            created from kwargs
-        **kwargs: Additional keyword arguments passed to OmniDiffusionConfig
-
-    Example:
-        >>> async_diffusion = AsyncOmniDiffusion(model="Qwen/Qwen-Image")
-        >>> result = await async_diffusion.generate(
-        ...     prompt="A beautiful sunset over the ocean",
-        ...     request_id="req-1",
-        ... )
-        >>> print(result.images)
-    """
-
-    def __init__(
-        self,
-        model: str,
-        od_config: OmniDiffusionConfig | None = None,
-        batch_size: int = 1,
-        **kwargs: Any,
-    ):
-        self.model = model
-
-        # Set batch size (default 1 for backward compatibility)
-        self._batch_size = max(1, batch_size)
-
-        # Capture stage info from kwargs before they might be filtered out
-        stage_id = kwargs.get("stage_id")
-        engine_input_source = kwargs.get("engine_input_source")
-        cfg_kv_collect_func = kwargs.pop("cfg_kv_collect_func", None)
-
-        # Build config
-        if od_config is None:
-            od_config = OmniDiffusionConfig.from_kwargs(model=model, **kwargs)
-        elif isinstance(od_config, dict):
-            # If config is dict, check it too (priority to kwargs if both exist)
-            if stage_id is None:
-                stage_id = od_config.get("stage_id")
-            if engine_input_source is None:
-                engine_input_source = od_config.get("engine_input_source")
-            od_config = OmniDiffusionConfig.from_kwargs(**od_config)
-
-        self.od_config = od_config
-
-        # Inject stage info into omni_kv_config if present
-        if stage_id is not None:
-            self.od_config.omni_kv_config.setdefault("stage_id", stage_id)
-        if engine_input_source is not None:
-            self.od_config.omni_kv_config.setdefault("engine_input_source", engine_input_source)
-
-        # Diffusers-style models expose `model_index.json` with `_class_name`.
-        # Non-diffusers models (e.g. Bagel, NextStep) only have `config.json`,
-        # so we fall back to reading that and mapping model_type manually.
-        try:
-            config_dict = get_hf_file_to_dict("model_index.json", od_config.model)
-            if config_dict is not None:
-                if od_config.model_class_name is None:
-                    od_config.model_class_name = config_dict.get("_class_name", None)
-                od_config.update_multimodal_support()
-
-                tf_config_dict = get_hf_file_to_dict("transformer/config.json", od_config.model)
-                od_config.set_tf_model_config(TransformerConfig.from_dict(tf_config_dict))
-            else:
-                raise FileNotFoundError("model_index.json not found")
-        except (AttributeError, OSError, ValueError, FileNotFoundError):
-            cfg = get_hf_file_to_dict("config.json", od_config.model)
-            if cfg is None:
-                raise ValueError(f"Could not find config.json or model_index.json for model {od_config.model}")
-
-            od_config.set_tf_model_config(TransformerConfig.from_dict(cfg))
-            model_type = cfg.get("model_type")
-            architectures = cfg.get("architectures") or []
-            # Bagel/NextStep models don't have a model_index.json, so we set the pipeline class name manually
-            if model_type == "bagel" or "BagelForConditionalGeneration" in architectures:
-                od_config.model_class_name = "BagelPipeline"
-                od_config.tf_model_config = TransformerConfig()
-                od_config.update_multimodal_support()
-            elif model_type == "nextstep":
-                if od_config.model_class_name is None:
-                    od_config.model_class_name = "NextStep11Pipeline"
-                od_config.tf_model_config = TransformerConfig()
-                od_config.update_multimodal_support()
-            elif architectures and len(architectures) == 1:
-                od_config.model_class_name = architectures[0]
-            else:
-                raise
-
-        if cfg_kv_collect_func is not None:
-            od_config.cfg_kv_collect_func = cfg_kv_collect_func
-
-        # Initialize engine
-        self.engine: DiffusionEngine = DiffusionEngine.make_engine(od_config)
-
-        # Thread pool for running sync engine in async context
-        self._executor = ThreadPoolExecutor(max_workers=1)
-        self._closed = False
-        self._weak_finalizer = weakref.finalize(
-            self,
-            _weak_close_async_omni_diffusion,
-            self.engine,
-            self._executor,
-        )
-
-        logger.info("AsyncOmniDiffusion initialized with model: %s, batch_size: %d", model, self._batch_size)
-
-    # ------------------------------------------------------------------
-    # batch_size property
-    # ------------------------------------------------------------------
-
-    @property
-    def batch_size(self) -> int:
-        """Return the configured batch size for request batching."""
-        return self._batch_size
-
-    @batch_size.setter
-    def batch_size(self, value: int) -> None:
-        if not isinstance(value, int) or value < 1:
-            raise ValueError("batch_size must be a positive integer")
-        self._batch_size = value
-
-    # ------------------------------------------------------------------
-    # Public batch generation API
-    # ------------------------------------------------------------------
-
-    async def generate_batch(
-        self,
-        prompts: list[OmniPromptType],
-        sampling_params: OmniDiffusionSamplingParams,
-        request_id: str | None = None,
-        lora_request: LoRARequest | None = None,
-    ) -> OmniRequestOutput:
-        """Generate images from multiple prompts in a single engine call.
-
-        Batches the given prompts into **one** ``DiffusionEngine.step()``
-        call and returns a single ``OmniRequestOutput`` containing all
-        generated images.  Called by ``StageDiffusionClient._run_batch``
-        when the orchestrator receives a list-prompt request.
-
-        Args:
-            prompts: List of text prompts describing the desired images.
-            sampling_params: Shared sampling parameters for all prompts.
-            request_id: Optional unique identifier. Auto-generated when *None*.
-            lora_request: Optional LoRA adapter to apply.
-
-        Returns:
-            A single ``OmniRequestOutput`` with all images combined.
-        """
-        if request_id is None:
-            request_id = f"diff-batch-{uuid.uuid4().hex[:8]}"
-        return await self._generate_batch(prompts, sampling_params, request_id, lora_request)
-
-    # ------------------------------------------------------------------
-    # Internal batch generation
-    # ------------------------------------------------------------------
-
-    async def _generate_batch(
-        self,
-        prompts: list[OmniPromptType],
-        sampling_params: OmniDiffusionSamplingParams,
-        request_id: str,
-        lora_request: LoRARequest | None = None,
-    ) -> OmniRequestOutput:
-        """Generate images from multiple prompts in a single engine call."""
-        if not prompts:
-            return OmniRequestOutput(request_id=request_id, images=[], final_output_type="image")
-
-        if sampling_params.guidance_scale:
-            sampling_params.guidance_scale_provided = True
-
-        if lora_request is not None:
-            sampling_params.lora_request = lora_request
-
-        request = OmniDiffusionRequest(
-            prompts=prompts,
-            sampling_params=sampling_params,
-            request_ids=[f"{request_id}-{i}" for i in range(len(prompts))],
-        )
-
-        logger.debug("Starting batch generation for %d prompts, request_id=%s", len(prompts), request_id)
-
-        loop = asyncio.get_event_loop()
-        try:
-            results = await loop.run_in_executor(
-                self._executor,
-                self.engine.step,
-                request,
-            )
-        except Exception as e:
-            logger.error("Batch generation failed for request %s: %s", request_id, e)
-            raise RuntimeError(f"Diffusion batch generation failed: {e}") from e
-
-        # Combine all per-prompt results into a single OmniRequestOutput
-        all_images = []
-        for result in results:
-            all_images.extend(result.images)
-
-        return OmniRequestOutput(
-            request_id=request_id,
-            images=all_images,
-            final_output_type="image",
-            finished=True,
-        )
-
-    def get_diffusion_od_config(self) -> OmniDiffusionConfig:
-        """Return the diffusion config used by this engine."""
-        return self.od_config
-
-    # ------------------------------------------------------------------
-    # Public generate API
-    # ------------------------------------------------------------------
-
-    async def generate(
-        self,
-        prompt: OmniPromptType,
-        sampling_params: OmniDiffusionSamplingParams,
-        request_id: str | None = None,
-        lora_request: LoRARequest | None = None,
-    ) -> OmniRequestOutput:
-        """Generate images asynchronously from a single text prompt.
-
-        For batched generation (multiple prompts in one engine call), use
-        :meth:`generate_batch` instead.  This method always processes
-        exactly one prompt per call.
-
-        Args:
-            prompt: Text prompt describing the desired image
-            sampling_params: Sampling parameters
-            request_id: Optional unique identifier for tracking the request
-            lora_request: Optional LoRA adapter to apply
-
-        Returns:
-            OmniRequestOutput containing generated images
-
-        Raises:
-            RuntimeError: If generation fails
-        """
-        if request_id is None:
-            request_id = f"diff-{uuid.uuid4().hex[:16]}"
-        if sampling_params.guidance_scale:
-            sampling_params.guidance_scale_provided = True
-
-        if lora_request is not None:
-            sampling_params.lora_request = lora_request
-
-        request = OmniDiffusionRequest(
-            prompts=[prompt],
-            sampling_params=sampling_params,
-            request_ids=[request_id],
-        )
-
-        logger.debug("Starting generation for request %s", request_id)
-
-        loop = asyncio.get_event_loop()
-        try:
-            result = await loop.run_in_executor(
-                self._executor,
-                self.engine.step,
-                request,
-            )
-            result = result[0]
-        except asyncio.CancelledError:
-            self.engine.abort(request_id)
-            raise
-        except DiffusionRequestAbortedError:
-            raise
-        except Exception as e:
-            logger.error("Generation failed for request %s: %s", request_id, e)
-            raise RuntimeError(f"Diffusion generation failed: {e}") from e
-
-        if not result.request_id:
-            result.request_id = request_id
-        return result
-
-    async def generate_stream(
-        self,
-        prompt: str,
-        request_id: str | None = None,
-        **kwargs: Any,
-    ) -> AsyncGenerator[OmniRequestOutput, None]:
-        """Generate images with streaming progress updates.
-
-        Currently, diffusion models don't support true streaming, so this
-        yields a single result after generation completes. Future implementations
-        may support step-by-step progress updates.
-
-        Args:
-            prompt: Text prompt describing the desired image
-            request_id: Optional unique identifier for tracking the request
-            **kwargs: Additional generation parameters
-
-        Yields:
-            OmniRequestOutput with generation progress/results
-        """
-        result = await self.generate(prompt=prompt, request_id=request_id, **kwargs)
-        yield result
-
-    def close(self) -> None:
-        """Close the engine and release resources.
-
-        Should be called when done using the AsyncOmniDiffusion instance.
-        """
-        if self._closed:
-            return
-        self._closed = True
-
-        finalizer = getattr(self, "_weak_finalizer", None)
-        if finalizer is not None and finalizer.alive:
-            finalizer.detach()
-
-        try:
-            self.engine.close()
-        except Exception as e:
-            logger.warning("Error closing diffusion engine: %s", e)
-
-        try:
-            self._executor.shutdown(wait=False)
-        except Exception as e:
-            logger.warning("Error shutting down executor: %s", e)
-
-        logger.info("AsyncOmniDiffusion closed")
-
-    def shutdown(self) -> None:
-        """Alias for close() method."""
-        self.close()
-
-    async def abort(self, request_id: str | Iterable[str]) -> None:
-        """Abort a request."""
-        self.engine.abort(request_id)
-
-    @property
-    def is_running(self) -> bool:
-        """Check if the engine is running."""
-        return not self._closed
-
-    @property
-    def is_stopped(self) -> bool:
-        """Check if the engine is stopped."""
-        return self._closed
-
-    async def remove_lora(self, adapter_id: int) -> bool:
-        """Remove a LoRA"""
-        loop = asyncio.get_event_loop()
-        results = await loop.run_in_executor(
-            self._executor,
-            self.engine.collective_rpc,
-            "remove_lora",
-            None,
-            (adapter_id,),
-            {},
-            None,
-        )
-        return all(results) if isinstance(results, list) else results
-
-    async def add_lora(self, lora_request: LoRARequest) -> bool:
-        """Add a LoRA adapter"""
-        loop = asyncio.get_event_loop()
-        results = await loop.run_in_executor(
-            self._executor,
-            self.engine.collective_rpc,
-            "add_lora",
-            None,
-            (),
-            {"lora_request": lora_request},
-            None,
-        )
-        return all(results) if isinstance(results, list) else results
-
-    async def list_loras(self) -> list[int]:
-        """List all registered LoRA adapter IDs."""
-        loop = asyncio.get_event_loop()
-        results = await loop.run_in_executor(
-            self._executor,
-            self.engine.collective_rpc,
-            "list_loras",
-            None,
-            (),
-            {},
-            None,
-        )
-        # collective_rpc returns list from workers; flatten unique ids
-        if not isinstance(results, list):
-            return results or []
-        merged: set[int] = set()
-        for part in results:
-            merged.update(part or [])
-        return sorted(merged)
-
-    async def pin_lora(self, lora_id: int) -> bool:
-        """Prevent an adapter from being evicted."""
-        loop = asyncio.get_event_loop()
-        results = await loop.run_in_executor(
-            self._executor,
-            self.engine.collective_rpc,
-            "pin_lora",
-            None,
-            (),
-            {"adapter_id": lora_id},
-            None,
-        )
-        return all(results) if isinstance(results, list) else results
-
-    async def profile(self, is_start: bool = True, profile_prefix: str | None = None) -> None:
-        """Start or stop profiling for the diffusion model.
-
-        Args:
-            is_start: True to start profiling, False to stop.
-            profile_prefix: Optional prefix for trace filename (vLLM compat).
-
-        Note:
-            Matches vLLM's worker.profile() signature for consistency.
-            Traces are saved automatically via on_trace_ready callback.
-        """
-        loop = asyncio.get_event_loop()
-        await loop.run_in_executor(
-            self._executor,
-            self.engine.profile,
-            is_start,
-            profile_prefix,
-        )

vllm_omni/entrypoints/openai/serving_chat.py

@@ -6,7 +6,7 @@ import uuid
 from collections.abc import AsyncGenerator, AsyncIterator, Callable
 from datetime import datetime, timedelta, timezone
 from io import BytesIO
-from typing import TYPE_CHECKING, Any, Final, Optional, cast
+from typing import Any, Final, cast
 
 import jinja2
 import torch
@@ -89,9 +89,6 @@ from vllm_omni.entrypoints.openai.utils import parse_lora_request
 from vllm_omni.lora.request import LoRARequest
 from vllm_omni.outputs import OmniRequestOutput
 
-if TYPE_CHECKING:
-    from vllm_omni.entrypoints.async_omni_diffusion import AsyncOmniDiffusion
-
 logger = init_logger(__name__)
 
 
@@ -107,13 +104,13 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
 
     # Diffusion mode attributes
     _diffusion_mode: bool = False
-    _diffusion_engine: Optional["AsyncOmniDiffusion"] = None
+    _diffusion_engine: AsyncOmni | None = None
     _diffusion_model_name: str = ""
 
     @classmethod
     def for_diffusion(
         cls,
-        diffusion_engine: "AsyncOmniDiffusion",
+        diffusion_engine: AsyncOmni,
         model_name: str,
     ) -> "OmniOpenAIServingChat":
         """Create a chat serving instance for diffusion models.
@@ -2153,7 +2150,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
             if resolution is not None:
                 gen_params.resolution = resolution
 
-            # Parse per-request LoRA (works for both AsyncOmniDiffusion and AsyncOmni).
+            # Parse per-request LoRA.
             if lora_body and isinstance(lora_body, dict):
                 try:
                     lora_req, lora_scale = parse_lora_request(lora_body)
@@ -2187,26 +2184,16 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                         )
 
             # Generate image
-            # Handle both AsyncOmniDiffusion (returns OmniRequestOutput) and AsyncOmni (returns AsyncGenerator)
-            if isinstance(self._diffusion_engine, AsyncOmni):
-                diffusion_engine = cast(AsyncOmni, self._diffusion_engine)
-                result = None
-                async for output in diffusion_engine.generate(
-                    prompt=gen_prompt,
-                    sampling_params_list=[gen_params],  # Pass as single-stage params
-                    request_id=request_id,
-                ):
-                    result = output
-                if result is None:
-                    return self._create_error_response("No output generated from AsyncOmni")
-            else:
-                # AsyncOmniDiffusion: direct call
-                diffusion_engine = cast(AsyncOmniDiffusion, self._diffusion_engine)
-                result = await diffusion_engine.generate(
-                    prompt=gen_prompt,
-                    sampling_params=gen_params,
-                    request_id=request_id,
-                )
+            diffusion_engine = cast(AsyncOmni, self._diffusion_engine)
+            result = None
+            async for output in diffusion_engine.generate(
+                prompt=gen_prompt,
+                sampling_params_list=[gen_params],  # Pass as single-stage params
+                request_id=request_id,
+            ):
+                result = output
+            if result is None:
+                return self._create_error_response("No output generated from AsyncOmni")
             # Extract images from result
             # Handle nested OmniRequestOutput structure where images might be in request_output
             images = getattr(result.request_output, "images", [])