[Bugfix][Refactor] Migrate Voxtral TTS config and parser registry (#3065)

Signed-off-by: Yuanheng Zhao <jonathan.zhaoyh@gmail.com>

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
 

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),
     ]:

vllm_omni/model_executor/models/voxtral_tts/configuration_voxtral_tts.py

@@ -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

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

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"]

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

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"]