Skip to content

vllm_omni.config.config_factory

Config factories for vllm-omni, e.g., StageConfigFactory.

logger module-attribute 

logger = init_logger(__name__)

StageConfigFactory

Factory that loads pipeline YAML and merges CLI overrides.

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.

create_default_diffusion classmethod 

create_default_diffusion(
    kwargs: dict[str, Any],
) -> list[dict[str, Any]]

Single-stage diffusion - no YAML needed.

Creates a default diffusion stage configuration for single-stage diffusion models. Returns a legacy OmegaConf-compatible dict for backward compatibility with OmniStage.

Parameters:

Name Type Description Default
kwargs dict[str, Any]

Engine arguments from CLI/API.

 required

Returns:

Type Description
list[dict[str, Any]]

List containing a single config dict for the diffusion stage.

create_from_model classmethod 

create_from_model(
    model: str,
    *,
    trust_remote_code: bool = False,
    cli_overrides: dict[str, Any] | None = None,
    deploy_config_path: str | None = None,
    strategy_specs: Mapping[Any, Any] | None = None,
    **deprecated_kwargs: Any,
) -> tuple[list[StageConfig] | None, str | None]

Load pipeline + deploy config, merge with CLI overrides.

Checks OMNI_PIPELINES first, since supported models should be explicitly registered. If a model is not registered in OMNI_PIPELINES, tries to fall back to using the Transformers config & finding pipelines that have overlapping supported architectures.

When strategy_specs is provided (a mapping of role -> list of StrategySpec), the derived parallel sizing is overlaid onto the merged stages (see vllm_omni.config.composable_parallel). This is opt-in: omitting it leaves the existing merge path untouched.

Returns (stages, omni_lb_policy): the merged stages (None when the model is not in the pipeline registry and the caller should fall back to the legacy YAML path) and the strategy-derived, pipeline-wide omni_lb_policy (None when no stage_replica axis set one). The policy is returned rather than threaded through a mutable out-param so the engine can apply it without every intermediate call carrying the dict.

get_hf_config cached classmethod 

get_hf_config(
    model: str, trust_remote_code: bool
) -> PretrainedConfig | None

Fetch the HF config (if it exists) from the model directory.

Parameters:

Name Type Description Default
model str

Model name or path.

 required
trust_remote_code bool

Whether to trust remote code for HF config loading.

 required

Returns:

Type Description
PretrainedConfig | None

the model's config or None.

get_pipeline_config cached classmethod 

get_pipeline_config(
    model: str,
    trust_remote_code: bool,
    deploy_config_path: str | None = None,
) -> PipelineConfig | None

Resolve the PipelineConfig for a model path/name.

get_pipeline_endpoint_restrictions classmethod 

get_pipeline_endpoint_restrictions(
    model: str,
    trust_remote_code: bool,
    deploy_config_path: str | None,
) -> tuple[EndpointRestriction, ...]

Given a model string, determine the corresponding endpoint restrictions.

Parameters:

Name Type Description Default
model str

Model name or path.

 required
trust_remote_code bool

Whether to trust remote code for HF config loading.

 required
deploy_config_path str | None

Optional path to the deploy config for the pipeline.

 required

Returns:

Type Description
tuple[EndpointRestriction, ...]

A tuple of model specific endpoint restrictions.

resolve_pipeline_config staticmethod 

resolve_pipeline_config(
    model_type: str,
    hf_config: PretrainedConfig | None = None,
) -> PipelineConfig | None

Given a model type, resolve to the pipeline to be used. If the pipeline maps to a callable we resolve based on the HF config.

try_infer_model_type cached classmethod 

try_infer_model_type(
    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.

Parameters:

Name Type Description Default
model str

Model name or path.

 required
trust_remote_code bool

Whether to trust remote code for HF config loading.

 required

Returns:

Type Description
str | None

model_type as a string; may be None on failure.