From 6b52db9e2b08621bd6625460edd7cd98f299b389 Mon Sep 17 00:00:00 2001 From: Yuanheng Zhao <54058983+yuanheng-zhao@users.noreply.github.com> Date: Sat, 25 Apr 2026 18:17:25 +0800 Subject: [PATCH] [Bugfix][Refactor] Migrate Voxtral TTS config and parser registry (#3065) Signed-off-by: Yuanheng Zhao --- vllm_omni/__init__.py | 1 + vllm_omni/engine/arg_utils.py | 4 - .../voxtral_tts/configuration_voxtral_tts.py | 108 ------------------ .../transformers_utils/configs/__init__.py | 3 + .../transformers_utils/configs/voxtral_tts.py | 36 ++++++ .../transformers_utils/parsers/__init__.py | 29 +++++ .../transformers_utils/parsers/voxtral_tts.py | 106 +++++++++++++++++ 7 files changed, 175 insertions(+), 112 deletions(-) delete mode 100644 vllm_omni/model_executor/models/voxtral_tts/configuration_voxtral_tts.py create mode 100644 vllm_omni/transformers_utils/configs/voxtral_tts.py create mode 100644 vllm_omni/transformers_utils/parsers/__init__.py create mode 100644 vllm_omni/transformers_utils/parsers/voxtral_tts.py diff --git a/vllm_omni/__init__.py b/vllm_omni/__init__.py index 65ad79c72..f3c9d9afd 100644 --- a/vllm_omni/__init__.py +++ b/vllm_omni/__init__.py @@ -28,6 +28,7 @@ except ModuleNotFoundError as exc: # pragma: no cover - optional dependency # Register custom configs (AutoConfig, AutoTokenizer) as early as possible. from vllm_omni.transformers_utils import configs as _configs # noqa: F401, E402 +from vllm_omni.transformers_utils import parsers as _parsers # noqa: F401, E402 from .config import OmniModelConfig diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index a6fff97e7..3f16c329e 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -40,9 +40,6 @@ def _register_omni_hf_configs() -> None: from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import ( Qwen3TTSConfig, ) - from vllm_omni.model_executor.models.voxtral_tts.configuration_voxtral_tts import ( - VoxtralTTSConfig, - ) from vllm_omni.transformers_utils.configs.voxcpm import VoxCPMConfig from vllm_omni.transformers_utils.configs.voxcpm2 import VoxCPM2Config except Exception as exc: # pragma: no cover - best-effort optional registration @@ -61,7 +58,6 @@ def _register_omni_hf_configs() -> None: ("qwen3_tts", Qwen3TTSConfig), ("cosyvoice3", CosyVoice3Config), ("omnivoice", OmniVoiceConfig), - ("voxtral_tts", VoxtralTTSConfig), ("voxcpm", VoxCPMConfig), ("voxcpm2", VoxCPM2Config), ]: diff --git a/vllm_omni/model_executor/models/voxtral_tts/configuration_voxtral_tts.py b/vllm_omni/model_executor/models/voxtral_tts/configuration_voxtral_tts.py deleted file mode 100644 index 0f22c764a..000000000 --- a/vllm_omni/model_executor/models/voxtral_tts/configuration_voxtral_tts.py +++ /dev/null @@ -1,108 +0,0 @@ -from __future__ import annotations - -from pathlib import Path -from typing import Any - -from transformers import PretrainedConfig -from vllm.logger import init_logger -from vllm.transformers_utils.config import MistralConfigParser, register_config_parser - -logger = init_logger(__name__) - - -class VoxtralTTSConfig(PretrainedConfig): - """HuggingFace-style config for Voxtral TTS models.""" - - model_type = "voxtral_tts" - - def __init__( - self, - text_config: PretrainedConfig | dict | None = None, - audio_config: dict[str, Any] | None = None, - **kwargs: Any, - ) -> None: - super().__init__(**kwargs) - - if isinstance(text_config, PretrainedConfig): - self.text_config = text_config - elif isinstance(text_config, dict): - self.text_config = PretrainedConfig.from_dict(text_config) - else: - self.text_config = PretrainedConfig() - - self.audio_config = audio_config or {} - - def get_text_config(self, **kwargs: Any) -> PretrainedConfig: - return self.text_config - - -@register_config_parser("mistral") -class VoxtralTTSConfigParser(MistralConfigParser): - """Config parser that extends the base Mistral parser with TTS support. - - This only support voxtral_tts for now. - """ - - def _remap_mistral_audio_args(self, config_dict: dict) -> dict: - encoder_args = config_dict["multimodal"].pop("audio_model_args") - audio_tokenizer_args = config_dict["multimodal"].pop("audio_tokenizer_args", None) - audio_config = {} - if encoder_args is not None: - # Default n_decoding_steps if not provided - acoustic_args = encoder_args.get("acoustic_transformer_args", {}) - if acoustic_args.get("n_decoding_steps") is None: - logger.warning( - "n_decoding_steps not provided in acoustic_transformer_args, defaulting to 7. " - "Please add 'n_decoding_steps' to params.json under acoustic_transformer_args." - ) - acoustic_args["n_decoding_steps"] = 7 - - audio_config = { - "sampling_rate": encoder_args["audio_encoding_args"]["sampling_rate"], - "codec_args": audio_tokenizer_args, - "audio_model_args": encoder_args, - "speaker_id": audio_tokenizer_args.get("voice", {}), - } - return audio_config - - def parse( - self, - model: str | Path, - trust_remote_code: bool, - revision: str | None = None, - code_revision: str | None = None, - **kwargs: Any, - ) -> tuple[dict, PretrainedConfig]: - from vllm.transformers_utils.config import ( - _download_mistral_config_file, - ) - - config_dict = _download_mistral_config_file(model, revision) - - from vllm.transformers_utils.configs.mistral import ( - _remap_general_mistral_args, - _remap_mistral_quantization_args, - ) - - # Extract audio config before building text config - audio_config = {} - if (config_dict.get("multimodal") or {}).get("audio_model_args"): - audio_config = self._remap_mistral_audio_args(config_dict) - - # Build text_config from the top-level keys - non_text_keys = {"multimodal"} - text_config = {k: v for k, v in config_dict.items() if k not in non_text_keys} - text_config = _remap_general_mistral_args(text_config) - if text_config.get("quantization"): - text_config = _remap_mistral_quantization_args(text_config) - - # The text sub-model is a plain MistralForCausalLM - text_config.setdefault("architectures", ["MistralForCausalLM"]) - - config = VoxtralTTSConfig( - text_config=PretrainedConfig.from_dict(text_config), - audio_config=audio_config, - architectures=config_dict.get("architectures", ["VoxtralTTSForConditionalGeneration"]), - ) - - return config_dict, config diff --git a/vllm_omni/transformers_utils/configs/__init__.py b/vllm_omni/transformers_utils/configs/__init__.py index 598ac3a96..cb9b8418a 100644 --- a/vllm_omni/transformers_utils/configs/__init__.py +++ b/vllm_omni/transformers_utils/configs/__init__.py @@ -19,6 +19,7 @@ _CLASS_TO_MODULE: dict[str, str] = { "FishSpeechFastARConfig": "vllm_omni.transformers_utils.configs.fish_speech", "VoxCPMConfig": "vllm_omni.transformers_utils.configs.voxcpm", "VoxCPM2Config": "vllm_omni.transformers_utils.configs.voxcpm2", + "VoxtralTTSConfig": "vllm_omni.transformers_utils.configs.voxtral_tts", "BailingMoeV2Config": "vllm_omni.transformers_utils.configs.ming_flash_omni", "BailingMM2Config": "vllm_omni.transformers_utils.configs.ming_flash_omni", "MingFlashOmniConfig": "vllm_omni.transformers_utils.configs.ming_flash_omni", @@ -36,6 +37,7 @@ __all__ = [ "FishSpeechFastARConfig", "VoxCPMConfig", "VoxCPM2Config", + "VoxtralTTSConfig", "BailingMoeV2Config", "BailingMM2Config", "MingFlashOmniConfig", @@ -64,3 +66,4 @@ from vllm_omni.transformers_utils.configs import mammoth_moda2 as _mammoth_moda2 from vllm_omni.transformers_utils.configs import ming_flash_omni as _ming_flash_omni # noqa: F401, E402 from vllm_omni.transformers_utils.configs import voxcpm as _voxcpm # noqa: F401, E402 from vllm_omni.transformers_utils.configs import voxcpm2 as _voxcpm2 # noqa: F401, E402 +from vllm_omni.transformers_utils.configs import voxtral_tts as _voxtral_tts # noqa: F401, E402 diff --git a/vllm_omni/transformers_utils/configs/voxtral_tts.py b/vllm_omni/transformers_utils/configs/voxtral_tts.py new file mode 100644 index 000000000..19af7b020 --- /dev/null +++ b/vllm_omni/transformers_utils/configs/voxtral_tts.py @@ -0,0 +1,36 @@ +from __future__ import annotations + +from typing import Any + +from transformers import AutoConfig, PretrainedConfig + + +class VoxtralTTSConfig(PretrainedConfig): + """HuggingFace-style config for Voxtral TTS models.""" + + model_type = "voxtral_tts" + + def __init__( + self, + text_config: PretrainedConfig | dict | None = None, + audio_config: dict[str, Any] | None = None, + **kwargs: Any, + ) -> None: + super().__init__(**kwargs) + + if isinstance(text_config, PretrainedConfig): + self.text_config = text_config + elif isinstance(text_config, dict): + self.text_config = PretrainedConfig.from_dict(text_config) + else: + self.text_config = PretrainedConfig() + + self.audio_config = audio_config or {} + + def get_text_config(self, **kwargs: Any) -> PretrainedConfig: + return self.text_config + + +AutoConfig.register("voxtral_tts", VoxtralTTSConfig) + +__all__ = ["VoxtralTTSConfig"] diff --git a/vllm_omni/transformers_utils/parsers/__init__.py b/vllm_omni/transformers_utils/parsers/__init__.py new file mode 100644 index 000000000..eed3d3f7d --- /dev/null +++ b/vllm_omni/transformers_utils/parsers/__init__.py @@ -0,0 +1,29 @@ +"""Custom vLLM config parsers for vllm-omni.""" + +from __future__ import annotations + +import importlib + +_CLASS_TO_MODULE: dict[str, str] = { + "VoxtralTTSConfigParser": "vllm_omni.transformers_utils.parsers.voxtral_tts", +} + +__all__ = ["VoxtralTTSConfigParser"] + + +def __getattr__(name: str): + if name in _CLASS_TO_MODULE: + module_name = _CLASS_TO_MODULE[name] + module = importlib.import_module(module_name) + return getattr(module, name) + + raise AttributeError(f"module 'vllm_omni.transformers_utils.parsers' has no attribute {name!r}") + + +def __dir__(): + return sorted(list(__all__)) + + +# Eagerly import parser modules so their registry side-effects run as soon as +# `vllm_omni.transformers_utils.parsers` is imported. +from vllm_omni.transformers_utils.parsers import voxtral_tts as _voxtral_tts # noqa: F401, E402 diff --git a/vllm_omni/transformers_utils/parsers/voxtral_tts.py b/vllm_omni/transformers_utils/parsers/voxtral_tts.py new file mode 100644 index 000000000..d4669258a --- /dev/null +++ b/vllm_omni/transformers_utils/parsers/voxtral_tts.py @@ -0,0 +1,106 @@ +from __future__ import annotations + +from pathlib import Path +from typing import Any + +from transformers import PretrainedConfig +from vllm.logger import init_logger +from vllm.transformers_utils.config import ( + _CONFIG_FORMAT_TO_CONFIG_PARSER, + MistralConfigParser, + _download_mistral_config_file, +) + +from vllm_omni.transformers_utils.configs.voxtral_tts import VoxtralTTSConfig + +logger = init_logger(__name__) + +_VOXTRAL_TTS_ARCHS = frozenset({"VoxtralTTSForConditionalGeneration"}) +_VOXTRAL_TTS_MODEL_TYPE = "voxtral_tts" + + +def _is_voxtral_tts_params(config_dict: dict) -> bool: + """Return True if the Mistral params.json describes a Voxtral-TTS model""" + if config_dict.get("model_type") == _VOXTRAL_TTS_MODEL_TYPE: + return True + architectures = set(config_dict.get("architectures") or []) + return bool(architectures & _VOXTRAL_TTS_ARCHS) + + +def _remap_voxtral_tts_audio_args(config_dict: dict) -> dict: + encoder_args = config_dict["multimodal"].pop("audio_model_args") + audio_tokenizer_args = config_dict["multimodal"].pop("audio_tokenizer_args", None) + if encoder_args is None: + return {} + + acoustic_args = encoder_args.get("acoustic_transformer_args", {}) + if acoustic_args.get("n_decoding_steps") is None: + logger.warning( + "n_decoding_steps not provided in acoustic_transformer_args, defaulting to 7. " + "Please add 'n_decoding_steps' to params.json under acoustic_transformer_args." + ) + acoustic_args["n_decoding_steps"] = 7 + + return { + "sampling_rate": encoder_args["audio_encoding_args"]["sampling_rate"], + "codec_args": audio_tokenizer_args, + "audio_model_args": encoder_args, + "speaker_id": (audio_tokenizer_args or {}).get("voice", {}), + } + + +def _parse_voxtral_tts(config_dict: dict) -> tuple[dict, PretrainedConfig]: + from vllm.transformers_utils.configs.mistral import ( + _remap_general_mistral_args, + _remap_mistral_quantization_args, + ) + + audio_config: dict[str, Any] = {} + if (config_dict.get("multimodal") or {}).get("audio_model_args"): + audio_config = _remap_voxtral_tts_audio_args(config_dict) + + text_config = {k: v for k, v in config_dict.items() if k != "multimodal"} + text_config = _remap_general_mistral_args(text_config) + if text_config.get("quantization"): + text_config = _remap_mistral_quantization_args(text_config) + text_config.setdefault("architectures", ["MistralForCausalLM"]) + + config = VoxtralTTSConfig( + text_config=PretrainedConfig.from_dict(text_config), + audio_config=audio_config, + architectures=config_dict.get("architectures", ["VoxtralTTSForConditionalGeneration"]), + ) + return config_dict, config + + +class VoxtralTTSConfigParser(MistralConfigParser): + """Mistral parser that also recognizes Voxtral-TTS checkpoints.""" + + def parse( + self, + model: str | Path, + trust_remote_code: bool, + revision: str | None = None, + code_revision: str | None = None, + **kwargs: Any, + ) -> tuple[dict, PretrainedConfig]: + config_dict = _download_mistral_config_file(model, revision) + + if _is_voxtral_tts_params(config_dict): + return _parse_voxtral_tts(config_dict) + + return super().parse( + model, + trust_remote_code, + revision=revision, + code_revision=code_revision, + **kwargs, + ) + + +# Replace the default "mistral" slot directly. +# Any non-Voxtral-TTS Mistral ckpt still goes through +# the upstream code path via super().parse(). +_CONFIG_FORMAT_TO_CONFIG_PARSER["mistral"] = VoxtralTTSConfigParser + +__all__ = ["VoxtralTTSConfigParser"]