Aluode/PerceptionLabPortable

1# coding=utf-82# Copyright 2023 the HuggingFace Inc. team. All rights reserved.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8#     http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15 16from typing import Optional, Union17 18import torch19from torch import nn20 21from transformers.models.llava.modeling_llava import (22    LlavaCausalLMOutputWithPast,23    LlavaForConditionalGeneration,24    LlavaModel,25    LlavaModelOutputWithPast,26    LlavaPreTrainedModel,27)28 29from ...activations import ACT2FN30from ...cache_utils import Cache31from ...utils import auto_docstring, logging32from .configuration_vipllava import VipLlavaConfig33 34 35logger = logging.get_logger(__name__)36 37 38class VipLlavaModelOutputWithPast(LlavaModelOutputWithPast):39    pass40 41 42class VipLlavaCausalLMOutputWithPast(LlavaCausalLMOutputWithPast):43    pass44 45 46class VipLlavaMultiModalProjector(nn.Module):47    def __init__(self, config: VipLlavaConfig):48        super().__init__()49        num_feature_layers = 1 if isinstance(config.vision_feature_layers, int) else len(config.vision_feature_layers)50        self.projector_layernorm = nn.LayerNorm(51            num_feature_layers * config.vision_config.hidden_size, eps=config.projector_layernorm_eps52        )53 54        self.linear_1 = nn.Linear(55            num_feature_layers * config.vision_config.hidden_size,56            config.text_config.hidden_size,57            bias=True,58        )59        self.act = ACT2FN[config.projector_hidden_act]60        self.linear_2 = nn.Linear(config.text_config.hidden_size, config.text_config.hidden_size, bias=True)61 62    def forward(self, hidden_states):63        hidden_states = self.projector_layernorm(hidden_states)64        hidden_states = self.linear_1(hidden_states)65        hidden_states = self.act(hidden_states)66        hidden_states = self.linear_2(hidden_states)67        return hidden_states68 69 70class VipLlavaPreTrainedModel(LlavaPreTrainedModel):71    pass72 73 74class VipLlavaModel(LlavaModel):75    def get_image_features(76        self, pixel_values: torch.FloatTensor, vision_feature_layers: Optional[Union[int, list[int]]] = None77    ):78        """79        Obtains image last hidden states from the vision tower and apply multimodal projection.80 81        Args:82            pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`)83               The tensors corresponding to the input images.84            vision_feature_layers (`Union[int, list[int]]`):85                The vision feature layer, or the list of indexes of the layers to select86                the vision feature.87        Returns:88            image_features (`torch.Tensor`): Image feature tensor of shape `(num_images, image_length, embed_dim)`).89        """90        vision_feature_layers = (91            vision_feature_layers if vision_feature_layers is not None else self.config.vision_feature_layers92        )93        image_outputs = self.vision_tower(pixel_values, output_hidden_states=True)94 95        # If multiple feature layers are provided (which is usually the case)96        # then the image features are concatenated after the CLS is removed.97        if isinstance(vision_feature_layers, int):98            image_features = image_outputs.hidden_states[vision_feature_layers][:, 1:]99        else:100            # Usually, we select the features from index 1: the layers -2, -5, -8, -11 and 6101            image_features = [image_outputs.hidden_states[index][:, 1:] for index in vision_feature_layers]102            image_features = torch.cat(image_features, dim=-1)103        image_features = self.multi_modal_projector(image_features)104        return image_features105 106    @auto_docstring107    def forward(108        self,109        input_ids: Optional[torch.LongTensor] = None,110        pixel_values: Optional[torch.FloatTensor] = None,111        attention_mask: Optional[torch.Tensor] = None,112        position_ids: Optional[torch.LongTensor] = None,113        past_key_values: Optional[Cache] = None,114        inputs_embeds: Optional[torch.FloatTensor] = None,115        vision_feature_layers: Optional[Union[int, list[int]]] = None,116        use_cache: Optional[bool] = None,117        output_attentions: Optional[bool] = None,118        output_hidden_states: Optional[bool] = None,119        return_dict: Optional[bool] = None,120        cache_position: Optional[torch.LongTensor] = None,121        **lm_kwargs,122    ) -> Union[tuple, VipLlavaModelOutputWithPast]:123        r"""124        vision_feature_layers (`Union[int, list[int]]`, *optional*):125            The vision feature layer, or the list of indexes of the layers to select126            the vision feature.127        """128        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions129        output_hidden_states = (130            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states131        )132        return_dict = return_dict if return_dict is not None else self.config.use_return_dict133        vision_feature_layers = (134            vision_feature_layers if vision_feature_layers is not None else self.config.vision_feature_layers135        )136 137        if (input_ids is None) ^ (inputs_embeds is not None):138            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")139 140        if inputs_embeds is None:141            inputs_embeds = self.get_input_embeddings()(input_ids)142 143        if pixel_values is not None:144            image_features = self.get_image_features(145                pixel_values=pixel_values, vision_feature_layers=vision_feature_layers146            )147            image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype)148            special_image_mask = self.get_placeholder_mask(149                input_ids, inputs_embeds=inputs_embeds, image_features=image_features150            )151            inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_features)152 153        outputs = self.language_model(154            attention_mask=attention_mask,155            position_ids=position_ids,156            past_key_values=past_key_values,157            inputs_embeds=inputs_embeds,158            use_cache=use_cache,159            output_attentions=output_attentions,160            output_hidden_states=output_hidden_states,161            return_dict=True,162            cache_position=cache_position,163            **lm_kwargs,164        )165 166        output = VipLlavaModelOutputWithPast(167            last_hidden_state=outputs.last_hidden_state,168            past_key_values=outputs.past_key_values,169            hidden_states=outputs.hidden_states,170            attentions=outputs.attentions,171            image_hidden_states=image_features if pixel_values is not None else None,172        )173        return output if return_dict else output.to_tuple()174 175 176class VipLlavaForConditionalGeneration(LlavaForConditionalGeneration):177    def get_image_features(178        self, pixel_values: torch.FloatTensor, vision_feature_layers: Optional[Union[int, list[int]]] = None179    ):180        return self.model.get_image_features(pixel_values=pixel_values, vision_feature_layers=vision_feature_layers)181 182    def forward(183        self,184        input_ids: Optional[torch.LongTensor] = None,185        pixel_values: Optional[torch.FloatTensor] = None,186        attention_mask: Optional[torch.Tensor] = None,187        position_ids: Optional[torch.LongTensor] = None,188        past_key_values: Optional[Cache] = None,189        inputs_embeds: Optional[torch.FloatTensor] = None,190        vision_feature_layers: Optional[Union[int, list[int]]] = None,191        labels: Optional[torch.LongTensor] = None,192        use_cache: Optional[bool] = None,193        output_attentions: Optional[bool] = None,194        output_hidden_states: Optional[bool] = None,195        return_dict: Optional[bool] = None,196        cache_position: Optional[torch.LongTensor] = None,197        logits_to_keep: Union[int, torch.Tensor] = 0,198        **lm_kwargs,199    ) -> Union[tuple, VipLlavaCausalLMOutputWithPast]:200        r"""201        vision_feature_layers (`Union[int, list[int]]`, *optional*):202            The vision feature layer, or the list of indexes of the layers to select203            the vision feature.204        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):205            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,206            config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored207            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.208 209        Example:210 211        ```python212        >>> import torch213        >>> from PIL import Image214        >>> import requests215        >>> from transformers import AutoProcessor, VipLlavaForConditionalGeneration216 217        >>> model = VipLlavaForConditionalGeneration.from_pretrained("llava-hf/vip-llava-7b-hf", device_map="auto", dtype=torch.float16)218        >>> processor = AutoProcessor.from_pretrained("llava-hf/vip-llava-7b-hf")219 220        >>> prompt = "A chat between a curious human and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the human's questions.###Human: <image>\n{}###Assistant:"221        >>> question = "Can you please describe this image?"222        >>> prompt = prompt.format(question)223        >>> url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/compel-neg.png"224        >>> image = Image.open(requests.get(url, stream=True).raw)225 226        >>> inputs = processor(text=text, images=image, return_tensors="pt").to(0, torch.float16)227 228        >>> # Generate229        >>> generate_ids = model.generate(**inputs, max_new_tokens=20)230        >>> processor.decode(generate_ids[0][len(inputs["input_ids"][0]):], skip_special_tokens=True)231        The image features a brown and white cat sitting on a green surface, with a red ball in its232        ```"""233 234        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions235        output_hidden_states = (236            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states237        )238        return_dict = return_dict if return_dict is not None else self.config.use_return_dict239        vision_feature_layers = (240            vision_feature_layers if vision_feature_layers is not None else self.config.vision_feature_layers241        )242 243        outputs = self.model(244            input_ids=input_ids,245            pixel_values=pixel_values,246            attention_mask=attention_mask,247            position_ids=position_ids,248            past_key_values=past_key_values,249            inputs_embeds=inputs_embeds,250            use_cache=use_cache,251            vision_feature_layers=vision_feature_layers,252            output_attentions=output_attentions,253            output_hidden_states=output_hidden_states,254            return_dict=True,255            cache_position=cache_position,256            **lm_kwargs,257        )258 259        hidden_states = outputs[0]260        # Only compute necessary logits, and do not upcast them to float if we are not computing the loss261        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep262        logits = self.lm_head(hidden_states[:, slice_indices, :])263 264        loss = None265        if labels is not None:266            loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.text_config.vocab_size)267 268        return VipLlavaCausalLMOutputWithPast(269            loss=loss,270            logits=logits,271            past_key_values=outputs.past_key_values,272            hidden_states=outputs.hidden_states,273            attentions=outputs.attentions,274            image_hidden_states=outputs.image_hidden_states,275        )276 277 278__all__ = ["VipLlavaModel", "VipLlavaForConditionalGeneration", "VipLlavaPreTrainedModel"]279