[Core] Simplify OmniModelConfig Initialization (#1768)

Signed-off-by: Alex Brooks <albrooks@redhat.com>

tests/engine/test_arg_utils.py

@@ -0,0 +1,134 @@
+"""
+Tests for Omni config utils. For stability, these tests should largely be
+invariant to the specific attributes of vLLM config except in cases where we
+explicitly patch values that differ from vLLM.
+"""
+
+import inspect
+from unittest.mock import Mock
+
+import pytest
+from pydantic import ValidationError
+from transformers import PretrainedConfig
+from vllm.engine.arg_utils import EngineArgs
+
+from vllm_omni.config.model import OmniModelConfig
+from vllm_omni.engine.arg_utils import OmniEngineArgs
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_sync_config_is_omni():
+    """Ensure create_model_config gives the right type."""
+    cfg = OmniEngineArgs().create_model_config()
+    assert isinstance(cfg, OmniModelConfig)
+
+
+def test_multimodal_kwarg_overrides():
+    """Ensure that overrides in the multimodal config are preserved."""
+    # Get a different value than the default for a multimodal field
+    sig = inspect.signature(OmniEngineArgs)
+    default_mm_cache = sig.parameters["mm_processor_cache_gb"].default
+    override_val = default_mm_cache + 1
+
+    # NOTE: This needs to be a model that resolves to supports_multimodal=True
+    # in vLLM, otherwise we won't have an MM config
+    cfg = OmniEngineArgs(
+        model="Qwen/Qwen2-VL-2B-Instruct",
+        mm_processor_cache_gb=override_val,
+    ).create_model_config()
+
+    # Ensure that the override was applied correctly
+    assert cfg.multimodal_config is not None
+    assert cfg.multimodal_config.mm_processor_cache_gb == override_val
+
+
+def test_from_vllm_config_validates_invalid_omni_kwargs():
+    """Ensure omni-specific field validation catches invalid keys."""
+    model_config = EngineArgs().create_model_config()
+    with pytest.raises(ValueError, match="Unexpected omni kwarg"):
+        OmniModelConfig.from_vllm_model_config(model_config, foo="bar")
+
+
+def test_from_vllm_config_validates_bad_omni_kwarg_types():
+    """Ensure omni-specific field validation catches type errors."""
+    model_config = EngineArgs().create_model_config()
+    with pytest.raises(ValidationError):
+        OmniModelConfig.from_vllm_model_config(model_config, stage_id="not_an_int")
+
+
+def test_default_all_values_are_initialized():
+    """Ensure omni-specific field initializes all fields"""
+    model_config = EngineArgs().create_model_config()
+    cfg = OmniModelConfig.from_vllm_model_config(model_config)
+
+    # Test a primitive
+    assert cfg.model_stage == "thinker"
+    # Test a field initialized with a default factory
+    assert cfg.stage_connector_config == {
+        "name": "SharedMemoryConnector",
+        "extra": {},
+    }
+
+    # Ensure that hf_config is initialized on model_config in the vLLM by ModelConfig's
+    # __post_init__, and that the hf_config is copied over to the OmniModelConfig;
+    # we explicitly set this since the field sets init=False
+    assert isinstance(model_config.hf_config, PretrainedConfig)
+    assert cfg.hf_config is model_config.hf_config
+
+    # Ensure that we can convert it to a string; this will convert
+    # all attributes, so should raise if we have attributes that are
+    # not initialized correctly, e.g., due to default factories
+    str(cfg)
+
+
+def test_qwen3_tts_codec_frame_rate_patching():
+    """Ensure the patch for qwen3 tts is applied correctly when creating the omni config."""
+    # Create a vLLM ModelConfig
+    vllm_config = EngineArgs().create_model_config()
+
+    # Create a mock talking config with a dummy value for position_id_per_seconds
+    mock_talker_config = Mock()
+    mock_talker_config.position_id_per_seconds = 12.3
+    vllm_config.hf_config.talker_config = mock_talker_config
+
+    # Ensure creating the config for a Qwen3TTSTalkerForConditionalGenerationARVLLM
+    # model calls the patch func to apply position_id_per_seconds from the talker
+    # config to the config's codec_frame_rate_hz
+    omni_config = OmniModelConfig.from_vllm_model_config(
+        vllm_config,
+        model_arch="Qwen3TTSTalkerForConditionalGenerationARVLLM",
+    )
+
+    # Verify codec_frame_rate_hz was patched
+    assert omni_config.codec_frame_rate_hz == 12.3
+
+
+def test_stage_specific_text_config_override():
+    """Ensure dependent attributes are updated when using stage-specific config."""
+    vllm_config = EngineArgs().create_model_config()
+    vllm_config.disable_sliding_window = True
+
+    # Switch the created hf text config with a mock whose
+    # values we want to pull through the text config helper
+    stage_text_config = vllm_config.hf_text_config
+    vllm_config.hf_text_config = Mock()
+    stage_text_config.sliding_window = 4096
+    stage_text_config.attention_chunk_size = 2048
+
+    # Move the stage config's text config getter & thinker config
+    mock_stage_config = Mock()
+    mock_stage_config.get_text_config.return_value = stage_text_config
+    vllm_config.hf_config.thinker_config = mock_stage_config
+
+    # Ensure that create from a vLLM config correctly pulls the
+    # expected values off of the thinker config & swaps the text config
+    omni_config = OmniModelConfig.from_vllm_model_config(
+        vllm_config,
+        hf_config_name="thinker_config",
+    )
+
+    assert omni_config.hf_text_config is stage_text_config
+    assert omni_config.attention_chunk_size == 2048
+    assert omni_config.max_model_len == 4096
+    assert omni_config.hf_text_config.sliding_window is None

vllm_omni/config/model.py

@@ -1,13 +1,11 @@
-from dataclasses import field
+from dataclasses import MISSING, field
 from typing import Any
 
-from pydantic import ConfigDict
+from pydantic import ConfigDict, TypeAdapter
 from vllm.config import ModelConfig
-from vllm.config.multimodal import MMCacheType, MMEncoderTPMode
 from vllm.config.utils import config
 from vllm.logger import init_logger
 from vllm.transformers_utils.config import get_hf_text_config
-from vllm.v1.attention.backends.registry import AttentionBackendEnum
 
 import vllm_omni.model_executor.models as me_models
 
@@ -18,31 +16,38 @@ logger = init_logger(__name__)
 class OmniModelConfig(ModelConfig):
     """Configuration for Omni models, extending the base ModelConfig.
 
-    This configuration class extends the base vLLM ModelConfig with
-    omni-specific fields for multi-stage pipeline processing.
+     This configuration class extends the base vLLM ModelConfig with
+     omni-specific fields for multi-stage pipeline processing.
 
-    Attributes:
-        stage_id: Identifier for the stage in a multi-stage pipeline (default: 0)
-        async_chunk: If set to True, perform async chunk
-        model_stage: Stage type identifier, e.g., "thinker" or "talker"
-            (default: "thinker")
-        model_arch: Model architecture name
-            (default: "Qwen2_5OmniForConditionalGeneration")
-        worker_type: Model Type, e.g., "ar" or "generation"
-        engine_output_type: Optional output type specification for the engine.
-            Used to route outputs to appropriate processors (e.g., "image",
-            "audio", "latents"). If None, output type is inferred.
-        stage_connector_config: Stage connector configuration dictionary.
-            Contains "name" (connector name), "extra" (extra connector config).
-        task_type: Default task type for TTS models (CustomVoice, VoiceDesign, or Base).
-            If not specified, will be inferred from model path.
+     Attributes:
+         hf_config: The model's HF Transformers config (default: None)
+         hf_text_config: The sub text_config of the model's hf_config (default: None)
+         stage_id: Identifier for the stage in a multi-stage pipeline (default: 0)
+         async_chunk: If set to True, perform async chunk
+         model_stage: Stage type identifier, e.g., "thinker" or "talker"
+             (default: "thinker")
+         model_arch: Model architecture name
+             (default: "Qwen2_5OmniForConditionalGeneration")
+         worker_type: Model Type, e.g., "ar" or "generation"
+         engine_output_type: Optional output type specification for the engine.
+             Used to route outputs to appropriate processors (e.g., "image",
+             "audio", "latents"). If None, output type is inferred.
+         stage_connector_config: Stage connector configuration dictionary.
+             Contains "name" (connector name), "extra" (extra connector config).
+         task_type: Default task type for TTS models (CustomVoice, VoiceDesign, or Base).
+             If not specified, will be inferred from model path.
 
-    Example:
-        >>> config = OmniModelConfig(
-        ...     stage_id=0,
-        ...     model_stage="thinker",
-        ...     model_arch="Qwen2_5OmniForConditionalGeneration"
-        ... )
+
+    The correct way to initialize this class is via vLLM config, as most
+    of the logic for handling values is in the ModelConfig's __post_init__.
+
+       Example:
+         >>> config = OmniModelConfig.from_vllm_model_config(
+         ...     vllm_config,
+         ...     stage_id=0,
+         ...     model_stage="thinker",
+         ...     model_arch="Qwen2_5OmniForConditionalGeneration"
+         ... )
     """
 
     stage_id: int = 0
@@ -100,59 +105,27 @@ class OmniModelConfig(ModelConfig):
             )
             return get_hf_text_config(self.hf_config)
 
-    def __post_init__(
-        self,
-        # Multimodal config init vars
-        language_model_only: bool,
-        limit_mm_per_prompt: dict[str, int | dict[str, int]] | None,
-        enable_mm_embeds: bool | None,
-        media_io_kwargs: dict[str, dict[str, Any]] | None,
-        mm_processor_kwargs: dict[str, Any] | None,
-        mm_processor_cache_gb: float | None,
-        mm_processor_cache_type: MMCacheType | None,
-        mm_shm_cache_max_object_size_mb: int | None,
-        mm_encoder_only: bool | None,
-        mm_encoder_tp_mode: MMEncoderTPMode | None,
-        mm_encoder_attn_backend: AttentionBackendEnum | str | None,
-        interleave_mm_strings: bool | None,
-        skip_mm_profiling: bool | None,
-        video_pruning_rate: float | None,
-    ) -> None:
-        # Call parent's __post_init__ to handle all standard ModelConfig initialization
-        super().__post_init__(
-            language_model_only=language_model_only,
-            limit_mm_per_prompt=limit_mm_per_prompt,
-            enable_mm_embeds=enable_mm_embeds,
-            media_io_kwargs=media_io_kwargs,
-            mm_processor_kwargs=mm_processor_kwargs,
-            mm_processor_cache_gb=mm_processor_cache_gb,
-            mm_processor_cache_type=mm_processor_cache_type,
-            mm_shm_cache_max_object_size_mb=mm_shm_cache_max_object_size_mb,
-            mm_encoder_only=mm_encoder_only,
-            mm_encoder_tp_mode=mm_encoder_tp_mode,
-            mm_encoder_attn_backend=mm_encoder_attn_backend,
-            interleave_mm_strings=interleave_mm_strings,
-            skip_mm_profiling=skip_mm_profiling,
-            video_pruning_rate=video_pruning_rate,
-        )
+    def _patch_qwen3_tts(self):
+        """Patches the value of `position_id_per_seconds` in Qwen3's
+        TTS's talker_config into the this class's codec_frame_rate_hz.
+        """
+        talker_cfg = getattr(self.hf_config, "talker_config", None)
+        if isinstance(talker_cfg, dict):
+            pos_per_sec = talker_cfg.get("position_id_per_seconds")
+        else:
+            pos_per_sec = getattr(talker_cfg, "position_id_per_seconds", None)
+        if pos_per_sec is not None:
+            try:
+                fps = float(pos_per_sec)
+            except Exception:
+                fps = None
+            if fps is not None and fps > 0:
+                self.codec_frame_rate_hz = fps
 
-        # Qwen3-TTS: infer codec frame rate from the model config for online serving.
-        if self.codec_frame_rate_hz is None and self.model_arch == "Qwen3TTSTalkerForConditionalGenerationARVLLM":
-            talker_cfg = getattr(self.hf_config, "talker_config", None)
-            if isinstance(talker_cfg, dict):
-                pos_per_sec = talker_cfg.get("position_id_per_seconds")
-            else:
-                pos_per_sec = getattr(talker_cfg, "position_id_per_seconds", None)
-            if pos_per_sec is not None:
-                try:
-                    fps = float(pos_per_sec)
-                except Exception:
-                    fps = None
-                if fps is not None and fps > 0:
-                    self.codec_frame_rate_hz = fps
-
-        # Override hf_text_config with omni-specific logic for multi-stage models
-        # (e.g., thinker_config, talker_config)
+    def _maybe_override_text_config(self):
+        """Override hf_text_config with omni-specific logic for multi-stage
+        models (e.g., thinker_config, talker_config).
+        """
         new_hf_text_config = self.draw_hf_text_config()
         if new_hf_text_config is not self.hf_text_config:
             self.hf_text_config = new_hf_text_config
@@ -161,5 +134,84 @@ class OmniModelConfig(ModelConfig):
             # Recalculate max_model_len since it depends on hf_text_config
             self.max_model_len = self.get_and_verify_max_len(self.original_max_model_len)
             # Reset sliding_window if needed
-            if self.disable_sliding_window:
+            if self.disable_sliding_window and self.hf_text_config is not None:
                 self.hf_text_config.sliding_window = None
+
+    @classmethod
+    def from_vllm_model_config(cls, model_config: ModelConfig, **omni_kwargs):
+        """Create OmniModelConfig from an existing vLLM ModelConfig
+        and additional Omni specific kwargs.
+
+        NOTE: The validation and __post_init__ for ModelConfig is expensive;
+        to avoid calling it a second time, we explicitly retrieve defaults
+        from dataclass attributes for values not passed to omni_kwargs,
+        and use that to initialize a __new__ instance. This is significantly
+        faster than creating the OmniModelConfig directly from the ModelConfig,
+        and saves us from having to pass all kwargs to the OmniModelConfig.
+        """
+        # Add missing defaults to the omni kwargs and ensure values are valid
+        cls.add_defaults_to_omni_kwargs(omni_kwargs)
+        cls._validate_omni_fields(**omni_kwargs)
+
+        # Allocate the new omni config and copy the model config & omni fields
+        omni_cfg = object.__new__(cls)
+        omni_cfg.__dict__.update(model_config.__dict__)
+        omni_cfg.__dict__.update(omni_kwargs)
+
+        # Apply any model specific patches or necessary overrides
+        if (
+            omni_cfg.codec_frame_rate_hz is None
+            and omni_cfg.model_arch == "Qwen3TTSTalkerForConditionalGenerationARVLLM"
+        ):
+            omni_cfg._patch_qwen3_tts()
+
+        omni_cfg._maybe_override_text_config()
+
+        if omni_cfg.hf_config is not None:
+            omni_cfg.hf_config.architectures = omni_cfg.architectures
+
+        return omni_cfg
+
+    @classmethod
+    def _validate_omni_fields(cls, **omni_kwargs):
+        """Validate omni-specific fields; we use TypeAdapters here to quickly
+        validate only omni kwargs to avoid rerunning validation on the
+        ModelConfig.
+
+        NOTE: This assumes add_defaults_to_omni_kwargs has already been called,
+        so that all omni fields are present in the provided omni_kwargs.
+        """
+        omni_fields = set(cls.__dataclass_fields__) - set(ModelConfig.__dataclass_fields__)
+
+        for key, value in omni_kwargs.items():
+            if key not in omni_fields:
+                raise ValueError(f"Unexpected omni kwarg: {key}")
+
+            field_type = cls.__dataclass_fields__[key].type
+            if field_type is not Any:
+                TypeAdapter(field_type).validate_python(value)
+
+        # We should not have any uninitialized keys
+        uninitialized_fields = omni_fields - omni_kwargs.keys()
+        if len(uninitialized_fields):
+            logger.error(f"The following OmniModelConfig keys were not initialized: {uninitialized_fields}")
+
+    @classmethod
+    def add_defaults_to_omni_kwargs(cls, omni_kwargs):
+        """Because we init the OmniModelConfig with __new__ to sidestep expensive
+        validation, we need to be careful to ensure fields with default factories
+        are initialized, otherwise we will get an AttributeError when we use it.
+
+        To work around this issue, we explicitly add defaults to the omni_kwargs
+        dict provided to ensure all fields are defined correctly.
+
+        NOTE: omni_kwargs are mutated in place.
+        """
+        omni_fields = set(cls.__dataclass_fields__) - set(ModelConfig.__dataclass_fields__)
+
+        for field_name in omni_fields - set(omni_kwargs.keys()):
+            field_def = cls.__dataclass_fields__[field_name]
+            if field_def.default_factory is not MISSING:
+                omni_kwargs[field_name] = field_def.default_factory()
+            elif field_def.default is not MISSING:
+                omni_kwargs[field_name] = field_def.default

vllm_omni/engine/arg_utils.py

@@ -3,10 +3,8 @@ import dataclasses
 from dataclasses import dataclass, field
 from typing import Any
 
-import vllm.envs as envs
 from vllm.engine.arg_utils import EngineArgs
 from vllm.logger import init_logger
-from vllm.transformers_utils.gguf_utils import is_gguf
 
 from vllm_omni.config import OmniModelConfig
 from vllm_omni.plugins import load_omni_general_plugins
@@ -68,6 +66,9 @@ class OmniEngineArgs(EngineArgs):
         engine_output_type: Optional output type specification for the engine.
             Used to route outputs to appropriate processors (e.g., "image",
             "audio", "latents"). If None, output type is inferred.
+        hf_config_name: Optional key for HF config subkey to be extracted
+            for this stage, e.g., talker_config; If None, the default
+            HF config will be used.
         custom_process_next_stage_input_func: Optional path to a custom function for processing
             inputs from previous stages
             If None, default processing is used.
@@ -113,38 +114,9 @@ class OmniEngineArgs(EngineArgs):
         Returns:
             OmniModelConfig instance with all configuration fields set
         """
-        # GGUF files need a specific model loader path in vLLM.
-        if is_gguf(self.model):
-            self.quantization = self.load_format = "gguf"
-
-        if not envs.VLLM_ENABLE_V1_MULTIPROCESSING:
-            logger.warning(
-                "The global random seed is set to %d. Since "
-                "VLLM_ENABLE_V1_MULTIPROCESSING is set to False, this may "
-                "affect the random state of the Python process that "
-                "launched vLLM.",
-                self.seed,
-            )
-
         # register omni models to avoid model not found error
         self._ensure_omni_models_registered()
 
-        # Keep compatibility when async args are constructed from partial payloads.
-        language_model_only = getattr(self, "language_model_only", False)
-        limit_mm_per_prompt = getattr(self, "limit_mm_per_prompt", {})
-        enable_mm_embeds = getattr(self, "enable_mm_embeds", False)
-        interleave_mm_strings = getattr(self, "interleave_mm_strings", False)
-        media_io_kwargs = getattr(self, "media_io_kwargs", {})
-        skip_mm_profiling = getattr(self, "skip_mm_profiling", False)
-        mm_processor_kwargs = getattr(self, "mm_processor_kwargs", None)
-        mm_processor_cache_gb = getattr(self, "mm_processor_cache_gb", 4)
-        mm_processor_cache_type = getattr(self, "mm_processor_cache_type", None)
-        mm_shm_cache_max_object_size_mb = getattr(self, "mm_shm_cache_max_object_size_mb", 128)
-        mm_encoder_only = getattr(self, "mm_encoder_only", False)
-        mm_encoder_tp_mode = getattr(self, "mm_encoder_tp_mode", "weights")
-        mm_encoder_attn_backend = getattr(self, "mm_encoder_attn_backend", None)
-        video_pruning_rate = getattr(self, "video_pruning_rate", 0.0)
-
         # Build stage_connector_config from stage_connector_spec
         stage_connector_config = {
             "name": self.stage_connector_spec.get("name", "SharedMemoryConnector"),
@@ -152,63 +124,12 @@ class OmniEngineArgs(EngineArgs):
         }
         stage_connector_config["extra"]["stage_id"] = self.stage_id
 
-        # Create OmniModelConfig directly from engine args
-        # Note: We pass the actual init parameters matching vLLM's EngineArgs.create_model_config()
-        omni_config = OmniModelConfig(
-            # Base ModelConfig fields (matching vLLM's EngineArgs.create_model_config)
-            model=self.model,
-            model_weights=self.model_weights,
-            hf_config_path=self.hf_config_path,
-            runner=self.runner,
-            convert=self.convert,
-            tokenizer=self.tokenizer,
-            tokenizer_mode=self.tokenizer_mode,
-            trust_remote_code=self.trust_remote_code,
-            allowed_local_media_path=self.allowed_local_media_path,
-            allowed_media_domains=self.allowed_media_domains,
-            dtype=self.dtype,
-            seed=self.seed,
-            revision=self.revision,
-            code_revision=self.code_revision,
-            hf_token=self.hf_token,
-            hf_overrides=self.hf_overrides,
-            tokenizer_revision=self.tokenizer_revision,
-            max_model_len=self.max_model_len,
-            quantization=self.quantization,
-            allow_deprecated_quantization=self.allow_deprecated_quantization,
-            enforce_eager=self.enforce_eager,
-            enable_return_routed_experts=self.enable_return_routed_experts,
-            max_logprobs=self.max_logprobs,
-            logprobs_mode=self.logprobs_mode,
-            disable_sliding_window=self.disable_sliding_window,
-            disable_cascade_attn=self.disable_cascade_attn,
-            skip_tokenizer_init=self.skip_tokenizer_init,
-            enable_prompt_embeds=self.enable_prompt_embeds,
-            served_model_name=self.served_model_name,
-            language_model_only=language_model_only,
-            limit_mm_per_prompt=limit_mm_per_prompt,
-            enable_mm_embeds=enable_mm_embeds,
-            interleave_mm_strings=interleave_mm_strings,
-            media_io_kwargs=media_io_kwargs,
-            skip_mm_profiling=skip_mm_profiling,
-            config_format=self.config_format,
-            mm_processor_kwargs=mm_processor_kwargs,
-            mm_processor_cache_gb=mm_processor_cache_gb,
-            mm_processor_cache_type=mm_processor_cache_type,
-            mm_shm_cache_max_object_size_mb=mm_shm_cache_max_object_size_mb,
-            mm_encoder_only=mm_encoder_only,
-            mm_encoder_tp_mode=mm_encoder_tp_mode,
-            mm_encoder_attn_backend=mm_encoder_attn_backend,
-            pooler_config=self.pooler_config,
-            generation_config=self.generation_config,
-            override_generation_config=self.override_generation_config,
-            enable_sleep_mode=self.enable_sleep_mode,
-            model_impl=self.model_impl,
-            override_attention_dtype=self.override_attention_dtype,
-            logits_processors=self.logits_processors,
-            video_pruning_rate=video_pruning_rate,
-            io_processor_plugin=self.io_processor_plugin,
-            # Omni-specific fields
+        # Build the vLLM config first, then use it to create the Omni config.
+        model_config = super().create_model_config()
+
+        omni_config = OmniModelConfig.from_vllm_model_config(
+            model_config=model_config,
+            # All kwargs below are Omni specific
             stage_id=self.stage_id,
             async_chunk=self.async_chunk,
             model_stage=self.model_stage,
@@ -221,6 +142,4 @@ class OmniEngineArgs(EngineArgs):
             omni_kv_config=self.omni_kv_config,
             task_type=self.task_type,
         )
-        omni_config.hf_config.architectures = omni_config.architectures
-
         return omni_config