387 lines
17 KiB
Python
387 lines
17 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Config factories for vllm-omni, e.g., StageConfigFactory."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import dataclasses
|
|
from dataclasses import asdict
|
|
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_omni.config.pipeline_registry import OMNI_PIPELINES
|
|
from vllm_omni.config.stage_config import (
|
|
_DEPLOY_DIR,
|
|
DeployConfig,
|
|
PipelineConfig,
|
|
StageConfig,
|
|
StageType,
|
|
build_stage_runtime_overrides,
|
|
load_deploy_config,
|
|
merge_pipeline_deploy,
|
|
)
|
|
from vllm_omni.config.yaml_util import create_config
|
|
from vllm_omni.diffusion.utils.hf_utils import _looks_like_dreamzero
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
class 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``.
|
|
"""
|
|
|
|
@classmethod
|
|
def create_from_model(
|
|
cls,
|
|
model: str,
|
|
cli_overrides: dict[str, Any] | None = None,
|
|
deploy_config_path: str | None = None,
|
|
**deprecated_kwargs: Any,
|
|
) -> list[StageConfig] | 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.
|
|
"""
|
|
if cli_overrides is None:
|
|
cli_overrides = {}
|
|
|
|
trust_remote_code = cli_overrides.get("trust_remote_code", True)
|
|
if trust_remote_code is None:
|
|
trust_remote_code = False
|
|
|
|
model_type, hf_config = cls._auto_detect_model_type(model, trust_remote_code=trust_remote_code)
|
|
if model_type == "vla":
|
|
if _looks_like_dreamzero(model):
|
|
model_type = "dreamzero"
|
|
|
|
# --- 1. Explicit deploy-config pipeline override (highest precedence) ---
|
|
# A deploy YAML may set ``pipeline: <model_type>`` to force routing for
|
|
# models whose HF config has a generic/colliding ``model_type`` or no
|
|
# matching architectures. Ming-omni-tts, for example, reports
|
|
# model_type="dense" with arch BailingMMNative…, which matches no
|
|
# pipeline — without this it falls through to the diffusion default and
|
|
# dies with "Model class BailingMMNativeForConditionalGeneration not
|
|
# found in diffusion model registry". Honor the key before auto-detection.
|
|
explicit_pipeline = None
|
|
if deploy_config_path:
|
|
try:
|
|
explicit_pipeline = load_deploy_config(deploy_config_path).pipeline
|
|
except Exception:
|
|
logger.exception("Failed to read 'pipeline' key from deploy config %s", deploy_config_path)
|
|
if explicit_pipeline:
|
|
pipeline_cfg = cls.resolve_pipeline_config(explicit_pipeline, hf_config)
|
|
if pipeline_cfg is not None:
|
|
return cls._create_from_registry(
|
|
explicit_pipeline,
|
|
pipeline_cfg,
|
|
cli_overrides,
|
|
deploy_config_path,
|
|
)
|
|
logger.warning(
|
|
"Deploy config %s requested pipeline %r which is not in OMNI_PIPELINES; "
|
|
"falling back to auto-detection.",
|
|
deploy_config_path,
|
|
explicit_pipeline,
|
|
)
|
|
|
|
# --- 2. Auto-detected model_type registered in OMNI_PIPELINES ---
|
|
if model_type and model_type in OMNI_PIPELINES:
|
|
pipeline_cfg = cls.resolve_pipeline_config(model_type, hf_config)
|
|
if pipeline_cfg is not None:
|
|
return cls._create_from_registry(
|
|
model_type,
|
|
pipeline_cfg,
|
|
cli_overrides,
|
|
deploy_config_path,
|
|
)
|
|
|
|
# --- HF architecture fallback: some models report a generic
|
|
# model_type that collides with another model. Match by the
|
|
# hf_architectures declared on each registered PipelineConfig.
|
|
if hf_config 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)
|
|
# Resolvers that get configs of the incorrect type should return None
|
|
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 cls._create_from_registry(
|
|
pipeline_cfg.model_type,
|
|
pipeline_cfg,
|
|
cli_overrides,
|
|
deploy_config_path,
|
|
)
|
|
|
|
# --- Explicit deploy-config pipeline ---
|
|
# When auto-detection above resolves nothing (generic/missing HF model_type
|
|
# and no matching architecture), honor an explicit pipeline key in the deploy config
|
|
if deploy_config_path is not None:
|
|
deploy_path = Path(deploy_config_path)
|
|
if deploy_path.exists():
|
|
deploy_cfg = load_deploy_config(deploy_path)
|
|
if deploy_cfg.pipeline:
|
|
pipeline_cfg = cls.resolve_pipeline_config(deploy_cfg.pipeline, hf_config)
|
|
if pipeline_cfg is not None:
|
|
return cls._create_from_registry(
|
|
pipeline_cfg.model_type,
|
|
pipeline_cfg,
|
|
cli_overrides,
|
|
deploy_config_path,
|
|
)
|
|
|
|
# Not in the pipeline registry — let the caller fall back to the
|
|
# legacy ``stage_configs/*.yaml`` path (resolve_model_config_path).
|
|
return None
|
|
|
|
@classmethod
|
|
def _create_from_registry(
|
|
cls,
|
|
model_type: str,
|
|
pipeline_cfg: PipelineConfig,
|
|
cli_overrides: dict[str, Any],
|
|
deploy_config_path: str | None = None,
|
|
**deprecated_kwargs: Any,
|
|
) -> list[StageConfig]:
|
|
"""Create StageConfigs from pipeline registry + deploy YAML.
|
|
|
|
Precedence: caller-typed (non-None) value > deploy YAML >
|
|
StageDeployConfig dataclass default.
|
|
"""
|
|
# Resolve deploy config path
|
|
if deploy_config_path is None:
|
|
deploy_path = _DEPLOY_DIR / f"{model_type}.yaml"
|
|
else:
|
|
deploy_path = Path(deploy_config_path)
|
|
|
|
if not deploy_path.exists():
|
|
logger.warning(
|
|
"Deploy config not found: %s — using pipeline defaults only",
|
|
deploy_path,
|
|
)
|
|
deploy_cfg = DeployConfig()
|
|
else:
|
|
deploy_cfg = load_deploy_config(deploy_path)
|
|
# Fallback to using the deploy config pipeline class if it's a mismatch
|
|
if deploy_cfg.pipeline and deploy_cfg.pipeline != model_type:
|
|
resolved = cls.resolve_pipeline_config(deploy_cfg.pipeline)
|
|
if resolved is None:
|
|
raise KeyError(
|
|
f"Pipeline {deploy_cfg.pipeline!r} from {deploy_path.name!r} "
|
|
f"not found in OMNI_PIPELINES. Available: "
|
|
f"{sorted(OMNI_PIPELINES.keys())}"
|
|
)
|
|
pipeline_cfg = resolved
|
|
|
|
cli_async_chunk = cli_overrides.get("async_chunk")
|
|
if cli_async_chunk is not None:
|
|
deploy_cfg.async_chunk = bool(cli_async_chunk)
|
|
|
|
stages = merge_pipeline_deploy(pipeline_cfg, deploy_cfg, cli_overrides)
|
|
|
|
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)
|
|
|
|
return stages
|
|
|
|
@classmethod
|
|
def create_default_diffusion(cls, 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.
|
|
|
|
Args:
|
|
kwargs: Engine arguments from CLI/API.
|
|
|
|
Returns:
|
|
List containing a single config dict for the diffusion stage.
|
|
"""
|
|
# 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}"
|
|
|
|
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")
|
|
engine_args["model_stage"] = "diffusion"
|
|
|
|
# Convert dtype to string for OmegaConf
|
|
if "dtype" in engine_args:
|
|
engine_args["dtype"] = str(engine_args["dtype"])
|
|
|
|
engine_args.setdefault("max_num_seqs", 1)
|
|
|
|
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": "image",
|
|
}
|
|
|
|
return [config_dict]
|
|
|
|
@classmethod
|
|
def _auto_detect_model_type(cls, model: str, trust_remote_code: bool = True) -> tuple[str | None, Any]:
|
|
"""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:
|
|
Tuple of (model_type, hf_config). Both may be None on failure.
|
|
"""
|
|
hf_config = None
|
|
|
|
try:
|
|
hf_config = get_config(model, trust_remote_code=trust_remote_code)
|
|
return hf_config.model_type, hf_config
|
|
except Exception as e:
|
|
logger.debug(f"`get_config` failed for {e}; Falling back to raw config.json path")
|
|
|
|
# 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", model, revision=None)
|
|
if config_dict:
|
|
if "model_type" in config_dict:
|
|
return config_dict["model_type"], None
|
|
# 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"], None
|
|
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", model, 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 pipeline_cfg.diffusers_class_name == class_name:
|
|
logger.info(
|
|
"Detected pipeline %r from model_index.json (_class_name=%r)",
|
|
pipeline_cfg.model_type,
|
|
class_name,
|
|
)
|
|
return pipeline_cfg.model_type, None
|
|
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).
|
|
model_lower = 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, None
|
|
|
|
return None, None
|
|
|
|
@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)
|
|
|
|
@staticmethod
|
|
def 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."""
|
|
if model_type not in OMNI_PIPELINES:
|
|
return None
|
|
obj = OMNI_PIPELINES[model_type]
|
|
return obj(hf_config) if callable(obj) else obj
|