arc2face_models.py

import torch
import torch.nn as nn
from transformers import CLIPTextModel
from transformers.models.clip.modeling_clip import CLIPAttention
from typing import Any, Callable, Dict, Optional, Tuple, Union, List
from transformers.modeling_outputs import BaseModelOutputWithPooling
from transformers.modeling_attn_mask_utils import AttentionMaskConverter
# from transformers.models.clip.modeling_clip import _make_causal_mask, _expand_mask
_make_causal_mask = AttentionMaskConverter._make_causal_mask
_expand_mask = AttentionMaskConverter._expand_mask

from adaface.util import add_noise_to_tensor

# Extend CLIPAttention by using multiple k_proj and v_proj in each head.
# To avoid too much increase of computation, we don't extend q_proj.
class CLIPAttentionMKV(nn.Module):
    """Multi-headed attention from 'Attention Is All You Need' paper"""

    def __init__(self, config, multiplier=2):
        super().__init__()
        self.config = config
        self.embed_dim = config.hidden_size
        self.num_heads = config.num_attention_heads
        self.head_dim = self.embed_dim // self.num_heads
        if self.head_dim * self.num_heads != self.embed_dim:
            raise ValueError(
                f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:"
                f" {self.num_heads})."
            )
        self.scale = self.head_dim**-0.5
        self.dropout = config.attention_dropout
        self.multiplier  = multiplier

        self.k_proj   = nn.Linear(self.embed_dim, self.embed_dim * self.multiplier)
        self.v_proj   = nn.Linear(self.embed_dim, self.embed_dim * self.multiplier)
        self.q_proj   = nn.Linear(self.embed_dim, self.embed_dim)
        self.out_proj = nn.Linear(self.embed_dim, self.embed_dim)

    # The (approximately) repeated token features are repeated along the last dim in tensor
    # (multiplier * num_heads * head_dim), and then reshaped to (bsz, -1, num_heads, head_dim).
    # Therefore, the "multiplier" dim is tucked into the seq_len dim, which looks like
    # [token1_emb, token1_emb, token2_emb, token2_emb, ..., tokenN_emb, tokenN_emb].
    def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
        return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()

    def extend_weights(self, clip_attn_layer, layer_idx, multiplier, noise_std=0.1, 
                       noise_std_is_relative=True, keep_norm=False, verbose=False):
        self.multiplier *= multiplier
        # q_proj and out_proj are the same as the original CLIPAttention.
        self.q_proj.weight.data   = clip_attn_layer.q_proj.weight.data.clone()
        self.q_proj.bias.data     = clip_attn_layer.q_proj.bias.data.clone()
        self.out_proj.weight.data = clip_attn_layer.out_proj.weight.data.clone()
        self.out_proj.bias.data   = clip_attn_layer.out_proj.bias.data.clone()

        # bias doesn't need noise perturbation, as after the weights are noised, 
        # different copies of the weight/bias will receive different gradients, 
        # making the bias terms diverge and identifiable after training.
        self.v_proj.bias.data     = clip_attn_layer.v_proj.bias.data.repeat(multiplier)
        self.k_proj.bias.data     = clip_attn_layer.k_proj.bias.data.repeat(multiplier)

        self.v_proj.weight.data   = clip_attn_layer.v_proj.weight.data.repeat(multiplier, 1)
        self.k_proj.weight.data   = clip_attn_layer.k_proj.weight.data.repeat(multiplier, 1)

        if noise_std > 0:
            ORIG_V_SHAPE    = list(clip_attn_layer.v_proj.weight.shape)
            ORIG_V_SHAPE_D0 = ORIG_V_SHAPE[0]
            # Adding noise to the extra copies of the weights (keep the first copy unchanged).
            self.v_proj.weight.data[ORIG_V_SHAPE_D0:] = \
                add_noise_to_tensor(self.v_proj.weight.data[ORIG_V_SHAPE_D0:], 
                                    noise_std, noise_std_is_relative, keep_norm)
            if verbose:
                NEW_V_SHAPE     = list(self.v_proj.weight.shape)
                NOISED_V_SHAPE  = list(self.v_proj.weight.data[ORIG_V_SHAPE_D0:].shape)
                print(f"Layer {layer_idx}: {NOISED_V_SHAPE} in {NEW_V_SHAPE} of v_proj is added with {noise_std} noise")

            ORIG_K_SHAPE    = list(clip_attn_layer.k_proj.weight.shape)
            ORIG_K_SHAPE_D0 = ORIG_K_SHAPE[0]
            # Adding noise to the extra copies of the weights.
            self.k_proj.weight.data[ORIG_K_SHAPE_D0:] = \
                add_noise_to_tensor(self.k_proj.weight.data[ORIG_K_SHAPE_D0:], 
                                    noise_std, noise_std_is_relative, keep_norm)
            if verbose:
                NEW_K_SHAPE     = list(self.k_proj.weight.shape)
                NOISED_K_SHAPE  = list(self.k_proj.weight.data[ORIG_K_SHAPE_D0:].shape)
                print(f"Layer {layer_idx}: {NOISED_K_SHAPE} in {NEW_K_SHAPE} of k_proj is added with {noise_std} noise")

    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,
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
        """Input shape: Batch x Time x Channel"""

        bsz, tgt_len, embed_dim = hidden_states.size()

        query_states = self.q_proj(hidden_states) * self.scale
        # For key_states and value_states, the multiplier is absorbed into the seq_len (dim 1, shape specified as -1).
        # [token0_head_emb, token0_head_emb, token1_head_emb, token1_head_emb, ..., tokenN-1_head_emb, tokenN-1_head_emb].
        key_states   = self._shape(self.k_proj(hidden_states), -1, bsz)
        value_states = self._shape(self.v_proj(hidden_states), -1, bsz)

        proj_shape = (bsz * self.num_heads, -1, self.head_dim)
        query_states = self._shape(query_states, tgt_len, bsz).view(*proj_shape)
        key_states   = key_states.view(*proj_shape)
        value_states = value_states.view(*proj_shape)

        src_len = key_states.size(1)
        # src_len0 is the original src_len without the multiplier.
        src_len0 = src_len // self.multiplier
        attn_weights = torch.bmm(query_states, key_states.transpose(1, 2))

        if attn_weights.size() != (bsz * self.num_heads, tgt_len, src_len):
            raise ValueError(
                f"Attention weights should be of size {(bsz * self.num_heads, tgt_len, src_len)}, but is"
                f" {attn_weights.size()}"
            )

        # apply the causal_attention_mask first
        if causal_attention_mask is not None:
            if causal_attention_mask.size() != (bsz, 1, tgt_len, src_len0):
                raise ValueError(
                    f"Attention mask should be of size {(bsz, 1, tgt_len, src_len0)}, but is"
                    f" {causal_attention_mask.size()}"
                )
            # The last dim of attn_weights corresponds to [token0, token0, token1, token1, ..., tokenN-1, tokenN-1].
            # If reshaping it as (self.multiplier, src_len0), it will become 
            # [[token0, token0, token1, token1, ..., tokenN//2], [tokenN//2+1, tokenN//2+1, ..., tokenN-1, tokenN-1]],
            # and the mask will be applied to wrong elements.
            # If reshaping it as (src_len0, self.multiplier), it will become 
            # [[token0, token1, ..., tokenN-1], [token0, token1, ..., tokenN-1]], and then
            # the mask at element i will mask all the multiplier elements at i, which is desired.
            attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len0, self.multiplier) + causal_attention_mask.unsqueeze(4)
            attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)

        if attention_mask is not None:
            if attention_mask.size() != (bsz, 1, tgt_len, src_len0):
                raise ValueError(
                    f"Attention mask should be of size {(bsz, 1, tgt_len, src_len0)}, but is {attention_mask.size()}"
                )
            attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len0, self.multiplier) + attention_mask.unsqueeze(4)
            attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)

        attn_weights = nn.functional.softmax(attn_weights, dim=-1)

        if output_attentions:
            # this operation is a bit awkward, but it's required to
            # make sure that attn_weights keeps its gradient.
            # In order to do so, attn_weights have to reshaped
            # twice and have to be reused in the following
            attn_weights_reshaped = attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
            attn_weights = attn_weights_reshaped.view(bsz * self.num_heads, tgt_len, src_len)
        else:
            attn_weights_reshaped = None

        attn_probs = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training)

        attn_output = torch.bmm(attn_probs, 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.view(bsz, self.num_heads, tgt_len, self.head_dim)
        attn_output = attn_output.transpose(1, 2)
        attn_output = attn_output.reshape(bsz, tgt_len, embed_dim)

        attn_output = self.out_proj(attn_output)

        return attn_output, attn_weights_reshaped
    
class CLIPTextModelWrapper(CLIPTextModel):
    # Adapted from https://github.com/huggingface/transformers/blob/v4.34.1/src/transformers/models/clip/modeling_clip.py#L812
    # Modified to accept precomputed token embeddings "input_token_embs" as input or calculate them from input_ids and return them.
    def forward(
        self,
        input_ids: Optional[torch.Tensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.Tensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        input_token_embs: Optional[torch.Tensor] = None,
        hidden_state_layer_weights: Optional[torch.Tensor] = None,
        return_token_embs: Optional[bool] = False,
    ) -> Union[Tuple, torch.Tensor, BaseModelOutputWithPooling]:

        if return_token_embs:
            return self.text_model.embeddings.token_embedding(input_ids)
    
        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.text_model.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.text_model.config.output_hidden_states
        )
        if hidden_state_layer_weights is not None:
            output_hidden_states = True
        return_dict = return_dict if return_dict is not None else self.text_model.config.use_return_dict
    
        if input_ids is None:
            raise ValueError("You have to specify input_ids")
    
        input_shape = input_ids.size()
        input_ids = input_ids.view(-1, input_shape[-1])
    
        hidden_states = self.text_model.embeddings(input_ids=input_ids, position_ids=position_ids, inputs_embeds=input_token_embs)
    
        # CLIP's text model uses causal mask, prepare it here.
        # https://github.com/openai/CLIP/blob/cfcffb90e69f37bf2ff1e988237a0fbe41f33c04/clip/model.py#L324
        causal_attention_mask = _make_causal_mask(input_shape, hidden_states.dtype, device=hidden_states.device)
        # expand attention_mask
        if attention_mask is not None:
            # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
            attention_mask = _expand_mask(attention_mask, hidden_states.dtype)
    
        encoder_outputs = self.text_model.encoder(
            inputs_embeds=hidden_states,
            attention_mask=attention_mask,
            causal_attention_mask=causal_attention_mask,
            output_attentions=output_attentions,
            # output_hidden_states is False by default, and only True if hidden_state_layer_weights is provided.
            output_hidden_states=output_hidden_states,      
            return_dict=return_dict,
        )
    
        # If output_hidden_states is True, then encoder_outputs[0] is last_hidden_state [1, 22, 768].
        # encoder_outputs[1] is hidden_states, which is a tuple of 13 hidden states, each being [1, 22, 768].
        # encoder_outputs[0] == encoder_outputs[1][12].
        if hidden_state_layer_weights is None:
            last_hidden_state = encoder_outputs[0]
        else:
            num_hidden_state_layers = len(hidden_state_layer_weights)
            last_hidden_states = encoder_outputs[1][-num_hidden_state_layers:]
            hidden_state_layer_weights = hidden_state_layer_weights.to(last_hidden_states[0].dtype)
            # Normalize the weights of to sum to 1 across layers.
            # hidden_state_layer_weights: [3, 1] or [3, 768].
            hidden_state_layer_weights = hidden_state_layer_weights / hidden_state_layer_weights.sum(dim=0, keepdim=True)
            # [3, 1/768] -> [3, 1, 1, 1/768]
            hidden_state_layer_weights = hidden_state_layer_weights.unsqueeze(1).unsqueeze(1)
            # A weighted sum of last_hidden_states.
            # [3, 1, 22, 768] * [3, 1, 1, 1/768] -> [3, 1, 22, 768] -> [1, 22, 768]
            last_hidden_state = (torch.stack(last_hidden_states, dim=0) * hidden_state_layer_weights).sum(dim=0)
            
        last_hidden_state = self.text_model.final_layer_norm(last_hidden_state)

        # self.text_model.eos_token_id == 2 is True.
        if self.text_model.eos_token_id == 2:
            # The `eos_token_id` was incorrect before PR #24773: Let's keep what have been done here.
            # A CLIP model with such `eos_token_id` in the config can't work correctly with extra new tokens added
            # ------------------------------------------------------------
            # text_embeds.shape = [batch_size, sequence_length, transformer.width]
            # take features from the eot embedding (eot_token is the highest number in each sequence)
            # casting to torch.int for onnx compatibility: argmax doesn't support int64 inputs with opset 14
            pooled_output = last_hidden_state[
                torch.arange(last_hidden_state.shape[0], device=last_hidden_state.device),
                input_ids.to(dtype=torch.int, device=last_hidden_state.device).argmax(dim=-1),
            ]
        else:
            # The config gets updated `eos_token_id` from PR #24773 (so the use of exta new tokens is possible)
            pooled_output = last_hidden_state[
                torch.arange(last_hidden_state.shape[0], device=last_hidden_state.device),
                # We need to get the first position of `eos_token_id` value (`pad_token_ids` might equal to `eos_token_id`)
                (input_ids.to(dtype=torch.int, device=last_hidden_state.device) == self.text_model.eos_token_id)
                .int()
                .argmax(dim=-1),
            ]
    
        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,
        )
    
    # Applied to layers [begin_layer_idx, end_layer_idx) in the encoder.
    # The layer indexed by end_layer_idx is not included.
    # If both layer indices are -1, then apply to all layers (0-11).
    def extend_clip_attention_MKV_multiplier(self, begin_layer_idx=-1, end_layer_idx=-1, multiplier=2, noise_std=0.1):
        num_extended_layers = 0

        for layer_idx, layer in enumerate(self.text_model.encoder.layers):
            if begin_layer_idx >= 0 and layer_idx < begin_layer_idx:
                continue
            if end_layer_idx >= 0 and layer_idx >= end_layer_idx:
                break
            # This shouldn't happen, unless self_attn has already been extended as CLIPAttentionMKV.
            if not isinstance(layer.self_attn, (CLIPAttention, CLIPAttentionMKV)):
                breakpoint()
            old_attn_layer = layer.self_attn
            if not isinstance(old_attn_layer, CLIPAttentionMKV):
                layer.self_attn = CLIPAttentionMKV(old_attn_layer.config, 1)
            layer.self_attn.extend_weights(old_attn_layer, layer_idx, multiplier, noise_std, verbose=True)
            num_extended_layers += 1
    
        return num_extended_layers