[Perf] Eliminate Hop 3 IPC overhead for single-stage diffusion via inline execution (#2736)

Signed-off-by: samithuang <285365963@qq.com>
Signed-off-by: Samit <285365963@qq.com>

tests/diffusion/test_inline_stage_diffusion_client.py

@@ -0,0 +1,96 @@
+from __future__ import annotations
+
+import asyncio
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.inline_stage_diffusion_client import InlineStageDiffusionClient
+from vllm_omni.engine.stage_init_utils import StageMetadata
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.outputs import OmniRequestOutput
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.fixture
+def mock_engine():
+    with patch("vllm_omni.diffusion.inline_stage_diffusion_client.DiffusionEngine") as mock:
+        engine_instance = MagicMock()
+        mock.make_engine.return_value = engine_instance
+        yield engine_instance
+
+
+@pytest.fixture
+def client(mock_engine):
+    metadata = StageMetadata(
+        stage_id=0,
+        stage_type="diffusion",
+        engine_output_type="image",
+        is_comprehension=False,
+        requires_multimodal_data=False,
+        engine_input_source="prompt",
+        final_output=True,
+        final_output_type="image",
+        default_sampling_params={},
+        custom_process_input_func=None,
+        model_stage=None,
+        runtime_cfg=None,
+    )
+    with patch.object(InlineStageDiffusionClient, "_enrich_config"):
+        od_config = MagicMock(spec=OmniDiffusionConfig)
+        c = InlineStageDiffusionClient(model="test_model", od_config=od_config, metadata=metadata, batch_size=1)
+        yield c
+        c.shutdown()
+
+
+@pytest.mark.asyncio
+async def test_inline_dispatch_request_success(client, mock_engine):
+    # Setup mock engine step to return a successful result
+    mock_result = OmniRequestOutput.from_diffusion(request_id="req-1", images=[MagicMock()])
+    mock_engine.step.return_value = [mock_result]
+
+    sampling_params = OmniDiffusionSamplingParams()
+    await client.add_request_async("req-1", "A test prompt", sampling_params)
+
+    # Wait for the task to be processed
+    for _ in range(10):
+        output = client.get_diffusion_output_nowait()
+        if output is not None:
+            break
+        await asyncio.sleep(0.01)
+
+    assert output is not None
+    assert output.request_id == "req-1"
+    mock_engine.step.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_inline_dispatch_request_error(client, mock_engine):
+    # Setup mock engine step to raise an exception
+    mock_engine.step.side_effect = RuntimeError("Engine failure")
+
+    sampling_params = OmniDiffusionSamplingParams()
+    await client.add_request_async("req-err", "A test prompt", sampling_params)
+
+    for _ in range(10):
+        output = client.get_diffusion_output_nowait()
+        if output is not None:
+            break
+        await asyncio.sleep(0.01)
+
+    assert output is not None
+    assert output.request_id == "req-err"
+    assert output.error == "Engine failure"
+    assert not output.images
+
+
+def test_inline_shutdown(client, mock_engine):
+    assert not client._shutting_down
+
+    # Shutting down should cleanly cancel anything queued and close engine
+    client.shutdown()
+
+    assert client._shutting_down
+    mock_engine.close.assert_called_once()

tests/engine/test_async_omni_engine_stage_init.py

@@ -100,10 +100,11 @@ def test_initialize_stages_passes_stage_init_timeout_to_diffusion_handshake(monk
     engine.log_stats = False
     engine.model = "dummy-model"
     engine.config_path = "dummy-config"
-    engine.num_stages = 1
+    engine.num_stages = 2
     engine.async_chunk = False
     engine.diffusion_batch_size = 1
     engine.single_stage_mode = False
+    engine._omni_master_server = None
     engine.stage_configs = [types.SimpleNamespace(stage_id=0, stage_type="diffusion", engine_args={})]
 
     metadata = types.SimpleNamespace(

vllm_omni/diffusion/data.py

@@ -666,6 +666,49 @@ class OmniDiffusionConfig:
     def update_multimodal_support(self) -> None:
         self.supports_multimodal_inputs = self.model_class_name in {"QwenImageEditPlusPipeline"}
 
+    def enrich_config(self) -> None:
+        """Load model metadata from HuggingFace and populate 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.
+        """
+        from vllm.transformers_utils.config import get_hf_file_to_dict
+
+        try:
+            config_dict = get_hf_file_to_dict("model_index.json", self.model)
+            if config_dict is not None:
+                if self.model_class_name is None:
+                    self.model_class_name = config_dict.get("_class_name", None)
+                self.update_multimodal_support()
+
+                tf_config_dict = get_hf_file_to_dict("transformer/config.json", self.model)
+                self.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", self.model)
+            if cfg is None:
+                raise ValueError(f"Could not find config.json or model_index.json for model {self.model}")
+
+            self.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:
+                self.model_class_name = "BagelPipeline"
+                self.tf_model_config = TransformerConfig()
+                self.update_multimodal_support()
+            elif model_type == "nextstep":
+                if self.model_class_name is None:
+                    self.model_class_name = "NextStep11Pipeline"
+                self.tf_model_config = TransformerConfig()
+                self.update_multimodal_support()
+            elif architectures and len(architectures) == 1:
+                self.model_class_name = architectures[0]
+            else:
+                raise
+
     @classmethod
     def from_kwargs(cls, **kwargs: Any) -> "OmniDiffusionConfig":
         # Backwards-compatibility: older callers may use a diffusion-specific

vllm_omni/diffusion/inline_stage_diffusion_client.py

@@ -0,0 +1,348 @@
+"""Inline Stage Diffusion Client for vLLM-Omni multi-stage runtime.
+
+Runs DiffusionEngine in a ThreadPoolExecutor inside the Orchestrator process
+instead of spawning a separate StageDiffusionProc subprocess, eliminating ZMQ
+IPC overhead. Used when there is only a single diffusion stage.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import time
+from concurrent.futures import ThreadPoolExecutor
+from typing import TYPE_CHECKING, Any
+
+import torch
+from PIL import Image
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.data import DiffusionRequestAbortedError
+from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.engine.stage_init_utils import StageMetadata
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.outputs import OmniRequestOutput
+
+if TYPE_CHECKING:
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+    from vllm_omni.inputs.data import OmniPromptType
+
+logger = init_logger(__name__)
+
+
+class InlineStageDiffusionClient:
+    """Runs DiffusionEngine in a thread executor inside the Orchestrator."""
+
+    stage_type: str = "diffusion"
+
+    def __init__(
+        self,
+        model: str,
+        od_config: OmniDiffusionConfig,
+        metadata: StageMetadata,
+        batch_size: int = 1,
+    ) -> None:
+        self.model = model
+        self.od_config = od_config
+        self.stage_id = metadata.stage_id
+        self.final_output = metadata.final_output
+        self.final_output_type = metadata.final_output_type
+        self.default_sampling_params = metadata.default_sampling_params
+        self.custom_process_input_func = metadata.custom_process_input_func
+        self.engine_input_source = metadata.engine_input_source
+        self.batch_size = batch_size
+
+        self._enrich_config()
+        self._engine = DiffusionEngine.make_engine(self.od_config)
+        self._executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="inline-diffusion")
+
+        self._output_queue: asyncio.Queue[OmniRequestOutput] = asyncio.Queue()
+        self._tasks: dict[str, asyncio.Task] = {}
+        self._shutting_down = False
+
+        logger.info(
+            "[InlineStageDiffusionClient] Stage-%s initialized inline (batch_size=%d)",
+            self.stage_id,
+            self.batch_size,
+        )
+
+    def _enrich_config(self) -> None:
+        """Load model metadata from HuggingFace and populate od_config fields."""
+        self.od_config.enrich_config()
+
+    # ------------------------------------------------------------------
+    # Request processing
+    # ------------------------------------------------------------------
+
+    async def add_request_async(
+        self,
+        request_id: str,
+        prompt: OmniPromptType,
+        sampling_params: OmniDiffusionSamplingParams,
+        kv_sender_info: dict[int, dict[str, Any]] | None = None,
+    ) -> None:
+        task = asyncio.create_task(
+            self._dispatch_request(
+                request_id,
+                prompt,
+                sampling_params,
+                kv_sender_info,
+            )
+        )
+        self._tasks[request_id] = task
+
+    async def _dispatch_request(
+        self,
+        request_id: str,
+        prompt: Any,
+        sampling_params: OmniDiffusionSamplingParams,
+        kv_sender_info: dict[str, Any] | None = None,
+    ) -> None:
+        try:
+            request = OmniDiffusionRequest(
+                prompts=[prompt],
+                sampling_params=sampling_params,
+                request_ids=[request_id],
+                request_id=request_id,
+                kv_sender_info=kv_sender_info,
+            )
+
+            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
+
+            self._output_queue.put_nowait(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)
+            error_output = OmniRequestOutput.from_diffusion(
+                request_id=request_id,
+                images=[],
+            )
+            error_output.error = str(e)
+            self._output_queue.put_nowait(error_output)
+        finally:
+            self._tasks.pop(request_id, None)
+
+    async def add_batch_request_async(
+        self,
+        request_id: str,
+        prompts: list[OmniPromptType],
+        sampling_params: OmniDiffusionSamplingParams,
+        kv_sender_info: dict[int, dict[str, Any]] | None = None,
+    ) -> None:
+        task = asyncio.create_task(
+            self._dispatch_batch(
+                request_id,
+                prompts,
+                sampling_params,
+                kv_sender_info,
+            )
+        )
+        self._tasks[request_id] = task
+
+    async def _dispatch_batch(
+        self,
+        request_id: str,
+        prompts: list[Any],
+        sampling_params: OmniDiffusionSamplingParams,
+        kv_sender_info: dict[str, Any] | None = None,
+    ) -> None:
+        try:
+            request = OmniDiffusionRequest(
+                prompts=prompts,
+                sampling_params=sampling_params,
+                request_ids=[f"{request_id}-{i}" for i in range(len(prompts))],
+                request_id=request_id,
+                kv_sender_info=kv_sender_info,
+            )
+
+            loop = asyncio.get_running_loop()
+            results = await loop.run_in_executor(self._executor, self._engine.step, request)
+
+            all_images: list = []
+            merged_mm: dict[str, Any] = {}
+            merged_metrics: dict[str, Any] = {}
+            merged_durations: dict[str, float] = {}
+            merged_custom: dict[str, Any] = {}
+            peak_mem = 0.0
+            latents = None
+            trajectory_latents: list[torch.Tensor] | None = None
+            trajectory_timesteps: list[torch.Tensor] | None = None
+            trajectory_log_probs: torch.Tensor | None = None
+            trajectory_decoded: list[Image.Image] | None = 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)
+                merged_custom.update(r._custom_output)
+                peak_mem = max(peak_mem, r.peak_memory_mb)
+                if latents is None and r.latents is not None:
+                    latents = r.latents
+                if trajectory_latents is None:
+                    trajectory_latents = r.trajectory_latents
+                if trajectory_timesteps is None:
+                    trajectory_timesteps = r.trajectory_timesteps
+                if trajectory_log_probs is None:
+                    trajectory_log_probs = r.trajectory_log_probs
+                if trajectory_decoded is None:
+                    trajectory_decoded = r.trajectory_decoded
+                if r.final_output_type != "image":
+                    final_output_type = r.final_output_type
+
+            result = OmniRequestOutput.from_diffusion(
+                request_id=request_id,
+                images=all_images,
+                prompt=prompts[0] if len(prompts) == 1 else None,
+                metrics=merged_metrics,
+                latents=latents,
+                trajectory_latents=trajectory_latents,
+                trajectory_timesteps=trajectory_timesteps,
+                trajectory_log_probs=trajectory_log_probs,
+                trajectory_decoded=trajectory_decoded,
+                custom_output=merged_custom or None,
+                multimodal_output=merged_mm or None,
+                final_output_type=final_output_type,
+                stage_durations=merged_durations,
+                peak_memory_mb=peak_mem,
+            )
+
+            self._output_queue.put_nowait(result)
+        except DiffusionRequestAbortedError as e:
+            logger.info("request_id: %s aborted: %s", request_id, str(e))
+        except Exception as e:
+            logger.exception("Batch diffusion request %s failed: %s", request_id, e)
+            error_output = OmniRequestOutput.from_diffusion(
+                request_id=request_id,
+                images=[],
+            )
+            error_output.error = str(e)
+            self._output_queue.put_nowait(error_output)
+        finally:
+            self._tasks.pop(request_id, None)
+
+    def get_diffusion_output_nowait(self) -> OmniRequestOutput | None:
+        try:
+            return self._output_queue.get_nowait()
+        except asyncio.QueueEmpty:
+            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()
+            self._engine.abort(rid)
+
+    async def collective_rpc_async(
+        self,
+        method: str,
+        timeout: float | None = None,
+        args: tuple[Any, ...] = (),
+        kwargs: dict[str, Any] | None = None,
+    ) -> Any:
+        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
+            if is_start and profile_prefix is None:
+                profile_prefix = f"stage_{self.stage_id}_diffusion_{int(time.time())}"
+            return await loop.run_in_executor(
+                self._executor,
+                self._engine.profile,
+                is_start,
+                profile_prefix,
+            )
+
+        kwargs = kwargs or {}
+
+        # LoRA methods
+        if method == "add_lora":
+            lora_request = args[0] if args else kwargs.get("lora_request")
+            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,
+                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
+
+        return await loop.run_in_executor(
+            self._executor,
+            self._engine.collective_rpc,
+            method,
+            timeout,
+            args,
+            kwargs,
+            None,
+        )
+
+    def shutdown(self) -> None:
+        self._shutting_down = True
+
+        # Cancel all pending tasks
+        for task in self._tasks.values():
+            task.cancel()
+
+        try:
+            # Cancel queued futures and wait for the running one to complete deterministically
+            self._executor.shutdown(wait=True, cancel_futures=True)
+        except Exception:
+            pass
+
+        try:
+            self._engine.close()
+        except Exception:
+            pass

vllm_omni/diffusion/stage_diffusion_client.py

@@ -34,6 +34,24 @@ if TYPE_CHECKING:
 logger = init_logger(__name__)
 
 
+def create_diffusion_client(
+    model: str,
+    od_config: OmniDiffusionConfig,
+    metadata: StageMetadata,
+    stage_init_timeout: int,
+    batch_size: int = 1,
+    use_inline: bool = False,
+) -> Any:
+    """Factory to create either an inline or out-of-process diffusion client."""
+    if use_inline:
+        from vllm_omni.diffusion.inline_stage_diffusion_client import InlineStageDiffusionClient
+
+        return InlineStageDiffusionClient(model, od_config, metadata, batch_size=batch_size)
+    return StageDiffusionClient(
+        model, od_config, metadata, stage_init_timeout=stage_init_timeout, batch_size=batch_size
+    )
+
+
 class StageDiffusionClient:
     """Communicates with StageDiffusionProc via ZMQ for use inside the Orchestrator.
 
@@ -154,6 +172,13 @@ class StageDiffusionClient:
                         "error": True,
                         "reason": error_msg,
                     }
+                elif req_id is not None:
+                    error_output = OmniRequestOutput.from_diffusion(
+                        request_id=req_id,
+                        images=[],
+                    )
+                    error_output.error = error_msg
+                    self._output_queue.put_nowait(error_output)
 
     # Fields that are subprocess-local and cannot be serialized across
     # process boundaries.  They are recreated in the subprocess with

vllm_omni/diffusion/stage_diffusion_proc.py

@@ -19,12 +19,11 @@ import zmq
 import zmq.asyncio
 from PIL import Image
 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.data import DiffusionRequestAbortedError
 from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.distributed.omni_connectors.utils.serialization import (
@@ -66,47 +65,8 @@ class StageDiffusionProc:
         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
+        """Load model metadata from HuggingFace and populate od_config fields."""
+        self._od_config.enrich_config()
 
     # ------------------------------------------------------------------
     # Request processing

vllm_omni/engine/async_omni_engine.py

@@ -759,12 +759,14 @@ class AsyncOmniEngine:
                                         self._omni_master_server,
                                     )
                                 else:
+                                    use_inline = True if self.num_stages == 1 else False
                                     stage_clients[stage_idx] = initialize_diffusion_stage(
                                         self.model,
                                         stage_cfg,
                                         metadata,
                                         stage_init_timeout=stage_init_timeout,
                                         batch_size=self.diffusion_batch_size,
+                                        use_inline=use_inline,
                                     )
                                 logger.info(
                                     "[AsyncOmniEngine] Stage %s initialized (diffusion, batch_size=%d)",

vllm_omni/engine/orchestrator.py

@@ -246,6 +246,23 @@ class Orchestrator:
                         idle = False
                         req_state = self.request_states.get(output.request_id)
                         if req_state is not None:
+                            if getattr(output, "error", None) is not None:
+                                parent_id = self._companion_to_parent.get(output.request_id, output.request_id)
+                                await self.output_async_queue.put(
+                                    {
+                                        "type": "error",
+                                        "request_id": parent_id,
+                                        "stage_id": stage_id,
+                                        "error": output.error,
+                                    }
+                                )
+                                role_map = self._companion_map.get(parent_id, {})
+                                for cid in role_map.values():
+                                    self.request_states.pop(cid, None)
+                                self._cleanup_companion_state(parent_id)
+                                self.request_states.pop(parent_id, None)
+                                continue
+
                             stage_metrics = self._build_stage_metrics(stage_id, output.request_id, [output], req_state)
                             await self._route_output(stage_id, output, req_state, stage_metrics)
                     continue

vllm_omni/engine/stage_init_utils.py

@@ -530,6 +530,7 @@ def initialize_diffusion_stage(
     metadata: StageMetadata,
     stage_init_timeout: int,
     batch_size: int = 1,
+    use_inline: bool = False,
 ) -> Any:
     """Build a diffusion stage client.
 
@@ -541,13 +542,12 @@ def initialize_diffusion_stage(
         batch_size: Maximum number of requests to batch together in the
             diffusion engine.  Passed through to ``StageDiffusionClient``
             and ultimately to ``AsyncOmni``.
+        use_inline: If True, uses the inline diffusion client instead of subprocess.
     """
-    from vllm_omni.diffusion.stage_diffusion_client import StageDiffusionClient
+    from vllm_omni.diffusion.stage_diffusion_client import create_diffusion_client
 
     od_config = build_diffusion_config(model, stage_cfg, metadata)
-    return StageDiffusionClient(
-        model, od_config, metadata, stage_init_timeout=stage_init_timeout, batch_size=batch_size
-    )
+    return create_diffusion_client(model, od_config, metadata, stage_init_timeout, batch_size, use_inline)
 
 
 def _shutdown_or_close_resource(resource: Any, resource_name: str, stage_id: int) -> None:

vllm_omni/entrypoints/openai/video_api_utils.py

@@ -227,6 +227,9 @@ def _encode_video_bytes(
         frames_np *= 255.0
         frames_u8 = np.round(frames_np).astype(np.uint8)
 
+    # Ensure contiguous memory layout for faster PyAV muxing
+    frames_u8 = np.ascontiguousarray(frames_u8)
+
     audio_np = _coerce_audio_to_numpy(audio) if audio is not None else None
 
     return mux_video_audio_bytes(

vllm_omni/outputs.py

@@ -100,6 +100,9 @@ class OmniRequestOutput:
     # memory usage info
     peak_memory_mb: float = 0.0
 
+    # error handling
+    error: str | None = None
+
     @classmethod
     def from_pipeline(
         cls,