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