add: UniViTAR models

.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
+pytorch_model.bin filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,3 +1,85 @@
----
-license: apache-2.0
----
+---
+license: apache-2.0
+datasets:
+- mlfoundations/datacomp_1b
+- kakaobrain/coyo-700m
+- laion/laion400m
+language:
+- en
+- zh
+metrics:
+- accuracy
+- recall
+pipeline_tag: feature-extraction
+library_name: transformers
+---
+
+<h1 align="center">Unified Vision Transformer with Native Resolution</h1>
+
+
+## 🌠 Introduction
+
+We present **UniViTAR**, a family of homogeneous vision foundation models tailored **for unified visual modality and native resolution scenario** in the era of multimodal. We train our UniViTAR family across multiple model scales from **0.3B to 1.4B** exclusively on public accessible image-caption data (14.6B), and observe a trend of performance increasing with parameter scaling. UniViTAR is a Transformer-based encoder model that inherits the original architecture of the conventional Vision Transformer but incorporates the following advanced modifications: *Unified Patchify for Native Image and Video Modality, 2D RoPE, SwiGLU, RMSNorm, and QK-Norm*. 
+
+
+## 🛠️ Environment   
+   ```bash
+   conda create -n univitar python=3.11 -y
+   conda activate univitar
+   pip3 install einops==0.8.0 ninja==1.11.1.1 numpy==1.26.4 pillow==10.4.0 psutil==6.0.0 torch==2.2.2 torchvision==0.17.2 transformers==4.49.0 timm==1.0.14
+   pip3 install flash-attn==2.6.3
+   ```
+
+
+## 🗝️ Model Usage
+
+```python
+import torch
+import numpy as np
+from PIL import Image
+from modeling_univitar import UniViTARVisionModel
+
+# Prepare Model
+model = UniViTARVisionModel("config.json")
+_ = model.load_state_dict(torch.load(f"pytorch_model.bin", map_location="cpu"))
+model = model.to(torch.bfloat16).cuda()
+
+# Prepare Data: [(3, H1, W1), ..., (3, Hn, Wn)] --> (N1+...+Nn, P)
+images = [Image.open(f"xx1.jpg"), Image.open(f"xx2.jpg")]
+data_inputs, grid_shapes = [], []
+for image in images:
+    data_item = model.image_transform(image)
+    input_data, grid_shape = model.data_patchify(data_item)
+    data_inputs.append(input_data.to(torch.bfloat16).cuda())
+    grid_shapes.append(grid_shape)
+data_inputs = torch.concatenate(data_inputs, dim=0)
+
+# Forward: (N1+...+Nn, P) --> [(N1, D), ..., (Nn, D)]
+data_embeds = model(pixel_values=data_inputs, grid_shapes=grid_shapes)
+data_embeds = data_embeds.split([np.prod(grid_shape) for grid_shape in grid_shapes])
+print(data_embeds[0].shape, data_embeds[1].shape)
+```
+
+## 📈 Evaluation
+
+| Model | Size | \#Seen | IN1K<sup>ZS<sup> | IN1K<sup>LP<sup> | Flickr<sup>T2I<sup> | Flickr<sup>I2T<sup> | K400<sup>ZS<sup>  | ADE20K | 
+|--------|-----|----|------|------|------|------|------|------|
+| [UniViTAR-0.3B](https://huggingface.co/MM-MVR/UniViTAR-0.3B) | 310M | 14.6B  | 81.5  | 87.7 | 84.0 | 95.1 | 66.0 | 54.6 |
+| [UniViTAR-0.6B](https://huggingface.co/MM-MVR/UniViTAR-0.6B) | 637M | 14.6B | 82.3  | 88.3 | 84.1 | 95.5 | 68.6 | 55.1 |
+| [UniViTAR-1B](https://huggingface.co/MM-MVR/UniViTAR-1B) | 1419M | 14.6B  | 82.9  | 89.2 | 83.5 | 95.1 | 69.0 | 56.2 |
+
+<font size=1>*ZS: Zero-shot Classification, LP: Linear-Probe Classification, T2I/I2T: Text-to-Image/Image-to-Text Retrieval*</font>
+
+
+## ✏️ Reference
+
+If you find UniViTAR useful in your research or applications, please consider citing the following BibTeX:
+
+```
+@article{qiao2025univitar,
+title={UniViTAR: Unified Vision Transformer with Native Resolution},
+author={Qiao, Limeng and Gan, Yiyang and Wang, Bairui and Qin, Jie and Xu, Shuang and Yang, Siqi and Ma, Lin},
+journal={arXiv preprint arXiv:2504.01792},
+year={2025}
+}
+```

config.json

@@ -0,0 +1,19 @@
+{
+    "resolution_mode": "native",
+    "min_tokens": 256,
+    "max_tokens": 16384,
+    "patch_size": 14,
+    "resize_factor": 2,
+    "spatial_merge_size": 1,
+    "temporal_patch_size": 2,
+    "num_hidden_layers": 24,
+    "num_attention_heads": 16,
+    "hidden_size": 1024,
+    "intermediate_size": 4224,
+    "pe_type": "rope2d",
+    "norm_type": "RMSNorm",
+    "hidden_act": "SwiGLU",
+    "init_method": "xavier",
+    "image_mean": [0.485, 0.456, 0.406],
+    "image_std": [0.229, 0.224, 0.225]
+}

modeling_univitar.py

@@ -0,0 +1,604 @@
+from typing import Iterable, Optional, Tuple, Union, List 
+
+import os
+import math
+import json
+import torch
+import numpy as np
+import torch.nn as nn
+import torch.utils.checkpoint
+import torch.nn.functional as F
+
+from PIL import Image
+from einops import rearrange
+from functools import partial
+from timm.layers import DropPath
+from dataclasses import dataclass
+from torchvision import transforms
+from transformers.utils import logging
+from transformers.activations import ACT2FN
+from transformers.modeling_utils import PreTrainedModel
+from transformers.configuration_utils import PretrainedConfig
+from transformers.modeling_outputs import BaseModelOutput, ModelOutput
+from flash_attn.bert_padding import pad_input
+from flash_attn.flash_attn_interface import flash_attn_varlen_qkvpacked_func as flash_attn_unpadded_qkvpacked_func
+
+logger = logging.get_logger(__name__)
+
+
+def rotate_half(x):
+    """Rotates half the hidden dims of the input."""
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb_vision(tensor: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor:
+    orig_dtype = tensor.dtype
+    tensor = tensor.float()
+    cos = freqs.cos()
+    sin = freqs.sin()
+    cos = cos.unsqueeze(1).repeat(1, 1, 2).unsqueeze(0).float()
+    sin = sin.unsqueeze(1).repeat(1, 1, 2).unsqueeze(0).float()
+    output = (tensor * cos) + (rotate_half(tensor) * sin)
+    output = output.to(orig_dtype)
+    return output
+
+
+class VisionRotaryEmbedding2D(nn.Module):
+    def __init__(self, dim: int, theta: float = 10000.0) -> None:
+        super().__init__()
+        inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float) / dim))
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+    def forward_(self, seqlen: int) -> torch.Tensor:
+        seq = torch.arange(seqlen, device=self.inv_freq.device, dtype=self.inv_freq.dtype)
+        freqs = torch.outer(seq, self.inv_freq)
+        return freqs
+    
+    def forward(self, grid_shapes, spatial_merge_size=2):
+        pos_ids = []
+        s = spatial_merge_size
+        for t, h, w in grid_shapes:
+            hpos_ids = torch.arange(h).unsqueeze(1).expand(-1, w)
+            hpos_ids = hpos_ids.reshape(h // s, s, w // s, s)
+            hpos_ids = hpos_ids.permute(0, 2, 1, 3)
+            hpos_ids = hpos_ids.flatten()
+            wpos_ids = torch.arange(w).unsqueeze(0).expand(h, -1)
+            wpos_ids = wpos_ids.reshape(h // s, s, w // s, s)
+            wpos_ids = wpos_ids.permute(0, 2, 1, 3)
+            wpos_ids = wpos_ids.flatten()
+            pos_ids.append(torch.stack([hpos_ids, wpos_ids], dim=-1).repeat(t, 1))
+        pos_ids = torch.cat(pos_ids, dim=0)
+        max_grid_size = torch.tensor(grid_shapes).max()
+        rotary_pos_emb_full = self.forward_(max_grid_size)
+        rotary_pos_emb = rotary_pos_emb_full[pos_ids].flatten(1)
+        return rotary_pos_emb
+
+
+class FlashAttention(nn.Module):
+    # https://github.com/Dao-AILab/flash-attention/blob/v0.2.8/flash_attn/flash_attention.py
+    """Implement the scaled dot product attention with softmax.
+    Arguments
+    ---------
+        softmax_scale: The temperature to use for the softmax attention.
+                      (default: 1/sqrt(d_keys) where d_keys is computed at
+                      runtime)
+        attention_dropout: The dropout rate to apply to the attention
+                           (default: 0.0)
+    """
+
+    def __init__(self, softmax_scale=None, attention_dropout=0.0, device=None, dtype=None):
+        super().__init__()
+        self.softmax_scale = softmax_scale
+        self.dropout_p = attention_dropout
+        self._deterministic = True
+
+    def forward(self, qkv, key_padding_mask=None, causal=False, cu_seqlens=None,
+                max_s=None, need_weights=False):
+        """Implements the multihead softmax attention.
+        Arguments
+        ---------
+            qkv: The tensor containing the query, key, and value. (B, S, 3, H, D) if key_padding_mask is None
+                if unpadded: (nnz, 3, h, d)
+            key_padding_mask: a bool tensor of shape (B, S)
+        """
+        assert not need_weights
+        assert qkv.dtype in [torch.float16, torch.bfloat16]
+        assert qkv.is_cuda
+
+        if cu_seqlens is None:
+            batch_size = qkv.shape[0]
+            seqlen = qkv.shape[1]
+            if key_padding_mask is None:
+                qkv = rearrange(qkv, 'b s ... -> (b s) ...')
+                max_s = seqlen
+                cu_seqlens = torch.arange(0, (batch_size + 1) * seqlen, step=seqlen, dtype=torch.int32,
+                                          device=qkv.device)
+                output = flash_attn_unpadded_qkvpacked_func(
+                    qkv, cu_seqlens, max_s, self.dropout_p if self.training else 0.0,
+                    softmax_scale=self.softmax_scale, causal=causal
+                )
+                output = rearrange(output, '(b s) ... -> b s ...', b=batch_size)
+            else:
+                qkv = qkv.squeeze()  # [1, n, h, d] -> [n, h, d]
+                seqlens_in_batch = key_padding_mask.sum(dim=-1, dtype=torch.int32)
+                max_seqlen_in_batch = seqlens_in_batch.max().item()
+                cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.torch.int32), (1, 0))
+                output = flash_attn_unpadded_qkvpacked_func(
+                    qkv, cu_seqlens, max_seqlen_in_batch, self.dropout_p if self.training else 0.0,
+                    softmax_scale=self.softmax_scale, causal=causal, deterministic=self._deterministic
+                )
+                output = output.unsqueeze(0)
+        else:
+            assert max_s is not None
+            output = flash_attn_unpadded_qkvpacked_func(
+                qkv, cu_seqlens, max_s, self.dropout_p if self.training else 0.0,
+                softmax_scale=self.softmax_scale, causal=causal
+            )
+
+        return output, None
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states):
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return self.weight * hidden_states.to(input_dtype)
+
+
+try:
+    from apex.normalization import FusedRMSNorm
+    RMSNorm = FusedRMSNorm  # noqa
+    logger.info('Discovered apex.normalization.FusedRMSNorm - will use it instead of RMSNorm')
+except ImportError:  # using the normal RMSNorm
+    pass
+except Exception:
+    logger.warning('discovered apex but it failed to load, falling back to RMSNorm')
+    pass
+
+
+@dataclass
+class BaseModelOutputWithKwargs(ModelOutput):
+    last_hidden_state: torch.FloatTensor = None
+    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+    kwargs: Optional[dict] = None
+
+
+class UniViTARVisionConfig(PretrainedConfig):
+    def __init__(
+            self,
+            resolution_mode="native",
+            init_method="xavier",
+            num_channels=3,
+            patch_size=14,
+            temporal_patch_size=2,
+            image_size=1792,
+            patch_dropout=0.0,
+            attention_dropout=0.0,
+            dropout=0.0,
+            drop_path_rate=0.0,
+            initializer_range=1e-10,
+            num_hidden_layers=24,
+            num_attention_heads=16,
+            hidden_size=1024,
+            intermediate_size=4224,
+            patch_embedding_bias=True,
+            qk_normalization=True,
+            qkv_bias=False,
+            initializer_factor=0.1,
+            use_pre_norm=False,
+            pe_type="rope2d",
+            rope_theta=10000,
+            spatial_merge_size=1,
+            norm_type="RMSNorm",
+            hidden_act='SwiGLU',
+            use_flash_attn=True,
+            layer_norm_eps=1e-6,
+            min_tokens=576,
+            max_tokens=16384,
+            image_mean=(0.485, 0.456, 0.406),
+            image_std=(0.229, 0.224, 0.225),
+            relarge_ratio=1.0,
+            **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.resolution_mode = resolution_mode
+        self.init_method = init_method
+        self.pe_type = pe_type
+        self.rope_theta = rope_theta
+        self.temporal_patch_size = temporal_patch_size
+        self.num_channels = num_channels
+        self.patch_size = patch_size
+        self.image_size = image_size
+        self.patch_dropout = patch_dropout
+        self.attention_dropout = attention_dropout
+        self.dropout = dropout
+        self.drop_path_rate = drop_path_rate
+        self.initializer_range = initializer_range
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads 
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.patch_embedding_bias = patch_embedding_bias
+        self.qk_normalization = qk_normalization
+        self.qkv_bias = qkv_bias
+        self.initializer_factor = initializer_factor
+        self.use_pre_norm = use_pre_norm
+        self.norm_type = norm_type
+        self.hidden_act = hidden_act
+        self.use_flash_attn = use_flash_attn
+        self.layer_norm_eps = layer_norm_eps
+        self.spatial_merge_size = spatial_merge_size
+        self.min_tokens = min_tokens
+        self.max_tokens = max_tokens
+        self.image_mean = image_mean
+        self.image_std = image_std
+        self.relarge_ratio = relarge_ratio
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path: Union[str, os.PathLike], **kwargs) -> 'PretrainedConfig':
+        config_dict, kwargs = cls.get_config_dict(pretrained_model_name_or_path, **kwargs)
+
+        if 'vision_config' in config_dict:
+            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 UniViTARImageTransform(object):
+    def __init__(self, config):
+        self.config = config
+        self.resolution_mode = config.resolution_mode
+        
+        self.image_mean, self.image_std = config.image_mean, config.image_std
+        self.patch_size = config.patch_size
+        self.temporal_patch_size = config.temporal_patch_size
+        self.spatial_merge_size = config.spatial_merge_size
+        self.resize_factor = config.patch_size * config.spatial_merge_size * config.resize_factor
+        self.relarge_ratio = config.relarge_ratio
+
+        self.forced_transform = None
+        self.min_pixels, self.max_pixels = None, None
+        assert self.resolution_mode in ["native", "224", "378", "756"]
+        if self.resolution_mode == "native":
+            self.min_pixels = config.min_tokens * config.patch_size * config.patch_size
+            self.max_pixels = config.max_tokens * config.patch_size * config.patch_size
+        else:
+            image_size = int(self.resolution_mode)
+            self.forced_transform = transforms.Compose([
+                    transforms.Resize((image_size, image_size), interpolation=transforms.InterpolationMode.BICUBIC),
+                    self.convert_to_rgb,
+                    transforms.ToTensor(),
+                    transforms.Normalize(mean=self.image_mean, std=self.image_std)
+                ]
+            )
+
+    def __call__(self, images):
+        
+        if not isinstance(images, List):
+            images = [images]  # shape of each image is [h, w, c]
+        assert len(images) == 1 or len(images) % self.temporal_patch_size == 0
+
+        if self.resolution_mode == "native":
+            sample_num = 1 if len(images) == 1 else len(images) // self.temporal_patch_size
+            min_pixels, max_pixels = self.min_pixels // sample_num, self.max_pixels // sample_num
+            width, height = images[0].size  # (w, h)
+            if self.relarge_ratio > 0 and self.relarge_ratio != 1:
+                height, width = int(height * self.relarge_ratio), int(width * self.relarge_ratio)
+            resized_height, resized_width = self.smart_resize(height, width, self.resize_factor, min_pixels, max_pixels)
+            processed_images = []
+            for image in images:
+                image = self.convert_to_rgb(image)
+                image = self.resize(image, size=(resized_height, resized_width), resample=Image.Resampling.BICUBIC)
+                image = self.rescale(image, scale=1/255)
+                image = self.normalize(image=image, mean=self.image_mean, std=self.image_std)
+                processed_images.append(image)
+            processed_images = np.array(processed_images)  # (num, h, w, c)
+            processed_images = processed_images.transpose(0, 3, 1, 2)  # (num, c, h, w)
+        else:
+            processed_images = [self.forced_transform(image).numpy() for image in images]
+            processed_images = np.array(processed_images)
+
+        if processed_images.shape[0] == 1:
+            processed_images = np.tile(processed_images, (self.temporal_patch_size, 1, 1, 1))
+
+        return torch.from_numpy(processed_images)    
+    
+    @staticmethod
+    def convert_to_rgb(image):
+        if not isinstance(image, Image.Image):
+            return image
+        # `image.convert("RGB")` would only work for .jpg images, as it creates a wrong background
+        # for transparent images. The call to `alpha_composite` handles this case
+        if image.mode == "RGB":
+            return image
+        image_rgba = image.convert("RGBA")
+        background = Image.new("RGBA", image_rgba.size, (255, 255, 255))
+        alpha_composite = Image.alpha_composite(background, image_rgba)
+        alpha_composite = alpha_composite.convert("RGB")
+        return alpha_composite
+    
+    @staticmethod
+    def resize(image, size, resample, return_numpy: bool = True) -> np.ndarray:
+        """
+        Resizes `image` to `(height, width)` specified by `size` using the PIL library.
+        """
+        if not len(size) == 2:
+            raise ValueError("size must have 2 elements")
+        assert isinstance(image, Image.Image)
+        height, width = size
+        resample = resample if resample is not None else Image.Resampling.BILINEAR
+        # PIL images are in the format (width, height)
+        resized_image = image.resize((width, height), resample=resample, reducing_gap=None)
+        if return_numpy:
+            resized_image = np.array(resized_image)
+            resized_image = np.expand_dims(resized_image, axis=-1) if resized_image.ndim == 2 else resized_image
+        return resized_image
+
+    @staticmethod
+    def rescale(image: np.ndarray, scale: float, dtype: np.dtype = np.float32) -> np.ndarray:
+        if not isinstance(image, np.ndarray):
+            raise TypeError(f"Input image must be of type np.ndarray, got {type(image)}")
+        rescaled_image = image * scale
+        rescaled_image = rescaled_image.astype(dtype)
+        return rescaled_image
+
+    @staticmethod
+    def normalize(image, mean, std) -> np.ndarray:
+        """
+        Normalizes `image` using the mean and standard deviation specified by `mean` and `std`.
+        image = (image - mean) / std
+        """
+        if not isinstance(image, np.ndarray):
+            raise ValueError("image must be a numpy array")
+        num_channels = image.shape[-1]
+        # We cast to float32 to avoid errors that can occur when subtracting uint8 values.
+        # We preserve the original dtype if it is a float type to prevent upcasting float16.
+        if not np.issubdtype(image.dtype, np.floating):
+            image = image.astype(np.float32)
+        if isinstance(mean, Iterable):
+            if len(mean) != num_channels:
+                raise ValueError(f"mean must have {num_channels} elements if it is an iterable, got {len(mean)}")
+        else:
+            mean = [mean] * num_channels
+        mean = np.array(mean, dtype=image.dtype)
+        if isinstance(std, Iterable):
+            if len(std) != num_channels:
+                raise ValueError(f"std must have {num_channels} elements if it is an iterable, got {len(std)}")
+        else:
+            std = [std] * num_channels
+        std = np.array(std, dtype=image.dtype)
+        image = (image - mean) / std
+        return image
+    
+    @staticmethod
+    def smart_resize(height, width, factor, min_pixels, max_pixels):
+        """ 
+        1. Both dimensions (height and width) are divisible by 'factor'.
+        2. The total number of pixels is within the range ['min_pixels', 'max_pixels'].
+        3. The aspect ratio of the image is maintained as closely as possible.
+        """
+        if height < factor or width < factor:
+            if height < factor:
+                ratio = factor / height
+                height, width = factor, int(ratio * width) + 1
+            if width < factor:
+                ratio = factor / width
+                width, height = factor, int(ratio * height) + 1
+        h_bar = round(height / factor) * factor
+        w_bar = round(width / factor) * factor
+        if h_bar * w_bar > max_pixels:
+            beta = math.sqrt((height * width) / max_pixels)
+            h_bar = math.floor(height / beta / factor) * factor
+            w_bar = math.floor(width / beta / factor) * factor
+        elif h_bar * w_bar < min_pixels:
+            beta = math.sqrt(min_pixels / (height * width))
+            h_bar = math.ceil(height * beta / factor) * factor
+            w_bar = math.ceil(width * beta / factor) * factor
+        return h_bar, w_bar
+
+
+class SwiGLU(nn.Module):
+    def __init__(self, config: UniViTARVisionConfig):
+        super().__init__()
+        self.config = config
+        self.inner_hidden_size = int(config.intermediate_size * 2 / 3)
+        self.act = ACT2FN['silu']
+        self.fc1 = nn.Linear(config.hidden_size, self.inner_hidden_size)
+        self.fc2 = nn.Linear(self.inner_hidden_size, config.hidden_size)
+        self.fc3 = nn.Linear(config.hidden_size, self.inner_hidden_size)
+        self.norm = RMSNorm(self.inner_hidden_size, eps=config.layer_norm_eps)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.fc1(x)
+        hidden_states = self.act(hidden_states)
+        hidden_states = self.fc2(self.norm(hidden_states * self.fc3(x)))
+        return hidden_states
+
+
+class Attention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+    def __init__(self, config: UniViTARVisionConfig):
+        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
+        assert config.use_flash_attn is True,  "FlashAttention must be used!"
+        assert self.head_dim * self.num_heads == self.embed_dim
+        
+        self.qkv = nn.Linear(self.embed_dim, 3 * self.embed_dim, bias=config.qkv_bias)
+        self.inner_attn = FlashAttention(attention_dropout=config.attention_dropout)
+        self.proj = nn.Linear(self.embed_dim, self.embed_dim)
+        self.proj_drop = nn.Dropout(config.dropout)
+        if self.config.qk_normalization:
+            self.q_norm = RMSNorm(self.embed_dim, eps=config.layer_norm_eps)
+            self.k_norm = RMSNorm(self.embed_dim, eps=config.layer_norm_eps) 
+
+    def forward(self, hidden_states: torch.Tensor, **kwargs) -> torch.Tensor: 
+        key_padding_mask = kwargs.get("key_padding_mask", None)
+        rotary_pos_emb = kwargs["rotary_pos_emb"]
+
+        qkv = self.qkv(hidden_states)
+        qkv = rearrange(qkv, '... (three h d) -> ... three h d', three=3, h=self.num_heads)
+        bind_dim = qkv.dim() - 3
+        target_dtype = qkv.dtype
+        q, k, v = qkv.unbind(bind_dim)
+        q = apply_rotary_pos_emb_vision(q.unsqueeze(0), rotary_pos_emb).squeeze(0)
+        k = apply_rotary_pos_emb_vision(k.unsqueeze(0), rotary_pos_emb).squeeze(0)
+        if self.config.qk_normalization:
+            q = self.q_norm(q.flatten(-2, -1)).view(q.shape)
+            k = self.k_norm(k.flatten(-2, -1)).view(k.shape)
+        qkv = torch.stack([q, k, v], dim=bind_dim).to(target_dtype)
+        context, _ = self.inner_attn(qkv, key_padding_mask=key_padding_mask, causal=False)
+        
+        outs = self.proj(rearrange(context, '... h d -> ... (h d)')) # input expected to be: [b s h d] or [s h d]
+        outs = self.proj_drop(outs)
+
+        return outs
+
+
+class UniViTARVisionEmbeddings(nn.Module):
+    def __init__(self, config: UniViTARVisionConfig):
+        super().__init__()
+        self.config = config
+        self.embed_dim = config.hidden_size
+        self.patch_size = config.patch_size
+        self.temporal_patch_size = config.temporal_patch_size
+        self.kernel_size = [self.temporal_patch_size, self.patch_size, self.patch_size]
+        self.use_bias = config.patch_embedding_bias
+        self.patch_embedding = nn.Conv3d(
+            in_channels=3, out_channels=self.embed_dim, kernel_size=self.kernel_size, stride=self.kernel_size, bias=self.use_bias)
+
+    def forward(self, pixel_values: torch.FloatTensor, **kwargs) -> torch.Tensor:  
+        pixel_values = pixel_values.view(-1, 3, *self.kernel_size)
+        patch_embeds = self.patch_embedding(pixel_values)
+        embeddings = patch_embeds.view(1, -1, self.embed_dim)
+        self.num_patches = embeddings.shape[1]
+        return embeddings
+
+
+class UniViTARVisionEncoderLayer(nn.Module):
+    def __init__(self, config: UniViTARVisionConfig, drop_path_rate: float):
+        super().__init__()
+        self.embed_dim = config.hidden_size
+        assert config.hidden_act == "SwiGLU"
+
+        self.attn = Attention(config)
+        self.norm1 = RMSNorm(self.embed_dim, eps=config.layer_norm_eps)
+        self.norm2 = RMSNorm(self.embed_dim, eps=config.layer_norm_eps)
+        self.mlp = SwiGLU(config)
+        
+        self.ls1 = nn.Parameter(config.initializer_factor * torch.ones(self.embed_dim))
+        self.ls2 = nn.Parameter(config.initializer_factor * torch.ones(self.embed_dim))
+        self.drop_path1 = DropPath(drop_path_rate) if drop_path_rate > 0. else nn.Identity()
+        self.drop_path2 = DropPath(drop_path_rate) if drop_path_rate > 0. else nn.Identity()
+
+    def forward(self, hidden_states: torch.Tensor, **kwargs):
+        hidden_states = hidden_states + self.drop_path1(self.attn(self.norm1(hidden_states), **kwargs) * self.ls1)
+        hidden_states = hidden_states + self.drop_path2(self.mlp(self.norm2(hidden_states)) * self.ls2)
+        return hidden_states
+
+
+class UniViTARVisionEncoder(nn.Module):
+    """ Transformer encoder consisting of `config.num_hidden_layers` self attention layers. """
+    def __init__(self, config: UniViTARVisionConfig):
+        super().__init__()
+        self.config = config
+        self.gradient_checkpointing = True
+
+        # stochastic depth decay rule
+        dpr = [x.item() for x in torch.linspace(0, config.drop_path_rate, config.num_hidden_layers)]
+        self.layers = nn.ModuleList([UniViTARVisionEncoderLayer(config, dpr[idx]) for idx in range(config.num_hidden_layers)])
+        if self.config.pe_type == "rope2d":
+            head_dim = config.hidden_size // config.num_attention_heads
+            self.rotary_pos_emb = VisionRotaryEmbedding2D(head_dim // 2, theta=self.config.rope_theta)
+        else:
+            raise NotImplementedError
+
+    def forward(self, inputs_embeds, output_hidden_states = False, **kwargs):
+        kwargs["rotary_pos_emb"] = self.rotary_pos_emb(kwargs["grid_shapes"], self.config.spatial_merge_size)
+        
+        encoder_states = () if output_hidden_states else None
+        hidden_states = inputs_embeds
+        for idx, encoder_layer in enumerate(self.layers):
+            if output_hidden_states:
+                encoder_states = encoder_states + (hidden_states,)
+            if self.gradient_checkpointing and self.training:
+                encoder_layer_forward = partial(encoder_layer, **kwargs)
+                layer_outputs = torch.utils.checkpoint.checkpoint(encoder_layer_forward, hidden_states, use_reentrant=True)
+            else:
+                layer_outputs = encoder_layer(hidden_states, **kwargs)
+            hidden_states = layer_outputs
+        if output_hidden_states:
+            encoder_states = encoder_states + (hidden_states,)
+
+        return BaseModelOutputWithKwargs(last_hidden_state=hidden_states, hidden_states=encoder_states, kwargs=kwargs)
+
+
+class UniViTARVisionModel(PreTrainedModel):
+    main_input_name = 'pixel_values'
+    config_class = UniViTARVisionConfig
+    _no_split_modules = ['UniViTARVisionEncoderLayer']
+
+    def __init__(self, model_config_path, *args, **kwargs):
+
+        model_config_dict = json.load(open(model_config_path, "r", encoding="utf8"))
+        config = UniViTARVisionConfig.from_dict(model_config_dict)
+
+        super().__init__(config)
+        self.config = config
+        self.image_transform = UniViTARImageTransform(config)
+
+        self.embeddings = UniViTARVisionEmbeddings(config)
+        self.encoder = UniViTARVisionEncoder(config)
+
+    def get_input_embeddings(self):
+        return self.embeddings
+        
+    def get_padding_mask(self, grid_shapes):
+        seq_len = torch.tensor([int((np.prod(thw) - 1) + 1) for thw in grid_shapes])
+        max_len = torch.max(seq_len)
+        batch_size = len(grid_shapes)
+        mask = torch.zeros((batch_size, max_len), dtype=torch.long)
+        range_matrix = torch.arange(max_len).expand(batch_size, max_len)
+        mask = (range_matrix < seq_len.unsqueeze(1))
+        return mask.cuda()
+
+    def forward(self, pixel_values, output_hidden_states = False, **kwargs):
+        assert len(pixel_values.shape) == 2, "(batch_num_tokens, hidden_size)"
+        assert "grid_shapes" in kwargs, "grid_shapes: [(t, h, w), ..., (t, h, w)]"
+        kwargs["key_padding_mask"] = self.get_padding_mask(kwargs["grid_shapes"])
+        hidden_states = self.embeddings(pixel_values, **kwargs)
+        encoder_outputs = self.encoder(hidden_states, output_hidden_states, **kwargs)
+        last_hidden_state = encoder_outputs.last_hidden_state
+        return last_hidden_state.squeeze(0)
+    
+    def data_patchify(self, input_data):
+        t, c, h, w = input_data.shape
+        grid_t, grid_h, grid_w = t // self.config.temporal_patch_size, h // self.config.patch_size, w // self.config.patch_size
+        grid_size = c * self.config.temporal_patch_size * self.config.patch_size * self.config.patch_size
+        input_data = input_data.reshape(
+            grid_t, self.config.temporal_patch_size, c, 
+            grid_h // self.config.spatial_merge_size, self.config.spatial_merge_size, self.config.patch_size, 
+            grid_w // self.config.spatial_merge_size, self.config.spatial_merge_size, self.config.patch_size
+        )
+        input_data = input_data.permute(0, 3, 6, 4, 7, 2, 1, 5, 8)
+        input_data = input_data.reshape(grid_t * grid_h * grid_w, grid_size).contiguous()
+        grid_shape = (grid_t, grid_h, grid_w)
+        return input_data, grid_shape

pytorch_model.bin

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