model/llava/multimodal_encoder/custom_clip.py

from transformers import GenerationMixin, GenerationConfig, PretrainedConfig, CLIPVisionModel
from transformers.models.clip.modeling_clip import CLIPVisionTransformer
from transformers.utils import add_start_docstrings_to_model_forward, replace_return_docstrings
import torch.nn as nn
from transformers.modeling_outputs import BaseModelOutputWithPooling
from typing import Optional, Tuple, Union
from transformers import CLIPConfig, CLIPTextConfig, CLIPVisionConfig
import torch
import torch.nn.functional as F
import math  
from dataclasses import dataclass
from transformers.models.clip.modeling_clip import (  
    CLIPEncoder,   
    CLIPEncoderLayer,  
    CLIPAttention,  
    CLIPMLP  
)
@dataclass
class BaseModelOutputWithPoolingAndKeys(BaseModelOutputWithPooling):
    keys: Optional[Tuple[torch.FloatTensor]] = None

class _CLIPAttention(CLIPAttention):  
    def forward(  
        self,  
        hidden_states: torch.Tensor,  
        attention_mask: Optional[torch.Tensor] = None,  
        causal_attention_mask: Optional[torch.Tensor] = None,  
        output_attentions: Optional[bool] = False,  
        output_keys: Optional[bool] = False,  
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]:  
        """自定义 CLIP Attention,支持返回 key 向量"""  
          
        bsz, tgt_len, embed_dim = hidden_states.size()  
  
        query_states = self.q_proj(hidden_states)  
        key_states = self.k_proj(hidden_states)  
        value_states = self.v_proj(hidden_states)  
  
        query_states = query_states.view(bsz, tgt_len, self.num_heads, self.head_dim).transpose(1, 2)  
        key_states = key_states.view(bsz, tgt_len, self.num_heads, self.head_dim).transpose(1, 2)  
        value_states = value_states.view(bsz, tgt_len, self.num_heads, self.head_dim).transpose(1, 2)  
  
        output_key_states = None  
        if output_keys:  
            output_key_states = key_states.detach()  
  
        attn_weights = torch.matmul(query_states, key_states.transpose(2, 3))  
          
        if attn_weights.size() != (bsz, self.num_heads, tgt_len, tgt_len):  
            raise ValueError(  
                f"Attention weights should be of size {(bsz, self.num_heads, tgt_len, tgt_len)}, but is"  
                f" {attn_weights.size()}"  
            )  
  
        attn_weights = attn_weights / math.sqrt(self.head_dim)  
  
        if causal_attention_mask is not None:  
            if causal_attention_mask.size() != (bsz, 1, tgt_len, tgt_len):  
                raise ValueError(  
                    f"Attention mask should be of size {(bsz, 1, tgt_len, tgt_len)}, but is"  
                    f" {causal_attention_mask.size()}"  
                )  
            attn_weights = attn_weights + causal_attention_mask  
  
        if attention_mask is not None:  
            if attention_mask.size() != (bsz, 1, tgt_len, tgt_len):  
                raise ValueError(  
                    f"Attention mask should be of size {(bsz, 1, tgt_len, tgt_len)}, but is"  
                    f" {attention_mask.size()}"  
                )  
            attn_weights = attn_weights + attention_mask  
  
        attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)  
  
        attn_output = torch.matmul(attn_weights, value_states)  
  
        if attn_output.size() != (bsz, self.num_heads, tgt_len, self.head_dim):  
            raise ValueError(  
                f"`attn_output` should be of size {(bsz, self.num_heads, tgt_len, self.head_dim)}, but is"  
                f" {attn_output.size()}"  
            )  
  
        attn_output = attn_output.transpose(1, 2).contiguous()  
        attn_output = attn_output.reshape(bsz, tgt_len, embed_dim)  
  
        attn_output = self.out_proj(attn_output)  
  
        return attn_output, attn_weights if output_attentions else None, output_key_states

class _CLIPEncoderLayer(CLIPEncoderLayer):  
    def __init__(self, config):  
        super().__init__(config)  
        self.self_attn = _CLIPAttention(config)  
      
    def forward(  
        self,  
        hidden_states: torch.Tensor,  
        attention_mask: torch.Tensor,  
        causal_attention_mask: torch.Tensor,  
        output_attentions: Optional[bool] = False,  
        output_keys: Optional[bool] = False,  
    ) -> Tuple[torch.FloatTensor]:  
        """  
        Args:  
            hidden_states: [bsz, seq_len, embed_dim]  
            attention_mask: [bsz, 1, tgt_len, src_len]  
            output_attentions: 是否返回 attention 权重  
            output_keys: 是否返回 key 向量  
        """  
        residual = hidden_states  
  
        hidden_states = self.layer_norm1(hidden_states)  
        hidden_states, attn_weights, key_states = self.self_attn(  
            hidden_states=hidden_states,  
            attention_mask=attention_mask,  
            causal_attention_mask=causal_attention_mask,  
            output_attentions=output_attentions,  
            output_keys=output_keys,  
        )  
        hidden_states = residual + hidden_states  
  
        residual = hidden_states  
        hidden_states = self.layer_norm2(hidden_states)  
        hidden_states = self.mlp(hidden_states)  
        hidden_states = residual + hidden_states  
  
        outputs = (hidden_states,)  
  
        if output_attentions:  
            outputs += (attn_weights,)  
  
        if output_keys:  
            outputs += (key_states,)  
  
        return outputs


class _CLIPEncoder(CLIPEncoder):  
    def __init__(self, config):  
        super().__init__(config)  
        self.layers = nn.ModuleList([_CLIPEncoderLayer(config) for _ in range(config.num_hidden_layers)])  
      
    def forward(  
        self,  
        inputs_embeds,  
        attention_mask: Optional[torch.Tensor] = None,  
        causal_attention_mask: Optional[torch.Tensor] = None,  
        output_attentions: Optional[bool] = None,  
        output_hidden_states: Optional[bool] = None,  
        return_dict: Optional[bool] = None,  
        output_keys: Optional[bool] = False,  
    ):  
        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  
        all_keys = () if output_keys else None  
  
        hidden_states = inputs_embeds  
        for idx, encoder_layer in enumerate(self.layers):  
            if output_hidden_states:  
                encoder_states = encoder_states + (hidden_states,)  
  
            layer_outputs = encoder_layer(  
                hidden_states,  
                attention_mask,  
                causal_attention_mask,  
                output_attentions=output_attentions,  
                output_keys=output_keys,  
            )  
  
            hidden_states = layer_outputs[0]  
  
            if output_attentions:  
                all_attentions = all_attentions + (layer_outputs[1],)  
  
            if output_keys:  
                key_idx = 2 if output_attentions else 1  
                all_keys = all_keys + (layer_outputs[key_idx],)  
  
        if output_hidden_states:  
            encoder_states = encoder_states + (hidden_states,)  
  
        if not return_dict:  
            outputs = (hidden_states,)  
            if output_hidden_states:  
                outputs += (encoder_states,)  
            if output_attentions:  
                outputs += (all_attentions,)  
            if output_keys:  
                outputs += (all_keys,)  
            return outputs  
  
        return {  
            'last_hidden_state': hidden_states,  
            'hidden_states': encoder_states,  
            'attentions': all_attentions,  
            'keys': all_keys,  
        }

CLIP_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
            [`AutoImageProcessor`]. See [`CLIPImageProcessor.__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.
"""

def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):
    """
    Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`.
    """
    bsz, src_len = mask.size()
    tgt_len = tgt_len if tgt_len is not None else src_len

    expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype)

    inverted_mask = 1.0 - expanded_mask

    return inverted_mask.masked_fill(inverted_mask.to(torch.bool), torch.finfo(dtype).min)

def visualize_attn_mask(mask):
    import cv2
    import numpy as np
    mask = mask[0].squeeze().float()
    fg = mask >= 0
    mask_show = torch.zeros_like(mask)
    mask_show[fg] = 255
    mask_show = mask_show.cpu().numpy()
    cv2.imwrite('test.jpg', mask_show.astype(np.uint8))

class _CLIPVisionTransformer(CLIPVisionTransformer):
    def __init__(self, config: CLIPVisionConfig):  
        super().__init__(config)  
        self.encoder = _CLIPEncoder(config)  

    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,
        attention_mask: Optional[torch.FloatTensor] = None,
        output_keys: Optional[bool] = False,
    ) -> 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)
        hidden_states = self.pre_layrnorm(hidden_states)

        if attention_mask is not None:
            attention_mask = _expand_mask(attention_mask, hidden_states.dtype)
        encoder_outputs = self.encoder(
            attention_mask=attention_mask,
            inputs_embeds=hidden_states,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            output_keys=output_keys,  
        )

        last_hidden_state = encoder_outputs[0]
        pooled_output = last_hidden_state[:, 0, :]
        pooled_output = self.post_layernorm(pooled_output)

        if not return_dict:  
            outputs = (last_hidden_state, pooled_output)  
            if output_hidden_states:  
                outputs += (encoder_outputs['hidden_states'],)  
            if output_attentions:  
                outputs += (encoder_outputs['attentions'],)  
            if output_keys:  
                outputs += (encoder_outputs['keys'],)  
            return outputs  
        return BaseModelOutputWithPoolingAndKeys(
            last_hidden_state=last_hidden_state,
            pooler_output=pooled_output,
            hidden_states=encoder_outputs['hidden_states'],
            attentions=encoder_outputs['attentions'],
            keys=encoder_outputs['keys'] if output_keys else None
        )

class _CLIPVisionModel(CLIPVisionModel):
    def __init__(self, config: CLIPVisionConfig):
        super().__init__(config)
        self.vision_model = _CLIPVisionTransformer(config)
        self.post_init()
    @add_start_docstrings_to_model_forward(CLIP_VISION_INPUTS_DOCSTRING)
    @replace_return_docstrings(output_type=BaseModelOutputWithPooling, config_class=CLIPVisionConfig)
    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,
        attention_mask: Optional[torch.FloatTensor] = None,
        output_keys: Optional[bool] = False,
    ) -> Union[Tuple, BaseModelOutputWithPooling]:
        r"""
        Returns:

        Examples:

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

        >>> model = CLIPVisionModel.from_pretrained("openai/clip-vit-base-patch32")
        >>> processor = AutoProcessor.from_pretrained("openai/clip-vit-base-patch32")

        >>> 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(**inputs)
        >>> last_hidden_state = outputs.last_hidden_state
        >>> pooled_output = outputs.pooler_output  # pooled CLS states
        ```"""
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        return self.vision_model(
            pixel_values=pixel_values,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            attention_mask=attention_mask,
            output_keys=output_keys 
        )