# Copyright (c) ModelScope Contributors. All rights reserved. import torch from transformers import PretrainedConfig, PreTrainedModel from transformers.dynamic_module_utils import get_class_from_dynamic_module from types import MethodType from swift.template import TemplateType from swift.utils import Processor, get_logger from ..constant import MLLMModelType from ..model_arch import ModelArch from ..model_meta import Model, ModelGroup, ModelMeta from ..patcher import patch_output_clone from ..register import ModelLoader, register_model from ..utils import use_submodel_func from .qwen import Qwen2VLLoader, patch_qwen_vl_utils logger = get_logger() class Idefics3Loader(ModelLoader): def get_model(self, model_dir: str, *args, **kwargs) -> PreTrainedModel: from transformers import AutoModelForVision2Seq self.auto_model_cls = self.auto_model_cls or AutoModelForVision2Seq return super().get_model(model_dir, *args, **kwargs) register_model( ModelMeta( MLLMModelType.idefics3, [ ModelGroup([ Model('AI-ModelScope/Idefics3-8B-Llama3', 'HuggingFaceM4/Idefics3-8B-Llama3'), ]), ], Idefics3Loader, template=TemplateType.idefics3, model_arch=ModelArch.idefics3, architectures=['Idefics3ForConditionalGeneration'], tags=['vision'], requires=['transformers>=4.45'], )) class PixtralLoader(ModelLoader): def get_model(self, model_dir: str, *args, **kwargs) -> PreTrainedModel: from transformers import LlavaForConditionalGeneration self.auto_model_cls = self.auto_model_cls or LlavaForConditionalGeneration return super().get_model(model_dir, *args, **kwargs) register_model( ModelMeta( MLLMModelType.pixtral, [ ModelGroup([ Model('AI-ModelScope/pixtral-12b', 'mistral-community/pixtral-12b'), ]), ], PixtralLoader, template=TemplateType.pixtral, model_arch=ModelArch.llava_hf, architectures=['LlavaForConditionalGeneration'], requires=['transformers>=4.45'], tags=['vision'], )) class MolMoeLoader(ModelLoader): def get_model(self, model_dir: str, *args, **kwargs) -> PreTrainedModel: model = super().get_model(model_dir, *args, **kwargs) # fix bug for molmoe-1b def to_dict(self, *args, **kwargs): res = self._to_dict(*args, **kwargs) res['vision_backbone'] = self.vision_backbone.__dict__ res.pop('to_dict') res.pop('_to_dict') return res model.config._to_dict = model.config.to_dict model.config.to_dict = MethodType(to_dict, model.config) patch_output_clone(model.model.transformer.wte) return model register_model( ModelMeta( MLLMModelType.molmoe, [ ModelGroup([ Model('LLM-Research/MolmoE-1B-0924', 'allenai/MolmoE-1B-0924'), ]), ], MolMoeLoader, template=TemplateType.molmo, model_arch=ModelArch.molmo, torch_dtype=torch.float32, architectures=['OLMoForCausalLM'], tags=['vision'], requires=['transformers>=4.45'], )) class MolmoLoader(ModelLoader): def get_model(self, model_dir: str, *args, **kwargs) -> PreTrainedModel: model_cls = get_class_from_dynamic_module('modeling_molmo.MolmoForCausalLM', model_dir) model_cls._no_split_modules = ['MolmoSequentialBlock'] model = super().get_model(model_dir, *args, **kwargs) patch_output_clone(model.model.transformer.wte) return model register_model( ModelMeta( MLLMModelType.molmo, [ ModelGroup([ Model('LLM-Research/Molmo-7B-O-0924', 'allenai/Molmo-7B-O-0924'), Model('LLM-Research/Molmo-7B-D-0924', 'allenai/Molmo-7B-D-0924'), Model('LLM-Research/Molmo-72B-0924', 'allenai/Molmo-72B-0924'), ]), ], MolmoLoader, template=TemplateType.molmo, model_arch=ModelArch.molmo, architectures=['MolmoForCausalLM'], tags=['vision'], requires=['transformers>=4.45'], )) class Molmo2Loader(ModelLoader): def get_model(self, model_dir: str, *args, **kwargs) -> PreTrainedModel: from transformers import AutoModelForImageTextToText model_cls = get_class_from_dynamic_module('modeling_molmo2.Molmo2ForConditionalGeneration', model_dir) no_split_modules = getattr(model_cls, '_no_split_modules', []) or [] if 'MolmoSequentialBlock' not in no_split_modules: model_cls._no_split_modules = no_split_modules + ['MolmoSequentialBlock'] self.auto_model_cls = self.auto_model_cls or AutoModelForImageTextToText model = super().get_model(model_dir, *args, **kwargs) patch_output_clone(model.model.transformer.wte) return model register_model( ModelMeta( MLLMModelType.molmo2, [ ModelGroup([ Model('allenai/Molmo2-4B', 'allenai/Molmo2-4B'), Model('allenai/Molmo2-8B', 'allenai/Molmo2-8B'), Model('allenai/Molmo2-O-7B', 'allenai/Molmo2-O-7B'), ]), ], Molmo2Loader, template=TemplateType.molmo2, model_arch=ModelArch.molmo, architectures=['Molmo2ForConditionalGeneration'], tags=['vision', 'video'], requires=['transformers>=4.57.1,<5', 'decord'], )) class MegrezOmniLoader(ModelLoader): def get_model(self, model_dir: str, *args, **kwargs) -> PreTrainedModel: model_cls = get_class_from_dynamic_module('modeling_megrezo.MegrezO', model_dir) model_cls._no_split_modules = ['ResidualAttentionBlock', 'LlamaDecoderLayer'] model_cls = get_class_from_dynamic_module('modeling_megrezo.SiglipVisionTransformer', model_dir) model_cls._no_split_modules = ['SiglipEncoderLayer'] model = super().get_model(model_dir, *args, **kwargs) patch_output_clone(model.llm.model.embed_tokens) use_submodel_func(model, 'llm') return model def _get_model_processor(self, model_dir, config): model, processor = super().get_processor(model_dir, config) if model: processor = model._get_or_init_processor() return model, processor register_model( ModelMeta( MLLMModelType.megrez_omni, [ ModelGroup([ Model('InfiniAI/Megrez-3B-Omni', 'Infinigence/Megrez-3B-Omni'), ]), ], MegrezOmniLoader, template=TemplateType.megrez_omni, model_arch=ModelArch.megrez_omni, architectures=['MegrezO'], tags=['vision', 'audio'], )) register_model( ModelMeta( MLLMModelType.qwen2_gme, [ ModelGroup([ Model('iic/gme-Qwen2-VL-2B-Instruct', 'Alibaba-NLP/gme-Qwen2-VL-2B-Instruct'), Model('iic/gme-Qwen2-VL-7B-Instruct', 'Alibaba-NLP/gme-Qwen2-VL-7B-Instruct'), ]), ], Qwen2VLLoader, template=TemplateType.qwen2_gme, model_arch=ModelArch.qwen2_vl, architectures=['Qwen2VLForConditionalGeneration'], tags=['vision'])) class JinaRerankerM0Loader(ModelLoader): def get_model(self, model_dir: str, *args, **kwargs) -> PreTrainedModel: # Use AutoModel to respect the model repo's dynamic class mapping # and load the custom Jina reranker head via trust_remote_code. from transformers import AutoModel from transformers.modeling_outputs import SequenceClassifierOutputWithPast self.auto_model_cls = self.auto_model_cls or AutoModel model = super().get_model(model_dir, *args, **kwargs) # Patch forward to return a sequence-classification-style output with `.logits` # Use the model's own head (already present in jina-reranker-m0), just wrap outputs. if not hasattr(model, '_forward_origin'): model._forward_origin = model.forward model.logit_bias = 2.65 def forward(self, input_ids=None, attention_mask=None, position_ids=None, inputs_embeds=None, pixel_values=None, image_grid_thw=None, video_grid_thw=None, output_attentions=None, output_hidden_states=None, return_dict=None, **kwargs): # Remove labels to avoid upstream asserts in ranking models kwargs.pop('labels', None) if return_dict is None: return_dict = True out = self._forward_origin( input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, inputs_embeds=inputs_embeds, pixel_values=pixel_values, image_grid_thw=image_grid_thw, video_grid_thw=video_grid_thw, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, **kwargs) logits = out.unsqueeze(-1) - self.logit_bias if not return_dict: return (logits, ) return SequenceClassifierOutputWithPast(logits=logits) model.forward = MethodType(forward, model) def padding_free_fn(self, output, kwargs, padding_side): return_dict = kwargs.get('return_dict', None) output.logits = output['last_hidden_state'][:, -1] logits = self.score(output.logits) logits = logits - self.logit_bias if not return_dict: return (logits, ) return SequenceClassifierOutputWithPast(logits=logits) model.padding_free_fn = MethodType(padding_free_fn, model) return model register_model( ModelMeta( MLLMModelType.jina_reranker_m0, [ModelGroup([Model('JinaAI/jina-reranker-m0', 'JinaAI/jina-reranker-m0')])], JinaRerankerM0Loader, template=TemplateType.jina_reranker_m0, model_arch=ModelArch.qwen2_vl, architectures=['JinaRerankerM0ForConditionalGeneration'], task_type='reranker', tags=['reranker', 'vision'], )) class KeyeVLLoader(ModelLoader): def get_processor(self, model_dir: str, config: PretrainedConfig) -> Processor: processor = super().get_processor(model_dir, config) from keye_vl_utils import vision_process global_vars = patch_qwen_vl_utils(vision_process) processor.global_vars = global_vars return processor register_model( ModelMeta( MLLMModelType.keye_vl, [ ModelGroup([ Model('Kwai-Keye/Keye-VL-8B-Preview', 'Kwai-Keye/Keye-VL-8B-Preview'), ]), ], KeyeVLLoader, template=TemplateType.keye_vl, model_arch=ModelArch.keye_vl, architectures=['KeyeForConditionalGeneration'], tags=['vision'], requires=['keye_vl_utils'], )) register_model( ModelMeta( MLLMModelType.keye_vl_1_5, [ ModelGroup([ Model('Kwai-Keye/Keye-VL-1_5-8B', 'Kwai-Keye/Keye-VL-1_5-8B'), ]), ], KeyeVLLoader, template=TemplateType.keye_vl_1_5, model_arch=ModelArch.keye_vl, architectures=['KeyeVL1_5ForConditionalGeneration'], tags=['vision'], requires=['keye_vl_utils>=1.5.2', 'transformers==4.52.4'], )) class DotsOCRLoader(ModelLoader): def get_model(self, model_dir: str, *args, **kwargs) -> PreTrainedModel: model_cls = get_class_from_dynamic_module('modeling_dots_vision.DotsVisionTransformer', model_dir) model_cls._no_split_modules = ['DotsVisionBlock'] return super().get_model(model_dir, *args, **kwargs) register_model( ModelMeta( MLLMModelType.dots_ocr, [ModelGroup([ Model('rednote-hilab/dots.ocr', 'rednote-hilab/dots.ocr'), ])], DotsOCRLoader, template=TemplateType.dots_ocr, model_arch=ModelArch.dots_ocr, architectures=['DotsOCRForCausalLM'], requires=['transformers>=4.51.0'], )) class Sail2VLLoader(ModelLoader): def get_model(self, model_dir: str, *args, **kwargs) -> PreTrainedModel: model = super().get_model(model_dir, *args, **kwargs) use_submodel_func(model, 'language_model') return model register_model( ModelMeta( MLLMModelType.sail_vl2, [ ModelGroup([ Model('BytedanceDouyinContent/SAIL-VL2-2B', 'BytedanceDouyinContent/SAIL-VL2-2B'), Model('BytedanceDouyinContent/SAIL-VL2-2B-Thinking', 'BytedanceDouyinContent/SAIL-VL2-2B-Thinking'), Model('BytedanceDouyinContent/SAIL-VL2-8B', 'BytedanceDouyinContent/SAIL-VL2-8B'), Model('BytedanceDouyinContent/SAIL-VL2-8B-Thinking', 'BytedanceDouyinContent/SAIL-VL2-8B-Thinking'), ]) ], Sail2VLLoader, template=TemplateType.sail_vl2, model_arch=ModelArch.internvl, architectures=['SAILVLModel'], requires=['transformers<=4.51.3'], tags=['vision']))