from typing import Optional, Tuple, Union, List import torch from transformers import CLIPTextConfig from transformers.modeling_outputs import BaseModelOutputWithPooling from transformers.models.clip.modeling_clip import ( _expand_mask, CLIPTextEmbeddings, CLIPTextTransformer, CLIPTextModel, ) from custom_nodes.KepPromptLang.lib.action.base import Action from custom_nodes.KepPromptLang.lib.actions.types import SegOrAction def slerp(val, low, high): low = low.unsqueeze(0) high = high.unsqueeze(0) low_norm = low/torch.norm(low, dim=1, keepdim=True) high_norm = high/torch.norm(high, dim=1, keepdim=True) omega = torch.acos((low_norm*high_norm).sum(1)) so = torch.sin(omega) res = (torch.sin((1.0-val)*omega)/so).unsqueeze(1)*low + (torch.sin(val*omega)/so).unsqueeze(1) * high return res class PromptLangCLIPTextEmbeddings(CLIPTextEmbeddings): def __init__(self, config: CLIPTextConfig): super().__init__(config) def forward( self, input_dicts: Optional[List[List[SegOrAction]]] = None, input_ids: Optional[torch.LongTensor] = None, position_ids: Optional[torch.LongTensor] = None, inputs_embeds: Optional[torch.FloatTensor] = None, ) -> torch.Tensor: if input_dicts is None: raise ValueError("You have to specify input_dicts") batches = [] for batch_idx, batch in enumerate(input_dicts): results = [] for seg_or_action in batch: if isinstance(seg_or_action, Action): results.append(seg_or_action.get_result(self.token_embedding)) else: results.append(seg_or_action.get_embeddings(self.token_embedding)) batches.append(results) seq_length = batches[0][0].shape[-2] if position_ids is None: position_ids = self.position_ids[:, :seq_length] embeds = [] for batch in batches: if len(batch) == 1: embeds.append(batch[0]) else: embeds.append(torch.cat(batch, dim=-2)) position_embeddings = self.position_embedding(position_ids) embeddings = torch.cat(embeds, dim=0) + position_embeddings return embeddings class PrompLangCLIPTextTransformer(CLIPTextTransformer): def __init__(self, config: CLIPTextConfig): super().__init__(config) self.embeddings = PromptLangCLIPTextEmbeddings(config) def forward( self, input_ids: Optional[List[List[SegOrAction]]] = 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, ) -> 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 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.embeddings(input_dicts=input_ids) bsz = len(input_ids) # TODO: Properly gather this seq_len = 77 input_shape = torch.Size([bsz, seq_len]) # CLIP's text model uses causal mask, prepare it here. # https://github.com/openai/CLIP/blob/cfcffb90e69f37bf2ff1e988237a0fbe41f33c04/clip/model.py#L324 ## VERSION DIFF ## # transformers < 4.30.0 if hasattr(self, "_build_causal_attention_mask"): print("Using transformers < 4.30.0") causal_attention_mask = self._build_causal_attention_mask(bsz, seq_len, hidden_states.dtype).to(hidden_states.device) else: # transformers >= 4.30.0 print("Using transformers >= 4.30.0") from transformers.models.clip.modeling_clip import _make_causal_mask 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.encoder( inputs_embeds=hidden_states, attention_mask=attention_mask, causal_attention_mask=causal_attention_mask, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, ) last_hidden_state = encoder_outputs[0] last_hidden_state = self.final_layer_norm(last_hidden_state) # Hacky way to get idx of first EOT token eot_idx = [1] for batch in input_ids[1:]: idx = 0 for seg_or_action in batch: if isinstance(seg_or_action, Action): idx += seg_or_action.token_length() else: if seg_or_action.text == '__PAD__': break eot_idx.append(idx) # 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 # TODO: Get the index of the first EOT token pooled_output = last_hidden_state[ torch.arange(last_hidden_state.shape[0], device=last_hidden_state.device), eot_idx ] 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, ) # This is necessary to pass the PromptLangCLIPTextTransformer class PromptLangTextModel(CLIPTextModel): def __init__(self, config: CLIPTextConfig): super().__init__(config) self.text_model = PrompLangCLIPTextTransformer(config) def forward( self, input_ids: Optional[List[List[SegOrAction]]] = 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, ) -> Union[Tuple, BaseModelOutputWithPooling]: return_dict = return_dict if return_dict is not None else self.config.use_return_dict return self.text_model( input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, )