vllm-omni/vllm_omni/config/config_factory.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Config factories for vllm-omni, e.g., StageConfigFactory."""

from __future__ import annotations

import functools
import json
from collections.abc import Mapping
from dataclasses import asdict, dataclass, fields
from pathlib import Path
from typing import Any

from transformers import PretrainedConfig
from vllm.logger import init_logger
from vllm.transformers_utils.config import get_config
from vllm.transformers_utils.repo_utils import get_hf_file_to_dict
from vllm.transformers_utils.runai_utils import ObjectStorageModel, is_runai_obj_uri

from vllm_omni.config.endpoint_policy import EndpointRestriction
from vllm_omni.config.omni_config import (
    VllmOmniConfig,
    normalize_and_validate_diffusion_engine_ingress_kwargs,
)
from vllm_omni.config.pipeline_registry import OMNI_PIPELINES, resolve_pipeline_config
from vllm_omni.config.stage_config import (
    _DEPLOY_DIR,
    DeployConfig,
    PipelineConfig,
    StageConfig,
    StageDeployConfig,
    StageExecutionType,
    StagePipelineConfig,
    StageType,
    build_stage_runtime_overrides,
    load_deploy_config,
    merge_pipeline_deploy,
    normalize_pipeline_cli_overrides,
)
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

logger = init_logger(__name__)


# Default degree for any parallel axis / replica count that isn't set anywhere
# (CLI, deploy YAML, or pipeline default): a single, un-parallelized rank.
# TODO(composable_parallel): this "1" default and the per-axis device-layout
# fallbacks are currently re-derived in the merge, apply, and reconcile layers.
# Centralize them in one schema so the pre-spawn device guard and the engine-args
# defaults can't drift apart. This is the light slice; the full device-layout
# centralization is tracked as a follow-up.
_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.

    vLLM's Run:AI streamer keeps ``s3://``/``gs://``/``az://`` URIs opaque until
    each stage builds its ``ModelConfig``; parent-process resolution (HF config
    lookup, pipeline/pipeline-key matching) would instead hand the URI to
    ``huggingface_hub`` helpers, which reject it with ``HFValidationError``.
    Pull the lightweight files once into vLLM's deterministic
    ``model_streamer/<hash>`` directory so config reads work here, and so the
    stage processes' own pull lands in that same directory.

    Returns the input unchanged for non object-storage paths.
    """
    if not is_runai_obj_uri(model):
        return model
    object_storage_model = ObjectStorageModel(url=model)
    object_storage_model.pull_files(model, allow_pattern=["*.model", "*.py", "*.json"])
    logger.info("Materialized object-storage configs for %s at %s", model, object_storage_model.dir)
    return object_storage_model.dir


def _name_match_candidate(model: str) -> str:
    """Last path component of a model reference, used for name-based matching.

    Object-storage URIs and HF repo ids carry non-model segments (bucket name,
    organization) that must not participate in substring matching; e.g. a
    bucket named ``qwen3-tts-models`` holding a ``Qwen3-Omni`` checkpoint must
    not resolve to the ``qwen3_tts`` pipeline.
    """
    return model.rstrip("/").rsplit("/", 1)[-1]


def with_trust_remote_code_override(
    overrides: Mapping[str, Any],
    trust_remote_code: bool | None,
) -> dict[str, Any]:
    """Merge the tri-state ``trust_remote_code`` into an override mapping.

    Single home for the precedence rule (explicit caller value > deploy
    yaml per-stage value > vLLM default False): a non-None value becomes an
    explicit override; ``None`` means "not specified" and leaves the deploy
    yaml's per-stage setting in effect. The serve ``--trust-remote-code``
    flag is store_true — its absent-False must be mapped to ``None`` at the
    CLI boundary before reaching here, since it cannot express an explicit
    False.
    """
    merged = dict(overrides)
    if trust_remote_code is not None:
        merged["trust_remote_code"] = trust_remote_code
    return merged


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
    where keys in OMNI_PIPELINES map to either a PipelineConfig, or a callable
    which accepts a Transformers config as an arg & resolves to a PipelineConfig.

    NOTE: Models with generic HF ``model_type`` collisions (e.g. MiMo Audio
    reports ``qwen2``) should declare ``hf_architectures=(...)`` on their
    ``PipelineConfig`` so the factory can disambiguate via ``hf_config.architectures``.
    """

    @classmethod
    def get_pipeline_endpoint_restrictions(
        cls,
        model: str,
        trust_remote_code: bool,
        deploy_config_path: str | None,
    ) -> tuple[EndpointRestriction, ...]:
        """Given a model string, determine the corresponding endpoint restrictions.

        Args:
            model: Model name or path.
            trust_remote_code: Whether to trust remote code for HF config loading.
            deploy_config_path: Optional path to the deploy config for the pipeline.

        Returns:
            A tuple of model specific endpoint restrictions.
        """
        pipeline_cfg = StageConfigFactory.get_pipeline_config(
            model=model,
            trust_remote_code=trust_remote_code,
            deploy_config_path=deploy_config_path,
        )
        return pipeline_cfg.endpoint_restrictions if pipeline_cfg else ()

    @classmethod
    @functools.cache
    def get_hf_config(cls, model: str, trust_remote_code: bool) -> PretrainedConfig | None:
        """Fetch the HF config (if it exists) from the model directory.

        Args:
            model: Model name or path.
            trust_remote_code: Whether to trust remote code for HF config loading.

        Returns:
            the model's config or None.
        """
        hf_config = None
        try:
            return get_config(_materialize_object_storage_configs(model), trust_remote_code=trust_remote_code)
        except Exception as e:
            logger.debug(f"`get_config` failed with exception {e}; inferred HF config is None")
        return hf_config

    @classmethod
    @functools.cache
    def try_infer_model_type(cls, model: str, trust_remote_code: bool) -> str | None:
        """Auto-detect model_type from model directory and apply any model
        specific patches to get the correct model_type str. If we are unable
        to infer it from the model directory, we fall back to the PipelineConfig.

        Args:
            model: Model name or path.
            trust_remote_code: Whether to trust remote code for HF config loading.

        Returns:
            model_type as a string; may be None on failure.
        """
        model_type = cls._try_infer_model_type(
            model=model,
            trust_remote_code=trust_remote_code,
        )
        if model_type == "vla":
            if _looks_like_dreamzero(model):
                model_type = "dreamzero"
        return model_type

    @classmethod
    def _try_infer_model_type(cls, model: str, trust_remote_code: bool) -> str | None:
        """Auto-detect model_type from model directory.

        Args:
            model: Model name or path.
            trust_remote_code: Whether to trust remote code for HF config loading.

        Returns:
            model_type as a string; may be None on failure.
        """
        hf_config = cls.get_hf_config(
            model=model,
            trust_remote_code=trust_remote_code,
        )
        if hf_config is not None:
            return hf_config.model_type

        config_source = _materialize_object_storage_configs(model)

        # Fallback: read config.json directly for custom model types that
        # are not registered with transformers (e.g. qwen3_tts).
        try:
            config_dict = get_hf_file_to_dict("config.json", config_source, revision=None)
            if config_dict:
                if "model_type" in config_dict:
                    return config_dict["model_type"]
                # VoxCPM2-style configs use singular ``architecture`` rather
                # than HF's standard ``model_type`` / ``architectures``. Accept
                # it as a fallback so the pipeline registry can still match.
                if "architecture" in config_dict and isinstance(config_dict["architecture"], str):
                    return config_dict["architecture"]
        except Exception as e:
            logger.debug(f"Failed to auto-detect model type for {model}: {e}")

        # Fallback for diffusers-style models: check model_index.json.
        # Some models (e.g. GLM-Image) have no root config.json but ship a
        # model_index.json with _class_name that maps to a pipeline key via
        # PipelineConfig.diffusers_class_name.
        try:
            model_index = get_hf_file_to_dict("model_index.json", config_source, revision=None)
            if model_index and "_class_name" in model_index:
                class_name = model_index["_class_name"]
                for obj in OMNI_PIPELINES.values():
                    # If we have a resolver, call it with the optional hf_config
                    # to get the default pipeline config for this key
                    pipeline_cfg = obj(hf_config) if callable(obj) else obj
                    if pipeline_cfg is not None and class_name in (
                        pipeline_cfg.diffusers_class_name,
                        *pipeline_cfg.diffusers_class_aliases,
                    ):
                        logger.info(
                            "Detected pipeline %r from model_index.json (_class_name=%r)",
                            pipeline_cfg.model_type,
                            class_name,
                        )
                        return pipeline_cfg.model_type
        except Exception as e:
            logger.debug(f"Failed to detect model type for diffusers-style models: {e}")

        # Final fallback: some models (e.g. CosyVoice3) ship an empty
        # config.json and rely on naming conventions. Match the model path
        # basename against registered pipeline keys — longest match wins
        # so "cosyvoice3" (length 10) beats "cosyvoice" (length 9). Only
        # the basename is scanned so URI segments such as the bucket name
        # cannot select an unrelated pipeline.
        model_lower = _name_match_candidate(model).lower().replace("-", "").replace("_", "")
        best: str | None = None
        best_len = 0
        for registered_key in OMNI_PIPELINES.keys():
            candidate = registered_key.lower().replace("-", "").replace("_", "")
            if candidate and candidate in model_lower and len(candidate) > best_len:
                best = registered_key
                best_len = len(candidate)
        if best is not None:
            return best

        return None

    @classmethod
    def get_pipeline_config(
        cls,
        model: str,
        trust_remote_code: bool,
        deploy_config_path: str | None = None,
        user_deploy_config: DeployConfig | None = None,
    ) -> PipelineConfig | None:
        """Resolve the PipelineConfig for a model path/name."""
        model_type = cls.try_infer_model_type(model=model, trust_remote_code=trust_remote_code)
        hf_config = cls.get_hf_config(model=model, trust_remote_code=trust_remote_code)

        # Resolve the deploy config & check if the user set the pipeline;
        # If the pipeline is explicitly set, it takes highest priority
        if user_deploy_config is None:
            user_deploy_config = cls._load_user_deploy_config(deploy_config_path)
        deploy_config_pipe = cls._get_deploy_override_pipe_config(hf_config, user_deploy_config)
        if deploy_config_pipe is not None:
            return deploy_config_pipe

        # Pipeline isn't set in the yaml spec, so we need infer it ourselves.
        if model_type and model_type in OMNI_PIPELINES:
            pipeline_cfg = resolve_pipeline_config(model_type, hf_config)
            if pipeline_cfg is not None:
                return pipeline_cfg

        if hf_config is not None:
            if model_type is not None:
                logger.warning("Inferred model type %s is not registered to an Omni pipeline", model_type)
            hf_archs = set(getattr(hf_config, "architectures", []) or [])
            if hf_archs:
                for registered in OMNI_PIPELINES.values():
                    pipeline_cfg = registered if isinstance(registered, PipelineConfig) else registered(hf_config)
                    if pipeline_cfg is None:
                        continue
                    predicate = pipeline_cfg.hf_config_predicate
                    if predicate is not None:
                        try:
                            if not predicate(hf_config):
                                logger.debug(
                                    "Pipeline %r matched on architectures %s but its "
                                    "hf_config_predicate rejected the loaded config; "
                                    "continuing fallback search.",
                                    pipeline_cfg.model_type,
                                    sorted(hf_archs.intersection(pipeline_cfg.hf_architectures)),
                                )
                                continue
                        except Exception:
                            logger.exception(
                                "Pipeline %r hf_config_predicate raised; skipping.",
                                pipeline_cfg.model_type,
                            )
                            continue
                    if isinstance(pipeline_cfg, PipelineConfig) and hf_archs.intersection(
                        pipeline_cfg.hf_architectures
                    ):
                        return pipeline_cfg
        return None

    @classmethod
    def _get_deploy_override_pipe_config(
        cls,
        hf_config: PretrainedConfig | None,
        deploy_config: DeployConfig | None,
    ) -> PipelineConfig | None:
        """Resolve an explicit pipeline override from a loaded deploy config."""
        if deploy_config is None or deploy_config.pipeline is None:
            return None

        pipeline_cfg = resolve_pipeline_config(deploy_config.pipeline, hf_config)
        if pipeline_cfg is None:
            raise KeyError(
                f"Pipeline {deploy_config.pipeline!r} from deploy config is not registered "
                f"to OMNI_PIPELINES. Available: {sorted(OMNI_PIPELINES)}"
            )
        return pipeline_cfg

    @staticmethod
    def _load_user_deploy_config(deploy_config_path: str | None) -> DeployConfig | None:
        """Load an explicit deploy YAML once for resolution and construction."""
        if deploy_config_path is None:
            return None
        deploy_path = Path(deploy_config_path)
        if not deploy_path.exists() and deploy_path.parent == Path("."):
            candidate = _DEPLOY_DIR / deploy_path
            if candidate.exists():
                deploy_path = candidate
        if not deploy_path.exists():
            raise FileNotFoundError(f"Deploy config not found: {deploy_path}")
        return load_deploy_config(deploy_path)

    @classmethod
    def create_from_model(
        cls,
        model: str,
        *,
        trust_remote_code: bool | None,
        cli_overrides: dict[str, Any],
        deploy_config_path: str | None,
        strategy_specs: Mapping[Any, Any] | None = None,
    ) -> VllmOmniConfig | None:
        """Build the structured Omni config for a model/deploy pair."""
        user_deploy_config = cls._load_user_deploy_config(deploy_config_path)
        pipeline_cfg = cls.get_pipeline_config(
            model=model,
            # HF config resolution needs a real bool: transformers treats
            # None as "prompt for consent", which blocks non-interactively.
            trust_remote_code=bool(trust_remote_code),
            deploy_config_path=deploy_config_path,
            user_deploy_config=user_deploy_config,
        )
        if pipeline_cfg is None:
            return None

        # Keep the typed representation's topology in lockstep with the
        # compatibility launch representation. The aligner is a real extra
        # stage, not an engine-only setting, so it must be injected before
        # either representation is built.
        if user_deploy_config is None:
            if pipeline_cfg.default_deploy_config_name is not None:
                default_deploy_path = _DEPLOY_DIR / pipeline_cfg.default_deploy_config_name
                user_deploy_config = load_deploy_config(default_deploy_path)
                deploy_config_path = str(default_deploy_path)
            else:
                user_deploy_config = DeployConfig()
        from vllm_omni.utils.forced_aligner import inject_forced_aligner_stage

        pipeline_cfg, user_deploy_config = inject_forced_aligner_stage(pipeline_cfg, user_deploy_config, cli_overrides)
        registry_cli_overrides = with_trust_remote_code_override(
            {**cli_overrides, "model": model},
            trust_remote_code,
        )
        return VllmOmniConfig.from_pipeline_config(
            pipeline_cfg,
            user_deploy_config=user_deploy_config,
            deploy_config_path=deploy_config_path,
            cli_overrides=registry_cli_overrides,
            strategy_specs=strategy_specs,
        )

    @classmethod
    def create_legacy_stage_configs_from_model(
        cls,
        model: str,
        *,
        trust_remote_code: bool | None,
        cli_overrides: dict[str, Any],
        deploy_config_path: str | None,
        strategy_specs: Mapping[Any, Any] | None = None,
    ) -> tuple[list[StageConfig] | None, str | None]:
        """Build the migration-only runtime ABI from the shared resolution.

        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(
            model=model,
            # See create_from_model: HF resolution needs a real bool.
            trust_remote_code=bool(trust_remote_code),
            deploy_config_path=deploy_config_path,
            user_deploy_config=user_deploy_config,
        )
        if pipeline_cfg is None:
            return None, None

        legacy_cli_overrides = with_trust_remote_code_override(cli_overrides, trust_remote_code)
        return cls._create_legacy_from_registry(
            pipeline_cfg,
            legacy_cli_overrides,
            deploy_config_path,
            user_deploy_config=user_deploy_config,
            strategy_specs=strategy_specs,
        )

    @classmethod
    def _create_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,
    ) -> tuple[list[StageConfig], str | None]:
        """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.
        """
        cli_overrides = normalize_pipeline_cli_overrides(pipeline_cfg, cli_overrides)
        deploy_cfg: DeployConfig | None
        if user_deploy_config is not None:
            deploy_cfg = user_deploy_config
        elif deploy_config_path is not None:
            deploy_cfg = cls._load_user_deploy_config(deploy_config_path)
            assert deploy_cfg is not None
        elif pipeline_cfg.default_deploy_config_name is not None:
            deploy_cfg = load_deploy_config(_DEPLOY_DIR / pipeline_cfg.default_deploy_config_name)
            assert deploy_cfg is not None
        else:
            deploy_cfg = DeployConfig()

        assert deploy_cfg is not None

        cli_async_chunk = cli_overrides.get("async_chunk")
        if cli_async_chunk is not None:
            deploy_cfg.async_chunk = bool(cli_async_chunk)

        from vllm_omni.utils.forced_aligner import inject_forced_aligner_stage

        pipeline_cfg, deploy_cfg = inject_forced_aligner_stage(pipeline_cfg, deploy_cfg, cli_overrides)

        stages = merge_pipeline_deploy(pipeline_cfg, deploy_cfg, cli_overrides)

        # Overlay declarative parallel strategies (opt-in) before CLI overrides.
        applied = cls._apply_strategy_specs(stages, strategy_specs)

        explicit_overrides = {k: v for k, v in cli_overrides.items() if v is not None}

        for stage in stages:
            stage.runtime_overrides = cls._merge_cli_overrides(stage, explicit_overrides)

        # Re-validate the resolved layout now that CLI overrides are on top.
        cls._reconcile_strategy_with_cli(stages, applied)

        omni_lb_policy = applied.omni_lb_policy if applied is not None else None
        return _LegacyConfigResolution(
            stage_configs=stages,
            pipeline_config=pipeline_cfg,
            omni_lb_policy=omni_lb_policy,
        )

    @staticmethod
    def _apply_strategy_specs(
        stages: list[StageConfig],
        strategy_specs: Mapping[Any, Any] | None,
    ) -> Any:
        """Overlay derived parallel sizing onto merged stages (opt-in).

        ``omni_lb_policy`` cannot be set from stage configs (omni reads it once
        at orchestrator construction), so a derived non-default policy is logged
        here and carried on the returned ``StrategyApplyResult`` for the caller
        to hand to the orchestrator.

        Returns the ``StrategyApplyResult`` (or ``None`` when no strategy was
        supplied) so the caller can re-validate the resolved layout once CLI
        overrides have been merged on top, and read its ``omni_lb_policy``.
        """
        if not strategy_specs:
            return None
        from vllm_omni.config.composable_parallel import apply_strategy_specs

        applied = apply_strategy_specs(stages, strategy_specs)
        if applied.omni_lb_policy is not None:
            logger.info(
                "[composable_parallel] strategy derived omni_lb_policy=%r; it will be applied "
                "to the orchestrator unless an explicit --omni-lb-policy was given.",
                applied.omni_lb_policy,
            )
        return applied

    @staticmethod
    def _reconcile_strategy_with_cli(
        stages: list[StageConfig],
        applied: Any,
    ) -> None:
        """Reconcile CLI overrides applied *after* a strategy overlay.

        CLI overrides are applied last (in ``to_omegaconf``), so for any stage a
        strategy touched we must (a) warn loudly when a CLI arg overrides a
        strategy-declared axis — the strategy is meant to be the single writer
        for the axes it declares, so a silent CLI win is surprising — and
        (b) re-run the device-count check against the *effective* (post-CLI)
        ``tp``/``dp``/``pp``/``num_replicas``/``devices`` so a CLI
        ``--devices``/``--tensor-parallel-size``/``--num-replicas`` cannot slip
        past the pre-spawn guard that exists to prevent silent OOMs.
        """
        if applied is None:
            return
        from vllm_omni.config.composable_parallel import check_device_layout

        # Axis kind -> (engine-arg field, strategy-derived attribute on the cfg).
        axis_fields = {
            "tp": ("tensor_parallel_size", "tensor_parallel_size"),
            "dp": ("data_parallel_size", "data_parallel_size"),
            "pp": ("pipeline_parallel_size", "pipeline_parallel_size"),
        }

        for stage in stages:
            cfg = applied.per_stage_config.get(stage.stage_id)
            if cfg is None:
                continue
            overrides = stage.runtime_overrides or {}
            declared = set(cfg.l1_owners.keys())

            for kind, (field_name, attr) in axis_fields.items():
                if kind in declared and overrides.get(field_name) is not None:
                    cli_val = overrides[field_name]
                    derived = getattr(cfg, attr)
                    if cli_val != derived:
                        logger.warning(
                            "[composable_parallel] stage %s: CLI %s=%s overrides the "
                            "strategy-derived %s=%s. The CLI value wins; remove one to avoid ambiguity.",
                            stage.stage_id,
                            field_name,
                            cli_val,
                            field_name,
                            derived,
                        )
            if "stage_replica" in declared and overrides.get("num_replicas") is not None:
                cli_val = overrides["num_replicas"]
                if cli_val != cfg.stage_replica_size:
                    logger.warning(
                        "[composable_parallel] stage %s: CLI num_replicas=%s overrides the "
                        "strategy-derived num_replicas=%s. The CLI value wins; remove one to avoid ambiguity.",
                        stage.stage_id,
                        cli_val,
                        cfg.stage_replica_size,
                    )

            def _eff(field_name: str, fallback: Any) -> Any:
                val = overrides.get(field_name)
                return val if val is not None else fallback

            def _eff_degree(field_name: str, source: dict[str, Any]) -> int:
                # Single place the per-axis "default to 1" fallback is applied,
                # for both the override and the YAML-default sides (see the
                # _DEFAULT_PARALLEL_DEGREE TODO).
                value = _eff(field_name, source.get(field_name, _DEFAULT_PARALLEL_DEGREE))
                return int(value or _DEFAULT_PARALLEL_DEGREE)

            check_device_layout(
                _eff("devices", stage.yaml_runtime.get("devices")),
                tensor_parallel_size=_eff_degree("tensor_parallel_size", stage.yaml_engine_args),
                data_parallel_size=_eff_degree("data_parallel_size", stage.yaml_engine_args),
                pipeline_parallel_size=_eff_degree("pipeline_parallel_size", stage.yaml_engine_args),
                num_replicas=_eff_degree("num_replicas", stage.yaml_runtime),
                role=stage.model_stage,
            )

    @classmethod
    def _normalize_default_diffusion(
        cls,
        kwargs: dict[str, Any],
    ) -> tuple[dict[str, Any], DiffusionParallelConfig, dict[str, Any], str]:
        """Normalize inputs shared by typed and compatibility diffusion builders."""
        # Preserve the default builder's explicit top-level override of model
        # extras before the strict same-source alias validator runs.
        if kwargs.get("auxiliary_text_encoder") is not None:
            kwargs = dict(kwargs)
            extras = dict(kwargs.get("extras") or {})
            extras.pop("auxiliary_text_encoder", None)
            kwargs["extras"] = extras
        kwargs = normalize_and_validate_diffusion_engine_ingress_kwargs(kwargs, stage_id=0)
        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 {}

        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"]))
        # Ingress keeps the engine spelling; canonicalize it before filtering
        # to terminal diffusion fields. Conflicts were validated at ingress.
        quantization = kwargs.pop("quantization", None)
        if quantization is not None:
            kwargs["quantization_config"] = quantization
        diffusion_config_fields = frozenset(config_field.name for config_field in fields(OmniDiffusionConfig))
        engine_args = OmniDiffusionConfig.normalize_init_kwargs(
            {name: value for name, value in kwargs.items() if name in diffusion_config_fields}
        )

        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

        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,
        )
        return engine_args, parallel_config, default_sampling_params, final_output_type

    @classmethod
    def create_default_diffusion(cls, kwargs: dict[str, Any]) -> list[dict[str, Any]]:
        """Build the temporary runtime ABI for a generic diffusion stage.

        The terminal diffusion config owns engine defaults and normalization;
        this compatibility builder only adds Omni stage topology, request
        defaults, and device placement.
        """
        engine_args, parallel_config, default_sampling_params, final_output_type = cls._normalize_default_diffusion(
            kwargs
        )
        engine_args["parallel_config"] = asdict(parallel_config)
        engine_args["model_stage"] = "diffusion"

        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 create_typed_default_diffusion(
        cls,
        model: str,
        kwargs: dict[str, Any],
    ) -> VllmOmniConfig:
        """Build generic diffusion directly into the structured runtime config."""
        engine_overrides, parallel_config, default_sampling_params, final_output_type = (
            cls._normalize_default_diffusion(kwargs)
        )
        engine_overrides.update(asdict(parallel_config))
        model_class_name = engine_overrides.get("model_class_name")
        pipeline = PipelineConfig(
            model_type="generic_diffusion",
            model_arch=str(model_class_name or ""),
            stages=(
                StagePipelineConfig(
                    stage_id=0,
                    model_stage="diffusion",
                    execution_type=StageExecutionType.DIFFUSION,
                    final_output=True,
                    final_output_type=final_output_type,
                ),
            ),
        )

        # These values have already been normalized for this diffusion stage.
        # Scope them explicitly so diffusion-only fields such as engine_backend
        # and extras are not filtered by the global LLM CLI argument surface.
        stage_overrides = {
            f"stage_0_{key}": value for key, value in engine_overrides.items() if key not in {"model", "stage_id"}
        }
        stage_overrides["stage_0_devices"] = (
            kwargs.get("stage_0_devices")
            or kwargs.get("devices")
            or ",".join(str(i) for i in range(parallel_config.world_size))
        )
        # Caller defaults belong to deployment, not immutable pipeline constraints:
        # explicit request sampling parameters must still be able to override them.
        deploy = DeployConfig(stages=[StageDeployConfig(stage_id=0, default_sampling_params=default_sampling_params)])
        return VllmOmniConfig.from_pipeline_config(
            pipeline, user_deploy_config=deploy, cli_overrides={"model": model, **stage_overrides}
        )

    @classmethod
    def _merge_cli_overrides(
        cls,
        stage: StageConfig,
        cli_overrides: dict[str, Any],
    ) -> dict[str, Any]:
        """Merge global and per-stage (``stage_N_*``) CLI overrides.

        Orchestrator-owned keys are filtered by ``build_stage_runtime_overrides``
        using ``OrchestratorArgs`` as the single source of truth; unknown
        server/uvicorn keys are dropped downstream by
        ``filter_dataclass_kwargs(OmniEngineArgs, ...)``.
        """
        return build_stage_runtime_overrides(stage.stage_id, cli_overrides)