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