186 lines
8.3 KiB
Python
186 lines
8.3 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
|
|
from transformers import PretrainedConfig
|
|
|
|
from vllm_omni.model_executor.models.common.ming.audio_vae import AudioVAEConfig
|
|
|
|
from .constants import (
|
|
AGGREGATOR_HIDDEN_SIZE,
|
|
HISTORY_PATCH_SIZE,
|
|
LATENT_DIM,
|
|
LLM_HIDDEN_SIZE,
|
|
LLM_VOCAB_SIZE,
|
|
MOE_AUDIO_DUMMY_TOKEN_ID,
|
|
MOE_AUDIO_EOS_TOKEN_ID,
|
|
MOE_TEXT_EOS_TOKEN_ID,
|
|
PATCH_SIZE,
|
|
SAMPLE_RATE,
|
|
)
|
|
|
|
|
|
def _to_plain_dict(obj: Any) -> dict[str, Any]:
|
|
"""Normalize nested config objects into plain dicts when possible."""
|
|
if obj is None:
|
|
return {}
|
|
if isinstance(obj, dict):
|
|
return dict(obj)
|
|
if isinstance(obj, PretrainedConfig):
|
|
return obj.to_dict()
|
|
if hasattr(obj, "to_dict") and callable(obj.to_dict):
|
|
try:
|
|
return dict(obj.to_dict())
|
|
except Exception:
|
|
pass
|
|
try:
|
|
return dict(vars(obj))
|
|
except Exception:
|
|
return {}
|
|
|
|
|
|
def _coerce_audio_vae_config(atc_raw: Any) -> AudioVAEConfig | None:
|
|
if atc_raw is None:
|
|
return None
|
|
if isinstance(atc_raw, AudioVAEConfig):
|
|
return atc_raw
|
|
if isinstance(atc_raw, PretrainedConfig):
|
|
atc_dict = atc_raw.to_dict()
|
|
elif isinstance(atc_raw, dict):
|
|
atc_dict = dict(atc_raw)
|
|
elif hasattr(atc_raw, "to_dict") and callable(atc_raw.to_dict):
|
|
atc_dict = dict(atc_raw.to_dict())
|
|
else:
|
|
raise TypeError(f"Unsupported audio_tokenizer_config type for Ming dense config: {type(atc_raw)!r}")
|
|
|
|
return AudioVAEConfig(**atc_dict)
|
|
|
|
|
|
def _nested_get(obj: Any, *keys: str, default: Any = None) -> Any:
|
|
"""Safe nested attribute/key access for dicts and config-like objects."""
|
|
cur = obj
|
|
for key in keys:
|
|
if cur is None:
|
|
return default
|
|
if isinstance(cur, dict):
|
|
cur = cur.get(key)
|
|
else:
|
|
cur = getattr(cur, key, None)
|
|
return cur if cur is not None else default
|
|
|
|
|
|
def validate_ming_tts_config(cfg: Any) -> None:
|
|
"""Run before GPU allocation/weight loading. Raises ValueError on mismatches."""
|
|
is_moe = getattr(cfg, "model_variant", "dense") == "moe"
|
|
exp_audio_dummy = MOE_AUDIO_DUMMY_TOKEN_ID if is_moe else 151705
|
|
exp_audio_eos = MOE_AUDIO_EOS_TOKEN_ID if is_moe else 151704
|
|
exp_text_eos = MOE_TEXT_EOS_TOKEN_ID if is_moe else 151669
|
|
if cfg.audio_dummy_token_id != exp_audio_dummy:
|
|
raise ValueError(
|
|
f"audio_dummy_token_id={cfg.audio_dummy_token_id}, expected {exp_audio_dummy} (<audioPatch>). "
|
|
"Wrong tokenizer/checkpoint?"
|
|
)
|
|
if cfg.audio_eos_token_id != exp_audio_eos:
|
|
raise ValueError(
|
|
f"audio_eos_token_id={cfg.audio_eos_token_id}, expected {exp_audio_eos} (<end_of_audio>). "
|
|
"Wrong tokenizer/checkpoint?"
|
|
)
|
|
if cfg.text_eos_token_id != exp_text_eos:
|
|
raise ValueError(
|
|
f"text_eos_token_id={cfg.text_eos_token_id}, expected {exp_text_eos}. Wrong tokenizer/checkpoint?"
|
|
)
|
|
|
|
if cfg.audio_tokenizer_config is None:
|
|
raise ValueError("audio_tokenizer_config is None. Nested AudioVAE config was not deserialized correctly.")
|
|
|
|
if cfg.latent_dim != LATENT_DIM:
|
|
raise ValueError(
|
|
f"latent_dim mismatch: got {cfg.latent_dim}, expected {LATENT_DIM}. "
|
|
"Check audio_tokenizer_config.enc_kwargs.latent_dim."
|
|
)
|
|
if cfg.patch_size != PATCH_SIZE:
|
|
raise ValueError(
|
|
f"patch_size mismatch: got {cfg.patch_size}, expected {PATCH_SIZE}. Check ditar_config.patch_size."
|
|
)
|
|
if cfg.history_patch_size != HISTORY_PATCH_SIZE:
|
|
raise ValueError(
|
|
f"history_patch_size mismatch: got {cfg.history_patch_size}, expected {HISTORY_PATCH_SIZE}. "
|
|
"Check ditar_config.history_patch_size."
|
|
)
|
|
# Absolute hidden/vocab pins are dense-only; the MoE backbone carries its
|
|
# own (larger) dims, cross-checked against llm_config.hidden_size below.
|
|
if not is_moe:
|
|
if cfg.llm_hidden_size != LLM_HIDDEN_SIZE:
|
|
raise ValueError(
|
|
f"llm_hidden_size mismatch: got {cfg.llm_hidden_size}, expected {LLM_HIDDEN_SIZE}. "
|
|
"Check llm_config.hidden_size."
|
|
)
|
|
if cfg.llm_vocab_size != LLM_VOCAB_SIZE:
|
|
raise ValueError(f"llm_vocab_size mismatch: got {cfg.llm_vocab_size}, expected {LLM_VOCAB_SIZE}.")
|
|
if cfg.sample_rate != SAMPLE_RATE:
|
|
raise ValueError(f"sample_rate mismatch: got {cfg.sample_rate}, expected {SAMPLE_RATE}.")
|
|
|
|
if cfg.vae_patch_size != cfg.patch_size:
|
|
raise ValueError(f"VAE patch size ({cfg.vae_patch_size}) != flow/DiT patch size ({cfg.patch_size}).")
|
|
|
|
llm_hidden_from_cfg = cfg.llm_config.get("hidden_size")
|
|
if llm_hidden_from_cfg is not None and llm_hidden_from_cfg != cfg.llm_hidden_size:
|
|
raise ValueError(f"llm_hidden_size ({cfg.llm_hidden_size}) != llm_config.hidden_size ({llm_hidden_from_cfg}).")
|
|
|
|
agg_h = cfg.aggregator_config.get("hidden_size")
|
|
dit_h = cfg.ditar_config.get("hidden_size")
|
|
if agg_h is not None and dit_h is not None and agg_h != dit_h:
|
|
raise ValueError(f"aggregator_config.hidden_size ({agg_h}) != ditar_config.hidden_size ({dit_h}).")
|
|
if agg_h is not None and agg_h != AGGREGATOR_HIDDEN_SIZE:
|
|
raise ValueError(f"aggregator hidden_size mismatch: got {agg_h}, expected {AGGREGATOR_HIDDEN_SIZE}.")
|
|
if dit_h is not None and dit_h != AGGREGATOR_HIDDEN_SIZE:
|
|
raise ValueError(f"ditar hidden_size mismatch: got {dit_h}, expected {AGGREGATOR_HIDDEN_SIZE}.")
|
|
|
|
atc = cfg.audio_tokenizer_config
|
|
semantic_module_kwargs = _nested_get(atc, "semantic_module_kwargs", default=None)
|
|
if semantic_module_kwargs is not None:
|
|
raise ValueError("Ming dense 0.5B expects audio_tokenizer_config.semantic_module_kwargs to be null.")
|
|
|
|
enc_latent = _nested_get(atc, "enc_kwargs", "latent_dim", default=None)
|
|
dec_latent = _nested_get(atc, "dec_kwargs", "latent_dim", default=None)
|
|
if enc_latent is not None and enc_latent != cfg.latent_dim:
|
|
raise ValueError(f"audio enc latent_dim ({enc_latent}) != Ming latent_dim ({cfg.latent_dim}).")
|
|
if dec_latent is not None and dec_latent != cfg.latent_dim:
|
|
raise ValueError(f"audio dec latent_dim ({dec_latent}) != Ming latent_dim ({cfg.latent_dim}).")
|
|
|
|
atc_patch = _nested_get(atc, "patch_size", default=None)
|
|
if atc_patch is not None and atc_patch != cfg.vae_patch_size:
|
|
raise ValueError(f"audio_tokenizer_config.patch_size ({atc_patch}) != vae_patch_size ({cfg.vae_patch_size}).")
|
|
|
|
atc_sr = _nested_get(atc, "sample_rate", default=None)
|
|
if atc_sr is not None and atc_sr != cfg.sample_rate:
|
|
raise ValueError(f"audio_tokenizer_config.sample_rate ({atc_sr}) != sample_rate ({cfg.sample_rate}).")
|
|
|
|
enc_input_dim = _nested_get(atc, "enc_kwargs", "input_dim", default=None)
|
|
enc_hop_size = _nested_get(atc, "enc_kwargs", "hop_size", default=None)
|
|
dec_output_dim = _nested_get(atc, "dec_kwargs", "output_dim", default=None)
|
|
|
|
if enc_input_dim is not None and enc_hop_size is not None and enc_input_dim != enc_hop_size:
|
|
raise ValueError(f"AudioVAE encoder input_dim ({enc_input_dim}) != hop_size ({enc_hop_size}).")
|
|
if enc_hop_size is not None and dec_output_dim is not None and enc_hop_size != dec_output_dim:
|
|
raise ValueError(
|
|
f"AudioVAE encoder hop_size ({enc_hop_size}) != decoder output_dim ({dec_output_dim}). "
|
|
"Expected 882 in this checkpoint family."
|
|
)
|
|
|
|
if cfg.latent_chunk_size <= 0:
|
|
raise ValueError(f"latent_chunk_size must be > 0, got {cfg.latent_chunk_size}.")
|
|
initial_latent_chunk_size = getattr(cfg, "initial_latent_chunk_size", 0)
|
|
if initial_latent_chunk_size < 0:
|
|
raise ValueError(f"initial_latent_chunk_size must be >= 0, got {initial_latent_chunk_size}.")
|
|
if cfg.latent_left_context < 0:
|
|
raise ValueError(f"latent_left_context must be >= 0, got {cfg.latent_left_context}.")
|
|
if cfg.max_decode_steps <= 0:
|
|
raise ValueError(f"max_decode_steps must be > 0, got {cfg.max_decode_steps}.")
|
|
if not (0.0 <= cfg.stop_head_threshold <= 1.0):
|
|
raise ValueError(f"stop_head_threshold must be in [0,1], got {cfg.stop_head_threshold}.")
|
|
if cfg.stop_head_min_steps < 0:
|
|
raise ValueError(f"stop_head_min_steps must be >= 0, got {cfg.stop_head_min_steps}.")
|