rhymes-ai/Aria-sequential_mlp

1# Copyright 2024 Rhymes AI. All rights reserved.2#3# Licensed to the Apache Software Foundation (ASF) under one4# or more contributor license agreements.  See the NOTICE file5# distributed with this work for additional information6# regarding copyright ownership.  The ASF licenses this file7# to you under the Apache License, Version 2.0 (the8# "License"); you may not use this file except in compliance9# with the License.  You may obtain a copy of the License at10#11#   http://www.apache.org/licenses/LICENSE-2.012#13# Unless required by applicable law or agreed to in writing,14# software distributed under the License is distributed on an15# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY16# KIND, either express or implied.  See the License for the17# specific language governing permissions and limitations18# under the License.19 20"""PyTorch Aria vision transformer."""21 22from typing import Optional, Tuple, Union23 24import torch25import torch.utils.checkpoint26from transformers import SiglipVisionConfig, SiglipVisionModel27from transformers.modeling_outputs import BaseModelOutputWithPooling28from transformers.models.idefics2.modeling_idefics2 import Idefics2VisionTransformer29 30 31class AriaVisionConfig(SiglipVisionConfig):32    """Configuration class for AriaVisionModel."""33 34    model_type = "aria_vision_model"35 36    def __init__(37        self,38        **kwargs,39    ):40        super().__init__(**kwargs)41 42 43class IdentityOp(torch.nn.Module):44    """45    An identity operation that returns the input unchanged.46 47    This can be used as a placeholder or to maintain architectural consistency48    when a specific operation is not needed.49    """50 51    def __init__(self, *args, **kwargs):52        super().__init__()53 54    def forward(self, x, *args, **kwargs):55        return x56 57 58class AriaVisionTransformer(Idefics2VisionTransformer):59    """60    Aria Vision Transformer model based on Idefics2VisionTransformer.61 62    This class extends the original Idefics2VisionTransformer by removing the post-layernorm operation.63    """64 65    def __init__(self, config: AriaVisionConfig):66        super().__init__(config)67        self.post_layernorm = IdentityOp()68 69 70class AriaVisionModel(SiglipVisionModel):71    """72    Aria Vision Model extends SiglipVisionModel to support pixel_mask.73 74    The pixel_mask is a 2D boolean tensor that indicates which pixels in the input75    image are actual content and which are padding. It has the same height and width76    as the input image, where:77    - True (1) values represent pixels from the original image78    - False (0) values represent padding pixels79 80    This mask helps the model focus on the relevant parts of the image during processing.81    """82 83    config_class = AriaVisionConfig84    main_input_name = "pixel_values"85    _supports_sdpa = False86 87    def __init__(self, config: AriaVisionConfig):88        super().__init__(config)89        self.vision_model = AriaVisionTransformer(config)90 91        # Initialize weights and apply final processing92        self.post_init()93 94    def forward(95        self,96        pixel_values: torch.Tensor,97        pixel_mask: Optional[torch.BoolTensor] = None,98        output_attentions: Optional[bool] = None,99        output_hidden_states: Optional[bool] = None,100        return_dict: Optional[bool] = None,101    ) -> Union[Tuple, BaseModelOutputWithPooling]:102        """103        Forward pass of the AriaVisionModel.104 105        Args:106            pixel_values (torch.Tensor): The pixel values of the input images.107            pixel_mask (Optional[torch.BoolTensor]): Mask for the pixel values.108            output_attentions (Optional[bool]): Whether to output attentions.109            output_hidden_states (Optional[bool]): Whether to output hidden states.110            return_dict (Optional[bool]): Whether to return a ModelOutput object.111 112        Returns:113            Union[Tuple, BaseModelOutputWithPooling]: The model's output.114        """115        return_dict = (116            return_dict if return_dict is not None else self.config.use_return_dict117        )118        patch_attention_mask = self._create_patch_attention_mask(pixel_mask)119 120        vit_oup = self.vision_model(121            pixel_values=pixel_values,122            patch_attention_mask=patch_attention_mask,123            output_attentions=output_attentions,124            output_hidden_states=output_hidden_states,125            return_dict=return_dict,126        )127 128        image_atts = self._create_image_attention_mask(patch_attention_mask)129 130        return vit_oup, image_atts131 132    def _create_patch_attention_mask(self, pixel_mask):133        if pixel_mask is None:134            return None135 136        patches_subgrid = pixel_mask.unfold(137            dimension=1,138            size=self.vision_model.config.patch_size,139            step=self.vision_model.config.patch_size,140        ).unfold(141            dimension=2,142            size=self.vision_model.config.patch_size,143            step=self.vision_model.config.patch_size,144        )145        return (patches_subgrid.sum(dim=(-1, -2)) > 0).bool()146 147    def _create_image_attention_mask(self, patch_attention_mask):148        if patch_attention_mask is None:149            return None150 151        flattened_mask = patch_attention_mask.flatten(1)152        return torch.logical_not(flattened_mask)153