# Copyright (c) ModelScope Contributors. All rights reserved. import inspect import torch import torch.distributed as dist import transformers from packaging import version from PIL import Image from transformers import PreTrainedModel from types import MethodType from swift.template import TemplateType from swift.utils import is_deepspeed_enabled, to_device from ..constant import LLMModelType, MLLMModelType from ..model_arch import ModelArch from ..model_meta import Model, ModelGroup, ModelMeta from ..patcher import patch_output_to_input_device from ..register import ModelLoader, SentenceTransformersLoader, register_model transformers_5_9 = version.parse(transformers.__version__) >= version.parse('5.9') class PaligemmaVisionLoader(ModelLoader): def get_model(self, model_dir: str, *args, **kwargs) -> PreTrainedModel: from transformers import PaliGemmaForConditionalGeneration self.auto_model_cls = self.auto_model_cls or PaliGemmaForConditionalGeneration return super().get_model(model_dir, *args, **kwargs) register_model( ModelMeta( MLLMModelType.paligemma, [ ModelGroup([ Model('AI-ModelScope/paligemma-3b-pt-224', 'google/paligemma-3b-pt-224'), Model('AI-ModelScope/paligemma-3b-pt-448', 'google/paligemma-3b-pt-448'), Model('AI-ModelScope/paligemma-3b-pt-896', 'google/paligemma-3b-pt-896'), ]), ModelGroup([ Model('AI-ModelScope/paligemma-3b-mix-224', 'google/paligemma-3b-mix-224'), Model('AI-ModelScope/paligemma-3b-mix-448', 'google/paligemma-3b-mix-448'), ]), ModelGroup([ Model('AI-ModelScope/paligemma2-3b-pt-224', 'google/paligemma2-3b-pt-224'), Model('AI-ModelScope/paligemma2-3b-pt-448', 'google/paligemma2-3b-pt-448'), Model('AI-ModelScope/paligemma2-3b-pt-896', 'google/paligemma2-3b-pt-896'), Model('AI-ModelScope/paligemma2-10b-pt-224', 'google/paligemma2-10b-pt-224'), Model('AI-ModelScope/paligemma2-10b-pt-448', 'google/paligemma2-10b-pt-448'), Model('AI-ModelScope/paligemma2-10b-pt-896', 'google/paligemma2-10b-pt-896'), Model('AI-ModelScope/paligemma2-28b-pt-224', 'google/paligemma2-28b-pt-224'), Model('AI-ModelScope/paligemma2-28b-pt-448', 'google/paligemma2-28b-pt-448'), Model('AI-ModelScope/paligemma2-28b-pt-896', 'google/paligemma2-28b-pt-896'), ]), ModelGroup([ Model('AI-ModelScope/paligemma2-3b-ft-docci-448', 'google/paligemma2-3b-ft-docci-448'), Model('AI-ModelScope/paligemma2-10b-ft-docci-448', 'google/paligemma2-10b-ft-docci-448'), ]), ], PaligemmaVisionLoader, template=TemplateType.paligemma, architectures=['PaliGemmaForConditionalGeneration'], model_arch=ModelArch.llava_hf, requires=['transformers>=4.41'], tags=['vision'], )) register_model( ModelMeta( LLMModelType.gemma, [ ModelGroup([ Model('AI-ModelScope/gemma-2b-it', 'google/gemma-2b-it'), Model('AI-ModelScope/gemma-2b', 'google/gemma-2b'), Model('AI-ModelScope/gemma-7b', 'google/gemma-7b'), Model('AI-ModelScope/gemma-7b-it', 'google/gemma-7b-it'), ], ), ], template=TemplateType.gemma, architectures=['GemmaForCausalLM'], model_arch=ModelArch.llama, requires=['transformers>=4.38'], )) register_model( ModelMeta( LLMModelType.gemma2, [ ModelGroup([ Model('LLM-Research/gemma-2-2b-it', 'google/gemma-2-2b-it'), Model('LLM-Research/gemma-2-2b', 'google/gemma-2-2b'), Model('LLM-Research/gemma-2-9b', 'google/gemma-2-9b'), Model('LLM-Research/gemma-2-9b-it', 'google/gemma-2-9b-it'), Model('LLM-Research/gemma-2-27b', 'google/gemma-2-27b'), Model('LLM-Research/gemma-2-27b-it', 'google/gemma-2-27b-it'), ], ), ], template=TemplateType.gemma, architectures=['Gemma2ForCausalLM'], model_arch=ModelArch.llama, requires=['transformers>=4.42'], )) class Gemma3TextLoader(ModelLoader): def get_config(self, model_dir): # It is strongly recommended to train Gemma3 models with the `eager` attention implementation instead of `sdpa`. self.attn_impl = self.attn_impl or 'eager' return super().get_config(model_dir) register_model( ModelMeta( LLMModelType.gemma3_text, [ ModelGroup([ Model('LLM-Research/gemma-3-1b-pt', 'google/gemma-3-1b-pt'), Model('LLM-Research/gemma-3-1b-it', 'google/gemma-3-1b-it'), Model('google/gemma-3-270m', 'google/gemma-3-270m'), Model('google/gemma-3-270m-it', 'google/gemma-3-270m-it'), Model('google/medgemma-27b-text-it', 'google/medgemma-27b-text-it'), ], ), ], Gemma3TextLoader, template=TemplateType.gemma3_text, architectures=['Gemma3ForCausalLM'], model_arch=ModelArch.llama, requires=['transformers>=4.49'], )) class Gemma3VisionLoader(ModelLoader): def get_config(self, model_dir): # It is strongly recommended to train Gemma3 models with the `eager` attention implementation instead of `sdpa`. self.attn_impl = self.attn_impl or 'eager' return super().get_config(model_dir) def get_model(self, model_dir: str, *args, **kwargs) -> PreTrainedModel: from transformers import Gemma3ForConditionalGeneration self.auto_model_cls = self.auto_model_cls or Gemma3ForConditionalGeneration return super().get_model(model_dir, *args, **kwargs) register_model( ModelMeta( MLLMModelType.gemma3_vision, [ ModelGroup([ Model('LLM-Research/gemma-3-4b-pt', 'google/gemma-3-4b-pt'), Model('LLM-Research/gemma-3-4b-it', 'google/gemma-3-4b-it'), Model('LLM-Research/gemma-3-12b-pt', 'google/gemma-3-12b-pt'), Model('LLM-Research/gemma-3-12b-it', 'google/gemma-3-12b-it'), Model('LLM-Research/gemma-3-27b-pt', 'google/gemma-3-27b-pt'), Model('LLM-Research/gemma-3-27b-it', 'google/gemma-3-27b-it'), Model('google/medgemma-4b-pt', 'google/medgemma-4b-pt'), Model('google/medgemma-4b-it', 'google/medgemma-4b-it'), Model('google/medgemma-27b-it', 'google/medgemma-27b-it'), ], ), ], Gemma3VisionLoader, template=TemplateType.gemma3_vision, architectures=['Gemma3ForConditionalGeneration'], model_arch=ModelArch.llava_hf, requires=['transformers>=4.49'], )) class Gemma3nLoader(ModelLoader): def get_model(self, model_dir: str, *args, **kwargs) -> PreTrainedModel: from transformers import Gemma3nForConditionalGeneration self.auto_model_cls = self.auto_model_cls or Gemma3nForConditionalGeneration model = super().get_model(model_dir, *args, **kwargs) patch_output_to_input_device(model.model.embed_vision) patch_output_to_input_device(model.model.embed_audio) return model register_model( ModelMeta( MLLMModelType.gemma3n, [ ModelGroup([ Model('google/gemma-3n-E2B', 'google/gemma-3n-E2B'), Model('google/gemma-3n-E4B', 'google/gemma-3n-E4B'), Model('google/gemma-3n-E2B-it', 'google/gemma-3n-E2B-it'), Model('google/gemma-3n-E4B-it', 'google/gemma-3n-E4B-it'), ], ), ], Gemma3nLoader, template=TemplateType.gemma3n, architectures=['Gemma3nForConditionalGeneration'], model_arch=ModelArch.gemma3n, requires=['transformers>=4.53.1'], )) register_model( ModelMeta( LLMModelType.gemma_emb, [ ModelGroup([ Model('google/embeddinggemma-300m', 'google/embeddinggemma-300m'), ], ), ], SentenceTransformersLoader, template=TemplateType.dummy, architectures=['Gemma3TextModel'], )) def _patch_gemma4_forward(model, processor, is_gemma4_unified: bool = False): if is_gemma4_unified: from transformers.models.gemma4_unified.modeling_gemma4_unified import \ Gemma4UnifiedModelOutputWithPast as Gemma4ModelOutputWithPast from transformers.models.gemma4_unified.modeling_gemma4_unified import (create_masks_for_generate, torch_compilable_check) else: from transformers.models.gemma4.modeling_gemma4 import (Gemma4ModelOutputWithPast, create_masks_for_generate, torch_compilable_check) if hasattr(model, 'origin_forward'): return def _forward_dummy_image(model, inputs_embeds): images = [Image.new('RGB', (32, 32), (0, 0, 0))] image_inputs = processor.image_processor(images=images, return_tensors='pt') image_inputs = to_device(image_inputs, inputs_embeds.device) dummy_pixel = image_inputs['pixel_values'].to(model.dtype) dummy_pos_ids = image_inputs.get('image_position_ids') image_features = model.get_image_features(dummy_pixel, dummy_pos_ids, return_dict=True).pooler_output inputs_embeds = inputs_embeds + image_features.mean() * 0. return inputs_embeds # transformers 5.6.2 def forward( self, input_ids: torch.LongTensor | None = None, pixel_values: torch.FloatTensor | None = None, pixel_values_videos: torch.FloatTensor | None = None, input_features: torch.FloatTensor | None = None, attention_mask: torch.Tensor | None = None, input_features_mask: torch.Tensor | None = None, position_ids: torch.LongTensor | None = None, past_key_values=None, mm_token_type_ids: torch.LongTensor | None = None, inputs_embeds: torch.FloatTensor | None = None, use_cache: bool | None = None, image_position_ids: torch.LongTensor | None = None, video_position_ids: torch.LongTensor | None = None, per_layer_inputs: torch.Tensor | None = None, **kwargs, ) -> Gemma4ModelOutputWithPast: r""" input_features_mask (`torch.FloatTensor]` of shape `(num_images, seq_length)`): The attention mask for the input audio. image_position_ids (`torch.LongTensor` of shape `(batch_size, max_patches, 2)`, *optional*): 2D patch position coordinates from the image processor, with `(-1, -1)` indicating padding. Passed through to the vision encoder for positional embedding computation. video_position_ids (`torch.LongTensor` of shape `(num_videos, num_frames, max_patches, 2)`, *optional*): 2D patch position coordinates from the video processor, with `(-1, -1)` indicating padding. Passed through to the vision encoder for positional embedding computation. """ if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError('You must specify exactly one of input_ids or inputs_embeds') image_mask, video_mask, audio_mask = self.get_placeholder_mask(input_ids, inputs_embeds) multimodal_mask = image_mask | video_mask | audio_mask # Replace image id with PAD if the image token if OOV, to avoid index-errors llm_input_ids = None if inputs_embeds is None: llm_input_ids = input_ids.clone() llm_input_ids[multimodal_mask] = self.config.text_config.pad_token_id inputs_embeds = self.get_input_embeddings()(llm_input_ids) if per_layer_inputs is None and self.config.get_text_config().hidden_size_per_layer_input: pad_embedding = self.language_model.embed_tokens.weight[self.config.text_config.pad_token_id, :] pad_embedding = pad_embedding.to(device=multimodal_mask.device) llm_inputs_embeds = torch.where(multimodal_mask[..., None], pad_embedding.view(1, 1, -1), inputs_embeds) per_layer_inputs = self.language_model.get_per_layer_inputs(llm_input_ids, llm_inputs_embeds) else: per_layer_inputs = None state = input_ids.new_tensor( [pixel_values is not None or pixel_values_videos is not None, input_features is not None], dtype=torch.bool) if dist.is_initialized() and is_deepspeed_enabled(): dist.all_reduce(state, dist.ReduceOp.MAX) has_image, has_audio = state.tolist() # Mixed modality training with both images and videos is not currently supported. if pixel_values is None and pixel_values_videos is None and has_image: inputs_embeds = _forward_dummy_image(self, inputs_embeds) # Merge text and images if pixel_values is not None: image_features = self.get_image_features(pixel_values, image_position_ids, return_dict=True).pooler_output image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype) # Confirm the number of soft tokens from the vision tower matches the number of slots in the embeddings. n_image_tokens = image_mask.sum() image_mask = image_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device) torch_compilable_check( inputs_embeds[image_mask].numel() == image_features.numel(), f'Image features and image tokens do not match, tokens: {n_image_tokens}, features:' f' {image_features.shape[0]}', ) inputs_embeds = inputs_embeds.masked_scatter( image_mask.to(inputs_embeds.device), image_features.to(inputs_embeds.device)) if pixel_values_videos is not None: video_features = self.get_video_features( pixel_values_videos, video_position_ids, return_dict=True).pooler_output video_features = video_features.to(inputs_embeds.device, inputs_embeds.dtype) # Confirm the number of soft tokens from the vision tower matches the number of slots in the embeddings. n_video_tokens = video_mask.sum() video_mask = video_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device) torch_compilable_check( inputs_embeds[video_mask].numel() == video_features.numel(), f'Video features and video tokens do not match, tokens: {n_video_tokens}, features:' f' {video_features.shape[0]}', ) inputs_embeds = inputs_embeds.masked_scatter( video_mask.to(inputs_embeds.device), video_features.to(inputs_embeds.device)) # Merge text and audio if input_features is not None and input_features_mask is not None: audio_output = self.get_audio_features(input_features, input_features_mask, return_dict=True) audio_features = audio_output.pooler_output audio_mask_from_encoder = audio_output.attention_mask # True = valid # Strip padding tokens: only keep real (non-padding) audio soft tokens. # audio_mask_from_encoder is True for valid positions, False for padding tokens. # This mirrors the vision encoder's padding stripping (see Gemma4VisionEncoder.forward). audio_features = audio_features[audio_mask_from_encoder] n_audio_tokens = audio_mask.sum() audio_mask = audio_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device) torch_compilable_check( inputs_embeds[audio_mask].numel() == audio_features.numel(), f'Audio features and audio tokens do not match, tokens: {n_audio_tokens}, features:' f' {audio_features.shape[0] * audio_features.shape[1]}', ) inputs_embeds = inputs_embeds.masked_scatter( audio_mask.to(inputs_embeds.device), audio_features.to(inputs_embeds.device)) elif has_audio and self.audio_tower is not None: feature_size = processor.feature_extractor.feature_size dummy_features = input_ids.new_zeros([1, 128, feature_size], dtype=self.audio_tower.dtype) dummy_mask = input_ids.new_ones([1, 128], dtype=torch.bool) audio_output = self.get_audio_features(dummy_features, dummy_mask, return_dict=True) audio_features = audio_output.pooler_output inputs_embeds = inputs_embeds + audio_features.mean() * 0. # It may already have been prepared by, e.g., `generate` if position_ids is None: past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens position_ids = position_ids.unsqueeze(0) bi_vision_attn = self.config.get_text_config().use_bidirectional_attention == 'vision' if not isinstance(causal_mask_mapping := attention_mask, dict): if bi_vision_attn and not transformers_5_9: from transformers.models.gemma4.modeling_gemma4 import create_causal_mask_mapping # Larger Gemma 4 models use Gemma 3's bidirectional attention mask for vision inputs causal_mask_mapping = create_causal_mask_mapping( self.config, inputs_embeds=inputs_embeds, attention_mask=attention_mask, past_key_values=past_key_values, position_ids=position_ids, mm_token_type_ids=mm_token_type_ids, ) else: mask_kwargs = { 'config': self.config, 'inputs_embeds': inputs_embeds, 'attention_mask': attention_mask, 'past_key_values': past_key_values, 'position_ids': position_ids, } if bi_vision_attn: from transformers.models.gemma4.modeling_gemma4 import get_block_sequence_ids_for_mask block_sequence_ids = torch.full([*inputs_embeds.size()[:-1]], -1, device=inputs_embeds.device) if mm_token_type_ids is not None: kwargs = { 'device': inputs_embeds.device } if 'device' in inspect.signature(get_block_sequence_ids_for_mask).parameters else {} block_sequence_ids = get_block_sequence_ids_for_mask(mm_token_type_ids, **kwargs) mask_kwargs['block_sequence_ids'] = block_sequence_ids causal_mask_mapping = create_masks_for_generate(**mask_kwargs) kwargs.pop('return_dict', None) outputs = self.language_model( per_layer_inputs=per_layer_inputs, attention_mask=causal_mask_mapping, position_ids=position_ids, past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, return_dict=True, **kwargs, ) return Gemma4ModelOutputWithPast( last_hidden_state=outputs.last_hidden_state, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states, attentions=outputs.attentions, image_hidden_states=image_features if pixel_values is not None else None, audio_hidden_states=audio_features if input_features is not None else None, ) model.origin_forward = model.forward model.forward = MethodType(forward, model) class Gemma4Loader(ModelLoader): def get_model(self, model_dir: str, config, processor, model_kwargs) -> PreTrainedModel: from transformers import Gemma4ForConditionalGeneration self.auto_model_cls = self.auto_model_cls or Gemma4ForConditionalGeneration model = super().get_model(model_dir, config, processor, model_kwargs) _patch_gemma4_forward(model.model, processor) return model register_model( ModelMeta( MLLMModelType.gemma4, [ ModelGroup([ Model('google/gemma-4-E2B', 'google/gemma-4-E2B'), Model('google/gemma-4-E2B-it', 'google/gemma-4-E2B-it'), Model('google/gemma-4-E4B', 'google/gemma-4-E4B'), Model('google/gemma-4-E4B-it', 'google/gemma-4-E4B-it'), ], template=TemplateType.gemma4_nothinking), ModelGroup([ Model('google/gemma-4-31B', 'google/gemma-4-31B'), Model('google/gemma-4-31B-it', 'google/gemma-4-31B-it'), Model('google/gemma-4-26B-A4B', 'google/gemma-4-26B-A4B'), Model('google/gemma-4-26B-A4B-it', 'google/gemma-4-26B-A4B-it'), ], template=TemplateType.gemma4), ], Gemma4Loader, architectures=['Gemma4ForConditionalGeneration'], model_arch=ModelArch.gemma3n, )) class Gemma4UnifiedLoader(ModelLoader): def get_model(self, model_dir: str, config, processor, model_kwargs) -> PreTrainedModel: from transformers import Gemma4UnifiedForConditionalGeneration self.auto_model_cls = self.auto_model_cls or Gemma4UnifiedForConditionalGeneration model = super().get_model(model_dir, config, processor, model_kwargs) _patch_gemma4_forward(model.model, processor, is_gemma4_unified=True) return model register_model( ModelMeta( MLLMModelType.gemma4_unified, [ ModelGroup([ Model('google/gemma-4-12B', 'google/gemma-4-12B'), Model('google/gemma-4-12B-it', 'google/gemma-4-12B-it'), ], template=TemplateType.gemma4), ], Gemma4UnifiedLoader, architectures=['Gemma4UnifiedForConditionalGeneration'], model_arch=ModelArch.gemma4_unified, requires=['transformers>=5.10.1'], )) class DiffusionGemmaLoader(ModelLoader): def get_model(self, model_dir: str, config, processor, model_kwargs) -> PreTrainedModel: from transformers import DiffusionGemmaForBlockDiffusion self.auto_model_cls = self.auto_model_cls or DiffusionGemmaForBlockDiffusion model = super().get_model(model_dir, config, processor, model_kwargs) model.prepare_inputs_for_generation = None model.config.use_cache = True return model register_model( ModelMeta( MLLMModelType.diffusion_gemma, [ ModelGroup([ Model('google/diffusiongemma-26B-A4B-it', 'google/diffusiongemma-26B-A4B-it'), ], template=TemplateType.diffusion_gemma), ], DiffusionGemmaLoader, architectures=['DiffusionGemmaForBlockDiffusion'], model_arch=ModelArch.diffusion_gemma, requires=['transformers>=5.11'], ))