From 465f1464d25cb5e20d61d7bd18d8a8f0c5d82c9a Mon Sep 17 00:00:00 2001 From: Alex Brooks Date: Tue, 23 Jun 2026 21:02:29 -0600 Subject: [PATCH] [Feat] Support Pipeline Resolution Based on HF Config & Make Pipeline Registry Eager (#3760) Signed-off-by: Alex Brooks --- tests/config/test_pipeline_registry.py | 184 ++----- tests/conftest.py | 1 + .../e2e/offline_inference/test_qwen3_omni.py | 18 +- tests/engine/test_async_omni_engine_abort.py | 4 +- .../test_resolve_dreamzero_config.py | 2 +- tests/helpers/fixtures/config.py | 14 + .../models/glm_image/test_glm_image_init.py | 76 --- .../models/minicpmo_4_5/test_pipeline.py | 15 +- tests/test_config_factory.py | 325 +++++++++--- tests/test_diffusion_config_propagation.py | 2 +- .../higgs_audio_v3/test_higgs_audio_v3.py | 40 +- tests/utils/test_tracking_parser.py | 13 +- vllm_omni/config/__init__.py | 4 - vllm_omni/config/config_factory.py | 338 +++++++++++++ vllm_omni/config/pipeline_registry.py | 290 +++++------ vllm_omni/config/stage_config.py | 471 +----------------- vllm_omni/entrypoints/utils.py | 2 +- .../models/aura_omni/__init__.py | 8 + .../model_executor/models/bagel/__init__.py | 9 +- .../models/dreamzero/__init__.py | 8 +- .../models/dynin_omni/__init__.py | 65 +-- .../models/glm_image/__init__.py | 22 +- .../model_executor/models/glm_tts/__init__.py | 20 +- .../model_executor/models/gr00t/__init__.py | 6 + .../models/higgs_audio_v2/__init__.py | 11 +- .../models/higgs_audio_v3/__init__.py | 11 +- .../higgs_audio_v3/higgs_audio_v3_code2wav.py | 6 +- .../higgs_audio_v3/higgs_audio_v3_talker.py | 4 +- .../models/hunyuan_image3/__init__.py | 8 +- .../models/indextts2/__init__.py | 6 + .../model_executor/models/lance/__init__.py | 11 +- .../models/mammoth_moda2/__init__.py | 23 +- .../models/ming_flash_omni/__init__.py | 24 +- .../models/ming_tts/__init__.py | 15 +- .../models/minicpmo_4_5/__init__.py | 15 +- .../models/moss_tts/__init__.py | 19 +- .../models/moss_tts_nano/__init__.py | 7 +- .../models/qwen3_omni/__init__.py | 9 +- .../models/qwen3_omni/pipeline.py | 38 ++ .../qwen3_omni/qwen3_omni_moe_talker.py | 3 - .../model_executor/models/voxcpm2/__init__.py | 7 +- .../transformers_utils/configs/__init__.py | 3 + .../configs/higgs_audio_v3.py} | 0 .../configs/mammoth_moda2.py | 7 +- 44 files changed, 1038 insertions(+), 1126 deletions(-) create mode 100644 tests/helpers/fixtures/config.py delete mode 100644 tests/model_executor/models/glm_image/test_glm_image_init.py create mode 100644 vllm_omni/config/config_factory.py create mode 100644 vllm_omni/model_executor/models/aura_omni/__init__.py rename vllm_omni/{model_executor/models/higgs_audio_v3/configuration_higgs_audio_v3.py => transformers_utils/configs/higgs_audio_v3.py} (100%) diff --git a/tests/config/test_pipeline_registry.py b/tests/config/test_pipeline_registry.py index ecd739857..4bc19b116 100644 --- a/tests/config/test_pipeline_registry.py +++ b/tests/config/test_pipeline_registry.py @@ -1,146 +1,68 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Tests for the central pipeline registry (2.5/N).""" - -from __future__ import annotations - -import logging +"""Tests for out of tree registration to OMNI_PIPELINES.""" import pytest +from transformers import PretrainedConfig -from vllm_omni.config.pipeline_registry import _OMNI_PIPELINES -from vllm_omni.config.stage_config import ( - _PIPELINE_REGISTRY, - PipelineConfig, - StageExecutionType, - StagePipelineConfig, - register_pipeline, -) +from vllm_omni.config.pipeline_registry import OMNI_PIPELINES, register_pipeline +from vllm_omni.config.stage_config import PipelineConfig, pipeline_cfg_resolver + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] -class TestCentralRegistryDeclarations: - """Every in-tree pipeline must be declared exactly once in the central registry.""" +@pytest.fixture +def custom_resolver(): + """Build a reusable custom resolver for PipelineConfigs.""" - def test_omni_entries_visible_in_registry(self): - for key in _OMNI_PIPELINES: - assert key in _PIPELINE_REGISTRY + class CustomConfigType(PretrainedConfig): + pass - def test_expected_omni_pipelines_present(self): - # Guard against accidental removal during future refactors. - assert "qwen2_5_omni" in _OMNI_PIPELINES - assert "qwen2_5_omni_thinker_only" in _OMNI_PIPELINES - assert "qwen3_omni_moe" in _OMNI_PIPELINES - assert "aura_omni" in _OMNI_PIPELINES - assert "qwen3_tts" in _OMNI_PIPELINES + @pipeline_cfg_resolver(config_type=CustomConfigType) + def custom_resolver( + hf_config: CustomConfigType, + ) -> PipelineConfig: + return PipelineConfig(model_type="resolved_type") + + return custom_resolver -class TestLazyLoading: - """Pipelines are imported only on first access.""" - - def test_contains_without_import(self): - # ``in`` hits the lazy map, not the loaded cache. - assert "qwen3_omni_moe" in _PIPELINE_REGISTRY - - def test_getitem_loads_correct_pipeline(self): - pipeline = _PIPELINE_REGISTRY["qwen3_omni_moe"] - assert pipeline.model_type == "qwen3_omni_moe" - assert pipeline.model_arch == "Qwen3OmniMoeForConditionalGeneration" - - def test_getitem_loads_aura_omni_pipeline(self): - pipeline = _PIPELINE_REGISTRY["aura_omni"] - assert pipeline.model_type == "aura_omni" - assert [stage.model_stage for stage in pipeline.stages] == [ - "asr", - "aura", - "qwen3_tts", - "code2wav", - ] - assert pipeline.stages[0].model_arch == "Qwen3ASRForConditionalGeneration" - assert pipeline.stages[1].model_arch == "AuraQwen3VLForConditionalGeneration" - assert pipeline.stages[2].model_arch == "Qwen3TTSTalkerForConditionalGeneration" - assert pipeline.stages[3].model_arch == "Qwen3TTSCode2Wav" - - def test_unknown_model_type_returns_none_via_get(self): - assert _PIPELINE_REGISTRY.get("not_a_real_pipeline") is None - - def test_unknown_model_type_raises_keyerror_via_getitem(self): - with pytest.raises(KeyError): - _PIPELINE_REGISTRY["not_a_real_pipeline"] - - def test_iteration_yields_registered_pipelines(self): - keys = set(_PIPELINE_REGISTRY) - assert "qwen2_5_omni" in keys - assert "qwen3_omni_moe" in keys - - def test_iteration_resilience(self, monkeypatch, caplog): - # Synthetic failure to verify iteration robustness against load errors. - target_logger = logging.getLogger("vllm_omni.config.stage_config") - target_logger.addHandler(caplog.handler) - prev_level = target_logger.level - target_logger.setLevel(logging.ERROR) - try: - monkeypatch.setattr( - _PIPELINE_REGISTRY, - "_lazy_map", - { - "healthy": ("vllm_omni.model_executor.models.qwen2_5_omni.pipeline", "QWEN2_5_OMNI_PIPELINE"), - "broken": ("vllm_omni.non_existent_module", "SomeConfig"), - "attr_fail": ("vllm_omni.config.stage_config", "NON_EXISTENT_VAR"), - }, - ) - monkeypatch.setattr(_PIPELINE_REGISTRY, "_loaded", {}) - - # Assert that list(_PIPELINE_REGISTRY.values()) and dict(_PIPELINE_REGISTRY.items()) - # both yield only healthy entries. - healthy_instance = _PIPELINE_REGISTRY["healthy"] - assert list(_PIPELINE_REGISTRY.values()) == [healthy_instance] - assert dict(_PIPELINE_REGISTRY.items()) == {"healthy": healthy_instance} - assert "Failed to import pipeline module" in caplog.text - assert "Pipeline variable 'NON_EXISTENT_VAR' not found" in caplog.text - - finally: - target_logger.removeHandler(caplog.handler) - target_logger.setLevel(prev_level) +def test_register_pipeline_config(clean_pipeline_registry): + """Ensure that we can register a custom pipeline config to OMNI_PIPELINES.""" + new_model_type = "new_model_type" + pipe_cfg = PipelineConfig(model_type=new_model_type) + assert new_model_type not in OMNI_PIPELINES + register_pipeline(pipe_cfg) + assert new_model_type in OMNI_PIPELINES + assert OMNI_PIPELINES[new_model_type] is pipe_cfg -class TestDynamicRegistration: - """``register_pipeline()`` still works for plugins and tests.""" +def test_register_pipeline_config_with_model_type(clean_pipeline_registry): + """Ensure that we can register a custom pipeline config with an explicit model_type to OMNI_PIPELINES.""" + new_model_type = "new_model_type" + unused_model_type = "foo" + pipe_cfg = PipelineConfig(model_type=unused_model_type) + assert new_model_type not in OMNI_PIPELINES + assert unused_model_type not in OMNI_PIPELINES - def test_register_adds_to_registry(self): - custom = PipelineConfig( - model_type="_test_dynamic_registration", - model_arch="DynamicTestModel", - stages=( - StagePipelineConfig( - stage_id=0, - model_stage="test", - execution_type=StageExecutionType.LLM_AR, - input_sources=(), - final_output=True, - ), - ), - ) - register_pipeline(custom) - try: - assert "_test_dynamic_registration" in _PIPELINE_REGISTRY - assert _PIPELINE_REGISTRY["_test_dynamic_registration"] is custom - finally: - # Don't leak the test registration into other tests. - if "_test_dynamic_registration" in _PIPELINE_REGISTRY: - del _PIPELINE_REGISTRY["_test_dynamic_registration"] + # Registering with an explicitly provided model_type uses + # the passed value instead of the pipeline_cfg.model_type + register_pipeline(pipe_cfg, new_model_type) + assert new_model_type in OMNI_PIPELINES + assert unused_model_type not in OMNI_PIPELINES + assert OMNI_PIPELINES[new_model_type] is pipe_cfg - def test_dynamic_registration_overrides_lazy_entry(self): - # Build a substitute for qwen3_omni_moe that we can distinguish. - original = _PIPELINE_REGISTRY["qwen3_omni_moe"] - override = PipelineConfig( - model_type="qwen3_omni_moe", - model_arch="OverriddenArch", - stages=original.stages, - ) - register_pipeline(override) - try: - assert _PIPELINE_REGISTRY["qwen3_omni_moe"].model_arch == "OverriddenArch" - finally: - # Remove the dynamic override so later tests see the original. - if "qwen3_omni_moe" in _PIPELINE_REGISTRY._loaded: - del _PIPELINE_REGISTRY["qwen3_omni_moe"] + +def test_register_resolver(custom_resolver, clean_pipeline_registry): + """Ensure that we can register a custom resolver to OMNI_PIPELINES.""" + new_model_type = "new_model_type" + assert new_model_type not in OMNI_PIPELINES + register_pipeline(custom_resolver, new_model_type) + assert new_model_type in OMNI_PIPELINES + assert OMNI_PIPELINES[new_model_type] is custom_resolver + + +def test_register_resolver_requires_model_type(custom_resolver, clean_pipeline_registry): + """Ensure that registering a custom resolver to OMNI_PIPELINES requires an explicit model_type.""" + with pytest.raises(ValueError): + register_pipeline(custom_resolver) diff --git a/tests/conftest.py b/tests/conftest.py index 163309b8c..11619b782 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -17,6 +17,7 @@ from tests.model_executor.helpers import bootstrap_vllm_layer_custom_op_modules bootstrap_vllm_layer_custom_op_modules() pytest_plugins = ( + "tests.helpers.fixtures.config", "tests.helpers.fixtures.env", "tests.helpers.fixtures.log", "tests.helpers.fixtures.run_args", diff --git a/tests/e2e/offline_inference/test_qwen3_omni.py b/tests/e2e/offline_inference/test_qwen3_omni.py index 22225e8b2..fd3a27d2e 100644 --- a/tests/e2e/offline_inference/test_qwen3_omni.py +++ b/tests/e2e/offline_inference/test_qwen3_omni.py @@ -14,7 +14,7 @@ from tests.helpers.stage_config import get_deploy_config_path, modify_stage_conf from vllm_omni.platforms import current_omni_platform models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"] - +thinker_only_models = ["Qwen/Qwen3-Omni-30B-A3B-Captioner"] # Single CI deploy YAML; rocm/xpu deltas are picked automatically via the # platforms: section. Only CUDA needs an extra enforce_eager tweak. @@ -40,6 +40,10 @@ else: # Create parameter combinations for model and stage config test_params = [(model, stage_config) for model in models for stage_config in stage_configs] +# we can use the same config for a model that only has thinker (i.e., does not +# enable audio output) because the resolver should figure out that it doesn't +# need the full pipeline based on the HF config. +thinker_test_params = [(model, stage_config) for model in thinker_only_models for stage_config in stage_configs] def get_question(prompt_type="video"): @@ -61,3 +65,15 @@ def test_video_to_audio(omni_runner, omni_runner_handler) -> None: # Test single completion omni_runner_handler.send_omni_request(request_config) + + +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=1) +@pytest.mark.parametrize("omni_runner", thinker_test_params, indirect=True) +def test_thinker_only_model_request(omni_runner, omni_runner_handler) -> None: + """Test that we can load and run a request through a model that only has the thinker stage.""" + request_config = {"prompts": "what color is the sky?", "modalities": ["text"]} + + # Test single completion + omni_runner_handler.send_omni_request(request_config) diff --git a/tests/engine/test_async_omni_engine_abort.py b/tests/engine/test_async_omni_engine_abort.py index f810bb80e..794475388 100644 --- a/tests/engine/test_async_omni_engine_abort.py +++ b/tests/engine/test_async_omni_engine_abort.py @@ -9,13 +9,13 @@ from vllm.inputs import PromptType from tests.helpers.mark import hardware_test from tests.helpers.stage_config import get_deploy_config_path -from vllm_omni.config.stage_config import _PIPELINE_REGISTRY +from vllm_omni.config.pipeline_registry import OMNI_PIPELINES from vllm_omni.entrypoints.async_omni import AsyncOmni os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" SEED = 42 -pipeline = _PIPELINE_REGISTRY["qwen2_5_omni_thinker_only"] +pipeline = OMNI_PIPELINES["qwen2_5_omni_thinker_only"] # Single-stage thinker-only deploy, materialized from tests.helpers.stage_config._CI_OVERLAYS. stage_config = get_deploy_config_path("ci/qwen2_5_omni_thinker_only.yaml") diff --git a/tests/entrypoints/test_resolve_dreamzero_config.py b/tests/entrypoints/test_resolve_dreamzero_config.py index a962c5262..374dcf7fa 100644 --- a/tests/entrypoints/test_resolve_dreamzero_config.py +++ b/tests/entrypoints/test_resolve_dreamzero_config.py @@ -24,7 +24,7 @@ def test_dreamzero_vla_resolves_to_dreamzero_config(monkeypatch): def test_dreamzero_config_sets_model_class_and_policy_config(monkeypatch): monkeypatch.setattr( - "vllm_omni.config.stage_config.StageConfigFactory._auto_detect_model_type", + "vllm_omni.config.config_factory.StageConfigFactory._auto_detect_model_type", classmethod(lambda _cls, _model, trust_remote_code=True: ("vla", None)), ) monkeypatch.setattr( diff --git a/tests/helpers/fixtures/config.py b/tests/helpers/fixtures/config.py new file mode 100644 index 000000000..2e309b672 --- /dev/null +++ b/tests/helpers/fixtures/config.py @@ -0,0 +1,14 @@ +"""Fixtures config creation & out of tree registry management.""" + +import pytest + +from vllm_omni.config.pipeline_registry import OMNI_PIPELINES + + +@pytest.fixture +def clean_pipeline_registry(): + """Ensure the OMNI_PIPELINES are in a clean state for a test that mutates it.""" + snapshot = dict(OMNI_PIPELINES) + yield + OMNI_PIPELINES.clear() + OMNI_PIPELINES.update(snapshot) diff --git a/tests/model_executor/models/glm_image/test_glm_image_init.py b/tests/model_executor/models/glm_image/test_glm_image_init.py deleted file mode 100644 index eb42e9649..000000000 --- a/tests/model_executor/models/glm_image/test_glm_image_init.py +++ /dev/null @@ -1,76 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Unit tests for GLM-Image model __init__.py lazy import pattern. - -The ``__init__.py`` uses ``__getattr__`` for lazy loading to avoid importing -``transformers.models.glm_image`` at module init, which may not be available -in all environments. -""" - -import pytest - -pytestmark = [pytest.mark.core_model, pytest.mark.cpu] - - -class TestLazyImport: - """Test the __getattr__ lazy import pattern in glm_image __init__.py.""" - - def test_getattr_exists_and_is_callable(self): - """Verify __getattr__ exists and is callable.""" - import vllm_omni.model_executor.models.glm_image as glm_image_pkg - - assert hasattr(glm_image_pkg, "__getattr__") - assert callable(glm_image_pkg.__getattr__) - - def test_getattr_returns_class_for_known_attribute(self): - """Verify __getattr__ returns GlmImageForConditionalGeneration for known attribute.""" - import vllm_omni.model_executor.models.glm_image as glm_image_pkg - - # Call __getattr__ directly to test the lazy import logic - result = glm_image_pkg.__getattr__("GlmImageForConditionalGeneration") - - # Verify we get a class (the actual GlmImageForConditionalGeneration) - assert result is not None - assert isinstance(result, type) - assert result.__name__ == "GlmImageForConditionalGeneration" - - def test_getattr_raises_for_unknown_attribute(self): - """Verify __getattr__ raises AttributeError for unknown attributes.""" - import vllm_omni.model_executor.models.glm_image as glm_image_pkg - - # Test unknown attribute via __getattr__ directly - with pytest.raises(AttributeError, match="has no attribute"): - glm_image_pkg.__getattr__("UnknownClass") - - def test___all___exports_correct_symbols(self): - """Verify __all__ contains the expected exported symbols.""" - import vllm_omni.model_executor.models.glm_image as glm_image_pkg - - assert hasattr(glm_image_pkg, "__all__") - assert "GlmImageForConditionalGeneration" in glm_image_pkg.__all__ - - -class TestLazyImportDoesNotImportTransformersAtInit: - """Verify that importing the package does not eagerly load transformers.""" - - def test_glm_image_module_has_getattr(self): - """Test that the module has __getattr__ for lazy loading.""" - import vllm_omni.model_executor.models.glm_image as glm_image_pkg - - # The module should have __getattr__ - assert hasattr(glm_image_pkg, "__getattr__") - assert callable(glm_image_pkg.__getattr__) - - # And __all__ - assert hasattr(glm_image_pkg, "__all__") - - def test_module_has_proper_structure(self): - """Test that the module has proper Python module structure.""" - import vllm_omni.model_executor.models.glm_image as glm_image_pkg - - # Should be a module - assert hasattr(glm_image_pkg, "__name__") - assert glm_image_pkg.__name__ == "vllm_omni.model_executor.models.glm_image" - - # Should have __file__ (even if it's a package) - assert hasattr(glm_image_pkg, "__file__") diff --git a/tests/model_executor/models/minicpmo_4_5/test_pipeline.py b/tests/model_executor/models/minicpmo_4_5/test_pipeline.py index c51db9fcc..48db0d497 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_pipeline.py +++ b/tests/model_executor/models/minicpmo_4_5/test_pipeline.py @@ -19,9 +19,8 @@ from types import SimpleNamespace import pytest -from vllm_omni.config.pipeline_registry import _OMNI_PIPELINES +from vllm_omni.config.pipeline_registry import OMNI_PIPELINES from vllm_omni.config.stage_config import ( - _PIPELINE_REGISTRY, PipelineConfig, StageExecutionType, ) @@ -34,13 +33,13 @@ _PIPELINE_KEY = "minicpmo_4_5" class TestRegistryDeclaration: def test_declared_in_omni_pipelines(self) -> None: - assert _PIPELINE_KEY in _OMNI_PIPELINES + assert _PIPELINE_KEY in OMNI_PIPELINES def test_visible_in_central_registry(self) -> None: - assert _PIPELINE_KEY in _PIPELINE_REGISTRY + assert _PIPELINE_KEY in OMNI_PIPELINES def test_lazy_load_returns_pipeline_config(self) -> None: - pipeline = _PIPELINE_REGISTRY[_PIPELINE_KEY] + pipeline = OMNI_PIPELINES[_PIPELINE_KEY] assert isinstance(pipeline, PipelineConfig) assert pipeline.model_type == _PIPELINE_KEY assert pipeline.model_arch == "MiniCPMO45OmniForConditionalGeneration" @@ -49,7 +48,7 @@ class TestRegistryDeclaration: class TestPipelineTopology: @pytest.fixture(scope="class") def pipeline(self) -> PipelineConfig: - return _PIPELINE_REGISTRY[_PIPELINE_KEY] + return OMNI_PIPELINES[_PIPELINE_KEY] def test_two_stages(self, pipeline: PipelineConfig) -> None: assert len(pipeline.stages) == 2 @@ -99,7 +98,7 @@ class TestArchAliases: @pytest.fixture(scope="class") def pipeline(self) -> PipelineConfig: - return _PIPELINE_REGISTRY[_PIPELINE_KEY] + return OMNI_PIPELINES[_PIPELINE_KEY] def test_shared_minicpmo_alias_present(self, pipeline: PipelineConfig) -> None: # MiniCPM-o 4.5 ships ``architectures=["MiniCPMO"]`` in its HF config. @@ -123,7 +122,7 @@ class TestHfConfigPredicate: @pytest.fixture(scope="class") def predicate(self): - pipeline = _PIPELINE_REGISTRY[_PIPELINE_KEY] + pipeline = OMNI_PIPELINES[_PIPELINE_KEY] assert pipeline.hf_config_predicate is not None, ( "MiniCPM-o 4.5 pipeline must declare hf_config_predicate to " "avoid misrouting MiniCPM-o 2.6 checkpoints into the 4.5 path." diff --git a/tests/test_config_factory.py b/tests/test_config_factory.py index 45a4ccfab..865315ad4 100644 --- a/tests/test_config_factory.py +++ b/tests/test_config_factory.py @@ -12,14 +12,15 @@ from pathlib import Path from unittest.mock import patch import pytest +from transformers import PretrainedConfig, Qwen3OmniMoeConfig from tests.helpers.stage_config import get_deploy_config_path +from vllm_omni.config.config_factory import StageConfigFactory +from vllm_omni.config.pipeline_registry import OMNI_PIPELINES, register_pipeline from vllm_omni.config.stage_config import ( - _PIPELINE_REGISTRY, DeployConfig, PipelineConfig, StageConfig, - StageConfigFactory, StageDeployConfig, StageExecutionType, StagePipelineConfig, @@ -30,14 +31,16 @@ from vllm_omni.config.stage_config import ( build_stage_runtime_overrides, load_deploy_config, merge_pipeline_deploy, - register_pipeline, + pipeline_cfg_resolver, strip_parent_engine_args, ) from vllm_omni.engine.arg_utils import SHARED_FIELDS, EngineArgs, internal_blacklist_keys pytestmark = [pytest.mark.core_model, pytest.mark.cpu] -pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +Q3_OMNI_ALL_STAGES_HF_CONFIG = Qwen3OmniMoeConfig(enable_audio_output=True) +Q3_OMNI_THINKER_HF_CONFIG = Qwen3OmniMoeConfig(enable_audio_output=False) class TestStageType: @@ -511,26 +514,39 @@ class TestStageResolutionHelpers: class TestPipelineDiscovery: - """Tests for the central pipeline registry (``pipeline_registry._OMNI_PIPELINES``).""" + """Tests for the central pipeline registry (``OMNI_PIPELINES``).""" def test_registry_has_known_models(self): - """Built-in pipelines are lazy-loaded from the central declaration - on first access; no eager import or discovery walk needed.""" - # ``in`` triggers the lazy-map lookup without forcing a load. - assert "qwen2_5_omni" in _PIPELINE_REGISTRY - assert "qwen3_omni_moe" in _PIPELINE_REGISTRY - assert "qwen3_tts" in _PIPELINE_REGISTRY + """Check that specific models are in OMNI_PIPELINES.""" + assert "qwen2_5_omni" in OMNI_PIPELINES + assert "qwen3_omni_moe" in OMNI_PIPELINES + assert "qwen3_tts" in OMNI_PIPELINES - def test_registry_loads_pipeline_on_getitem(self): - """Looking up a registered model_type returns the matching PipelineConfig.""" - pipeline = _PIPELINE_REGISTRY["qwen3_omni_moe"] + def test_registry_resolver_qwen3_omni_all_stages(self): + """Test that providing the HF config for qwen3 omni with audio enabled uses all stages.""" + pipeline = StageConfigFactory.resolve_pipeline_config( + "qwen3_omni_moe", + Q3_OMNI_ALL_STAGES_HF_CONFIG, + ) + assert isinstance(pipeline, PipelineConfig) assert pipeline.model_type == "qwen3_omni_moe" assert len(pipeline.stages) == 3 # thinker + talker + code2wav + def test_registry_resolver_qwen3_omni_thinker_only(self): + """Test that providing the HF config for qwen3 omni without audio is thinker only.""" + pipeline = StageConfigFactory.resolve_pipeline_config( + "qwen3_omni_moe", + Q3_OMNI_THINKER_HF_CONFIG, + ) + assert isinstance(pipeline, PipelineConfig) + assert pipeline.model_type == "qwen3_omni_moe_thinker_only" + assert len(pipeline.stages) == 1 # thinker only + def test_registry_returns_none_for_unknown(self): - """Unknown model_types aren't found; ``get()`` returns None.""" - assert "definitely_not_a_real_model" not in _PIPELINE_REGISTRY - assert _PIPELINE_REGISTRY.get("definitely_not_a_real_model") is None + """Unknown model_types aren't found and resolve to `None`.""" + assert "definitely_not_a_real_model" not in OMNI_PIPELINES + assert OMNI_PIPELINES.get("definitely_not_a_real_model") is None + assert StageConfigFactory.resolve_pipeline_config("definitely_not_a_real_model") is None def test_pipeline_config_supports_hf_architectures(self): """PipelineConfig accepts hf_architectures for HF-arch fallback @@ -579,6 +595,67 @@ class TestPipelineConfigNew: assert any("no stages" in e.lower() for e in p.validate()) +class TestPipelineRegistration: + def test_pipeline_registration(self, clean_pipeline_registry): + """Ensure that we can register and create a custom pipeline config.""" + new_model_type = "new_model_type" + pipe_cfg = PipelineConfig(model_type=new_model_type) + + # Register the new PipelineConfig + assert new_model_type not in OMNI_PIPELINES + register_pipeline(pipe_cfg) + assert new_model_type in OMNI_PIPELINES + assert OMNI_PIPELINES[new_model_type] is pipe_cfg + + class FakeConfig(PretrainedConfig): + model_type = new_model_type + + # Create the model + with ( + patch("vllm_omni.config.config_factory.get_config", return_value=FakeConfig()), + patch.object(StageConfigFactory, "_create_from_registry") as mock_create, + ): + StageConfigFactory.create_from_model("fake/model") + mock_create.assert_called_once() + # Ensure that the passed pipeline_config is the right type + pipeline_cfg = mock_create.call_args[0][1] + assert isinstance(pipeline_cfg, PipelineConfig) + assert pipe_cfg.model_type == new_model_type + + def test_resolver_registration(self, clean_pipeline_registry): + """Ensure that we can register and create a custom resolver.""" + new_model_type = "new_model_type" + resolved_type = "resolved_type" + + class FakeConfig(PretrainedConfig): + model_type = new_model_type + + @pipeline_cfg_resolver(config_type=FakeConfig) + def custom_resolver( + hf_config: FakeConfig, + ) -> PipelineConfig: + return PipelineConfig(model_type=resolved_type) + + # Register the new PipelineConfig + assert new_model_type not in OMNI_PIPELINES + + register_pipeline(custom_resolver, model_type=new_model_type) + assert new_model_type in OMNI_PIPELINES + assert OMNI_PIPELINES[new_model_type] is custom_resolver + + # Create the model + with ( + patch("vllm_omni.config.config_factory.get_config", return_value=FakeConfig()), + patch.object(StageConfigFactory, "_create_from_registry") as mock_create, + ): + StageConfigFactory.create_from_model("fake/model") + mock_create.assert_called_once() + # Ensure that the passed pipeline_config is the right type + pipeline_cfg = mock_create.call_args[0][1] + assert isinstance(pipeline_cfg, PipelineConfig) + assert pipeline_cfg.model_type == resolved_type + + class TestResolveScheduler: def test_all_execution_types_handled(self): for et in StageExecutionType: @@ -603,18 +680,6 @@ class TestResolveScheduler: assert _resolve_scheduler(StageExecutionType.DIFFUSION) is None -class TestPipelineRegistry: - def test_register_and_lookup(self): - p = PipelineConfig( - model_type="__test_only__", - model_arch="A", - stages=(StagePipelineConfig(stage_id=0, model_stage="a"),), - ) - register_pipeline(p) - assert _PIPELINE_REGISTRY["__test_only__"] is p - del _PIPELINE_REGISTRY["__test_only__"] - - class TestDeployConfigLoading: def test_custom_voice_dir_is_pipeline_wide_engine_arg(self, tmp_path): deploy_path = tmp_path / "qwen3_tts_custom_voice.yaml" @@ -633,8 +698,9 @@ stages: ) deploy = load_deploy_config(deploy_path) - pipeline = _PIPELINE_REGISTRY["qwen3_tts"] + pipeline = StageConfigFactory.resolve_pipeline_config("qwen3_tts") stages = merge_pipeline_deploy(pipeline, deploy) + assert isinstance(pipeline, PipelineConfig) assert deploy.custom_voice_dir == str(custom_voice_dir) assert {s.yaml_engine_args.get("custom_voice_dir") for s in stages} == {str(custom_voice_dir)} @@ -664,11 +730,16 @@ stages: assert "compilation_config" not in deploy.stages[0].engine_extras def test_merge_pipeline_deploy(self): - pipeline = _PIPELINE_REGISTRY["qwen3_omni_moe"] deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml" if not deploy_path.exists(): pytest.skip("Deploy config not found") + pipeline = StageConfigFactory.resolve_pipeline_config( + "qwen3_omni_moe", + Q3_OMNI_ALL_STAGES_HF_CONFIG, + ) + assert isinstance(pipeline, PipelineConfig) + deploy = load_deploy_config(deploy_path) stages = merge_pipeline_deploy(pipeline, deploy) @@ -680,7 +751,12 @@ stages: assert s0.yaml_extras["default_sampling_params"]["detokenize"] is True def test_merge_pipeline_deploy_preserves_num_replicas(self, tmp_path): - pipeline = _PIPELINE_REGISTRY["qwen3_omni_moe"] + pipeline = StageConfigFactory.resolve_pipeline_config( + "qwen3_omni_moe", + Q3_OMNI_ALL_STAGES_HF_CONFIG, + ) + assert isinstance(pipeline, PipelineConfig) + base = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml" if not base.exists(): pytest.skip("Deploy config not found") @@ -894,14 +970,24 @@ stages: class TestQwen3OmniPipeline: def test_registered(self): - p = _PIPELINE_REGISTRY.get("qwen3_omni_moe") - assert p is not None + p = StageConfigFactory.resolve_pipeline_config( + "qwen3_omni_moe", + Q3_OMNI_ALL_STAGES_HF_CONFIG, + ) + assert isinstance(p, PipelineConfig) assert p.model_arch == "Qwen3OmniMoeForConditionalGeneration" assert len(p.stages) == 3 assert p.validate() == [] def test_thinker(self): - s = _PIPELINE_REGISTRY["qwen3_omni_moe"].get_stage(0) + p = StageConfigFactory.resolve_pipeline_config( + "qwen3_omni_moe", + Q3_OMNI_ALL_STAGES_HF_CONFIG, + ) + assert isinstance(p, PipelineConfig) + + s = p.get_stage(0) + assert isinstance(s, StagePipelineConfig) assert s.model_stage == "thinker" assert s.execution_type == StageExecutionType.LLM_AR assert s.owns_tokenizer is True @@ -909,14 +995,28 @@ class TestQwen3OmniPipeline: assert s.sampling_constraints["detokenize"] is True def test_talker(self): - s = _PIPELINE_REGISTRY["qwen3_omni_moe"].get_stage(1) + p = StageConfigFactory.resolve_pipeline_config( + "qwen3_omni_moe", + Q3_OMNI_ALL_STAGES_HF_CONFIG, + ) + assert isinstance(p, PipelineConfig) + + s = p.get_stage(1) + assert isinstance(s, StagePipelineConfig) assert s.input_sources == (0,) assert s.sampling_constraints["stop_token_ids"] == [2150] assert s.custom_process_input_func is not None assert s.custom_process_next_stage_input_func is not None def test_code2wav(self): - s = _PIPELINE_REGISTRY["qwen3_omni_moe"].get_stage(2) + p = StageConfigFactory.resolve_pipeline_config( + "qwen3_omni_moe", + Q3_OMNI_ALL_STAGES_HF_CONFIG, + ) + assert isinstance(p, PipelineConfig) + + s = p.get_stage(2) + assert isinstance(s, StagePipelineConfig) assert s.execution_type == StageExecutionType.LLM_GENERATION assert s.final_output_type == "audio" assert s.custom_process_input_func is not None @@ -924,14 +1024,18 @@ class TestQwen3OmniPipeline: class TestQwen2_5OmniPipeline: def test_registered(self): - p = _PIPELINE_REGISTRY.get("qwen2_5_omni") - assert p is not None + p = StageConfigFactory.resolve_pipeline_config("qwen2_5_omni") + assert isinstance(p, PipelineConfig) assert p.model_arch == "Qwen2_5OmniForConditionalGeneration" assert len(p.stages) == 3 assert p.validate() == [] def test_thinker(self): - s = _PIPELINE_REGISTRY["qwen2_5_omni"].get_stage(0) + p = StageConfigFactory.resolve_pipeline_config("qwen2_5_omni") + assert isinstance(p, PipelineConfig) + + s = p.get_stage(0) + assert isinstance(s, StagePipelineConfig) assert s.model_stage == "thinker" assert s.execution_type == StageExecutionType.LLM_AR assert s.owns_tokenizer is True @@ -939,13 +1043,21 @@ class TestQwen2_5OmniPipeline: assert s.requires_multimodal_data is True def test_talker(self): - s = _PIPELINE_REGISTRY["qwen2_5_omni"].get_stage(1) + p = StageConfigFactory.resolve_pipeline_config("qwen2_5_omni") + assert isinstance(p, PipelineConfig) + + s = p.get_stage(1) + assert isinstance(s, StagePipelineConfig) assert s.input_sources == (0,) assert s.sampling_constraints["stop_token_ids"] == [8294] assert s.custom_process_input_func is not None def test_code2wav(self): - s = _PIPELINE_REGISTRY["qwen2_5_omni"].get_stage(2) + p = StageConfigFactory.resolve_pipeline_config("qwen2_5_omni") + assert isinstance(p, PipelineConfig) + + s = p.get_stage(2) + assert isinstance(s, StagePipelineConfig) assert s.execution_type == StageExecutionType.LLM_GENERATION assert s.final_output_type == "audio" assert s.engine_output_type == "audio" @@ -953,14 +1065,19 @@ class TestQwen2_5OmniPipeline: class TestQwen3TTSPipeline: def test_registered(self): - p = _PIPELINE_REGISTRY.get("qwen3_tts") + p = StageConfigFactory.resolve_pipeline_config("qwen3_tts") + assert isinstance(p, PipelineConfig) assert p is not None assert p.model_arch == "Qwen3TTSTalkerForConditionalGeneration" assert len(p.stages) == 2 assert p.validate() == [] def test_talker_stage(self): - s = _PIPELINE_REGISTRY["qwen3_tts"].get_stage(0) + p = StageConfigFactory.resolve_pipeline_config("qwen3_tts") + assert isinstance(p, PipelineConfig) + + s = p.get_stage(0) + assert isinstance(s, StagePipelineConfig) assert s.model_stage == "qwen3_tts" assert s.execution_type == StageExecutionType.LLM_AR assert s.owns_tokenizer is True @@ -970,7 +1087,11 @@ class TestQwen3TTSPipeline: assert s.model_arch is None def test_code2wav_stage_has_per_stage_model_arch(self): - s = _PIPELINE_REGISTRY["qwen3_tts"].get_stage(1) + p = StageConfigFactory.resolve_pipeline_config("qwen3_tts") + assert isinstance(p, PipelineConfig) + + s = p.get_stage(1) + assert isinstance(s, StagePipelineConfig) assert s.execution_type == StageExecutionType.LLM_GENERATION assert s.final_output_type == "audio" assert s.engine_output_type == "audio" @@ -986,7 +1107,7 @@ class TestQwen3TTSPipeline: pytest.skip("qwen3_tts deploy yaml not found") deploy = load_deploy_config(deploy_path) - pipeline = _PIPELINE_REGISTRY["qwen3_tts"] + pipeline = OMNI_PIPELINES["qwen3_tts"] stages = merge_pipeline_deploy(pipeline, deploy) # Stage 0 inherits pipeline-level model_arch @@ -1025,14 +1146,18 @@ class TestQwen3TTSPipeline: class TestMingFlashOmniPipeline: def test_registered(self): - p = _PIPELINE_REGISTRY.get("ming_flash_omni") - assert p is not None + p = StageConfigFactory.resolve_pipeline_config("ming_flash_omni") + assert isinstance(p, PipelineConfig) assert p.model_arch == "MingFlashOmniForConditionalGeneration" assert len(p.stages) == 2 assert p.validate() == [] def test_thinker_stage(self): - s = _PIPELINE_REGISTRY["ming_flash_omni"].get_stage(0) + p = StageConfigFactory.resolve_pipeline_config("ming_flash_omni") + assert isinstance(p, PipelineConfig) + + s = p.get_stage(0) + assert isinstance(s, StagePipelineConfig) assert s.model_stage == "thinker" assert s.execution_type == StageExecutionType.LLM_AR assert s.owns_tokenizer is True @@ -1042,7 +1167,11 @@ class TestMingFlashOmniPipeline: assert s.sampling_constraints["detokenize"] is True def test_talker_stage(self): - s = _PIPELINE_REGISTRY["ming_flash_omni"].get_stage(1) + p = StageConfigFactory.resolve_pipeline_config("ming_flash_omni") + assert isinstance(p, PipelineConfig) + + s = p.get_stage(1) + assert isinstance(s, StagePipelineConfig) assert s.model_stage == "ming_tts" assert s.execution_type == StageExecutionType.LLM_GENERATION assert s.input_sources == (0,) @@ -1060,20 +1189,28 @@ class TestMingFlashOmniPipeline: Lazy string references only fail at first inference otherwise — this catches typos in the pipeline declaration at import / registration time. """ - s = _PIPELINE_REGISTRY["ming_flash_omni"].get_stage(1) + p = StageConfigFactory.resolve_pipeline_config("ming_flash_omni") + assert isinstance(p, PipelineConfig) + + s = p.get_stage(1) + assert isinstance(s, StagePipelineConfig) module_path, _, attr = s.custom_process_input_func.rpartition(".") module = importlib.import_module(module_path) assert callable(getattr(module, attr)) def test_tts_pipeline_registered(self): - p = _PIPELINE_REGISTRY.get("ming_flash_omni_tts") - assert p is not None + p = StageConfigFactory.resolve_pipeline_config("ming_flash_omni_tts") + assert isinstance(p, PipelineConfig) assert p.model_arch == "MingFlashOmniTalkerForConditionalGeneration" assert len(p.stages) == 1 assert p.validate() == [] def test_tts_stage(self): - s = _PIPELINE_REGISTRY["ming_flash_omni_tts"].get_stage(0) + p = StageConfigFactory.resolve_pipeline_config("ming_flash_omni_tts") + assert isinstance(p, PipelineConfig) + + s = p.get_stage(0) + assert isinstance(s, StagePipelineConfig) assert s.model_stage == "ming_tts" assert s.execution_type == StageExecutionType.LLM_GENERATION assert s.input_sources == () @@ -1096,7 +1233,8 @@ class TestMingFlashOmniPipeline: # We won't test stage 0/1 colocation contract here, # as there could exist more variant of custom device setup - pipeline = _PIPELINE_REGISTRY["ming_flash_omni"] + pipeline = StageConfigFactory.resolve_pipeline_config("ming_flash_omni") + assert isinstance(pipeline, PipelineConfig) stages = merge_pipeline_deploy(pipeline, deploy) assert len(stages) == 2 assert stages[0].yaml_engine_args["model_arch"] == "MingFlashOmniForConditionalGeneration" @@ -1112,20 +1250,24 @@ class TestMingFlashOmniPipeline: assert len(deploy.stages) == 1 assert deploy.pipeline == "ming_flash_omni_tts" - pipeline = _PIPELINE_REGISTRY["ming_flash_omni_tts"] + pipeline = OMNI_PIPELINES["ming_flash_omni_tts"] stages = merge_pipeline_deploy(pipeline, deploy) assert len(stages) == 1 assert stages[0].yaml_engine_args["model_arch"] == "MingFlashOmniTalkerForConditionalGeneration" def test_thinker_only_pipeline_registered(self): - p = _PIPELINE_REGISTRY.get("ming_flash_omni_thinker_only") - assert p is not None + p = StageConfigFactory.resolve_pipeline_config("ming_flash_omni_thinker_only") + assert isinstance(p, PipelineConfig) assert p.model_arch == "MingFlashOmniForConditionalGeneration" assert len(p.stages) == 1 assert p.validate() == [] def test_thinker_only_stage(self): - s = _PIPELINE_REGISTRY["ming_flash_omni_thinker_only"].get_stage(0) + p = StageConfigFactory.resolve_pipeline_config("ming_flash_omni_thinker_only") + assert isinstance(p, PipelineConfig) + + s = p.get_stage(0) + assert isinstance(s, StagePipelineConfig) assert s.model_stage == "thinker" assert s.execution_type == StageExecutionType.LLM_AR assert s.input_sources == () @@ -1146,20 +1288,21 @@ class TestMingFlashOmniPipeline: assert len(deploy.stages) == 1 assert deploy.pipeline == "ming_flash_omni_thinker_only" - pipeline = _PIPELINE_REGISTRY["ming_flash_omni_thinker_only"] + pipeline = StageConfigFactory.resolve_pipeline_config("ming_flash_omni_thinker_only") + assert isinstance(pipeline, PipelineConfig) stages = merge_pipeline_deploy(pipeline, deploy) assert len(stages) == 1 assert stages[0].yaml_engine_args["model_arch"] == "MingFlashOmniForConditionalGeneration" def test_image_pipeline_registered(self): - p = _PIPELINE_REGISTRY.get("ming_flash_omni_image") + p = OMNI_PIPELINES.get("ming_flash_omni_image") assert p is not None assert p.model_arch == "MingFlashOmniForConditionalGeneration" assert len(p.stages) == 2 assert p.validate() == [] def test_image_thinker_stage(self): - s = _PIPELINE_REGISTRY["ming_flash_omni_image"].get_stage(0) + s = StageConfigFactory.resolve_pipeline_config("ming_flash_omni_image").get_stage(0) assert s.model_stage == "thinker" assert s.execution_type == StageExecutionType.LLM_AR assert s.input_sources == () @@ -1173,7 +1316,7 @@ class TestMingFlashOmniPipeline: assert s.prompt_expand_func is not None def test_image_dit_stage(self): - s = _PIPELINE_REGISTRY["ming_flash_omni_image"].get_stage(1) + s = StageConfigFactory.resolve_pipeline_config("ming_flash_omni_image").get_stage(1) assert s.model_stage == "dit" assert s.execution_type == StageExecutionType.DIFFUSION assert s.input_sources == (0,) @@ -1185,8 +1328,11 @@ class TestMingFlashOmniPipeline: def test_image_processor_wiring_resolves(self): """The prompt_expand_func and custom_process_input_func strings must point to real callables.""" - thinker = _PIPELINE_REGISTRY["ming_flash_omni_image"].get_stage(0) - dit = _PIPELINE_REGISTRY["ming_flash_omni_image"].get_stage(1) + pipeline = StageConfigFactory.resolve_pipeline_config("ming_flash_omni_image") + assert isinstance(pipeline, PipelineConfig) + + thinker = pipeline.get_stage(0) + dit = pipeline.get_stage(1) for ref in (thinker.prompt_expand_func, dit.custom_process_input_func): module_path, _, attr = ref.rpartition(".") module = importlib.import_module(module_path) @@ -1205,7 +1351,8 @@ class TestMingFlashOmniPipeline: assert deploy.connectors is not None assert "shared_memory_connector" in deploy.connectors - pipeline = _PIPELINE_REGISTRY["ming_flash_omni_image"] + pipeline = StageConfigFactory.resolve_pipeline_config("ming_flash_omni_image") + assert isinstance(pipeline, PipelineConfig) stages = merge_pipeline_deploy(pipeline, deploy) assert len(stages) == 2 # Stage 0 thinker: AR worker that emits latents. @@ -1356,12 +1503,16 @@ class TestCLIOverrideFlow: """Test --stage-overrides JSON merge into StageConfig.""" def test_stage_overrides_merge(self): - pipeline = _PIPELINE_REGISTRY["qwen3_omni_moe"] deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml" if not deploy_path.exists(): pytest.skip("Deploy config not found") deploy = load_deploy_config(deploy_path) + pipeline = StageConfigFactory.resolve_pipeline_config( + "qwen3_omni_moe", + Q3_OMNI_ALL_STAGES_HF_CONFIG, + ) + assert isinstance(pipeline, PipelineConfig) stages = merge_pipeline_deploy(pipeline, deploy) # Simulate --stage-overrides '{"0": {"gpu_memory_utilization": 0.5}}' @@ -1370,12 +1521,16 @@ class TestCLIOverrideFlow: assert stages[0].runtime_overrides["gpu_memory_utilization"] == 0.5 def test_global_override_applies_to_all(self): - pipeline = _PIPELINE_REGISTRY["qwen3_omni_moe"] deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml" if not deploy_path.exists(): pytest.skip("Deploy config not found") deploy = load_deploy_config(deploy_path) + pipeline = StageConfigFactory.resolve_pipeline_config( + "qwen3_omni_moe", + Q3_OMNI_ALL_STAGES_HF_CONFIG, + ) + assert isinstance(pipeline, PipelineConfig) stages = merge_pipeline_deploy(pipeline, deploy) overrides = {"enforce_eager": True} @@ -1392,8 +1547,11 @@ class TestAuraOmniDeploy: assert deploy.pipeline == "aura_omni" def test_aura_omni_deploy_resolves_four_native_stages(self): + pipeline_cfg = StageConfigFactory.resolve_pipeline_config("aura_omni") + stages = StageConfigFactory._create_from_registry( "qwen3_tts", + pipeline_cfg, cli_overrides={}, deploy_config_path=str(Path(__file__).parent.parent / "vllm_omni" / "deploy" / "aura_omni.yaml"), ) @@ -1417,8 +1575,14 @@ class TestSentinelDefaultPrecedence: to YAML / dataclass defaults (#3035).""" def _stages(self, cli_overrides): + model_type = "qwen3_omni_moe" + pipeline_cfg = StageConfigFactory.resolve_pipeline_config( + model_type, + Q3_OMNI_ALL_STAGES_HF_CONFIG, + ) return StageConfigFactory._create_from_registry( "qwen3_omni_moe", + pipeline_cfg, cli_overrides=cli_overrides, ) @@ -1483,7 +1647,8 @@ class TestSentinelDefaultPrecedence: """A single ``qwen3_tts`` pipeline picks per-chunk vs end-to-end processors based on ``deploy.async_chunk``, without needing a separate variant pipeline registration.""" - pipeline = _PIPELINE_REGISTRY["qwen3_tts"] + pipeline = StageConfigFactory.resolve_pipeline_config("qwen3_tts") + assert isinstance(pipeline, PipelineConfig) # async_chunk=True → stage 0's per-chunk processor wires up, stage 1 # has no sync input processor. @@ -1552,9 +1717,8 @@ class TestSentinelDefaultPrecedence: ``thinker2talker_token_only`` placeholder (sync_process_input_func). Merge under either async_chunk mode must not re-introduce a stage-0 full-payload hook.""" - from vllm_omni.config.stage_config import DeployConfig, merge_pipeline_deploy - - pipeline = _PIPELINE_REGISTRY["ming_flash_omni"] + pipeline = StageConfigFactory.resolve_pipeline_config("ming_flash_omni") + assert isinstance(pipeline, PipelineConfig) stage0, stage1 = pipeline.stages assert stage0.custom_process_next_stage_input_func is None, ( @@ -1587,15 +1751,30 @@ class TestSamplingConstraintsPrecedence: """Test that pipeline sampling_constraints override deploy defaults.""" def test_constraints_win(self): - pipeline = _PIPELINE_REGISTRY["qwen3_omni_moe"] deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml" if not deploy_path.exists(): pytest.skip("Deploy config not found") deploy = load_deploy_config(deploy_path) + pipeline = StageConfigFactory.resolve_pipeline_config( + "qwen3_omni_moe", + Q3_OMNI_ALL_STAGES_HF_CONFIG, + ) + assert isinstance(pipeline, PipelineConfig) stages = merge_pipeline_deploy(pipeline, deploy) # Pipeline says detokenize=True for thinker, deploy can't override assert stages[0].yaml_extras["default_sampling_params"]["detokenize"] is True # Pipeline says stop_token_ids=[2150] for talker assert stages[1].yaml_extras["default_sampling_params"]["stop_token_ids"] == [2150] + + +class TestPipelineConfigResolvers: + @pytest.mark.parametrize("resolver", [obj for obj in OMNI_PIPELINES.values() if callable(obj)]) + def test_all_resolvers_reject_bad_types(self, resolver): + """Ensure that all resolvers registered reject incorrect config types.""" + + class NotTheRightHfConfig(PretrainedConfig): + pass + + assert resolver(NotTheRightHfConfig()) is None diff --git a/tests/test_diffusion_config_propagation.py b/tests/test_diffusion_config_propagation.py index 0991c6601..10fe00168 100644 --- a/tests/test_diffusion_config_propagation.py +++ b/tests/test_diffusion_config_propagation.py @@ -10,7 +10,7 @@ from collections.abc import Mapping import pytest import torch -from vllm_omni.config.stage_config import StageConfigFactory +from vllm_omni.config.config_factory import StageConfigFactory from vllm_omni.diffusion.data import ( DiffusionParallelConfig, OmniDiffusionConfig, diff --git a/tests/unit/higgs_audio_v3/test_higgs_audio_v3.py b/tests/unit/higgs_audio_v3/test_higgs_audio_v3.py index 33be067db..b701d9f8b 100644 --- a/tests/unit/higgs_audio_v3/test_higgs_audio_v3.py +++ b/tests/unit/higgs_audio_v3/test_higgs_audio_v3.py @@ -19,7 +19,7 @@ import torch class TestHiggsAudioV3Config: def test_default_config_loads(self): - from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import ( + from vllm_omni.transformers_utils.configs.higgs_audio_v3 import ( HiggsAudioV3Config, ) @@ -34,7 +34,7 @@ class TestHiggsAudioV3Config: assert config.tie_modality_embeddings is True def test_custom_audio_encoder_config(self): - from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import ( + from vllm_omni.transformers_utils.configs.higgs_audio_v3 import ( HiggsAudioV3Config, ) @@ -51,7 +51,7 @@ class TestHiggsAudioV3Config: assert config.tie_modality_embeddings is False def test_invalid_num_codebooks_rejected(self): - from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import ( + from vllm_omni.transformers_utils.configs.higgs_audio_v3 import ( HiggsAudioV3Config, ) @@ -59,7 +59,7 @@ class TestHiggsAudioV3Config: HiggsAudioV3Config(audio_encoder_config={"num_codebooks": 0, "vocab_size": 1026}) def test_negative_num_codebooks_rejected(self): - from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import ( + from vllm_omni.transformers_utils.configs.higgs_audio_v3 import ( HiggsAudioV3Config, ) @@ -67,7 +67,7 @@ class TestHiggsAudioV3Config: HiggsAudioV3Config(audio_encoder_config={"num_codebooks": -1, "vocab_size": 1026}) def test_special_token_ids_initially_none(self): - from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import ( + from vllm_omni.transformers_utils.configs.higgs_audio_v3 import ( HiggsAudioV3Config, ) @@ -77,7 +77,7 @@ class TestHiggsAudioV3Config: assert config.audio_continuation_id is None def test_hidden_size_from_text_config(self): - from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import ( + from vllm_omni.transformers_utils.configs.higgs_audio_v3 import ( HiggsAudioV3Config, ) @@ -797,12 +797,12 @@ class TestAudioFeedback: class TestCodecStrictness: def test_bundled_missing_quantizer_key_raises(self): """Bundled codec load must fail when a quantizer codebook key is missing.""" - from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import ( - HiggsAudioV3Config, - ) from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_code2wav import ( HiggsAudioV3Code2Wav, ) + from vllm_omni.transformers_utils.configs.higgs_audio_v3 import ( + HiggsAudioV3Config, + ) config = HiggsAudioV3Config() c2w = HiggsAudioV3Code2Wav(config=config) @@ -821,12 +821,12 @@ class TestCodecStrictness: def test_bundled_missing_fc2_raises(self): """Bundled codec load must fail when fc2 keys are missing.""" - from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import ( - HiggsAudioV3Config, - ) from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_code2wav import ( HiggsAudioV3Code2Wav, ) + from vllm_omni.transformers_utils.configs.higgs_audio_v3 import ( + HiggsAudioV3Config, + ) config = HiggsAudioV3Config() c2w = HiggsAudioV3Code2Wav(config=config) @@ -844,12 +844,12 @@ class TestCodecStrictness: def test_bundled_unknown_decoder_key_rejected(self, monkeypatch): """Unknown decoder-side key should be rejected as unconsumed.""" from vllm_omni.model_executor.models.higgs_audio_v3 import higgs_audio_v3_code2wav as c2w_mod - from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import ( - HiggsAudioV3Config, - ) from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_code2wav import ( HiggsAudioV3Code2Wav, ) + from vllm_omni.transformers_utils.configs.higgs_audio_v3 import ( + HiggsAudioV3Config, + ) class FakeDAC(torch.nn.Module): def __init__(self): @@ -871,12 +871,12 @@ class TestCodecStrictness: def test_bundled_encoder_side_keys_accepted(self, monkeypatch): """Known encoder-side keys should not trigger unconsumed-key rejection.""" from vllm_omni.model_executor.models.higgs_audio_v3 import higgs_audio_v3_code2wav as c2w_mod - from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import ( - HiggsAudioV3Config, - ) from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_code2wav import ( HiggsAudioV3Code2Wav, ) + from vllm_omni.transformers_utils.configs.higgs_audio_v3 import ( + HiggsAudioV3Config, + ) class FakeDAC(torch.nn.Module): def __init__(self): @@ -1199,9 +1199,9 @@ class TestRegistry: assert "HiggsAudioV3Code2WavForConditionalGeneration" in _OMNI_MODELS def test_pipeline_registered(self): - from vllm_omni.config.pipeline_registry import _OMNI_PIPELINES + from vllm_omni.config.pipeline_registry import OMNI_PIPELINES - assert "higgs_multimodal_qwen3" in _OMNI_PIPELINES + assert "higgs_multimodal_qwen3" in OMNI_PIPELINES def test_deploy_yaml_exists(self): import os diff --git a/tests/utils/test_tracking_parser.py b/tests/utils/test_tracking_parser.py index 537454661..437fbd50a 100644 --- a/tests/utils/test_tracking_parser.py +++ b/tests/utils/test_tracking_parser.py @@ -8,11 +8,11 @@ import pytest import yaml from vllm.utils.argparse_utils import FlexibleArgumentParser +from vllm_omni.config.config_factory import StageConfigFactory +from vllm_omni.config.pipeline_registry import OMNI_PIPELINES from vllm_omni.config.stage_config import ( - _PIPELINE_REGISTRY, DeployConfig, PipelineConfig, - StageConfigFactory, StageDeployConfig, StagePipelineConfig, ) @@ -44,9 +44,9 @@ _TEST_DEPLOY = DeployConfig( @pytest.fixture() def mock_stages(monkeypatch): """Register a fake pipeline and mock deploy YAML loading.""" - monkeypatch.setitem(_PIPELINE_REGISTRY._loaded, _TEST_MODEL, _TEST_PIPELINE) + monkeypatch.setitem(OMNI_PIPELINES, _TEST_MODEL, _TEST_PIPELINE) monkeypatch.setattr( - "vllm_omni.config.stage_config.load_deploy_config", + "vllm_omni.config.config_factory.load_deploy_config", lambda _path: _TEST_DEPLOY, ) return __file__ @@ -462,6 +462,7 @@ def test_explicit_cli_arg_reaches_runtime_overrides(mock_stages): explicit_kwargs = ns.get_explicit_kwargs_dict() stages = StageConfigFactory._create_from_registry( _TEST_MODEL, + _TEST_PIPELINE, explicit_kwargs, deploy_config_path=mock_stages, ) @@ -478,6 +479,7 @@ def test_omitted_default_not_in_runtime_overrides(mock_stages): explicit_kwargs = ns.get_explicit_kwargs_dict() stages = StageConfigFactory._create_from_registry( _TEST_MODEL, + _TEST_PIPELINE, explicit_kwargs, deploy_config_path=mock_stages, ) @@ -499,6 +501,7 @@ def test_config_file_args_reach_runtime_overrides(mock_stages): explicit_kwargs = ns.get_explicit_kwargs_dict() stages = StageConfigFactory._create_from_registry( _TEST_MODEL, + _TEST_PIPELINE, explicit_kwargs, deploy_config_path=mock_stages, ) @@ -516,6 +519,7 @@ def test_per_stage_override_routes_correctly(mock_stages): explicit_kwargs = ns.get_explicit_kwargs_dict() stages = StageConfigFactory._create_from_registry( _TEST_MODEL, + _TEST_PIPELINE, explicit_kwargs, deploy_config_path=mock_stages, ) @@ -537,6 +541,7 @@ def test_explicit_args_omitted_from_yaml(mock_stages): explicit_kwargs = ns.get_explicit_kwargs_dict() stages = StageConfigFactory._create_from_registry( _TEST_MODEL, + _TEST_PIPELINE, explicit_kwargs, deploy_config_path=mock_stages, ) diff --git a/vllm_omni/config/__init__.py b/vllm_omni/config/__init__.py index e5762456f..a43cc41d9 100644 --- a/vllm_omni/config/__init__.py +++ b/vllm_omni/config/__init__.py @@ -8,14 +8,12 @@ from vllm_omni.config.stage_config import ( DeployConfig, PipelineConfig, StageConfig, - StageConfigFactory, StageDeployConfig, StageExecutionType, StagePipelineConfig, StageType, load_deploy_config, merge_pipeline_deploy, - register_pipeline, ) from vllm_omni.config.yaml_util import ( create_config, @@ -28,7 +26,6 @@ __all__ = [ "OmniModelConfig", "LoRAConfig", "StageConfig", - "StageConfigFactory", "StageType", "StageExecutionType", "StagePipelineConfig", @@ -37,7 +34,6 @@ __all__ = [ "DeployConfig", "load_deploy_config", "merge_pipeline_deploy", - "register_pipeline", "create_config", "load_yaml_config", "merge_configs", diff --git a/vllm_omni/config/config_factory.py b/vllm_omni/config/config_factory.py new file mode 100644 index 000000000..3b3f0d506 --- /dev/null +++ b/vllm_omni/config/config_factory.py @@ -0,0 +1,338 @@ +# 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 + + # --- New path: check pipeline registry by model_type first --- + 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" + + 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, + ) + # 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 diff --git a/vllm_omni/config/pipeline_registry.py b/vllm_omni/config/pipeline_registry.py index 63548ef0f..dc5843a26 100644 --- a/vllm_omni/config/pipeline_registry.py +++ b/vllm_omni/config/pipeline_registry.py @@ -1,176 +1,142 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Central declarative registry of all vllm-omni pipelines. +"""Pipeline registry and factory for vllm-omni. -Mirrors the pattern in ``vllm/model_executor/models/registry.py``: each entry -is ``model_type -> (module_path, variable_name)``, and the module is imported -lazily on first lookup (see ``_LazyPipelineRegistry`` in -``vllm_omni/config/stage_config.py``). Keeping every pipeline declared in one -file makes it easy to spot a missing registration, which was the original -motivation in https://github.com/vllm-project/vllm-omni/issues/2887 (item 4). +``OMNI_PIPELINES`` maps each ``model_type`` to either a ``PipelineConfig`` +instance or a resolver callable that accepts an optional HF config and returns +a ``PipelineConfig``. -Per-model ``pipeline.py`` modules still define the ``PipelineConfig`` instance; -they just no longer need to self-register via ``register_pipeline(...)``. - -Adding a new pipeline: +To add a new pipeline: 1. Define the ``PipelineConfig`` instance as a module-level variable in ``vllm_omni/.../pipeline.py``. - 2. Add one line to ``_OMNI_PIPELINES`` below. + 2. If the model needs to support several configurations, e.g., because some + stages are optional, implement a resolver that consumes the HF config + and returns a ``PipelineConfig``. + 3. Update the registry to map the key to the new config object (in the case + of new keys) or to the resolver func. -Plain single-stage diffusion models continue to use the -``_create_default_diffusion_stage_cfg`` fallback in ``async_omni_engine.py``. -The empty ``_DIFFUSION_PIPELINES`` placeholder previously here (#2915) was -removed once #2987 (which would have populated it) was deferred. +Out of tree pipeline configs or resolvers can also be registered with register_pipeline. -``register_pipeline(config)`` in ``stage_config`` is still supported for -out-of-tree plugins and tests that create pipelines at runtime; those override -the entries declared here. +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. """ from __future__ import annotations +from collections.abc import Callable +from typing import TypeAlias + +from transformers import PreTrainedConfig +from vllm.logger import init_logger + +from vllm_omni.config.stage_config import ( + PipelineConfig, +) +from vllm_omni.model_executor.models.aura_omni.pipeline import AURA_OMNI_PIPELINE +from vllm_omni.model_executor.models.bagel.pipeline import ( + BAGEL_PIPELINE, + BAGEL_SINGLE_STAGE_PIPELINE, + BAGEL_THINK_PIPELINE, +) +from vllm_omni.model_executor.models.cosyvoice3.pipeline import COSYVOICE3_PIPELINE +from vllm_omni.model_executor.models.covo_audio.pipeline import COVO_AUDIO_PIPELINE +from vllm_omni.model_executor.models.dreamzero.pipeline import DREAMZERO_PIPELINE +from vllm_omni.model_executor.models.dynin_omni.pipeline import DYNIN_OMNI_PIPELINE +from vllm_omni.model_executor.models.fish_speech.pipeline import FISH_SPEECH_PIPELINE +from vllm_omni.model_executor.models.glm_image.pipeline import GLM_IMAGE_PIPELINE +from vllm_omni.model_executor.models.glm_tts.pipeline import GLM_TTS_PIPELINE +from vllm_omni.model_executor.models.gr00t.pipeline import GR00T_N1D7_PIPELINE +from vllm_omni.model_executor.models.higgs_audio_v2.pipeline import HIGGS_AUDIO_V2_PIPELINE +from vllm_omni.model_executor.models.higgs_audio_v3.pipeline import HIGGS_AUDIO_V3_PIPELINE +from vllm_omni.model_executor.models.hunyuan_image3.pipeline import ( + HUNYUAN_IMAGE3_AR_PIPELINE, + HUNYUAN_IMAGE3_DIT_PIPELINE, + HUNYUAN_IMAGE3_PIPELINE, +) +from vllm_omni.model_executor.models.indextts2.pipeline import INDEXTTS2_PIPELINE +from vllm_omni.model_executor.models.lance.pipeline import LANCE_PIPELINE +from vllm_omni.model_executor.models.mimo_audio.pipeline import MIMO_AUDIO_PIPELINE +from vllm_omni.model_executor.models.ming_flash_omni.pipeline import ( + MING_FLASH_OMNI_IMAGE_PIPELINE, + MING_FLASH_OMNI_PIPELINE, + MING_FLASH_OMNI_THINKER_ONLY_PIPELINE, + MING_FLASH_OMNI_TTS_PIPELINE, +) +from vllm_omni.model_executor.models.ming_tts.pipeline import MING_TTS_PIPELINE +from vllm_omni.model_executor.models.minicpmo_4_5.pipeline import MINICPMO_4_5_PIPELINE +from vllm_omni.model_executor.models.moss_tts.pipeline import MOSS_TTS_PIPELINE, MOSS_TTS_REALTIME_PIPELINE +from vllm_omni.model_executor.models.moss_tts_nano.pipeline import MOSS_TTS_NANO_PIPELINE +from vllm_omni.model_executor.models.qwen2_5_omni.pipeline import ( + QWEN2_5_OMNI_PIPELINE, + QWEN2_5_OMNI_THINKER_ONLY_PIPELINE, +) +from vllm_omni.model_executor.models.qwen3_omni.pipeline import resolve_qwen3_omni_pipeline +from vllm_omni.model_executor.models.qwen3_tts.pipeline import QWEN3_TTS_PIPELINE +from vllm_omni.model_executor.models.voxcpm2.pipeline import VOXCPM2_PIPELINE +from vllm_omni.model_executor.models.voxtral_tts.pipeline import VOXTRAL_TTS_PIPELINE + +logger = init_logger(__name__) + +PipelineResolverFunc: TypeAlias = Callable[[PreTrainedConfig | None], PipelineConfig | None] + # --- Multi-stage omni pipelines (LLM-centric; audio / video I/O) --- -_OMNI_PIPELINES: dict[str, tuple[str, str]] = { - # model_type -> (module_path, variable_name) - "qwen2_5_omni": ( - "vllm_omni.model_executor.models.qwen2_5_omni.pipeline", - "QWEN2_5_OMNI_PIPELINE", - ), - "qwen2_5_omni_thinker_only": ( - "vllm_omni.model_executor.models.qwen2_5_omni.pipeline", - "QWEN2_5_OMNI_THINKER_ONLY_PIPELINE", - ), - "qwen3_omni_moe": ( - "vllm_omni.model_executor.models.qwen3_omni.pipeline", - "QWEN3_OMNI_PIPELINE", - ), - "aura_omni": ( - "vllm_omni.model_executor.models.aura_omni.pipeline", - "AURA_OMNI_PIPELINE", - ), - "qwen3_tts": ( - "vllm_omni.model_executor.models.qwen3_tts.pipeline", - "QWEN3_TTS_PIPELINE", - ), - "covo_audio": ( - "vllm_omni.model_executor.models.covo_audio.pipeline", - "COVO_AUDIO_PIPELINE", - ), - "bagel": ( - "vllm_omni.model_executor.models.bagel.pipeline", - "BAGEL_PIPELINE", - ), - "bagel_think": ( - "vllm_omni.model_executor.models.bagel.pipeline", - "BAGEL_THINK_PIPELINE", - ), - "bagel_single_stage": ( - "vllm_omni.model_executor.models.bagel.pipeline", - "BAGEL_SINGLE_STAGE_PIPELINE", - ), - # Lance (ByteDance) — BAGEL-lineage unified AR+diffusion, single-stage DiT. - "lance": ( - "vllm_omni.model_executor.models.lance.pipeline", - "LANCE_PIPELINE", - ), - "dreamzero": ( - "vllm_omni.model_executor.models.dreamzero.pipeline", - "DREAMZERO_PIPELINE", - ), - "glm_image": ( - "vllm_omni.model_executor.models.glm_image.pipeline", - "GLM_IMAGE_PIPELINE", - ), - "Gr00tN1d7": ( - "vllm_omni.model_executor.models.gr00t.pipeline", - "GR00T_N1D7_PIPELINE", - ), - "hunyuan_image_3_moe": ( - "vllm_omni.model_executor.models.hunyuan_image3.pipeline", - "HUNYUAN_IMAGE3_PIPELINE", - ), - "hunyuan_image3_ar": ( - "vllm_omni.model_executor.models.hunyuan_image3.pipeline", - "HUNYUAN_IMAGE3_AR_PIPELINE", - ), - "hunyuan_image3_dit": ( - "vllm_omni.model_executor.models.hunyuan_image3.pipeline", - "HUNYUAN_IMAGE3_DIT_PIPELINE", - ), - "voxcpm2": ( - "vllm_omni.model_executor.models.voxcpm2.pipeline", - "VOXCPM2_PIPELINE", - ), - "cosyvoice3": ( - "vllm_omni.model_executor.models.cosyvoice3.pipeline", - "COSYVOICE3_PIPELINE", - ), - "mimo_audio": ( - "vllm_omni.model_executor.models.mimo_audio.pipeline", - "MIMO_AUDIO_PIPELINE", - ), - "ming_tts": ( - "vllm_omni.model_executor.models.ming_tts.pipeline", - "MING_TTS_PIPELINE", - ), - "voxtral_tts": ( - "vllm_omni.model_executor.models.voxtral_tts.pipeline", - "VOXTRAL_TTS_PIPELINE", - ), - "glm_tts": ( - "vllm_omni.model_executor.models.glm_tts.pipeline", - "GLM_TTS_PIPELINE", - ), - "fish_qwen3_omni": ( - "vllm_omni.model_executor.models.fish_speech.pipeline", - "FISH_SPEECH_PIPELINE", - ), - "ming_flash_omni": ( - "vllm_omni.model_executor.models.ming_flash_omni.pipeline", - "MING_FLASH_OMNI_PIPELINE", - ), - "ming_flash_omni_tts": ( - "vllm_omni.model_executor.models.ming_flash_omni.pipeline", - "MING_FLASH_OMNI_TTS_PIPELINE", - ), - "ming_flash_omni_thinker_only": ( - "vllm_omni.model_executor.models.ming_flash_omni.pipeline", - "MING_FLASH_OMNI_THINKER_ONLY_PIPELINE", - ), - "ming_flash_omni_image": ( - "vllm_omni.model_executor.models.ming_flash_omni.pipeline", - "MING_FLASH_OMNI_IMAGE_PIPELINE", - ), - "moss_tts_nano": ( - "vllm_omni.model_executor.models.moss_tts_nano.pipeline", - "MOSS_TTS_NANO_PIPELINE", - ), - "moss_tts_delay": ( - "vllm_omni.model_executor.models.moss_tts.pipeline", - "MOSS_TTS_PIPELINE", - ), - "moss_tts_realtime": ( - "vllm_omni.model_executor.models.moss_tts.pipeline", - "MOSS_TTS_REALTIME_PIPELINE", - ), - "minicpmo_4_5": ( - "vllm_omni.model_executor.models.minicpmo_4_5.pipeline", - "MINICPMO_4_5_PIPELINE", - ), - "higgs_audio_v2": ( - "vllm_omni.model_executor.models.higgs_audio_v2.pipeline", - "HIGGS_AUDIO_V2_PIPELINE", - ), - "higgs_multimodal_qwen3": ( - "vllm_omni.model_executor.models.higgs_audio_v3.pipeline", - "HIGGS_AUDIO_V3_PIPELINE", - ), - "dynin_omni": ( - "vllm_omni.model_executor.models.dynin_omni.pipeline", - "DYNIN_OMNI_PIPELINE", - ), - "indextts2": ( - "vllm_omni.model_executor.models.indextts2.pipeline", - "INDEXTTS2_PIPELINE", - ), +OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = { + "aura_omni": AURA_OMNI_PIPELINE, + "qwen2_5_omni": QWEN2_5_OMNI_PIPELINE, + "qwen2_5_omni_thinker_only": QWEN2_5_OMNI_THINKER_ONLY_PIPELINE, + "qwen3_omni_moe": resolve_qwen3_omni_pipeline, + "qwen3_tts": QWEN3_TTS_PIPELINE, + "covo_audio": COVO_AUDIO_PIPELINE, + "bagel": BAGEL_PIPELINE, + "bagel_think": BAGEL_THINK_PIPELINE, + "bagel_single_stage": BAGEL_SINGLE_STAGE_PIPELINE, + "lance": LANCE_PIPELINE, + "dreamzero": DREAMZERO_PIPELINE, + "Gr00tN1d7": GR00T_N1D7_PIPELINE, + "glm_image": GLM_IMAGE_PIPELINE, + "hunyuan_image_3_moe": HUNYUAN_IMAGE3_PIPELINE, + "hunyuan_image3_ar": HUNYUAN_IMAGE3_AR_PIPELINE, + "hunyuan_image3_dit": HUNYUAN_IMAGE3_DIT_PIPELINE, + "voxcpm2": VOXCPM2_PIPELINE, + "cosyvoice3": COSYVOICE3_PIPELINE, + "mimo_audio": MIMO_AUDIO_PIPELINE, + "ming_tts": MING_TTS_PIPELINE, + "voxtral_tts": VOXTRAL_TTS_PIPELINE, + "glm_tts": GLM_TTS_PIPELINE, + "fish_qwen3_omni": FISH_SPEECH_PIPELINE, + "ming_flash_omni": MING_FLASH_OMNI_PIPELINE, + "ming_flash_omni_tts": MING_FLASH_OMNI_TTS_PIPELINE, + "ming_flash_omni_thinker_only": MING_FLASH_OMNI_THINKER_ONLY_PIPELINE, + "ming_flash_omni_image": MING_FLASH_OMNI_IMAGE_PIPELINE, + "moss_tts_nano": MOSS_TTS_NANO_PIPELINE, + "moss_tts_delay": MOSS_TTS_PIPELINE, + "moss_tts_realtime": MOSS_TTS_REALTIME_PIPELINE, + "minicpmo_4_5": MINICPMO_4_5_PIPELINE, + "higgs_audio_v2": HIGGS_AUDIO_V2_PIPELINE, + "higgs_multimodal_qwen3": HIGGS_AUDIO_V3_PIPELINE, + "dynin_omni": DYNIN_OMNI_PIPELINE, + "indextts2": INDEXTTS2_PIPELINE, } + + +def register_pipeline(pipeline: PipelineConfig | PipelineResolverFunc, model_type: str | None = None): + """Register an out of tree pipeline or PipelineResolverFunc to a model_type key. + If a PipelineConfig is provided, model_type is optional, and pipeline.model_type + will be used by default. If a callable is provided, model_type must be provided, + since resolvers can return multiple different PipelineConfigs depending on the + consumed config. + """ + errors: list[str] = [] + if isinstance(pipeline, PipelineConfig): + errors = pipeline.validate() + model_type = model_type if model_type is not None else pipeline.model_type + else: + if model_type is None: + raise ValueError("Model type must be explicitly provided when registering a pipeline resolver") + + if model_type in OMNI_PIPELINES: + errors.append(f"Model type {model_type} is already registered; the old mapping will be clobbered") + if errors: + logger.warning("Registration for pipeline of type %s produced the following issues: %s", model_type, errors) + OMNI_PIPELINES[model_type] = pipeline diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 1bff5ef94..edabee9a9 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -5,6 +5,7 @@ from __future__ import annotations import dataclasses +import functools import re import warnings from collections.abc import Callable @@ -13,6 +14,7 @@ from enum import Enum from pathlib import Path from typing import Any, NamedTuple +from transformers import PretrainedConfig from vllm.logger import init_logger from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler @@ -22,10 +24,26 @@ from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationSchedul logger = init_logger(__name__) +_DEPLOY_DIR = Path(__file__).resolve().parent.parent / "deploy" _STAGE_OVERRIDE_PATTERN = re.compile(r"^stage_(\d+)_(.+)$") +def pipeline_cfg_resolver(config_type: type[PretrainedConfig]): + """Wraps a resolver such that we return None if a hf_config of the wrong type is provided.""" + + def resolver_builder(func): + @functools.wraps(func) + def wrapper(hf_config: PretrainedConfig | None): + if hf_config is None or not isinstance(hf_config, config_type): + return None + return func(hf_config) + + return wrapper + + return resolver_builder + + def build_stage_runtime_overrides( stage_id: int, cli_overrides: dict[str, Any], @@ -273,153 +291,6 @@ class PipelineConfig: return errors -class _LazyPipelineRegistry: - """Dict-like registry that lazy-loads pipelines from the central declaration. - - In-tree pipelines are declared once in - ``vllm_omni/config/pipeline_registry.py`` as - ``model_type -> (module_path, variable_name)`` entries; the module is - imported only when the pipeline is first looked up. This mirrors the - pattern in ``vllm/model_executor/models/registry.py`` and addresses - https://github.com/vllm-project/vllm-omni/issues/2887 (item 4): having - every registration in one file makes a missing entry easy to spot. - - Out-of-tree / dynamic registrations via ``register_pipeline()`` are stored - directly in ``_loaded`` and take precedence over the lazy-map entry with - the same ``model_type``. - - The class exposes the subset of ``dict`` operations the rest of this - module relies on (``__contains__``, ``__getitem__``, ``__setitem__``, - ``get``, ``keys``, ``values``, ``items``, ``__iter__``), so existing call - sites don't need to change. - """ - - def __init__(self) -> None: - self._loaded: dict[str, PipelineConfig] = {} - # Populated lazily to avoid a circular import at module init time. - self._lazy_map: dict[str, tuple[str, str]] | None = None - - def _get_lazy_map(self) -> dict[str, tuple[str, str]]: - if self._lazy_map is None: - from vllm_omni.config.pipeline_registry import _OMNI_PIPELINES - - self._lazy_map = _OMNI_PIPELINES - return self._lazy_map - - def _load_lazy(self, model_type: str) -> PipelineConfig | None: - entry = self._get_lazy_map().get(model_type) - if entry is None: - return None - module_path, var_name = entry - import importlib - - try: - module = importlib.import_module(module_path) - pipeline = getattr(module, var_name, None) - if pipeline is None: - logger.error( - "Pipeline variable %r not found in module %r (registered for %r)", - var_name, - module_path, - model_type, - ) - return None - errors = pipeline.validate() - if errors: - logger.warning("Pipeline %s has issues: %s", pipeline.model_type, errors) - self._loaded[model_type] = pipeline - return pipeline - except Exception: - logger.exception("Failed to import pipeline module %r for %r", module_path, model_type) - return None - - def __contains__(self, model_type: str) -> bool: - if model_type in self._loaded: - return True - return model_type in self._get_lazy_map() - - def __getitem__(self, model_type: str) -> PipelineConfig: - if model_type in self._loaded: - return self._loaded[model_type] - pipeline = self._load_lazy(model_type) - if pipeline is None: - raise KeyError(model_type) - return pipeline - - def get(self, model_type: str, default: PipelineConfig | None = None) -> PipelineConfig | None: - if model_type in self._loaded: - return self._loaded[model_type] - pipeline = self._load_lazy(model_type) - return pipeline if pipeline is not None else default - - def __setitem__(self, model_type: str, pipeline: PipelineConfig) -> None: - self._loaded[model_type] = pipeline - - def __delitem__(self, model_type: str) -> None: - """Remove a dynamically-registered pipeline. - - Only the dynamic-cache side of the registry can be mutated; the - central declarative registry is immutable at runtime. Calling ``del`` - on a model_type that only exists in the central registry raises - ``KeyError``. - """ - if model_type in self._loaded: - del self._loaded[model_type] - return - if model_type in self._get_lazy_map(): - raise KeyError( - f"{model_type!r} is declared in the central pipeline_registry and " - "cannot be removed at runtime. Edit " - "vllm_omni/config/pipeline_registry.py to delete a built-in entry." - ) - raise KeyError(model_type) - - def keys(self) -> set[str]: - return set(self._get_lazy_map().keys()) | set(self._loaded.keys()) - - def _safe_get(self, key: str) -> PipelineConfig | None: - try: - return self[key] - except Exception: - logger.warning("Skipping pipeline %r because it failed to load.", key) - return None - - def values(self): - # Iterating forces a lazy import for each pipeline; failures are logged and skipped. - for key in self.keys(): - if (p := self._safe_get(key)) is not None: - yield p - - def items(self): - for key in self.keys(): - if (p := self._safe_get(key)) is not None: - yield key, p - - def __iter__(self): - return iter(self.keys()) - - -_PIPELINE_REGISTRY = _LazyPipelineRegistry() - - -def register_pipeline(pipeline: PipelineConfig) -> None: - """Register a pipeline config dynamically. - - In-tree pipelines are declared in ``pipeline_registry._OMNI_PIPELINES`` - and loaded lazily; calling ``register_pipeline`` is only needed for - out-of-tree plugins or tests that build a ``PipelineConfig`` at runtime. - A dynamic registration overrides the central-registry entry with the same - ``model_type``. - """ - errors = pipeline.validate() - if errors: - logger.warning("Pipeline %s has issues: %s", pipeline.model_type, errors) - _PIPELINE_REGISTRY[pipeline.model_type] = pipeline - - -_DEPLOY_DIR = Path(__file__).resolve().parent.parent / "deploy" - - @dataclass class StageDeployConfig: """Per-stage deployment knobs. @@ -905,6 +776,10 @@ def merge_pipeline_deploy( deploy = _apply_platform_overrides(deploy) deploy_by_id = {s.stage_id: s for s in deploy.stages} + # async_chunk is irrelevant for single-stage pipelines, so we always disable it + if len(pipeline.stages) <= 1: + deploy.async_chunk = False + # A pipeline supports async_chunk if any stage has either an explicit # async-chunk-only processor slot OR a custom next-stage processor (some # pipelines like qwen3_omni wire async-chunk processing directly through @@ -1052,305 +927,3 @@ class StageConfig: config_dict.update(self.yaml_extras) return create_config(config_dict) - - -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 - loaded lazily via ``_PIPELINE_REGISTRY``; no hardcoded model-type → - directory mapping is maintained here. 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, - ) -> list[StageConfig] | None: - """Load pipeline + deploy config, merge with CLI overrides. - - Checks _PIPELINE_REGISTRY first (new path), falls back to legacy YAML. - """ - 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 - - # --- New path: check pipeline registry by model_type first --- - model_type, hf_config = cls._auto_detect_model_type(model, trust_remote_code=trust_remote_code) - if model_type == "vla": - from vllm_omni.diffusion.utils.hf_utils import _looks_like_dreamzero - - if _looks_like_dreamzero(model): - model_type = "dreamzero" - if model_type and model_type in _PIPELINE_REGISTRY: - return cls._create_from_registry(model_type, 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. - # When a pipeline declares ``hf_config_predicate`` (e.g. to - # disambiguate generation-revision collisions like MiniCPM-o 2.6 vs - # 4.5 which share ``architectures=["MiniCPMO"]``), the predicate - # must also accept the loaded ``hf_config``. - if hf_config is not None: - hf_archs = set(getattr(hf_config, "architectures", []) or []) - if hf_archs: - for registered in _PIPELINE_REGISTRY.values(): - if not hf_archs.intersection(registered.hf_architectures): - continue - predicate = registered.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.", - registered.model_type, - sorted(hf_archs.intersection(registered.hf_architectures)), - ) - continue - except Exception: - logger.exception( - "Pipeline %r hf_config_predicate raised; skipping.", - registered.model_type, - ) - continue - return cls._create_from_registry(registered.model_type, cli_overrides, deploy_config_path) - - # An explicit deploy config can select a pipeline when the model's HF - # metadata is too generic for registry auto-detection. - if deploy_config_path is not None: - deploy_path = Path(deploy_config_path) - if deploy_path.exists(): - deploy_config = load_deploy_config(deploy_path) - if deploy_config.pipeline in _PIPELINE_REGISTRY: - return cls._create_from_registry( - deploy_config.pipeline, - 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, - cli_overrides: dict[str, Any], - deploy_config_path: str | None = None, - ) -> 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) - # Bare-name lookup: if the value isn't a real path, try the - # built-in deploy directory. This lets ``deploy_config="moss_tts"`` - # resolve to ``vllm_omni/deploy/moss_tts.yaml`` when several model - # variants share an HF ``model_type`` and need different YAMLs. - if not deploy_path.exists() and deploy_path.parent == Path("."): - bare_name = deploy_path.name - if not bare_name.endswith(".yaml"): - bare_name = f"{bare_name}.yaml" - candidate = _DEPLOY_DIR / bare_name - if candidate.exists(): - deploy_path = candidate - - 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) - - cli_async_chunk = cli_overrides.get("async_chunk") - if cli_async_chunk is not None: - deploy_cfg.async_chunk = bool(cli_async_chunk) - - pipeline_key = deploy_cfg.pipeline or model_type - if pipeline_key not in _PIPELINE_REGISTRY: - raise KeyError( - f"Pipeline {pipeline_key!r} not in registry " - f"(resolved from {deploy_path.name!r}). Available: " - f"{sorted(_PIPELINE_REGISTRY.keys())}" - ) - pipeline_cfg = _PIPELINE_REGISTRY[pipeline_key] - - 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. - """ - try: - from vllm.transformers_utils.config import get_config - - 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: - from vllm.transformers_utils.config import get_hf_file_to_dict - - 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: - from vllm.transformers_utils.config import get_hf_file_to_dict - - 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 pipeline_cfg in _PIPELINE_REGISTRY.values(): - if 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 _PIPELINE_REGISTRY.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) diff --git a/vllm_omni/entrypoints/utils.py b/vllm_omni/entrypoints/utils.py index 2ce51a749..961638e14 100644 --- a/vllm_omni/entrypoints/utils.py +++ b/vllm_omni/entrypoints/utils.py @@ -11,7 +11,7 @@ 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.stage_config import StageConfigFactory +from vllm_omni.config.config_factory import StageConfigFactory from vllm_omni.config.yaml_util import create_config, load_yaml_config, merge_configs from vllm_omni.diffusion.utils.hf_utils import _looks_like_dreamzero from vllm_omni.entrypoints.stage_utils import _to_dict diff --git a/vllm_omni/model_executor/models/aura_omni/__init__.py b/vllm_omni/model_executor/models/aura_omni/__init__.py new file mode 100644 index 000000000..fb6738a29 --- /dev/null +++ b/vllm_omni/model_executor/models/aura_omni/__init__.py @@ -0,0 +1,8 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# NOTE: Do not import model classes in this file. Importing any +# submodule in this package triggers __init__.py execution, and +# both the model registry and pipeline registry import submodules +# directly — heavy imports here would be loaded as a side effect +# even though nothing depends on these re-exports. diff --git a/vllm_omni/model_executor/models/bagel/__init__.py b/vllm_omni/model_executor/models/bagel/__init__.py index 11e4c6667..fb6738a29 100644 --- a/vllm_omni/model_executor/models/bagel/__init__.py +++ b/vllm_omni/model_executor/models/bagel/__init__.py @@ -1,3 +1,8 @@ -from .bagel import OmniBagelForConditionalGeneration +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -__all__ = ["OmniBagelForConditionalGeneration"] +# NOTE: Do not import model classes in this file. Importing any +# submodule in this package triggers __init__.py execution, and +# both the model registry and pipeline registry import submodules +# directly — heavy imports here would be loaded as a side effect +# even though nothing depends on these re-exports. diff --git a/vllm_omni/model_executor/models/dreamzero/__init__.py b/vllm_omni/model_executor/models/dreamzero/__init__.py index 0897383fc..fb6738a29 100644 --- a/vllm_omni/model_executor/models/dreamzero/__init__.py +++ b/vllm_omni/model_executor/models/dreamzero/__init__.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from vllm_omni.model_executor.models.dreamzero.pipeline import DREAMZERO_PIPELINE - -__all__ = ["DREAMZERO_PIPELINE"] +# NOTE: Do not import model classes in this file. Importing any +# submodule in this package triggers __init__.py execution, and +# both the model registry and pipeline registry import submodules +# directly — heavy imports here would be loaded as a side effect +# even though nothing depends on these re-exports. diff --git a/vllm_omni/model_executor/models/dynin_omni/__init__.py b/vllm_omni/model_executor/models/dynin_omni/__init__.py index 2a3bae8a9..fb6738a29 100644 --- a/vllm_omni/model_executor/models/dynin_omni/__init__.py +++ b/vllm_omni/model_executor/models/dynin_omni/__init__.py @@ -1,59 +1,8 @@ -from __future__ import annotations +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from typing import TYPE_CHECKING, Any - -from .dynin_omni import DyninOmniForConditionalGeneration -from .dynin_omni_common import ( - get_dynin_magvit_attr, - get_dynin_modeling_attr, - get_dynin_sampling_attr, -) - -if TYPE_CHECKING: - from .dynin_omni_token2audio import DyninOmniToken2Audio - from .dynin_omni_token2image import DyninOmniToken2Image - from .dynin_omni_token2text import DyninOmniToken2Text - - -_STAGE_EXPORTS = { - "DyninOmniToken2Audio": (".dynin_omni_token2audio", "DyninOmniToken2Audio"), - "DyninOmniToken2Image": (".dynin_omni_token2image", "DyninOmniToken2Image"), - "DyninOmniToken2Text": (".dynin_omni_token2text", "DyninOmniToken2Text"), -} - -_MODELING_EXPORTS = {"DyninOmniConfig", "DyninOmniModelLM", "VideoTokenMerger"} -_MAGVIT_EXPORTS = {"VQGANEncoder", "VQGANDecoder", "LFQuantizer", "MAGVITv2"} - - -def __getattr__(name: str) -> Any: - if name in _STAGE_EXPORTS: - module_name, attr_name = _STAGE_EXPORTS[name] - module = __import__(module_name, globals(), locals(), [attr_name], 1) - return getattr(module, attr_name) - - if name in _MODELING_EXPORTS: - return get_dynin_modeling_attr(name) - - if name in _MAGVIT_EXPORTS: - return get_dynin_magvit_attr(name) - - if name == "get_mask_schedule": - return get_dynin_sampling_attr("get_mask_schedule") - - raise AttributeError(f"module {__name__!r} has no attribute {name!r}") - - -__all__ = [ - "DyninOmniForConditionalGeneration", - "DyninOmniToken2Audio", - "DyninOmniToken2Image", - "DyninOmniToken2Text", - "DyninOmniConfig", - "DyninOmniModelLM", - "VideoTokenMerger", - "VQGANEncoder", - "VQGANDecoder", - "LFQuantizer", - "MAGVITv2", - "get_mask_schedule", -] +# NOTE: Do not import model classes in this file. Importing any +# submodule in this package triggers __init__.py execution, and +# both the model registry and pipeline registry import submodules +# directly — heavy imports here would be loaded as a side effect +# even though nothing depends on these re-exports. diff --git a/vllm_omni/model_executor/models/glm_image/__init__.py b/vllm_omni/model_executor/models/glm_image/__init__.py index 6153f07c0..fb6738a29 100644 --- a/vllm_omni/model_executor/models/glm_image/__init__.py +++ b/vllm_omni/model_executor/models/glm_image/__init__.py @@ -1,16 +1,8 @@ -def __getattr__(name: str): - """Lazy import to avoid importing transformers.models.glm_image at module init. +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - The AR model depends on ``transformers.models.glm_image`` which is only - available when the model weights are present. Importing it eagerly in - ``__init__.py`` breaks the pipeline registry lookup for environments that - don't have this custom transformers extension installed. - """ - if name == "GlmImageForConditionalGeneration": - from .glm_image_ar import GlmImageForConditionalGeneration - - return GlmImageForConditionalGeneration - raise AttributeError(f"module {__name__!r} has no attribute {name!r}") - - -__all__ = ["GlmImageForConditionalGeneration"] +# NOTE: Do not import model classes in this file. Importing any +# submodule in this package triggers __init__.py execution, and +# both the model registry and pipeline registry import submodules +# directly — heavy imports here would be loaded as a side effect +# even though nothing depends on these re-exports. diff --git a/vllm_omni/model_executor/models/glm_tts/__init__.py b/vllm_omni/model_executor/models/glm_tts/__init__.py index 9618227ca..fb6738a29 100644 --- a/vllm_omni/model_executor/models/glm_tts/__init__.py +++ b/vllm_omni/model_executor/models/glm_tts/__init__.py @@ -1,18 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""GLM-TTS model support for vLLM-Omni. -GLM-TTS is a two-stage text-to-speech system: - - Stage 0 (LLM): Llama-based AR model generates speech tokens from text - - Stage 1 (DiT): Flow matching model converts speech tokens to mel-spectrogram - -Reference: https://github.com/zai-org/GLM-TTS -""" - -from vllm_omni.model_executor.models.glm_tts.glm_tts import ( - GLMTTSForConditionalGeneration, -) - -__all__ = [ - "GLMTTSForConditionalGeneration", -] +# NOTE: Do not import model classes in this file. Importing any +# submodule in this package triggers __init__.py execution, and +# both the model registry and pipeline registry import submodules +# directly — heavy imports here would be loaded as a side effect +# even though nothing depends on these re-exports. diff --git a/vllm_omni/model_executor/models/gr00t/__init__.py b/vllm_omni/model_executor/models/gr00t/__init__.py index 208f01a7c..fb6738a29 100644 --- a/vllm_omni/model_executor/models/gr00t/__init__.py +++ b/vllm_omni/model_executor/models/gr00t/__init__.py @@ -1,2 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# NOTE: Do not import model classes in this file. Importing any +# submodule in this package triggers __init__.py execution, and +# both the model registry and pipeline registry import submodules +# directly — heavy imports here would be loaded as a side effect +# even though nothing depends on these re-exports. diff --git a/vllm_omni/model_executor/models/higgs_audio_v2/__init__.py b/vllm_omni/model_executor/models/higgs_audio_v2/__init__.py index 1a7c23eba..fb6738a29 100644 --- a/vllm_omni/model_executor/models/higgs_audio_v2/__init__.py +++ b/vllm_omni/model_executor/models/higgs_audio_v2/__init__.py @@ -1,9 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""vllm-omni integration for boson-ai's higgs-audio v2 (two-stage TTS).""" -from vllm_omni.model_executor.models.higgs_audio_v2.configuration_higgs_audio_v2 import ( - HiggsAudioV2Config, -) - -__all__ = ["HiggsAudioV2Config"] +# NOTE: Do not import model classes in this file. Importing any +# submodule in this package triggers __init__.py execution, and +# both the model registry and pipeline registry import submodules +# directly — heavy imports here would be loaded as a side effect +# even though nothing depends on these re-exports. diff --git a/vllm_omni/model_executor/models/higgs_audio_v3/__init__.py b/vllm_omni/model_executor/models/higgs_audio_v3/__init__.py index 052d51eef..fb6738a29 100644 --- a/vllm_omni/model_executor/models/higgs_audio_v3/__init__.py +++ b/vllm_omni/model_executor/models/higgs_audio_v3/__init__.py @@ -1,9 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""vllm-omni integration for boson-ai's higgs-audio v3 (two-stage TTS).""" -from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import ( - HiggsAudioV3Config, -) - -__all__ = ["HiggsAudioV3Config"] +# NOTE: Do not import model classes in this file. Importing any +# submodule in this package triggers __init__.py execution, and +# both the model registry and pipeline registry import submodules +# directly — heavy imports here would be loaded as a side effect +# even though nothing depends on these re-exports. diff --git a/vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_code2wav.py b/vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_code2wav.py index bc82c474b..45b5d68be 100644 --- a/vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_code2wav.py +++ b/vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_code2wav.py @@ -23,11 +23,11 @@ from vllm_omni.model_executor.models.higgs_audio_v2.higgs_audio_decoder import ( build_higgs_audio_acoustic_decoder, load_higgs_audio_codec, ) -from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import ( - HiggsAudioV3Config, -) from vllm_omni.model_executor.models.output_templates import OmniOutput from vllm_omni.platforms import current_omni_platform +from vllm_omni.transformers_utils.configs.higgs_audio_v3 import ( + HiggsAudioV3Config, +) __all__ = [ "HiggsAudioV3Code2Wav", diff --git a/vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_talker.py b/vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_talker.py index 79c612dc1..f51adafa9 100644 --- a/vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_talker.py +++ b/vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_talker.py @@ -36,10 +36,10 @@ from vllm.model_executor.models.qwen3 import Qwen3Model from vllm.platforms import current_platform from vllm.v1.outputs import SamplerOutput -from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import ( +from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.transformers_utils.configs.higgs_audio_v3 import ( HiggsAudioV3Config, ) -from vllm_omni.model_executor.models.output_templates import OmniOutput __all__ = ["HiggsAudioV3TalkerForConditionalGeneration"] diff --git a/vllm_omni/model_executor/models/hunyuan_image3/__init__.py b/vllm_omni/model_executor/models/hunyuan_image3/__init__.py index a3618aefc..fb6738a29 100644 --- a/vllm_omni/model_executor/models/hunyuan_image3/__init__.py +++ b/vllm_omni/model_executor/models/hunyuan_image3/__init__.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from .hunyuan_image3 import HunyuanImage3ForConditionalGeneration - -__all__ = ["HunyuanImage3ForConditionalGeneration"] +# NOTE: Do not import model classes in this file. Importing any +# submodule in this package triggers __init__.py execution, and +# both the model registry and pipeline registry import submodules +# directly — heavy imports here would be loaded as a side effect +# even though nothing depends on these re-exports. diff --git a/vllm_omni/model_executor/models/indextts2/__init__.py b/vllm_omni/model_executor/models/indextts2/__init__.py index 208f01a7c..fb6738a29 100644 --- a/vllm_omni/model_executor/models/indextts2/__init__.py +++ b/vllm_omni/model_executor/models/indextts2/__init__.py @@ -1,2 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# NOTE: Do not import model classes in this file. Importing any +# submodule in this package triggers __init__.py execution, and +# both the model registry and pipeline registry import submodules +# directly — heavy imports here would be loaded as a side effect +# even though nothing depends on these re-exports. diff --git a/vllm_omni/model_executor/models/lance/__init__.py b/vllm_omni/model_executor/models/lance/__init__.py index 2ecd5df1f..fb6738a29 100644 --- a/vllm_omni/model_executor/models/lance/__init__.py +++ b/vllm_omni/model_executor/models/lance/__init__.py @@ -1,9 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Lance (ByteDance) — vLLM-Omni model_executor package. -Single-stage topology: a self-contained diffusion stage (``LancePipeline`` -defined in ``vllm_omni.diffusion.models.lance``) that owns the Qwen2-MoT LLM, -the Qwen2.5-VL ViT, the Wan2.2 VAE and the tokenizer — directly analogous to -``bagel_single_stage``. The ``PipelineConfig`` lives in :mod:`.pipeline`. -""" +# NOTE: Do not import model classes in this file. Importing any +# submodule in this package triggers __init__.py execution, and +# both the model registry and pipeline registry import submodules +# directly — heavy imports here would be loaded as a side effect +# even though nothing depends on these re-exports. diff --git a/vllm_omni/model_executor/models/mammoth_moda2/__init__.py b/vllm_omni/model_executor/models/mammoth_moda2/__init__.py index 80e7c2263..fb6738a29 100644 --- a/vllm_omni/model_executor/models/mammoth_moda2/__init__.py +++ b/vllm_omni/model_executor/models/mammoth_moda2/__init__.py @@ -1,17 +1,8 @@ -from transformers import AutoTokenizer +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from vllm_omni.tokenizers.mammoth_moda2_tokenizer import MammothUTokenizer -from vllm_omni.transformers_utils.configs.mammoth_moda2 import ( # noqa: F401 registers AutoConfig - Mammothmoda2Config, - Mammothmoda2Qwen2_5_VLConfig, -) - -from .mammoth_moda2 import MammothModa2ARForConditionalGeneration - -# AutoConfig.register is done inside transformers_utils/configs/mammoth_moda2.py -AutoTokenizer.register(config_class=Mammothmoda2Config, slow_tokenizer_class=MammothUTokenizer) -AutoTokenizer.register(config_class=Mammothmoda2Qwen2_5_VLConfig, slow_tokenizer_class=MammothUTokenizer) - -__all__ = [ - "MammothModa2ARForConditionalGeneration", -] +# NOTE: Do not import model classes in this file. Importing any +# submodule in this package triggers __init__.py execution, and +# both the model registry and pipeline registry import submodules +# directly — heavy imports here would be loaded as a side effect +# even though nothing depends on these re-exports. diff --git a/vllm_omni/model_executor/models/ming_flash_omni/__init__.py b/vllm_omni/model_executor/models/ming_flash_omni/__init__.py index 4cd086e64..fb6738a29 100644 --- a/vllm_omni/model_executor/models/ming_flash_omni/__init__.py +++ b/vllm_omni/model_executor/models/ming_flash_omni/__init__.py @@ -1,20 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 -# Copyright 2025 The vLLM-Omni team. +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from .ming_flash_omni import MingFlashOmniForConditionalGeneration -from .ming_flash_omni_talker import MingFlashOmniTalkerForConditionalGeneration -from .ming_flash_omni_thinker import ( - MingFlashOmniThinkerDummyInputsBuilder, - MingFlashOmniThinkerForConditionalGeneration, - MingFlashOmniThinkerMultiModalProcessor, - MingFlashOmniThinkerProcessingInfo, -) - -__all__ = [ - "MingFlashOmniForConditionalGeneration", - "MingFlashOmniTalkerForConditionalGeneration", - "MingFlashOmniThinkerForConditionalGeneration", - "MingFlashOmniThinkerProcessingInfo", - "MingFlashOmniThinkerMultiModalProcessor", - "MingFlashOmniThinkerDummyInputsBuilder", -] +# NOTE: Do not import model classes in this file. Importing any +# submodule in this package triggers __init__.py execution, and +# both the model registry and pipeline registry import submodules +# directly — heavy imports here would be loaded as a side effect +# even though nothing depends on these re-exports. diff --git a/vllm_omni/model_executor/models/ming_tts/__init__.py b/vllm_omni/model_executor/models/ming_tts/__init__.py index 35afa2ce5..fb6738a29 100644 --- a/vllm_omni/model_executor/models/ming_tts/__init__.py +++ b/vllm_omni/model_executor/models/ming_tts/__init__.py @@ -1,13 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from .config_ming_tts import MingDenseConfig -from .ming_tts import MingTTSForConditionalGeneration -from .ming_tts_audio_vae import MingAudioVAEModel -from .ming_tts_llm import MingLLMModel -__all__ = [ - "MingDenseConfig", - "MingTTSForConditionalGeneration", - "MingLLMModel", - "MingAudioVAEModel", -] +# NOTE: Do not import model classes in this file. Importing any +# submodule in this package triggers __init__.py execution, and +# both the model registry and pipeline registry import submodules +# directly — heavy imports here would be loaded as a side effect +# even though nothing depends on these re-exports. diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/__init__.py b/vllm_omni/model_executor/models/minicpmo_4_5/__init__.py index 7284c5b7f..fb6738a29 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/__init__.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/__init__.py @@ -1,9 +1,8 @@ -from .minicpmo_4_5_omni import MiniCPMO45OmniForConditionalGeneration -from .minicpmo_4_5_omni_llm import MiniCPMO45OmniLLMForConditionalGeneration -from .minicpmo_4_5_omni_tts import MiniCPMO45OmniTTSForConditionalGeneration +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -__all__ = [ - "MiniCPMO45OmniForConditionalGeneration", - "MiniCPMO45OmniLLMForConditionalGeneration", - "MiniCPMO45OmniTTSForConditionalGeneration", -] +# NOTE: Do not import model classes in this file. Importing any +# submodule in this package triggers __init__.py execution, and +# both the model registry and pipeline registry import submodules +# directly — heavy imports here would be loaded as a side effect +# even though nothing depends on these re-exports. diff --git a/vllm_omni/model_executor/models/moss_tts/__init__.py b/vllm_omni/model_executor/models/moss_tts/__init__.py index 23a73b3fe..fb6738a29 100644 --- a/vllm_omni/model_executor/models/moss_tts/__init__.py +++ b/vllm_omni/model_executor/models/moss_tts/__init__.py @@ -1,13 +1,8 @@ -from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_codec import ( - MossTTSCodecDecoder, -) -from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_talker import ( - MossTTSDelayTalkerForGeneration, - MossTTSRealtimeTalkerForGeneration, -) +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -__all__ = [ - "MossTTSDelayTalkerForGeneration", - "MossTTSRealtimeTalkerForGeneration", - "MossTTSCodecDecoder", -] +# NOTE: Do not import model classes in this file. Importing any +# submodule in this package triggers __init__.py execution, and +# both the model registry and pipeline registry import submodules +# directly — heavy imports here would be loaded as a side effect +# even though nothing depends on these re-exports. diff --git a/vllm_omni/model_executor/models/moss_tts_nano/__init__.py b/vllm_omni/model_executor/models/moss_tts_nano/__init__.py index 896ead57d..fb6738a29 100644 --- a/vllm_omni/model_executor/models/moss_tts_nano/__init__.py +++ b/vllm_omni/model_executor/models/moss_tts_nano/__init__.py @@ -1,5 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from .modeling_moss_tts_nano import MossTTSNanoForGeneration -__all__ = ["MossTTSNanoForGeneration"] +# NOTE: Do not import model classes in this file. Importing any +# submodule in this package triggers __init__.py execution, and +# both the model registry and pipeline registry import submodules +# directly — heavy imports here would be loaded as a side effect +# even though nothing depends on these re-exports. diff --git a/vllm_omni/model_executor/models/qwen3_omni/__init__.py b/vllm_omni/model_executor/models/qwen3_omni/__init__.py index a469704f6..fb6738a29 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/__init__.py +++ b/vllm_omni/model_executor/models/qwen3_omni/__init__.py @@ -1,3 +1,8 @@ -from .qwen3_omni import Qwen3OmniMoeForConditionalGeneration +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -__all__ = ["Qwen3OmniMoeForConditionalGeneration"] +# NOTE: Do not import model classes in this file. Importing any +# submodule in this package triggers __init__.py execution, and +# both the model registry and pipeline registry import submodules +# directly — heavy imports here would be loaded as a side effect +# even though nothing depends on these re-exports. diff --git a/vllm_omni/model_executor/models/qwen3_omni/pipeline.py b/vllm_omni/model_executor/models/qwen3_omni/pipeline.py index 1f3df4206..07debc724 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/pipeline.py +++ b/vllm_omni/model_executor/models/qwen3_omni/pipeline.py @@ -7,10 +7,13 @@ Stage 1: Talker — text embeddings → RVQ codec codes Stage 2: Code2Wav — RVQ codes → audio waveform """ +from transformers import Qwen3OmniMoeConfig + from vllm_omni.config.stage_config import ( PipelineConfig, StageExecutionType, StagePipelineConfig, + pipeline_cfg_resolver, ) _PROC = "vllm_omni.model_executor.stage_input_processors.qwen3_omni" @@ -64,3 +67,38 @@ QWEN3_OMNI_PIPELINE = PipelineConfig( ), ), ) + +QWEN3_OMNI_THINKER_ONLY_PIPELINE = PipelineConfig( + model_type="qwen3_omni_moe_thinker_only", + model_arch="Qwen3OmniMoeForConditionalGeneration", + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="thinker", + execution_type=StageExecutionType.LLM_AR, + input_sources=(), + final_output=True, + final_output_type="text", + owns_tokenizer=True, + requires_multimodal_data=True, + hf_config_name="thinker_config", + engine_output_type="latent", + sampling_constraints={"detokenize": True}, + ), + ), +) + + +@pipeline_cfg_resolver(config_type=Qwen3OmniMoeConfig) +def resolve_qwen3_omni_pipeline( + hf_config: Qwen3OmniMoeConfig, +) -> PipelineConfig: + """Select the right pipeline variant based on the HF config, since some variants, + e.g., Qwen3-Omni-30B-A3B-Captioner, are thinker only. + + By default, we load the full pipeline, as this is the common case. + """ + # If we have a config and it explicitly disabled audio input, load thinker only + if not hf_config.enable_audio_output: + return QWEN3_OMNI_THINKER_ONLY_PIPELINE + return QWEN3_OMNI_PIPELINE diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py index 9fd929036..4b27e4ba0 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py @@ -106,9 +106,6 @@ class Qwen3OmniMoeTalkerForConditionalGeneration( self.hidden_projection = Qwen3OmniMoeTalkerResizeMLP(self.config) self.codec_head = nn.Linear(self.config.text_config.hidden_size, self.config.text_config.vocab_size, bias=False) - self.rope_deltas = None - self.spatial_merge_size = self.config.spatial_merge_size - self.vocab_size = self.config.code_predictor_config.vocab_size self.num_code_groups = self.config.code_predictor_config.num_code_groups diff --git a/vllm_omni/model_executor/models/voxcpm2/__init__.py b/vllm_omni/model_executor/models/voxcpm2/__init__.py index 77bd8dfb5..fb6738a29 100644 --- a/vllm_omni/model_executor/models/voxcpm2/__init__.py +++ b/vllm_omni/model_executor/models/voxcpm2/__init__.py @@ -1,5 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from .voxcpm2_talker import VoxCPM2TalkerForConditionalGeneration -__all__ = ["VoxCPM2TalkerForConditionalGeneration"] +# NOTE: Do not import model classes in this file. Importing any +# submodule in this package triggers __init__.py execution, and +# both the model registry and pipeline registry import submodules +# directly — heavy imports here would be loaded as a side effect +# even though nothing depends on these re-exports. diff --git a/vllm_omni/transformers_utils/configs/__init__.py b/vllm_omni/transformers_utils/configs/__init__.py index ba6d329ea..88a23dc0c 100644 --- a/vllm_omni/transformers_utils/configs/__init__.py +++ b/vllm_omni/transformers_utils/configs/__init__.py @@ -10,6 +10,7 @@ from __future__ import annotations import importlib _CLASS_TO_MODULE: dict[str, str] = { + "HiggsAudioV3Config": "vllm_omni.transformers_utils.configs.higgs_audio_v3", "Mammothmoda2Config": "vllm_omni.transformers_utils.configs.mammoth_moda2", "Mammothmoda2Qwen2_5_VLConfig": "vllm_omni.transformers_utils.configs.mammoth_moda2", "Mammothmoda2Qwen2_5_VLTextConfig": "vllm_omni.transformers_utils.configs.mammoth_moda2", @@ -31,6 +32,7 @@ _CLASS_TO_MODULE: dict[str, str] = { } __all__ = [ + "HiggsAudioV3Config", "Mammothmoda2Config", "Mammothmoda2Qwen2_5_VLConfig", "Mammothmoda2Qwen2_5_VLTextConfig", @@ -70,6 +72,7 @@ def __dir__(): from vllm_omni.transformers_utils.configs import cosyvoice3 as _cosyvoice3 # noqa: F401, E402 from vllm_omni.transformers_utils.configs import fish_speech as _fish_speech # noqa: F401, E402 from vllm_omni.transformers_utils.configs import glm_tts as _glm_tts # noqa: F401, E402 +from vllm_omni.transformers_utils.configs import higgs_audio_v3 as _higgs_audio_v3 # noqa: F401, E402 from vllm_omni.transformers_utils.configs import mammoth_moda2 as _mammoth_moda2 # noqa: F401, E402 from vllm_omni.transformers_utils.configs import ming_flash_omni as _ming_flash_omni # noqa: F401, E402 from vllm_omni.transformers_utils.configs import omnivoice as _omnivoice # noqa: F401, E402 diff --git a/vllm_omni/model_executor/models/higgs_audio_v3/configuration_higgs_audio_v3.py b/vllm_omni/transformers_utils/configs/higgs_audio_v3.py similarity index 100% rename from vllm_omni/model_executor/models/higgs_audio_v3/configuration_higgs_audio_v3.py rename to vllm_omni/transformers_utils/configs/higgs_audio_v3.py diff --git a/vllm_omni/transformers_utils/configs/mammoth_moda2.py b/vllm_omni/transformers_utils/configs/mammoth_moda2.py index 4acd098e2..93cddd735 100644 --- a/vllm_omni/transformers_utils/configs/mammoth_moda2.py +++ b/vllm_omni/transformers_utils/configs/mammoth_moda2.py @@ -2,13 +2,15 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from typing import ClassVar, Literal -from transformers import AutoConfig, PretrainedConfig +from transformers import AutoConfig, AutoTokenizer, PretrainedConfig from transformers.models.qwen2_5_vl.configuration_qwen2_5_vl import ( Qwen2_5_VLConfig, Qwen2_5_VLTextConfig, Qwen2_5_VLVisionConfig, ) +from vllm_omni.tokenizers.mammoth_moda2_tokenizer import MammothUTokenizer + __all__ = [ "Mammothmoda2Config", "Mammothmoda2Qwen2_5_VLConfig", @@ -281,3 +283,6 @@ AutoConfig.register(Mammothmoda2Config.model_type, Mammothmoda2Config) AutoConfig.register(Mammothmoda2Qwen2_5_VLConfig.model_type, Mammothmoda2Qwen2_5_VLConfig) AutoConfig.register(Mammothmoda2Qwen2_5_VLTextConfig.model_type, Mammothmoda2Qwen2_5_VLTextConfig) AutoConfig.register(Mammothmoda2Qwen2_5_VLVisionConfig.model_type, Mammothmoda2Qwen2_5_VLVisionConfig) +# Register tokenizer_type -> the configs & AutoTokenizer +AutoTokenizer.register(config_class=Mammothmoda2Config, slow_tokenizer_class=MammothUTokenizer) +AutoTokenizer.register(config_class=Mammothmoda2Qwen2_5_VLConfig, slow_tokenizer_class=MammothUTokenizer)