[Refactor 1/N] Centralize Omni config resolution (#5140)

Signed-off-by: Chenguang ZHENG <645327136@qq.com>

docs/design/architecture_overview.md

@@ -323,18 +323,24 @@ inputs, fields, and ownership rules are described below.
 ```mermaid
 flowchart TB
     layer1["Layer 1 · Authoring inputs<br/>PipelineConfig + DeployConfig"]
-    layer2["Layer 2 · Resolve once<br/>StageConfigFactory.create_from_model()<br/>VllmOmniConfig.from_pipeline_config()"]
-    layer3["Layer 3 · Transport-safe control plane<br/>VllmOmniConfig"]
+    layer2["Layer 2 · Production resolution boundary<br/>resolve_omni_config()<br/>(StageConfigFactory + VllmOmniConfig internally)"]
+    layer3["Layer 3 · Startup hand-off<br/>OmniConfigResolution<br/>PipelineConfig + temporary OmegaConf stage bridge"]
     layer4["Layer 4 · Runtime launch planning<br/>StageRuntime"]
     layer5["Layer 5 · Engine materialization<br/>VllmConfig / OmniDiffusionConfig"]
 
     layer1 --> layer2 --> layer3 --> layer4 --> layer5
 ```
 
-The single resolution boundary is implemented by
-`StageConfigFactory.create_from_model()` and
-`VllmOmniConfig.from_pipeline_config()` in
+The production resolution boundary is `resolve_omni_config()` in
 [`vllm_omni/config`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/config).
+It delegates typed construction to `StageConfigFactory.create_from_model()` and
+`VllmOmniConfig.from_pipeline_config()`, then returns an
+`OmniConfigResolution` consumed by both `AsyncOmniEngine` and headless startup.
+Until `StageRuntime` consumes typed stage configs directly, this envelope carries
+the effective `PipelineConfig` alongside OmegaConf-compatible `stage_configs` as
+a temporary runtime bridge; both views describe the same resolved topology,
+including injected stages.
+
 The legacy `stage_args` YAML path has been removed. Model topology now resolves
 through `PipelineConfig`, with runtime overrides supplied by `DeployConfig`.
 
@@ -354,9 +360,12 @@ The important ownership rules are:
    deploy-time defaults; it does not redefine the model graph.
 3. CLI and Python overrides are applied at the resolution boundary, with
    per-stage overrides taking precedence over global values where supported.
-4. `StageRuntime` owns launch planning and replica lifecycle. `ReplicaInitPlan`
+4. `OmniConfigResolution` is the sole startup hand-off. Its `pipeline_config`
+   and temporary `stage_configs` compatibility view must describe the same
+   topology.
+5. `StageRuntime` owns launch planning and replica lifecycle. `ReplicaInitPlan`
    is runtime-private state, not a user configuration object.
-5. `VllmConfig` and the enriched `OmniDiffusionConfig` are materialized in the
+6. `VllmConfig` and the enriched `OmniDiffusionConfig` are materialized in the
    process that owns the corresponding engine.
 
 ## Main features

docs/mkdocs/hooks/generate_argparse.py

@@ -124,6 +124,7 @@ def extract_omni_serve_subparser_init():
                         "make_arg_parser": lambda parser: parser,  # no-op for doc
                         "_ensure_vllm_platform": lambda: None,  # no-op for doc
                         "_nonneg_finite_float": float,
+                        "_parse_stage_overrides": __import__("json").loads,
                         "VLLM_SUBCMD_PARSER_EPILOG": "",
                         "logger": logger,
                         "DummySubparsers": DummySubparsers,

tests/config/test_config_factory.py

@@ -6,7 +6,6 @@ Unit tests for StageConfigFactory and related classes.
 
 import importlib
 import warnings
-from dataclasses import dataclass
 from pathlib import Path
 from unittest.mock import patch
 
@@ -37,6 +36,7 @@ from vllm_omni.config.stage_config import (
     normalize_pipeline_cli_overrides,
     pipeline_cfg_resolver,
 )
+from vllm_omni.diffusion.data import DiffusionParallelConfig
 from vllm_omni.engine.arg_utils import SHARED_FIELDS, internal_blacklist_keys
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -451,13 +451,8 @@ class TestStageConfigFactory:
 
     def test_default_diffusion_with_parallel_config(self):
         """Test diffusion config calculates devices from parallel_config."""
-
-        @dataclass
-        class MockParallelConfig:
-            world_size: int = 4
-
         kwargs = {
-            "parallel_config": MockParallelConfig(),
+            "parallel_config": DiffusionParallelConfig(tensor_parallel_size=4),
             "cache_backend": "tea_cache",
         }
         configs = StageConfigFactory.create_default_diffusion(kwargs)

tests/config/test_omni_config.py

@@ -1253,6 +1253,70 @@ def test_diffusion_config_from_kwargs_reuses_legacy_normalization(monkeypatch):
     assert cfg.diffusers_call_kwargs == {}
 
 
+def test_diffusion_config_none_values_preserve_dataclass_defaults():
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+    normalized = OmniDiffusionConfig.normalize_init_kwargs(
+        {
+            "lora_scale": None,
+            "enable_sleep_mode": None,
+            "diffusers_load_kwargs": None,
+        }
+    )
+
+    assert "lora_scale" not in normalized
+    assert "enable_sleep_mode" not in normalized
+    assert normalized["diffusers_load_kwargs"] == {}
+
+    config = OmniDiffusionConfig.from_kwargs(
+        lora_scale=None,
+        enable_sleep_mode=None,
+    )
+    assert config.lora_scale == 1.0
+    assert config.enable_sleep_mode is False
+
+
+@pytest.mark.parametrize(
+    ("canonical_key", "alias_key", "canonical_value", "alias_value"),
+    [
+        ("lora_scale", "static_lora_scale", 0.75, 0.25),
+        (
+            "quantization_config",
+            "diffusion_quantization_config",
+            {"method": "canonical"},
+            {"method": "diffusion-alias"},
+        ),
+        (
+            "quantization_config",
+            "quantization",
+            {"method": "canonical"},
+            "legacy-alias",
+        ),
+        ("diffusion_kv_cache_dtype", "kv_cache_dtype", "fp8", "fp16"),
+        ("diffusion_kv_cache_skip_steps", "kv_cache_skip_steps", "0-1", "2-3"),
+        ("diffusion_kv_cache_skip_layers", "kv_cache_skip_layers", "1-2", "3-4"),
+        ("streaming_output", "diffusion_streaming_output", False, True),
+    ],
+)
+def test_diffusion_alias_conflicts_prefer_canonical_key(
+    canonical_key,
+    alias_key,
+    canonical_value,
+    alias_value,
+):
+    from vllm_omni.diffusion.data import normalize_omni_diffusion_kwargs
+
+    normalized = normalize_omni_diffusion_kwargs(
+        {
+            canonical_key: canonical_value,
+            alias_key: alias_value,
+        }
+    )
+
+    assert normalized[canonical_key] == canonical_value
+    assert alias_key not in normalized
+
+
 def test_from_pipeline_config_normalizes_diffusion_config_aliases_from_engine_args(tmp_path, monkeypatch):
     from vllm_omni.platforms import current_omni_platform
 

tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py

@@ -177,9 +177,9 @@ def _write_partition_index(path, *, partition, tasks):
 
 
 def test_modular_diffusers_index_is_resolved_generically(tmp_path):
+    from vllm_omni.config.resolver import resolve_omni_config
     from vllm_omni.diffusion.data import OmniDiffusionConfig, resolve_model_class_name
     from vllm_omni.diffusion.utils.hf_utils import is_diffusion_model
-    from vllm_omni.entrypoints.utils import resolve_model_config_path
 
     (tmp_path / "modular_model_index.json").write_text(
         json.dumps(
@@ -193,7 +193,16 @@ def test_modular_diffusers_index_is_resolved_generically(tmp_path):
 
     assert is_diffusion_model(str(tmp_path))
     assert resolve_model_class_name(str(tmp_path)) == "MiniMaxH3ModularPipeline"
-    assert resolve_model_config_path(str(tmp_path)) is None
+    resolved = resolve_omni_config(
+        str(tmp_path),
+        trust_remote_code=False,
+        deploy_config_path=None,
+        cli_overrides=None,
+        stage_overrides=None,
+        strategy_config_path=None,
+    )
+    assert resolved.config_path is None
+    assert resolved.stage_by_id(0).engine_args.model_class_name == "MiniMaxH3ModularPipeline"
 
     config = OmniDiffusionConfig(model=str(tmp_path))
     config.enrich_config()

tests/diffusion/test_diffusion_config_propagation.py

@@ -131,7 +131,7 @@ class TestCreateDefaultDiffusion:
 
     def test_dtype_serialized_as_string(self):
         stages = StageConfigFactory.create_default_diffusion({"dtype": torch.float16, "model": "x"})
-        assert stages[0]["engine_args"]["dtype"] == "torch.float16"
+        assert stages[0]["engine_args"]["dtype"] == "float16"
 
     def test_cache_backend_defaults_to_none(self):
         stages = StageConfigFactory.create_default_diffusion({"model": "x"})

tests/diffusion/test_diffusion_step_pipeline.py

@@ -15,6 +15,7 @@ from pytest_mock import MockerFixture
 
 import vllm_omni.diffusion.worker.diffusion_model_runner as model_runner_module
 from tests.helpers.mark import hardware_test
+from vllm_omni.config.config_factory import StageConfigFactory
 from vllm_omni.diffusion.data import DiffusionOutput
 from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
 from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
@@ -46,7 +47,6 @@ from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunn
 from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker
 from vllm_omni.diffusion.worker.input_batch import InputBatch
 from vllm_omni.diffusion.worker.utils import RunnerOutput, StepRequestState
-from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
 from vllm_omni.errors import OmniClientError
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.platforms import current_omni_platform
@@ -1182,7 +1182,7 @@ class TestSupportedPipelines:
     """Step-execution protocol checks for supported pipelines."""
 
     def test_default_stage_config_includes_step_execution(self):
-        stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
+        stage_cfg = StageConfigFactory.create_default_diffusion(
             {
                 "step_execution": True,
             }

tests/engine/test_async_omni_engine_stage_init.py

@@ -11,7 +11,7 @@ import types
 
 import pytest
 
-from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, normalize_omni_diffusion_kwargs
+from vllm_omni.diffusion.data import AttentionConfig
 from vllm_omni.engine import async_omni_engine as async_omni_engine_module
 from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
 from vllm_omni.engine.stage_init_utils import (
@@ -1732,133 +1732,6 @@ def test_inject_kv_stage_info_infers_receiver_tp_topology():
     assert stage1.engine_args["omni_kv_config"]["rank_mapping"] == {"from_tp": 4, "to_tp": 2}
 
 
-def test_resolve_stage_configs_injects_global_diffusion_attention_when_missing(monkeypatch):
-    import vllm_omni.engine.async_omni_engine as engine_mod
-
-    engine = object.__new__(AsyncOmniEngine)
-    stage_cfg = types.SimpleNamespace(
-        stage_type="diffusion",
-        engine_args=types.SimpleNamespace(
-            diffusion_attention_config=None,
-            lora_path=None,
-            lora_scale=None,
-            enable_sleep_mode=None,
-            quantization_config=None,
-        ),
-    )
-
-    monkeypatch.setattr(
-        engine_mod,
-        "load_and_resolve_stage_configs",
-        lambda *args, **kwargs: ("dummy-config", [stage_cfg], None),
-    )
-
-    _config_path, stage_configs = engine._resolve_stage_configs(
-        model="dummy-model",
-        kwargs={"diffusion_attention_backend": "FLASH_ATTN"},
-        trust_remote_code=False,
-    )
-
-    diffusion_attention_config = stage_configs[0].engine_args.diffusion_attention_config
-    assert isinstance(diffusion_attention_config, AttentionConfig)
-    assert diffusion_attention_config.default is not None
-    assert diffusion_attention_config.default.backend == "FLASH_ATTN"
-
-
-def test_resolve_stage_configs_preserves_stage_diffusion_attention(monkeypatch):
-    import vllm_omni.engine.async_omni_engine as engine_mod
-
-    engine = object.__new__(AsyncOmniEngine)
-    existing_attention = AttentionConfig(default=AttentionSpec(backend="TORCH_SDPA"))
-    stage_cfg = types.SimpleNamespace(
-        stage_type="diffusion",
-        engine_args=types.SimpleNamespace(
-            diffusion_attention_config=existing_attention,
-            lora_path=None,
-            lora_scale=None,
-            enable_sleep_mode=None,
-            quantization_config=None,
-        ),
-    )
-
-    monkeypatch.setattr(
-        engine_mod,
-        "load_and_resolve_stage_configs",
-        lambda *args, **kwargs: ("dummy-config", [stage_cfg], None),
-    )
-
-    _config_path, stage_configs = engine._resolve_stage_configs(
-        model="dummy-model",
-        kwargs={"diffusion_attention_backend": "FLASH_ATTN"},
-        trust_remote_code=False,
-    )
-
-    assert stage_configs[0].engine_args.diffusion_attention_config is existing_attention
-
-
-def test_resolve_stage_configs_does_not_inject_over_stage_diffusion_attention_backend(monkeypatch):
-    import vllm_omni.engine.async_omni_engine as engine_mod
-
-    engine = object.__new__(AsyncOmniEngine)
-    stage_cfg = types.SimpleNamespace(
-        stage_type="diffusion",
-        engine_args=types.SimpleNamespace(
-            diffusion_attention_backend="TORCH_SDPA",
-            diffusion_attention_config=None,
-            lora_path=None,
-            lora_scale=None,
-            enable_sleep_mode=None,
-            quantization_config=None,
-        ),
-    )
-
-    monkeypatch.setattr(
-        engine_mod,
-        "load_and_resolve_stage_configs",
-        lambda *args, **kwargs: ("dummy-config", [stage_cfg], None),
-    )
-
-    _config_path, stage_configs = engine._resolve_stage_configs(
-        model="dummy-model",
-        kwargs={"diffusion_attention_backend": "FLASH_ATTN"},
-        trust_remote_code=False,
-    )
-
-    engine_args = stage_configs[0].engine_args
-    assert engine_args.diffusion_attention_backend == "TORCH_SDPA"
-    assert engine_args.diffusion_attention_config is None
-    normalized = normalize_omni_diffusion_kwargs(vars(engine_args))
-    assert normalized["diffusion_attention_config"].default.backend == "TORCH_SDPA"
-
-
-def test_resolve_stage_configs_does_not_inject_diffusion_attention_into_llm_stage(monkeypatch):
-    import vllm_omni.engine.async_omni_engine as engine_mod
-
-    engine = object.__new__(AsyncOmniEngine)
-    stage_cfg = types.SimpleNamespace(
-        stage_type="llm",
-        engine_args=types.SimpleNamespace(
-            attention_config={"backend": "FLASH_ATTN"},
-            enable_sleep_mode=None,
-        ),
-    )
-
-    monkeypatch.setattr(
-        engine_mod,
-        "load_and_resolve_stage_configs",
-        lambda *args, **kwargs: ("dummy-config", [stage_cfg], None),
-    )
-
-    _config_path, stage_configs = engine._resolve_stage_configs(
-        model="dummy-model",
-        kwargs={"diffusion_attention_backend": "TORCH_SDPA"},
-        trust_remote_code=False,
-    )
-
-    assert stage_configs[0].engine_args.attention_config == {"backend": "FLASH_ATTN"}
-    assert not hasattr(stage_configs[0].engine_args, "diffusion_attention_config")
-
-
 def test_extract_legacy_stage_metadata_rocm_does_not_inject_diffusion_attention(monkeypatch):
     """ROCm default attention logic only applies to LLM stages, not diffusion."""
     from vllm_omni.engine.stage_init_utils import extract_legacy_stage_metadata

tests/entrypoints/test_async_omni_diffusion_config.py

@@ -4,8 +4,12 @@
 from types import SimpleNamespace
 
 import pytest
+import torch
+from pydantic import ValidationError
 
-from vllm_omni.diffusion.data import AttentionConfig
+from vllm_omni.config.config_factory import StageConfigFactory
+from vllm_omni.config.resolver import OmniConfigResolution
+from vllm_omni.diffusion.data import AttentionConfig, OmniDiffusionConfig
 from vllm_omni.engine import stage_init_utils
 from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
 from vllm_omni.entrypoints.cli.serve import OmniServeCommand
@@ -14,9 +18,13 @@ from vllm_omni.utils.tracking_parser import TrackingArgumentParser
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
+def _terminal_config(stage_cfg: dict) -> OmniDiffusionConfig:
+    return OmniDiffusionConfig.from_kwargs(**stage_cfg["engine_args"])
+
+
 def test_default_stage_config_includes_cache_backend():
     """Ensure cache knobs survive the default diffusion-stage builder."""
-    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
+    stage_cfg = StageConfigFactory.create_default_diffusion(
         {
             "cache_backend": "cache_dit",
             "cache_config": '{"Fn_compute_blocks": 2}',
@@ -30,12 +38,12 @@ def test_default_stage_config_includes_cache_backend():
     assert engine_args["cache_backend"] == "cache_dit"
     assert engine_args["cache_config"]["Fn_compute_blocks"] == 2
     assert engine_args["vae_use_slicing"] is True
-    assert engine_args["parallel_config"].ulysses_degree == 2
+    assert engine_args["parallel_config"]["ulysses_degree"] == 2
     assert engine_args["model_stage"] == "diffusion"
 
 
 def test_default_stage_config_preserves_ulysses_a2a_permute() -> None:
-    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
+    stage_cfg = StageConfigFactory.create_default_diffusion(
         {
             "ulysses_degree": 4,
             "ulysses_a2a_permute": True,
@@ -43,18 +51,18 @@ def test_default_stage_config_preserves_ulysses_a2a_permute() -> None:
     )[0]
 
     parallel_config = stage_cfg["engine_args"]["parallel_config"]
-    assert parallel_config.ulysses_degree == 4
-    assert parallel_config.ulysses_a2a_permute is True
+    assert parallel_config["ulysses_degree"] == 4
+    assert parallel_config["ulysses_a2a_permute"] is True
 
 
 def test_default_stage_config_preserves_model_extras():
-    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg({"extras": {"ltx2_use_conv_vae": True}})[0]
+    stage_cfg = StageConfigFactory.create_default_diffusion({"extras": {"ltx2_use_conv_vae": True}})[0]
 
     assert stage_cfg["engine_args"]["extras"]["ltx2_use_conv_vae"] is True
 
 
 def test_default_stage_config_preserves_and_overrides_promoted_extras():
-    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
+    stage_cfg = StageConfigFactory.create_default_diffusion(
         {
             "extras": {
                 "auxiliary_text_encoder": "/models/extras-llama",
@@ -68,7 +76,7 @@ def test_default_stage_config_preserves_and_overrides_promoted_extras():
     assert extras["auxiliary_text_encoder"] == "/models/extras-llama"
     assert extras["default_llama_model_id"] == "extras/default-llama"
 
-    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
+    stage_cfg = StageConfigFactory.create_default_diffusion(
         {
             "extras": {
                 "auxiliary_text_encoder": "/models/extras-llama",
@@ -84,43 +92,57 @@ def test_default_stage_config_preserves_and_overrides_promoted_extras():
     assert extras["default_llama_model_id"] == "top-level/default-llama"
 
 
-def test_stage_override_preserves_model_extras_for_default_diffusion_stage(mocker):
+@pytest.mark.parametrize(
+    "stage_overrides",
+    [
+        {"0": {"extras": {"ltx2_use_conv_vae": True}}},
+        '{"0":{"extras":{"ltx2_use_conv_vae":true}}}',
+    ],
+)
+def test_stage_override_preserves_model_extras_for_default_diffusion_stage(mocker, stage_overrides):
     """Local/unregistered Diffusers checkpoints still honor stage-0 extras."""
-
-    def resolve_with_default(*_args, default_stage_cfg_factory, **_kwargs):
-        return None, default_stage_cfg_factory(), None
-
     mocker.patch(
-        "vllm_omni.engine.async_omni_engine.load_and_resolve_stage_configs",
-        side_effect=resolve_with_default,
+        "vllm_omni.config.resolver.StageConfigFactory.create_from_model",
+        return_value=None,
+    )
+    mocker.patch(
+        "vllm_omni.config.resolver._resolve_generic_diffusion_model_class",
+        return_value=(True, "LTX2Pipeline"),
     )
     engine = AsyncOmniEngine.__new__(AsyncOmniEngine)
 
     _, stage_configs = engine._resolve_stage_configs(
         "/models/LTX-2.5-Diffusers",
-        {"stage_overrides": '{"0":{"extras":{"ltx2_use_conv_vae":true}}}'},
+        {"stage_overrides": stage_overrides},
         trust_remote_code=False,
     )
 
     assert stage_configs[0]["engine_args"]["extras"]["ltx2_use_conv_vae"] is True
 
 
+def test_default_stage_rejects_unknown_nested_parallel_config_key():
+    unknown_key = "unknown_parallel_field"
+    with pytest.raises(ValidationError, match=unknown_key):
+        StageConfigFactory.create_default_diffusion(
+            {"parallel_config": {unknown_key: 2}},
+        )
+
+
 def test_default_cache_config_used_when_missing():
     """Ensure default cache_config is synthesized when only backend is given."""
-    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
+    stage_cfg = StageConfigFactory.create_default_diffusion(
         {
             "cache_backend": "cache_dit",
         }
     )[0]
 
-    cache_config = stage_cfg["engine_args"]["cache_config"]
-    assert cache_config is not None
-    assert cache_config["Fn_compute_blocks"] == 1
+    cache_config = _terminal_config(stage_cfg).cache_config
+    assert cache_config.Fn_compute_blocks == 1
 
 
 def test_default_stage_devices_from_sequence_parallel():
     """Ensure runtime devices reflect computed diffusion world size."""
-    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
+    stage_cfg = StageConfigFactory.create_default_diffusion(
         {
             "ulysses_degree": 2,
             "ring_degree": 2,
@@ -132,7 +154,7 @@ def test_default_stage_devices_from_sequence_parallel():
 
 def test_default_stage_devices_and_dp_from_num_gpus():
     """Resolve omitted DP before deriving devices from an explicit GPU count."""
-    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
+    stage_cfg = StageConfigFactory.create_default_diffusion(
         {
             "num_gpus": 8,
             "tensor_parallel_size": 2,
@@ -141,15 +163,14 @@ def test_default_stage_devices_and_dp_from_num_gpus():
     )[0]
 
     parallel_config = stage_cfg["engine_args"]["parallel_config"]
-    assert parallel_config.data_parallel_size == 2
-    assert parallel_config.world_size == 8
+    assert parallel_config["data_parallel_size"] == 2
     assert stage_cfg["engine_args"]["num_gpus"] == 8
     assert stage_cfg["runtime"]["devices"] == "0,1,2,3,4,5,6,7"
 
 
 def test_default_stage_config_uses_parallel_size_kwargs():
     """Ensure default diffusion parallel_config uses CLI/API parallel sizes."""
-    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
+    stage_cfg = StageConfigFactory.create_default_diffusion(
         {
             "pipeline_parallel_size": 2,
             "data_parallel_size": 3,
@@ -159,15 +180,15 @@ def test_default_stage_config_uses_parallel_size_kwargs():
     )[0]
 
     parallel_config = stage_cfg["engine_args"]["parallel_config"]
-    assert parallel_config.pipeline_parallel_size == 2
-    assert parallel_config.data_parallel_size == 3
-    assert parallel_config.tensor_parallel_size == 4
-    assert parallel_config.enable_expert_parallel is True
+    assert parallel_config["pipeline_parallel_size"] == 2
+    assert parallel_config["data_parallel_size"] == 3
+    assert parallel_config["tensor_parallel_size"] == 4
+    assert parallel_config["enable_expert_parallel"] is True
 
 
 def test_default_stage_config_preserves_omitted_dp_for_runtime_inference():
     """Keep omitted DP unresolved until the runtime WORLD size is known."""
-    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
+    stage_cfg = StageConfigFactory.create_default_diffusion(
         {
             "pipeline_parallel_size": None,
             "data_parallel_size": None,
@@ -180,18 +201,19 @@ def test_default_stage_config_preserves_omitted_dp_for_runtime_inference():
     )[0]
 
     parallel_config = stage_cfg["engine_args"]["parallel_config"]
-    assert parallel_config.pipeline_parallel_size == 1
-    assert parallel_config.data_parallel_size is None
-    assert parallel_config.tensor_parallel_size == 1
-    assert parallel_config.enable_expert_parallel is False
-    assert stage_cfg["engine_args"]["enforce_eager"] is False
-    assert stage_cfg["engine_args"]["diffusion_compile_granularity"] == "regional"
-    assert stage_cfg["engine_args"]["diffusion_compile_dynamic"] is True
+    assert parallel_config["pipeline_parallel_size"] == 1
+    assert parallel_config["data_parallel_size"] is None
+    assert parallel_config["tensor_parallel_size"] == 1
+    assert parallel_config["enable_expert_parallel"] is False
+    terminal_config = _terminal_config(stage_cfg)
+    assert terminal_config.enforce_eager is False
+    assert terminal_config.diffusion_compile_granularity == "regional"
+    assert terminal_config.diffusion_compile_dynamic is True
 
 
 def test_default_stage_config_propagates_ulysses_mode():
     """Ensure UAA mode survives default diffusion-stage creation."""
-    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
+    stage_cfg = StageConfigFactory.create_default_diffusion(
         {
             "ulysses_degree": 4,
             "ulysses_mode": "advanced_uaa",
@@ -199,13 +221,13 @@ def test_default_stage_config_propagates_ulysses_mode():
     )[0]
 
     parallel_config = stage_cfg["engine_args"]["parallel_config"]
-    assert parallel_config.ulysses_degree == 4
-    assert parallel_config.ulysses_mode == "advanced_uaa"
+    assert parallel_config["ulysses_degree"] == 4
+    assert parallel_config["ulysses_mode"] == "advanced_uaa"
 
 
 def test_default_stage_config_includes_default_sampling_params():
     """Ensure default sampling params survive the default diffusion-stage builder."""
-    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
+    stage_cfg = StageConfigFactory.create_default_diffusion(
         {
             "default_sampling_params": '{"0": {"generator_device":"cpu", "guidance_scale":7.5}}',
         }
@@ -219,13 +241,13 @@ def test_default_stage_config_includes_default_sampling_params():
 
 def test_default_stage_config_includes_diffusion_attention_backend():
     """Ensure diffusion attention shorthand lands in engine_args.diffusion_attention_config."""
-    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
+    stage_cfg = StageConfigFactory.create_default_diffusion(
         {
             "diffusion_attention_backend": "FLASH_ATTN",
         }
     )[0]
 
-    diffusion_attention_config = stage_cfg["engine_args"]["diffusion_attention_config"]
+    diffusion_attention_config = _terminal_config(stage_cfg).diffusion_attention_config
     assert isinstance(diffusion_attention_config, AttentionConfig)
     assert diffusion_attention_config.default is not None
     assert diffusion_attention_config.default.backend == "FLASH_ATTN"
@@ -233,7 +255,7 @@ def test_default_stage_config_includes_diffusion_attention_backend():
 
 def test_default_stage_config_includes_diffusion_attention_config():
     """Ensure structured diffusion attention config survives default stage creation."""
-    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
+    stage_cfg = StageConfigFactory.create_default_diffusion(
         {
             "diffusion_attention_config": {
                 "default": {"backend": "FLASH_ATTN"},
@@ -242,7 +264,7 @@ def test_default_stage_config_includes_diffusion_attention_config():
         }
     )[0]
 
-    diffusion_attention_config = stage_cfg["engine_args"]["diffusion_attention_config"]
+    diffusion_attention_config = _terminal_config(stage_cfg).diffusion_attention_config
     assert isinstance(diffusion_attention_config, AttentionConfig)
     assert diffusion_attention_config.default is not None
     assert diffusion_attention_config.default.backend == "FLASH_ATTN"
@@ -252,7 +274,7 @@ def test_default_stage_config_includes_diffusion_attention_config():
 def test_default_stage_config_rejects_conflicting_diffusion_attention_inputs():
     """Ensure shorthand and default.backend stay mutually exclusive."""
     with pytest.raises(ValueError, match="mutually exclusive"):
-        AsyncOmniEngine._create_default_diffusion_stage_cfg(
+        StageConfigFactory.create_default_diffusion(
             {
                 "diffusion_attention_backend": "FLASH_ATTN",
                 "diffusion_attention_config": {
@@ -264,7 +286,7 @@ def test_default_stage_config_rejects_conflicting_diffusion_attention_inputs():
 
 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(
+    stage_cfg = StageConfigFactory.create_default_diffusion(
         {
             "distributed_executor_backend": "ray",
             "boundary_ratio": 0.875,
@@ -283,7 +305,7 @@ def test_default_stage_config_engine_args():
 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(
+    stage_cfg = StageConfigFactory.create_default_diffusion(
         {
             # DeployConfig pipeline-wide
             "trust_remote_code": None,
@@ -294,12 +316,11 @@ def test_default_stage_config_whitelist_none_fallback():
         }
     )[0]
 
-    engine_args = stage_cfg["engine_args"]
-
-    assert engine_args["trust_remote_code"] is False
-    assert engine_args["distributed_executor_backend"] is None
-    assert engine_args["dtype"] == "auto"
-    assert engine_args["enforce_eager"] is False
+    terminal = _terminal_config(stage_cfg)
+    assert terminal.trust_remote_code is False
+    assert terminal.distributed_executor_backend is None
+    assert terminal.dtype == torch.bfloat16
+    assert terminal.enforce_eager is False
 
 
 def test_serve_cli_accepts_ulysses_mode():
@@ -321,12 +342,12 @@ def test_serve_cli_accepts_ulysses_mode():
     )
 
     explicit_kwargs = args.get_explicit_kwargs_dict()
-    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0]
+    stage_cfg = StageConfigFactory.create_default_diffusion(explicit_kwargs)[0]
     parallel_config = stage_cfg["engine_args"]["parallel_config"]
 
     assert args.ulysses_mode == "advanced_uaa"
-    assert parallel_config.ulysses_degree == 4
-    assert parallel_config.ulysses_mode == "advanced_uaa"
+    assert parallel_config["ulysses_degree"] == 4
+    assert parallel_config["ulysses_mode"] == "advanced_uaa"
 
 
 def test_serve_cli_accepts_text_encoder_tp_size():
@@ -345,11 +366,11 @@ def test_serve_cli_accepts_text_encoder_tp_size():
     )
 
     explicit_kwargs = args.get_explicit_kwargs_dict()
-    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0]
+    stage_cfg = StageConfigFactory.create_default_diffusion(explicit_kwargs)[0]
     parallel_config = stage_cfg["engine_args"]["parallel_config"]
 
     assert args.text_encoder_tp_size == 4
-    assert parallel_config.text_encoder_tp_size == 4
+    assert parallel_config["text_encoder_tp_size"] == 4
 
 
 def test_serve_cli_forwards_model_defined_task_type_to_diffusion_stage():
@@ -368,7 +389,7 @@ def test_serve_cli_forwards_model_defined_task_type_to_diffusion_stage():
     )
 
     explicit_kwargs = args.get_explicit_kwargs_dict()
-    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0]
+    stage_cfg = StageConfigFactory.create_default_diffusion(explicit_kwargs)[0]
 
     assert args.task_type == "fl2va"
     assert stage_cfg["engine_args"]["task_type"] == "fl2va"
@@ -390,7 +411,7 @@ def test_serve_cli_accepts_diffusion_pipeline_profiler_flag():
     )
 
     explicit_kwargs = args.get_explicit_kwargs_dict()
-    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0]
+    stage_cfg = StageConfigFactory.create_default_diffusion(explicit_kwargs)[0]
 
     assert args.enable_diffusion_pipeline_profiler is True
     assert stage_cfg["engine_args"]["enable_diffusion_pipeline_profiler"] is True
@@ -416,7 +437,7 @@ def test_serve_cli_forwards_distilled_lora_to_diffusion_stage():
     )
 
     explicit_kwargs = args.get_explicit_kwargs_dict()
-    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0]
+    stage_cfg = StageConfigFactory.create_default_diffusion(explicit_kwargs)[0]
     engine_args = stage_cfg["engine_args"]
 
     assert explicit_kwargs["lora_backend"] == "distill"
@@ -447,7 +468,7 @@ def test_serve_cli_forwards_compact_diffusion_offload_config():
     )
 
     explicit_kwargs = args.get_explicit_kwargs_dict()
-    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0]
+    stage_cfg = StageConfigFactory.create_default_diffusion(explicit_kwargs)[0]
     engine_args = stage_cfg["engine_args"]
 
     expected = {
@@ -513,7 +534,7 @@ def test_serve_cli_forwards_hwr_policy_for_no_allgather_dlo():
     )
 
     explicit_kwargs = args.get_explicit_kwargs_dict()
-    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0]
+    stage_cfg = StageConfigFactory.create_default_diffusion(explicit_kwargs)[0]
     engine_args = stage_cfg["engine_args"]
 
     assert explicit_kwargs["host_weight_runtime_mode"] == "preferred"
@@ -542,7 +563,7 @@ def test_serve_cli_accepts_diffusion_compile_controls():
     )
 
     explicit_kwargs = args.get_explicit_kwargs_dict()
-    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0]
+    stage_cfg = StageConfigFactory.create_default_diffusion(explicit_kwargs)[0]
 
     assert args.diffusion_compile_granularity == "full"
     assert args.diffusion_compile_dynamic is False
@@ -569,7 +590,7 @@ def test_serve_cli_accepts_diffusion_attention_backend():
     )
 
     explicit_kwargs = args.get_explicit_kwargs_dict()
-    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0]
+    stage_cfg = StageConfigFactory.create_default_diffusion(explicit_kwargs)[0]
     diffusion_attention_config = stage_cfg["engine_args"]["diffusion_attention_config"]
 
     assert args.diffusion_attention_backend == "FASTVIDEO_VSA"
@@ -597,7 +618,7 @@ def test_serve_cli_accepts_request_batch_max_wait_ms():
     )
 
     explicit_kwargs = args.get_explicit_kwargs_dict()
-    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0]
+    stage_cfg = StageConfigFactory.create_default_diffusion(explicit_kwargs)[0]
 
     assert args.request_batch_max_wait_ms == 250.0
     assert stage_cfg["engine_args"]["request_batch_max_wait_ms"] == 250.0
@@ -637,7 +658,7 @@ def test_serve_cli_accepts_additional_config():
         ]
     )
 
-    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(vars(args))[0]
+    stage_cfg = StageConfigFactory.create_default_diffusion(vars(args))[0]
 
     engine_args = stage_cfg["engine_args"]
 
@@ -645,79 +666,19 @@ def test_serve_cli_accepts_additional_config():
     assert engine_args["additional_config"] == {"torchair_graph_config": {"enabled": True}}
 
 
-def test_default_stage_resolves_video_output_from_checkpoint(mocker):
-    captured = {}
-
-    def resolve_with_default(*args, default_stage_cfg_factory, **kwargs):
-        del args, kwargs
-        captured.update(default_stage_cfg_factory()[0])
-        stage = SimpleNamespace(stage_type="diffusion", engine_args=SimpleNamespace())
-        return ("", [stage], None)
-
-    resolver = mocker.patch(
-        "vllm_omni.engine.async_omni_engine.resolve_model_class_name",
-        return_value="MiniMaxH3Pipeline",
-    )
-    mocker.patch(
-        "vllm_omni.engine.async_omni_engine.load_and_resolve_stage_configs",
-        side_effect=resolve_with_default,
-    )
-    engine = AsyncOmniEngine.__new__(AsyncOmniEngine)
-    engine._strip_single_engine_args = lambda kwargs: kwargs
-
-    engine._resolve_stage_configs(
-        "/models/MiniMax-H3/FL2VA",
-        {},
-        trust_remote_code=False,
-    )
-
-    resolver.assert_called_once_with("/models/MiniMax-H3/FL2VA", "default")
-    assert captured["final_output_type"] == "video"
-
-
-def test_default_diffusers_stage_preserves_video_model_identity(mocker):
-    captured = {}
-
-    def resolve_with_default(*args, default_stage_cfg_factory, **kwargs):
-        del args, kwargs
-        captured.update(default_stage_cfg_factory()[0])
-        stage = SimpleNamespace(stage_type="diffusion", engine_args=SimpleNamespace())
-        return ("", [stage], None)
-
-    mocker.patch(
-        "vllm_omni.diffusion.utils.hf_utils.get_diffusion_model_index",
-        return_value={"_class_name": "WanImageToVideoPipeline"},
-    )
-    mocker.patch(
-        "vllm_omni.engine.async_omni_engine.load_and_resolve_stage_configs",
-        side_effect=resolve_with_default,
-    )
-    engine = AsyncOmniEngine.__new__(AsyncOmniEngine)
-    engine._strip_single_engine_args = lambda kwargs: kwargs
-
-    engine._resolve_stage_configs(
-        "/models/Wan2.2-I2V",
-        {"diffusion_load_format": "diffusers"},
-        trust_remote_code=False,
-    )
-
-    assert captured["engine_args"]["model_class_name"] == "WanImageToVideoPipeline"
-    assert captured["final_output_type"] == "video"
-
-
-def test_resolve_stage_configs_injects_additional_config_into_diffusion_stage(mocker):
-    """Ensure YAML/deploy stage resolution forwards top-level additional_config."""
+def test_resolve_stage_configs_delegates_overrides_to_resolver(mocker):
+    """The engine consumes resolver output without a second merge pass."""
+    additional_config = {"torchair_graph_config": {"enabled": True}}
     fake_diffusion_stage = SimpleNamespace(
         stage_type="diffusion",
-        engine_args=SimpleNamespace(),
+        engine_args=SimpleNamespace(additional_config=additional_config),
     )
-    fake_llm_stage = SimpleNamespace(
-        stage_type="llm",
-        engine_args=SimpleNamespace(),
-    )
-    mocker.patch(
-        "vllm_omni.engine.async_omni_engine.load_and_resolve_stage_configs",
-        return_value=("dummy.yaml", [fake_llm_stage, fake_diffusion_stage], None),
+    resolve_config = mocker.patch(
+        "vllm_omni.engine.async_omni_engine.resolve_omni_config",
+        return_value=OmniConfigResolution(
+            config_path="dummy.yaml",
+            stage_configs=(fake_diffusion_stage,),
+        ),
     )
 
     engine = AsyncOmniEngine.__new__(AsyncOmniEngine)
@@ -726,13 +687,15 @@ def test_resolve_stage_configs_injects_additional_config_into_diffusion_stage(mo
         "dummy-model",
         {
             "deploy_config": "dummy.yaml",
-            "additional_config": {"torchair_graph_config": {"enabled": True}},
+            "additional_config": additional_config,
         },
         trust_remote_code=False,
     )
 
-    assert not hasattr(stage_configs[0].engine_args, "additional_config")
-    assert stage_configs[1].engine_args.additional_config == {"torchair_graph_config": {"enabled": True}}
+    assert stage_configs == [fake_diffusion_stage]
+    assert resolve_config.call_args.args == ("dummy-model",)
+    assert resolve_config.call_args.kwargs["deploy_config_path"] == "dummy.yaml"
+    assert resolve_config.call_args.kwargs["cli_overrides"]["additional_config"] is additional_config
 
 
 @pytest.mark.parametrize(
@@ -760,30 +723,6 @@ def test_default_stage_config_includes_quantization_config():
         "weights": "weights.bin",
     }
 
-    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg({"quantization_config": quantization_config})[0]
+    stage_cfg = StageConfigFactory.create_default_diffusion({"quantization_config": quantization_config})[0]
 
     assert stage_cfg["engine_args"]["quantization_config"] == quantization_config
-
-
-def test_resolve_stage_configs_injects_quantization_config_into_diffusion_stage(mocker):
-    fake_diffusion_stage = SimpleNamespace(
-        stage_type="diffusion",
-        engine_args=SimpleNamespace(quantization_config=None),
-    )
-    mocker.patch(
-        "vllm_omni.engine.async_omni_engine.load_and_resolve_stage_configs",
-        return_value=("dummy.yaml", [fake_diffusion_stage], None),
-    )
-
-    engine = AsyncOmniEngine.__new__(AsyncOmniEngine)
-
-    _, stage_configs = engine._resolve_stage_configs(
-        "dummy-model",
-        {
-            "deploy_config": "dummy.yaml",
-            "quantization_config": {"method": "bitsandbytes"},
-        },
-        trust_remote_code=False,
-    )
-
-    assert stage_configs[0].engine_args.quantization_config == {"method": "bitsandbytes"}

tests/entrypoints/test_pd_disaggregation.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 """Unit tests for PD (Prefill-Decode) disaggregation in the Omni orchestrator.
 
 Tests the PD detection, validation, config parsing, sampling param
@@ -74,7 +77,7 @@ class _FakeStageConfig:
 
 class _FakeQueue:
     def __init__(self, maxsize=0):
-        self._queue = Queue(maxsize=maxsize)
+        self._queue: Queue[Any] = Queue(maxsize=maxsize)
 
     def put(self, item):
         self._queue.put(item)
@@ -343,7 +346,7 @@ def mock_get_config(monkeypatch):
     monkeypatch.setattr(
         "vllm.transformers_utils.config.get_config", lambda model, **kwargs: fake_hf_config, raising=False
     )
-    monkeypatch.setattr("vllm_omni.entrypoints.utils.get_config", lambda model, **kwargs: fake_hf_config, raising=False)
+    monkeypatch.setattr("vllm_omni.config.config_factory.get_config", lambda model, **kwargs: fake_hf_config)
 
     def _mock_cached_file(path_or_repo_id, *args, **kwargs):
         import os

tests/entrypoints/test_resolve_dreamzero_config.py

@@ -1,28 +1,16 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 import pytest
 
+from vllm_omni.config.resolver import resolve_omni_config
 from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.stage_diffusion_proc import StageDiffusionProc
-from vllm_omni.entrypoints.utils import load_stage_configs_from_model, resolve_model_config_path
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
-def test_dreamzero_vla_resolves_to_dreamzero_config(monkeypatch):
-    monkeypatch.setattr(
-        "vllm_omni.entrypoints.utils.get_config",
-        lambda _model, trust_remote_code=True: type("Cfg", (), {"model_type": "vla"})(),
-    )
-    monkeypatch.setattr(
-        "vllm_omni.entrypoints.utils._looks_like_dreamzero",
-        lambda _model: True,
-    )
-    result = resolve_model_config_path("GEAR-Dreams/DreamZero-DROID")
-
-    assert result is not None
-    assert result.endswith("vllm_omni/deploy/dreamzero.yaml")
-
-
-def test_dreamzero_config_sets_model_class_and_policy_config(monkeypatch):
+def test_dreamzero_resolves_through_registry_with_model_defaults(monkeypatch):
     monkeypatch.setattr(
         "vllm_omni.config.config_factory.StageConfigFactory._try_infer_model_type",
         classmethod(lambda _cls, model, trust_remote_code=True: "vla"),
@@ -36,12 +24,18 @@ def test_dreamzero_config_sets_model_class_and_policy_config(monkeypatch):
         lambda _model: True,
     )
 
-    stage_configs, _ = load_stage_configs_from_model(
+    resolved = resolve_omni_config(
         "GEAR-Dreams/DreamZero-DROID",
         trust_remote_code=False,
+        deploy_config_path=None,
+        cli_overrides=None,
+        stage_overrides=None,
+        strategy_config_path=None,
     )
-    engine_args = stage_configs[0].engine_args
+    engine_args = resolved.stage_configs[0].engine_args
 
+    assert resolved.config_path is not None
+    assert resolved.config_path.endswith("vllm_omni/deploy/dreamzero.yaml")
     assert engine_args.model_class_name == "DreamZeroPipeline"
     assert engine_args.model_config.policy_server_config.action_space == "joint_position"
 

tests/entrypoints/test_serve.py

@@ -11,13 +11,22 @@ from types import SimpleNamespace
 import pytest
 from pytest_mock import MockerFixture
 
-from vllm_omni.entrypoints.cli.serve import OmniServeCommand, run_headless
+from vllm_omni.config.resolver import OmniConfigResolution
+from vllm_omni.entrypoints.cli.serve import (
+    OmniServeCommand,
+    _parse_stage_overrides,
+    run_headless,
+)
 from vllm_omni.entrypoints.utils import parse_stage_overrides
 from vllm_omni.utils.tracking_parser import TrackingArgumentParser, TrackingNamespace
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
+def _resolved(*stages: SimpleNamespace) -> OmniConfigResolution:
+    return OmniConfigResolution(config_path="/fake/stages.yaml", stage_configs=tuple(stages))
+
+
 def test_serve_parser_accepts_no_async_chunk_and_marks_it_explicit() -> None:
     """``--no-async-chunk`` should parse to ``async_chunk=False`` and mark the
     shared deploy-level dest as explicitly provided by the user."""
@@ -110,7 +119,7 @@ def test_serve_parser_accepts_diffusion_quantization_config() -> None:
     assert args.get_explicit_kwargs_dict()["diffusion_quantization_config"] == expected
 
 
-def _make_headless_args(**kwargs) -> TrackingNamespace:
+def _make_headless_args(*, explicit_keys: frozenset[str] | None = None, **kwargs) -> TrackingNamespace:
     defaults = {
         "model": "fake-model",
         "stage_id": 0,
@@ -130,7 +139,7 @@ def _make_headless_args(**kwargs) -> TrackingNamespace:
     ns = argparse.Namespace(**ns_kwargs)
     return TrackingNamespace(
         unfiltered_ns=ns,
-        explicit_keys=frozenset(ns.__dict__.keys()),
+        explicit_keys=frozenset(ns.__dict__.keys()) if explicit_keys is None else explicit_keys,
     )
 
 
@@ -159,41 +168,49 @@ def test_run_headless_rejects_non_multiprocess_worker_backend() -> None:
 
 
 # ---------------------------------------------------------------------------
-# --stage-overrides parsing parity (headless vs standard path)
+# --stage-overrides parsing at the serving boundary
 # ---------------------------------------------------------------------------
 
 
 def test_parse_stage_overrides_valid_json() -> None:
     """A valid JSON string is parsed into the nested per-stage dict."""
-    parsed = parse_stage_overrides('{"0": {"devices": "0,1"}, "1": {"devices": "2"}}')
+    parsed = _parse_stage_overrides('{"0": {"devices": "0,1"}, "1": {"devices": "2"}}')
     assert parsed == {"0": {"devices": "0,1"}, "1": {"devices": "2"}}
 
 
-def test_parse_stage_overrides_none_and_empty_return_none() -> None:
-    """No overrides (None / empty string) resolve to ``None``."""
-    assert parse_stage_overrides(None) is None
-    assert parse_stage_overrides("") is None
+def test_serve_parser_parses_stage_overrides_before_resolution() -> None:
+    parser = TrackingArgumentParser()
+    subparsers = parser.add_subparsers(dest="subcommand")
+    OmniServeCommand().subparser_init(subparsers)
+
+    args = parser.parse_args(
+        [
+            "serve",
+            "fake-model",
+            "--omni",
+            "--stage-overrides",
+            '{"0": {"devices": "0,1"}}',
+        ]
+    )
+
+    assert args.stage_overrides == {"0": {"devices": "0,1"}}
 
 
-def test_parse_stage_overrides_empty_dict_returns_none() -> None:
-    """An empty dict is falsy and must resolve to ``None``, locking in parity
-    with the original standard-path ``if stage_overrides_json:`` falsy check."""
-    assert parse_stage_overrides({}) is None
+def test_serve_parser_accepts_empty_stage_overrides_as_noop() -> None:
+    parser = TrackingArgumentParser()
+    subparsers = parser.add_subparsers(dest="subcommand")
+    OmniServeCommand().subparser_init(subparsers)
 
+    args = parser.parse_args(["serve", "fake-model", "--omni", "--stage-overrides", "{}"])
 
-def test_parse_stage_overrides_passes_through_non_str() -> None:
-    """An already-parsed mapping is returned unchanged (identity)."""
-    overrides = {"0": {"devices": "0"}}
-    assert parse_stage_overrides(overrides) is overrides
+    assert args.stage_overrides == {}
 
 
 def test_parse_stage_overrides_invalid_json_raises() -> None:
-    """Invalid JSON raises ValueError whose message matches the standard path
-    verbatim: the ``--stage-overrides is not valid JSON:`` prefix AND the
-    ``Got: <repr>`` suffix echoing the raw input."""
+    """Invalid JSON fails at the serving boundary with the raw input."""
     bad = "{not valid json}"
-    with pytest.raises(ValueError) as excinfo:
-        parse_stage_overrides(bad)
+    with pytest.raises(argparse.ArgumentTypeError) as excinfo:
+        _parse_stage_overrides(bad)
     message = str(excinfo.value)
     assert message.startswith("--stage-overrides is not valid JSON:")
     assert f"Got: {bad!r}" in message
@@ -271,11 +288,8 @@ def test_parse_stage_overrides_accepts_empty_inner_dict() -> None:
     assert parsed == {"0": {}, "1": {}}
 
 
-def test_run_headless_parses_and_forwards_stage_overrides(mocker: MockerFixture) -> None:
-    """Regression: the headless path must parse ``--stage-overrides`` (a JSON
-    string) and forward the parsed dict to ``load_and_resolve_stage_configs``,
-    mirroring the standard engine path. Previously it was dropped entirely,
-    silently producing a different per-stage device layout."""
+def test_run_headless_forwards_parsed_stage_overrides(mocker: MockerFixture) -> None:
+    """The headless resolver receives the mapping parsed by argparse."""
     captured: dict = {}
 
     def _fake_resolve(*args, **kwargs):
@@ -283,10 +297,10 @@ def test_run_headless_parses_and_forwards_stage_overrides(mocker: MockerFixture)
         captured.update(kwargs)
         # Return a stage that does NOT match stage_id=0 so run_headless stops
         # right after the resolver call (we only care about how it was called).
-        return ("/fake/stages.yaml", [SimpleNamespace(stage_id=99)], None)
+        return _resolved(SimpleNamespace(stage_id=99))
 
     mocker.patch(
-        "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
+        "vllm_omni.config.resolver.resolve_omni_config",
         side_effect=_fake_resolve,
     )
 
@@ -294,28 +308,22 @@ def test_run_headless_parses_and_forwards_stage_overrides(mocker: MockerFixture)
         stage_id=0,
         deploy_config="/tmp/deploy.yaml",
         strategy_config="/tmp/strategy.yaml",
-        stage_overrides='{"0": {"devices": "0,1"}, "1": {"devices": "2"}}',
+        stage_overrides={"0": {"devices": "0,1"}, "1": {"devices": "2"}},
     )
     with pytest.raises(ValueError, match="No stage config found for stage_id=0"):
         run_headless(args)
 
-    assert len(captured["args"]) == 2
+    assert captured["args"] == ("fake-model",)
     assert captured["deploy_config_path"] == "/tmp/deploy.yaml"
     assert captured["stage_overrides"] == {"0": {"devices": "0,1"}, "1": {"devices": "2"}}
     assert captured["strategy_config_path"] == "/tmp/strategy.yaml"
 
 
-def test_run_headless_invalid_stage_overrides_raises(mocker: MockerFixture) -> None:
-    """Invalid ``--stage-overrides`` JSON in headless mode fails fast with the
-    shared ValueError instead of being silently ignored."""
-    mocker.patch(
-        "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
-        return_value=("/fake/stages.yaml", [SimpleNamespace(stage_id=0)], None),
-    )
-
-    args = _make_headless_args(stage_id=0, stage_overrides="{not valid json}")
-    with pytest.raises(ValueError, match="--stage-overrides is not valid JSON"):
-        run_headless(args)
+def test_parse_stage_overrides_rejects_non_mapping_values() -> None:
+    with pytest.raises(argparse.ArgumentTypeError, match="JSON object"):
+        _parse_stage_overrides('["not", "a", "mapping"]')
+    with pytest.raises(argparse.ArgumentTypeError, match="must be an object"):
+        _parse_stage_overrides('{"0": "not a mapping"}')
 
 
 def test_run_headless_raises_when_stage_id_not_in_configs(mocker: MockerFixture) -> None:
@@ -323,8 +331,8 @@ def test_run_headless_raises_when_stage_id_not_in_configs(mocker: MockerFixture)
     fails fast when the launcher's --stage-id doesn't match any entry."""
     other_stage = SimpleNamespace(stage_id=99)
     mocker.patch(
-        "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
-        return_value=("/fake/stages.yaml", [other_stage], None),
+        "vllm_omni.config.resolver.resolve_omni_config",
+        return_value=_resolved(other_stage),
     )
 
     args = _make_headless_args(stage_id=0)
@@ -372,8 +380,8 @@ def test_run_headless_llm_registers_with_auto_assigned_replica_id(mocker: Mocker
     engine_manager = mocker.Mock()
 
     mocker.patch(
-        "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
-        return_value=("/fake/stages.yaml", [stage_cfg], None),
+        "vllm_omni.config.resolver.resolve_omni_config",
+        return_value=_resolved(stage_cfg),
     )
     mocker.patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment")
     mocker.patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=None)
@@ -452,8 +460,8 @@ def test_run_headless_llm_launches_one_manager_per_omni_dp_size_local(mocker: Mo
     manager_b = mocker.Mock()
 
     mocker.patch(
-        "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
-        return_value=("/fake/stages.yaml", [stage_cfg], None),
+        "vllm_omni.config.resolver.resolve_omni_config",
+        return_value=_resolved(stage_cfg),
     )
     mocker.patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment")
     mocker.patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=None)
@@ -511,8 +519,8 @@ def test_run_headless_diffusion_registers_and_spawns_proc(mocker: MockerFixture)
     proc.is_alive.return_value = False
 
     mocker.patch(
-        "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
-        return_value=("/fake/stages.yaml", [stage_cfg], None),
+        "vllm_omni.config.resolver.resolve_omni_config",
+        return_value=_resolved(stage_cfg),
     )
     mocker.patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment")
     mocker.patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=None)
@@ -591,8 +599,8 @@ def test_run_headless_diffusion_raises_on_nonzero_proc_exit(mocker: MockerFixtur
     proc.is_alive.return_value = False
 
     mocker.patch(
-        "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
-        return_value=("/fake/stages.yaml", [stage_cfg], None),
+        "vllm_omni.config.resolver.resolve_omni_config",
+        return_value=_resolved(stage_cfg),
     )
     mocker.patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment")
     mocker.patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=None)

tests/entrypoints/test_utils.py

@@ -4,10 +4,8 @@
 """Unit tests for vllm_omni.entrypoints.utils module."""
 
 import logging
-import os
 from collections import Counter
 from dataclasses import dataclass
-from pathlib import Path
 from types import SimpleNamespace
 
 import pytest
@@ -15,24 +13,20 @@ import torch
 from pytest_mock import MockerFixture
 from vllm.sampling_params import RequestOutputKind, SamplingParams
 
+from vllm_omni.config.pipeline_registry import OMNI_PIPELINES
+from vllm_omni.config.resolver import (
+    OmniConfigResolution,
+    _convert_dataclasses_to_dict,
+    _filter_dict_like_object,
+    resolve_omni_config,
+)
+from vllm_omni.config.stage_config import PipelineConfig
 from vllm_omni.config.yaml_util import create_config
 from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.engine.arg_utils import OmniEngineArgs
-from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
-from vllm_omni.engine.stage_init_utils import (
-    get_stage_connector_spec,
-    load_omni_transfer_config_for_model,
-)
 from vllm_omni.entrypoints.utils import (
-    _convert_dataclasses_to_dict,
-    _filter_dict_like_object,
-    _try_resolve_omni_model_type,
     coerce_param_message_types,
     filter_dataclass_kwargs,
-    filter_stages,
-    load_and_resolve_stage_configs,
-    load_stage_configs_from_model,
-    resolve_model_config_path,
 )
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -173,7 +167,7 @@ class TestFilterDictLikeObject:
             "normal": "value",
         }
 
-        mocker.patch("vllm_omni.entrypoints.utils.logger")
+        mocker.patch("vllm_omni.config.resolver.logger")
         result = _filter_dict_like_object(input_dict)
 
         # Normal key should exist
@@ -195,7 +189,7 @@ class TestFilterDictLikeObject:
             "normal": "value",
         }
 
-        mocker.patch("vllm_omni.entrypoints.utils.logger")
+        mocker.patch("vllm_omni.config.resolver.logger")
         result = _filter_dict_like_object(input_dict)
 
         # Callable should be filtered
@@ -215,7 +209,7 @@ class TestConvertDataclassesToDict:
             "callable": lambda x: x,
         }
 
-        mocker.patch("vllm_omni.entrypoints.utils.logger")
+        mocker.patch("vllm_omni.config.resolver.logger")
         result = _convert_dataclasses_to_dict(input_dict)
 
         # Callable should be filtered out by _filter_dict_like_object
@@ -300,254 +294,174 @@ class TestFilterDataclassKwargs:
         assert "extra_param" not in result["cache_config"]
 
 
-class TestResolveModelConfigPath:
-    """Test suite for resolve_model_config_path function with diffusers format models."""
-
-    def test_glm_image_diffusers_format_resolution(self, mocker: MockerFixture):
-        """Test GlmImagePipeline diffusers class resolves to glm_image config."""
-        mocker.patch(
-            "vllm_omni.entrypoints.utils.file_or_path_exists",
-            return_value=True,
-        )
-        mocker.patch(
-            "vllm_omni.entrypoints.utils._try_get_class_name_from_diffusers_config",
-            return_value="GlmImagePipeline",
-        )
-        mocker.patch(
-            "vllm_omni.entrypoints.utils.current_omni_platform.get_default_stage_config_path",
-            return_value="vllm_omni/model_executor/stage_configs",
+class TestResolveOmniConfig:
+    def test_stage_lookup_error_lists_resolved_ids(self):
+        resolved = OmniConfigResolution(
+            config_path=None,
+            stage_configs=(SimpleNamespace(stage_id=1), SimpleNamespace(stage_id=3)),
         )
 
-        original_exists = os.path.exists
+        with pytest.raises(KeyError, match=r"no stage 2; resolved stages: \[1, 3\]"):
+            resolved.stage_by_id(2)
 
-        def mock_exists(path):
-            if "glm_image.yaml" in str(path):
-                return True
-            return original_exists(path)
-
-        mocker.patch("os.path.exists", side_effect=mock_exists)
-
-        result = resolve_model_config_path("zai-org/GLM-Image")
-
-        assert result is not None
-        assert "glm_image.yaml" in result
-
-    def test_uses_registry_default_deploy_when_hf_model_type_differs(
-        self,
-        mocker: MockerFixture,
-        tmp_path,
-    ):
-        from vllm_omni.config.stage_config import _DEPLOY_DIR as _REAL_DEPLOY_DIR
-        from vllm_omni.entrypoints import utils as utils_mod
-
-        hf_config = SimpleNamespace(
-            model_type="minicpmo",
-            architectures=["MiniCPMO"],
-            version="4.5",
-        )
-        model = "local/minicpmo-registry-default-test"
-        mocker.patch(
-            "vllm_omni.entrypoints.utils.get_config",
-            return_value=hf_config,
-        )
-        mocker.patch(
-            "vllm_omni.config.config_factory.get_config",
-            return_value=hf_config,
-        )
-        # A bare deploy/<hf_model_type>.yaml must not win over the registered
-        # pipeline default when HF model_type is not an OMNI_PIPELINES key.
-        (tmp_path / "minicpmo.yaml").write_text("stages: []\n", encoding="utf-8")
-        (tmp_path / "minicpmo_4_5.yaml").write_bytes((_REAL_DEPLOY_DIR / "minicpmo_4_5.yaml").read_bytes())
-        mocker.patch.object(utils_mod, "_DEPLOY_DIR", tmp_path)
-
-        result = resolve_model_config_path(model)
-
-        assert result is not None
-        assert Path(result).as_posix().endswith("minicpmo_4_5.yaml")
-        assert Path(result).name == "minicpmo_4_5.yaml"
-        transfer_config = load_omni_transfer_config_for_model(model, result)
-        assert transfer_config is not None
-
-        sender = get_stage_connector_spec(
-            omni_transfer_config=transfer_config,
-            stage_id=1,
-            async_chunk=True,
-        )
-        receiver = get_stage_connector_spec(
-            omni_transfer_config=transfer_config,
-            stage_id=2,
-            async_chunk=True,
-        )
-        assert sender["extra"]["role"] == "sender"
-        assert receiver["extra"]["role"] == "receiver"
-        assert sender["extra"]["connector_get_sleep_s"] == 0.01
-        assert sender["extra"]["connector_get_max_wait_first_chunk"] == 3000
-        assert sender["extra"]["connector_get_max_wait"] == 300
-
-    def test_object_storage_uri_resolves_via_materialized_configs(
-        self,
-        mocker: MockerFixture,
-        tmp_path,
-    ):
-        """An object-storage URI must not reach HF helpers; resolution reads the
-        locally materialized config files instead."""
-        hf_config = SimpleNamespace(
-            model_type="qwen3_omni_moe",
-            architectures=["Qwen3OmniMoeForConditionalGeneration"],
-        )
-        uri = "s3://qwen3-tts-models/Qwen3-Omni-30B-A3B-Instruct"
-        materialize = mocker.patch(
-            "vllm_omni.entrypoints.utils._materialize_object_storage_configs",
-            return_value=str(tmp_path),
-        )
-        get_config = mocker.patch(
-            "vllm_omni.entrypoints.utils.get_config",
-            return_value=hf_config,
-        )
-
-        result = resolve_model_config_path(uri)
-
-        materialize.assert_called_once_with(uri)
-        get_config.assert_called_once()
-        assert get_config.call_args[0][0] == str(tmp_path)
-        assert result is not None
-        assert Path(result).as_posix().endswith("qwen3_omni_moe.yaml")
-
-
-class TestTryResolveOmniModelType:
-    """Name matching must only scan the last path component."""
-
-    def test_uri_bucket_name_cannot_hijack_pipeline(self):
-        # "qwen3-tts-models" normalizes to "qwen3ttsmodels", which contains the
-        # registered key "qwen3tts" — only the basename may participate, so a
-        # misleadingly named bucket selects nothing.
-        assert _try_resolve_omni_model_type("s3://qwen3-tts-models/plain-checkpoint") is None
-
-    def test_basename_still_matches_pipeline(self):
-        assert _try_resolve_omni_model_type("gs://any-bucket/Fun-CosyVoice3-0.5B-2512") == "cosyvoice3"
-
-
-class TestLoadAndResolveStageConfigs:
     def test_load_and_resolve_with_kwargs(self, mocker: MockerFixture):
-        """Ensure that kwargs survive default stage creation."""
+        """Ensure that generic diffusion overrides survive resolution."""
         engine_backend = "vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine"
-        kwargs = {"dtype": torch.float32, "engine_backend": engine_backend}
-        kwargs["revision"] = "pinned-revision"
-        mocker.patch("vllm_omni.entrypoints.utils.resolve_model_config_path", return_value=None)
-        mocker.patch("vllm_omni.entrypoints.utils.load_stage_configs_from_model", return_value=([], None))
-
-        config_path, stage_configs, _ = load_and_resolve_stage_configs(
-            model="black-forest-labs/FLUX.2-klein-4B",
-            kwargs=kwargs,
-            trust_remote_code=False,
-            default_stage_cfg_factory=lambda: AsyncOmniEngine._create_default_diffusion_stage_cfg(kwargs),
+        mocker.patch(
+            "vllm_omni.config.resolver.StageConfigFactory.create_from_model",
+            return_value=None,
         )
-        assert config_path is None
-        assert len(stage_configs) == 1
-        assert "dtype" in stage_configs[0]["engine_args"]
-        assert stage_configs[0]["engine_args"]["engine_backend"] == engine_backend
-        assert stage_configs[0]["engine_args"]["revision"] == "pinned-revision"
-
-    def test_deploy_config_preserves_cli_overrides_and_replicas(self, tmp_path, mocker: MockerFixture):
-        deploy_path = tmp_path / "qwen3_multi.yaml"
-        deploy_path.write_text(
-            'stages:\n  - stage_id: 0\n    devices: "0"\n  - stage_id: 1\n    devices: "1,2,3"\n    num_replicas: 3\n',
-            encoding="utf-8",
+        mocker.patch(
+            "vllm_omni.config.resolver._resolve_generic_diffusion_model_class",
+            return_value=(True, "FluxPipeline"),
         )
-
-        returned_stage_configs = [
-            create_config({"stage_id": 0, "runtime": {"devices": "0"}, "engine_args": {"model": "dummy"}}),
-            create_config(
-                {
-                    "stage_id": 1,
-                    "runtime": {"devices": "1,2,3", "num_replicas": 3},
-                    "engine_args": {"model": "dummy"},
-                }
-            ),
-        ]
-        load_stage_configs = mocker.patch(
-            "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
-            return_value=(returned_stage_configs, None),
-        )
-        cli_overrides = {
-            "interleave_mm_strings": True,
-            "media_io_kwargs": {"video": {"fps": 1, "num_frames": 128}},
-            "gpu_memory_utilization": 0.55,
+        kwargs = {
+            "dtype": torch.float32,
+            "engine_backend": engine_backend,
+            "revision": "pinned-revision",
         }
-
-        config_path, stage_configs, _ = load_and_resolve_stage_configs(
-            model="dummy-model",
-            kwargs=cli_overrides,
-            trust_remote_code=True,
-            deploy_config_path=str(deploy_path),
-        )
-
-        load_stage_configs.assert_called_once_with(
-            "dummy-model",
-            trust_remote_code=True,
-            base_engine_args=cli_overrides,
-            deploy_config_path=str(deploy_path),
+        resolved = resolve_omni_config(
+            "black-forest-labs/FLUX.2-klein-4B",
+            trust_remote_code=False,
+            deploy_config_path=None,
+            cli_overrides=kwargs,
             stage_overrides=None,
             strategy_config_path=None,
         )
-        assert config_path == str(deploy_path)
-        assert len(stage_configs) == 2
-        assert stage_configs[1].runtime.num_replicas == 3
-        assert stage_configs[1].runtime.devices == "1,2,3"
+        assert resolved.config_path is None
+        assert len(resolved.stage_configs) == 1
+        engine_args = resolved.stage_configs[0]["engine_args"]
+        assert "dtype" in engine_args
+        assert engine_args["engine_backend"] == engine_backend
+        assert engine_args["revision"] == "pinned-revision"
 
-    def test_filter_stages_selects_mode_stages_without_mutating_stage_config(self, tmp_path):
-        config_path = tmp_path / "deploy.yaml"
-        config_path.write_text(
-            """modes:
-  - mode: text-to-text
-    stages: [0]
-  - mode: text-to-image
-    stages: [0, 1]
-""",
-            encoding="utf-8",
-        )
-        stages = [
-            create_config(
-                {
-                    "stage_id": 0,
-                    "runtime": {"requires_multimodal_data": True},
-                    "final_output": False,
-                    "final_output_type": None,
-                }
-            ),
-            create_config(
-                {
-                    "stage_id": 1,
-                    "runtime": {"requires_multimodal_data": True},
-                    "final_output": True,
-                    "final_output_type": "image",
-                }
-            ),
-        ]
-
-        filtered = filter_stages(str(config_path), stages, {"mode": "text-to-text"})
-
-        assert len(filtered) == 1
-        assert filtered[0].stage_id == 0
-        assert filtered[0].runtime.requires_multimodal_data is True
-        assert filtered[0].final_output is False
-        assert filtered[0].final_output_type is None
-
-
-class TestLoadStageConfigsFromModel:
-    def test_unresolved_model_does_not_fall_back_to_yaml(self, mocker: MockerFixture, caplog):
+    def test_generic_diffusion_uses_registered_model_metadata(self, mocker: MockerFixture):
         mocker.patch(
-            "vllm_omni.entrypoints.utils.StageConfigFactory.create_legacy_stage_configs_from_model",
-            return_value=(None, None),
+            "vllm_omni.config.resolver.StageConfigFactory.create_from_model",
+            return_value=None,
+        )
+        resolve_model_class = mocker.patch(
+            "vllm_omni.config.resolver.resolve_model_class_name",
+            return_value="WanImageToVideoPipeline",
+        )
+        mocker.patch(
+            "vllm_omni.config.resolver.DiffusionModelRegistry.get_supported_archs",
+            return_value={"WanImageToVideoPipeline"},
         )
-        resolve_path = mocker.patch("vllm_omni.entrypoints.utils.resolve_model_config_path")
 
-        result = load_stage_configs_from_model("unregistered-model", trust_remote_code=False)
+        resolved = resolve_omni_config(
+            "/models/Wan2.2-I2V",
+            trust_remote_code=False,
+            deploy_config_path=None,
+            cli_overrides={
+                "diffusion_load_format": "diffusers",
+                "revision": "pinned-revision",
+            },
+            stage_overrides=None,
+            strategy_config_path=None,
+        )
 
-        assert result == ([], None)
-        resolve_path.assert_not_called()
-        assert "No registered PipelineConfig resolved" in caplog.text
-        assert "deploy_config" in caplog.text
+        stage = resolved.stage_configs[0]
+        resolve_model_class.assert_called_once_with(
+            "/models/Wan2.2-I2V",
+            "diffusers",
+            "pinned-revision",
+        )
+        assert stage.engine_args.model_class_name == "WanImageToVideoPipeline"
+        assert stage.final_output_type == "video"
+
+    def test_registered_pipeline_uses_structured_metadata_and_preserves_override_trust(self, mocker: MockerFixture):
+        endpoint_restriction = SimpleNamespace(name="chat")
+        structured_config = SimpleNamespace(
+            orchestrator_config=SimpleNamespace(deploy_config_path="/resolved/deploy.yaml"),
+            pipeline_config=SimpleNamespace(endpoint_restrictions=(endpoint_restriction,)),
+        )
+        runtime_stage = create_config(
+            {
+                "stage_id": 1,
+                "runtime": {"devices": "1,2,3", "num_replicas": 3},
+                "engine_args": {"model": "dummy-model"},
+            }
+        )
+        legacy_stage = SimpleNamespace(to_omegaconf=lambda: runtime_stage)
+        create_structured = mocker.patch(
+            "vllm_omni.config.resolver.StageConfigFactory.create_from_model",
+            return_value=structured_config,
+        )
+        create_legacy = mocker.patch(
+            "vllm_omni.config.resolver.StageConfigFactory._resolve_legacy_from_registry",
+            return_value=SimpleNamespace(
+                stage_configs=[legacy_stage],
+                pipeline_config=structured_config.pipeline_config,
+                omni_lb_policy="round_robin",
+            ),
+        )
+        strategy_specs = {"stage_1": {"dp": 3}}
+        load_strategy = mocker.patch(
+            "vllm_omni.config.resolver._load_strategy_specs",
+            return_value=strategy_specs,
+        )
+
+        resolved = resolve_omni_config(
+            "dummy-model",
+            trust_remote_code=None,
+            deploy_config_path="deploy.yaml",
+            cli_overrides={"dtype": "bfloat16", "trust_remote_code": True},
+            stage_overrides={"1": {"tensor_parallel_size": 2}},
+            strategy_config_path="strategy.yaml",
+        )
+
+        expected_overrides = {
+            "dtype": "bfloat16",
+            "trust_remote_code": True,
+            "stage_1_tensor_parallel_size": 2,
+        }
+        create_structured.assert_called_once_with(
+            "dummy-model",
+            trust_remote_code=None,
+            cli_overrides=expected_overrides,
+            deploy_config_path="deploy.yaml",
+        )
+        create_legacy.assert_called_once_with(
+            structured_config.pipeline_config,
+            expected_overrides,
+            "/resolved/deploy.yaml",
+            strategy_specs=strategy_specs,
+        )
+        load_strategy.assert_called_once_with("strategy.yaml")
+        assert resolved.config_path == "/resolved/deploy.yaml"
+        assert resolved.pipeline_config is structured_config.pipeline_config
+        assert resolved.omni_lb_policy == "round_robin"
+        assert resolved.endpoint_restrictions == (endpoint_restriction,)
+        assert resolved.stage_configs == (runtime_stage,)
+
+    def test_registered_resolution_exposes_forced_aligner_topology(self, mocker: MockerFixture):
+        pipeline = OMNI_PIPELINES["qwen3_tts"]
+        assert isinstance(pipeline, PipelineConfig)
+        structured_config = SimpleNamespace(
+            orchestrator_config=SimpleNamespace(deploy_config_path=None),
+            pipeline_config=pipeline,
+        )
+        mocker.patch(
+            "vllm_omni.config.resolver.StageConfigFactory.create_from_model",
+            return_value=structured_config,
+        )
+
+        resolved = resolve_omni_config(
+            "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice",
+            trust_remote_code=False,
+            deploy_config_path=None,
+            cli_overrides={"forced_aligner": "/models/Qwen3-ForcedAligner-0.6B"},
+            stage_overrides=None,
+            strategy_config_path=None,
+        )
+
+        assert resolved.pipeline_config is not None
+        pipeline_stage_ids = [stage.stage_id for stage in resolved.pipeline_config.stages]
+        runtime_stage_ids = [stage.stage_id for stage in resolved.stage_configs]
+        assert pipeline_stage_ids == runtime_stage_ids
+        assert resolved.pipeline_config.stages[-1].model_stage == "forced_aligner"
+        assert len(resolved.pipeline_config.stages) == len(pipeline.stages) + 1
 
 
 class TestCumulativeStreamingCoercion:

vllm_omni/config/__init__.py

@@ -45,8 +45,10 @@ from vllm_omni.config.yaml_util import (
 # `from vllm_omni.config.lora import LoRAConfig` (used while data.py is still
 # loading) cannot close a circular import through DiffusionOutput.
 _LAZY_ATTRS = {
+    "OmniConfigResolution": ("vllm_omni.config.resolver", "OmniConfigResolution"),
     "StageConfigFactory": ("vllm_omni.config.config_factory", "StageConfigFactory"),
     "register_pipeline": ("vllm_omni.config.pipeline_registry", "register_pipeline"),
+    "resolve_omni_config": ("vllm_omni.config.resolver", "resolve_omni_config"),
 }
 
 __all__ = [
@@ -60,6 +62,8 @@ __all__ = [
     "VllmOmniGenerationStageConfig",
     "VllmOmniDiffusionStageConfig",
     "StageConfigType",
+    "OmniConfigResolution",
+    "resolve_omni_config",
     # Structured Omni sub-configs.
     "OmniStageCacheConfig",
     "OmniStageConnectorConfig",

vllm_omni/config/config_factory.py

@@ -4,10 +4,10 @@
 
 from __future__ import annotations
 
-import dataclasses
 import functools
+import json
 from collections.abc import Mapping
-from dataclasses import asdict
+from dataclasses import asdict, dataclass
 from pathlib import Path
 from typing import Any
 
@@ -31,7 +31,7 @@ from vllm_omni.config.stage_config import (
     merge_pipeline_deploy,
     normalize_pipeline_cli_overrides,
 )
-from vllm_omni.config.yaml_util import create_config
+from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig
 from vllm_omni.diffusion.io_support import get_diffusion_output_type
 from vllm_omni.diffusion.utils.hf_utils import _looks_like_dreamzero
 
@@ -48,6 +48,15 @@ logger = init_logger(__name__)
 _DEFAULT_PARALLEL_DEGREE = 1
 
 
+@dataclass(frozen=True)
+class _LegacyConfigResolution:
+    """Factory-owned result for the temporary legacy runtime bridge."""
+
+    stage_configs: list[StageConfig]
+    pipeline_config: PipelineConfig
+    omni_lb_policy: str | None
+
+
 @functools.cache
 def _materialize_object_storage_configs(model: str) -> str:
     """Materialize an object-storage model URI's config files locally.
@@ -104,6 +113,11 @@ def with_trust_remote_code_override(
 class StageConfigFactory:
     """Factory that loads pipeline YAML and merges CLI overrides.
 
+    Production startup source selection is owned by
+    :func:`vllm_omni.config.resolver.resolve_omni_config`. Factory methods are
+    lower-level construction primitives for that resolver, config internals,
+    and focused tests; entrypoints and engines must not call them directly.
+
     Handles both single-stage and multi-stage models.
 
     Pipelines are declared in ``vllm_omni/config/pipeline_registry.py`` and
@@ -394,11 +408,13 @@ class StageConfigFactory:
         deploy_config_path: str | None,
         strategy_specs: Mapping[Any, Any] | None = None,
     ) -> tuple[list[StageConfig] | None, str | None]:
-        """Build current runtime stage configs from the shared resolution.
+        """Build the migration-only runtime ABI from the shared resolution.
 
-        The engine still consumes the legacy StageConfig/OmegaConf shape.
-        RFC #4021 will replace this transitional path as runtime consumers move
-        to VllmOmniConfig.
+        The engine still consumes the legacy ``StageConfig``/OmegaConf shape.
+        This method is the resolver's temporary compatibility bridge, not an
+        alternative production source-selection entrypoint and not a stable
+        public contract. RFC #4021 will remove it as runtime consumers move to
+        ``VllmOmniConfig``.
         """
         user_deploy_config = cls._load_user_deploy_config(deploy_config_path)
         pipeline_cfg = cls.get_pipeline_config(
@@ -429,14 +445,29 @@ class StageConfigFactory:
         user_deploy_config: DeployConfig | None = None,
         strategy_specs: Mapping[Any, Any] | None = None,
     ) -> tuple[list[StageConfig], str | None]:
-        """Create current runtime StageConfigs from registry + deploy YAML.
+        """Return the existing two-value legacy runtime ABI."""
+        resolved = cls._resolve_legacy_from_registry(
+            pipeline_cfg,
+            cli_overrides,
+            deploy_config_path,
+            user_deploy_config,
+            strategy_specs,
+        )
+        return resolved.stage_configs, resolved.omni_lb_policy
+
+    @classmethod
+    def _resolve_legacy_from_registry(
+        cls,
+        pipeline_cfg: PipelineConfig,
+        cli_overrides: dict[str, Any],
+        deploy_config_path: str | None = None,
+        user_deploy_config: DeployConfig | None = None,
+        strategy_specs: Mapping[Any, Any] | None = None,
+    ) -> _LegacyConfigResolution:
+        """Create runtime stages and retain their effective topology.
 
         Precedence: caller-typed (non-None) value > deploy YAML >
         StageDeployConfig dataclass default.
-
-        Returns ``(stages, omni_lb_policy)`` — the strategy-derived pipeline-wide
-        load-balancer policy (``None`` when no strategy set one) travels with the
-        stages instead of through a mutable out-param.
         """
         cli_overrides = normalize_pipeline_cli_overrides(pipeline_cfg, cli_overrides)
         deploy_cfg: DeployConfig | None
@@ -475,7 +506,11 @@ class StageConfigFactory:
         cls._reconcile_strategy_with_cli(stages, applied)
 
         omni_lb_policy = applied.omni_lb_policy if applied is not None else None
-        return stages, omni_lb_policy
+        return _LegacyConfigResolution(
+            stage_configs=stages,
+            pipeline_config=pipeline_cfg,
+            omni_lb_policy=omni_lb_policy,
+        )
 
     @staticmethod
     def _apply_strategy_specs(
@@ -587,65 +622,64 @@ class StageConfigFactory:
 
     @classmethod
     def create_default_diffusion(cls, kwargs: dict[str, Any]) -> list[dict[str, Any]]:
-        """Single-stage diffusion - no YAML needed.
+        """Build the temporary runtime ABI for a generic diffusion stage.
 
-        Creates a default diffusion stage configuration for single-stage
-        diffusion models. Returns a legacy OmegaConf-compatible dict for
-        backward compatibility with OmniStage.
-
-        Args:
-            kwargs: Engine arguments from CLI/API.
-
-        Returns:
-            List containing a single config dict for the diffusion stage.
+        The terminal diffusion config owns engine defaults and normalization;
+        this compatibility builder only adds Omni stage topology, request
+        defaults, and device placement.
         """
-        # Calculate devices based on parallel config
-        devices = "0"
-        if "parallel_config" in kwargs:
-            num_devices = kwargs["parallel_config"].world_size
-            for i in range(1, num_devices):
-                devices += f",{i}"
+        raw_sampling_params = kwargs.get("default_sampling_params")
+        if isinstance(raw_sampling_params, str):
+            try:
+                raw_sampling_params = json.loads(raw_sampling_params)
+            except json.JSONDecodeError:
+                logger.warning("Invalid default_sampling_params JSON, ignoring stage defaults.")
+                raw_sampling_params = None
+        if not isinstance(raw_sampling_params, Mapping):
+            raw_sampling_params = None
+        default_sampling_params = dict(raw_sampling_params.get("0", {})) if raw_sampling_params else {}
 
-        engine_args: dict[str, Any] = {}
-        for key, value in kwargs.items():
-            if key in ("parallel_config",):
-                continue
-            engine_args[key] = value
-
-        # Serialize parallel_config as dict for OmegaConf. Test helpers
-        # sometimes pass SimpleNamespace rather than a dataclass instance.
-        if "parallel_config" in kwargs:
-            parallel_config = kwargs["parallel_config"]
-            if dataclasses.is_dataclass(parallel_config) and not isinstance(parallel_config, type):
-                engine_args["parallel_config"] = asdict(parallel_config)
-            elif hasattr(parallel_config, "__dict__"):
-                engine_args["parallel_config"] = dict(vars(parallel_config))
-            else:
-                engine_args["parallel_config"] = parallel_config
-
-        engine_args.setdefault("cache_backend", "none")
+        parallel_config = DiffusionParallelConfig.from_stage_overrides(kwargs)
+        if kwargs.get("num_gpus") is not None:
+            parallel_config.resolve_data_parallel_size(int(kwargs["num_gpus"]))
+        engine_args = OmniDiffusionConfig.normalize_init_kwargs(kwargs)
+        engine_args["parallel_config"] = asdict(parallel_config)
         engine_args["model_stage"] = "diffusion"
 
-        # Convert dtype to string for OmegaConf
-        if "dtype" in engine_args:
-            engine_args["dtype"] = str(engine_args["dtype"])
+        extras = dict(engine_args.get("extras") or {})
+        for key, default in (
+            ("auxiliary_text_encoder", None),
+            ("default_llama_model_id", "meta-llama/Meta-Llama-3.1-8B-Instruct"),
+        ):
+            value = kwargs.get(key)
+            if value is not None:
+                extras[key] = value
+            else:
+                extras.setdefault(key, default)
+        engine_args["extras"] = extras
 
-        engine_args.setdefault("max_num_seqs", 1)
         model_class_name = engine_args.get("model_class_name")
+        final_output_type = get_diffusion_output_type(model_class_name)
+        logger.info(
+            "Resolved generic diffusion final_output_type=%r for model_class_name=%r.",
+            final_output_type,
+            model_class_name,
+        )
 
-        config_dict: dict[str, Any] = {
-            "stage_id": 0,
-            "stage_type": StageType.DIFFUSION.value,
-            "runtime": {
-                "process": True,
-                "devices": devices,
-            },
-            "engine_args": create_config(engine_args),
-            "final_output": True,
-            "final_output_type": get_diffusion_output_type(model_class_name),
-        }
-
-        return [config_dict]
+        return [
+            {
+                "stage_id": 0,
+                "stage_type": StageType.DIFFUSION.value,
+                "runtime": {
+                    "process": True,
+                    "devices": ",".join(str(i) for i in range(parallel_config.world_size)),
+                },
+                "engine_args": engine_args,
+                "default_sampling_params": default_sampling_params,
+                "final_output": True,
+                "final_output_type": final_output_type,
+            }
+        ]
 
     @classmethod
     def _merge_cli_overrides(

vllm_omni/config/pipeline_registry.py

@@ -17,9 +17,8 @@ To add a new pipeline:
 
 Out of tree pipeline configs or resolvers can also be registered with register_pipeline.
 
-NOTE: Single-stage diffusion models continue to use the
-``_create_default_diffusion_stage_cfg`` fallback in
-``async_omni_engine.py``; for now we do not add them to registry.
+NOTE: Generic single-stage diffusion is selected by ``config.resolver`` when
+no registered Omni pipeline matches.
 """
 
 from __future__ import annotations

vllm_omni/config/resolver.py

@@ -0,0 +1,266 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from dataclasses import dataclass, fields, is_dataclass
+from typing import Any
+
+from vllm.logger import init_logger
+
+from vllm_omni.config.composable_parallel.strategy_loader import load_strategy_specs
+from vllm_omni.config.config_factory import StageConfigFactory, with_trust_remote_code_override
+from vllm_omni.config.endpoint_policy import EndpointRestriction
+from vllm_omni.config.omni_config import VllmOmniConfig
+from vllm_omni.config.stage_config import PipelineConfig
+from vllm_omni.config.yaml_util import create_config
+from vllm_omni.diffusion.data import resolve_model_class_name
+from vllm_omni.diffusion.registry import DiffusionModelRegistry
+from vllm_omni.diffusion.utils.hf_utils import is_diffusion_model
+
+logger = init_logger(__name__)
+
+
+@dataclass(frozen=True)
+class OmniConfigResolution:
+    """Migration envelope returned by the production config resolver.
+
+    ``stage_configs`` intentionally carries the current OmegaConf-compatible
+    runtime ABI only until stage startup consumes ``VllmOmniConfig`` directly.
+    ``pipeline_config`` is the effective topology after runtime stage injection,
+    so both fields always describe the same set of stages.
+    It is not a stable authoring or extension API; new production callers
+    should resolve through :func:`resolve_omni_config` and must not construct or
+    merge this compatibility shape themselves.
+    """
+
+    config_path: str | None
+    stage_configs: tuple[Any, ...]  # Temporary StageConfig/OmegaConf bridge.
+    pipeline_config: PipelineConfig | None = None
+    omni_lb_policy: str | None = None
+
+    @property
+    def endpoint_restrictions(self) -> tuple[EndpointRestriction, ...]:
+        if self.pipeline_config is None:
+            return ()
+        return tuple(self.pipeline_config.endpoint_restrictions)
+
+    def stage_by_id(self, stage_id: int) -> Any:
+        for stage in self.stage_configs:
+            if stage.stage_id == stage_id:
+                return stage
+        available = [stage.stage_id for stage in self.stage_configs]
+        raise KeyError(f"no stage {stage_id}; resolved stages: {available}")
+
+
+def _filter_dict_like_object(obj: dict | Any) -> dict:
+    """Convert a dict-like object while dropping OmegaConf-incompatible callables."""
+    result = {}
+    filtered_keys = []
+    for key, value in obj.items():
+        # Preserve class objects as import paths for consumers such as
+        # custom_pipeline_args.pipeline_class.
+        if isinstance(value, type):
+            module = getattr(value, "__module__", None)
+            qualname = getattr(value, "__qualname__", getattr(value, "__name__", None))
+            result[key] = f"{module}.{qualname}" if module and qualname and module != "builtins" else qualname
+        elif callable(value):
+            filtered_keys.append(str(key))
+        else:
+            result[key] = _convert_dataclasses_to_dict(value)
+    if filtered_keys:
+        logger.warning(
+            "Filtered out %d callable object(s) from base_engine_args that are not compatible with OmegaConf: %s.",
+            len(filtered_keys),
+            filtered_keys,
+        )
+    return result
+
+
+def _convert_dataclasses_to_dict(obj: Any) -> Any:
+    """Recursively convert caller values to OmegaConf-compatible types."""
+    # Check by class name before dict to cover both collections.Counter and
+    # vllm.utils.Counter without importing either implementation.
+    if hasattr(obj, "__class__") and obj.__class__.__name__ == "Counter":
+        try:
+            return dict(obj)
+        except (TypeError, ValueError):
+            return {}
+    if isinstance(obj, set):
+        return list(obj)
+    if is_dataclass(obj) and not isinstance(obj, type):
+        result = {}
+        for config_field in fields(obj):
+            if not config_field.init:
+                continue
+            value = getattr(obj, config_field.name)
+            # At the CLI/deploy boundary, None means "unset" rather than
+            # "clear an inherited value". Preserve the dataclass default by
+            # omitting fields whose declared default is already None.
+            if value is None and config_field.default is None:
+                continue
+            result[config_field.name] = _convert_dataclasses_to_dict(value)
+        return result
+    if isinstance(obj, dict):
+        return _filter_dict_like_object(obj)
+    if isinstance(obj, type):
+        module = getattr(obj, "__module__", None)
+        qualname = getattr(obj, "__qualname__", getattr(obj, "__name__", None))
+        return f"{module}.{qualname}" if module and qualname and module != "builtins" else qualname
+    if callable(obj):
+        logger.warning(
+            "Cannot convert callable %r to an OmegaConf-compatible value.",
+            obj,
+        )
+        raise TypeError(f"callable {obj!r} is not an OmegaConf-compatible value")
+    if isinstance(obj, (list, tuple)):
+        converted = []
+        for item in obj:
+            if callable(item):
+                logger.warning(
+                    "Filtered callable %r from an OmegaConf-compatible sequence.",
+                    item,
+                )
+                continue
+            converted.append(_convert_dataclasses_to_dict(item))
+        return type(obj)(converted)
+    if hasattr(obj, "keys") and hasattr(obj, "values") and not isinstance(obj, (str, bytes)):
+        try:
+            return _filter_dict_like_object(obj)
+        except (TypeError, ValueError, AttributeError) as exc:
+            logger.warning(
+                "Failed to convert dict-like %s to an OmegaConf-compatible mapping.",
+                type(obj).__name__,
+            )
+            raise TypeError(f"cannot convert dict-like {type(obj).__name__}") from exc
+    return obj
+
+
+def _flatten_stage_overrides(
+    cli_overrides: dict[str, Any],
+    stage_overrides: Mapping[str, Mapping[str, Any]] | None,
+) -> None:
+    """Translate per-stage mappings to the factory's flat CLI override ABI."""
+    if not stage_overrides:
+        return
+    for stage_id, overrides in stage_overrides.items():
+        if not isinstance(overrides, Mapping):
+            raise TypeError(f"stage override {stage_id!r} must be a mapping, got {type(overrides).__name__}")
+        for key, value in overrides.items():
+            cli_overrides[f"stage_{stage_id}_{key}"] = value
+
+
+def _apply_generic_stage_overrides(
+    cli_overrides: dict[str, Any],
+    stage_overrides: Mapping[str, Mapping[str, Any]] | None,
+) -> None:
+    """Apply stage-zero overrides to the generic single-stage fallback."""
+    if not stage_overrides or not (stage_zero := stage_overrides.get("0")):
+        return
+    for key, value in stage_zero.items():
+        if key == "extras":
+            if not isinstance(value, Mapping):
+                raise TypeError(f"stage override '0'.extras must be a mapping, got {type(value).__name__}")
+            cli_overrides[key] = {**dict(cli_overrides.get(key) or {}), **value}
+        else:
+            cli_overrides[key] = value
+
+
+def _load_strategy_specs(strategy_config_path: str | None) -> Mapping[Any, Any] | None:
+    if strategy_config_path is None:
+        return None
+
+    return load_strategy_specs(strategy_config_path)
+
+
+def _build_registered_resolution(
+    structured_config: VllmOmniConfig,
+    *,
+    cli_overrides: dict[str, Any],
+    strategy_config_path: str | None,
+) -> OmniConfigResolution:
+    """Build the temporary runtime view for an already-resolved pipeline."""
+    # Runtime consumers have not yet moved to typed per-stage configs. Use the
+    # factory-owned compatibility bridge instead of reimplementing legacy YAML
+    # discovery and merging in this resolver.
+    effective_deploy_path = structured_config.orchestrator_config.deploy_config_path
+    legacy_resolution = StageConfigFactory._resolve_legacy_from_registry(
+        structured_config.pipeline_config,
+        cli_overrides,
+        effective_deploy_path,
+        strategy_specs=_load_strategy_specs(strategy_config_path),
+    )
+
+    return OmniConfigResolution(
+        config_path=effective_deploy_path,
+        stage_configs=tuple(stage.to_omegaconf() for stage in legacy_resolution.stage_configs),
+        pipeline_config=legacy_resolution.pipeline_config,
+        omni_lb_policy=legacy_resolution.omni_lb_policy,
+    )
+
+
+def _resolve_generic_diffusion_model_class(
+    model: str,
+    cli_overrides: Mapping[str, Any],
+) -> tuple[bool, str | None]:
+    """Detect generic diffusion support and its serving pipeline class."""
+    model_class_name = cli_overrides.get("model_class_name") or resolve_model_class_name(
+        model,
+        str(cli_overrides.get("diffusion_load_format") or "default"),
+        cli_overrides.get("revision"),
+    )
+    supported = bool(model_class_name and model_class_name in DiffusionModelRegistry.get_supported_archs())
+    if not supported:
+        supported = is_diffusion_model(model)
+    return supported, str(model_class_name) if model_class_name else None
+
+
+def resolve_omni_config(
+    model: str,
+    *,
+    trust_remote_code: bool | None,
+    deploy_config_path: str | None,
+    cli_overrides: Mapping[str, Any] | None,
+    stage_overrides: Mapping[str, Mapping[str, Any]] | None,
+    strategy_config_path: str | None,
+) -> OmniConfigResolution:
+    """Resolve registry/deploy inputs through the single public entrypoint."""
+    normalized_overrides = _convert_dataclasses_to_dict(dict(cli_overrides or {}))
+    normalized_overrides = with_trust_remote_code_override(normalized_overrides, trust_remote_code)
+    registry_overrides = dict(normalized_overrides)
+    _flatten_stage_overrides(registry_overrides, stage_overrides)
+
+    structured_config = StageConfigFactory.create_from_model(
+        model,
+        trust_remote_code=trust_remote_code,
+        cli_overrides=registry_overrides,
+        deploy_config_path=deploy_config_path,
+    )
+    if structured_config is not None:
+        return _build_registered_resolution(
+            structured_config,
+            cli_overrides=registry_overrides,
+            strategy_config_path=strategy_config_path,
+        )
+
+    _apply_generic_stage_overrides(normalized_overrides, stage_overrides)
+    supported, model_class_name = _resolve_generic_diffusion_model_class(model, normalized_overrides)
+    if not supported:
+        raise ValueError(
+            f"Model {model!r} did not resolve to a registered Omni pipeline or a supported diffusion model."
+        )
+    if model_class_name is not None:
+        normalized_overrides.setdefault("model_class_name", model_class_name)
+    default_stages = StageConfigFactory.create_default_diffusion(normalized_overrides)
+
+    return OmniConfigResolution(
+        config_path=deploy_config_path,
+        stage_configs=tuple(create_config(_convert_dataclasses_to_dict(default_stages))),
+    )
+
+
+__all__ = [
+    "OmniConfigResolution",
+    "resolve_omni_config",
+]

vllm_omni/diffusion/data.py

@@ -2,6 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 import copy
+import json
 import math
 import os
 import random
@@ -42,45 +43,57 @@ if TYPE_CHECKING:
 logger = init_logger(__name__)
 
 
-def normalize_omni_diffusion_kwargs(kwargs: Mapping[str, Any]) -> dict[str, Any]:
+def normalize_omni_diffusion_kwargs(raw_kwargs: Mapping[str, Any]) -> dict[str, Any]:
     """Normalize legacy diffusion kwargs before config construction."""
-    normalized = dict(kwargs)
+    config_kwargs = dict(raw_kwargs)
+
+    dtype = config_kwargs.get("dtype")
+    if dtype is None:
+        config_kwargs["dtype"] = "auto"
+    elif isinstance(dtype, torch.dtype):
+        config_kwargs["dtype"] = str(dtype).removeprefix("torch.")
+    elif not isinstance(dtype, str):
+        raise TypeError(f"Provided dtype must be a string or torch.dtype, got {type(dtype).__name__}")
 
     # Backwards-compatibility: older callers may use a diffusion-specific
     # "static_lora_scale" kwarg. Normalize it to the canonical "lora_scale".
-    if "static_lora_scale" in normalized:
-        if "lora_scale" not in normalized:
-            normalized["lora_scale"] = normalized["static_lora_scale"]
-        normalized.pop("static_lora_scale", None)
+    if "static_lora_scale" in config_kwargs:
+        if "lora_scale" not in config_kwargs:
+            config_kwargs["lora_scale"] = config_kwargs["static_lora_scale"]
+        config_kwargs.pop("static_lora_scale", None)
+
+    diffusion_quantization = config_kwargs.pop("diffusion_quantization_config", None)
+    if config_kwargs.get("quantization_config") is None and diffusion_quantization is not None:
+        config_kwargs["quantization_config"] = diffusion_quantization
 
     # Backwards-compatibility: map "quantization" to "quantization_config"
     # so callers using the old field name still work.
-    if "quantization" in normalized and normalized.get("quantization_config", None) is None:
-        normalized["quantization_config"] = normalized.pop("quantization")
+    if "quantization" in config_kwargs and config_kwargs.get("quantization_config", None) is None:
+        config_kwargs["quantization_config"] = config_kwargs.pop("quantization")
     else:
-        normalized.pop("quantization", None)
+        config_kwargs.pop("quantization", None)
 
     # Renamed from kv_cache_* to avoid clashing with vLLM's --kv-cache-dtype.
-    if normalized.get("diffusion_kv_cache_dtype") is None and "kv_cache_dtype" in normalized:
-        normalized["diffusion_kv_cache_dtype"] = normalized.pop("kv_cache_dtype")
+    if config_kwargs.get("diffusion_kv_cache_dtype") is None and "kv_cache_dtype" in config_kwargs:
+        config_kwargs["diffusion_kv_cache_dtype"] = config_kwargs.pop("kv_cache_dtype")
     else:
-        normalized.pop("kv_cache_dtype", None)
-    if normalized.get("diffusion_kv_cache_skip_steps") is None and "kv_cache_skip_steps" in normalized:
-        normalized["diffusion_kv_cache_skip_steps"] = normalized.pop("kv_cache_skip_steps")
+        config_kwargs.pop("kv_cache_dtype", None)
+    if config_kwargs.get("diffusion_kv_cache_skip_steps") is None and "kv_cache_skip_steps" in config_kwargs:
+        config_kwargs["diffusion_kv_cache_skip_steps"] = config_kwargs.pop("kv_cache_skip_steps")
     else:
-        normalized.pop("kv_cache_skip_steps", None)
-    if normalized.get("diffusion_kv_cache_skip_layers") is None and "kv_cache_skip_layers" in normalized:
-        normalized["diffusion_kv_cache_skip_layers"] = normalized.pop("kv_cache_skip_layers")
+        config_kwargs.pop("kv_cache_skip_steps", None)
+    if config_kwargs.get("diffusion_kv_cache_skip_layers") is None and "kv_cache_skip_layers" in config_kwargs:
+        config_kwargs["diffusion_kv_cache_skip_layers"] = config_kwargs.pop("kv_cache_skip_layers")
     else:
-        normalized.pop("kv_cache_skip_layers", None)
+        config_kwargs.pop("kv_cache_skip_layers", None)
 
     # Handle "diffusion_attention_backend" shorthand: merge into
     # diffusion_attention_config before field filtering.
-    diffusion_attn_backend = normalized.pop("diffusion_attention_backend", None)
-    fastvideo_vsa_topk = normalized.pop("fastvideo_vsa_topk", None)
+    diffusion_attn_backend = config_kwargs.pop("diffusion_attention_backend", None)
+    fastvideo_vsa_topk = config_kwargs.pop("fastvideo_vsa_topk", None)
     if diffusion_attn_backend is not None or fastvideo_vsa_topk is not None:
-        existing = normalized.get("diffusion_attention_config")
-        normalized["diffusion_attention_config"] = parse_attention_config(
+        existing = config_kwargs.get("diffusion_attention_config")
+        config_kwargs["diffusion_attention_config"] = parse_attention_config(
             existing,
             attention_backend=diffusion_attn_backend,
             fastvideo_vsa_topk=fastvideo_vsa_topk,
@@ -88,21 +101,33 @@ def normalize_omni_diffusion_kwargs(kwargs: Mapping[str, Any]) -> dict[str, Any]
 
     # Check environment variable as fallback for cache_backend.
     # Support both old DIFFUSION_CACHE_ADAPTER and new DIFFUSION_CACHE_BACKEND.
-    if "cache_backend" not in normalized:
+    if "cache_backend" not in config_kwargs:
         cache_backend = os.environ.get("DIFFUSION_CACHE_BACKEND") or os.environ.get("DIFFUSION_CACHE_ADAPTER")
-        normalized["cache_backend"] = cache_backend.lower() if cache_backend else "none"
-    elif normalized["cache_backend"] is None:
+        config_kwargs["cache_backend"] = cache_backend.lower() if cache_backend else "none"
+    elif config_kwargs["cache_backend"] is None:
         # Callers (e.g. example CLIs with `default=None`) pass an explicit
         # None for "no cache"; canonicalize it so every consumer sees the
         # declared `str` value instead of relying on per-model None handling.
-        normalized["cache_backend"] = "none"
+        config_kwargs["cache_backend"] = "none"
+
+    cache_config = config_kwargs.get("cache_config")
+    if isinstance(cache_config, str):
+        try:
+            config_kwargs["cache_config"] = json.loads(cache_config)
+        except json.JSONDecodeError:
+            logger.warning("Invalid cache_config JSON, using backend defaults.")
+            config_kwargs.pop("cache_config", None)
+
+    if config_kwargs.get("streaming_output") is None and config_kwargs.get("diffusion_streaming_output") is not None:
+        config_kwargs["streaming_output"] = config_kwargs["diffusion_streaming_output"]
+    config_kwargs.pop("diffusion_streaming_output", None)
 
     # Convert optional YAML null values to empty containers.
     for key in ("diffusers_load_kwargs", "diffusers_call_kwargs"):
-        if key in normalized and normalized[key] is None:
-            normalized[key] = {}
+        if key in config_kwargs and config_kwargs[key] is None:
+            config_kwargs[key] = {}
 
-    return normalized
+    return config_kwargs
 
 
 def validate_host_weight_runtime_options(*, mode: object, root: object) -> None:
@@ -439,6 +464,27 @@ class DiffusionParallelConfig:
             raise TypeError(f"Expected parallel config dict, got {type(data)!r}")
         return cls(**data)
 
+    @classmethod
+    def from_stage_overrides(cls, kwargs: Mapping[str, Any]) -> "DiffusionParallelConfig":
+        """Resolve nested or flat stage inputs through the config's own defaults."""
+        parallel_config = kwargs.get("parallel_config")
+        if isinstance(parallel_config, cls):
+            return parallel_config
+
+        valid_fields = {item.name for item in fields(cls) if item.init}
+        if isinstance(parallel_config, Mapping):
+            return cls.from_dict(dict(parallel_config))
+        elif parallel_config is not None and hasattr(parallel_config, "__dict__"):
+            return cls.from_dict(dict(vars(parallel_config)))
+        elif parallel_config is not None:
+            raise TypeError(
+                f"parallel_config must be a mapping or DiffusionParallelConfig, got {type(parallel_config).__name__}"
+            )
+        else:
+            values = {key: value for key, value in kwargs.items() if key in valid_fields}
+
+        return cls(**{key: value for key, value in values.items() if value is not None})
+
 
 @dataclass
 class TransformerConfig:
@@ -1550,15 +1596,18 @@ class OmniDiffusionConfig:
                     raise
 
     @classmethod
-    def from_kwargs(cls, **kwargs: Any) -> "OmniDiffusionConfig":
-        kwargs = normalize_omni_diffusion_kwargs(kwargs)
-
-        # Filter kwargs to only include valid fields
+    def normalize_init_kwargs(cls, raw_kwargs: Mapping[str, Any]) -> dict[str, Any]:
+        config_kwargs = normalize_omni_diffusion_kwargs(raw_kwargs)
         valid_fields = {f.name for f in fields(cls)}
-        filtered_kwargs = {k: v for k, v in kwargs.items() if k in valid_fields}
+        # Remaining ``None`` values mean "unset" at the CLI/deploy boundary.
+        # Drop them so non-optional dataclass defaults are not overwritten.
+        # Fields where ``None`` has normalization semantics (for example dtype
+        # and nullable container inputs) are handled above before this filter.
+        return {key: value for key, value in config_kwargs.items() if key in valid_fields and value is not None}
 
-        instance = cls(**filtered_kwargs)
-        return instance
+    @classmethod
+    def from_kwargs(cls, **kwargs: Any) -> "OmniDiffusionConfig":
+        return cls(**cls.normalize_init_kwargs(kwargs))
 
 
 @dataclass

vllm_omni/engine/arg_utils.py

@@ -518,7 +518,7 @@ class OrchestratorArgs:
 
     # === Config Files ===
     deploy_config: str | None = None
-    stage_overrides: str | None = None  # raw JSON string; parsed downstream
+    stage_overrides: dict[str, dict[str, Any]] | None = None
     # Optional composable-parallel strategy.yaml; orchestrator reads it, overlays
     # derived sizing onto merged stages, then drops it before per-stage engine args.
     strategy_config: str | None = None

vllm_omni/engine/async_omni_engine.py

@@ -13,7 +13,6 @@ from __future__ import annotations
 import asyncio
 import concurrent.futures
 import copy
-import json
 import queue
 import shutil
 import threading
@@ -21,30 +20,23 @@ import time
 import uuid
 import weakref
 from collections.abc import Mapping, Sequence
-from dataclasses import asdict
 from typing import Any, Literal, cast
 
 import janus
-import torch
-from omegaconf import OmegaConf
 from vllm import envs as vllm_envs
 from vllm.inputs import PromptType
 from vllm.logger import init_logger
 from vllm.v1.engine import EngineCoreRequest
 from vllm.v1.engine.input_processor import InputProcessor
 
-from vllm_omni.config.config_factory import StageConfigFactory, with_trust_remote_code_override
+from vllm_omni.config.config_factory import StageConfigFactory
+from vllm_omni.config.resolver import OmniConfigResolution, resolve_omni_config
 from vllm_omni.config.stage_config import (
     DuplexSessionRuntimeConfig,
+    PipelineConfig,
     load_deploy_config,
 )
 from vllm_omni.data_entry_keys import REQUEST_ARTIFACT_DIRS_KEY, TRANSFORM_OWNED_META_KEYS
-from vllm_omni.diffusion.data import (
-    DiffusionParallelConfig,
-    parse_attention_config,
-    resolve_model_class_name,
-)
-from vllm_omni.diffusion.io_support import get_diffusion_output_type
 from vllm_omni.engine import OmniEngineCoreRequest
 from vllm_omni.engine.async_engine_utils import (
     SHUTDOWN_ENQUEUE_TIMEOUT_S,
@@ -88,10 +80,7 @@ from vllm_omni.engine.stage_runtime import (
     create_stage_runtime,
 )
 from vllm_omni.entrypoints.pd_utils import PDDisaggregationMixin
-from vllm_omni.entrypoints.utils import (
-    load_and_resolve_stage_configs,
-    parse_stage_overrides,
-)
+from vllm_omni.entrypoints.utils import parse_stage_overrides
 from vllm_omni.inputs.data import OmniInteractionPrompt, OmniSamplingParams
 from vllm_omni.metrics.prometheus import OmniRequestCounter
 
@@ -99,6 +88,27 @@ logger = init_logger(__name__)
 
 _STARTUP_POLL_INTERVAL_S = 1.0
 _REQUEST_QUEUE_MAXSIZE = 256
+_ConfigResolutionResult = OmniConfigResolution | tuple[str | None, list[Any], str | None]
+
+
+def load_and_resolve_stage_configs(
+    model: str,
+    kwargs: dict[str, Any],
+    *,
+    trust_remote_code: bool | None,
+    deploy_config_path: str | None,
+    stage_overrides: Mapping[str, Mapping[str, Any]] | None,
+    strategy_config_path: str | None,
+) -> OmniConfigResolution:
+    """Compatibility seam delegating to the single config resolver."""
+    return resolve_omni_config(
+        model,
+        trust_remote_code=trust_remote_code,
+        cli_overrides=kwargs,
+        deploy_config_path=deploy_config_path,
+        stage_overrides=stage_overrides,
+        strategy_config_path=strategy_config_path,
+    )
 
 
 class AsyncOmniEngine:
@@ -198,41 +208,28 @@ class AsyncOmniEngine:
                 self._omni_master_port,
             )
 
-        # Stage resolution pops deploy_config, so get pipeline-wide settings beforehand.
+        # Keep the historical tuple return from _resolve_stage_configs while
+        # retaining the richer resolver result for pipeline-wide settings.
+        # Overrides of that private seam fall back to the factory below.
         deploy_config_path = kwargs.get("deploy_config")
-        # ``trust_remote_code`` is tri-state (bool | None): ``None`` means "not
-        # specified" so stage-config resolution can defer to the deploy yaml's
-        # per-stage value (see ``with_trust_remote_code_override``). The
-        # restriction path below loads the top-level HF config via vLLM's
-        # ``get_config``, which needs a real bool, so collapse ``None`` to the
-        # default ``False`` here (#5495).
-        pipeline_config = StageConfigFactory.get_pipeline_config(
-            model=model,
-            trust_remote_code=bool(trust_remote_code),
-            deploy_config_path=deploy_config_path,
-        )
-        self.endpoint_restrictions = pipeline_config.endpoint_restrictions if pipeline_config is not None else ()
-        self._duplex_runtime_extension_path = (
-            pipeline_config.duplex_runtime_extension if pipeline_config is not None else None
-        )
-        self.duplex_serving_adapter_path = (
-            pipeline_config.duplex_serving_adapter if pipeline_config is not None else None
-        )
-        self._duplex_control_enabled = bool(pipeline_config and pipeline_config.duplex_control_enabled)
-        self.duplex_session_config = DuplexSessionRuntimeConfig()
-        if deploy_config_path is not None:
-            self.duplex_session_config = load_deploy_config(deploy_config_path).duplex_session
-
-        # Tri-state: None means "not specified" — the deploy yaml's per-stage
-        # trust_remote_code stays in effect. An explicit True/False here is a
-        # global override (precedence: caller > deploy yaml > default False);
-        # the merge rule lives in with_trust_remote_code_override.
-        kwargs = with_trust_remote_code_override(kwargs, trust_remote_code)
+        self._config_resolution: OmniConfigResolution | None = None
         self.config_path, self.stage_configs = self._resolve_stage_configs(
             model,
             kwargs,
             trust_remote_code=trust_remote_code,
         )
+        if self._config_resolution is None:
+            pipeline_config = StageConfigFactory.get_pipeline_config(
+                model=model,
+                trust_remote_code=bool(trust_remote_code),
+                deploy_config_path=deploy_config_path,
+            )
+            self._set_pipeline_runtime_config(pipeline_config, deploy_config_path)
+        else:
+            self._set_pipeline_runtime_config(
+                self._config_resolution.pipeline_config,
+                self._config_resolution.config_path,
+            )
 
         self.num_stages = len(self.stage_configs)
         stage0_args = getattr(self.stage_configs[0], "engine_args", None) if self.num_stages > 0 else None
@@ -921,52 +918,6 @@ class AsyncOmniEngine:
             len(companions),
         )
 
-    @staticmethod
-    def _get_default_cache_config(cache_backend: str | None) -> dict[str, Any] | None:
-        if cache_backend == "cache_dit":
-            return {
-                "Fn_compute_blocks": 1,
-                "Bn_compute_blocks": 0,
-                "max_warmup_steps": 4,
-                "residual_diff_threshold": 0.24,
-                "max_continuous_cached_steps": 3,
-                "enable_taylorseer": False,
-                "taylorseer_order": 1,
-                "scm_steps_mask_policy": None,
-                "scm_steps_policy": "dynamic",
-            }
-        if cache_backend == "tea_cache":
-            return {
-                "rel_l1_thresh": 0.2,
-            }
-        if cache_backend == "mag_cache":
-            return {
-                "mag_threshold": 0.24,
-                "mag_max_skip_steps": 5,
-                "mag_retention_ratio": 0.1,
-            }
-        if cache_backend == "step_cache":
-            return {
-                "step_cache_dit_enabled": True,
-                "velocity_sim_thresholds": [0.95, 0.93],
-                "velocity_skip_countdowns": [4, 2],
-                "step_cache_dit_min_history": 2,
-                "step_cache_dit_max_history": 2,
-            }
-        return None
-
-    @staticmethod
-    def _normalize_cache_config(cache_backend: str | None, cache_config: Any | None) -> Any | None:
-        if isinstance(cache_config, str):
-            try:
-                cache_config = json.loads(cache_config)
-            except json.JSONDecodeError:
-                logger.warning("Invalid cache_config JSON, using defaults.")
-                cache_config = None
-        if cache_config is None and cache_backend not in (None, "", "none"):
-            cache_config = AsyncOmniEngine._get_default_cache_config(cache_backend)
-        return cache_config
-
     def _detect_pd_config(self) -> dict[str, Any] | None:
         """Detect PD (Prefill-Decode) disaggregation config from stage_configs.
         Returns a dict with 'pd_pair' and 'bootstrap_addr', or None.
@@ -1009,219 +960,6 @@ class AsyncOmniEngine:
             "prefill_engine_id": prefill_engine_id,
         }
 
-    @staticmethod
-    def _create_default_diffusion_stage_cfg(kwargs: dict[str, Any]) -> list:
-        """Create a default single-stage diffusion config from kwargs."""
-        # We temporally create a default config for diffusion stage.
-        # In the future, we should merge the default config with the user-provided config.
-        normalized_kwargs = dict(kwargs)
-        default_sampling_params = normalized_kwargs.get("default_sampling_params")
-        if isinstance(default_sampling_params, str):
-            try:
-                default_sampling_params = json.loads(default_sampling_params)
-            except json.JSONDecodeError:
-                logger.warning("Invalid default_sampling_params JSON, ignoring stage defaults.")
-                default_sampling_params = None
-        if not isinstance(default_sampling_params, dict):
-            default_sampling_params = None
-        stage_default_sampling_params = default_sampling_params.get("0", {}) if default_sampling_params else {}
-        if normalized_kwargs.get("dtype") is None:
-            normalized_kwargs["dtype"] = "auto"
-
-        # TODO: hack, convert dtype to string to avoid non-premitive omegaconf create error.
-        if "dtype" in normalized_kwargs and not isinstance(normalized_kwargs["dtype"], str):
-            if not isinstance(normalized_kwargs["dtype"], torch.dtype):
-                raise TypeError(
-                    f"Provided dtype must be a string or torch.dtype, got {type(normalized_kwargs['dtype']).__name__}"
-                )
-            normalized_kwargs["dtype"] = str(normalized_kwargs["dtype"]).removeprefix("torch.")
-
-        cache_backend = normalized_kwargs.get("cache_backend", "none")
-        cache_config = AsyncOmniEngine._normalize_cache_config(
-            cache_backend,
-            normalized_kwargs.get("cache_config", None),
-        )
-
-        parallel_config = normalized_kwargs.get("parallel_config")
-        if isinstance(parallel_config, dict):
-            parallel_config = DiffusionParallelConfig.from_dict(parallel_config)
-        if parallel_config is None:
-            ulysses_degree = normalized_kwargs.get("ulysses_degree") or 1
-            ring_degree = normalized_kwargs.get("ring_degree") or 1
-            allgather_degree = normalized_kwargs.get("allgather_degree") or 1
-            ulysses_mode = normalized_kwargs.get("ulysses_mode") or "strict"
-            ulysses_a2a_permute = bool(normalized_kwargs.get("ulysses_a2a_permute", False))
-            sequence_parallel_size = normalized_kwargs.get("sequence_parallel_size")
-            pipeline_parallel_size = normalized_kwargs.get("pipeline_parallel_size") or 1
-            data_parallel_size = normalized_kwargs.get("data_parallel_size")
-            tensor_parallel_size = normalized_kwargs.get("tensor_parallel_size") or 1
-            cfg_parallel_size = normalized_kwargs.get("cfg_parallel_size") or 1
-            pipeline_parallel_size = normalized_kwargs.get("pipeline_parallel_size") or 1
-            vae_patch_parallel_size = normalized_kwargs.get("vae_patch_parallel_size") or 1
-            vae_parallel_mode = normalized_kwargs.get("vae_parallel_mode") or "tile"
-            text_encoder_tp_size = normalized_kwargs.get("text_encoder_tp_size") or 1
-            enable_expert_parallel = normalized_kwargs.get("enable_expert_parallel") or False
-            use_hsdp = normalized_kwargs.get("use_hsdp", False)
-            hsdp_shard_size = normalized_kwargs.get("hsdp_shard_size", -1)
-            hsdp_replicate_size = normalized_kwargs.get("hsdp_replicate_size", 1)
-            if sequence_parallel_size is None:
-                sequence_parallel_size = allgather_degree if allgather_degree > 1 else ulysses_degree * ring_degree
-
-            parallel_config = DiffusionParallelConfig(
-                pipeline_parallel_size=pipeline_parallel_size,
-                data_parallel_size=data_parallel_size,
-                tensor_parallel_size=tensor_parallel_size,
-                enable_expert_parallel=enable_expert_parallel,
-                sequence_parallel_size=sequence_parallel_size,
-                ulysses_degree=ulysses_degree,
-                ring_degree=ring_degree,
-                allgather_degree=allgather_degree,
-                ulysses_mode=ulysses_mode,
-                ulysses_a2a_permute=ulysses_a2a_permute,
-                cfg_parallel_size=cfg_parallel_size,
-                vae_patch_parallel_size=vae_patch_parallel_size,
-                vae_parallel_mode=vae_parallel_mode,
-                text_encoder_tp_size=text_encoder_tp_size,
-                use_hsdp=use_hsdp,
-                hsdp_shard_size=hsdp_shard_size,
-                hsdp_replicate_size=hsdp_replicate_size,
-            )
-
-        num_gpus = normalized_kwargs.get("num_gpus")
-        if num_gpus is not None:
-            num_gpus = int(num_gpus)
-            parallel_config.resolve_data_parallel_size(num_gpus)
-
-        num_devices = max(1, int(parallel_config.world_size))
-        devices = ",".join(str(i) for i in range(num_devices))
-        model_class_name = kwargs.get("model_class_name", None)
-        final_output_type = get_diffusion_output_type(model_class_name)
-
-        attention_config = None
-        if (
-            kwargs.get("diffusion_attention_config") is not None
-            or kwargs.get("diffusion_attention_backend") is not None
-            or kwargs.get("fastvideo_vsa_topk") is not None
-        ):
-            attention_config = parse_attention_config(
-                kwargs.get("diffusion_attention_config"),
-                attention_backend=kwargs.get("diffusion_attention_backend"),
-                fastvideo_vsa_topk=kwargs.get("fastvideo_vsa_topk"),
-            )
-
-        extras = dict(kwargs.get("extras") or {})
-        for key, default in (
-            ("auxiliary_text_encoder", None),
-            ("default_llama_model_id", "meta-llama/Meta-Llama-3.1-8B-Instruct"),
-        ):
-            top_level_value = kwargs.get(key)
-            if top_level_value is not None:
-                extras[key] = top_level_value
-            else:
-                extras.setdefault(key, default)
-
-        stage_engine_args = {
-            "max_num_seqs": kwargs.get("max_num_seqs") or 1,
-            "parallel_config": parallel_config,
-            # Default-stage construction bypasses the structured projection.
-            # Runner selection remains owned by the selected engine/platform.
-            "engine_backend": kwargs.get("engine_backend", "default"),
-            "model_class_name": kwargs.get("model_class_name", None),
-            "task_type": kwargs.get("task_type", None),
-            "model_config": kwargs.get("model_config", None),
-            "additional_config": kwargs.get("additional_config", None),
-            "step_execution": kwargs.get("step_execution", False),
-            "request_batch_max_wait_ms": kwargs.get("request_batch_max_wait_ms", 0.0),
-            "vae_use_slicing": kwargs.get("vae_use_slicing", False),
-            "vae_use_tiling": kwargs.get("vae_use_tiling", False),
-            "cache_backend": cache_backend,
-            "cache_config": cache_config,
-            "enable_cache_dit_summary": kwargs.get("enable_cache_dit_summary", False),
-            "diffusion_offload_config": kwargs.get("diffusion_offload_config", None),
-            "enable_cpu_offload": kwargs.get("enable_cpu_offload", False),
-            "enable_layerwise_offload": kwargs.get("enable_layerwise_offload", False),
-            "enable_distributed_layerwise_offload": kwargs.get("enable_distributed_layerwise_offload", False),
-            "dlo_use_allgather": kwargs.get("dlo_use_allgather", True),
-            "dlo_resident_layers": kwargs.get("dlo_resident_layers", 0),
-            "host_weight_runtime_mode": kwargs.get("host_weight_runtime_mode", "disabled"),
-            "host_weight_runtime_root": kwargs.get("host_weight_runtime_root"),
-            "dlo_host_registration_limit_gib": kwargs.get("dlo_host_registration_limit_gib", 0.0),
-            "enforce_eager": False if kwargs.get("enforce_eager") is None else kwargs.get("enforce_eager"),
-            "diffusion_compile_granularity": (
-                "regional"
-                if kwargs.get("diffusion_compile_granularity") is None
-                else kwargs["diffusion_compile_granularity"]
-            ),
-            "diffusion_compile_dynamic": (
-                True if kwargs.get("diffusion_compile_dynamic") is None else kwargs["diffusion_compile_dynamic"]
-            ),
-            "fa_deterministic": bool(kwargs.get("fa_deterministic", False)),
-            "boundary_ratio": kwargs.get("boundary_ratio", None),
-            "flow_shift": kwargs.get("flow_shift", None),
-            "diffusion_load_format": kwargs.get("diffusion_load_format", "default"),
-            "lora_path": kwargs.get("lora_path", None),
-            "lora_scale": kwargs.get("lora_scale", 1.0),
-            "lora_backend": kwargs.get("lora_backend", "peft"),
-            "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": kwargs.get("distributed_executor_backend"),
-            "enable_sleep_mode": kwargs.get("enable_sleep_mode", False),
-            "enable_prompt_embed_cache": kwargs.get("enable_prompt_embed_cache", False),
-            "prompt_embed_cache_size": kwargs.get("prompt_embed_cache_size", 32),
-            "enable_multithread_weight_load": kwargs.get("enable_multithread_weight_load", True),
-            "num_weight_load_threads": kwargs.get("num_weight_load_threads", 4),
-            "quantization": kwargs.get("quantization", None),
-            "quantization_config": kwargs.get("quantization_config", None),
-            "diffusion_kv_cache_dtype": kwargs.get("diffusion_kv_cache_dtype", None),
-            "diffusion_kv_cache_skip_steps": kwargs.get("diffusion_kv_cache_skip_steps", None),
-            "diffusion_kv_cache_skip_layers": kwargs.get("diffusion_kv_cache_skip_layers", None),
-            **({"diffusion_attention_config": attention_config} if attention_config is not None else {}),
-            "force_cutlass_fp8": bool(kwargs.get("force_cutlass_fp8", False)),
-            "enable_diffusion_pipeline_profiler": kwargs.get("enable_diffusion_pipeline_profiler", False),
-            "streaming_output": kwargs.get("diffusion_streaming_output", False),
-            "enable_ar_profiler": kwargs.get("enable_ar_profiler", False),
-            "extras": extras,
-            **(
-                {
-                    "profiler_config": asdict(kwargs["profiler_config"])
-                    if hasattr(kwargs["profiler_config"], "__dataclass_fields__")
-                    else kwargs["profiler_config"]
-                }
-                if kwargs.get("profiler_config") is not None
-                else {}
-            ),
-        }
-        if num_gpus is not None:
-            stage_engine_args["num_gpus"] = num_gpus
-        # Only set dtype if it was already explicitly passed and normalized
-        if "dtype" in normalized_kwargs:
-            stage_engine_args["dtype"] = normalized_kwargs["dtype"]
-
-        # New split fields for diffusers adapter kwargs.
-        if kwargs.get("diffusers_load_kwargs") is not None:
-            stage_engine_args["diffusers_load_kwargs"] = kwargs["diffusers_load_kwargs"]
-        if kwargs.get("diffusers_call_kwargs") is not None:
-            stage_engine_args["diffusers_call_kwargs"] = kwargs["diffusers_call_kwargs"]
-
-        default_stage_cfg = [
-            {
-                "stage_id": 0,
-                "stage_type": "diffusion",
-                "runtime": {
-                    "process": True,
-                    "devices": devices,
-                },
-                "engine_args": stage_engine_args,
-                "engine_input_source": [],
-                "default_sampling_params": stage_default_sampling_params,
-                "final_output": True,
-                "final_output_type": final_output_type,
-            }
-        ]
-        default_stage_cfg[0]["engine_args"]["model_stage"] = "diffusion"
-        return default_stage_cfg
-
     def _apply_strategy_lb_policy(self, derived: str | None, kwargs: dict[str, Any]) -> None:
         """Apply a strategy-derived ``omni_lb_policy`` to the engine.
 
@@ -1251,6 +989,29 @@ class AsyncOmniEngine:
             )
             self._omni_lb_policy = str(derived)
 
+    @staticmethod
+    def _create_default_diffusion_stage_cfg(kwargs: dict[str, Any]) -> list[dict[str, Any]]:
+        """Compatibility seam for the factory-owned diffusion fallback."""
+        return StageConfigFactory.create_default_diffusion(kwargs)
+
+    def _set_pipeline_runtime_config(
+        self,
+        pipeline_config: PipelineConfig | None,
+        config_path: str | None,
+    ) -> None:
+        """Initialize engine-wide settings resolved from pipeline metadata."""
+        self.endpoint_restrictions = pipeline_config.endpoint_restrictions if pipeline_config is not None else ()
+        self._duplex_runtime_extension_path = (
+            pipeline_config.duplex_runtime_extension if pipeline_config is not None else None
+        )
+        self.duplex_serving_adapter_path = (
+            pipeline_config.duplex_serving_adapter if pipeline_config is not None else None
+        )
+        self._duplex_control_enabled = bool(pipeline_config and pipeline_config.duplex_control_enabled)
+        self.duplex_session_config = DuplexSessionRuntimeConfig()
+        if config_path is not None:
+            self.duplex_session_config = load_deploy_config(config_path).duplex_session
+
     def _resolve_stage_configs(
         self,
         model: str,
@@ -1266,159 +1027,43 @@ class AsyncOmniEngine:
 
         deploy_config_path = kwargs.pop("deploy_config", None)
         strategy_config_path = kwargs.pop("strategy_config", None)
-        stage_overrides_json = kwargs.pop("stage_overrides", None)
+        # CLI callers arrive pre-parsed; offline Python callers may use the
+        # JSON-string form documented in recipes.
+        stage_overrides = parse_stage_overrides(kwargs.pop("stage_overrides", None))
 
         # ``diffusion_streaming_output`` is the public AsyncOmni/serve kwarg;
-        # stage configs know the field as ``streaming_output``. The unregistered
-        # single-stage fallback translates it in
-        # ``_create_default_diffusion_stage_cfg``, but a registered pipeline
-        # resolves through StageConfigFactory, which only passes through keys the
-        # stage schema recognizes — so mirror the translation here. Only a truthy
-        # value is mirrored: the serve CLI always carries the flag's ``False``
-        # default, which must not override a deploy YAML's ``streaming_output``.
+        # stage configs know the field as ``streaming_output``. Mirror only a
+        # truthy value so the CLI's False default does not override deploy YAML.
         if kwargs.get("diffusion_streaming_output") and kwargs.get("streaming_output") is None:
             kwargs["streaming_output"] = True
 
-        # Parse --stage-overrides JSON string if provided
-        stage_overrides = parse_stage_overrides(stage_overrides_json)
-
-        # Unregistered diffusion checkpoints use the single-stage fallback
-        # below instead of StageConfigFactory, so fold stage-0 model extras
-        # into the fallback's input as well.  Registered pipelines still get
-        # the complete override mapping through load_and_resolve_stage_configs.
-        default_stage_kwargs = kwargs
-        stage_zero_overrides = (stage_overrides or {}).get("0", {})
-        if "extras" in stage_zero_overrides:
-            override_extras = stage_zero_overrides["extras"]
-            if not isinstance(override_extras, Mapping):
-                raise TypeError("stage 0 extras must be a mapping")
-            default_stage_kwargs = {
-                **kwargs,
-                "extras": {
-                    **(kwargs.get("extras") or {}),
-                    **override_extras,
-                },
-            }
-
-        def create_default_stage_config() -> list:
-            fallback_kwargs = dict(default_stage_kwargs)
-            if not fallback_kwargs.get("model_class_name"):
-                model_class_name = resolve_model_class_name(
-                    model,
-                    fallback_kwargs.get("diffusion_load_format", "default"),
-                )
-                if model_class_name is not None:
-                    fallback_kwargs["model_class_name"] = model_class_name
-            return self._create_default_diffusion_stage_cfg(fallback_kwargs)
-
-        config_path, stage_configs, strategy_lb_policy = load_and_resolve_stage_configs(
-            model,
-            kwargs,
-            trust_remote_code=trust_remote_code,
-            default_stage_cfg_factory=create_default_stage_config,
-            deploy_config_path=deploy_config_path,
-            stage_overrides=stage_overrides,
-            strategy_config_path=strategy_config_path,
+        resolution = cast(
+            _ConfigResolutionResult,
+            load_and_resolve_stage_configs(
+                model,
+                kwargs,
+                trust_remote_code=trust_remote_code,
+                deploy_config_path=deploy_config_path,
+                stage_overrides=stage_overrides,
+                strategy_config_path=strategy_config_path,
+            ),
         )
+        if isinstance(resolution, OmniConfigResolution):
+            self._config_resolution = resolution
+            config_path = resolution.config_path
+            stage_configs = list(resolution.stage_configs)
+            strategy_lb_policy = resolution.omni_lb_policy
+        else:
+            # Compatibility for overrides of the historical tuple-returning
+            # seam. Production always receives OmniConfigResolution above.
+            config_path, stage_configs, strategy_lb_policy = resolution
 
         # A strategy.yaml may derive a pipeline-wide load-balancer policy. It is
         # an orchestrator-level knob (read once at construction), so apply it here
         # rather than as a per-stage config field.
         self._apply_strategy_lb_policy(strategy_lb_policy, kwargs)
 
-        # Inject diffusion LoRA-related knobs from kwargs if not present in the stage config.
-        for cfg in stage_configs:
-            try:
-                if not hasattr(cfg, "engine_args") or cfg.engine_args is None:
-                    cfg.engine_args = OmegaConf.create({})
-                global_sleep_mode = kwargs.get("enable_sleep_mode")
-                if global_sleep_mode is not None:
-                    if not hasattr(cfg.engine_args, "enable_sleep_mode") or cfg.engine_args.enable_sleep_mode is None:
-                        cfg.engine_args.enable_sleep_mode = global_sleep_mode
-                if getattr(cfg, "stage_type", None) != "diffusion":
-                    continue
-                if not hasattr(cfg, "engine_args") or cfg.engine_args is None:
-                    cfg.engine_args = OmegaConf.create({})
-                additional_config = kwargs.get("additional_config")
-                if additional_config is not None:
-                    current_additional_config = getattr(cfg.engine_args, "additional_config", None)
-                    if current_additional_config in (None, {}):
-                        cfg.engine_args.additional_config = additional_config
-                if kwargs.get("lora_path") is not None:
-                    if not hasattr(cfg.engine_args, "lora_path") or cfg.engine_args.lora_path is None:
-                        cfg.engine_args.lora_path = kwargs["lora_path"]
-                lora_scale = kwargs.get("lora_scale")
-                if lora_scale is None:
-                    # Backwards compatibility for older callers.
-                    lora_scale = kwargs.get("static_lora_scale")
-                if lora_scale is not None:
-                    if not hasattr(cfg.engine_args, "lora_scale") or cfg.engine_args.lora_scale is None:
-                        cfg.engine_args.lora_scale = lora_scale
-                if kwargs.get("lora_backend") is not None:
-                    if not hasattr(cfg.engine_args, "lora_backend") or cfg.engine_args.lora_backend is None:
-                        cfg.engine_args.lora_backend = kwargs["lora_backend"]
-                if (
-                    kwargs.get("diffusion_attention_config") is not None
-                    or kwargs.get("diffusion_attention_backend") is not None
-                    or kwargs.get("fastvideo_vsa_topk") is not None
-                ):
-                    has_stage_attention = (
-                        getattr(cfg.engine_args, "diffusion_attention_config", None) is not None
-                        or getattr(cfg.engine_args, "diffusion_attention_backend", None) is not None
-                    )
-                    if not has_stage_attention:
-                        cfg.engine_args.diffusion_attention_config = parse_attention_config(
-                            kwargs.get("diffusion_attention_config"),
-                            attention_backend=kwargs.get("diffusion_attention_backend"),
-                            fastvideo_vsa_topk=kwargs.get("fastvideo_vsa_topk"),
-                        )
-                quantization_config = kwargs.get("diffusion_quantization_config") or kwargs.get("quantization_config")
-                if quantization_config is not None:
-                    if (
-                        not hasattr(cfg.engine_args, "quantization_config")
-                        or cfg.engine_args.quantization_config is None
-                    ):
-                        cfg.engine_args.quantization_config = quantization_config
-                # Inject profiler flags for diffusion stages
-                for profiler_key in (
-                    "enable_diffusion_pipeline_profiler",
-                    "enable_ar_profiler",
-                ):
-                    val = kwargs.get(profiler_key)
-                    if val:
-                        if not hasattr(cfg.engine_args, profiler_key) or not getattr(
-                            cfg.engine_args, profiler_key, False
-                        ):
-                            setattr(cfg.engine_args, profiler_key, val)
-                quantization = kwargs.get("quantization")
-                if quantization is not None:
-                    if not hasattr(cfg.engine_args, "quantization") or cfg.engine_args.quantization is None:
-                        cfg.engine_args.quantization = quantization
-                diffusion_kv_cache_dtype = kwargs.get("diffusion_kv_cache_dtype")
-                if diffusion_kv_cache_dtype is not None:
-                    if (
-                        not hasattr(cfg.engine_args, "diffusion_kv_cache_dtype")
-                        or cfg.engine_args.diffusion_kv_cache_dtype is None
-                    ):
-                        cfg.engine_args.diffusion_kv_cache_dtype = diffusion_kv_cache_dtype
-                diffusion_kv_cache_skip_steps = kwargs.get("diffusion_kv_cache_skip_steps")
-                if diffusion_kv_cache_skip_steps is not None:
-                    if (
-                        not hasattr(cfg.engine_args, "diffusion_kv_cache_skip_steps")
-                        or cfg.engine_args.diffusion_kv_cache_skip_steps is None
-                    ):
-                        cfg.engine_args.diffusion_kv_cache_skip_steps = diffusion_kv_cache_skip_steps
-                diffusion_kv_cache_skip_layers = kwargs.get("diffusion_kv_cache_skip_layers")
-                if diffusion_kv_cache_skip_layers is not None:
-                    if (
-                        not hasattr(cfg.engine_args, "diffusion_kv_cache_skip_layers")
-                        or cfg.engine_args.diffusion_kv_cache_skip_layers is None
-                    ):
-                        cfg.engine_args.diffusion_kv_cache_skip_layers = diffusion_kv_cache_skip_layers
-            except Exception as e:
-                logger.warning("Failed to inject LoRA config for stage: %s", e)
-
-        return config_path, stage_configs
+        return cast(str, config_path), stage_configs
 
     # ==================== Public API ====================
 
@@ -1944,7 +1589,7 @@ class AsyncOmniEngine:
             tokens generated before abort (empty for diffusion / no OP state).
         """
         if not request_ids or getattr(self, "_shutdown_called", False):
-            return
+            return []
         if self.request_queue is None:
             raise RuntimeError("request_queue is not initialized")
         transport = self._correlated_rpc_client
@@ -1971,7 +1616,7 @@ class AsyncOmniEngine:
             result_msg = await loop.run_in_executor(None, _wait)
         except Exception as exc:
             if getattr(self, "_shutdown_called", False) and is_abort_transport_shutdown(exc):
-                return
+                return []
             raise
         if not result_msg.success:
             raise RuntimeError(result_msg.error or "abort failed")

vllm_omni/engine/stage_init_utils.py

@@ -53,7 +53,7 @@ from vllm_omni.config.stage_config import StageType
 from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.engine.arg_utils import OmniEngineArgs
 from vllm_omni.entrypoints.stage_utils import _to_dict, set_stage_devices
-from vllm_omni.entrypoints.utils import filter_dataclass_kwargs, resolve_model_config_path
+from vllm_omni.entrypoints.utils import filter_dataclass_kwargs
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniSamplingParams
 from vllm_omni.inputs.preprocess import OmniInputPreprocessor
 from vllm_omni.outputs.output_processor import MultimodalOutputProcessor
@@ -1811,7 +1811,7 @@ def release_device_locks(lock_fds: list[int]) -> None:
 
 
 def load_omni_transfer_config_for_model(model: str, config_path: str | None) -> Any:
-    """Load omni transfer config from an explicit path or resolved model config.
+    """Load omni transfer config from the resolver-selected deploy config.
 
     Resolves ``base_config`` inheritance (CI overlay → base deploy YAML) so
     that connectors defined in the base config are visible to the transfer
@@ -1820,12 +1820,11 @@ def load_omni_transfer_config_for_model(model: str, config_path: str | None) ->
     from vllm_omni.distributed.omni_connectors import load_omni_transfer_config
 
     try:
-        resolved_config_path = config_path or resolve_model_config_path(model)
-        if resolved_config_path is None:
+        if config_path is None:
             return None
         from vllm_omni.config.stage_config import resolve_deploy_yaml
 
-        resolved_dict = resolve_deploy_yaml(resolved_config_path)
+        resolved_dict = resolve_deploy_yaml(config_path)
         return load_omni_transfer_config(config_dict=resolved_dict)
     except Exception as e:
         logger.warning("[stage_init] Failed to load transfer config: %s", e)

vllm_omni/entrypoints/cli/serve.py

@@ -14,6 +14,7 @@ import math
 import os
 import signal
 from types import FrameType
+from typing import Any, cast
 
 import uvloop
 from vllm.entrypoints.cli.types import CLISubcommand
@@ -23,6 +24,7 @@ from vllm.logger import init_logger
 
 from vllm_omni.entrypoints.cli.logo import log_logo
 from vllm_omni.entrypoints.openai.api_server import omni_run_server
+from vllm_omni.entrypoints.utils import parse_stage_overrides
 from vllm_omni.utils.tracking_parser import TrackingArgumentParser, TrackingNamespace
 
 logger = init_logger(__name__)
@@ -47,6 +49,17 @@ Search by using: `--help=<ConfigGroup>` to explore options by section (e.g.,
 """
 
 
+def _parse_stage_overrides(value: str) -> dict[str, dict[str, Any]]:
+    """Adapt shared stage-override validation to argparse's error type."""
+    try:
+        parsed = parse_stage_overrides(value)
+    except ValueError as exc:
+        raise argparse.ArgumentTypeError(str(exc)) from exc
+    if parsed is None:
+        raise argparse.ArgumentTypeError("--stage-overrides requires a JSON object")
+    return parsed
+
+
 def _nonneg_finite_float(value: str) -> float:
     """Argparse type for finite, non-negative floats (rejects nan/inf)."""
     try:
@@ -286,7 +299,7 @@ class OmniServeCommand(CLISubcommand):
         )
         omni_config_group.add_argument(
             "--stage-overrides",
-            type=str,
+            type=_parse_stage_overrides,
             default=None,
             help="Per-stage JSON overrides. Example: "
             '\'{"0": {"gpu_memory_utilization": 0.8}, "2": {"enforce_eager": true}}\'',
@@ -945,6 +958,7 @@ def run_headless(args: TrackingNamespace) -> None:
     from vllm.v1.executor.multiproc_executor import MultiprocExecutor
     from vllm.version import __version__ as VLLM_VERSION
 
+    from vllm_omni.config.resolver import resolve_omni_config
     from vllm_omni.distributed.omni_connectors.utils.initialization import resolve_omni_kv_config_for_stage
     from vllm_omni.engine.stage_engine_startup import (
         get_headless_replica_devices,
@@ -959,10 +973,6 @@ def run_headless(args: TrackingNamespace) -> None:
         load_omni_transfer_config_for_model,
         prepare_engine_environment,
     )
-    from vllm_omni.entrypoints.utils import (
-        load_and_resolve_stage_configs,
-        parse_stage_overrides,
-    )
 
     model = args.model
     stage_id: int | None = args.stage_id
@@ -984,6 +994,10 @@ def run_headless(args: TrackingNamespace) -> None:
     # Filter down to a dict of things explicitly requested by the user
     args_dict = args.get_explicit_kwargs_dict()
 
+    deploy_config_path = args_dict.pop("deploy_config", None)
+    strategy_config_path = args_dict.pop("strategy_config", None)
+    stage_overrides = args_dict.pop("stage_overrides", None)
+
     # ``--replica-id`` is deprecated and ignored — replica ids are
     # auto-assigned by ``OmniMasterServer`` so headless processes carry
     # no knowledge of their per-replica id at launch time. Warn (don't
@@ -996,33 +1010,27 @@ def run_headless(args: TrackingNamespace) -> None:
             "master server.",
             args.replica_id,
         )
+        args_dict.pop("replica_id")
 
-    # Parse --stage-overrides (raw JSON string) exactly like the standard
-    # engine path (AsyncOmniEngine._resolve_stage_configs) so headless and
-    # standard launches resolve to the same per-stage device layout.
-    stage_overrides = parse_stage_overrides(args_dict.get("stage_overrides"))
-
-    config_path, stage_configs, _ = load_and_resolve_stage_configs(
+    resolved = resolve_omni_config(
         model,
-        args_dict,
         # store_true cannot express an explicit False: absent maps to None
         # ("not specified") so the deploy yaml's per-stage value applies.
         trust_remote_code=getattr(args, "trust_remote_code", None) or None,
-        deploy_config_path=args_dict.get("deploy_config"),
+        cli_overrides=args_dict,
+        deploy_config_path=deploy_config_path,
         stage_overrides=stage_overrides,
-        strategy_config_path=args_dict.get("strategy_config"),
+        strategy_config_path=strategy_config_path,
     )
+    config_path = resolved.config_path
+    stage_configs = list(resolved.stage_configs)
 
-    # Locate the stage config that matches stage_id.
-    stage_cfg = None
-    for cfg in stage_configs:
-        if cfg.stage_id == stage_id:
-            stage_cfg = cfg
-            break
-    if stage_cfg is None:
+    try:
+        stage_cfg = resolved.stage_by_id(stage_id)
+    except KeyError:
         raise ValueError(
             f"No stage config found for stage_id={stage_id}. Available stage ids: {[c.stage_id for c in stage_configs]}"
-        )
+        ) from None
 
     prepare_engine_environment()
     per_replica_devices = get_headless_replica_devices(stage_cfg, stage_id, omni_dp_size_local)
@@ -1037,7 +1045,7 @@ def run_headless(args: TrackingNamespace) -> None:
             omni_master_port=omni_master_port,
             omni_dp_size_local=omni_dp_size_local,
             per_replica_devices=per_replica_devices,
-            config_path=config_path,
+            config_path=cast(str, config_path),
             replica_bind_address=omni_replica_address,
         )
         return

vllm_omni/entrypoints/utils.py

@@ -4,45 +4,18 @@
 import json
 import os
 import types
-from collections import Counter
 from dataclasses import fields, is_dataclass
-from pathlib import Path
 from typing import Any, get_args, get_origin
 
 from vllm.logger import init_logger
 from vllm.sampling_params import RequestOutputKind, SamplingParams
-from vllm.transformers_utils.config import get_config, get_hf_file_to_dict
-from vllm.transformers_utils.repo_utils import file_or_path_exists
 
-from vllm_omni.config.config_factory import (
-    StageConfigFactory,
-    _materialize_object_storage_configs,
-    _name_match_candidate,
-    with_trust_remote_code_override,
-)
-from vllm_omni.config.pipeline_registry import OMNI_PIPELINES
-from vllm_omni.config.stage_config import _DEPLOY_DIR
-from vllm_omni.config.yaml_util import create_config, load_yaml_config
-from vllm_omni.diffusion.utils.hf_utils import (
-    _looks_like_dreamzero,
-    get_diffusion_model_index,
-    resolve_native_diffusion_model_class,
-)
 from vllm_omni.entrypoints.stage_utils import _to_dict
 from vllm_omni.inputs.data import OmniSamplingParams
-from vllm_omni.platforms import current_omni_platform
-
-# Get the project root directory (2 levels up from this file)
-PROJECT_ROOT = Path(__file__).parent.parent.parent
 
 logger = init_logger(__name__)
 
 
-_DIFFUSERS_CLASS_TO_CONFIG: dict[str, str] = {
-    "GlmImagePipeline": "glm_image",
-}
-
-
 def inject_omni_kv_config(stage: Any, omni_conn_cfg: dict[str, Any], omni_from: str, omni_to: str) -> None:
     """Inject connector configuration into stage engine arguments."""
     # Prepare omni_kv_config dict
@@ -76,501 +49,9 @@ def inject_omni_kv_config(stage: Any, omni_conn_cfg: dict[str, Any], omni_from:
         logger.error(f"Failed to inject omni connector config into stage: {e}")
 
 
-def _try_get_class_name_from_diffusers_config(model: str) -> str | None:
-    """Try to get class name from diffusers model configuration files.
-
-    Args:
-        model: Model name or path
-
-    Returns:
-        Model type string if found, None otherwise
-    """
-    model_index = get_diffusion_model_index(model)
-    if model_index and isinstance(model_index, dict) and "_class_name" in model_index:
-        logger.debug(f"Found Diffusers model type '{model_index['_class_name']}'")
-        return model_index["_class_name"]
-
-    return None
-
-
-def _filter_dict_like_object(obj: dict | Any) -> dict:
-    """Filter dict-like object by removing callables and recursively converting values.
-
-    Converts dict-like objects to regular dicts while filtering out callable values
-    that are incompatible with OmegaConf. Recursively processes values through
-    _convert_dataclasses_to_dict for nested object conversion.
-
-    Args:
-        obj: Dict or dict-like object to filter
-
-    Returns:
-        Regular dict with callables filtered out and values recursively converted
-
-    Raises:
-        TypeError: If obj doesn't support .items() method
-        ValueError: If dict conversion fails unexpectedly
-    """
-
-    def _is_callable_value(value: Any) -> bool:
-        if callable(value):
-            return True
-        return isinstance(
-            value,
-            types.FunctionType | types.MethodType | types.BuiltinFunctionType | types.BuiltinMethodType,
-        )
-
-    result = {}
-    filtered_keys = []
-    for k, v in obj.items():
-        # Preserve class objects by converting to a fully qualified name string
-        # so callers that resolve via import path (e.g. custom_pipeline_args.pipeline_class)
-        # still work after OmegaConf round-trip.
-        if isinstance(v, type):
-            module = getattr(v, "__module__", None)
-            qualname = getattr(v, "__qualname__", getattr(v, "__name__", None))
-            if module and qualname and module != "builtins":
-                result[k] = f"{module}.{qualname}"
-            else:
-                result[k] = qualname
-        elif _is_callable_value(v):
-            filtered_keys.append(str(k))
-        else:
-            result[k] = _convert_dataclasses_to_dict(v)
-    if filtered_keys:
-        logger.warning(
-            f"Filtered out {len(filtered_keys)} callable object(s) from base_engine_args "
-            f"that are not compatible with OmegaConf: {filtered_keys}. "
-        )
-    return result
-
-
-def _convert_dataclasses_to_dict(obj: Any) -> Any:
-    """Recursively convert non-serializable objects to OmegaConf-compatible types.
-
-    This is needed because OmegaConf cannot handle:
-    - Dataclass objects with Literal type annotations (e.g., StructuredOutputsConfig)
-    - Counter objects (from collections or vllm.utils)
-    - Set objects
-    - Callable objects (functions, methods, etc.)
-    - Other non-primitive types
-    """
-    # IMPORTANT: Check Counter BEFORE dict, since Counter is a subclass of dict
-    # Handle Counter objects (convert to dict)
-    # Check by class name first to catch both collections.Counter and vllm.utils.Counter
-    if hasattr(obj, "__class__") and obj.__class__.__name__ == "Counter":
-        try:
-            return dict(obj)
-        except (TypeError, ValueError):
-            # If Counter can't be converted to dict, return empty dict
-            return {}
-    # Also check isinstance for collections.Counter (must be before dict check)
-    if isinstance(obj, Counter):
-        return dict(obj)
-    # Handle set objects (convert to list)
-    if isinstance(obj, set):
-        return list(obj)
-    # Handle dataclass objects
-    # Use field iteration instead of asdict() to:
-    # 1. Only include init fields (non-init fields cause "unexpected kwarg" errors)
-    # 2. Skip None values matching field defaults (avoids Pydantic validation
-    #    failures when None is explicitly passed for non-Optional typed fields,
-    #    e.g. CompilationConfig.cudagraph_capture_sizes: list[int] = None)
-    if is_dataclass(obj) and not isinstance(obj, type):
-        result = {}
-        for f in fields(obj):
-            if not f.init:
-                continue
-            value = getattr(obj, f.name)
-            if value is None and f.default is None:
-                continue
-            result[f.name] = _convert_dataclasses_to_dict(value)
-        return result
-    # Handle dictionaries (recurse into values) and filter out callables(cause error in OmegaConf.create)
-    # Note: This must come AFTER Counter check since Counter is a dict subclass
-    if isinstance(obj, dict):
-        return _filter_dict_like_object(obj)
-    # Preserve class objects by converting to a fully qualified name string.
-    if isinstance(obj, type):
-        module = getattr(obj, "__module__", None)
-        qualname = getattr(obj, "__qualname__", getattr(obj, "__name__", None))
-        if module and qualname and module != "builtins":
-            return f"{module}.{qualname}"
-        return qualname
-    # Handle callable objects (functions, methods, etc.) - skip them
-    # Note: This comes after dict/list checks to avoid misclassifying dict-like objects
-    if callable(obj):
-        return None
-    # Handle lists and tuples (recurse into items)
-    if isinstance(obj, list | tuple):
-        return type(obj)(_convert_dataclasses_to_dict(item) for item in obj if not callable(item))
-    # Try to convert any dict-like object (has keys/values methods) to dict
-    if hasattr(obj, "keys") and hasattr(obj, "values") and not isinstance(obj, str | bytes):
-        try:
-            return _filter_dict_like_object(obj)
-        except (TypeError, ValueError, AttributeError):
-            # If conversion fails, return as-is
-            return obj
-    # Primitive types and other objects that OmegaConf can handle
-    return obj
-
-
-def _try_resolve_omni_model_type(model: str) -> str | None:
-    """Try to resolve model_type for omni models with empty config.json.
-
-    Searches both the legacy ``stage_configs/*.yaml`` directory and the
-    migrated ``deploy/*.yaml`` directory for a stem that substring-matches
-    the model path basename (e.g. ``cosyvoice3`` in
-    ``FunAudioLLM/Fun-CosyVoice3-0.5B-2512``). Only the basename is scanned
-    so URI segments such as the bucket name cannot select an unrelated
-    pipeline. The longest match wins so ``cosyvoice3`` beats ``cosyvoice``
-    and ``bagel_single_stage`` beats ``bagel``.
-    """
-    model_lower = _name_match_candidate(model).lower().replace("-", "").replace("_", "")
-    best_match: str | None = None
-    best_len = 0
-    for subdir in ("model_executor/stage_configs", "deploy"):
-        config_dir = PROJECT_ROOT / "vllm_omni" / subdir
-        if not config_dir.exists():
-            continue
-        for config_file in sorted(config_dir.glob("*.yaml")):
-            candidate = config_file.stem.replace("-", "").replace("_", "")
-            if candidate and candidate in model_lower and len(candidate) > best_len:
-                best_match = config_file.stem
-                best_len = len(candidate)
-    return best_match
-
-
-def _registry_default_deploy_path(model: str) -> str | None:
-    """Return the registered pipeline's default deploy YAML path, if any.
-
-    A checkpoint's HF ``model_type`` is not always the registered Omni pipeline
-    key. MiniCPM-o 4.5, for example, reports ``minicpmo`` while its
-    predicate-selected pipeline is ``minicpmo_4_5``. Stage construction already
-    loads that pipeline's default deploy config, so connector discovery must
-    resolve the same path.
-    """
-    pipeline_config = StageConfigFactory.get_pipeline_config(
-        model=model,
-        trust_remote_code=True,
-    )
-    if pipeline_config is None or pipeline_config.default_deploy_config_name is None:
-        return None
-    default_deploy_path = _DEPLOY_DIR / pipeline_config.default_deploy_config_name
-    if default_deploy_path.is_file():
-        return str(default_deploy_path)
-    return None
-
-
-def resolve_model_config_path(model: str) -> str | None:
-    """Resolve the stage/deploy config file path from the model name.
-
-    Resolves configuration path based on the model type and device type.
-    Order:
-    1. Device-specific stage config under ``stage_configs/{device_type}/``
-    2. If HF ``model_type`` is not an ``OMNI_PIPELINES`` key, the registered
-       pipeline's ``default_deploy_config_name`` (keeps connectors aligned with
-       stage construction when HF type and pipeline key differ)
-    3. ``deploy/{model_type}.yaml``
-    4. Legacy ``stage_configs/{model_type}.yaml``
-    5. Registered pipeline default deploy config (final fallthrough)
-
-    Args:
-        model: Model name or path (used to determine model_type)
-
-    Returns:
-        String path to the stage/deploy configuration file, or ``None`` if no
-        matching config file exists.
-
-    Raises:
-        ValueError: If model_type cannot be determined
-    """
-    # Object-storage URIs are streamed by vLLM's Run:AI streamer only after
-    # each stage builds its ModelConfig, so config resolution here reads the
-    # local copy that vLLM-Omni materializes for such URIs. Name-based
-    # fallbacks keep the original string (the materialized path is a hash).
-    config_source = _materialize_object_storage_configs(model)
-    # Try to get config from standard transformers format first
-    try:
-        hf_config = get_config(config_source, trust_remote_code=True)
-        model_type = hf_config.model_type
-    except (ValueError, Exception):
-        # If standard transformers format fails, try diffusers format
-        native_model_class = resolve_native_diffusion_model_class(config_source)
-        if native_model_class is not None:
-            model_type = native_model_class
-        elif get_diffusion_model_index(config_source) is not None:
-            model_type = _try_get_class_name_from_diffusers_config(config_source)
-            if model_type is None:
-                raise ValueError(
-                    f"Could not determine model_type for diffusers model: {model}. "
-                    "Please ensure its Diffusers pipeline index contains '_class_name'"
-                )
-        elif file_or_path_exists(config_source, "config.json", revision=None):
-            # Try to read config.json manually for custom models like Bagel that fail get_config
-            # but have a valid config.json with model_type
-            try:
-                config_dict = get_hf_file_to_dict("config.json", config_source, revision=None)
-                if config_dict and "model_type" in config_dict:
-                    model_type = config_dict["model_type"]
-                else:
-                    # For models with empty config.json (e.g. CosyVoice3),
-                    # try matching against registered omni stage configs.
-                    model_type = _try_resolve_omni_model_type(model)
-                    if model_type is None:
-                        raise ValueError(f"config.json found but missing 'model_type' for model: {model}")
-            except Exception as e:
-                raise ValueError(f"Failed to read config.json for model: {model}. Error: {e}") from e
-        else:
-            # No config.json at repo root (e.g. GLM-TTS stores configs in
-            # subdirectories only).  Try matching against registered deploy
-            # YAML filenames before giving up.
-            model_type = _try_resolve_omni_model_type(model)
-            if model_type is None:
-                raise ValueError(
-                    f"Could not determine model_type for model: {model}. "
-                    "Model is not in standard transformers or Diffusers format. "
-                    f"Please ensure the model has proper configuration files with 'model_type' field"
-                )
-
-    default_config_path = current_omni_platform.get_default_stage_config_path()
-    if model_type == "vla" and _looks_like_dreamzero(config_source):
-        model_type = "dreamzero"
-
-    if model_type in _DIFFUSERS_CLASS_TO_CONFIG:
-        normalized_model_type = _DIFFUSERS_CLASS_TO_CONFIG[model_type]
-    else:
-        normalized_model_type = model_type.replace("-", "_")
-    model_type_str = f"{normalized_model_type}.yaml"
-    complete_config_path = PROJECT_ROOT / default_config_path / model_type_str
-    if os.path.exists(complete_config_path):
-        return str(complete_config_path)
-
-    # Prefer the registry default before deploy/<hf_model_type>.yaml when the
-    # HF type is not itself a pipeline key. Otherwise a future
-    # deploy/minicpmo.yaml would desync connectors from stages for MiniCPM-o 4.5.
-    if normalized_model_type not in OMNI_PIPELINES:
-        registry_path = _registry_default_deploy_path(model)
-        if registry_path is not None:
-            return registry_path
-
-    deploy_config_path = _DEPLOY_DIR / model_type_str
-    if os.path.exists(deploy_config_path):
-        return str(deploy_config_path)
-
-    stage_config_file = f"vllm_omni/model_executor/stage_configs/{normalized_model_type}.yaml"
-    stage_config_path = PROJECT_ROOT / stage_config_file
-    if os.path.exists(stage_config_path):
-        return str(stage_config_path)
-
-    return _registry_default_deploy_path(model)
-
-
-def load_stage_configs_from_model(
-    model: str,
-    *,
-    trust_remote_code: bool | None,
-    base_engine_args: dict | None = None,
-    deploy_config_path: str | None = None,
-    stage_overrides: dict[str, dict[str, Any]] | None = None,
-    strategy_config_path: str | None = None,
-) -> tuple[list, str | None]:
-    """Load stage configurations from model's default config file.
-
-    For models registered in the pipeline registry, uses
-    ``StageConfigFactory.create_legacy_stage_configs_from_model()`` which merges
-    PipelineConfig + DeployConfig + CLI overrides.
-
-    Models that cannot be resolved through the registry return no stages so the
-    caller can apply its default stage configuration, when available.
-
-    Args:
-        model: Model name or path (used to determine model_type)
-        trust_remote_code: Whether to trust remote code while resolving the model config.
-        base_engine_args: Base engine args to merge as CLI overrides.
-        deploy_config_path: Optional explicit deploy config path.
-        stage_overrides: Per-stage overrides from --stage-overrides.
-        strategy_config_path: Optional path to a composable-parallel
-            ``strategy.yaml`` whose derived sizing is overlaid onto the
-            registry-merged stages.
-
-    Returns:
-        ``(stage_configs, omni_lb_policy)``: the list of stage configuration
-        dictionaries plus the strategy-derived pipeline-wide ``omni_lb_policy``
-        (``None`` when no strategy set one). The policy is returned rather than
-        written into a caller-provided mutable dict.
-    """
-    if base_engine_args is None:
-        base_engine_args = {}
-
-    cli_overrides = _convert_dataclasses_to_dict(dict(base_engine_args))
-    # A False inherited from the engine-args dump is the store_true flag's
-    # default, not an explicit choice — drop it so only the tri-state
-    # parameter below decides (see with_trust_remote_code_override).
-    if not cli_overrides.get("trust_remote_code"):
-        cli_overrides.pop("trust_remote_code", None)
-    cli_overrides = with_trust_remote_code_override(cli_overrides, trust_remote_code)
-    if stage_overrides:
-        for stage_id_str, overrides in stage_overrides.items():
-            for key, val in overrides.items():
-                cli_overrides[f"stage_{stage_id_str}_{key}"] = val
-
-    # Current runtime initialization still consumes legacy OmegaConf stage
-    # configs. ``StageConfigFactory.create_from_model`` now produces the
-    # structured ``VllmOmniConfig`` object; future RFC #4021 changes will
-    # migrate the engine/runtime consumers and replace this legacy resolver.
-    strategy_specs = None
-    if strategy_config_path is not None:
-        from vllm_omni.config.composable_parallel.strategy_loader import load_strategy_specs
-
-        strategy_specs = load_strategy_specs(strategy_config_path)
-
-    stages, omni_lb_policy = StageConfigFactory.create_legacy_stage_configs_from_model(
-        model,
-        trust_remote_code=trust_remote_code,
-        cli_overrides=cli_overrides,
-        deploy_config_path=deploy_config_path,
-        strategy_specs=strategy_specs,
-    )
-    if stages is not None:
-        # Convert StageConfig objects to OmegaConf for backward compat
-        return [stage.to_omegaconf() for stage in stages], omni_lb_policy
-
-    strategy_note = ""
-    if strategy_config_path is not None:
-        strategy_note = f" Strategy config {strategy_config_path!r} was not applied."
-    logger.warning(
-        "No registered PipelineConfig resolved for model %r. Legacy `stage_args` "
-        "YAMLs are no longer supported; register the pipeline and provide deployment "
-        "overrides through `deploy_config`.%s",
-        model,
-        strategy_note,
-    )
-    return [], None
-
-
-def filter_stages(
-    config_path: str | None,
-    stage_configs: list,
-    kwargs: dict | None,
-) -> list:
-    """Filter stage configs by mode when YAML defines a `modes` section.
-
-    The YAML can define, e.g.:
-
-        modes:
-          - mode: text-to-image
-            stages: [1]
-          - mode: image-to-text
-            stages: [0]
-
-    When users pass `mode="image-to-text"` into Omni(**kwargs), only the stages
-    listed for that mode are returned. If no mode is provided, defaults to
-    "text-to-image". If no modes are defined or filtering fails, returns the
-    original stage_configs unchanged.
-
-    Args:
-        config_path: Path to the YAML config (used to read `modes`).
-        stage_configs: Loaded list of stage configs.
-        kwargs: Engine/caller kwargs; may contain "mode".
-
-    Returns:
-        Filtered list of stage configs (or original list if filtering not applied).
-    """
-    if not stage_configs or config_path is None:
-        return stage_configs
-
-    try:
-        cfg = load_yaml_config(config_path)
-        yaml_modes = getattr(cfg, "modes", None)
-        if yaml_modes is None:
-            return stage_configs
-
-        mode_to_stage_ids: dict[str, list[int]] = {}
-        if yaml_modes is not None:
-            for entry in yaml_modes:
-                mode_name = None
-                stages = None
-                if hasattr(entry, "mode") or hasattr(entry, "stages"):
-                    mode_name = getattr(entry, "mode", None)
-                    stages = getattr(entry, "stages", None)
-                elif isinstance(entry, dict):
-                    mode_name = entry.get("mode")
-                    stages = entry.get("stages")
-
-                if mode_name is None or stages is None:
-                    continue
-
-                if isinstance(stages, int):
-                    stage_list = [stages]
-                else:
-                    stage_list = list(stages)
-
-                mode_to_stage_ids[str(mode_name)] = [int(sid) for sid in stage_list]
-
-        # No modes section or empty mapping: use all stages and return early.
-        active_mode: str | None = None
-        if isinstance(kwargs, dict):
-            active_mode = kwargs.get("mode")
-
-        if active_mode is None:
-            active_mode = "text-to-image"
-
-        if active_mode not in mode_to_stage_ids:
-            logger.warning(
-                "Requested mode '%s' not found in config '%s'; available modes: %s. Using all stages.",
-                active_mode,
-                config_path,
-                sorted(mode_to_stage_ids.keys()),
-            )
-            return stage_configs
-
-        allowed_ids = set(mode_to_stage_ids[active_mode])
-        filtered_stage_configs = [sc for sc in stage_configs if getattr(sc, "stage_id", None) in allowed_ids]
-        if not filtered_stage_configs:
-            logger.warning(
-                "Mode '%s' in config '%s' resolved to stage ids %s, but none matched loaded stage_args. "
-                "Falling back to all stages.",
-                active_mode,
-                config_path,
-                sorted(allowed_ids),
-            )
-            return stage_configs
-
-        return filtered_stage_configs
-    except Exception as e:
-        logger.warning("Failed to apply mode-based stage filtering: %s", e)
-        return stage_configs
-
-
 def parse_stage_overrides(value: Any) -> dict[str, dict[str, Any]] | None:
-    """Parse the ``--stage-overrides`` value into a per-stage override dict.
-
-    ``value`` may be a raw JSON string (as supplied on the CLI) or an
-    already-parsed mapping. Returns ``None`` when no overrides are given.
-
-    Only the **shape** of the override mapping is validated here:
-
-    - top-level must be a dict
-    - keys must be non-negative ASCII integer strings (stage ids)
-    - values must be dicts
-
-    Field-name / type / range validation is intentionally not done at this
-    layer. The downstream resolver forwards every surviving key as
-    ``stage_<id>_<key>`` into ``cli_overrides`` (see
-    ``load_stage_configs_from_model``); field semantics live in
-    ``StageConfigFactory`` for registered pipelines, and the
-    default-diffusion fallback in ``async_omni_engine.py`` reads
-    ``stage_zero_overrides["extras"]`` for unregistered ones. Unknown
-    engine-arg keys that don't match ``OmniEngineArgs`` are dropped with
-    a warning at ``filter_dataclass_kwargs`` rather than failing here,
-    so a typo surfaces in the log instead of being silently lost.
-
-    Raises:
-        ValueError: when ``value`` is not a valid per-stage override
-            mapping (invalid JSON, non-dict top level, non-ASCII-digit
-            stage id, non-dict value).
-    """
-    if not value:
+    """Parse and validate the shape of per-stage JSON overrides."""
+    if value is None:
         return None
     if isinstance(value, str):
         try:
@@ -582,93 +63,22 @@ def parse_stage_overrides(value: Any) -> dict[str, dict[str, Any]] | None:
 
     if not isinstance(parsed, dict):
         raise ValueError(
-            f"--stage-overrides must be a JSON object mapping stage_id -> overrides, "
+            "--stage-overrides must be a JSON object mapping stage_id -> overrides, "
             f"got {type(parsed).__name__}: {parsed!r}"
         )
-    if not parsed:
-        return None
-
-    for stage_id_str, overrides in parsed.items():
-        # ``str.isdigit()`` accepts Unicode digit classes (e.g. fullwidth
-        # ``"０"``); downstream's ``_STAGE_OVERRIDE_PATTERN`` is ASCII-only
-        # (``^stage_(\d+)_(.+)$``), so a non-ASCII-digit key would parse
-        # here and silently misroute as a global key. ``isascii()`` closes
-        # that hole.
-        if not isinstance(stage_id_str, str) or not stage_id_str.isdigit() or not stage_id_str.isascii():
+    for stage_id, overrides in parsed.items():
+        if not isinstance(stage_id, str) or not stage_id.isascii() or not stage_id.isdigit():
             raise ValueError(
-                f"--stage-overrides keys must be non-negative integer stage ids (as strings), got {stage_id_str!r}"
+                f"--stage-overrides keys must be non-negative integer stage ids (as strings), got {stage_id!r}"
             )
         if not isinstance(overrides, dict):
             raise ValueError(
-                f"--stage-overrides[{stage_id_str!r}] must be an object, got {type(overrides).__name__}: {overrides!r}"
+                f"--stage-overrides[{stage_id!r}] must be an object, got {type(overrides).__name__}: {overrides!r}"
             )
 
     return parsed
 
 
-def load_and_resolve_stage_configs(
-    model: str,
-    kwargs: dict | None,
-    *,
-    trust_remote_code: bool | None,
-    default_stage_cfg_factory: Any = None,
-    deploy_config_path: str | None = None,
-    stage_overrides: dict[str, dict[str, Any]] | None = None,
-    strategy_config_path: str | None = None,
-) -> tuple[str, list, str | None]:
-    """Load stage configurations from a deploy YAML or model defaults.
-
-    Args:
-        model: Model name or path
-        kwargs: Engine arguments to merge with stage configs
-        trust_remote_code: Whether to trust remote code while resolving the model config.
-        default_stage_cfg_factory: Optional callable that takes no args and returns
-            default stage config list when no configs are found
-        deploy_config_path: Optional path to a deploy YAML.
-        stage_overrides: Per-stage overrides from ``--stage-overrides`` JSON.
-            Keys are stage_id strings, values are dicts of overrides.
-        strategy_config_path: Optional path to a composable-parallel
-            ``strategy.yaml`` overlaid onto the registry-merged stages.
-
-    Returns:
-        Tuple of ``(config_path, stage_configs, omni_lb_policy)`` — the last is
-        the strategy-derived pipeline-wide load-balancer policy (``None`` when no
-        strategy set one), returned for the engine to apply.
-    """
-    config_path = deploy_config_path if deploy_config_path is not None else resolve_model_config_path(model)
-    stage_configs, omni_lb_policy = load_stage_configs_from_model(
-        model,
-        trust_remote_code=trust_remote_code,
-        base_engine_args=kwargs,
-        deploy_config_path=deploy_config_path,
-        stage_overrides=stage_overrides,
-        strategy_config_path=strategy_config_path,
-    )
-    if not stage_configs:
-        if default_stage_cfg_factory is not None:
-            default_stage_cfg = default_stage_cfg_factory()
-            stage_configs = create_config(_convert_dataclasses_to_dict(default_stage_cfg))
-        else:
-            stage_configs = []
-
-    stage_configs = filter_stages(config_path, stage_configs, kwargs)
-
-    # ``revision`` selects the model source for every stage. Keep the global
-    # CLI pin on each stage's engine args so diffusion config construction can
-    # pass it through to all component loaders. The registry-to-legacy-stage
-    # conversion currently drops this inherited vLLM EngineArgs field.
-    revision = (kwargs or {}).get("revision")
-    if revision is not None:
-        for stage_config in stage_configs:
-            engine_args = stage_config.get("engine_args") if hasattr(stage_config, "get") else None
-            if engine_args is not None and engine_args.get("revision") is None:
-                engine_args["revision"] = revision
-
-    logger.debug(f"stage_configs: {stage_configs}")
-
-    return config_path, stage_configs, omni_lb_policy
-
-
 def get_final_stage_id_for_e2e(
     output_modalities: list[str] | None, default_modalities: list[str], stage_list: list
 ) -> int:
@@ -839,7 +249,7 @@ def detect_pid_host() -> bool:
     if not ic:
         return True
 
-    return has_pid_host()
+    return has_pid_host() is True
 
 
 ### Helpers for handling delta messages