model/blip/modeling_blip.py

import warnings
from dataclasses import dataclass
from typing import Any, Optional, Tuple, Dict, Union

import torch
from torch import nn
from transformer_lens.hook_points import HookedRootModule, HookPoint

from transformers.modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling
from transformers.utils import (
    add_start_docstrings_to_model_forward,
    replace_return_docstrings,
)
from .configuration_blip import BlipConfig, BlipVisionConfig
from .modeling_blip_text import BlipTextLMHeadModel

from transformers.models.blip.modeling_blip import (
    BlipVisionEmbeddings, BlipForConditionalGenerationModelOutput, 
    BlipPreTrainedModel, BlipEncoder, BlipEncoderLayer
)

from transformer_lens import HookedTransformerConfig

    
class HookedBlipEncoderLayer(BlipEncoderLayer):
    def __init__(self, config: BlipConfig, layer_idx: int):
        super().__init__(config)
        self.layer_idx = layer_idx
        self.hook_attn_out = HookPoint()
        self.hook_mlp_out = HookPoint()
        self.hook_resid_pre = HookPoint()
        self.hook_resid_post = HookPoint()
        
    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 size
                `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
                `(config.encoder_attention_heads,)`.
            output_attentions (`bool`, *optional*):
                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
                returned tensors for more detail.
        """
        
        # pre-residual hook
        self.hook_resid_pre.layer_idx = self.layer_idx
        hidden_states = self.hook_resid_pre(hidden_states)
        
        residual = hidden_states
        
        hidden_states = self.layer_norm1(hidden_states)
        hidden_states, attn_weights = self.self_attn(
            hidden_states=hidden_states,
            head_mask=attention_mask,
            output_attentions=output_attentions,
        )
        
        # attn output hook
        self.hook_attn_out.layer_idx = self.layer_idx
        hidden_states = self.hook_attn_out(hidden_states)
        
        hidden_states = hidden_states + residual
        residual = hidden_states
        
        hidden_states = self.layer_norm2(hidden_states)
        hidden_states = self.mlp(hidden_states)
        
        # mlp output hook
        self.hook_mlp_out.layer_idx = self.layer_idx
        hidden_states = self.hook_mlp_out(hidden_states)
        
        hidden_states = hidden_states + residual
        
        # post-residual hook
        self.hook_resid_post.layer_idx = self.layer_idx
        hidden_states = self.hook_resid_post(hidden_states)
        
        outputs = (hidden_states,)

        if output_attentions:
            outputs += (attn_weights,)

        return outputs

BLIP_VISION_INPUTS_DOCSTRING = r"""
    Args:
        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
            Pixel values. Padding will be ignored by default should you provide it. Pixel values can be obtained using
            [`BlipImageProcessor`]. See [`BlipImageProcessor.__call__`] for details.
        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.
"""

class HookedBlipEncoder(BlipEncoder, HookedRootModule):
    def __init__(self, config: BlipConfig):
        HookedRootModule.__init__(self)
        self.config = config
        self.gradient_checkpointing = False
        self.layers = nn.ModuleList([
            HookedBlipEncoderLayer(config, layer_idx=i) for i in range(config.num_hidden_layers)
        ])
        
        
    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)`):
                Embedded representation of the inputs. Should be float, not int tokens.
            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 idx, encoder_layer in enumerate(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 BlipVisionModel(BlipPreTrainedModel):
    main_input_name = "pixel_values"
    config_class = BlipVisionConfig

    def __init__(self, config: BlipVisionConfig):
        super().__init__(config)
        self.config = config
        embed_dim = config.hidden_size

        self.embeddings = BlipVisionEmbeddings(config)
        self.encoder = HookedBlipEncoder(config)
        self.post_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)

        self.post_init()

    @add_start_docstrings_to_model_forward(BLIP_VISION_INPUTS_DOCSTRING)
    @replace_return_docstrings(output_type=BaseModelOutputWithPooling, config_class=BlipVisionConfig)
    def forward(
        self,
        pixel_values: Optional[torch.FloatTensor] = None,
        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

        if pixel_values is None:
            raise ValueError("You have to specify pixel_values")

        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 = last_hidden_state[:, 0, :]
        pooled_output = self.post_layernorm(pooled_output)

        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,
        )

    def get_input_embeddings(self):
        return self.embeddings

class BlipForConditionalGeneration(BlipPreTrainedModel, HookedRootModule):
    config_class = BlipConfig
    _tied_weights_keys = ["text_decoder.cls.predictions.decoder.bias"]
    main_input_name = "pixel_values"

    def __init__(self, config: BlipConfig):
        HookedRootModule.__init__(self)
        BlipPreTrainedModel.__init__(self, config)

        self.vision_model = BlipVisionModel(config.vision_config)

        self.text_decoder = BlipTextLMHeadModel(config.text_config)

        self.decoder_input_ids = config.text_config.bos_token_id
        self.decoder_pad_token_id = config.text_config.pad_token_id

        self.cfg = HookedTransformerConfig(
            n_layers=config.text_config.num_hidden_layers,
            d_model=config.text_config.hidden_size,
            d_head=config.text_config.num_attention_heads,
            d_mlp=config.text_config.intermediate_size,
            d_vocab=config.text_config.vocab_size,
            n_ctx=config.text_config.max_position_embeddings,
            act_fn=config.text_config.hidden_act,
            device='cuda:0',
        )
        
        
        # Initialize weights and apply final processing
        self.post_init()
        self.setup()

    def get_input_embeddings(self) -> nn.Module:
        return self.vision_model.embeddings.patch_embedding

    @add_start_docstrings_to_model_forward(BLIP_VISION_INPUTS_DOCSTRING)
    @replace_return_docstrings(output_type=BlipForConditionalGenerationModelOutput, config_class=BlipVisionConfig)
    
    # Commented out the original forward function 
    
    # def forward(
    #     self,
    #     pixel_values: torch.FloatTensor,
    #     input_ids: Optional[torch.LongTensor] = None,
    #     attention_mask: Optional[torch.LongTensor] = None,
    #     output_attentions: Optional[bool] = None,
    #     output_hidden_states: Optional[bool] = None,
    #     labels: Optional[torch.LongTensor] = None,
    #     return_dict: Optional[bool] = None,
    # ) -> Union[Tuple, BlipForConditionalGenerationModelOutput]:
    #     r"""
    #     Returns:

    #     Examples:

    #     ```python
    #     >>> from PIL import Image
    #     >>> import requests
    #     >>> from transformers import AutoProcessor, BlipForConditionalGeneration

    #     >>> processor = AutoProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
    #     >>> model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")

    #     >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
    #     >>> image = Image.open(requests.get(url, stream=True).raw)
    #     >>> text = "A picture of"

    #     >>> inputs = processor(images=image, text=text, return_tensors="pt")

    #     >>> outputs = model(**inputs)
    #     ```"""

    #     return_dict = return_dict if return_dict is not None else self.config.use_return_dict
    #     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
    #     )

    #     vision_outputs = self.vision_model(
    #         pixel_values=pixel_values,
    #         output_attentions=output_attentions,
    #         output_hidden_states=output_hidden_states,
    #         return_dict=return_dict,
    #     )

    #     image_embeds = vision_outputs[0]

    #     outputs = self.text_decoder(
    #         input_ids=input_ids,
    #         attention_mask=attention_mask,
    #         encoder_hidden_states=image_embeds,
    #         labels=labels,
    #         return_dict=return_dict,
    #         reduction="mean",
    #     )

    #     if not return_dict:
    #         outputs = (outputs[0], outputs[1], image_embeds, vision_outputs[0]) + vision_outputs[2:]
    #         return tuple(output for output in outputs if output is not None)

    #     return BlipForConditionalGenerationModelOutput(
    #         loss=outputs.loss,
    #         logits=outputs.logits,
    #         image_embeds=image_embeds,
    #         last_hidden_state=vision_outputs.last_hidden_state,
    #         hidden_states=vision_outputs.hidden_states,
    #         attentions=vision_outputs.attentions,
    #     )
    
    
    
    def forward(
        self,
        inputs: Dict[str, torch.Tensor], 
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        ) -> torch.Tensor:
            r"""
            Returns:

            Examples:

            ```python
            >>> from PIL import Image
            >>> import requests
            >>> from transformers import AutoProcessor, BlipForConditionalGeneration

            >>> processor = AutoProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
            >>> model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")

            >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
            >>> image = Image.open(requests.get(url, stream=True).raw)
            >>> text = "A picture of"

            >>> inputs = processor(images=image, text=text, return_tensors="pt")

            >>> outputs = model(inputs)  
            ```"""

            pixel_values = inputs.get('pixel_values')
            input_ids = inputs.get('input_ids')
            attention_mask = inputs.get('attention_mask')
            labels = inputs.get('labels')

            if pixel_values is None:
                raise ValueError("`pixel_values` must be provided in the inputs dict.")

            return_dict = return_dict if return_dict is not None else self.config.use_return_dict
            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
            )

            vision_outputs = self.vision_model(
                pixel_values=pixel_values,
                output_attentions=output_attentions,
                output_hidden_states=output_hidden_states,
                return_dict=return_dict,
            )

            image_embeds = vision_outputs[0]

            outputs = self.text_decoder(
                input_ids=input_ids,
                attention_mask=attention_mask,
                encoder_hidden_states=image_embeds,
                labels=labels,
                return_dict=return_dict,
                reduction="mean",
            )

            if not return_dict:
                outputs = (outputs[0], outputs[1], image_embeds, vision_outputs[0]) + vision_outputs[2:]
                return tuple(output for output in outputs if output is not None)

            
            # Commented out the original output 
            
            # return BlipForConditionalGenerationModelOutput(
            #     loss=outputs.loss,
            #     logits=outputs.logits,
            #     image_embeds=image_embeds,
            #     last_hidden_state=vision_outputs.last_hidden_state,
            #     hidden_states=vision_outputs.hidden_states,
            #     attentions=vision_outputs.attentions,
            # )
            
            
            return outputs.logits
    
        
    

    @torch.no_grad()
    def generate(
        self,
        pixel_values: torch.FloatTensor,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.LongTensor] = None,
        **generate_kwargs,
    ) -> torch.LongTensor:
        r"""
        Overrides *generate* function to be able to use the model as a conditional generator

        Parameters:
            pixel_values (*torch.FloatTensor* of shape *(batch_size, num_channels, image_height, image_width)*:
                Input image to be processed
            input_ids (*torch.LongTensor* of shape *(batch_size, sequence_length)*, *optional*):
                The sequence used as a prompt for the generation.
            attention_mask (*torch.LongTensor* of shape *(batch_size, sequence_length)*, *optional*):
                Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:


        Examples:
        ```python
        >>> from PIL import Image
        >>> import requests
        >>> from transformers import AutoProcessor, BlipForConditionalGeneration

        >>> model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")
        >>> processor = AutoProcessor.from_pretrained("Salesforce/blip-image-captioning-base")

        >>> 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.generate(**inputs)
        >>> print(processor.decode(outputs[0], skip_special_tokens=True))
        two cats sleeping on a couch
        ```
        """

        batch_size = pixel_values.shape[0]
        vision_outputs = self.vision_model(pixel_values=pixel_values)

        image_embeds = vision_outputs[0]

        image_attention_mask = torch.ones(image_embeds.size()[:-1], dtype=torch.long).to(image_embeds.device)

        if isinstance(input_ids, list):
            input_ids = torch.LongTensor(input_ids)
        elif input_ids is None:
            input_ids = (
                torch.LongTensor([[self.decoder_input_ids, self.config.text_config.eos_token_id]])
                .repeat(batch_size, 1)
                .to(image_embeds.device)
            )

        input_ids[:, 0] = self.config.text_config.bos_token_id
        attention_mask = attention_mask[:, :-1] if attention_mask is not None else None

        outputs = self.text_decoder.generate(
            input_ids=input_ids[:, :-1],
            eos_token_id=self.config.text_config.sep_token_id,
            pad_token_id=self.config.text_config.pad_token_id,
            attention_mask=attention_mask,
            encoder_hidden_states=image_embeds,
            encoder_attention_mask=image_attention_mask,
            **generate_kwargs,
        )

        return outputs
    
    def to(self, device=None, dtype=None, non_blocking=False):
        nn.Module.to(self, device, dtype, non_blocking)
        if device is not None:
            self.cfg.device = device
            
        return self