Upload ImpForCausalLM

config.json

@@ -1,13 +1,13 @@
 {
-  "_name_or_path": "MILVLG/imp-v1-3b",
+  "_name_or_path": "/workspace/imp/sparrow",
   "activation_function": "gelu_new",
   "architectures": [
     "ImpForCausalLM"
   ],
   "attn_pdrop": 0.0,
   "auto_map": {
-    "AutoConfig": "MILVLG/imp-v1-3b--configuration_imp.ImpConfig",
-    "AutoModelForCausalLM": "MILVLG/imp-v1-3b--modeling_imp.ImpForCausalLM"
+    "AutoConfig": "configuration_imp.ImpConfig",
+    "AutoModelForCausalLM": "modeling_imp.ImpForCausalLM"
   },
   "embd_pdrop": 0.0,
   "eos_token_id": 50295,

configuration_imp.py

@@ -0,0 +1,183 @@
+# Copyright (c) MILVLG team.
+# Licensed under the Apache 2.0 license.
+#
+# Some code here is copied from the project Phi-2 (https://huggingface.co/microsoft/phi-2),
+# SigLIP@transformers==4.37.0.dev0 (https://huggingface.co/google/siglip-so400m-patch14-384),
+# and Llava (https://github.com/haotian-liu/LLaVA), and modified by 
+# Zhenwei Shao ([email protected]) @ MILVLG. We thank them for their great works.
+#
+# We keep their original copyright statements as follows, which should be inherited:
+# ------------------------------- Phi-2 ---------------------------------------------
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+# https://huggingface.co/google/siglip-so400m-patch14-384
+# 
+# Copyright (c) 2022, Tri Dao, [email protected].
+# Licensed under the BSD 3-Clause License.
+# ------------------------------- SigLIP --------------------------------------------
+# Copyright 2024 Google AI and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ------------------------------- Llava ---------------------------------------------
+# Copyright 2023 Haotian Liu
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# -----------------------------------------------------------------------------------
+
+
+import os
+import math
+from typing import Optional, Union
+
+from transformers import PretrainedConfig
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+
+class PhiConfig(PretrainedConfig):
+    """Phi configuration."""
+
+    model_type = "phi-msft"
+    attribute_map = {
+        "max_position_embeddings": "n_positions",
+        "hidden_size": "n_embd",
+        "num_attention_heads": "n_head",
+        "num_hidden_layers": "n_layer",
+    }
+
+    def __init__(
+        self,
+        vocab_size: int = 50304,
+        n_positions: int = 2048,
+        n_embd: int = 1024,
+        n_layer: int = 20,
+        n_inner: Optional[int] = None,
+        n_head: int = 16,
+        n_head_kv: Optional[int] = None,
+        rotary_dim: Optional[int] = 32,
+        activation_function: Optional[str] = "gelu_new",
+        flash_attn: bool = False,
+        flash_rotary: bool = False,
+        fused_dense: bool = False,
+        attn_pdrop: float = 0.0,
+        embd_pdrop: float = 0.0,
+        resid_pdrop: float = 0.0,
+        layer_norm_epsilon: float = 1e-5,
+        initializer_range: float = 0.02,
+        tie_word_embeddings: bool = False,
+        pad_vocab_size_multiple: int = 64,
+        **kwargs
+    ) -> None:
+        self.vocab_size = int(math.ceil(vocab_size / pad_vocab_size_multiple) * pad_vocab_size_multiple)
+        self.n_positions = n_positions
+        self.n_embd = n_embd
+        self.n_layer = n_layer
+        self.n_inner = n_inner
+        self.n_head = n_head
+        self.n_head_kv = n_head_kv
+        self.rotary_dim = min(rotary_dim, n_embd // n_head)
+        self.activation_function = activation_function
+        self.flash_attn = flash_attn
+        self.flash_rotary = flash_rotary
+        self.fused_dense = fused_dense
+        self.attn_pdrop = attn_pdrop
+        self.embd_pdrop = embd_pdrop
+        self.resid_pdrop = resid_pdrop
+        self.layer_norm_epsilon = layer_norm_epsilon
+        self.initializer_range = initializer_range
+
+        super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
+
+
+
+class SiglipVisionConfig(PretrainedConfig):
+
+    model_type = "siglip_vision_model"
+
+    def __init__(
+        self,
+        hidden_size=768,
+        intermediate_size=3072,
+        num_hidden_layers=12,
+        num_attention_heads=12,
+        num_channels=3,
+        image_size=224,
+        patch_size=16,
+        hidden_act="gelu_pytorch_tanh",
+        layer_norm_eps=1e-6,
+        attention_dropout=0.0,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_channels = num_channels
+        self.patch_size = patch_size
+        self.image_size = image_size
+        self.attention_dropout = attention_dropout
+        self.layer_norm_eps = layer_norm_eps
+        self.hidden_act = hidden_act
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path: Union[str, os.PathLike], **kwargs) -> "PretrainedConfig":
+        cls._set_token_in_kwargs(kwargs)
+
+        config_dict, kwargs = cls.get_config_dict(pretrained_model_name_or_path, **kwargs)
+
+        # get the vision config dict if we are loading from SiglipConfig
+        if config_dict.get("model_type") == "siglip":
+            config_dict = config_dict["vision_config"]
+
+        if "model_type" in config_dict and hasattr(cls, "model_type") and config_dict["model_type"] != cls.model_type:
+            logger.warning(
+                f"You are using a model of type {config_dict['model_type']} to instantiate a model of type "
+                f"{cls.model_type}. This is not supported for all configurations of models and can yield errors."
+            )
+
+        return cls.from_dict(config_dict, **kwargs)
+
+
+class ImpConfig(PhiConfig):
+    model_type = "imp"
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.image_token_index = getattr(self, "image_token_index", 50296)
+        self.image_token = getattr(self, "image_token", "<image>")
+
+        if not hasattr(self, "vision_tower_config") and hasattr(self, "mm_vision_tower"):
+            vision_tower_config = SiglipVisionConfig.from_pretrained(self.mm_vision_tower)
+            self.vision_tower_config = vision_tower_config.to_diff_dict()
+    
+    @property
+    def vision_tower_cfg(self):
+        cfg = SiglipVisionConfig.from_dict(self.vision_tower_config)
+        # imp-v1 only supports `patch` feature for now w/o cls token
+        # cfg.mm_vision_select_feature = self.mm_vision_select_feature
+        cfg.mm_vision_select_layer = self.mm_vision_select_layer
+        cfg.mm_vision_tower = self.mm_vision_tower
+        return cfg

modeling_imp.py

@@ -0,0 +1,1262 @@
+# Copyright (c) MILVLG team.
+# Licensed under the Apache 2.0 license.
+#
+# Some code here is copied from the project Phi-2 (https://huggingface.co/microsoft/phi-2),
+# SigLIP@transformers==4.37.0.dev0 (https://huggingface.co/google/siglip-so400m-patch14-384),
+# and Llava (https://github.com/haotian-liu/LLaVA), and modified by 
+# Zhenwei Shao ([email protected]) @ MILVLG. We thank them for their great works.
+# And their original licenses and copyright should be inherited (see the statements
+# in `configuration_imp.py` for more details).
+
+
+# Be careful: The way how `past_key_values.seqlen_offset` is updated is modified from
+# the implementation of original Phi-2. See the comments below for details.
+
+from __future__ import annotations
+import os
+import math
+import re
+from dataclasses import dataclass, field
+from typing import Any, Dict, Optional, Tuple, Union, List
+from abc import ABC, abstractmethod
+
+import torch
+import torch.nn as nn
+from einops import rearrange, repeat
+from transformers import (
+    PretrainedConfig, 
+    PreTrainedModel,
+    AutoConfig,
+    AutoModelForCausalLM
+)    
+from transformers.activations import ACT2FN
+from transformers.modeling_outputs import CausalLMOutputWithPast
+import sys
+from .configuration_imp import PhiConfig, ImpConfig
+from .vision_encoder import VisionTower
+
+try:
+    from flash_attn.bert_padding import pad_input, unpad_input
+    from flash_attn.layers.rotary import RotaryEmbedding as FlashRotaryEmbedding
+    from flash_attn.modules.mha import FlashCrossAttention, FlashSelfAttention
+    from flash_attn.ops.fused_dense import FusedDense
+except:
+    pad_input, unpad_input = None, None
+    FlashRotaryEmbedding = None
+    FlashSelfAttention, FlashCrossAttention = None, None
+    FusedDense = None
+
+
+@dataclass
+class InferenceParams:
+    """Inference parameters passed to model to efficiently calculate
+    and store context during inference.
+
+    Reference:
+        https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/utils/generation.py.
+
+    Args:
+        max_seqlen: Maximum sequence length.
+        max_batch_size: Maximum batch size.
+        seqlen_offset: Sequence length offset.
+        batch_size_offset: Batch size offset.
+        key_value_memory_dict: Key value memory dictionary.
+        lengths_per_sample: Lengths per sample.
+
+    """
+
+    max_seqlen: int = field(metadata={"help": "Maximum sequence length."})
+
+    max_batch_size: int = field(metadata={"help": "Maximum batch size."})
+
+    seqlen_offset: int = field(default=0, metadata={"help": "Sequence length offset."})
+
+    batch_size_offset: int = field(default=0, metadata={"help": "Batch size offset."})
+
+    key_value_memory_dict: Dict[str, Any] = field(
+        default_factory=dict, metadata={"help": "Key value memory dictionary."}
+    )
+
+    lengths_per_sample: torch.Tensor = field(default=None, metadata={"help": "Lengths per sample."})
+
+
+class Embedding(nn.Module):
+    """Token embedding with dropout."""
+
+    def __init__(self, config: PretrainedConfig) -> None:
+        super().__init__()
+
+        self.wte = nn.Embedding(config.vocab_size, config.n_embd)
+        self.drop = nn.Dropout(config.embd_pdrop)
+
+    def forward(self, input_ids: torch.LongTensor) -> torch.FloatTensor:
+        input_shape = input_ids.size()
+        input_ids = input_ids.view(-1, input_shape[-1])
+
+        hidden_states = self.wte(input_ids)
+        hidden_states = self.drop(hidden_states)
+
+        return hidden_states
+
+
+
+def _apply_rotary_emb(
+    x: torch.FloatTensor,
+    cos: torch.FloatTensor,
+    sin: torch.FloatTensor,
+) -> torch.FloatTensor:
+    _, seqlen, _, _ = x.shape
+    _, rotary_dim = cos.shape
+    rotary_dim *= 2
+
+    x_rot = x[:, :, :, :rotary_dim]
+    x_pass = x[:, :, :, rotary_dim:]
+
+    x1, x2 = x_rot.chunk(2, dim=-1)
+    c, s = rearrange(cos[:seqlen], "s d -> s 1 d"), rearrange(sin[:seqlen], "s d -> s 1 d")
+    x1, x2, c, s = [t.to(dtype=torch.float32) for t in [x1, x2, c, s]]
+
+    x_rot = torch.cat([x1 * c - x2 * s, x1 * s + x2 * c], axis=-1).to(x.dtype)
+
+    return torch.cat([x_rot, x_pass], axis=-1)
+
+
+def _apply_rotary_emb_kv(
+    kv: torch.FloatTensor,
+    cos: torch.FloatTensor,
+    sin: torch.FloatTensor,
+    cos_k: Optional[torch.FloatTensor] = None,
+    sin_k: Optional[torch.FloatTensor] = None,
+) -> torch.FloatTensor:
+    _, seqlen, _, _, _ = kv.shape
+    _, rotary_dim = cos.shape
+    rotary_dim *= 2
+
+    k_rot = kv[:, :, 0, :, :rotary_dim]
+    k_pass = kv[:, :, 0, :, rotary_dim:]
+
+    k1, k2 = k_rot.chunk(2, dim=-1)
+    c, s = rearrange(cos[:seqlen], "s d -> s 1 d"), rearrange(sin[:seqlen], "s d -> s 1 d")
+    k1, k2, c, s = [t.to(dtype=torch.float32) for t in [k1, k2, c, s]]
+
+    k_rot = torch.cat([k1 * c - k2 * s, k1 * s + k2 * c], axis=-1).to(kv.dtype)
+
+    return torch.cat(
+        [
+            torch.cat([k_rot, k_pass], axis=-1).unsqueeze(2),
+            kv[:, :, 1:2, :, :],
+        ],
+        axis=2,
+    )
+
+
+def _apply_rotary_emb_qkv(
+    qkv: torch.FloatTensor,
+    cos: torch.FloatTensor,
+    sin: torch.FloatTensor,
+    cos_k: Optional[torch.FloatTensor] = None,
+    sin_k: Optional[torch.FloatTensor] = None,
+) -> torch.FloatTensor:
+    _, seqlen, _, _, _ = qkv.shape
+    _, rotary_dim = cos.shape
+    rotary_dim *= 2
+
+    q_rot = qkv[:, :, 0, :, :rotary_dim]
+    q_pass = qkv[:, :, 0, :, rotary_dim:]
+
+    k_rot = qkv[:, :, 1, :, :rotary_dim]
+    k_pass = qkv[:, :, 1, :, rotary_dim:]
+
+    q1, q2 = q_rot.chunk(2, dim=-1)
+    k1, k2 = k_rot.chunk(2, dim=-1)
+    c, s = rearrange(cos[:seqlen], "s d -> s 1 d"), rearrange(sin[:seqlen], "s d -> s 1 d")
+    q1, q2, k1, k2, c, s = [t.to(dtype=torch.float32) for t in [q1, q2, k1, k2, c, s]]
+
+    q_rot = torch.cat([q1 * c - q2 * s, q1 * s + q2 * c], axis=-1).to(qkv.dtype)
+    k_rot = torch.cat([k1 * c - k2 * s, k1 * s + k2 * c], axis=-1).to(qkv.dtype)
+
+    return torch.cat(
+        [
+            torch.cat([q_rot, q_pass], axis=-1).unsqueeze(2),
+            torch.cat([k_rot, k_pass], axis=-1).unsqueeze(2),
+            qkv[:, :, 2:3, :, :],
+        ],
+        axis=2,
+    )
+
+
+class RotaryEmbedding(nn.Module):
+    """Rotary positional embedding (RoPE).
+
+    Reference:
+        RoFormer: Enhanced Transformer with Rotary Position Embedding.
+        https://arxiv.org/pdf/2104.09864.pdf.
+
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        base: int = 10000,
+        scale_base: Optional[float] = None,
+        pos_idx_in_fp32: bool = True,
+        max_position_embeddings: int = 2048,
+        device: Optional[str] = None,
+        **kwargs,
+    ) -> None:
+        super().__init__()
+
+        if scale_base is not None:
+            raise NotImplementedError
+
+        self.dim = dim
+        self.base = float(base)
+        self.scale_base = scale_base
+        self.pos_idx_in_fp32 = pos_idx_in_fp32
+        self.max_position_embeddings = max_position_embeddings
+        self.device = device
+
+        # Generate and save the inverse frequency buffer (non-trainable)
+        inv_freq = self._compute_inv_freq(device)
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+        # Generate and save the scale buffer (non-trainable)
+        scale = (
+            (torch.arange(0, dim, 2, device=device, dtype=torch.float32) + 0.4 * dim) / (1.4 * dim)
+            if scale_base is not None
+            else None
+        )
+        self.register_buffer("scale", scale, persistent=False)
+
+        # Initialize cached attributes since ONNX can't rely on dynamic initialization
+        self._update_cos_sin_cache(max_position_embeddings, device=device, dtype=torch.float32)
+
+    def _compute_inv_freq(self, device: Optional[str] = None) -> torch.FloatTensor:
+        return 1.0 / (self.base ** (torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) / self.dim))
+
+    def _update_cos_sin_cache(
+        self,
+        seqlen: int,
+        device: Optional[str] = None,
+        dtype: Optional[torch.dtype] = None,
+    ) -> None:
+        self._seq_len_cached = seqlen
+
+        # fp32 is preferred since the output of `torch.arange` can be quite large
+        # and bf16 would lose a lot of precision
+        if self.pos_idx_in_fp32:
+            t = torch.arange(seqlen, device=device, dtype=torch.float32)
+            if self.inv_freq.dtype != torch.float32:
+                inv_freq = self._compute_inv_freq(device=device)
+            else:
+                inv_freq = self.inv_freq
+        else:
+            t = torch.arange(seqlen, device=device, dtype=self.inv_freq.dtype)
+            inv_freq = self.inv_freq
+
+        # `torch.outer` is preferred since `torch.einsum` converts from fp32 to fp16 if used with AMP
+        freqs = torch.outer(t, inv_freq)
+        if self.scale is None:
+            self._cos_cached = torch.cos(freqs).to(dtype)
+            self._sin_cached = torch.sin(freqs).to(dtype)
+        else:
+            power = (
+                torch.arange(seqlen, dtype=self.scale.dtype, device=self.scale.device) - seqlen // 2
+            ) / self.scale_base
+            scale = self.scale.to(device=power.device) ** rearrange(power, "s -> s 1")
+
+            # Force the scale multiplication to happen in fp32
+            self._cos_cached = (torch.cos(freqs) * scale).to(dtype)
+            self._sin_cached = (torch.sin(freqs) * scale).to(dtype)
+            self._cos_k_cached = (torch.cos(freqs) / scale).to(dtype)
+            self._sin_k_cached = (torch.sin(freqs) / scale).to(dtype)
+
+    def forward(
+        self,
+        qkv: torch.Tensor,
+        kv: Optional[torch.Tensor] = None,
+        seqlen_offset: int = 0,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        if (
+            self._seq_len_cached < qkv.shape[1] + seqlen_offset
+            or self._cos_cached.device != qkv.device
+            or self._cos_cached.dtype != qkv.dtype
+            or (self.training and self._cos_cached.is_inference())
+        ):
+            self._update_cos_sin_cache(qkv.shape[1] + seqlen_offset, device=qkv.device, dtype=qkv.dtype)
+
+        if kv is None:
+            return _apply_rotary_emb_qkv(
+                qkv,
+                self._cos_cached[seqlen_offset:],
+                self._sin_cached[seqlen_offset:],
+            )
+        else:
+            q = _apply_rotary_emb(
+                qkv,
+                self._cos_cached[seqlen_offset:],
+                self._sin_cached[seqlen_offset:],
+            )
+            kv = _apply_rotary_emb_kv(
+                kv,
+                self._cos_cached[seqlen_offset:],
+                self._sin_cached[seqlen_offset:],
+            )
+
+            return q, kv
+
+
+class MLP(nn.Module):
+    """Multi-Layer Perceptron.
+
+    Reference:
+        Attention Is All You Need.
+        https://arxiv.org/pdf/1706.03762.pdf.
+
+    """
+
+    def __init__(
+        self,
+        config: PretrainedConfig,
+        n_inner: Optional[int] = None,
+        act_fn: Optional[str] = None,
+    ) -> None:
+        super().__init__()
+
+        act_fn = config.activation_function if act_fn is None else act_fn
+
+        n_inner = getattr(config, "n_inner", None) if n_inner is None else n_inner
+        n_inner = n_inner if n_inner is not None else 4 * config.n_embd
+
+        self.fc1 = nn.Linear(config.n_embd, n_inner)
+        self.fc2 = nn.Linear(n_inner, config.n_embd)
+        self.act = ACT2FN[act_fn]
+
+    def forward(self, hidden_states: torch.FloatTensor) -> torch.FloatTensor:
+        hidden_states = self.fc1(hidden_states)
+        hidden_states = self.act(hidden_states)
+        hidden_states = self.fc2(hidden_states)
+
+        return hidden_states
+
+
+class SelfAttention(nn.Module):
+    """Self-attention layer (compatible with PyTorch).
+
+    Reference:
+        https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/modules/mha.py.
+
+    """
+
+    def __init__(
+        self,
+        causal: bool = True,
+        softmax_scale: Optional[float] = None,
+        attention_dropout: float = 0.0,
+    ) -> None:
+        super().__init__()
+
+        self.causal = causal
+        self.softmax_scale = softmax_scale
+        self.drop = nn.Dropout(attention_dropout)
+
+    @torch.autocast("cpu", enabled=False)
+    @torch.autocast("cuda", enabled=False)
+    def forward(
+        self,
+        qkv: torch.FloatTensor,
+        causal: bool = None,
+        key_padding_mask: Optional[torch.BoolTensor] = None,
+        **kwargs,
+    ) -> torch.FloatTensor:
+        batch_size, seqlen = qkv.shape[0], qkv.shape[1]
+        q, k, v = qkv.unbind(dim=2)
+
+        q = q.to(torch.float32)
+        k = k.to(torch.float32)
+
+        causal = self.causal if causal is None else causal
+        softmax_scale = self.softmax_scale or 1.0 / math.sqrt(q.shape[-1])
+
+        # Autocast is manually disabled to avoid `torch.einsum` performing the operation
+        # using float16, which might lead to overflow
+        scores = torch.einsum("bthd,bshd->bhts", q, k * softmax_scale)
+
+        if key_padding_mask is not None:
+            padding_mask = torch.full((batch_size, seqlen), -10000.0, dtype=scores.dtype, device=scores.device)
+            padding_mask.masked_fill_(key_padding_mask, 0.0)
+
+            scores = scores + rearrange(padding_mask, "b s -> b 1 1 s")
+
+        if causal:
+            causal_mask = torch.triu(torch.full((seqlen, seqlen), -10000.0, device=scores.device), 1)
+            scores = scores + causal_mask.to(dtype=scores.dtype)
+
+        attention = torch.softmax(scores, dim=-1).to(v.dtype)
+        attention = self.drop(attention)
+
+        output = torch.einsum("bhts,bshd->bthd", attention, v)
+
+        return output
+
+
+class CrossAttention(nn.Module):
+    """Cross-attention layer (compatible with PyTorch).
+
+    Reference:
+        https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/modules/mha.py.
+
+    """
+
+    def __init__(
+        self,
+        causal: bool = True,
+        softmax_scale: Optional[float] = None,
+        attention_dropout: float = 0.0,
+    ) -> None:
+        super().__init__()
+
+        self.causal = causal
+        self.softmax_scale = softmax_scale
+        self.drop = nn.Dropout(attention_dropout)
+
+    @torch.autocast("cpu", enabled=False)
+    @torch.autocast("cuda", enabled=False)
+    def forward(
+        self,
+        q: torch.FloatTensor,
+        kv: torch.FloatTensor,
+        causal: bool = None,
+        key_padding_mask: Optional[torch.BoolTensor] = None,
+        **kwargs,
+    ) -> torch.FloatTensor:
+        batch_size, seqlen_q = q.shape[0], q.shape[1]
+        seqlen_k = kv.shape[1]
+
+        if kv.shape[3] != q.shape[2]:
+            kv = repeat(kv, "... hkv d -> ... (hkv g) d", g=q.shape[2] // kv.shape[3])
+        k, v = kv.unbind(dim=2)
+
+        q = q.to(torch.float32)
+        k = k.to(torch.float32)
+
+        causal = self.causal if causal is None else causal
+        softmax_scale = self.softmax_scale or 1.0 / math.sqrt(q.shape[-1])
+
+        # Autocast is manually disabled to avoid `torch.einsum` performing the operation
+        # using float16, which might lead to overflow
+        scores = torch.einsum("bthd,bshd->bhts", q, k * softmax_scale)
+
+        if key_padding_mask is not None:
+            padding_mask = torch.full(
+                (batch_size, seqlen_k),
+                -10000.0,
+                dtype=scores.dtype,
+                device=scores.device,
+            )
+            padding_mask.masked_fill_(key_padding_mask, 0.0)
+
+            scores = scores + rearrange(padding_mask, "b s -> b 1 1 s")
+
+        if causal:
+            rows = rearrange(torch.arange(seqlen_q, device=q.device, dtype=torch.long), "s -> s 1")
+            cols = torch.arange(seqlen_k, device=k.device, dtype=torch.long)
+            causal_mask = cols > rows + seqlen_k - seqlen_q
+
+            scores = scores.masked_fill(causal_mask, -10000.0)
+
+        attention = torch.softmax(scores, dim=-1).to(v.dtype)
+        attention = self.drop(attention)
+
+        output = torch.einsum("bhts,bshd->bthd", attention, v)
+
+        return output
+
+
+def _find_mha_dims(
+    config: PretrainedConfig,
+    n_head: Optional[int] = None,
+    n_head_kv: Optional[int] = None,
+    head_dim: Optional[int] = None,
+) -> Tuple[int, int]:
+    if n_head is None and head_dim is None:
+        head_dim = config.n_embd // config.n_head
+        n_head = config.n_head
+    elif n_head is None or head_dim is None:
+        raise ValueError("`n_head` and `head_dim` must be both specified or `None`.")
+
+    if n_head_kv is None:
+        n_head_kv = getattr(config, "n_head_kv", None) or n_head
+
+    return n_head, n_head_kv, head_dim
+
+
+def _update_kv_cache(kv: torch.FloatTensor, inference_params: InferenceParams, layer_idx: int) -> torch.FloatTensor:
+    num_heads, head_dim = kv.shape[-2:]
+
+    if layer_idx not in inference_params.key_value_memory_dict:
+        inference_params.key_value_memory_dict[layer_idx] = torch.empty(
+            inference_params.max_batch_size,
+            inference_params.max_seqlen,
+            2,
+            num_heads,
+            head_dim,
+            dtype=kv.dtype,
+            device=kv.device,
+        )
+
+    batch_start = inference_params.batch_size_offset
+    batch_end = batch_start + kv.shape[0]
+
+    sequence_start = inference_params.seqlen_offset
+    sequence_end = sequence_start + kv.shape[1]
+
+    # When the current sequence length is equal to or larger than the maximum sequence length,
+    # we need to concatenate the current `kv` with the cached `kv` to expand its length
+    if sequence_end >= inference_params.max_seqlen:
+        inference_params.key_value_memory_dict[layer_idx] = torch.concatenate((inference_params.key_value_memory_dict[layer_idx], kv), dim=1)
+
+    inference_params.key_value_memory_dict[layer_idx][batch_start:batch_end, sequence_start:sequence_end, ...] = kv
+    kv = inference_params.key_value_memory_dict[layer_idx][batch_start:batch_end, :sequence_end, ...]
+        
+    return kv
+
+
+class MHA(nn.Module):
+    """Multi-head attention layer."""
+
+    def __init__(
+        self,
+        config: PretrainedConfig,
+        dtype: Optional[torch.dtype] = None,
+        device: Optional[str] = None,
+        rotary_dim: Optional[int] = None,
+        rotary_base: float = 10000.0,
+        rotary_scale_base: Optional[float] = None,
+        n_head: Optional[int] = None,
+        n_head_kv: Optional[int] = None,
+        head_dim: Optional[int] = None,
+        bias: bool = True,
+        causal: bool = True,
+        softmax_scale: Optional[float] = None,
+        layer_idx: Optional[int] = None,
+        return_residual: bool = False,
+        checkpointing: bool = False,
+    ) -> None:
+        super().__init__()
+
+        # Rotary embedding
+        self.rotary_dim = rotary_dim if rotary_dim is not None else getattr(config, "rotary_dim", 0)
+        if self.rotary_dim > 0:
+            rotary_cls = FlashRotaryEmbedding if config.flash_rotary else RotaryEmbedding
+            if rotary_cls is None:
+                rotary_cls = RotaryEmbedding
+
+            rotary_kwargs = {}
+            if rotary_cls is RotaryEmbedding:
+                rotary_kwargs["max_position_embeddings"] = config.n_positions
+
+            self.rotary_emb = rotary_cls(
+                self.rotary_dim,
+                base=rotary_base,
+                scale_base=rotary_scale_base,
+                device=device,
+                **rotary_kwargs,
+            )
+
+        # MLP
+        self.n_head, self.n_head_kv, self.head_dim = _find_mha_dims(
+            config, n_head=n_head, n_head_kv=n_head_kv, head_dim=head_dim
+        )
+        op_size = self.head_dim * (self.n_head + 2 * self.n_head_kv)
+        hidden_size = config.n_embd
+
+        linear_cls = FusedDense if config.fused_dense else nn.Linear
+        if linear_cls is None:
+            linear_cls = nn.Linear
+
+        self.Wqkv = linear_cls(hidden_size, op_size, bias=bias, device=device, dtype=dtype)
+        self.out_proj = linear_cls(hidden_size, hidden_size, bias=bias, device=device, dtype=dtype)
+
+        # Attention
+        attn_cls = FlashSelfAttention if config.flash_attn else SelfAttention
+        if attn_cls is None:
+            attn_cls = SelfAttention
+
+        cross_attn_cls = FlashCrossAttention if config.flash_attn else CrossAttention
+        if cross_attn_cls is None:
+            cross_attn_cls = CrossAttention
+
+        self.inner_attn = attn_cls(
+            causal=causal,
+            softmax_scale=softmax_scale,
+            attention_dropout=config.attn_pdrop,
+        )
+        self.inner_cross_attn = cross_attn_cls(
+            causal=causal,
+            softmax_scale=softmax_scale,
+            attention_dropout=config.attn_pdrop,
+        )
+
+        self.flash_attn = config.flash_attn and attn_cls is FlashSelfAttention
+        self.layer_idx = layer_idx
+        self.return_residual = return_residual
+        self.checkpointing = checkpointing
+
+    def _forward_self_attn(
+        self, x: torch.FloatTensor, key_padding_mask: Optional[torch.BoolTensor]
+    ) -> torch.FloatTensor:
+        qkv = self.Wqkv(x)
+        qkv = rearrange(qkv, "... (three h d) -> ... three h d", three=3, d=self.head_dim)
+
+        if self.rotary_dim > 0:
+            qkv = self.rotary_emb(qkv)
+
+        if self.flash_attn:
+            batch_size, seqlen = qkv.shape[0], qkv.shape[1]
+
+            cu_seqlens, max_seqlen = None, None
+            if key_padding_mask is not None:
+                # If `key_padding_mask` is supplied, we need to unpad the input and retrieve
+                # the `cu_seqlens` and `max_seqlen` to be used by `flash-attn`
+                qkv, indices, cu_seqlens, max_seqlen = unpad_input(qkv, key_padding_mask)
+
+            if self.checkpointing:
+                attn_output = torch.utils.checkpoint.checkpoint(
+                    self.inner_attn, qkv, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen
+                )
+            else:
+                attn_output = self.inner_attn(qkv, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen).to(qkv.device)
+
+            # If `key_padding_mask` is supplied, we need to pad the output back to the original shape
+            return pad_input(attn_output, indices, batch_size, seqlen) if key_padding_mask is not None else attn_output
+
+        if self.checkpointing:
+            return torch.utils.checkpoint.checkpoint(self.inner_attn, qkv, key_padding_mask=key_padding_mask)
+
+        return self.inner_attn(qkv, key_padding_mask=key_padding_mask)
+
+    def _forward_cross_attn(
+        self,
+        x: torch.FloatTensor,
+        past_key_values: Optional[InferenceParams],
+        key_padding_mask: Optional[torch.BoolTensor],
+    ) -> torch.FloatTensor:
+        batch_size = x.shape[0]
+
+        qkv = self.Wqkv(x)
+
+        q = qkv[..., : self.n_head * self.head_dim]
+        q = rearrange(q, "... (h d) -> ... h d", d=self.head_dim)
+
+        kv = qkv[..., self.n_head * self.head_dim :]
+        kv = rearrange(kv, "... (two hkv d) -> ... two hkv d", two=2, d=self.head_dim)
+
+        seqlen_offset = past_key_values.seqlen_offset if past_key_values is not None else 0
+        causal = None if seqlen_offset == 0 else False
+        if self.rotary_dim > 0:
+            q, kv = self.rotary_emb(q, kv=kv, seqlen_offset=seqlen_offset)
+
+        if past_key_values is not None:
+            kv = _update_kv_cache(kv, past_key_values, self.layer_idx)
+
+        if self.flash_attn:
+            batch_size, seqlen_q = q.shape[0], q.shape[1]
+            seqlen_k = kv.shape[1]
+
+            cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k = (
+                None,
+                None,
+                None,
+                None,
+            )
+            if key_padding_mask is not None:
+                kv, _, cu_seqlens_k, max_seqlen_k = unpad_input(kv, key_padding_mask)
+
+                if seqlen_q == 1:
+                    key_padding_mask = torch.ones(batch_size, 1, device=q.device)
+                elif seqlen_q != seqlen_k:
+                    key_padding_mask = key_padding_mask[:, -seqlen_q:]
+
+                q, indices_q, cu_seqlens_q, max_seqlen_q = unpad_input(q, key_padding_mask)
+
+            if self.checkpointing:
+                attn_output = torch.utils.checkpoint.checkpoint(
+                    self.inner_cross_attn,
+                    q,
+                    kv,
+                    causal=causal,
+                    cu_seqlens=cu_seqlens_q,
+                    max_seqlen=max_seqlen_q,
+                    cu_seqlens_k=cu_seqlens_k,
+                    max_seqlen_k=max_seqlen_k,
+                )
+            else:
+                attn_output = self.inner_cross_attn(
+                    q,
+                    kv,
+                    causal=causal,
+                    cu_seqlens=cu_seqlens_q,
+                    max_seqlen=max_seqlen_q,
+                    cu_seqlens_k=cu_seqlens_k,
+                    max_seqlen_k=max_seqlen_k,
+                )
+
+            return (
+                pad_input(attn_output, indices_q, batch_size, max_seqlen_q)
+                if key_padding_mask is not None
+                else attn_output
+            )
+
+        if self.checkpointing:
+            return torch.utils.checkpoint.checkpoint(
+                self.inner_cross_attn,
+                q,
+                kv,
+                key_padding_mask=key_padding_mask,
+                causal=causal,
+            )
+
+        return self.inner_cross_attn(q, kv, key_padding_mask=key_padding_mask, causal=causal)
+
+    def forward(
+        self,
+        x: torch.FloatTensor,
+        past_key_values: Optional[InferenceParams] = None,
+        attention_mask: Optional[Union[torch.LongTensor, torch.BoolTensor]] = None,
+        **kwargs,
+    ) -> Tuple[torch.FloatTensor, torch.FloatTensor]:
+        if attention_mask is not None:
+            attention_mask = attention_mask.bool()
+        else:
+            attention_mask = None
+
+        # MHA
+        if self.n_head == self.n_head_kv:
+            if past_key_values is None:
+                # If `past_key_values` are not supplied, we run self-attention
+                attn_output = self._forward_self_attn(x, attention_mask)
+            else:
+                # If `past_key_values` are supplied, it means that we might have cached values and
+                # could take advantage of cross-attention
+                attn_output = self._forward_cross_attn(x, past_key_values, attention_mask)
+        # MQA / GQA
+        else:
+            # Regardless of `past_key_values` being supplied or not, it always use cross-attention
+            # because `q` and `kv` lengths might be different
+            attn_output = self._forward_cross_attn(x, past_key_values, attention_mask)
+
+        output = rearrange(attn_output, "... h d -> ... (h d)")
+        output = self.out_proj(output)
+
+        return output if not self.return_residual else (output, x)
+
+
+class ParallelBlock(nn.Module):
+    """Parallel block.
+
+    This block applies parallel mixer and MLP layers to the input (used in GPT-J and CodeGen).
+
+    """
+
+    def __init__(
+        self,
+        config: PretrainedConfig,
+        block_idx: Optional[int] = None,
+    ) -> None:
+        super().__init__()
+
+        self.ln = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
+        self.resid_dropout = nn.Dropout(config.resid_pdrop)
+        self.block_idx = block_idx
+
+        self.mixer = MHA(config, layer_idx=block_idx)
+        self.mlp = MLP(config)
+
+    def forward(
+        self,
+        hidden_states: torch.FloatTensor,
+        past_key_values: Optional[Union[torch.FloatTensor, InferenceParams]] = None,
+        attention_mask: Optional[torch.BoolTensor] = None,
+        **kwargs,
+    ) -> torch.FloatTensor:
+        residual = hidden_states
+        hidden_states = self.ln(hidden_states)
+
+        attn_outputs = self.mixer(
+            hidden_states,
+            past_key_values=past_key_values,
+            attention_mask=attention_mask,
+        )
+        if isinstance(attn_outputs, tuple):
+            attn_outputs = attn_outputs[0]
+
+        attn_outputs = self.resid_dropout(attn_outputs)
+        feed_forward_hidden_states = self.resid_dropout(self.mlp(hidden_states))
+
+        hidden_states = attn_outputs + feed_forward_hidden_states + residual
+
+        return hidden_states
+
+
+class CausalLMHead(nn.Module):
+    """Causal Language Modeling head.
+
+    Reference:
+        Improving Language Understanding by Generative Pre-Training.
+        https://cdn.openai.com/research-covers/language-unsupervised/language_understanding_paper.pdf.
+
+    """
+
+    def __init__(self, config: PretrainedConfig) -> None:
+        super().__init__()
+
+        self.ln = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
+        self.linear = nn.Linear(config.n_embd, config.vocab_size)
+
+    def forward(self, hidden_states: torch.FloatTensor) -> torch.FloatTensor:
+        hidden_states = self.ln(hidden_states)
+        logits = self.linear(hidden_states).to(torch.float32)
+
+        return logits
+
+
+class PhiPreTrainedModel(PreTrainedModel):
+    """Phi pre-trained model."""
+
+    config_class = PhiConfig
+    base_model_prefix = "transformer"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["ParallelBlock", "CLIPEncoderLayer", "Block"]
+
+    def __init__(self, *inputs, **kwargs) -> None:
+        super().__init__(*inputs, **kwargs)
+
+    def _init_weights(self, module: nn.Module) -> None:
+        if isinstance(module, (nn.Linear,)):
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+        elif isinstance(module, nn.LayerNorm):
+            if module.bias is not None:
+                module.bias.data.zero_()
+            module.weight.data.fill_(1.0)
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids: torch.LongTensor,
+        past_key_values: Optional[Union[torch.FloatTensor, InferenceParams]] = None,
+        attention_mask: Optional[Union[torch.LongTensor, torch.BoolTensor]] = None,
+        **kwargs,
+    ) -> Dict[str, Any]:
+        if past_key_values is None or not (isinstance(past_key_values, InferenceParams)):
+            past_key_values = InferenceParams(
+                max_seqlen=self.config.n_positions,
+                max_batch_size=input_ids.shape[0],
+                seqlen_offset=0,
+                batch_size_offset=0,
+                key_value_memory_dict={},
+                lengths_per_sample=None,
+            )
+        else:
+            # ======================================================================
+            # Assume that `past_key_values` has cached all tokens up to the last token in `input_ids`
+            # inference_params.key_value_memory_dict[layer_idx][batch_start:batch_end, sequence_start:sequence_end, ...]
+            # past_key_values.seqlen_offset = input_ids.shape[1] - 1
+            # ======================================================================
+            # I change the way of updating `past_key_values.seqlen_offset` to make the inference of imp work.
+            # [Edited by zhenwei - 2024-01-20 21:15]
+            input_ids = input_ids[:, -1].unsqueeze(-1)
+
+        return {
+            "input_ids": input_ids,
+            "past_key_values": past_key_values,
+            "attention_mask": attention_mask,
+        }
+
+
+class LlavaMetaModel(ABC):
+    """
+    Define the APIs for building components that are related to image perceiving.
+    This implementation is based on the implementation from the Llave project.
+    """
+
+    def get_vision_tower(self):
+        vision_tower = getattr(self, 'vision_tower', None)
+        if type(vision_tower) is list:
+            vision_tower = vision_tower[0]
+        return vision_tower
+    
+    def build_vision_tower(self, config):
+        self.vision_tower = VisionTower(config.vision_tower_cfg)
+
+    def build_vision_projector(self, config):
+        projector_type = getattr(config, 'mm_projector_type', 'linear')
+
+        if projector_type == 'linear':
+            self.mm_projector = nn.Linear(config.mm_hidden_size, config.hidden_size)
+            return
+
+        mlp_gelu_match = re.match(r'^mlp(\d+)x_gelu$', projector_type)
+        if mlp_gelu_match:
+            mlp_depth = int(mlp_gelu_match.group(1))
+            modules = [nn.Linear(config.mm_hidden_size, config.hidden_size)]
+            for _ in range(1, mlp_depth):
+                modules.append(nn.GELU())
+                modules.append(nn.Linear(config.hidden_size, config.hidden_size))
+            self.mm_projector = nn.Sequential(*modules)
+            return
+
+        if projector_type == 'identity':
+            self.mm_projector = nn.Identity()
+            return
+
+        raise ValueError(f'Unknown projector type: {projector_type}')
+
+
+class ImpModel(PhiPreTrainedModel, LlavaMetaModel):
+    """Imp model. This implementation is modified from the implementation of Phi-2"""
+
+    config_class = ImpConfig
+    # _keys_to_ignore_on_load_missing = [""]
+    # _keys_to_ignore_on_load_unexpected = [r"h\.\d+\.mlp.(fc_in|fc_out)\.(weight|bias)"]
+
+    def __init__(self, config: ImpConfig) -> None:
+        super().__init__(config)
+
+        self.embd = Embedding(config)
+        self.h = nn.ModuleList([ParallelBlock(config, block_idx=i) for i in range(config.n_layer)])
+        self.gradient_checkpointing = False
+
+        if hasattr(config, "mm_vision_tower"):
+            self.build_vision_tower(config)
+            self.build_vision_projector(config)
+
+        self.post_init()
+
+    def embed_tokens(self, input_ids: torch.LongTensor) -> torch.FloatTensor:
+        return self.embd(input_ids)[0]
+    
+    def get_input_embeddings(self) -> nn.Embedding:
+        return self.embd.wte
+    
+    def set_input_embeddings(self, new_embeddings: nn.Embedding) -> None:
+        self.embd.wte = new_embeddings
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor,
+        past_key_values: Optional[Union[torch.FloatTensor, InferenceParams]] = None,
+        attention_mask: Optional[torch.BoolTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None
+    ) -> torch.FloatTensor:
+
+        if inputs_embeds is None:
+            hidden_states = self.embd(input_ids)
+        else:
+            hidden_states = inputs_embeds
+
+        for layer in self.h:
+            if self.gradient_checkpointing and self.training:
+
+                def create_custom_forward(module):
+                    def custom_forward(*inputs):
+                        # None for past_key_value
+                        return module(*inputs)
+
+                    return custom_forward
+
+                hidden_states = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(layer),
+                    hidden_states,
+                    None,
+                    attention_mask,
+                )
+            else:
+                hidden_states = layer(
+                    hidden_states,
+                    past_key_values=past_key_values,
+                    attention_mask=attention_mask,
+                )
+
+        # I change the way of updating `past_key_values.seqlen_offset` to make the inference of imp work.
+        # [Edited by zhenwei - 2024-01-20 21:15]
+        if past_key_values is not None: # FIXME: when multi-batch inference, it is a bug
+            past_key_values.seqlen_offset += hidden_states.shape[1]
+
+        return hidden_states
+
+
+class LlavaMetaForCausalLM(ABC):
+    """This implementation is based on the implementation from the Llave project."""
+
+    def init_constants(self, config):
+        self.IGNORE_INDEX = getattr(config, 'ignore_index', -100)
+        self.IMAGE_TOKEN_INDEX = getattr(config, 'image_token_index', 50296)
+        self.DEFAULT_IMAGE_TOKEN = getattr(config, 'image_token', "<image>")
+
+    @abstractmethod
+    def get_model(self):
+        pass
+
+    def get_vision_tower(self):
+        return self.get_model().get_vision_tower()
+
+    def encode_images(self, images):
+        image_features = self.get_model().get_vision_tower()(images)
+        image_features = self.get_model().mm_projector(image_features)
+        return image_features
+
+    def prepare_inputs_labels_for_multimodal(
+        self, input_ids, position_ids, attention_mask, past_key_values, labels, images
+    ):
+        vision_tower = self.get_vision_tower()
+        # if vision_tower is None or images is None or past_key_values.seqlen_offset != 0:
+        if vision_tower is None or images is None or input_ids.shape[1] == 1:
+            if past_key_values is not None and vision_tower is not None and images is not None and input_ids.shape[1] == 1:
+                target_shape = past_key_values.seqlen_offset + 1
+                # inference_params.key_value_memory_dict[layer_idx][batch_start:batch_end, sequence_start:sequence_end, ...]
+                attention_mask = torch.cat((attention_mask, torch.ones(
+                    (attention_mask.shape[0], target_shape - attention_mask.shape[1]),
+                    dtype=attention_mask.dtype,
+                    device=attention_mask.device
+                )), dim=1)
+                position_ids = torch.sum(attention_mask, dim=1).unsqueeze(-1) - 1
+            return input_ids, position_ids, attention_mask, past_key_values, None, labels
+
+        if type(images) is list or images.ndim == 5:
+            concat_images = torch.cat([image for image in images], dim=0)
+            concat_images = concat_images.to(device=self.device, dtype=vision_tower.dtype)
+            image_features = self.encode_images(concat_images)
+            split_sizes = [image.shape[0] for image in images]
+            image_features = torch.split(image_features, split_sizes, dim=0)
+            image_features = [x.flatten(0, 1).to(self.device) for x in image_features]
+        else:
+            images = images.to(device=self.device, dtype=vision_tower.dtype)
+            image_features = self.encode_images(images).to(self.device)
+
+        # TODO: image start / end is not implemented here to support pretraining.
+        if getattr(self.config, 'tune_mm_mlp_adapter', False) and getattr(self.config, 'mm_use_im_start_end', False):
+            raise NotImplementedError
+
+        # Let's just add dummy tensors if they do not exist,
+        # it is a headache to deal with None all the time.
+        # But it is not ideal, and if you have a better idea,
+        # please open an issue / submit a PR, thanks.
+        _labels = labels
+        _position_ids = position_ids
+        _attention_mask = attention_mask
+        if attention_mask is None:
+            attention_mask = torch.ones_like(input_ids, dtype=torch.bool)
+        else:
+            attention_mask = attention_mask.bool()
+        if position_ids is None:
+            position_ids = torch.arange(0, input_ids.shape[1], dtype=torch.long, device=input_ids.device)
+        if labels is None:
+            labels = torch.full_like(input_ids, self.IGNORE_INDEX)
+
+        # remove the padding using attention_mask -- TODO: double check
+        input_ids = [cur_input_ids[cur_attention_mask] for cur_input_ids, cur_attention_mask in zip(input_ids, attention_mask)]
+        labels = [cur_labels[cur_attention_mask] for cur_labels, cur_attention_mask in zip(labels, attention_mask)]
+
+        new_input_embeds = []
+        new_labels = []
+        cur_image_idx = 0
+        for batch_idx, cur_input_ids in enumerate(input_ids):
+            num_images = (cur_input_ids == self.IMAGE_TOKEN_INDEX).sum()
+            if num_images == 0:
+                cur_image_features = image_features[cur_image_idx]
+                cur_input_embeds_1 = self.get_model().embed_tokens(cur_input_ids)
+                cur_input_embeds = torch.cat([cur_input_embeds_1, cur_image_features[0:0]], dim=0)
+                new_input_embeds.append(cur_input_embeds)
+                new_labels.append(labels[batch_idx])
+                cur_image_idx += 1
+                continue
+
+            image_token_indices = [-1] + torch.where(cur_input_ids == self.IMAGE_TOKEN_INDEX)[0].tolist() + [cur_input_ids.shape[0]]
+            cur_input_ids_noim = []
+            cur_labels = labels[batch_idx]
+            cur_labels_noim = []
+            for i in range(len(image_token_indices) - 1):
+                cur_input_ids_noim.append(cur_input_ids[image_token_indices[i]+1:image_token_indices[i+1]])
+                cur_labels_noim.append(cur_labels[image_token_indices[i]+1:image_token_indices[i+1]])
+            split_sizes = [x.shape[0] for x in cur_labels_noim]
+            cur_input_embeds = self.get_model().embed_tokens(torch.cat(cur_input_ids_noim))
+            # print(cur_input_embeds.shape)
+            cur_input_embeds_no_im = torch.split(cur_input_embeds, split_sizes, dim=0)
+            cur_new_input_embeds = []
+            cur_new_labels = []
+
+            for i in range(num_images + 1):
+                cur_new_input_embeds.append(cur_input_embeds_no_im[i])
+                cur_new_labels.append(cur_labels_noim[i])
+                if i < num_images:
+                    cur_image_features = image_features[cur_image_idx]
+                    cur_image_idx += 1
+                    cur_new_input_embeds.append(cur_image_features)
+                    cur_new_labels.append(torch.full((cur_image_features.shape[0],), self.IGNORE_INDEX, device=cur_labels.device, dtype=cur_labels.dtype))
+
+            cur_new_input_embeds = torch.cat(cur_new_input_embeds)
+            cur_new_labels = torch.cat(cur_new_labels)
+
+            new_input_embeds.append(cur_new_input_embeds)
+            new_labels.append(cur_new_labels)
+
+        # Truncate sequences to max length as image embeddings can make the sequence longer
+        tokenizer_model_max_length = getattr(self.config, 'tokenizer_model_max_length', None)
+        if tokenizer_model_max_length is not None:
+            new_input_embeds = [x[:tokenizer_model_max_length] for x in new_input_embeds]
+            new_labels = [x[:tokenizer_model_max_length] for x in new_labels]
+
+        # Combine them
+        max_len = max(x.shape[0] for x in new_input_embeds)
+        batch_size = len(new_input_embeds)
+
+        new_input_embeds_padded = []
+        new_labels_padded = torch.full((batch_size, max_len), self.IGNORE_INDEX, dtype=new_labels[0].dtype, device=new_labels[0].device)
+        attention_mask = torch.zeros((batch_size, max_len), dtype=attention_mask.dtype, device=attention_mask.device)
+        position_ids = torch.zeros((batch_size, max_len), dtype=position_ids.dtype, device=position_ids.device)
+
+        for i, (cur_new_embed, cur_new_labels) in enumerate(zip(new_input_embeds, new_labels)):
+            cur_len = cur_new_embed.shape[0]
+            if getattr(self.config, 'tokenizer_padding_side', 'right') == "left":
+                new_input_embeds_padded.append(torch.cat((
+                    torch.zeros((max_len - cur_len, cur_new_embed.shape[1]), dtype=cur_new_embed.dtype, device=cur_new_embed.device),
+                    cur_new_embed
+                ), dim=0))
+                if cur_len > 0:
+                    new_labels_padded[i, -cur_len:] = cur_new_labels
+                    attention_mask[i, -cur_len:] = True
+                    position_ids[i, -cur_len:] = torch.arange(0, cur_len, dtype=position_ids.dtype, device=position_ids.device)
+            else:
+                new_input_embeds_padded.append(torch.cat((
+                    cur_new_embed,
+                    torch.zeros((max_len - cur_len, cur_new_embed.shape[1]), dtype=cur_new_embed.dtype, device=cur_new_embed.device)
+                ), dim=0))
+                if cur_len > 0:
+                    new_labels_padded[i, :cur_len] = cur_new_labels
+                    attention_mask[i, :cur_len] = True
+                    position_ids[i, :cur_len] = torch.arange(0, cur_len, dtype=position_ids.dtype, device=position_ids.device)
+
+        new_input_embeds = torch.stack(new_input_embeds_padded, dim=0)
+
+        if _labels is None:
+            new_labels = None
+        else:
+            new_labels = new_labels_padded
+
+        if _attention_mask is None:
+            attention_mask = None
+        else:
+            attention_mask = attention_mask.to(dtype=_attention_mask.dtype)
+
+        if _position_ids is None:
+            position_ids = None
+
+        return None, position_ids, attention_mask, past_key_values, new_input_embeds, new_labels
+
+
+class ImpForCausalLM(PhiPreTrainedModel, LlavaMetaForCausalLM):
+    """Imp for Causal Language Modeling."""
+
+    # _keys_to_ignore_on_load_missing = [""]
+    # _keys_to_ignore_on_load_unexpected = [r"transformer\.h\.\d+\.mlp.(fc_in|fc_out)\.(weight|bias)"]
+    config_class = ImpConfig
+
+    def __init__(self, config: ImpConfig) -> None:
+        super().__init__(config)
+
+        self.transformer = ImpModel(config)
+        self.lm_head = CausalLMHead(config)
+        
+        self.post_init()
+        self.init_constants(config)
+
+    def get_output_embeddings(self) -> nn.Linear:
+        return self.lm_head.linear
+
+    def set_output_embeddings(self, new_embeddings: nn.Linear) -> None:
+        self.lm_head.linear = new_embeddings
+
+    def get_model(self):
+        return self.transformer
+    
+    def image_preprocess(self, images):
+        return self.get_vision_tower().image_processor(images)['pixel_values']
+    
+    def backbone_forward(
+        self,
+        input_ids: torch.LongTensor,
+        past_key_values: Optional[Union[torch.FloatTensor, InferenceParams]] = None,
+        attention_mask: Optional[torch.BoolTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        **kwargs,
+    ) -> CausalLMOutputWithPast:
+        hidden_states = self.transformer(input_ids, past_key_values=past_key_values, attention_mask=attention_mask, inputs_embeds=inputs_embeds)
+        lm_logits = self.lm_head(hidden_states)
+
+        return CausalLMOutputWithPast(loss=None, logits=lm_logits, past_key_values=past_key_values)
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = 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,
+        images: Optional[torch.FloatTensor] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+
+        if inputs_embeds is None:
+            (
+                input_ids,
+                position_ids,
+                attention_mask,
+                past_key_values,
+                inputs_embeds,
+                labels
+            ) = self.prepare_inputs_labels_for_multimodal(
+                input_ids,
+                position_ids,
+                attention_mask,
+                past_key_values,
+                labels,
+                images
+            )
+
+        return self.backbone_forward(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            labels=labels,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict
+        )
+
+    def prepare_inputs_for_generation(self, input_ids, past_key_values=None, inputs_embeds=None, **kwargs):
+        images = kwargs.pop("images", None)
+        _inputs = super().prepare_inputs_for_generation(
+            input_ids, past_key_values=past_key_values, inputs_embeds=inputs_embeds, **kwargs
+        )
+        if images is not None:
+            _inputs['images'] = images
+        return _inputs
+
+
+AutoConfig.register("imp", ImpConfig)
+AutoModelForCausalLM.register(ImpConfig, ImpForCausalLM)

pytorch_model.bin

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

vision_encoder.py

@@ -0,0 +1,593 @@
+# Copyright (c) MILVLG team.
+# Licensed under the Apache 2.0 license.
+#
+# Some code here is copied from the project Phi-2 (https://huggingface.co/microsoft/phi-2),
+# SigLIP@transformers==4.37.0.dev0 (https://huggingface.co/google/siglip-so400m-patch14-384),
+# and Llava (https://github.com/haotian-liu/LLaVA), and modified by 
+# Zhenwei Shao ([email protected]) @ MILVLG. We thank them for their great works.
+# And their original licenses and copyright should be inherited (see the statements
+# in `configuration_imp.py` for more details).
+
+
+from typing import Any, Optional, Tuple, Union, List, Dict
+from dataclasses import dataclass
+import math
+import warnings
+from functools import partial, reduce
+
+
+import numpy as np
+from PIL import Image
+import torch
+import torch.utils.checkpoint
+from torch import nn
+
+from transformers.image_processing_utils import BatchFeature
+from transformers.image_transforms import (
+    convert_to_rgb,
+    normalize,
+    rescale,
+    resize,
+    to_channel_dimension_format,
+)
+from transformers.image_utils import (
+    ChannelDimension,
+    PILImageResampling,
+    to_numpy_array,
+)
+from transformers.activations import ACT2FN
+from transformers.modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling
+from transformers.modeling_utils import PreTrainedModel
+from transformers.utils import ModelOutput
+
+from .configuration_imp import SiglipVisionConfig
+
+
+# ============================================================================
+# A simple image preprocessor for SigLIP models. 
+# ============================================================================
+
+def simple_image_processor(
+        images, 
+        image_mean=(0.5, 0.5, 0.5), 
+        image_std=(0.5, 0.5, 0.5), 
+        size=(384, 384), 
+        resample=PILImageResampling.BICUBIC, 
+        rescale_factor=1 / 255, 
+        data_format=ChannelDimension.FIRST,
+        return_tensors="pt"
+    ):
+
+    if isinstance(images, Image.Image):
+        images = [images]
+    else:
+        assert isinstance(images, list)
+    
+    transforms = [
+        convert_to_rgb,
+        to_numpy_array,
+        partial(resize, size=size, resample=resample, data_format=data_format),
+        partial(rescale, scale=rescale_factor, data_format=data_format),
+        partial(normalize, mean=image_mean, std=image_std, data_format=data_format),
+        partial(to_channel_dimension_format, channel_dim=data_format, input_channel_dim=data_format),
+    ]
+
+    images = reduce(lambda x, f: [*map(f, x)], transforms, images)
+    data = {"pixel_values": images}
+    
+    return BatchFeature(data=data, tensor_type=return_tensors)
+
+# ============================================================================
+# Definitions for SigLIP models. 
+# ============================================================================
+
+@dataclass
+# Copied from transformers.models.clip.modeling_clip.CLIPVisionModelOutput with CLIP->Siglip
+class SiglipVisionModelOutput(ModelOutput):
+    """
+    Base class for vision model's outputs that also contains image embeddings of the pooling of the last hidden states.
+
+    Args:
+        image_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim)` *optional* returned when model is initialized with `with_projection=True`):
+            The image embeddings obtained by applying the projection layer to the pooler_output.
+        last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
+            Sequence of hidden-states at the output of the last layer of the model.
+        hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
+            Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +
+            one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`.
+
+            Hidden-states of the model at the output of each layer plus the optional initial embedding outputs.
+        attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
+            Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
+            sequence_length)`.
+
+            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
+            heads.
+    """
+
+    image_embeds: Optional[torch.FloatTensor] = None
+    last_hidden_state: torch.FloatTensor = None
+    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+    attentions: Optional[Tuple[torch.FloatTensor]] = None
+
+
+class SiglipVisionEmbeddings(nn.Module):
+    def __init__(self, config: SiglipVisionConfig):
+        super().__init__()
+        self.config = config
+        self.embed_dim = config.hidden_size
+        self.image_size = config.image_size
+        self.patch_size = config.patch_size
+
+        self.patch_embedding = nn.Conv2d(
+            in_channels=config.num_channels,
+            out_channels=self.embed_dim,
+            kernel_size=self.patch_size,
+            stride=self.patch_size,
+            padding="valid",
+        )
+
+        self.num_patches = (self.image_size // self.patch_size) ** 2
+        self.num_positions = self.num_patches
+        self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim)
+        self.register_buffer("position_ids", torch.arange(self.num_positions).expand((1, -1)), persistent=False)
+
+    def forward(self, pixel_values: torch.FloatTensor) -> torch.Tensor:
+        patch_embeds = self.patch_embedding(pixel_values)  # shape = [*, width, grid, grid]
+        embeddings = patch_embeds.flatten(2).transpose(1, 2)
+
+        embeddings = embeddings + self.position_embedding(self.position_ids)
+        return embeddings
+
+
+
+class SiglipAttention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    # Copied from transformers.models.clip.modeling_clip.CLIPAttention.__init__
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.embed_dim = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.head_dim = self.embed_dim // self.num_heads
+        if self.head_dim * self.num_heads != self.embed_dim:
+            raise ValueError(
+                f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:"
+                f" {self.num_heads})."
+            )
+        self.scale = self.head_dim**-0.5
+        self.dropout = config.attention_dropout
+
+        self.k_proj = nn.Linear(self.embed_dim, self.embed_dim)
+        self.v_proj = nn.Linear(self.embed_dim, self.embed_dim)
+        self.q_proj = nn.Linear(self.embed_dim, self.embed_dim)
+        self.out_proj = nn.Linear(self.embed_dim, self.embed_dim)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        output_attentions: Optional[bool] = False,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        """Input shape: Batch x Time x Channel"""
+
+        batch_size, q_len, _ = hidden_states.size()
+
+        query_states = self.q_proj(hidden_states)
+        key_states = self.k_proj(hidden_states)
+        value_states = self.v_proj(hidden_states)
+
+        query_states = query_states.view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+        key_states = key_states.view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+        value_states = value_states.view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+
+        k_v_seq_len = key_states.shape[-2]
+        attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) * self.scale
+
+        if attn_weights.size() != (batch_size, self.num_heads, q_len, k_v_seq_len):
+            raise ValueError(
+                f"Attention weights should be of size {(batch_size, self.num_heads, q_len, k_v_seq_len)}, but is"
+                f" {attn_weights.size()}"
+            )
+
+        if attention_mask is not None:
+            if attention_mask.size() != (batch_size, 1, q_len, k_v_seq_len):
+                raise ValueError(
+                    f"Attention mask should be of size {(batch_size, 1, q_len, k_v_seq_len)}, but is {attention_mask.size()}"
+                )
+            attn_weights = attn_weights + attention_mask
+
+        # upcast attention to fp32
+        attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
+        attn_weights = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training)
+        attn_output = torch.matmul(attn_weights, value_states)
+
+        if attn_output.size() != (batch_size, self.num_heads, q_len, self.head_dim):
+            raise ValueError(
+                f"`attn_output` should be of size {(batch_size, self.num_heads, q_len, self.head_dim)}, but is"
+                f" {attn_output.size()}"
+            )
+
+        attn_output = attn_output.transpose(1, 2).contiguous()
+        attn_output = attn_output.reshape(batch_size, q_len, self.embed_dim)
+
+        attn_output = self.out_proj(attn_output)
+
+        return attn_output, attn_weights
+
+
+# Copied from transformers.models.clip.modeling_clip.CLIPMLP with CLIP->Siglip
+class SiglipMLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.activation_fn = ACT2FN[config.hidden_act]
+        self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size)
+        self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.fc1(hidden_states)
+        hidden_states = self.activation_fn(hidden_states)
+        hidden_states = self.fc2(hidden_states)
+        return hidden_states
+
+
+# Copied from transformers.models.clip.modeling_clip.CLIPEncoderLayer with CLIP->Siglip
+class SiglipEncoderLayer(nn.Module):
+    def __init__(self, config: SiglipVisionConfig):
+        super().__init__()
+        self.embed_dim = config.hidden_size
+        self.self_attn = SiglipAttention(config)
+        self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
+        self.mlp = SiglipMLP(config)
+        self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
+
+    # Ignore copy
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor,
+        output_attentions: Optional[bool] = False,
+    ) -> Tuple[torch.FloatTensor]:
+        """
+        Args:
+            hidden_states (`torch.FloatTensor`):
+                Input to the layer of shape `(batch, seq_len, embed_dim)`.
+            attention_mask (`torch.FloatTensor`):
+                Attention mask of shape `(batch, 1, q_len, k_v_seq_len)` where padding elements are indicated by very large negative values.
+            output_attentions (`bool`, *optional*, defaults to `False`):
+                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
+                returned tensors for more detail.
+        """
+        residual = hidden_states
+
+        hidden_states = self.layer_norm1(hidden_states)
+        hidden_states, attn_weights = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            output_attentions=output_attentions,
+        )
+        hidden_states = residual + hidden_states
+
+        residual = hidden_states
+        hidden_states = self.layer_norm2(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+
+        outputs = (hidden_states,)
+
+        if output_attentions:
+            outputs += (attn_weights,)
+
+        return outputs
+
+
+class SiglipPreTrainedModel(PreTrainedModel):
+    """
+    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
+    models.
+    """
+
+    config_class = SiglipVisionConfig
+    base_model_prefix = "siglip"
+    supports_gradient_checkpointing = True
+
+    def _init_weights(self, module):
+        """Initialize the weights"""
+        pass
+
+# Copied from transformers.models.clip.modeling_clip.CLIPEncoder with CLIP->Siglip
+class SiglipEncoder(nn.Module):
+    """
+    Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a
+    [`SiglipEncoderLayer`].
+
+    Args:
+        config: SiglipVisionConfig
+    """
+
+    def __init__(self, config: SiglipVisionConfig):
+        super().__init__()
+        self.config = config
+        self.layers = nn.ModuleList([SiglipEncoderLayer(config) for _ in range(config.num_hidden_layers)])
+        self.gradient_checkpointing = False
+
+    # Ignore copy
+    def forward(
+        self,
+        inputs_embeds,
+        attention_mask: Optional[torch.Tensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, BaseModelOutput]:
+        r"""
+        Args:
+            inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
+                Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation.
+                This is useful if you want more control over how to convert `input_ids` indices into associated vectors
+                than the model's internal embedding lookup matrix.
+            attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
+                Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
+
+                - 1 for tokens that are **not masked**,
+                - 0 for tokens that are **masked**.
+
+                [What are attention masks?](../glossary#attention-mask)
+            output_attentions (`bool`, *optional*):
+                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
+                returned tensors for more detail.
+            output_hidden_states (`bool`, *optional*):
+                Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors
+                for more detail.
+            return_dict (`bool`, *optional*):
+                Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
+        """
+        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
+
+        encoder_states = () if output_hidden_states else None
+        all_attentions = () if output_attentions else None
+
+        hidden_states = inputs_embeds
+        for encoder_layer in self.layers:
+            if output_hidden_states:
+                encoder_states = encoder_states + (hidden_states,)
+            if self.gradient_checkpointing and self.training:
+                layer_outputs = self._gradient_checkpointing_func(
+                    encoder_layer.__call__,
+                    hidden_states,
+                    attention_mask,
+                    output_attentions,
+                )
+            else:
+                layer_outputs = encoder_layer(
+                    hidden_states,
+                    attention_mask,
+                    output_attentions=output_attentions,
+                )
+
+            hidden_states = layer_outputs[0]
+
+            if output_attentions:
+                all_attentions = all_attentions + (layer_outputs[1],)
+
+        if output_hidden_states:
+            encoder_states = encoder_states + (hidden_states,)
+
+        if not return_dict:
+            return tuple(v for v in [hidden_states, encoder_states, all_attentions] if v is not None)
+        return BaseModelOutput(
+            last_hidden_state=hidden_states, hidden_states=encoder_states, attentions=all_attentions
+        )
+
+
+class SiglipVisionTransformer(nn.Module):
+    def __init__(self, config: SiglipVisionConfig):
+        super().__init__()
+        self.config = config
+        embed_dim = config.hidden_size
+
+        self.embeddings = SiglipVisionEmbeddings(config)
+        self.encoder = SiglipEncoder(config)
+        self.post_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
+        self.head = SiglipMultiheadAttentionPoolingHead(config)
+
+    def forward(
+        self,
+        pixel_values,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, BaseModelOutputWithPooling]:
+        r"""
+        Returns:
+
+        """
+        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
+
+        hidden_states = self.embeddings(pixel_values)
+
+        encoder_outputs = self.encoder(
+            inputs_embeds=hidden_states,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        last_hidden_state = encoder_outputs[0]
+        last_hidden_state = self.post_layernorm(last_hidden_state)
+
+        pooled_output = self.head(last_hidden_state)
+
+        if not return_dict:
+            return (last_hidden_state, pooled_output) + encoder_outputs[1:]
+
+        return BaseModelOutputWithPooling(
+            last_hidden_state=last_hidden_state,
+            pooler_output=pooled_output,
+            hidden_states=encoder_outputs.hidden_states,
+            attentions=encoder_outputs.attentions,
+        )
+
+
+class SiglipMultiheadAttentionPoolingHead(nn.Module):
+    """Multihead Attention Pooling."""
+
+    def __init__(self, config: SiglipVisionConfig):
+        super().__init__()
+
+        self.probe = nn.Parameter(torch.randn(1, 1, config.hidden_size))
+        self.attention = torch.nn.MultiheadAttention(config.hidden_size, config.num_attention_heads, batch_first=True)
+        self.layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.mlp = SiglipMLP(config)
+
+    def forward(self, hidden_state):
+        batch_size = hidden_state.shape[0]
+        probe = self.probe.repeat(batch_size, 1, 1)
+
+        hidden_state = self.attention(probe, hidden_state, hidden_state)[0]
+
+        residual = hidden_state
+        hidden_state = self.layernorm(hidden_state)
+        hidden_state = residual + self.mlp(hidden_state)
+
+        return hidden_state[:, 0]
+
+
+class SiglipVisionModel(SiglipPreTrainedModel):
+    config_class = SiglipVisionConfig
+    main_input_name = "pixel_values"
+    _no_split_modules = ["SiglipEncoderLayer"]
+
+    def __init__(self, config: SiglipVisionConfig):
+        super().__init__(config)
+
+        self.vision_model = SiglipVisionTransformer(config)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self) -> nn.Module:
+        return self.vision_model.embeddings.patch_embedding
+
+    def forward(
+        self,
+        pixel_values,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, BaseModelOutputWithPooling]:
+        r"""
+        Returns:
+
+        Examples:
+
+        ```python
+        >>> from PIL import Image
+        >>> import requests
+        >>> from transformers import AutoProcessor, SiglipVisionModel
+
+        >>> model = SiglipVisionModel.from_pretrained("google/siglip-base-patch16-224")
+        >>> processor = AutoProcessor.from_pretrained("google/siglip-base-patch16-224")
+
+        >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+        >>> image = Image.open(requests.get(url, stream=True).raw)
+
+        >>> inputs = processor(images=image, return_tensors="pt")
+
+        >>> outputs = model(**inputs)
+        >>> last_hidden_state = outputs.last_hidden_state
+        >>> pooled_output = outputs.pooler_output  # pooled features
+        ```"""
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        return self.vision_model(
+            pixel_values=pixel_values,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+
+# ============================================================================
+# VisionTower module for Imp
+# ============================================================================
+
+class VisionTower(nn.Module):
+    def __init__(self, vision_tower_cfg, delay_load=False):
+        super().__init__()
+
+        self.is_loaded = False
+
+        self.config = vision_tower_cfg
+        self.vision_tower_name = vision_tower_cfg.mm_vision_tower
+        self.select_layer = vision_tower_cfg.mm_vision_select_layer
+        # self.select_feature = getattr(vision_tower_cfg, 'mm_vision_select_feature', 'patch')
+
+        self.image_processor = simple_image_processor
+
+        if not delay_load:
+            self.load_model()
+        else:
+            raise NotImplementedError("delay load is not implemented yet.")
+
+    def load_model(self):
+        if self.is_loaded:
+            return
+
+        # "google/siglip-so400m-patch14-384"
+        # self.vision_tower = SiglipVisionModel.from_pretrained(self.vision_tower_name)
+        self.vision_tower = SiglipVisionModel(self.config)
+        del self.vision_tower.vision_model.encoder.layers[(self.select_layer + 1):]
+        self.vision_tower.vision_model.head = nn.Identity()
+        self.vision_tower.requires_grad_(False)
+        self.vision_tower.eval()
+
+        self.is_loaded = True
+
+    @torch.no_grad()
+    def forward(self, images):
+        if type(images) is list:
+            image_features = []
+            for image in images:
+                image_forward_out = self.vision_tower(image.to(device=self.device, dtype=self.dtype).unsqueeze(0), output_hidden_states=True)
+                image_feature = image_forward_out.hidden_states[-1].to(image.dtype)
+                assert image_features.shape[-2] == 729
+                image_features.append(image_feature)
+        else:
+            image_forward_outs = self.vision_tower(images.to(device=self.device, dtype=self.dtype), output_hidden_states=True)
+            image_features = image_forward_outs.hidden_states[-1].to(images.dtype)
+            assert image_features.shape[-2] == 729
+
+        return image_features
+
+    @property
+    def dummy_feature(self):
+        return torch.zeros(1, self.hidden_size, device=self.device, dtype=self.dtype)
+
+    @property
+    def dtype(self):
+        for p in self.vision_tower.parameters():
+            return p.dtype
+
+    @property
+    def device(self):
+        for p in self.vision_tower.parameters():
+            return p.device
+
+    @property
+    def hidden_size(self):
+        return self.config.hidden_size
+
+    @property
+    def num_patches(self):
+        return (self.config.image_size // self.config.patch_size) ** 2