Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

added_tokens.json

@@ -0,0 +1,28 @@
+{
+  "</think>": 151668,
+  "</tool_call>": 151658,
+  "</tool_response>": 151666,
+  "<think>": 151667,
+  "<tool_call>": 151657,
+  "<tool_response>": 151665,
+  "<|box_end|>": 151649,
+  "<|box_start|>": 151648,
+  "<|endoftext|>": 151643,
+  "<|file_sep|>": 151664,
+  "<|fim_middle|>": 151660,
+  "<|fim_pad|>": 151662,
+  "<|fim_prefix|>": 151659,
+  "<|fim_suffix|>": 151661,
+  "<|im_end|>": 151645,
+  "<|im_start|>": 151644,
+  "<|image_pad|>": 151655,
+  "<|object_ref_end|>": 151647,
+  "<|object_ref_start|>": 151646,
+  "<|quad_end|>": 151651,
+  "<|quad_start|>": 151650,
+  "<|repo_name|>": 151663,
+  "<|video_pad|>": 151656,
+  "<|vision_end|>": 151653,
+  "<|vision_pad|>": 151654,
+  "<|vision_start|>": 151652
+}

chat_template.jinja

@@ -0,0 +1,7 @@
+{% for message in messages %}{% if loop.first and message['role'] != 'system' %}<|im_start|>system
+You are a helpful assistant.<|im_end|>
+{% endif %}<|im_start|>{{ message['role'] }}
+{% if message['role'] == 'assistant' %}{% generation %}{{ message['content'][0]['text'] }}<|im_end|>
+{% endgeneration %}{% else %}{% for content in message['content'] %}{% if content['type'] == 'image' or 'image' in content or 'image_url' in content %}<|vision_start|><|image_pad|><|vision_end|>{% elif 'text' in content %}{{ content['text'] }}{% endif %}{% endfor %}<|im_end|>
+{% endif %}{% endfor %}{% if add_generation_prompt %}<|im_start|>assistant
+{% endif %}

config.json

@@ -0,0 +1,190 @@
+{
+  "architectures": [
+    "SmallVLMForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_smallvlm.SmallVLMConfig",
+    "AutoModelForCausalLM": "modeling_smallvlm.SmallVLMForCausalLM"
+  },
+  "image_token_id": 151655,
+  "language_model_config": {
+    "_name_or_path": "pretrained/Qwen/Qwen3-1.7B",
+    "add_cross_attention": false,
+    "architectures": [
+      "Qwen3ForCausalLM"
+    ],
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "bad_words_ids": null,
+    "begin_suppress_tokens": null,
+    "bos_token_id": 151643,
+    "chunk_size_feed_forward": 0,
+    "cross_attention_hidden_size": null,
+    "decoder_start_token_id": null,
+    "diversity_penalty": 0.0,
+    "do_sample": false,
+    "early_stopping": false,
+    "encoder_no_repeat_ngram_size": 0,
+    "eos_token_id": 151645,
+    "exponential_decay_length_penalty": null,
+    "finetuning_task": null,
+    "forced_bos_token_id": null,
+    "forced_eos_token_id": null,
+    "head_dim": 128,
+    "hidden_act": "silu",
+    "hidden_size": 2048,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "initializer_range": 0.02,
+    "intermediate_size": 6144,
+    "is_decoder": false,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "length_penalty": 1.0,
+    "max_length": 20,
+    "max_position_embeddings": 40960,
+    "max_window_layers": 28,
+    "min_length": 0,
+    "model_type": "qwen3",
+    "no_repeat_ngram_size": 0,
+    "num_attention_heads": 16,
+    "num_beam_groups": 1,
+    "num_beams": 1,
+    "num_hidden_layers": 28,
+    "num_key_value_heads": 8,
+    "num_return_sequences": 1,
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "output_scores": false,
+    "pad_token_id": null,
+    "prefix": null,
+    "problem_type": null,
+    "pruned_heads": {},
+    "remove_invalid_values": false,
+    "repetition_penalty": 1.0,
+    "return_dict": true,
+    "return_dict_in_generate": false,
+    "rms_norm_eps": 1e-06,
+    "rope_scaling": null,
+    "rope_theta": 1000000,
+    "sep_token_id": null,
+    "sliding_window": null,
+    "suppress_tokens": null,
+    "task_specific_params": null,
+    "temperature": 1.0,
+    "tf_legacy_loss": false,
+    "tie_encoder_decoder": false,
+    "tie_word_embeddings": true,
+    "tokenizer_class": null,
+    "top_k": 50,
+    "top_p": 1.0,
+    "torch_dtype": "bfloat16",
+    "torchscript": false,
+    "typical_p": 1.0,
+    "use_bfloat16": false,
+    "use_cache": true,
+    "use_sliding_window": false,
+    "vocab_size": 151936
+  },
+  "model_type": "smallvlm",
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.52.1",
+  "video_token_id": 151656,
+  "vision_abstractor_config": null,
+  "vision_model_config": {
+    "_name_or_path": "",
+    "add_cross_attention": false,
+    "architectures": null,
+    "auto_map": {
+      "AutoConfig": "configuration_moonvit.MoonViTConfig",
+      "AutoModel": "modeling_moonvit.MoonVitPretrainedModel"
+    },
+    "bad_words_ids": null,
+    "begin_suppress_tokens": null,
+    "bos_token_id": null,
+    "chunk_size_feed_forward": 0,
+    "cross_attention_hidden_size": null,
+    "decoder_start_token_id": null,
+    "diversity_penalty": 0.0,
+    "do_sample": false,
+    "early_stopping": false,
+    "encoder_no_repeat_ngram_size": 0,
+    "eos_token_id": null,
+    "exponential_decay_length_penalty": null,
+    "finetuning_task": null,
+    "forced_bos_token_id": null,
+    "forced_eos_token_id": null,
+    "hidden_size": 1152,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "init_pos_emb_height": 64,
+    "init_pos_emb_width": 64,
+    "intermediate_size": 4304,
+    "is_decoder": false,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "length_penalty": 1.0,
+    "max_length": 20,
+    "max_position_embeddings": 128000,
+    "merge_kernel_size": [
+      2,
+      2
+    ],
+    "min_length": 0,
+    "model_type": "moonvit",
+    "no_repeat_ngram_size": 0,
+    "num_attention_heads": 16,
+    "num_beam_groups": 1,
+    "num_beams": 1,
+    "num_hidden_layers": 27,
+    "num_return_sequences": 1,
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "output_scores": false,
+    "pad_token_id": null,
+    "patch_size": 14,
+    "prefix": null,
+    "problem_type": null,
+    "pruned_heads": {},
+    "remove_invalid_values": false,
+    "repetition_penalty": 1.0,
+    "return_dict": true,
+    "return_dict_in_generate": false,
+    "rope_scaling": {
+      "mrope_section": [
+        12,
+        12,
+        12
+      ],
+      "rope_type": "default",
+      "type": "default"
+    },
+    "rope_theta": 1000000.0,
+    "sep_token_id": null,
+    "suppress_tokens": null,
+    "task_specific_params": null,
+    "temperature": 1.0,
+    "text_hidden_size": 2048,
+    "tf_legacy_loss": false,
+    "tie_encoder_decoder": false,
+    "tie_word_embeddings": true,
+    "tokenizer_class": null,
+    "top_k": 50,
+    "top_p": 1.0,
+    "torch_dtype": null,
+    "torchscript": false,
+    "typical_p": 1.0,
+    "use_bfloat16": false
+  },
+  "vision_start_token_id": 151652
+}

configuration_moonvit.py

@@ -0,0 +1,39 @@
+from transformers.configuration_utils import PretrainedConfig
+
+
+class MoonViTConfig(PretrainedConfig):
+    model_type = "moonvit"
+
+    def __init__(
+        self,
+        patch_size: int = 14,
+        init_pos_emb_height: int = 64,
+        init_pos_emb_width: int = 64,
+        num_attention_heads: int = 16,
+        num_hidden_layers: int = 27,
+        hidden_size: int = 1152,
+        text_hidden_size: int = 2048,
+        intermediate_size: int = 4304,
+        merge_kernel_size: tuple[int, int] = (2, 2),
+        rope_theta: float = 1000000.0,
+        max_position_embeddings: int = 128000,
+        rope_scaling: dict = {'type': 'default', 'mrope_section': [12, 12, 12], 'rope_type': 'default'},
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.patch_size = patch_size
+        # Positional embedding config
+        self.init_pos_emb_height = init_pos_emb_height
+        self.init_pos_emb_width = init_pos_emb_width
+        # Transformer config
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.hidden_size = hidden_size
+        self.text_hidden_size = text_hidden_size
+        self.intermediate_size = intermediate_size
+        # Patch merger config
+        self.merge_kernel_size = merge_kernel_size
+
+        self.rope_theta = rope_theta
+        self.max_position_embeddings = max_position_embeddings
+        self.rope_scaling = rope_scaling

configuration_smallvlm.py

@@ -0,0 +1,52 @@
+from transformers import PretrainedConfig, AutoConfig, CONFIG_MAPPING
+from transformers.dynamic_module_utils import get_class_from_dynamic_module
+
+
+class SmallVLMConfig(PretrainedConfig):
+    model_type = "smallvlm"
+    is_composition = True
+
+    def __init__(
+            self,
+            language_model_config=None,
+            vision_model_config=None,
+            image_token_id=None,
+            **kwargs):
+        super().__init__(**kwargs)
+        if isinstance(language_model_config, dict):
+            if '_name_or_path' not in language_model_config:
+                language_model_config['_name_or_path'] = self._name_or_path
+            language_model_type = language_model_config.get('model_type', '')
+            is_remote_code = '.' in language_model_config.get('auto_map', {}).get('AutoConfig', '')
+            if language_model_type in CONFIG_MAPPING and not is_remote_code:
+                language_model_config = AutoConfig.for_model(**language_model_config)
+            elif language_model_type:
+                Config = get_class_from_dynamic_module(language_model_config["auto_map"]["AutoConfig"], language_model_config['_name_or_path'])
+                language_model_config = Config(**language_model_config)
+        self.language_model_config = language_model_config
+
+        if isinstance(vision_model_config, dict):
+            # if '_name_or_path' not in vision_model_config:
+            vision_model_config['_name_or_path'] = self._name_or_path
+            vision_model_type = vision_model_config.get('model_type', '')
+            is_remote_code = '.' in vision_model_config.get('auto_map', {}).get('AutoConfig', '')
+            if vision_model_type in CONFIG_MAPPING and not is_remote_code:
+                vision_model_config = AutoConfig.for_model(**vision_model_config)
+            elif vision_model_type:
+                Config = get_class_from_dynamic_module(vision_model_config["auto_map"]["AutoConfig"], vision_model_config['_name_or_path'])
+                vision_model_config = Config(**vision_model_config)
+        self.vision_model_config = vision_model_config
+
+        self.image_token_id = image_token_id
+        self.video_token_id = 151656
+        self.vision_start_token_id = 151652
+
+    @property
+    def hidden_size(self):
+        return self.language_model_config.hidden_size
+
+    @classmethod
+    def from_dict(cls, config_dict, **kwargs):
+        if 'name_or_path' in kwargs:
+            config_dict['_name_or_path'] = kwargs.pop('name_or_path')
+        return super().from_dict(config_dict, **kwargs)

generation_config.json

@@ -0,0 +1,13 @@
+{
+  "bos_token_id": 151643,
+  "do_sample": true,
+  "eos_token_id": [
+    151645,
+    151643
+  ],
+  "pad_token_id": 151643,
+  "temperature": 0.6,
+  "top_k": 20,
+  "top_p": 0.95,
+  "transformers_version": "4.52.1"
+}

image_processing_moonvit.py

@@ -0,0 +1,126 @@
+"""Image processor class for KimiVL."""
+
+import math
+import numpy as np
+from PIL import Image
+from typing import Optional, Union
+
+import torch
+from torchvision.transforms import functional as TF
+from transformers.image_utils import ImageInput, make_list_of_images, valid_images
+from transformers.image_processing_utils import BaseImageProcessor, BatchFeature
+from transformers.utils import TensorType
+
+
+OPENAI_DATASET_MEAN = (0.48145466, 0.4578275, 0.40821073)
+OPENAI_DATASET_STD = (0.26862954, 0.26130258, 0.27577711)
+
+
+class MoonViTImageProcessor(BaseImageProcessor):
+    model_type = "moonvit"
+
+    def __init__(
+        self,
+        patch_size: int = 14,
+        pad_input: bool = False,
+        image_mean: tuple[float, float, float] = OPENAI_DATASET_MEAN,
+        image_std: tuple[float, float, float] = OPENAI_DATASET_STD,
+        in_token_limit: int = 4096,
+        merge_kernel_size: list[int, int] = [2, 2],
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.in_token_limit = in_token_limit
+        self.patch_size = patch_size
+        self.pad_input = pad_input
+        self.image_mean = image_mean
+        self.image_std = image_std
+        self.merge_kernel_size = merge_kernel_size
+
+    def rescale(
+        self, image: Image.Image, merge_kernel_size: list[int, int] = [2, 2]
+    ) -> Image.Image:
+        w, h = image.size
+        patch_size = self.patch_size
+
+        if (w // patch_size) * (h // patch_size) > self.in_token_limit:
+            scale = math.sqrt(self.in_token_limit / ((w // patch_size) * (h // patch_size)))
+            new_w, new_h = int(w * scale), int(h * scale)
+            image = image.resize((new_w, new_h), Image.Resampling.BICUBIC)
+        if self.pad_input:
+            new_w, new_h = image.size
+            pad_size_h = merge_kernel_size[0] * patch_size
+            pad_size_w = merge_kernel_size[1] * patch_size
+
+            pad_h = (pad_size_h - new_h % pad_size_h) % pad_size_h
+            pad_w = (pad_size_w - new_w % pad_size_w) % pad_size_w
+
+            image = TF.pad(image, (0, 0, pad_w, pad_h))
+        else:
+            new_w, new_h = image.size
+            new_w = new_w - new_w % patch_size
+            new_h = new_h - new_h % patch_size
+            image = TF.center_crop(image, (new_h, new_w))
+
+        w, h = image.size
+        if w // patch_size >= 512 or h // patch_size >= 512:
+            raise ValueError("Exceed pos emb")
+
+        return image
+
+    def to_tensor(self, image: Image.Image) -> torch.Tensor:
+        return TF.to_tensor(image.convert("RGB"))
+
+    def normalize(self, image: torch.Tensor) -> torch.Tensor:
+        return TF.normalize(image, self.image_mean, self.image_std)
+
+    def patchify(self, image: torch.Tensor) -> tuple[torch.Tensor, list[int, int]]:
+        patch_size = self.patch_size
+        C, H, W = image.shape
+        patches = image.reshape(C, H // patch_size, patch_size, W // patch_size, patch_size)
+        patches = patches.permute(1, 3, 0, 2, 4)
+        patches = patches.contiguous().view(-1, C, patch_size, patch_size)
+        grid_hw = (H // patch_size, W // patch_size)
+        return patches, grid_hw
+
+    def _preprocess(self, image: ImageInput) -> tuple[torch.Tensor, list[int, int]]:
+        """
+        Preprocess image and patchify it.
+
+        Args:
+            image (`ImageInput`):
+                Image to preprocess. Expects pixel values ranging from 0 to 255. If pixel values range from 0 to 1, set `do_rescale=False`.
+
+        Returns:
+            patches: torch.Tensor
+            grid_hw: list[int, int]
+        """
+        image = self.rescale(image, self.merge_kernel_size)
+        image = self.to_tensor(image)
+        image = self.normalize(image)
+        patches, grid_hw = self.patchify(image)
+        return patches, grid_hw
+
+    def preprocess(
+        self,
+        images: ImageInput,
+        return_tensors: Optional[Union[str, TensorType]] = None,
+    ) -> BatchFeature:
+        images = make_list_of_images(images)
+
+        if not valid_images(images):
+            raise ValueError(
+                "Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, "
+                "torch.Tensor, tf.Tensor or jax.ndarray."
+            )
+
+        pixel_values, image_grid_hws = [], []
+        for image in images:
+            patches, image_grid_hw = self._preprocess(image)
+            pixel_values.append(patches)
+            image_grid_hws.append(image_grid_hw)
+        pixel_values = torch.concat(pixel_values, dim=0)
+        image_grid_hws = np.array(image_grid_hws)
+        data = {"pixel_values": pixel_values, "image_grid_hws": image_grid_hws}
+
+        return BatchFeature(data=data, tensor_type=return_tensors)

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:574587a763e7adc19186a15a9da4969aaa0563310c1d695f732c06e9831cd95c
+size 4938445104

modeling_moonvit.py

@@ -0,0 +1,729 @@
+
+
+
+import math
+from copy import deepcopy
+from typing import Union, Tuple, Sequence, Optional, List
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.checkpoint
+from transformers.activations import GELUActivation, ACT2FN, PytorchGELUTanh
+
+from transformers.activations import PytorchGELUTanh
+from transformers.modeling_utils import PreTrainedModel
+from transformers.utils import is_flash_attn_2_available
+
+from .configuration_moonvit import MoonViTConfig
+
+if is_flash_attn_2_available():
+    from flash_attn import flash_attn_varlen_func
+else:
+    flash_attn_varlen_func = None
+
+
+def rotate_half(x):
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+def apply_multimodal_rotary_pos_emb(q, k, cos, sin, mrope_section=[12, 12, 12], unsqueeze_dim=1):
+    mrope_section = mrope_section * 2
+    cos = torch.cat([m[i % 3] for i, m in enumerate(cos.split(mrope_section, dim=-1))], dim=-1).unsqueeze(unsqueeze_dim)
+    sin = torch.cat([m[i % 3] for i, m in enumerate(sin.split(mrope_section, dim=-1))], dim=-1).unsqueeze(unsqueeze_dim)
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    return q_embed, k_embed
+
+def get_rope_index(
+    image_token_id,
+    video_token_id,
+    vision_start_token_id,
+    spatial_merge_size: int = 2,
+    input_ids: Optional[torch.LongTensor] = None,
+    image_grid_thw: Optional[torch.LongTensor] = None,
+    video_grid_thw: Optional[torch.LongTensor] = None,
+    second_per_grid_ts: Optional[torch.Tensor] = None,
+    attention_mask: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+
+    mrope_position_deltas = []
+    if input_ids is not None and (image_grid_thw is not None or video_grid_thw is not None):
+        total_input_ids = input_ids
+        if attention_mask is None:
+            attention_mask = torch.ones_like(total_input_ids)
+        position_ids = torch.ones(3,input_ids.shape[0],input_ids.shape[1],dtype=input_ids.dtype,device=input_ids.device)
+        image_index, video_index = 0, 0
+        attention_mask = attention_mask.to(total_input_ids.device)
+        for i, input_ids in enumerate(total_input_ids):
+            input_ids = input_ids[attention_mask[i] == 1]
+            image_nums, video_nums = 0, 0
+            vision_start_indices = torch.argwhere(input_ids == vision_start_token_id).squeeze(1)
+            vision_tokens = input_ids[vision_start_indices + 1]
+            image_nums = (vision_tokens == image_token_id).sum()
+            video_nums = (vision_tokens == video_token_id).sum()
+            input_tokens = input_ids.tolist()
+            llm_pos_ids_list: list = []
+            st = 0
+            remain_images, remain_videos = image_nums, video_nums
+            for _ in range(image_nums + video_nums):
+                if image_token_id in input_tokens and remain_images > 0:
+                    ed_image = input_tokens.index(image_token_id, st)
+                else:
+                    ed_image = len(input_tokens) + 1
+                if video_token_id in input_tokens and remain_videos > 0:
+                    ed_video = input_tokens.index(video_token_id, st)
+                else:
+                    ed_video = len(input_tokens) + 1
+                if ed_image < ed_video:
+                    t, h, w = (image_grid_thw[image_index][0],image_grid_thw[image_index][1],image_grid_thw[image_index][2])
+                    second_per_grid_t = 0
+                    image_index += 1
+                    remain_images -= 1
+                    ed = ed_image
+
+                else:
+                    t, h, w = (video_grid_thw[video_index][0],video_grid_thw[video_index][1],video_grid_thw[video_index][2])
+                    if second_per_grid_ts is not None:
+                        second_per_grid_t = second_per_grid_ts[video_index]
+                    else:
+                        second_per_grid_t = 1.0
+                    video_index += 1
+                    remain_videos -= 1
+                    ed = ed_video
+                llm_grid_t, llm_grid_h, llm_grid_w = (t.item(),h.item() // spatial_merge_size,w.item() // spatial_merge_size)
+                text_len = ed - st
+
+                st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+                llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)
+
+                range_tensor = torch.arange(llm_grid_t).view(-1, 1)
+                expanded_range = range_tensor.expand(-1, llm_grid_h * llm_grid_w)
+
+                ## normalize type, send to device.
+                second_per_grid_t = torch.as_tensor(second_per_grid_t, dtype=range_tensor.dtype, device=range_tensor.device)
+
+                time_tensor = expanded_range * second_per_grid_t * 2
+
+                time_tensor_long = time_tensor.long()
+                t_index = time_tensor_long.flatten()
+
+                h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(llm_grid_t, -1, llm_grid_w).flatten()
+                w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(llm_grid_t, llm_grid_h, -1).flatten()
+                llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + text_len + st_idx)
+                st = ed + llm_grid_t * llm_grid_h * llm_grid_w
+
+            if st < len(input_tokens):
+                st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+                text_len = len(input_tokens) - st
+                llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)
+
+            llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)
+            position_ids[..., i, attention_mask[i] == 1] = llm_positions.to(position_ids.device)
+            mrope_position_deltas.append(llm_positions.max() + 1 - len(total_input_ids[i]))
+        mrope_position_deltas = torch.tensor(mrope_position_deltas, device=input_ids.device).unsqueeze(1)
+        return position_ids, mrope_position_deltas
+    else:
+        if attention_mask is not None:
+            position_ids = attention_mask.long().cumsum(-1) - 1
+            position_ids.masked_fill_(attention_mask == 0, 1)
+            position_ids = position_ids.unsqueeze(0).expand(3, -1, -1).to(attention_mask.device)
+            max_position_ids = position_ids.max(0, keepdim=False)[0].max(-1, keepdim=True)[0]
+            mrope_position_deltas = max_position_ids + 1 - attention_mask.shape[-1]
+        else:
+            position_ids = (torch.arange(input_ids.shape[1], device=input_ids.device).view(1, 1, -1).expand(3, input_ids.shape[0], -1))
+            mrope_position_deltas = torch.zeros([input_ids.shape[0], 1],device=input_ids.device,dtype=input_ids.dtype,)
+
+        return position_ids, mrope_position_deltas
+
+def multihead_attention(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    q_cu_seqlens: Optional[torch.Tensor] = None,
+    k_cu_seqlens: Optional[torch.Tensor] = None,
+):
+    """Multi-head attention using flash attention 2.
+    Args:
+        q, k, v: tensor of shape (batch_size, seqlen, num_heads, head_dim),
+            or (tot_seqlens, num_heads, head_dim) if packing.
+        q_cu_seqlens (torch.Tensor): cumulative sequence lengths of q.
+            The first element should be 0 and the last element should be q.shape[0].
+        k_cu_seqlens (torch.Tensor): cumulative sequence lengths of k.
+            The first element should be 0 and the last element should be k.shape[0].
+    Returns:
+        output: shape (batch_size, seqlen, dim) or (tot_seqlens, dim) if packing,
+            where dim = num_heads * head_dim
+    """
+    # Unified format legal check
+    assert q.dim() == k.dim() == v.dim() == 3, "q, k, v must have 3 dims"
+    assert q_cu_seqlens[-1] == q.shape[0], "q_cu_seqlens must sum to q.shape[0]"
+    assert (
+        k_cu_seqlens[-1] == k.shape[0] == v.shape[0]
+    ), "k_cu_seqlens must sum to k.shape[0]"
+    assert q.dtype in [
+        torch.bfloat16,
+        torch.float16,
+    ], f"unsupported dtype {q.dtype} for multihead attn"
+
+    max_seqlen_q = (q_cu_seqlens[1:] - q_cu_seqlens[:-1]).max().item()
+    max_seqlen_k = (k_cu_seqlens[1:] - k_cu_seqlens[:-1]).max().item()
+    attn_out = flash_attn_varlen_func(
+        q,
+        k,
+        v,
+        q_cu_seqlens,
+        k_cu_seqlens,
+        max_seqlen_q,
+        max_seqlen_k,
+        causal=False,
+    )
+    attn_out = attn_out.flatten(start_dim=-2)
+
+    return attn_out
+
+
+def sdpa_attention(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    attention_mask: torch.Tensor,
+) -> torch.Tensor:
+    """SDPA attention.
+    Args:
+        q, k, v: tensor of shape (batch_size, num_heads, seqlen, head_dim),
+            or (batch_size, seqlen, num_heads, head_dim) if packing.
+    """
+    # bs, num_heads, seq_length, head_dim = q.shape
+    # attention_mask = attention_mask.repeat(1, num_heads, 1, 1)
+    attn_output = F.scaled_dot_product_attention(q, k, v, attention_mask, dropout_p=0.0)
+    attn_output = attn_output.transpose(1, 2)
+    return attn_output
+
+
+def eager_attention(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    q_cu_seqlens: Optional[torch.Tensor] = None,
+    k_cu_seqlens: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    seq_length = q.shape[0]
+    attention_mask = torch.zeros(
+        [1, seq_length, seq_length], device=q.device, dtype=torch.bool
+    )
+    for i in range(1, len(q_cu_seqlens)):
+        attention_mask[
+            ...,
+            q_cu_seqlens[i - 1] : q_cu_seqlens[i],
+            q_cu_seqlens[i - 1] : q_cu_seqlens[i],
+        ] = True
+    q = q.transpose(0, 1)
+    k = k.transpose(0, 1)
+    v = v.transpose(0, 1)
+
+    attn_weight = q @ k.transpose(-2, -1) / math.sqrt(q.shape[-1])
+    attn_weight += attention_mask
+    attn_weight = torch.softmax(attn_weight, dim=-1, dtype=torch.float32).to(q.dtype)
+
+    attn_output = attn_weight @ v
+    attn_output = attn_output.transpose(0, 1)
+    attn_output = attn_output.reshape(seq_length, -1)
+    return attn_output
+
+
+VL_VISION_ATTENTION_FUNCTIONS = {
+    "flash_attention_2": multihead_attention,
+    "sdpa": sdpa_attention,
+    "eager": eager_attention,
+}
+
+
+def _apply_rope_input_validation(x, freqs_cis):
+    assert x.ndim == freqs_cis.ndim + 1, (x.shape, freqs_cis.shape)
+    assert x.shape[:-2] == freqs_cis.shape[:-1], (x.shape, freqs_cis.shape)
+    assert x.shape[-1] == 2 * freqs_cis.shape[-1], (x.shape, freqs_cis.shape)
+    assert freqs_cis.dtype == torch.complex64, freqs_cis.dtype
+
+
+def apply_rope(
+    xq: torch.Tensor, xk: torch.Tensor, freqs_cis: torch.Tensor
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    Args: (The leading dimensions of all inputs should be the same)
+        xq: query, tensor of shape (..., num_heads, head_dim)
+        xk: key, tensor of shape (..., num_heads, head_dim)
+        freqs_cis: tensor of shape (..., head_dim/2), dtype=torch.complex64. It contains the precomputed cis(freqs) for each position in the 2D grid.
+    Returns:
+        xq_out, xk_out: tensors of shape (..., num_heads, head_dim)
+    """
+    _apply_rope_input_validation(xq, freqs_cis)
+    _apply_rope_input_validation(xk, freqs_cis)
+
+    freqs_cis = freqs_cis.unsqueeze(-2)  # ..., 1, head_dim/2
+    # ..., num_heads, head_dim/2
+    xq_ = torch.view_as_complex(xq.float().view(*xq.shape[:-1], -1, 2))
+    xk_ = torch.view_as_complex(xk.float().view(*xq.shape[:-1], -1, 2))
+    xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(-2)  # ..., num_heads, head_dim
+    xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(-2)  # ..., num_heads, head_dim
+    return xq_out.type_as(xq), xk_out.type_as(xk)
+
+
+class Learnable2DInterpPosEmb(nn.Module):
+    def __init__(
+        self, height: int, width: int, dim: int, interpolation_mode: str = "bicubic"
+    ) -> None:
+        super().__init__()
+        self.height = height
+        self.width = width
+        self.interpolation_mode = interpolation_mode
+        self.weight = nn.Parameter(torch.empty(height, width, dim))
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        nn.init.normal_(self.weight)
+
+    def forward(self, x: torch.Tensor, grid_hws: torch.Tensor) -> torch.Tensor:
+        pos_embs = []
+        for shape in grid_hws.tolist():
+            if shape == self.weight.shape[:-1]:
+                pos_embs.append(self.weight.flatten(end_dim=1))
+            else:
+                pos_embs.append(
+                    F.interpolate(
+                        self.weight.permute((2, 0, 1)).unsqueeze(0),
+                        size=shape,
+                        mode=self.interpolation_mode,
+                    )
+                    .squeeze(0)
+                    .permute((1, 2, 0))
+                    .flatten(end_dim=1)
+                )
+        out = x + torch.cat(pos_embs)
+        return out
+
+
+class MoonVisionPatchEmbed(nn.Module):
+
+    def __init__(
+        self,
+        out_dim: int,
+        in_dim: int = 3,
+        patch_size: Union[int, Tuple[int, int]] = (14, 14),
+        pos_emb_height: int = 14,
+        pos_emb_width: int = 14,
+    ):
+        super().__init__()
+        assert isinstance(
+            patch_size, (int, Sequence)
+        ), f"Invalid patch_size type: {type(patch_size)}"
+        if isinstance(patch_size, int):
+            patch_size = (patch_size, patch_size)
+        assert (
+            len(patch_size) == 2
+        ), f"Expected patch_size to be a tuple of 2, got {patch_size}"
+        self.patch_size = patch_size
+
+        self.proj = nn.Conv2d(
+            in_dim, out_dim, kernel_size=patch_size, stride=patch_size
+        )
+
+        self.pos_emb = Learnable2DInterpPosEmb(
+            height=pos_emb_height, width=pos_emb_width, dim=out_dim
+        )
+
+    def forward(self, x: torch.Tensor, grid_hws: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x (L, Channels): input tensor
+            grid_hws (N, 2): grid height and width
+        Returns:
+            (L, Cout) tensor
+        """
+        x = self.proj(x).view(x.size(0), -1)
+        # apply positional embedding
+        x = self.pos_emb(x, grid_hws)
+        return x
+
+
+class Rope2DPosEmb(nn.Module):
+    """2D rotary position embedding with multi-resolution support.
+    This class is intended to be used in the following way:
+    1. Before training, create an instance of Rope2DPosEmb. This instance will hold the precomputed cis.
+    2. Before each forward pass, call `get_freqs_cis_by_*` to get the `freqs_cis` tensor for this iteration.
+    3. During the forward pass, pass the `freqs_cis` tensor to each attention layer, and call `apply` just before each attention operation.
+        The rope is shared across all attention layers and all heads.
+    Refs:
+    - RoFormer: https://arxiv.org/abs/2104.09864
+    - VisionLLaMA: https://arxiv.org/abs/2403.00522
+    - https://github.com/Meituan-AutoML/VisionLLaMA/blob/main/dit/models.py
+    Args:
+        dim (int): usually the multi-head attention dimension, should be divisible by 4 (TODO: relax this constraint if needed)
+        max_height (int): the maximum height of the 2D grid
+        max_width (int): the maximum width of the 2D grid
+        theta_base (float): the base of the theta
+        device (str): the device to store the precomputed cis
+    """
+
+    def __init__(self, dim: int, max_height: int, max_width: int, theta_base=10000):
+        super().__init__()
+        self.dim = dim
+        assert self.dim % 4 == 0, "dim must be divisible by 4"
+        self.max_height = max_height
+        self.max_width = max_width
+        self.theta_base = theta_base
+
+        self.freqs_cis = None
+
+    def extra_repr(self):
+        return f"dim={self.dim}, max_height={self.max_height}, max_width={self.max_width}, theta_base={self.theta_base}"
+
+    def _precompute_freqs_cis(self, device: torch.device) -> torch.Tensor:
+        """Calculate the cis(freqs) for each position in the 2D grid.
+        Return: complex tensor of shape (max_height, max_width, dim//2) and value:
+            height axis: ret[h, w, 2*i] = cis(h * theta_base**(-4*i/dim))
+            weight axis: ret[h, w, 2*i+1] = cis(w * theta_base**(-4*i/dim))   with (i in [0, dim//4))
+            note: `cis` is a mathematical notation defined by cis x = cos x + i sin x,
+        """
+        N = self.max_height * self.max_width
+        flat_pos = torch.arange(0, N).float().to(device)
+        x_pos = flat_pos % self.max_width
+        y_pos = flat_pos // self.max_width
+        dim_range = (
+            torch.arange(0, self.dim, 4)[: (self.dim // 4)].float().to(device)
+        )  # C/4
+        freqs = 1.0 / (self.theta_base ** (dim_range / self.dim))
+        x_freqs = torch.outer(x_pos, freqs).float()  # N, C/4
+        y_freqs = torch.outer(y_pos, freqs).float()  # N, C/4
+        x_cis = torch.polar(torch.ones_like(x_freqs), x_freqs)  # N, C/4
+        y_cis = torch.polar(torch.ones_like(y_freqs), y_freqs)  # N, C/4
+        # N, C/4, 2
+        freqs_cis = torch.cat(
+            [x_cis.unsqueeze(dim=-1), y_cis.unsqueeze(dim=-1)], dim=-1
+        )
+        # max_height, max_width, C/2
+        freqs_cis = freqs_cis.reshape(self.max_height, self.max_width, -1)
+        return freqs_cis
+
+    def get_freqs_cis(self, grid_hws: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            grid_hws (torch.Tensor): grid height and width
+        Returns:
+            freqs_cis: tensor of shape (sum(t * height * width), dim//2)
+        """
+        if self.freqs_cis is None:
+            self.freqs_cis = self._precompute_freqs_cis(grid_hws.device)
+
+        shapes = grid_hws.tolist()
+        assert all(
+            1 <= h <= self.max_height and 1 <= w <= self.max_width for h, w in shapes
+        ), (
+            shapes,
+            self.max_height,
+            self.max_width,
+        )
+        # freqs_cis = torch.cat(
+        #     [self.freqs_cis[:h, :w].reshape(-1, self.dim // 2) for h, w in shapes],
+        #     dim=0,
+        # )
+        max_h, max_w = grid_hws.max(dim=0).values.tolist()
+        max_h, max_w = max_h // 2, max_w // 2
+        freqs_cis = self.freqs_cis[:max_h, :max_w].reshape(-1, self.dim // 2).repeat(len(shapes), 1, 1)
+        return freqs_cis
+
+
+class MLP2(nn.Module):
+    """
+    Args:
+        dims: [in_dim, hidden_dim, out_dim]
+        bias: whether to use bias in linear layer.
+    """
+
+    def __init__(self, dims: list[int], activation, bias=True):
+        super().__init__()
+        assert len(dims) == 3
+        self.fc0 = nn.Linear(dims[0], dims[1], bias=bias)
+        self.fc1 = nn.Linear(dims[1], dims[2], bias=bias)
+        self.activation = activation
+        for m in [self.fc0, self.fc1]:
+            nn.init.trunc_normal_(m.weight, std=math.sqrt(2 / m.in_features))
+            if m.bias is not None:
+                nn.init.zeros_(m.bias)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.fc0(x)
+        x = self.activation(x)
+        return self.fc1(x)
+
+
+class MoonVitEncoderLayer(nn.Module):
+
+    def __init__(
+        self,
+        layer_idx: int,
+        num_heads: int,
+        hidden_dim: int,
+        mlp_dim: int,
+        attn_implementation: str = "eager",
+        activation=F.gelu,
+        attn_bias: bool = False,
+    ):
+        super().__init__()
+        self.layer_idx = layer_idx
+        self.num_heads = num_heads
+        self.hidden_dim = hidden_dim
+        self.hidden_size_per_attention_head = self.hidden_dim // self.num_heads
+        self.attn_implementation = attn_implementation
+
+        self.norm0 = nn.LayerNorm(hidden_dim)
+        self.norm1 = nn.LayerNorm(hidden_dim)
+        self.mlp = MLP2([hidden_dim, mlp_dim, hidden_dim], activation)
+        self.wqkv = nn.Linear(hidden_dim, hidden_dim * 3, bias=attn_bias)
+        self.wo = nn.Linear(hidden_dim, hidden_dim, bias=attn_bias)
+
+    def attention_qkvpacked(
+        self,
+        x: torch.Tensor,
+        attention_mask: torch.Tensor,
+        rope_freqs_cis: Optional[torch.Tensor] = None,
+        past_key_value = None
+    ):
+        """
+        Args:
+            x (torch.Tensor): (batch_size, seqlen, hidden_dim)
+            cu_seqlens (torch.Tensor):
+        """
+        batch_size, seqlen, hidden_dim = x.shape
+        xqkv = self.wqkv(x)
+        xqkv = xqkv.view(batch_size, seqlen, 3, self.num_heads, self.hidden_size_per_attention_head)
+        xq, xk, xv = torch.unbind(xqkv, dim=-3)
+        
+        xq = xq.transpose(1, 2)
+        xk = xk.transpose(1, 2)
+        xv = xv.transpose(1, 2)
+
+        # xq, xk = apply_rope(xq, xk, rope_freqs_cis)
+        cos, sin = rope_freqs_cis
+        xq, xk = apply_multimodal_rotary_pos_emb(xq, xk, cos, sin)
+
+        if past_key_value is not None:
+            xk, xv = past_key_value.update(xk, xv, self.layer_idx)
+
+        attn_func = VL_VISION_ATTENTION_FUNCTIONS[self.attn_implementation]
+        attn_out = attn_func(xq, xk, xv, attention_mask)
+        attn_out = attn_out.reshape(batch_size, seqlen, hidden_dim).contiguous()
+        attn_out = self.wo(attn_out)
+        return attn_out
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor,
+        rope_freqs_cis: Union[torch.Tensor, None] = None,
+        past_key_value = None
+    ) -> torch.Tensor:
+        """
+        Args:
+            hidden_states: non-packed (B, N, D) or packed (L, D). if non-packed, seqlens should be None, if packed, seqlens should be set
+        Returns:
+            output: same shape of input, non-packed (B, N, D) for non-packed input, (L, D) for packed input
+        """
+        residual = hidden_states
+        hidden_states = self.norm0(hidden_states)
+        attn_out = self.attention_qkvpacked(
+            hidden_states, attention_mask, rope_freqs_cis=rope_freqs_cis, past_key_value=past_key_value,
+        )
+        hidden_states = residual + attn_out
+
+        residual = hidden_states
+        hidden_states = self.mlp(self.norm1(hidden_states))
+        hidden_states = residual + hidden_states
+        return hidden_states
+
+
+class MoonVitEncoder(nn.Module):
+
+    def __init__(
+        self,
+        hidden_dim: int,
+        num_layers: int,
+        block_cfg: dict,
+    ) -> None:
+        super().__init__()
+        self.blocks = nn.ModuleList(
+            [MoonVitEncoderLayer(layer_idx, **block_cfg) for layer_idx in range(num_layers)]
+        )
+        self.final_layernorm = nn.LayerNorm(hidden_dim)
+        self.gradient_checkpointing = False
+
+    def forward(self, hidden_states, attention_mask, rope_freqs_cis, past_key_value=None) -> torch.Tensor:
+
+        for _, block in enumerate(self.blocks):
+            if self.gradient_checkpointing and self.training:
+                # hidden_states = self._gradient_checkpointing_func(
+                #     block.__call__, hidden_states, attention_mask, rope_freqs_cis
+                # )
+                hidden_states = torch.utils.checkpoint.checkpoint(
+                    block.__call__, hidden_states, attention_mask, rope_freqs_cis, past_key_value
+                )
+            else:
+                hidden_states = block(
+                    hidden_states, attention_mask, rope_freqs_cis=rope_freqs_cis, past_key_value=past_key_value,
+                )
+
+        hidden_states = self.final_layernorm(hidden_states)
+
+        return hidden_states
+
+
+def patch_merger(
+    x: torch.Tensor,
+    grid_hws: torch.Tensor,
+    merge_kernel_size: list[int, int] = (2, 2),
+) -> List[torch.Tensor]:
+    d_model = x.size(-1)
+
+    outputs = []
+    pre_sum = 0
+    for i, x_shape in enumerate(grid_hws.tolist()):
+        height, width = x_shape[0], x_shape[1]
+        # Get the current sequence
+        seq = x[pre_sum:pre_sum+height * width]
+        # Reshape along self.merge_kernel_size and concat to the last dimension
+        kernel_height, kernel_width = merge_kernel_size
+        new_height, new_width = height // kernel_height, width // kernel_width
+        reshaped_seq = seq.view(
+            new_height, kernel_height, new_width, kernel_width, d_model
+        )
+        reshaped_seq = reshaped_seq.permute(0, 2, 1, 3, 4).contiguous()
+        padded_seq = reshaped_seq.view(
+            new_height * new_width, kernel_height * kernel_width, -1
+        )
+        outputs.append(padded_seq)
+        pre_sum += height * width
+
+    return outputs
+
+
+class MultiModalProjector(nn.Module):
+
+    def __init__(self, config):
+        super().__init__()
+
+        self.hidden_size = (
+            config.hidden_size
+            * config.merge_kernel_size[0]
+            * config.merge_kernel_size[1]
+        )
+
+        self.pre_norm = torch.nn.LayerNorm(config.hidden_size, eps=1e-05)
+        self.linear_1 = nn.Linear(self.hidden_size, self.hidden_size, bias=True)
+        self.act = GELUActivation()
+        self.linear_2 = nn.Linear(self.hidden_size, config.text_hidden_size, bias=True)
+        # self.linear_2 = nn.Linear(self.hidden_size, config.hidden_size, bias=True)
+
+    def forward(self, image_features: list[torch.Tensor]) -> torch.Tensor:
+        # image_features = torch.cat(image_features, dim=0)
+        # hidden_states = self.pre_norm(image_features).view(-1, self.hidden_size)
+        hidden_states = self.pre_norm(image_features)
+        hidden_states = self.linear_1(hidden_states)
+        hidden_states = self.act(hidden_states)
+        hidden_states = self.linear_2(hidden_states)
+
+        return hidden_states
+
+
+class MoonVitPretrainedModel(PreTrainedModel):
+    config_class = MoonViTConfig
+    model_type = "moonvit"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["PackingTransformer"]
+    _supports_flash_attn_2 = True
+    _supports_sdpa = True
+
+    def __init__(self, config: MoonViTConfig, *inputs, **kwargs):
+        super().__init__(config, *inputs, **kwargs)
+        config = deepcopy(config)
+        self.merge_kernel_size = config.merge_kernel_size
+        self.patch_size = config.patch_size
+        self.patch_embed = MoonVisionPatchEmbed(
+            out_dim=config.hidden_size,
+            patch_size=config.patch_size,
+            pos_emb_height=config.init_pos_emb_height,
+            pos_emb_width=config.init_pos_emb_width,
+        )
+        self.rope_2d = Rope2DPosEmb(
+            config.hidden_size // config.num_attention_heads, 512, 512
+        )
+
+        self.encoder = MoonVitEncoder(
+            hidden_dim=config.hidden_size,
+            num_layers=config.num_hidden_layers,
+            block_cfg={
+                "num_heads": config.num_attention_heads,
+                "hidden_dim": config.hidden_size,
+                "mlp_dim": config.intermediate_size,
+                "activation": PytorchGELUTanh(),
+                "attn_bias": True,
+                "attn_implementation": config._attn_implementation,
+            },
+        )
+
+        self.pixel_merger = nn.Sequential(
+            nn.Linear(config.hidden_size*4, config.hidden_size),
+            nn.GELU(),
+            nn.Linear(config.hidden_size, config.hidden_size)
+        )
+            
+        self.projector = nn.Sequential(
+            nn.LayerNorm(config.hidden_size),
+            nn.Linear(config.hidden_size, config.hidden_size, bias=True),
+            nn.GELU(),
+            nn.Linear(config.hidden_size, config.text_hidden_size, bias=True),
+        )
+
+    def _init_weights(self, module):
+        """Initialize the weights"""
+        if isinstance(module, nn.Linear):
+            nn.init.xavier_uniform_(module.weight)
+            nn.init.normal_(module.bias, std=1e-6)
+        elif isinstance(module, nn.LayerNorm):
+            module.bias.data.zero_()
+            module.weight.data.fill_(1.0)
+        
+    def forward(
+        self, pixel_values: torch.Tensor, image_grid_hws: torch.Tensor
+    ) -> torch.Tensor:
+        """
+        Args:
+            pixel_values (torch.Tensor): The input pixel values.
+            grid_hws (torch.Tensor): The grid height and width.
+        Returns:
+            torch.Tensor: The output tokens.
+        """
+        hidden_states = self.patch_embed(pixel_values, image_grid_hws)
+        
+        hidden_states_list = patch_merger(
+            hidden_states, image_grid_hws, merge_kernel_size=self.merge_kernel_size
+        )
+        hidden_states = self.pixel_merger(torch.cat(hidden_states_list).view(-1, hidden_states.shape[-1] * 4))
+
+        num_tokens = (image_grid_hws.prod(dim=1) // 4).tolist()
+        hidden_states_list = hidden_states.split(num_tokens, dim=0)
+        max_length = max(num_tokens)
+        max_h, max_w = image_grid_hws.max(dim=0).values.tolist()
+        max_length = max_h * max_w // 4
+        hidden_states = torch.stack([F.pad(h, (0, 0, 0, max_length - h.shape[0])) for h in hidden_states_list])
+        attention_mask = torch.zeros(len(image_grid_hws), max_length, device=hidden_states.device, dtype=torch.bool)
+        for i in range(len(image_grid_hws)):
+            attention_mask[i][:num_tokens[i]] = True
+        from transformers.modeling_attn_mask_utils import _prepare_4d_attention_mask
+        attention_mask = _prepare_4d_attention_mask(attention_mask, hidden_states.dtype)
+        
+        rope_freqs_cis = self.rope_2d.get_freqs_cis(grid_hws=image_grid_hws)
+        hidden_states = self.encoder(hidden_states, attention_mask, rope_freqs_cis)
+        hidden_states = torch.cat([hidden_states[i][:num_tokens[i]] for i in range(len(image_grid_hws))])
+        # hidden_states = self.projector(hidden_states)
+        return hidden_states

modeling_smallvlm.py

@@ -0,0 +1,452 @@
+from typing import Optional, List
+import torch
+from torch import nn
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union
+from transformers import PreTrainedModel, AutoModel, AutoModelForCausalLM
+from transformers.modeling_outputs import ModelOutput
+from transformers.generation.utils import GenerationMixin
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
+from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+from transformers.models.qwen3.modeling_qwen3 import eager_attention_forward, BaseModelOutputWithPast
+
+from .modeling_moonvit import patch_merger, get_rope_index, apply_multimodal_rotary_pos_emb
+from .configuration_smallvlm import SmallVLMConfig
+
+class Qwen2_5_VLRotaryEmbedding(nn.Module):
+    def __init__(self, config, device=None):
+        super().__init__()
+        # BC: "rope_type" was originally "type"
+        if hasattr(config, "rope_scaling") and config.rope_scaling is not None:
+            self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type"))
+        else:
+            self.rope_type = "default"
+        self.max_seq_len_cached = config.max_position_embeddings
+        self.original_max_seq_len = config.max_position_embeddings
+
+        self.config = config
+        self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+
+        inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self.original_inv_freq = self.inv_freq
+
+    @torch.no_grad()
+    @dynamic_rope_update  # power user: used with advanced RoPE types (e.g. dynamic rope)
+    def forward(self, x, position_ids):
+        # In contrast to other models, Qwen2_5_VL has different position ids for the grids
+        # So we expand the inv_freq to shape (3, ...)
+        inv_freq_expanded = self.inv_freq[None, None, :, None].float().expand(3, position_ids.shape[1], -1, 1)
+        position_ids_expanded = position_ids[:, :, None, :].float()  # shape (3, bs, 1, positions)
+
+        device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
+        with torch.autocast(device_type=device_type, enabled=False):  # Force float32
+            freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(2, 3)
+            emb = torch.cat((freqs, freqs), dim=-1)
+            cos = emb.cos() * self.attention_scaling
+            sin = emb.sin() * self.attention_scaling
+
+        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+    
+
+def build_vision_model(config, model=None):
+    if model is None:
+        model = AutoModel.from_config(config, trust_remote_code=True)
+    return model
+
+def mrope_forward(
+    self,
+    hidden_states: torch.Tensor,
+    position_embeddings: Tuple[torch.Tensor, torch.Tensor],
+    attention_mask: Optional[torch.Tensor],
+    past_key_value: Optional[Cache] = None,
+    cache_position: Optional[torch.LongTensor] = None,
+    **kwargs,
+) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+    input_shape = hidden_states.shape[:-1]
+    hidden_shape = (*input_shape, -1, self.head_dim)
+
+    query_states = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+    key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+    value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+
+    cos, sin = position_embeddings
+    query_states, key_states = apply_multimodal_rotary_pos_emb(query_states, key_states, cos, sin, [16, 24, 24], unsqueeze_dim=1)
+
+    if past_key_value is not None:
+        # sin and cos are specific to RoPE models; cache_position needed for the static cache
+        cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+        key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
+    attention_interface: Callable = eager_attention_forward
+    if self.config._attn_implementation != "eager":
+        if self.config._attn_implementation == "sdpa" and kwargs.get("output_attentions", False):
+            pass
+        else:
+            attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+
+    attn_output, attn_weights = attention_interface(
+        self,
+        query_states,
+        key_states,
+        value_states,
+        attention_mask,
+        dropout=0.0 if not self.training else self.attention_dropout,
+        scaling=self.scaling,
+        sliding_window=self.sliding_window,  # diff with Llama
+        **kwargs,
+    )
+
+    attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+    attn_output = self.o_proj(attn_output)
+    return attn_output, attn_weights
+
+import transformers
+transformers.models.qwen3.modeling_qwen3.Qwen3Attention.forward = mrope_forward
+
+
+def forward(
+    self,
+    input_ids: Optional[torch.LongTensor] = None,
+    attention_mask: Optional[torch.Tensor] = None,
+    position_ids: Optional[torch.LongTensor] = None,
+    past_key_values: Optional[Cache] = None,
+    inputs_embeds: Optional[torch.FloatTensor] = None,
+    use_cache: Optional[bool] = None,
+    output_attentions: Optional[bool] = None,
+    output_hidden_states: Optional[bool] = None,
+    cache_position: Optional[torch.LongTensor] = None,
+    **flash_attn_kwargs,
+) -> BaseModelOutputWithPast:
+    output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+    output_hidden_states = (
+        output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+    )
+    use_cache = use_cache if use_cache is not None else self.config.use_cache
+
+    if (input_ids is None) ^ (inputs_embeds is not None):
+        raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+    if self.gradient_checkpointing and self.training and use_cache:
+        use_cache = False
+
+    # TODO (joao): remove this exception in v4.56 -- it exists for users that try to pass a legacy cache
+    if not isinstance(past_key_values, (type(None), Cache)):
+        raise ValueError("The `past_key_values` should be either a `Cache` object or `None`.")
+
+    if inputs_embeds is None:
+        inputs_embeds = self.embed_tokens(input_ids)
+
+    if use_cache and past_key_values is None:
+        past_key_values = DynamicCache()
+
+    if cache_position is None:
+        past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+        cache_position = torch.arange(
+            past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
+        )
+
+    if position_ids is None:
+        position_ids = cache_position.unsqueeze(0)
+
+    causal_mask = self._update_causal_mask(
+        attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions
+    )
+
+    hidden_states = inputs_embeds
+
+    # create position embeddings to be shared across the decoder layers
+    position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+    # decoder layers
+    all_hidden_states = () if output_hidden_states else None
+    all_self_attns = () if output_attentions else None
+
+    for decoder_layer in self.layers[: self.config.num_hidden_layers]:
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+
+        if self.gradient_checkpointing and self.training:
+            layer_outputs = self._gradient_checkpointing_func(
+                decoder_layer.__call__,
+                hidden_states,
+                causal_mask,
+                position_ids,
+                past_key_values,
+                output_attentions,
+                use_cache,
+                cache_position,
+                position_embeddings,
+            )
+        else:
+            layer_outputs = decoder_layer(
+                hidden_states,
+                attention_mask=causal_mask,
+                position_ids=position_ids,
+                past_key_value=past_key_values,
+                output_attentions=output_attentions,
+                use_cache=use_cache,
+                cache_position=cache_position,
+                position_embeddings=position_embeddings,
+                **flash_attn_kwargs,
+            )
+
+        hidden_states = layer_outputs[0]
+
+        if output_attentions:
+            all_self_attns += (layer_outputs[1],)
+
+    hidden_states = self.norm(hidden_states)
+
+    # add hidden states from the last decoder layer
+    if output_hidden_states:
+        all_hidden_states += (hidden_states,)
+
+    return BaseModelOutputWithPast(
+        last_hidden_state=hidden_states,
+        past_key_values=past_key_values if use_cache else None,
+        hidden_states=all_hidden_states,
+        attentions=all_self_attns,
+    )
+transformers.models.qwen3.modeling_qwen3.Qwen3Model.forward = forward
+
+class SmallVLMForCausalLM(PreTrainedModel, GenerationMixin):
+    config_class = SmallVLMConfig
+    supports_gradient_checkpointing = True
+    _skip_keys_device_placement = "past_key_values"
+    _supports_cache_class = True
+    _supports_flash_attn_2 = True
+    _supports_sdpa = True
+
+    def __init__(self, config, language_model=None, vision_model=None):
+        super().__init__(config)
+        self.rope_deltas = None  # cache rope_deltas here
+        
+        vision_model = build_vision_model(config.vision_model_config, vision_model)
+        if language_model is None:
+            kwargs_ = {}
+            if config._attn_implementation_internal is not None:
+                kwargs_['attn_implementation'] = config._attn_implementation_internal
+            language_model = AutoModelForCausalLM.from_config(config.language_model_config, trust_remote_code=True, **kwargs_)
+
+        self.vision_model = vision_model
+        
+        self.language_model = language_model
+
+        self.vision_to_text_proj = nn.Sequential( # map the text embeddings to vision encoder
+            nn.Linear(self.config.vision_model_config.hidden_size, self.config.language_model_config.hidden_size),
+            nn.GELU(),
+            nn.Linear(self.config.language_model_config.hidden_size, self.config.language_model_config.hidden_size)
+        )
+
+        self.text_to_vision_proj = nn.Sequential(
+            nn.Linear(self.config.language_model_config.hidden_size, self.config.vision_model_config.hidden_size),
+            nn.GELU(),
+            nn.Linear(self.config.vision_model_config.hidden_size, self.config.vision_model_config.hidden_size)
+        )
+        self.vision_rotary_emb = Qwen2_5_VLRotaryEmbedding(config.vision_model_config)
+        self.text_rotary_emb = Qwen2_5_VLRotaryEmbedding(config.language_model_config)
+        self.language_model.model.rotary_emb = self.text_rotary_emb
+
+        for layer in self.language_model.model.layers:
+            setattr(layer.self_attn, 'layer_idx', layer.self_attn.layer_idx + self.vision_model.config.num_hidden_layers)
+
+        self.gradient_checkpointing = False
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        pixel_values: Optional[torch.FloatTensor] = None,
+        grid_hws: Optional[torch.LongTensor] = None,
+    ):
+
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+        if self.gradient_checkpointing and self.training and use_cache:
+            use_cache = False
+
+        if use_cache and past_key_values is None:
+            past_key_values = DynamicCache()
+
+        inputs_embeds = self.get_input_embeddings()(input_ids)
+        inputs_embeds = self.text_to_vision_proj(inputs_embeds)
+        is_dummy_input = pixel_values is not None and pixel_values.size(0) == 0
+        if is_dummy_input:
+            pixel_values = torch.zeros((4,) + pixel_values.shape[1:], dtype=pixel_values.dtype, device=pixel_values.device)
+            grid_hws = torch.tensor([[1, 2, 2]], dtype=torch.int32).to(pixel_values.device)
+
+        if pixel_values is not None:
+            vision_embeds = self.vision_model.patch_embed(pixel_values, grid_hws[:, 1:])
+            vision_embeds_list = patch_merger(
+                vision_embeds, grid_hws[:, 1:], merge_kernel_size=self.vision_model.merge_kernel_size
+            )
+            vision_embeds = self.vision_model.pixel_merger(torch.cat(vision_embeds_list).view(-1, vision_embeds.shape[-1] * 4))
+            
+            vision_mask = (input_ids == self.config.image_token_id).to(inputs_embeds.device)
+            inputs_embeds[vision_mask] = vision_embeds
+
+            image_token_lens = (grid_hws.prod(dim=1) // 4)
+            bsz, src_len = attention_mask.size()
+            causal_mask = attention_mask[:, None, None, :].expand(bsz, 1, src_len, src_len).to(inputs_embeds.dtype)
+            causal_mask.tril_()
+            idx = 0
+            for i, _ in enumerate(causal_mask):
+                vision_mask = input_ids[i] == self.config.image_token_id
+                while (vision_mask.sum() > 0):
+                    start = torch.nonzero(vision_mask)[0][0]
+                    num = image_token_lens[idx]
+                    idx += 1
+                    causal_mask[i, 0, start:start+num, start:start+num] = 1
+                    vision_mask[start:start+num] = 0
+
+            causal_mask = 1.0 - causal_mask
+            causal_mask = causal_mask.masked_fill(causal_mask.to(torch.bool), torch.finfo(vision_embeds.dtype).min)
+        else:
+            causal_mask = None
+
+        if self.is_gradient_checkpointing and torch.is_grad_enabled() and self.training:
+            inputs_embeds.requires_grad_(True)
+
+        if cache_position is None:
+            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+            cache_position = torch.arange(
+                past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
+            )
+
+        if position_ids is None and (attention_mask is None or attention_mask.ndim == 2):
+            # calculate RoPE index once per generation in the pre-fill stage only
+            if (
+                (cache_position is not None and cache_position[0] == 0)
+                or self.rope_deltas is None
+                or (past_key_values is None or past_key_values.get_seq_length() == 0)
+            ):
+                position_ids, rope_deltas = get_rope_index(
+                    self.config.image_token_id,
+                    self.config.video_token_id,
+                    self.config.vision_start_token_id, 
+                    spatial_merge_size=2,
+                    input_ids=input_ids,
+                    image_grid_thw=grid_hws,
+                    video_grid_thw=None,
+                    attention_mask=attention_mask
+                )
+                self.rope_deltas = rope_deltas
+            # then use the prev pre-calculated rope-deltas to get the correct position ids
+            else:
+                batch_size, seq_length, _ = inputs_embeds.shape
+                delta = (
+                    (cache_position[0] + self.rope_deltas).to(inputs_embeds.device)
+                    if cache_position is not None
+                    else 0
+                )
+                position_ids = torch.arange(seq_length, device=inputs_embeds.device)
+                position_ids = position_ids.view(1, -1).expand(batch_size, -1)
+                if cache_position is not None:  # otherwise `deltas` is an int `0`
+                    delta = delta.repeat_interleave(batch_size // delta.shape[0], dim=0)
+                position_ids = position_ids.add(delta)
+                position_ids = position_ids.unsqueeze(0).expand(3, -1, -1)
+
+        position_embeddings = self.vision_rotary_emb(inputs_embeds, position_ids)
+        inputs_embeds = self.vision_model.encoder(inputs_embeds, causal_mask, position_embeddings, past_key_values)
+
+        # return ModelOutput(
+        #     last_hidden_state=self.vision_model.projector(inputs_embeds),
+        #     text_hidden_state=self.vision_to_text_proj(inputs_embeds),
+        # )
+    
+        inputs_embeds = self.vision_to_text_proj(inputs_embeds)
+
+        outputs = self.language_model(
+            input_ids=None,
+            labels=labels,
+            attention_mask=causal_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            cache_position=cache_position,
+            return_dict=True,
+        )
+
+        return ModelOutput(
+            loss=outputs.loss,
+            logits=outputs.logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+    def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None):
+        super().gradient_checkpointing_enable(gradient_checkpointing_kwargs)
+        self.language_model.enable_input_require_grads()
+
+    def get_input_embeddings(self):
+        return self.language_model.get_input_embeddings()
+
+    def set_input_embeddings(self, value):
+        self.language_model.set_input_embeddings(value)
+
+    def get_output_embeddings(self):
+        return self.language_model.get_output_embeddings()
+
+    def set_output_embeddings(self, new_embeddings):
+        self.language_model.set_output_embeddings(new_embeddings)
+
+    def set_decoder(self, decoder):
+        self.language_model.set_decoder(decoder)
+
+    def get_decoder(self):
+        return self.language_model.get_decoder()
+
+    def tie_weights(self):
+        return self.language_model.tie_weights()
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids,
+        past_key_values=None,
+        attention_mask=None,
+        inputs_embeds=None,
+        cache_position=None,
+        position_ids=None,
+        use_cache=True,
+        pixel_values=None,
+        **kwargs,
+    ):
+        # Overwritten -- in specific circumstances we don't want to forward image inputs to the model
+        model_inputs = super().prepare_inputs_for_generation(
+            input_ids,
+            past_key_values=past_key_values,
+            attention_mask=attention_mask,
+            inputs_embeds=inputs_embeds,
+            cache_position=cache_position,
+            position_ids=position_ids,
+            pixel_values=pixel_values,
+            use_cache=use_cache,
+            **kwargs,
+        )
+
+        # Qwen2-5-VL position_ids are prepareed with rope_deltas in forward
+        model_inputs["position_ids"] = None
+        if cache_position[0] != 0:
+            model_inputs["pixel_values"] = None
+
+        return model_inputs

preprocessor_config.json

@@ -0,0 +1,26 @@
+{
+  "auto_map": {
+    "AutoProcessor": "processing_smallvlm.SmallVLMProcessor",
+    "AutoImageProcessor": "image_processing_moonvit.MoonViTImageProcessor"
+  },
+  "image_mean": [
+    0.5,
+    0.5,
+    0.5
+  ],
+  "image_processor_type": "MoonViTImageProcessor",
+  "image_std": [
+    0.5,
+    0.5,
+    0.5
+  ],
+  "in_token_limit": 16384,
+  "merge_kernel_size": [
+    2,
+    2
+  ],
+  "num_pooled_tokens": 1024,
+  "pad_input": true,
+  "patch_size": 14,
+  "processor_class": "SmallVLMProcessor"
+}

processing_smallvlm.py

@@ -0,0 +1,182 @@
+import torch
+
+from collections import UserDict, OrderedDict
+from typing import Union, List, Dict, Any
+
+from transformers.processing_utils import ProcessorMixin
+from transformers.feature_extraction_utils import BatchFeature
+from transformers.utils.chat_template_utils import render_jinja_template
+
+from .image_processing_moonvit import MoonViTImageProcessor
+
+class SmallVLMProcessor(ProcessorMixin):
+    attributes = ["tokenizer", "image_processor"]
+    optional_attributes = ['chat_template']
+    model_input_names = ['input_ids', 'attention_mask', 'pixel_values']
+    image_processor_class = "AutoImageProcessor"
+    tokenizer_class = "AutoTokenizer"
+
+    image_token = '<|image_pad|>'
+
+    def __init__(self, tokenizer, image_processor, chat_template, **kwargs):
+        super().__init__(tokenizer=tokenizer, image_processor=image_processor, chat_template=chat_template)
+        self.tokenizer.add_special_tokens({'additional_special_tokens': [self.image_token]}, replace_additional_special_tokens=False)
+        self.image_token_id = self.tokenizer.convert_tokens_to_ids(self.image_token)
+
+    def __call__(self, inputs=None, images=[], text=None, **kwargs) -> BatchFeature:
+
+        truncation = kwargs.pop('truncation', False)
+        max_length = kwargs.pop('max_length', 1024)
+        padding = kwargs.pop('padding', False)
+
+        if inputs is None:
+            inputs = {}
+        if isinstance(inputs, UserDict):
+            inputs = inputs.data
+
+        if 'input_ids' not in inputs:
+            input_ids = self.tokenizer(text, padding=False, truncation=False, return_attention_mask=False, **kwargs)['input_ids'][0]
+            inputs['input_ids'] = input_ids.tolist()
+
+        inputs = self.process_images(images, inputs=inputs)
+
+        if 'attention_mask' not in inputs:
+            inputs['attention_mask'] = [1] * len(inputs['input_ids'])
+
+        if 'assistant_masks' in inputs:
+            inputs['prompt_mask'] = [1-x for x in inputs.pop('assistant_masks')]
+
+        inputs = self.process_inputs(inputs)
+
+
+        if truncation and len(inputs['input_ids']) > max_length:
+            inputs = self.truncate(inputs, max_length)
+
+        if padding and len(inputs['input_ids']) < max_length:
+            inputs = self.padding(inputs, max_length)
+
+        inputs = self.to_tensor(inputs)
+
+        self.check(inputs)
+
+        new_inputs = {
+            "input_ids": inputs["input_ids"],
+            "attention_mask": inputs["attention_mask"],
+        }
+        if "pixel_values" in inputs:
+            new_inputs['pixel_values'] = inputs['pixel_values']
+            new_inputs['grid_hws'] = torch.cat([torch.ones_like(inputs['image_grid_hws'])[:, :1], inputs['image_grid_hws']], dim=1)
+        if 'prompt_mask' in inputs:
+            new_inputs['prompt_mask'] = inputs['prompt_mask']
+        
+        return BatchFeature(new_inputs)
+
+    def process_images(self, images, inputs):
+        if len(images) > 0:
+            pixel_values, image_grid_hws = self.image_transform(images)
+        else:
+            pixel_values = torch.zeros((0, 3, 14, 14), dtype=torch.float32)
+            image_grid_hws = torch.zeros((0, 2), dtype=torch.int64)
+        
+        inputs['pixel_values'] = pixel_values
+        inputs['image_grid_hws'] = image_grid_hws
+        return inputs
+
+    def image_transform(self, images):
+        image_inputs = self.image_processor(images, return_tensors='pt')
+        return image_inputs['pixel_values'], image_inputs['image_grid_hws']
+
+    def truncate(self, inputs: Dict[str, Any], max_length: int):
+        assert self.image_token_id not in inputs['input_ids'][max_length:], f"Truncate image token is not allowed."
+
+        inputs['input_ids'] = inputs['input_ids'][:max_length]
+        inputs['attention_mask'] = inputs['attention_mask'][:max_length]
+        if 'prompt_mask' in inputs:
+            inputs['prompt_mask'] = inputs['prompt_mask'][:max_length]
+
+        return inputs
+
+    def get_image_token_length(self, inputs: Dict[str, Any]) -> List[int]:
+        image_grid_hws = inputs.get('image_grid_hws', None)
+        if image_grid_hws is None:
+            return []
+        image_token_lens = (image_grid_hws.prod(dim=1) // 4).tolist()
+        return image_token_lens
+
+    def process_inputs(self, inputs: Dict[str, Any]):
+        graft_token_lens = self._get_graft_token_length(inputs)
+
+        inputs['input_ids'] = self._graft_token(inputs['input_ids'], graft_token_lens, self.image_token_id)
+        inputs['attention_mask'] = self._graft_token(inputs['attention_mask'], graft_token_lens, 'replicate')
+        if 'prompt_mask' in inputs:
+            inputs['prompt_mask'] = self._graft_token(inputs['prompt_mask'], graft_token_lens, 'replicate')
+
+        return inputs
+
+    def _graft_token(self, seq, graft_token_lens, value):
+        if value == 'replicate':
+            for i in reversed(graft_token_lens.keys()):
+                seq[i:] = [seq[i]] * graft_token_lens[i] + seq[i+1:]
+        else:
+            for i in reversed(graft_token_lens.keys()):
+                assert value == seq[i]
+                seq[i:] = [value] * graft_token_lens[i] + seq[i+1:]
+        return seq
+
+    def _get_graft_token_length(self, inputs: Dict[str, Any]) -> Dict[int, int]:
+        image_token_pos = [i for i, x in enumerate(inputs['input_ids']) if x == self.image_token_id]
+        image_token_lens = self.get_image_token_length(inputs)
+
+        assert len(image_token_pos) == len(image_token_lens), \
+            "Wrong image token count, " \
+            f"image_token_count({len(image_token_pos)}) != image_count({len(image_token_lens)})"
+
+        graft_token_lens = OrderedDict(item for item in zip(image_token_pos, image_token_lens))
+
+        return graft_token_lens
+
+    def check(self, inputs: Dict[str, Any]):
+        image_embed_token_count = torch.count_nonzero(inputs['input_ids'] == self.image_token_id).item()
+        image_embed_count = sum(self.get_image_token_length(inputs))
+        assert image_embed_token_count == image_embed_count, "Wrong image embed token count"
+
+    def padding(self, inputs: Dict[str, Any], max_length: int):
+        padding_len = max_length - len(inputs['input_ids'])
+        inputs['input_ids'] += [self.pad_token_id] * padding_len
+        inputs['attention_mask'] += [0] * padding_len
+        if 'prompt_mask' in inputs:
+            inputs['prompt_mask'] += [0] * padding_len
+        return inputs
+
+    def decode(self, token_ids: Union[List[int], torch.Tensor], **kwargs):
+        if isinstance(token_ids, torch.Tensor):
+            token_ids = token_ids.tolist()
+        text = self.tokenizer.decode(token_ids, **kwargs)
+        return text
+
+    def batch_decode(self, sequences: Union[List[List[int]], torch.Tensor], **kwargs):
+        if isinstance(sequences, torch.Tensor):
+            sequences = sequences.tolist()
+        texts = self.tokenizer.batch_decode(sequences, **kwargs)
+        return texts
+
+    def to_tensor(self, inputs):
+        inputs['input_ids'] = torch.tensor([inputs['input_ids']], dtype=torch.long)
+        inputs['attention_mask'] = torch.tensor([inputs['attention_mask']], dtype=torch.bool)
+        if 'prompt_mask' in inputs:
+            inputs['prompt_mask'] = torch.tensor([inputs['prompt_mask']], dtype=torch.bool)
+        return inputs
+
+    @property
+    def pad_token_id(self):
+        return self.tokenizer.pad_token_id
+
+    @property
+    def special_tokens(self):
+        return [token.content for token in self.tokenizer.added_tokens_decoder.values()]
+
+    def __repr__(self):
+        pass
+
+    def __str__(self):
+        return ''

processor_config.json

@@ -0,0 +1,6 @@
+{
+  "auto_map": {
+    "AutoProcessor": "processing_smallvlm.SmallVLMProcessor"
+  },
+  "processor_class": "SmallVLMProcessor"
+}

special_tokens_map.json

@@ -0,0 +1,31 @@
+{
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>"
+  ],
+  "eos_token": {
+    "content": "<|im_end|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:aeb13307a71acd8fe81861d94ad54ab689df773318809eed3cbe794b4492dae4
+size 11422654

tokenizer_config.json

@@ -0,0 +1,243 @@
+{
+  "add_bos_token": false,
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "151643": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151644": {
+      "content": "<|im_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151645": {
+      "content": "<|im_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151646": {
+      "content": "<|object_ref_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151647": {
+      "content": "<|object_ref_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151648": {
+      "content": "<|box_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151649": {
+      "content": "<|box_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151650": {
+      "content": "<|quad_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151651": {
+      "content": "<|quad_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151652": {
+      "content": "<|vision_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151653": {
+      "content": "<|vision_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151654": {
+      "content": "<|vision_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151655": {
+      "content": "<|image_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151656": {
+      "content": "<|video_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151657": {
+      "content": "<tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151658": {
+      "content": "</tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151659": {
+      "content": "<|fim_prefix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151660": {
+      "content": "<|fim_middle|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151661": {
+      "content": "<|fim_suffix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151662": {
+      "content": "<|fim_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151663": {
+      "content": "<|repo_name|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151664": {
+      "content": "<|file_sep|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151665": {
+      "content": "<tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151666": {
+      "content": "</tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151667": {
+      "content": "<think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151668": {
+      "content": "</think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    }
+  },
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>"
+  ],
+  "auto_map": {
+    "AutoProcessor": "processing_smallvlm.SmallVLMProcessor"
+  },
+  "bos_token": null,
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "errors": "replace",
+  "extra_special_tokens": {},
+  "model_max_length": 131072,
+  "pad_token": "<|endoftext|>",
+  "processor_class": "SmallVLMProcessor",
+  "split_special_tokens": false,
+  "tokenizer_class": "Qwen2Tokenizer",
+  "unk_token": null
+}