From 1e8dc841503146bcb2b5af01b36d1eca94dd8e24 Mon Sep 17 00:00:00 2001 From: Haco <923390377@qq.com> Date: Wed, 6 May 2026 22:21:29 +0800 Subject: [PATCH] [Bugfix] Fix default diffusion stage config generator drops runtime engine args (#2559) Signed-off-by: xiaohajiayou <923390377@qq.com> Co-authored-by: reidliu41 --- .../test_async_omni_diffusion_config.py | 40 +++++++++++++++++++ vllm_omni/engine/async_omni_engine.py | 4 ++ 2 files changed, 44 insertions(+) diff --git a/tests/entrypoints/test_async_omni_diffusion_config.py b/tests/entrypoints/test_async_omni_diffusion_config.py index 235add572..f06ede401 100644 --- a/tests/entrypoints/test_async_omni_diffusion_config.py +++ b/tests/entrypoints/test_async_omni_diffusion_config.py @@ -131,6 +131,46 @@ def test_default_stage_config_includes_default_sampling_params(): } +def test_default_stage_config_engine_args(): + """Ensure default diffusion-stage builder sets and propagates engine_args.""" + stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg( + { + "distributed_executor_backend": "ray", + "boundary_ratio": 0.875, + "flow_shift": 5.0, + "trust_remote_code": True, + } + )[0] + + engine_args = stage_cfg["engine_args"] + assert engine_args["distributed_executor_backend"] == "ray" + assert engine_args["boundary_ratio"] == 0.875 + assert engine_args["flow_shift"] == 5.0 + assert engine_args["trust_remote_code"] is True + + +def test_default_stage_config_whitelist_none_fallback(): + """DeployConfig / StageDeployConfig whitelist fields with value None + fall back to OmniDiffusionConfig dataclass defaults.""" + stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg( + { + # DeployConfig pipeline-wide + "trust_remote_code": None, + "distributed_executor_backend": None, + "dtype": None, + # StageDeployConfig + "enforce_eager": None, + } + )[0] + + engine_args = stage_cfg["engine_args"] + + assert engine_args["trust_remote_code"] is False + assert engine_args["distributed_executor_backend"] == "mp" + assert engine_args["dtype"] == "auto" + assert engine_args["enforce_eager"] is False + + def test_serve_cli_accepts_ulysses_mode(): """Ensure diffusion serve CLI exposes ulysses_mode and wires it to parallel_config.""" parser = FlexibleArgumentParser() diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index bd98f9beb..130ed014a 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -1455,6 +1455,10 @@ class AsyncOmniEngine: "diffusion_load_format": kwargs.get("diffusion_load_format", "default"), "custom_pipeline_args": kwargs.get("custom_pipeline_args", None), "worker_extension_cls": kwargs.get("worker_extension_cls", None), + "trust_remote_code": (False if kwargs.get("trust_remote_code") is None else kwargs["trust_remote_code"]), + "distributed_executor_backend": ( + "mp" if kwargs.get("distributed_executor_backend") is None else kwargs["distributed_executor_backend"] + ), "enable_sleep_mode": kwargs.get("enable_sleep_mode", False), "enable_multithread_weight_load": kwargs.get("enable_multithread_weight_load", True), "num_weight_load_threads": kwargs.get("num_weight_load_threads", 4),