From 61a3cbdff5785290501d711717e2b2e526ffe34f Mon Sep 17 00:00:00 2001 From: Samit <285365963@qq.com> Date: Wed, 15 Apr 2026 11:46:06 +0800 Subject: [PATCH] [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> --- .../test_inline_stage_diffusion_client.py | 96 +++++ .../test_async_omni_engine_stage_init.py | 3 +- vllm_omni/diffusion/data.py | 43 +++ .../inline_stage_diffusion_client.py | 348 ++++++++++++++++++ vllm_omni/diffusion/stage_diffusion_client.py | 25 ++ vllm_omni/diffusion/stage_diffusion_proc.py | 46 +-- vllm_omni/engine/async_omni_engine.py | 2 + vllm_omni/engine/orchestrator.py | 17 + vllm_omni/engine/stage_init_utils.py | 8 +- .../entrypoints/openai/video_api_utils.py | 3 + vllm_omni/outputs.py | 3 + 11 files changed, 546 insertions(+), 48 deletions(-) create mode 100644 tests/diffusion/test_inline_stage_diffusion_client.py create mode 100644 vllm_omni/diffusion/inline_stage_diffusion_client.py diff --git a/tests/diffusion/test_inline_stage_diffusion_client.py b/tests/diffusion/test_inline_stage_diffusion_client.py new file mode 100644 index 000000000..385f39b12 --- /dev/null +++ b/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() diff --git a/tests/engine/test_async_omni_engine_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py index 7b995fe70..84b0cb0be 100644 --- a/tests/engine/test_async_omni_engine_stage_init.py +++ b/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( diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 56a891aa5..fca0a5bad 100644 --- a/vllm_omni/diffusion/data.py +++ b/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 diff --git a/vllm_omni/diffusion/inline_stage_diffusion_client.py b/vllm_omni/diffusion/inline_stage_diffusion_client.py new file mode 100644 index 000000000..a33a3e956 --- /dev/null +++ b/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 diff --git a/vllm_omni/diffusion/stage_diffusion_client.py b/vllm_omni/diffusion/stage_diffusion_client.py index 7e740dc89..480d113d1 100644 --- a/vllm_omni/diffusion/stage_diffusion_client.py +++ b/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 diff --git a/vllm_omni/diffusion/stage_diffusion_proc.py b/vllm_omni/diffusion/stage_diffusion_proc.py index cef697630..eced444fd 100644 --- a/vllm_omni/diffusion/stage_diffusion_proc.py +++ b/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 diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 9609cf6e2..054d5342d 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/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)", diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index 386b545eb..0fdab9c0d 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/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 diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index 158b4c547..bf40aa77c 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/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: diff --git a/vllm_omni/entrypoints/openai/video_api_utils.py b/vllm_omni/entrypoints/openai/video_api_utils.py index 193546979..3fb991225 100644 --- a/vllm_omni/entrypoints/openai/video_api_utils.py +++ b/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( diff --git a/vllm_omni/outputs.py b/vllm_omni/outputs.py index 2c2c1d21c..c02c0c142 100644 --- a/vllm_omni/outputs.py +++ b/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,