845 lines
24 KiB
Python
845 lines
24 KiB
Python
# Copyright (c) ModelScope Contributors. All rights reserved.
|
|
import transformers
|
|
from dataclasses import dataclass, field
|
|
from packaging import version
|
|
from typing import List, Optional, Union
|
|
|
|
transformers_ge_4_52 = version.parse(transformers.__version__) >= version.parse('4.52')
|
|
|
|
|
|
class LLMModelArch:
|
|
qwen = 'qwen'
|
|
llama = 'llama'
|
|
internlm2 = 'internlm2'
|
|
chatglm = 'chatglm'
|
|
deepseek_v2 = 'deepseek_v2'
|
|
baichuan = 'baichuan'
|
|
|
|
yuan = 'yuan'
|
|
codefuse = 'codefuse'
|
|
phi2 = 'phi2'
|
|
phi3 = 'phi3'
|
|
phi3_small = 'phi3_small'
|
|
telechat = 'telechat'
|
|
dbrx = 'dbrx'
|
|
|
|
|
|
class MLLMModelArch:
|
|
qwen_vl = 'qwen_vl'
|
|
qwen_audio = 'qwen_audio'
|
|
qwen2_vl = 'qwen2_vl'
|
|
qwen2_audio = 'qwen2_audio'
|
|
qwen2_5_omni = 'qwen2_5_omni'
|
|
qwen3_vl = 'qwen3_vl'
|
|
qwen3_omni = 'qwen3_omni'
|
|
qwen3_asr = 'qwen3_asr'
|
|
qwen3_tts = 'qwen3_tts'
|
|
|
|
cogvlm = 'cogvlm'
|
|
chatglm4v = 'chatglm4v'
|
|
glm4v = 'glm4v'
|
|
glm_edge_v = 'glm_edge_v'
|
|
|
|
llama3_1_omni = 'llama3_1_omni'
|
|
llama3_2_vision = 'llama3_2_vision'
|
|
llama4 = 'llama4'
|
|
|
|
llava_hf = 'llava_hf'
|
|
llava_hf_legacy = 'llava_hf_legacy' # transformers<4.52
|
|
llava_next_video_hf = 'llava_next_video_hf'
|
|
llava_onevision1_5 = 'llava_onevision1_5'
|
|
|
|
llava_llama = 'llava_llama'
|
|
llava_mistral = 'llava_mistral'
|
|
|
|
xcomposer = 'xcomposer'
|
|
internvl = 'internvl'
|
|
interns1 = 'interns1'
|
|
minicpmv = 'minicpmv'
|
|
minicpmo = 'minicpmo'
|
|
minicpmv4_6 = 'minicpmv4_6'
|
|
deepseek_vl = 'deepseek_vl'
|
|
deepseek_vl2 = 'deepseek_vl2'
|
|
deepseek_janus = 'deepseek_janus'
|
|
deepseek_ocr = 'deepseek_ocr'
|
|
deepseek_ocr2 = 'deepseek_ocr2'
|
|
unlimited_ocr = 'unlimited-ocr'
|
|
kimi_k25 = 'kimi_k25'
|
|
|
|
mplug_owl2 = 'mplug_owl2'
|
|
mplug_owl2_1 = 'mplug_owl2_1'
|
|
mplug_owl3 = 'mplug_owl3'
|
|
doc_owl2 = 'doc_owl2'
|
|
|
|
phi3_vision = 'phi3_vision'
|
|
phi4_multimodal = 'phi4_multimodal'
|
|
florence = 'florence'
|
|
idefics3 = 'idefics3'
|
|
|
|
got_ocr2 = 'got_ocr2'
|
|
dots_ocr = 'dots_ocr'
|
|
ernie_vl = 'ernie_vl'
|
|
|
|
ovis = 'ovis'
|
|
ovis2_5 = 'ovis2_5'
|
|
molmo = 'molmo'
|
|
emu3_chat = 'emu3_chat'
|
|
megrez_omni = 'megrez_omni'
|
|
valley = 'valley'
|
|
gemma3n = 'gemma3n'
|
|
gemma4_unified = 'gemma4_unified'
|
|
diffusion_gemma = 'diffusion_gemma'
|
|
keye_vl = 'keye_vl'
|
|
|
|
midashenglm = 'midashenglm'
|
|
step_audio2_mini = 'step_audio2_mini'
|
|
hunyuan_vl = 'hunyuan_vl'
|
|
step3_vl = 'step3_vl'
|
|
paddleocr_vl = 'paddleocr_vl'
|
|
minimax_m3_vl = 'minimax_m3_vl'
|
|
|
|
|
|
class ModelArch(LLMModelArch, MLLMModelArch):
|
|
# Multimodal models typically require specifying model_arch,
|
|
# while text-only models usually do not need to specify model_arch.
|
|
pass
|
|
|
|
|
|
@dataclass
|
|
class ModelKeys:
|
|
"""Used to support training of tuners such as llama-pro"""
|
|
arch_name: str = None
|
|
|
|
embedding: str = None
|
|
module_list: str = None
|
|
lm_head: str = None
|
|
|
|
q_proj: str = None
|
|
k_proj: str = None
|
|
v_proj: str = None
|
|
o_proj: str = None
|
|
attention: str = None
|
|
|
|
mlp: str = None
|
|
down_proj: str = None
|
|
|
|
qkv_proj: str = None
|
|
qk_proj: str = None
|
|
qa_proj: str = None
|
|
qb_proj: str = None
|
|
kv_proj: str = None
|
|
kva_proj: str = None
|
|
kvb_proj: str = None
|
|
|
|
|
|
@dataclass
|
|
class MultiModelKeys(ModelKeys):
|
|
"""Used to support freeze_vit/freeze_aligner/freeze_llm"""
|
|
language_model: Union[str, List[str]] = field(default_factory=list)
|
|
aligner: Union[str, List[str]] = field(default_factory=list)
|
|
vision_tower: Union[str, List[str]] = field(default_factory=list)
|
|
generator: Union[str, List[str]] = field(default_factory=list)
|
|
|
|
def __post_init__(self):
|
|
for key in ['language_model', 'aligner', 'vision_tower', 'generator']:
|
|
v = getattr(self, key)
|
|
if isinstance(v, str):
|
|
setattr(self, key, [v])
|
|
if v is None:
|
|
setattr(self, key, [])
|
|
|
|
|
|
MODEL_ARCH_MAPPING = {}
|
|
|
|
|
|
def register_model_arch(model_arch: ModelKeys, *, exist_ok: bool = False) -> None:
|
|
"""
|
|
model_type: The unique ID for the model type. Models with the same model_type share
|
|
the same architectures, template, get_function, etc.
|
|
"""
|
|
arch_name = model_arch.arch_name
|
|
if not exist_ok and arch_name in MODEL_ARCH_MAPPING:
|
|
raise ValueError(f'The `{arch_name}` has already been registered in the MODEL_ARCH_MAPPING.')
|
|
|
|
MODEL_ARCH_MAPPING[arch_name] = model_arch
|
|
|
|
|
|
register_model_arch(
|
|
ModelKeys(
|
|
LLMModelArch.llama,
|
|
module_list='model.layers',
|
|
mlp='model.layers.{}.mlp',
|
|
down_proj='model.layers.{}.mlp.down_proj',
|
|
attention='model.layers.{}.self_attn',
|
|
o_proj='model.layers.{}.self_attn.o_proj',
|
|
q_proj='model.layers.{}.self_attn.q_proj',
|
|
k_proj='model.layers.{}.self_attn.k_proj',
|
|
v_proj='model.layers.{}.self_attn.v_proj',
|
|
embedding='model.embed_tokens',
|
|
lm_head='lm_head',
|
|
))
|
|
|
|
register_model_arch(
|
|
ModelKeys(
|
|
LLMModelArch.internlm2,
|
|
module_list='model.layers',
|
|
mlp='model.layers.{}.feed_forward',
|
|
down_proj='model.layers.{}.feed_forward.w2',
|
|
attention='model.layers.{}.attention',
|
|
o_proj='model.layers.{}.attention.wo',
|
|
qkv_proj='model.layers.{}.attention.wqkv',
|
|
embedding='model.tok_embeddings',
|
|
lm_head='output',
|
|
))
|
|
|
|
register_model_arch(
|
|
ModelKeys(
|
|
LLMModelArch.chatglm,
|
|
module_list='transformer.encoder.layers',
|
|
mlp='transformer.encoder.layers.{}.mlp',
|
|
down_proj='transformer.encoder.layers.{}.mlp.dense_4h_to_h',
|
|
attention='transformer.encoder.layers.{}.self_attention',
|
|
o_proj='transformer.encoder.layers.{}.self_attention.dense',
|
|
qkv_proj='transformer.encoder.layers.{}.self_attention.query_key_value',
|
|
embedding='transformer.embedding',
|
|
lm_head='transformer.output_layer'))
|
|
|
|
register_model_arch(
|
|
ModelKeys(
|
|
LLMModelArch.telechat,
|
|
module_list='transformer.h',
|
|
mlp='transformer.h.{}.mlp',
|
|
down_proj='transformer.h.{}.mlp.down_proj',
|
|
attention='transformer.h.{}.self_attention',
|
|
o_proj='transformer.h.{}.self_attention.dense',
|
|
q_proj='transformer.h.{}.self_attention.query',
|
|
kv_proj='transformer.h.{}.self_attention.key_value',
|
|
embedding='transformer.word_embeddings',
|
|
lm_head='lm_head'))
|
|
|
|
register_model_arch(
|
|
ModelKeys(
|
|
LLMModelArch.baichuan,
|
|
module_list='model.layers',
|
|
mlp='model.layers.{}.mlp',
|
|
down_proj='model.layers.{}.mlp.down_proj',
|
|
attention='model.layers.{}.self_attn',
|
|
qkv_proj='model.layers.{}.self_attn.W_pack',
|
|
embedding='model.embed_tokens',
|
|
lm_head='lm_head',
|
|
))
|
|
|
|
register_model_arch(
|
|
ModelKeys(
|
|
LLMModelArch.yuan,
|
|
module_list='model.layers',
|
|
mlp='model.layers.{}.mlp',
|
|
down_proj='model.layers.{}.mlp.down_proj',
|
|
attention='model.layers.{}.self_attn',
|
|
qk_proj='model.layers.{}.self_attn.qk_proj',
|
|
o_proj='model.layers.{}.self_attn.o_proj',
|
|
q_proj='model.layers.{}.self_attn.q_proj',
|
|
k_proj='model.layers.{}.self_attn.k_proj',
|
|
v_proj='model.layers.{}.self_attn.v_proj',
|
|
embedding='model.embed_tokens',
|
|
lm_head='lm_head',
|
|
))
|
|
|
|
register_model_arch(
|
|
ModelKeys(
|
|
LLMModelArch.codefuse,
|
|
module_list='gpt_neox.layers',
|
|
mlp='gpt_neox.layers.{}.mlp',
|
|
down_proj='gpt_neox.layers.{}.mlp.dense_4h_to_h',
|
|
attention='gpt_neox.layers.{}.attention',
|
|
o_proj='gpt_neox.layers.{}.attention.dense',
|
|
qkv_proj='gpt_neox.layers.{}.attention.query_key_value',
|
|
embedding='gpt_neox.embed_in',
|
|
lm_head='gpt_neox.embed_out',
|
|
))
|
|
|
|
register_model_arch(
|
|
ModelKeys(
|
|
LLMModelArch.phi2,
|
|
module_list='model.layers',
|
|
mlp='model.layers.{}.mlp',
|
|
down_proj='model.layers.{}.mlp.fc2',
|
|
attention='model.layers.{}.self_attn',
|
|
o_proj='model.layers.{}.self_attn.dense',
|
|
q_proj='model.layers.{}.self_attn.q_proj',
|
|
k_proj='model.layers.{}.self_attn.k_proj',
|
|
v_proj='model.layers.{}.self_attn.v_proj',
|
|
embedding='model.embed_tokens',
|
|
lm_head='lm_head',
|
|
))
|
|
|
|
register_model_arch(
|
|
ModelKeys(
|
|
LLMModelArch.qwen,
|
|
module_list='transformer.h',
|
|
mlp='transformer.h.{}.mlp',
|
|
down_proj='transformer.h.{}.mlp.c_proj',
|
|
attention='transformer.h.{}.attn',
|
|
o_proj='transformer.h.{}.attn.c_proj',
|
|
qkv_proj='transformer.h.{}.attn.c_attn',
|
|
embedding='transformer.wte',
|
|
lm_head='lm_head',
|
|
))
|
|
|
|
register_model_arch(
|
|
ModelKeys(
|
|
LLMModelArch.dbrx,
|
|
module_list='transformer.blocks',
|
|
mlp='transformer.blocks.{}.ffn',
|
|
attention='transformer.blocks.{}.norm_attn_norm.attn',
|
|
o_proj='transformer.blocks.{}.norm_attn_norm.attn.out_proj',
|
|
qkv_proj='transformer.blocks.{}.norm_attn_norm.attn.Wqkv',
|
|
embedding='transformer.wte',
|
|
lm_head='lm_head',
|
|
))
|
|
|
|
register_model_arch(
|
|
ModelKeys(
|
|
LLMModelArch.phi3,
|
|
module_list='model.layers',
|
|
mlp='model.layers.{}.mlp',
|
|
down_proj='model.layers.{}.mlp.down_proj',
|
|
attention='model.layers.{}.self_attn',
|
|
o_proj='model.layers.{}.self_attn.o_proj',
|
|
qkv_proj='model.layers.{}.self_attn.qkv_proj',
|
|
embedding='model.embed_tokens',
|
|
lm_head='lm_head',
|
|
))
|
|
|
|
register_model_arch(
|
|
ModelKeys(
|
|
LLMModelArch.phi3_small,
|
|
module_list='model.layers',
|
|
mlp='model.layers.{}.mlp',
|
|
down_proj='model.layers.{}.mlp.down_proj',
|
|
attention='model.layers.{}.self_attn',
|
|
o_proj='model.layers.{}.self_attn.dense',
|
|
qkv_proj='model.layers.{}.self_attn.query_key_value',
|
|
embedding='model.embed_tokens',
|
|
lm_head='lm_head',
|
|
))
|
|
|
|
register_model_arch(
|
|
ModelKeys(
|
|
LLMModelArch.deepseek_v2,
|
|
module_list='model.layers',
|
|
mlp='model.layers.{}.mlp',
|
|
down_proj='model.layers.{}.mlp.down_proj',
|
|
attention='model.layers.{}.self_attn',
|
|
o_proj='model.layers.{}.self_attn.o_proj',
|
|
qa_proj='model.layers.{}.self_attn.q_a_proj',
|
|
qb_proj='model.layers.{}.self_attn.q_b_proj',
|
|
kva_proj='model.layers.{}.self_attn.kv_a_proj_with_mqa',
|
|
kvb_proj='model.layers.{}.self_attn.kv_b_proj',
|
|
embedding='model.embed_tokens',
|
|
lm_head='lm_head',
|
|
))
|
|
|
|
register_model_arch(
|
|
MultiModelKeys(
|
|
MLLMModelArch.llava_hf_legacy,
|
|
language_model='language_model',
|
|
aligner='multi_modal_projector',
|
|
vision_tower='vision_tower',
|
|
))
|
|
|
|
if transformers_ge_4_52:
|
|
register_model_arch(
|
|
MultiModelKeys(
|
|
MLLMModelArch.llava_hf,
|
|
language_model=['model.language_model', 'lm_head'],
|
|
aligner='model.multi_modal_projector',
|
|
vision_tower='model.vision_tower',
|
|
))
|
|
else:
|
|
register_model_arch(
|
|
MultiModelKeys(
|
|
MLLMModelArch.llava_hf,
|
|
language_model='language_model',
|
|
aligner='multi_modal_projector',
|
|
vision_tower='vision_tower',
|
|
))
|
|
|
|
register_model_arch(
|
|
MultiModelKeys(
|
|
MLLMModelArch.llava_mistral,
|
|
language_model='model.layers',
|
|
aligner='model.mm_projector',
|
|
vision_tower='model.vision_tower',
|
|
))
|
|
|
|
if transformers_ge_4_52:
|
|
register_model_arch(
|
|
MultiModelKeys(
|
|
MLLMModelArch.llava_next_video_hf,
|
|
language_model=['model.language_model', 'lm_head'],
|
|
aligner=['model.multi_modal_projector'],
|
|
vision_tower='model.vision_tower'))
|
|
else:
|
|
register_model_arch(
|
|
MultiModelKeys(
|
|
MLLMModelArch.llava_next_video_hf,
|
|
language_model='language_model',
|
|
aligner=['multi_modal_projector'],
|
|
vision_tower='vision_tower'))
|
|
|
|
register_model_arch(
|
|
MultiModelKeys(
|
|
MLLMModelArch.llava_llama,
|
|
language_model='model.layers',
|
|
aligner='model.mm_projector',
|
|
vision_tower='model.vision_tower',
|
|
))
|
|
|
|
register_model_arch(
|
|
MultiModelKeys(
|
|
MLLMModelArch.kimi_k25,
|
|
language_model='language_model',
|
|
aligner='mm_projector',
|
|
vision_tower='vision_tower',
|
|
))
|
|
|
|
register_model_arch(
|
|
MultiModelKeys(
|
|
MLLMModelArch.xcomposer,
|
|
language_model='model',
|
|
aligner='vision_proj',
|
|
vision_tower='vit',
|
|
))
|
|
|
|
register_model_arch(
|
|
MultiModelKeys(
|
|
MLLMModelArch.internvl,
|
|
language_model='language_model',
|
|
aligner='mlp1',
|
|
vision_tower='vision_model',
|
|
))
|
|
|
|
register_model_arch(
|
|
MultiModelKeys(
|
|
MLLMModelArch.interns1,
|
|
language_model=['model.language_model', 'lm_head'],
|
|
aligner='model.multi_modal_projector',
|
|
vision_tower='model.vision_tower',
|
|
))
|
|
|
|
register_model_arch(
|
|
MultiModelKeys(
|
|
MLLMModelArch.mplug_owl3,
|
|
language_model='language_model',
|
|
aligner='vision2text_model',
|
|
vision_tower='vision_model',
|
|
))
|
|
|
|
register_model_arch(
|
|
MultiModelKeys(
|
|
MLLMModelArch.doc_owl2,
|
|
language_model='model.layers',
|
|
aligner=['model.vision2text', 'model.hr_compressor'],
|
|
vision_tower='model.vision_model',
|
|
))
|
|
|
|
register_model_arch(
|
|
MultiModelKeys(
|
|
MLLMModelArch.deepseek_vl,
|
|
language_model='language_model',
|
|
aligner='aligner',
|
|
vision_tower='vision_model',
|
|
))
|
|
|
|
register_model_arch(
|
|
MultiModelKeys(
|
|
MLLMModelArch.deepseek_janus,
|
|
language_model='language_model',
|
|
vision_tower='vision_model',
|
|
aligner='aligner',
|
|
generator=['gen_vision_model', 'gen_aligner', 'gen_head', 'gen_embed']))
|
|
|
|
register_model_arch(
|
|
MultiModelKeys(
|
|
MLLMModelArch.deepseek_ocr,
|
|
language_model='model.layers',
|
|
vision_tower=['model.sam_model', 'model.vision_model'],
|
|
aligner='model.projector',
|
|
))
|
|
|
|
register_model_arch(
|
|
MultiModelKeys(
|
|
MLLMModelArch.deepseek_ocr2,
|
|
language_model='model.layers',
|
|
vision_tower=['model.sam_model', 'model.qwen2_model'],
|
|
aligner='model.projector',
|
|
))
|
|
|
|
register_model_arch(
|
|
MultiModelKeys(
|
|
MLLMModelArch.unlimited_ocr,
|
|
language_model=['model.model.embed_tokens', 'model.model.layers', 'model.model.norm', 'model.lm_head'],
|
|
vision_tower=['model.model.vision_model', 'model.model.sam_model'],
|
|
aligner=['model.model.projector'],
|
|
))
|
|
|
|
register_model_arch(
|
|
MultiModelKeys(
|
|
MLLMModelArch.deepseek_vl2,
|
|
language_model='language',
|
|
vision_tower='vision',
|
|
aligner='projector',
|
|
))
|
|
|
|
register_model_arch(
|
|
MultiModelKeys(
|
|
MLLMModelArch.minicpmv,
|
|
language_model='llm',
|
|
aligner='resampler',
|
|
vision_tower='vpm',
|
|
))
|
|
|
|
register_model_arch(
|
|
MultiModelKeys(
|
|
MLLMModelArch.minicpmv4_6,
|
|
language_model='model.language_model',
|
|
aligner='model.merger',
|
|
vision_tower='model.vision_tower',
|
|
))
|
|
|
|
register_model_arch(
|
|
MultiModelKeys(
|
|
MLLMModelArch.minicpmo,
|
|
language_model='llm',
|
|
aligner='resampler',
|
|
vision_tower=['vpm', 'apm'],
|
|
))
|
|
|
|
register_model_arch(
|
|
MultiModelKeys(
|
|
MLLMModelArch.phi3_vision,
|
|
language_model='model.layers',
|
|
aligner='model.vision_embed_tokens.img_projection',
|
|
vision_tower='model.vision_embed_tokens.img_processor',
|
|
))
|
|
|
|
register_model_arch(
|
|
MultiModelKeys(
|
|
MLLMModelArch.phi4_multimodal,
|
|
language_model='model.layers',
|
|
aligner=[
|
|
'model.embed_tokens_extend.image_embed.img_projection',
|
|
'model.embed_tokens_extend.audio_embed.audio_projection'
|
|
],
|
|
vision_tower=[
|
|
'model.embed_tokens_extend.image_embed.img_processor', 'model.embed_tokens_extend.audio_embed.encoder'
|
|
],
|
|
))
|
|
|
|
register_model_arch(MultiModelKeys(
|
|
MLLMModelArch.cogvlm,
|
|
language_model='model.layers',
|
|
vision_tower='model.vision',
|
|
))
|
|
|
|
register_model_arch(
|
|
MultiModelKeys(
|
|
MLLMModelArch.florence,
|
|
language_model='language_model',
|
|
vision_tower='vision_tower',
|
|
))
|
|
|
|
register_model_arch(
|
|
MultiModelKeys(
|
|
MLLMModelArch.qwen_vl,
|
|
language_model='transformer.h',
|
|
vision_tower='transformer.visual',
|
|
))
|
|
# TODO: check lm_head, ALL
|
|
register_model_arch(
|
|
MultiModelKeys(
|
|
MLLMModelArch.qwen_audio,
|
|
language_model='transformer.h',
|
|
vision_tower='transformer.audio',
|
|
))
|
|
|
|
register_model_arch(
|
|
MultiModelKeys(
|
|
MLLMModelArch.qwen2_audio,
|
|
language_model='language_model',
|
|
aligner='multi_modal_projector',
|
|
vision_tower='audio_tower',
|
|
))
|
|
|
|
if transformers_ge_4_52:
|
|
register_model_arch(
|
|
MultiModelKeys(
|
|
MLLMModelArch.qwen2_vl,
|
|
language_model=['model.language_model', 'lm_head'],
|
|
aligner='model.visual.merger',
|
|
vision_tower='model.visual',
|
|
))
|
|
else:
|
|
register_model_arch(
|
|
MultiModelKeys(
|
|
MLLMModelArch.qwen2_vl,
|
|
language_model=['model', 'lm_head'],
|
|
aligner='visual.merger',
|
|
vision_tower='visual',
|
|
))
|
|
|
|
register_model_arch(
|
|
MultiModelKeys(
|
|
MLLMModelArch.qwen3_vl,
|
|
language_model=['model.language_model', 'lm_head'],
|
|
aligner=['model.visual.merger', 'model.visual.deepstack_merger_list'],
|
|
vision_tower='model.visual',
|
|
))
|
|
|
|
register_model_arch(
|
|
MultiModelKeys(
|
|
MLLMModelArch.qwen2_5_omni,
|
|
language_model=['thinker.model', 'thinker.lm_head'],
|
|
vision_tower=['thinker.audio_tower', 'thinker.visual'],
|
|
aligner=['thinker.audio_tower.proj', 'thinker.visual.merger'],
|
|
generator=['talker', 'token2wav'],
|
|
))
|
|
|
|
register_model_arch(
|
|
MultiModelKeys(
|
|
MLLMModelArch.qwen3_omni,
|
|
language_model=['thinker.model', 'thinker.lm_head'],
|
|
vision_tower=['thinker.audio_tower', 'thinker.visual'],
|
|
aligner=[
|
|
'thinker.audio_tower.proj1', 'thinker.audio_tower.proj2', 'thinker.visual.merger',
|
|
'thinker.visual.merger_list'
|
|
],
|
|
generator=['talker', 'code2wav'],
|
|
))
|
|
|
|
register_model_arch(
|
|
MultiModelKeys(
|
|
MLLMModelArch.qwen3_asr,
|
|
language_model=['thinker.model', 'thinker.lm_head'],
|
|
vision_tower='thinker.audio_tower',
|
|
aligner=['thinker.audio_tower.proj1', 'thinker.audio_tower.proj2'],
|
|
))
|
|
|
|
register_model_arch(
|
|
MultiModelKeys(
|
|
MLLMModelArch.qwen3_tts,
|
|
language_model='talker',
|
|
generator='speaker_encoder', # no grad
|
|
))
|
|
|
|
register_model_arch(
|
|
MultiModelKeys(
|
|
MLLMModelArch.midashenglm,
|
|
language_model='decoder',
|
|
aligner=['audio_projector'],
|
|
vision_tower=['audio_encoder'],
|
|
))
|
|
|
|
register_model_arch(
|
|
MultiModelKeys(
|
|
MLLMModelArch.step_audio2_mini,
|
|
language_model=['model', 'lm_head'],
|
|
aligner=['adapter'],
|
|
vision_tower=['encoder'],
|
|
))
|
|
|
|
register_model_arch(
|
|
MultiModelKeys(
|
|
MLLMModelArch.chatglm4v,
|
|
language_model='transformer.encoder',
|
|
vision_tower='transformer.vision',
|
|
))
|
|
|
|
register_model_arch(
|
|
MultiModelKeys(
|
|
MLLMModelArch.glm4v,
|
|
language_model=['model.language_model', 'lm_head'],
|
|
aligner='model.visual.merger',
|
|
vision_tower='model.visual',
|
|
))
|
|
|
|
register_model_arch(
|
|
MultiModelKeys(
|
|
MLLMModelArch.idefics3,
|
|
language_model='model.text_model',
|
|
aligner='model.connector',
|
|
vision_tower='model.vision_model',
|
|
))
|
|
|
|
register_model_arch(
|
|
MultiModelKeys(
|
|
MLLMModelArch.llama3_1_omni,
|
|
language_model='model.layers',
|
|
aligner='model.speech_projector',
|
|
vision_tower='model.speech_encoder',
|
|
generator='speech_generator',
|
|
))
|
|
|
|
register_model_arch(
|
|
MultiModelKeys(
|
|
MLLMModelArch.got_ocr2,
|
|
language_model='model.layers',
|
|
aligner='model.mm_projector_vary',
|
|
vision_tower='model.vision_tower_high',
|
|
))
|
|
|
|
register_model_arch(
|
|
MultiModelKeys(
|
|
MLLMModelArch.ernie_vl,
|
|
language_model=['model', 'lm_head'],
|
|
aligner='model.resampler_model',
|
|
vision_tower='vision_model',
|
|
))
|
|
|
|
if transformers_ge_4_52:
|
|
register_model_arch(
|
|
MultiModelKeys(
|
|
MLLMModelArch.llama3_2_vision,
|
|
language_model=['model.language_model', 'lm_head'],
|
|
aligner='model.multi_modal_projector',
|
|
vision_tower='model.vision_model',
|
|
))
|
|
else:
|
|
register_model_arch(
|
|
MultiModelKeys(
|
|
MLLMModelArch.llama3_2_vision,
|
|
language_model='language_model',
|
|
aligner='multi_modal_projector',
|
|
vision_tower='vision_model',
|
|
))
|
|
|
|
register_model_arch(
|
|
MultiModelKeys(
|
|
MLLMModelArch.llama4,
|
|
language_model='language_model',
|
|
aligner='multi_modal_projector',
|
|
vision_tower='vision_model',
|
|
))
|
|
|
|
register_model_arch(
|
|
MultiModelKeys(
|
|
MLLMModelArch.ovis,
|
|
language_model='llm',
|
|
vision_tower=['visual_tokenizer', 'vte'],
|
|
))
|
|
|
|
register_model_arch(
|
|
MultiModelKeys(
|
|
MLLMModelArch.ovis2_5,
|
|
language_model='llm',
|
|
aligner='visual_tokenizer.head',
|
|
vision_tower=['visual_tokenizer.vit', 'vte'],
|
|
))
|
|
|
|
register_model_arch(
|
|
MultiModelKeys(
|
|
MLLMModelArch.molmo,
|
|
language_model='model.transformer',
|
|
vision_tower='model.vision_backbone',
|
|
aligner='model.vision_backbone.image_projector'))
|
|
|
|
register_model_arch(
|
|
MultiModelKeys(
|
|
MLLMModelArch.megrez_omni,
|
|
language_model='llm',
|
|
vision_tower=['vision', 'audio'],
|
|
))
|
|
|
|
register_model_arch(MultiModelKeys(MLLMModelArch.emu3_chat, language_model='model'))
|
|
|
|
register_model_arch(
|
|
MultiModelKeys(MLLMModelArch.glm_edge_v, language_model='model.layers', vision_tower='model.vision'))
|
|
|
|
register_model_arch(
|
|
MultiModelKeys(
|
|
MLLMModelArch.valley,
|
|
language_model='model',
|
|
vision_tower=['model.vision_tower', 'model.qwen2vl_vision_tower'],
|
|
))
|
|
|
|
register_model_arch(
|
|
MultiModelKeys(
|
|
MLLMModelArch.gemma3n,
|
|
language_model=['model.language_model', 'lm_head'],
|
|
aligner=['model.embed_vision', 'model.embed_audio'],
|
|
vision_tower=['model.vision_tower', 'model.audio_tower'],
|
|
))
|
|
|
|
register_model_arch(
|
|
MultiModelKeys(
|
|
MLLMModelArch.gemma4_unified,
|
|
language_model=['model.language_model', 'lm_head'],
|
|
aligner=['model.embed_vision', 'model.embed_audio'],
|
|
))
|
|
|
|
register_model_arch(
|
|
MultiModelKeys(
|
|
MLLMModelArch.diffusion_gemma,
|
|
language_model=['model.encoder.language_model', 'model.decoder', 'lm_head'],
|
|
vision_tower=['model.encoder.vision_tower'],
|
|
aligner=['model.encoder.embed_vision'],
|
|
))
|
|
|
|
register_model_arch(
|
|
MultiModelKeys(
|
|
MLLMModelArch.keye_vl,
|
|
language_model=['model', 'lm_head'],
|
|
aligner='mlp_AR',
|
|
vision_tower='visual',
|
|
))
|
|
|
|
register_model_arch(MultiModelKeys(
|
|
MLLMModelArch.dots_ocr,
|
|
language_model='model',
|
|
))
|
|
|
|
register_model_arch(
|
|
MultiModelKeys(
|
|
MLLMModelArch.llava_onevision1_5,
|
|
language_model=['model.language_model', 'lm_head'],
|
|
aligner='model.visual.merger',
|
|
vision_tower='model.visual',
|
|
))
|
|
|
|
register_model_arch(
|
|
MultiModelKeys(
|
|
MLLMModelArch.hunyuan_vl,
|
|
language_model='model',
|
|
aligner='vit.perceive',
|
|
vision_tower='vit',
|
|
))
|
|
|
|
register_model_arch(
|
|
MultiModelKeys(
|
|
MLLMModelArch.step3_vl,
|
|
language_model=['model.language_model', 'lm_head'],
|
|
aligner='model.vit_large_projector',
|
|
vision_tower='model.vision_model',
|
|
))
|
|
|
|
register_model_arch(
|
|
MultiModelKeys(
|
|
MLLMModelArch.paddleocr_vl,
|
|
language_model=['model.language_model', 'lm_head'],
|
|
aligner='model.projector',
|
|
vision_tower='model.visual',
|
|
))
|
|
|
|
register_model_arch(
|
|
MultiModelKeys(
|
|
MLLMModelArch.minimax_m3_vl,
|
|
language_model=['model.language_model', 'lm_head'],
|
|
aligner='model.multi_modal_projector',
|
|
vision_tower='model.vision_tower',
|
|
))
|
|
|
|
|
|
def get_model_arch(arch_name: Optional[str]) -> Optional[MultiModelKeys]:
|
|
return MODEL_ARCH_MAPPING.get(arch_name)
|