anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[Feat] Support Pipeline Resolution Based on HF Config & Make Pipeline Registry Eager (#3760)

Browse Source 
Signed-off-by: Alex Brooks <albrooks@redhat.com>
This commit is contained in:
Alex Brooks authored and GitHub committed 2026-06-24 11:02:29 +08:00
1 parent f6894f4308
commit 465f1464d2
44 files changed
+1038 -1126
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tests/config/test_pipeline_registry.py
+53 -131
View File
@@ -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)
tests/conftest.py
+1
View File
@@ -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",
tests/e2e/offline_inference/test_qwen3_omni.py
+17 -1
View File
@@ -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)
tests/engine/test_async_omni_engine_abort.py
+2 -2
View File
@@ -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")
tests/entrypoints/test_resolve_dreamzero_config.py
+1 -1
View File
@@ -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(
tests/helpers/fixtures/config.py
+14
View File
@@ -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)
tests/model_executor/models/glm_image/test_glm_image_init.py
-76
View File
@@ -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__")
tests/model_executor/models/minicpmo_4_5/test_pipeline.py
+7 -8
View File
@@ -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."
tests/test_config_factory.py
+252 -73
View File
@@ -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
tests/test_diffusion_config_propagation.py
+1 -1
View File
@@ -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,
tests/unit/higgs_audio_v3/test_higgs_audio_v3.py
+20 -20
View File
@@ -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
tests/utils/test_tracking_parser.py
+9 -4
View File
@@ -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,
)
vllm_omni/config/__init__.py
-4
View File
@@ -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",
vllm_omni/config/config_factory.py
+338
View File
@@ -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
vllm_omni/config/pipeline_registry.py
+128 -162
View File
@@ -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
vllm_omni/config/stage_config.py
+22 -449
View File
@@ -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)
vllm_omni/entrypoints/utils.py
+1 -1
View File
@@ -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
vllm_omni/model_executor/models/aura_omni/__init__.py
+8
View File
@@ -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.
vllm_omni/model_executor/models/bagel/__init__.py
+7 -2
View File
@@ -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.
vllm_omni/model_executor/models/dreamzero/__init__.py
+5 -3
View File
@@ -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.
vllm_omni/model_executor/models/dynin_omni/__init__.py
+7 -58
View File
@@ -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.
vllm_omni/model_executor/models/glm_image/__init__.py
+7 -15
View File
@@ -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.
vllm_omni/model_executor/models/glm_tts/__init__.py
+5 -15
View File
@@ -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.
vllm_omni/model_executor/models/gr00t/__init__.py
+6
View File
@@ -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.
vllm_omni/model_executor/models/higgs_audio_v2/__init__.py
+5 -6
View File
@@ -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.
vllm_omni/model_executor/models/higgs_audio_v3/__init__.py
+5 -6
View File
@@ -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.
vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_code2wav.py
+3 -3
View File
@@ -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",
vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_talker.py
+2 -2
View File
@@ -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"]
vllm_omni/model_executor/models/hunyuan_image3/__init__.py
+5 -3
View File
@@ -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.
vllm_omni/model_executor/models/indextts2/__init__.py
+6
View File
@@ -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.
vllm_omni/model_executor/models/lance/__init__.py
+5 -6
View File
@@ -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.
vllm_omni/model_executor/models/mammoth_moda2/__init__.py
+7 -16
View File
@@ -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.
vllm_omni/model_executor/models/ming_flash_omni/__init__.py
+6 -18
View File
@@ -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.
vllm_omni/model_executor/models/ming_tts/__init__.py
+5 -10
View File
@@ -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.
vllm_omni/model_executor/models/minicpmo_4_5/__init__.py
+7 -8
View File
@@ -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.
vllm_omni/model_executor/models/moss_tts/__init__.py
+7 -12
View File
@@ -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.
vllm_omni/model_executor/models/moss_tts_nano/__init__.py
+5 -2
View File
@@ -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.
vllm_omni/model_executor/models/qwen3_omni/__init__.py
+7 -2
View File
@@ -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.
vllm_omni/model_executor/models/qwen3_omni/pipeline.py
+38
View File
@@ -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
vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py
-3
View File
@@ -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
vllm_omni/model_executor/models/voxcpm2/__init__.py
+5 -2
View File
@@ -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.
vllm_omni/transformers_utils/configs/__init__.py
+3
View File
@@ -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
vllm_omni/model_executor/models/higgs_audio_v3/configuration_higgs_audio_v3.py → vllm_omni/transformers_utils/configs/higgs_audio_v3.py
View File
File renamed without changes.
vllm_omni/transformers_utils/configs/mammoth_moda2.py
+6 -1
View File
@@ -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)