Files
vllm-omni/vllm_omni/model_executor/models/registry.py
T
8939eb835d [Rebase] Rebase to vllm 0.12.0 (#335)
Signed-off-by: Rein Yang <ruiruyang2@gmail.com>
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>
Co-authored-by: Rein Yang <ruiruyang2@gmail.com>
Co-authored-by: Isotr0py <mozf@mail2.sysu.edu.cn>
2025-12-19 16:42:02 +08:00

77 lines
2.2 KiB
Python

from vllm.model_executor.models.registry import _VLLM_MODELS, _LazyRegisteredModel, _ModelRegistry
_OMNI_MODELS = {
"Qwen2_5OmniForConditionalGeneration": (
"qwen2_5_omni",
"qwen2_5_omni",
"Qwen2_5OmniForConditionalGeneration",
),
"Qwen2_5OmniThinkerModel": (
"qwen2_5_omni",
"qwen2_5_omni_thinker",
"Qwen2_5OmniThinkerForConditionalGeneration",
),
"Qwen2_5OmniTalkerModel": (
"qwen2_5_omni",
"qwen2_5_omni_talker",
"Qwen2_5OmniTalkerForConditionalGeneration",
),
"Qwen2_5OmniToken2WavModel": (
"qwen2_5_omni",
"qwen2_5_omni_token2wav",
"Qwen2_5OmniToken2WavForConditionalGenerationVLLM",
),
"Qwen2_5OmniToken2WavDiTModel": (
"qwen2_5_omni",
"qwen2_5_omni_token2wav",
"Qwen2_5OmniToken2WavModel",
),
"Qwen2ForCausalLM_old": ("qwen2_5_omni", "qwen2_old", "Qwen2ForCausalLM"), # need to discuss
# Qwen3 Omni MoE models
"Qwen3OmniMoeForConditionalGeneration": (
"qwen3_omni",
"qwen3_omni",
"Qwen3OmniMoeForConditionalGeneration",
),
"Qwen3OmniMoeThinkerForConditionalGeneration": (
"qwen3_omni",
"qwen3_omni_moe_thinker",
"Qwen3OmniMoeThinkerForConditionalGeneration",
),
"Qwen3OmniMoeTalkerForConditionalGeneration": (
"qwen3_omni",
"qwen3_omni_moe_talker",
"Qwen3OmniMoeTalkerForConditionalGeneration",
),
"Qwen3OmniMoeCode2Wav": (
"qwen3_omni",
"qwen3_omni_code2wav",
"Qwen3OmniMoeCode2Wav",
),
}
_VLLM_OMNI_MODELS = {
**_VLLM_MODELS,
**_OMNI_MODELS,
}
OmniModelRegistry = _ModelRegistry(
{
**{
model_arch: _LazyRegisteredModel(
module_name=f"vllm.model_executor.models.{mod_relname}",
class_name=cls_name,
)
for model_arch, (mod_relname, cls_name) in _VLLM_MODELS.items()
},
**{
model_arch: _LazyRegisteredModel(
module_name=f"vllm_omni.model_executor.models.{mod_folder}.{mod_relname}",
class_name=cls_name,
)
for model_arch, (mod_folder, mod_relname, cls_name) in _OMNI_MODELS.items()
},
}
)