From fcb4ca25fe3bfd7ea4d64cbf72da7b85087e2129 Mon Sep 17 00:00:00 2001 From: "Huang, Zeyu" <11222265+fhfuih@users.noreply.github.com> Date: Tue, 20 Jan 2026 03:13:37 +0000 Subject: [PATCH] [Bugfix] Diffusion model fails to load when stage config is present Signed-off-by: Huang, Zeyu <11222265+fhfuih@users.noreply.github.com> --- vllm_omni/entrypoints/utils.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/vllm_omni/entrypoints/utils.py b/vllm_omni/entrypoints/utils.py index eae3ea7afc4..78f20bc3ce9 100644 --- a/vllm_omni/entrypoints/utils.py +++ b/vllm_omni/entrypoints/utils.py @@ -195,7 +195,8 @@ def load_stage_configs_from_yaml(config_path: str, base_engine_args: dict | None # Update base_engine_args with stage-specific engine_args if they exist if hasattr(stage_arg, "engine_args") and stage_arg.engine_args is not None: base_engine_args_tmp = OmegaConf.merge(base_engine_args_tmp, stage_arg.engine_args) - if hasattr(stage_arg, "runtime") and stage_arg.runtime is not None: + stage_type = getattr(stage_arg, "stage_type", "llm") + if hasattr(stage_arg, "runtime") and stage_arg.runtime is not None and stage_type != "diffusion": runtime_cfg = stage_arg.runtime max_batch_size = int(runtime_cfg.get("max_batch_size", 1) or 1) base_engine_args_tmp["max_num_seqs"] = max_batch_size