codingutf8 2022 microsoft research and the huggingface team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license pytorch xclip model from copy import copy from dataclasses import dataclass from typing import any optional tuple union import torch import torch utils checkpoint from torch import nn from activations import act2fn from modelingattnmaskutils import create4dcausalattentionmask prepare4dattentionmask from modelingoutputs import basemodeloutput basemodeloutputwithpooling from modelingutils import pretrainedmodel from utils import modeloutput addstartdocstrings addstartdocstringstomodelforward logging replacereturndocstrings from configurationxclip import xclipconfig xcliptextconfig xclipvisionconfig logger logging getloggername checkpointfordoc microsoftxclipbasepatch32 xclippretrainedmodelarchivelist microsoftxclipbasepatch32 see all xclip models at https huggingface comodels filterxclip contrastive loss function adapted from https sachinruk github ioblogpytorchpytorch20lightningloss20functiongpu20210307clip html def contrastivelosslogits torch tensor torch tensor return nn functional crossentropylogits torch arangelenlogits devicelogits device copied from transformers models clip modelingclip cliploss with clipxclip def xcliplosssimilarity torch tensor torch tensor captionloss contrastivelosssimilarity imageloss contrastivelosssimilarity t return captionloss imageloss 2 0 dataclass class xclipoutputmodeloutput loss optionaltorch floattensor none logitspervideo torch floattensor none logitspertext torch floattensor none textembeds torch floattensor none videoembeds torch floattensor none textmodeloutput basemodeloutputwithpooling none visionmodeloutput basemodeloutputwithpooling none mitoutput basemodeloutputwithpooling none def totupleself tupleany return tuple selfk if k not in textmodeloutput visionmodeloutput mitoutput else getattrself k totuple for k in self keys copied from transformers models clip modelingclip clipvisionembeddings with clipxclip class xclipvisionembeddingsnn module def initself config xclipvisionconfig super init self config config self embeddim config hiddensize self imagesize config imagesize self patchsize config patchsize self classembedding nn parametertorch randnself embeddim self patchembedding nn conv2d inchannelsconfig numchannels outchannelsself embeddim kernelsizeself patchsize strideself patchsize biasfalse self numpatches self imagesize self patchsize 2 self numpositions self numpatches 1 self positionembedding nn embeddingself numpositions self embeddim self registerbufferpositionids torch arangeself numpositions expand1 1 persistentfalse def forwardself pixelvalues torch floattensor torch tensor batchsize pixelvalues shape0 targetdtype self patchembedding weight dtype patchembeds self patchembeddingpixelvalues todtypetargetdtype shape width grid grid patchembeds patchembeds flatten2 transpose1 2 classembeds self classembedding expandbatchsize 1 1 embeddings torch catclassembeds patchembeds dim1 embeddings embeddings self positionembeddingself positionids return embeddings copied from transformers models clip modelingclip cliptextembeddings with clipxclip class xcliptextembeddingsnn module def initself config xcliptextconfig super init embeddim config hiddensize self tokenembedding nn embeddingconfig vocabsize embeddim self positionembedding nn embeddingconfig maxpositionembeddings embeddim positionids 1 len position emb is contiguous in memory and exported when serialized self registerbuffer positionids torch arangeconfig maxpositionembeddings expand1 1 persistentfalse def forward self inputids optionaltorch longtensor none positionids optionaltorch longtensor none inputsembeds optionaltorch floattensor none torch tensor seqlength inputids shape1 if inputids is not none else inputsembeds shape2 if positionids is none positionids self positionids seqlength if inputsembeds is none inputsembeds self tokenembeddinginputids positionembeddings self positionembeddingpositionids embeddings inputsembeds positionembeddings return embeddings copied from transformers models clip modelingclip clipattention with clipxclip class xclipattentionnn module input shape batch x time x channel bsz tgtlen embeddim hiddenstates size get query proj querystates self qprojhiddenstates self scale keystates self shapeself kprojhiddenstates 1 bsz valuestates self shapeself vprojhiddenstates 1 bsz projshape bsz self numheads 1 self headdim querystates self shapequerystates tgtlen bsz viewprojshape keystates keystates viewprojshape valuestates valuestates viewprojshape srclen keystates size1 attnweights torch bmmquerystates keystates transpose1 2 if attnweights size bsz self numheads tgtlen srclen raise valueerror fattention weights should be of size bsz self numheads tgtlen srclen but is f attnweights size apply the causalattentionmask first if causalattentionmask is not none if causalattentionmask size bsz 1 tgtlen srclen raise valueerror fattention mask should be of size bsz 1 tgtlen srclen but is f causalattentionmask size attnweights attnweights viewbsz self numheads tgtlen srclen causalattentionmask attnweights attnweights viewbsz self numheads tgtlen srclen if attentionmask is not none if attentionmask size bsz 1 tgtlen srclen raise valueerror fattention mask should be of size bsz 1 tgtlen srclen but is attentionmask size attnweights attnweights viewbsz self numheads tgtlen srclen attentionmask attnweights attnweights viewbsz self numheads tgtlen srclen attnweights nn functional softmaxattnweights dim1 if outputattentions this operation is a bit akward but it s required to make sure that attnweights keeps its gradient in order to do so attnweights have to reshaped twice and have to be reused in the following attnweightsreshaped attnweights viewbsz self numheads tgtlen srclen attnweights attnweightsreshaped viewbsz self numheads tgtlen srclen else attnweightsreshaped none attnprobs nn functional dropoutattnweights pself dropout trainingself training attnoutput torch bmmattnprobs valuestates if attnoutput size bsz self numheads tgtlen self headdim raise valueerror fattnoutput should be of size bsz self numheads tgtlen self headdim but is f attnoutput size attnoutput attnoutput viewbsz self numheads tgtlen self headdim attnoutput attnoutput transpose1 2 attnoutput attnoutput reshapebsz tgtlen embeddim attnoutput self outprojattnoutput return attnoutput attnweightsreshaped copied from transformers models clip modelingclip clipmlp with clipxclip class xclipmlpnn module def initself config super init self config config self activationfn act2fnconfig hiddenact self fc1 nn linearconfig hiddensize config intermediatesize self fc2 nn linearconfig intermediatesize config hiddensize def forwardself hiddenstates torch tensor torch tensor hiddenstates self fc1hiddenstates hiddenstates self activationfnhiddenstates hiddenstates self fc2hiddenstates return hiddenstates copied from transformers models clip modelingclip clipencoderlayer with clipxclip class xclipencoderlayernn module def initself config xclipconfig super init self embeddim config hiddensize self selfattn xclipattentionconfig self layernorm1 nn layernormself embeddim epsconfig layernormeps self mlp xclipmlpconfig self layernorm2 nn layernormself embeddim epsconfig layernormeps def forward self hiddenstates torch tensor attentionmask torch tensor causalattentionmask torch tensor outputattentions optionalbool false tupletorch floattensor residual hiddenstates hiddenstates self layernorm1hiddenstates hiddenstates attnweights self selfattn hiddenstateshiddenstates attentionmaskattentionmask causalattentionmaskcausalattentionmask outputattentionsoutputattentions hiddenstates residual hiddenstates residual hiddenstates hiddenstates self layernorm2hiddenstates hiddenstates self mlphiddenstates hiddenstates residual hiddenstates outputs hiddenstates if outputattentions outputs attnweights return outputs copied from transformers models beit modelingbeit droppath def droppathinput torch tensor dropprob float 0 0 training bool false torch tensor if dropprob 0 0 or not training return input keepprob 1 dropprob shape input shape0 1 input ndim 1 work with diff dim tensors not just 2d convnets randomtensor keepprob torch randshape dtypeinput dtype deviceinput device randomtensor floor binarize output input divkeepprob randomtensor return output copied from transformers models beit modelingbeit beitdroppath with beitxclip class xclipdroppathnn module this corresponds to the crossframelattentionblock class in the original implementation args hiddenstates torch floattensor input to the layer of shape batch seqlen embeddim attentionmask torch floattensor attention mask of size batch 1 tgtlen srclen where padding elements are indicated by very large negative values config encoderattentionheads causalattentionmask torch tensor of shape batchsize sequencelength optional causal mask for the text model mask values selected in 0 1 1 for tokens that are not masked 0 for tokens that are masked what are attention masks glossaryattentionmask outputattentions bool optional whether or not to return the attentions tensors of all attention layers see attentions under returned tensors for more detail add dummy sequence dimension an abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models initialize the weights factor self config initializerfactor if isinstancemodule xcliptextembeddings module tokenembedding weight data normalmean0 0 stdfactor 0 02 module positionembedding weight data normalmean0 0 stdfactor 0 02 elif isinstancemodule xclipvisionembeddings factor self config initializerfactor nn init normalmodule classembedding mean0 0 stdmodule embeddim0 5 factor nn init normalmodule patchembedding weight stdmodule config initializerrange factor nn init normalmodule positionembedding weight stdmodule config initializerrange factor elif isinstancemodule xclipattention factor self config initializerfactor inprojstd module embeddim0 5 2 module config numhiddenlayers 0 5 factor outprojstd module embeddim0 5 factor nn init normalmodule qproj weight stdinprojstd nn init normalmodule kproj weight stdinprojstd nn init normalmodule vproj weight stdinprojstd nn init normalmodule outproj weight stdoutprojstd elif isinstancemodule xclipmlp factor self config initializerfactor inprojstd module config hiddensize0 5 2 module config numhiddenlayers 0 5 factor fcstd 2 module config hiddensize 0 5 factor nn init normalmodule fc1 weight stdfcstd nn init normalmodule fc2 weight stdinprojstd elif isinstancemodule xclipmodel factor self config initializerfactor nn init normal module textprojection weight stdmodule textembeddim0 5 factor nn init normal module visualprojection weight stdmodule visionembeddim0 5 factor nn init normalmodule promptsvisualprojection mean0 0 stdmodule visionembeddim0 5 factor elif isinstancemodule xclipmultiframeintegrationtransformer nn init normalmodule positionembedding stdself config initializerfactor if isinstancemodule nn layernorm module bias data zero module weight data fill1 0 if isinstancemodule nn linear module weight data normalmean0 0 stdself config initializerfactor if module bias is not none module bias data zero xclipstartdocstring r this model is a pytorch torch nn modulehttps pytorch orgdocsstablenn htmltorch nn module subclass use it as a regular pytorch module and refer to the pytorch documentation for all matter related to general usage and behavior parameters config xclipconfig model configuration class with all the parameters of the model initializing with a config file does not load the weights associated with the model only the configuration check out the pretrainedmodel frompretrained method to load the model weights xclipvisioninputsdocstring r args pixelvalues torch floattensor of shape batchsize numchannels 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 outputattentions bool optional whether or not to return the attentions tensors of all attention layers see attentions under returned tensors for more detail outputhiddenstates bool optional whether or not to return the hidden states of all layers see hiddenstates under returned tensors for more detail returndict bool optional whether or not to return a utils modeloutput instead of a plain tuple copied from transformers models clip modelingclip clipencoder with clipxclip class xclipencodernn module def initself config xclipconfig super init self config config self layers nn modulelistxclipencoderlayerconfig for in rangeconfig numhiddenlayers self gradientcheckpointing false def forward self inputsembeds attentionmask optionaltorch tensor none causalattentionmask optionaltorch tensor none outputattentions optionalbool none outputhiddenstates optionalbool none returndict optionalbool none uniontuple basemodeloutput r args inputsembeds torch floattensor of shape batchsize sequencelength hiddensize optionally instead of passing inputids you can choose to directly pass an embedded representation this is useful if you want more control over how to convert inputids indices into associated vectors than the model s internal embedding lookup matrix attentionmask torch tensor of shape batchsize sequencelength 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 glossaryattentionmask causalattentionmask torch tensor of shape batchsize sequencelength optional causal mask for the text model mask values selected in 0 1 1 for tokens that are not masked 0 for tokens that are masked what are attention masks glossaryattentionmask outputattentions bool optional whether or not to return the attentions tensors of all attention layers see attentions under returned tensors for more detail outputhiddenstates bool optional whether or not to return the hidden states of all layers see hiddenstates under returned tensors for more detail returndict bool optional whether or not to return a utils modeloutput instead of a plain tuple outputattentions outputattentions if outputattentions is not none else self config outputattentions outputhiddenstates outputhiddenstates if outputhiddenstates is not none else self config outputhiddenstates returndict returndict if returndict is not none else self config usereturndict if inputids is none raise valueerroryou have to specify either inputids inputshape inputids size inputids inputids view1 inputshape1 hiddenstates self embeddingsinputidsinputids positionidspositionids xclip s text model uses causal mask prepare it here https github comopenaiclipblobcfcffb90e69f37bf2ff1e988237a0fbe41f33c04clipmodel pyl324 causalattentionmask create4dcausalattentionmask inputshape hiddenstates dtype devicehiddenstates device expand attentionmask if attentionmask is not none batchsize seqlen batchsize 1 tgtseqlen srcseqlen attentionmask prepare4dattentionmaskattentionmask hiddenstates dtype encoderoutputs self encoder inputsembedshiddenstates attentionmaskattentionmask causalattentionmaskcausalattentionmask outputattentionsoutputattentions outputhiddenstatesoutputhiddenstates returndictreturndict lasthiddenstate encoderoutputs0 lasthiddenstate self finallayernormlasthiddenstate textembeds shape batchsize sequencelength transformer width take features from the eot embedding eottoken is the highest number in each sequence pooledoutput lasthiddenstatetorch arangelasthiddenstate shape0 inputids argmaxdim1 if not returndict return lasthiddenstate pooledoutput encoderoutputs1 return basemodeloutputwithpooling lasthiddenstatelasthiddenstate pooleroutputpooledoutput hiddenstatesencoderoutputs hiddenstates attentionsencoderoutputs attentions class xcliptextmodelxclippretrainedmodel configclass xcliptextconfig def initself config xcliptextconfig super initconfig self textmodel xcliptexttransformerconfig initialize weights and apply final processing self postinit def getinputembeddingsself nn module return self textmodel embeddings tokenembedding def setinputembeddingsself value self textmodel embeddings tokenembedding value addstartdocstringstomodelforwardxcliptextinputsdocstring replacereturndocstringsoutputtypebasemodeloutputwithpooling configclassxcliptextconfig def forward self inputids optionaltorch tensor none attentionmask optionaltorch tensor none positionids optionaltorch tensor none outputattentions optionalbool none outputhiddenstates optionalbool none returndict optionalbool none uniontuple basemodeloutputwithpooling r returns examples python from transformers import autotokenizer xcliptextmodel model xcliptextmodel frompretrainedmicrosoftxclipbasepatch32 tokenizer autotokenizer frompretrainedmicrosoftxclipbasepatch32 inputs tokenizera photo of a cat a photo of a dog paddingtrue returntensorspt outputs modelinputs lasthiddenstate outputs lasthiddenstate pooledoutput outputs pooleroutput pooled eos token states return self textmodel inputidsinputids attentionmaskattentionmask positionidspositionids outputattentionsoutputattentions outputhiddenstatesoutputhiddenstates returndictreturndict class xclipvisionencodernn module def initself config xclipconfig super init self config config self layers nn modulelistxclipvisionencoderlayerconfig for in rangeconfig numhiddenlayers self gradientcheckpointing false def forward self inputsembeds attentionmask optionaltorch tensor none causalattentionmask optionaltorch tensor none outputattentions optionalbool none outputhiddenstates optionalbool none returndict optionalbool none uniontuple basemodeloutput r args inputsembeds torch floattensor of shape batchsize sequencelength hiddensize optionally instead of passing inputids you can choose to directly pass an embedded representation this is useful if you want more control over how to convert inputids indices into associated vectors than the model s internal embedding lookup matrix attentionmask torch tensor of shape batchsize sequencelength 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 glossaryattentionmask causalattentionmask torch tensor of shape batchsize sequencelength optional causal mask for the text model mask values selected in 0 1 1 for tokens that are not masked 0 for tokens that are masked what are attention masks glossaryattentionmask outputattentions bool optional whether or not to return the attentions tensors of all attention layers see attentions under returned tensors for more detail outputhiddenstates bool optional whether or not to return the hidden states of all layers see hiddenstates under returned tensors for more detail returndict bool optional whether or not to return a utils modeloutput instead of a plain tuple this corresponds to the crossframecommunicationtransformer class in the original implementation outputattentions outputattentions if outputattentions is not none else self config outputattentions outputhiddenstates outputhiddenstates if outputhiddenstates is not none else self config outputhiddenstates returndict returndict if returndict is not none else self config usereturndict hiddenstates self embeddingspixelvalues hiddenstates self prelayernormhiddenstates encoderoutputs self encoder inputsembedshiddenstates outputattentionsoutputattentions outputhiddenstatesoutputhiddenstates returndictreturndict lasthiddenstate encoderoutputs0 pooledoutput lasthiddenstate 0 pooledoutput self postlayernormpooledoutput if not returndict return lasthiddenstate pooledoutput encoderoutputs1 return basemodeloutputwithpooling lasthiddenstatelasthiddenstate pooleroutputpooledoutput hiddenstatesencoderoutputs hiddenstates attentionsencoderoutputs attentions class xclipvisionmodelxclippretrainedmodel configclass xclipvisionconfig maininputname pixelvalues def initself config xclipvisionconfig super initconfig self visionmodel xclipvisiontransformerconfig initialize weights and apply final processing self postinit def getinputembeddingsself nn module return self visionmodel embeddings patchembedding addstartdocstringstomodelforwardxclipvisioninputsdocstring replacereturndocstringsoutputtypebasemodeloutputwithpooling configclassxclipvisionconfig def forward self pixelvalues optionaltorch floattensor none outputattentions optionalbool none outputhiddenstates optionalbool none returndict optionalbool none uniontuple basemodeloutputwithpooling r returns examples python import av import torch import numpy as np from transformers import autoprocessor xclipvisionmodel from huggingfacehub import hfhubdownload np random seed0 def readvideopyavcontainer indices decode the video with pyav decoder args container av container input inputcontainer pyav container indices listint list of frame indices to decode returns result np ndarray np array of decoded frames of shape numframes height width 3 frames container seek0 startindex indices0 endindex indices1 for i frame in enumeratecontainer decodevideo0 if i endindex break if i startindex and i in indices frames appendframe return np stackx tondarrayformatrgb24 for x in frames def sampleframeindicescliplen framesamplerate seglen sample a given number of frame indices from the video args cliplen int total number of frames to sample framesamplerate int sample every nth frame seglen int maximum allowed index of sample s last frame returns indices listint list of sampled frame indices convertedlen intcliplen framesamplerate endidx np random randintconvertedlen seglen startidx endidx convertedlen indices np linspacestartidx endidx numcliplen indices np clipindices startidx endidx 1 astypenp int64 return indices video clip consists of 300 frames 10 seconds at 30 fps filepath hfhubdownload repoidnielsrvideodemo filenameeatingspaghetti mp4 repotypedataset container av openfilepath sample 16 frames indices sampleframeindicescliplen8 framesamplerate1 seglencontainer streams video0 frames video readvideopyavcontainer indices processor autoprocessor frompretrainedmicrosoftxclipbasepatch32 model xclipvisionmodel frompretrainedmicrosoftxclipbasepatch32 pixelvalues processorvideoslistvideo returntensorspt pixelvalues batchsize numframes numchannels height width pixelvalues shape pixelvalues pixelvalues reshape1 numchannels height width outputs modelpixelvalues lasthiddenstate outputs lasthiddenstate return self visionmodel pixelvaluespixelvalues outputattentionsoutputattentions outputhiddenstatesoutputhiddenstates returndictreturndict class xclipmultiframeintegrationtransformernn module def initself config xclipvisionconfig super init self positionembedding nn parametertorch empty1 config numframes config hiddensize self encoder xclipencoderconfig def forward self hiddenstates outputattentions optionalbool none outputhiddenstates optionalbool none returndict optionalbool none uniontuple basemodeloutput residual hiddenstates add position embeddings hiddenstates hiddenstates self positionembedding encoderoutputs self encoder inputsembedshiddenstates outputattentionsoutputattentions outputhiddenstatesoutputhiddenstates returndictreturndict lasthiddenstate encoderoutputs0 lasthiddenstate lasthiddenstate typehiddenstates dtype residual pooledoutput lasthiddenstate meandim1 keepdimfalse if not returndict return lasthiddenstate pooledoutput encoderoutputs1 return basemodeloutputwithpooling lasthiddenstatelasthiddenstate pooleroutputpooledoutput hiddenstatesencoderoutputs hiddenstates attentionsencoderoutputs attentions class xclipcrossattentionnn module input shape batch x time x channel batchsize queryseqlen hiddensize queries shape batchsize keyseqlen hiddensize keys shape queries self qprojqueries reshapebatchsize queryseqlen self numheads hiddensize self numheads permute0 2 1 3 keys self kprojkeys reshapebatchsize keyseqlen self numheads hiddensize self numheads permute0 2 1 3 values self vprojvalues reshapebatchsize keyseqlen self numheads hiddensize self numheads permute0 2 1 3 attn queries keys transpose2 1 self scale attn attn softmaxdim1 attn self attndropattn x attn values transpose1 2 reshapebatchsize queryseqlen hiddensize x self projx x self projdropx return x class promptgeneratorlayernn module def initself config super init embeddim config projectiondim self crossattn xclipcrossattentionconfig self norm1 nn layernormembeddim epsconfig textconfig layernormeps self norm3 nn layernormembeddim epsconfig textconfig layernormeps self mlp nn sequential nn linearembeddim embeddim 4 act2fnconfig prompthiddenact nn dropoutconfig promptattentiondropout nn linearembeddim 4 embeddim def forwardself x visual x x self crossattnself norm1x visual visual x x self mlpself norm3x return x class xclippromptgeneratornn module initialize weights and apply final processing use xclip model s config for some fields if specified instead of those of vision text components use xclip model s config for some fields if specified instead of those of vision text components use xclip model s config for some fields if specified instead of those of vision text components normalized features cosine similarity as logits coding utf 8 2022 microsoft research and the huggingface team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license pytorch x clip model see all x clip models at https huggingface co models filter x clip contrastive loss function adapted from https sachinruk github io blog pytorch pytorch 20lightning loss 20function gpu 2021 03 07 clip html copied from transformers models clip modeling_clip clip_loss with clip x_clip args loss torch floattensor of shape 1 optional returned when return_loss is true contrastive loss for video text similarity logits_per_video torch floattensor of shape video_batch_size text_batch_size the scaled dot product scores between video_embeds and text_embeds this represents the video text similarity scores logits_per_text torch floattensor of shape text_batch_size video_batch_size the scaled dot product scores between text_embeds and video_embeds this represents the text video similarity scores text_embeds torch floattensor of shape batch_size output_dim the text embeddings obtained by applying the projection layer to the pooled output of xcliptextmodel video_embeds torch floattensor of shape batch_size output_dim the video embeddings obtained by applying the projection layer to the pooled output of xclipvisionmodel text_model_output basemodeloutputwithpooling the output of the xcliptextmodel vision_model_output basemodeloutputwithpooling the output of the xclipvisionmodel mit_output basemodeloutputwithpooling the output of xclipmultiframeintegrationtransformer mit for short copied from transformers models clip modeling_clip clipvisionembeddings with clip xclip shape width grid grid copied from transformers models clip modeling_clip cliptextembeddings with clip xclip position_ids 1 len position emb is contiguous in memory and exported when serialized copied from transformers models clip modeling_clip clipattention with clip xclip multi headed attention from attention is all you need paper input shape batch x time x channel get query proj apply the causal_attention_mask first this operation is a bit akward 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 copied from transformers models clip modeling_clip clipmlp with clip xclip copied from transformers models clip modeling_clip clipencoderlayer with clip xclip 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 copied from transformers models beit modeling_beit drop_path drop paths stochastic depth per sample when applied in main path of residual blocks comment by ross wightman this is the same as the dropconnect impl i created for efficientnet etc networks however the original name is misleading as drop connect is a different form of dropout in a separate paper see discussion https github com tensorflow tpu issues 494 issuecomment 532968956 i ve opted for changing the layer and argument names to drop path rather than mix dropconnect as a layer name and use survival rate as the argument work with diff dim tensors not just 2d convnets binarize copied from transformers models beit modeling_beit beitdroppath with beit xclip drop paths stochastic depth per sample when applied in main path of residual blocks this corresponds to the crossframelattentionblock class in the original implementation 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 causal_attention_mask torch tensor of shape batch_size sequence_length optional causal mask for the text model 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 add dummy sequence dimension an abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models initialize the weights this model is a pytorch torch nn module https pytorch org docs stable nn html torch nn module subclass use it as a regular pytorch module and refer to the pytorch documentation for all matter related to general usage and behavior parameters config xclipconfig model configuration class with all the parameters of the model initializing with a config file does not load the weights associated with the model only the configuration check out the pretrainedmodel from_pretrained method to load the model weights args input_ids torch longtensor of shape batch_size sequence_length indices of input sequence tokens in the vocabulary padding will be ignored by default should you provide it indices can be obtained using autotokenizer see pretrainedtokenizer encode and pretrainedtokenizer __call__ for details what are input ids glossary input ids 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 position_ids torch longtensor of shape batch_size sequence_length optional indices of positions of each input sequence tokens in the position embeddings selected in the range 0 config max_position_embeddings 1 what are position ids glossary position ids 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 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 args input_ids torch longtensor of shape batch_size sequence_length indices of input sequence tokens in the vocabulary padding will be ignored by default should you provide it indices can be obtained using autotokenizer see pretrainedtokenizer encode and pretrainedtokenizer __call__ for details what are input ids glossary input ids 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 position_ids torch longtensor of shape batch_size sequence_length optional indices of positions of each input sequence tokens in the position embeddings selected in the range 0 config max_position_embeddings 1 what are position ids glossary position ids 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 return_loss bool optional whether or not to return the contrastive loss 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 copied from transformers models clip modeling_clip clipencoder with clip xclip transformer encoder consisting of config num_hidden_layers self attention layers each layer is a xclipencoderlayer args config xclipconfig args inputs_embeds torch floattensor of shape batch_size sequence_length hidden_size optionally instead of passing input_ids you can choose to directly pass an embedded representation this is useful if you want more control over how to convert input_ids indices into associated vectors than the model s internal embedding lookup matrix 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 causal_attention_mask torch tensor of shape batch_size sequence_length optional causal mask for the text model 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 returns x_clip s text model uses causal mask prepare it here https github com openai clip blob cfcffb90e69f37bf2ff1e988237a0fbe41f33c04 clip model py l324 expand attention_mask batch_size seq_len batch_size 1 tgt_seq_len src_seq_len text_embeds shape batch_size sequence_length transformer width take features from the eot embedding eot_token is the highest number in each sequence initialize weights and apply final processing returns examples python from transformers import autotokenizer xcliptextmodel model xcliptextmodel from_pretrained microsoft xclip base patch32 tokenizer autotokenizer from_pretrained microsoft xclip base patch32 inputs tokenizer a photo of a cat a photo of a dog padding true return_tensors pt outputs model inputs last_hidden_state outputs last_hidden_state pooled_output outputs pooler_output pooled eos token states transformer encoder consisting of config num_hidden_layers self attention layers each layer is a xclipvisionencoderlayer args config xclipconfig args inputs_embeds torch floattensor of shape batch_size sequence_length hidden_size optionally instead of passing input_ids you can choose to directly pass an embedded representation this is useful if you want more control over how to convert input_ids indices into associated vectors than the model s internal embedding lookup matrix 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 causal_attention_mask torch tensor of shape batch_size sequence_length optional causal mask for the text model 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 this corresponds to the crossframecommunicationtransformer class in the original implementation returns initialize weights and apply final processing returns examples python import av import torch import numpy as np from transformers import autoprocessor xclipvisionmodel from huggingface_hub import hf_hub_download np random seed 0 def read_video_pyav container indices decode the video with pyav decoder args container av container input inputcontainer pyav container indices list int list of frame indices to decode returns result np ndarray np array of decoded frames of shape num_frames height width 3 frames container seek 0 start_index indices 0 end_index indices 1 for i frame in enumerate container decode video 0 if i end_index break if i start_index and i in indices frames append frame return np stack x to_ndarray format rgb24 for x in frames def sample_frame_indices clip_len frame_sample_rate seg_len sample a given number of frame indices from the video args clip_len int total number of frames to sample frame_sample_rate int sample every n th frame seg_len int maximum allowed index of sample s last frame returns indices list int list of sampled frame indices converted_len int clip_len frame_sample_rate end_idx np random randint converted_len seg_len start_idx end_idx converted_len indices np linspace start_idx end_idx num clip_len indices np clip indices start_idx end_idx 1 astype np int64 return indices video clip consists of 300 frames 10 seconds at 30 fps file_path hf_hub_download repo_id nielsr video demo filename eating_spaghetti mp4 repo_type dataset container av open file_path sample 16 frames indices sample_frame_indices clip_len 8 frame_sample_rate 1 seg_len container streams video 0 frames video read_video_pyav container indices processor autoprocessor from_pretrained microsoft xclip base patch32 model xclipvisionmodel from_pretrained microsoft xclip base patch32 pixel_values processor videos list video return_tensors pt pixel_values batch_size num_frames num_channels height width pixel_values shape pixel_values pixel_values reshape 1 num_channels height width outputs model pixel_values last_hidden_state outputs last_hidden_state this corresponds to the multiframeintegrationtransformer class in the original implementation add position embeddings multi headed attention from attention is all you need paper input shape batch x time x channel this corresponds to the videospecificprompt class in the original implementation initialize weights and apply final processing returns text_features torch floattensor of shape batch_size output_dim the text embeddings obtained by applying the projection layer to the pooled output of xcliptextmodel examples python from transformers import autotokenizer automodel tokenizer autotokenizer from_pretrained microsoft xclip base patch32 model automodel from_pretrained microsoft xclip base patch32 inputs tokenizer a photo of a cat a photo of a dog padding true return_tensors pt text_features model get_text_features inputs use x_clip model s config for some fields if specified instead of those of vision text components returns video_features torch floattensor of shape batch_size output_dim the video embeddings obtained by applying the projection layer to the pooled output of xclipvisionmodel and xclipmultiframeintegrationtransformer examples python import av import torch import numpy as np from transformers import autoprocessor automodel from huggingface_hub import hf_hub_download np random seed 0 def read_video_pyav container indices decode the video with pyav decoder args container av container input inputcontainer pyav container indices list int list of frame indices to decode returns result np ndarray np array of decoded frames of shape num_frames height width 3 frames container seek 0 start_index indices 0 end_index indices 1 for i frame in enumerate container decode video 0 if i end_index break if i start_index and i in indices frames append frame return np stack x to_ndarray format rgb24 for x in frames def sample_frame_indices clip_len frame_sample_rate seg_len sample a given number of frame indices from the video args clip_len int total number of frames to sample frame_sample_rate int sample every n th frame seg_len int maximum allowed index of sample s last frame returns indices list int list of sampled frame indices converted_len int clip_len frame_sample_rate end_idx np random randint converted_len seg_len start_idx end_idx converted_len indices np linspace start_idx end_idx num clip_len indices np clip indices start_idx end_idx 1 astype np int64 return indices video clip consists of 300 frames 10 seconds at 30 fps file_path hf_hub_download repo_id nielsr video demo filename eating_spaghetti mp4 repo_type dataset container av open file_path sample 8 frames indices sample_frame_indices clip_len 8 frame_sample_rate 1 seg_len container streams video 0 frames video read_video_pyav container indices processor autoprocessor from_pretrained microsoft xclip base patch32 model automodel from_pretrained microsoft xclip base patch32 inputs processor videos list video return_tensors pt video_features model get_video_features inputs use x_clip model s config for some fields if specified instead of those of vision text components returns examples python import av import torch import numpy as np from transformers import autoprocessor automodel from huggingface_hub import hf_hub_download np random seed 0 def read_video_pyav container indices decode the video with pyav decoder args container av container input inputcontainer pyav container indices list int list of frame indices to decode returns result np ndarray np array of decoded frames of shape num_frames height width 3 frames container seek 0 start_index indices 0 end_index indices 1 for i frame in enumerate container decode video 0 if i end_index break if i start_index and i in indices frames append frame return np stack x to_ndarray format rgb24 for x in frames def sample_frame_indices clip_len frame_sample_rate seg_len sample a given number of frame indices from the video args clip_len int total number of frames to sample frame_sample_rate int sample every n th frame seg_len int maximum allowed index of sample s last frame returns indices list int list of sampled frame indices converted_len int clip_len frame_sample_rate end_idx np random randint converted_len seg_len start_idx end_idx converted_len indices np linspace start_idx end_idx num clip_len indices np clip indices start_idx end_idx 1 astype np int64 return indices video clip consists of 300 frames 10 seconds at 30 fps file_path hf_hub_download repo_id nielsr video demo filename eating_spaghetti mp4 repo_type dataset container av open file_path sample 8 frames indices sample_frame_indices clip_len 8 frame_sample_rate 1 seg_len container streams video 0 frames video read_video_pyav container indices processor autoprocessor from_pretrained microsoft xclip base patch32 model automodel from_pretrained microsoft xclip base patch32 inputs processor text playing sports eating spaghetti go shopping videos list video return_tensors pt padding true forward pass with torch no_grad outputs model inputs logits_per_video outputs logits_per_video this is the video text similarity score probs logits_per_video softmax dim 1 we can take the softmax to get the label probabilities print probs tensor 1 9496e 04 9 9960e 01 2 0825e 04 use x_clip model s config for some fields if specified instead of those of vision text components normalized features cosine similarity as logits
from copy import copy
from dataclasses import dataclass
from typing import Any, Optional, Tuple, Union

import torch
import torch.utils.checkpoint
from torch import nn

from ...activations import ACT2FN
from ...modeling_attn_mask_utils import _create_4d_causal_attention_mask, _prepare_4d_attention_mask
from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling
from ...modeling_utils import PreTrainedModel
from ...utils import (
    ModelOutput,
    add_start_docstrings,
    add_start_docstrings_to_model_forward,
    logging,
    replace_return_docstrings,
)
from .configuration_x_clip import XCLIPConfig, XCLIPTextConfig, XCLIPVisionConfig


logger = logging.get_logger(__name__)

_CHECKPOINT_FOR_DOC = "microsoft/xclip-base-patch32"

XCLIP_PRETRAINED_MODEL_ARCHIVE_LIST = [
    "microsoft/xclip-base-patch32",
    
]




def contrastive_loss(logits: torch.Tensor) -> torch.Tensor:
    return nn.functional.cross_entropy(logits, torch.arange(len(logits), device=logits.device))



def x_clip_loss(similarity: torch.Tensor) -> torch.Tensor:
    caption_loss = contrastive_loss(similarity)
    image_loss = contrastive_loss(similarity.t())
    return (caption_loss + image_loss) / 2.0


@dataclass
class XCLIPOutput(ModelOutput):
    

    loss: Optional[torch.FloatTensor] = None
    logits_per_video: torch.FloatTensor = None
    logits_per_text: torch.FloatTensor = None
    text_embeds: torch.FloatTensor = None
    video_embeds: torch.FloatTensor = None
    text_model_output: BaseModelOutputWithPooling = None
    vision_model_output: BaseModelOutputWithPooling = None
    mit_output: BaseModelOutputWithPooling = None

    def to_tuple(self) -> Tuple[Any]:
        return tuple(
            self[k]
            if k not in ["text_model_output", "vision_model_output", "mit_output"]
            else getattr(self, k).to_tuple()
            for k in self.keys()
        )



class XCLIPVisionEmbeddings(nn.Module):
    def __init__(self, config: XCLIPVisionConfig):
        super().__init__()
        self.config = config
        self.embed_dim = config.hidden_size
        self.image_size = config.image_size
        self.patch_size = config.patch_size

        self.class_embedding = nn.Parameter(torch.randn(self.embed_dim))

        self.patch_embedding = nn.Conv2d(
            in_channels=config.num_channels,
            out_channels=self.embed_dim,
            kernel_size=self.patch_size,
            stride=self.patch_size,
            bias=False,
        )

        self.num_patches = (self.image_size // self.patch_size) ** 2
        self.num_positions = self.num_patches + 1
        self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim)
        self.register_buffer("position_ids", torch.arange(self.num_positions).expand((1, -1)), persistent=False)

    def forward(self, pixel_values: torch.FloatTensor) -> torch.Tensor:
        batch_size = pixel_values.shape[0]
        target_dtype = self.patch_embedding.weight.dtype
        patch_embeds = self.patch_embedding(pixel_values.to(dtype=target_dtype))  
        patch_embeds = patch_embeds.flatten(2).transpose(1, 2)

        class_embeds = self.class_embedding.expand(batch_size, 1, -1)
        embeddings = torch.cat([class_embeds, patch_embeds], dim=1)
        embeddings = embeddings + self.position_embedding(self.position_ids)
        return embeddings



class XCLIPTextEmbeddings(nn.Module):
    def __init__(self, config: XCLIPTextConfig):
        super().__init__()
        embed_dim = config.hidden_size

        self.token_embedding = nn.Embedding(config.vocab_size, embed_dim)
        self.position_embedding = nn.Embedding(config.max_position_embeddings, embed_dim)

        
        self.register_buffer(
            "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False
        )

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
    ) -> torch.Tensor:
        seq_length = input_ids.shape[-1] if input_ids is not None else inputs_embeds.shape[-2]

        if position_ids is None:
            position_ids = self.position_ids[:, :seq_length]

        if inputs_embeds is None:
            inputs_embeds = self.token_embedding(input_ids)

        position_embeddings = self.position_embedding(position_ids)
        embeddings = inputs_embeds + position_embeddings

        return embeddings



class XCLIPAttention(nn.Module):
    

    def __init__(self, config):
        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.k_proj = nn.Linear(self.embed_dim, self.embed_dim)
        self.v_proj = nn.Linear(self.embed_dim, self.embed_dim)
        self.q_proj = nn.Linear(self.embed_dim, self.embed_dim)
        self.out_proj = nn.Linear(self.embed_dim, self.embed_dim)

    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 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]]]:
        

        bsz, tgt_len, embed_dim = hidden_states.size()

        
        query_states = self.q_proj(hidden_states) * self.scale
        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)
        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()}"
            )

        
        if causal_attention_mask is not None:
            if causal_attention_mask.size() != (bsz, 1, tgt_len, src_len):
                raise ValueError(
                    f"Attention mask should be of size {(bsz, 1, tgt_len, src_len)}, but is"
                    f" {causal_attention_mask.size()}"
                )
            attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len) + causal_attention_mask
            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_len):
                raise ValueError(
                    f"Attention mask should be of size {(bsz, 1, tgt_len, src_len)}, but is {attention_mask.size()}"
                )
            attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len) + attention_mask
            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:
            
            
            
            
            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 XCLIPMLP(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.activation_fn = ACT2FN[config.hidden_act]
        self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size)
        self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        hidden_states = self.fc1(hidden_states)
        hidden_states = self.activation_fn(hidden_states)
        hidden_states = self.fc2(hidden_states)
        return hidden_states



class XCLIPEncoderLayer(nn.Module):
    def __init__(self, config: XCLIPConfig):
        super().__init__()
        self.embed_dim = config.hidden_size
        self.self_attn = XCLIPAttention(config)
        self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
        self.mlp = XCLIPMLP(config)
        self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: torch.Tensor,
        causal_attention_mask: torch.Tensor,
        output_attentions: Optional[bool] = False,
    ) -> Tuple[torch.FloatTensor]:
        
        residual = hidden_states

        hidden_states = self.layer_norm1(hidden_states)
        hidden_states, attn_weights = self.self_attn(
            hidden_states=hidden_states,
            attention_mask=attention_mask,
            causal_attention_mask=causal_attention_mask,
            output_attentions=output_attentions,
        )
        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,)

        return outputs



def drop_path(input: torch.Tensor, drop_prob: float = 0.0, training: bool = False) -> torch.Tensor:
    
    if drop_prob == 0.0 or not training:
        return input
    keep_prob = 1 - drop_prob
    shape = (input.shape[0],) + (1,) * (input.ndim - 1)  
    random_tensor = keep_prob + torch.rand(shape, dtype=input.dtype, device=input.device)
    random_tensor.floor_()  
    output = input.div(keep_prob) * random_tensor
    return output



class XCLIPDropPath(nn.Module):
    

    def __init__(self, drop_prob: Optional[float] = None) -> None:
        super().__init__()
        self.drop_prob = drop_prob

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        return drop_path(hidden_states, self.drop_prob, self.training)

    def extra_repr(self) -> str:
        return "p={}".format(self.drop_prob)


class XCLIPVisionEncoderLayer(nn.Module):
    

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

        self.message_fc = nn.Linear(self.embed_dim, self.embed_dim)
        self.message_ln = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
        self.message_attn = XCLIPAttention(config)

        self.drop_path = XCLIPDropPath(config.drop_path_rate) if config.drop_path_rate > 0.0 else nn.Identity()

        self.self_attn = XCLIPAttention(config)
        self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
        self.mlp = XCLIPMLP(config)
        self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: torch.Tensor,
        causal_attention_mask: torch.Tensor,
        output_attentions: Optional[bool] = False,
    ) -> Tuple[torch.FloatTensor]:
        
        batch_time, seq_length, hidden_size = hidden_states.size()
        batch_size = batch_time // self.num_frames
        msg_token = self.message_fc(hidden_states[:, 0, :])
        msg_token = msg_token.view(batch_size, self.num_frames, hidden_size)

        msg_token = msg_token + self.drop_path(self.message_attn(self.message_ln(msg_token))[0])
        
        msg_token = msg_token.view(-1, 1, hidden_size)

        hidden_states = torch.cat([hidden_states, msg_token], dim=1)

        residual = hidden_states

        hidden_states = self.layer_norm1(hidden_states)
        hidden_states, attn_weights = self.self_attn(
            hidden_states=hidden_states,
            attention_mask=attention_mask,
            causal_attention_mask=causal_attention_mask,
            output_attentions=output_attentions,
        )
        hidden_states = residual + hidden_states

        hidden_states = hidden_states[:, :seq_length, :]

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

        return outputs


class XCLIPPreTrainedModel(PreTrainedModel):
    

    config_class = XCLIPConfig
    base_model_prefix = "x_clip"
    supports_gradient_checkpointing = True

    def _init_weights(self, module):
        
        factor = self.config.initializer_factor
        if isinstance(module, XCLIPTextEmbeddings):
            module.token_embedding.weight.data.normal_(mean=0.0, std=factor * 0.02)
            module.position_embedding.weight.data.normal_(mean=0.0, std=factor * 0.02)
        elif isinstance(module, XCLIPVisionEmbeddings):
            factor = self.config.initializer_factor
            nn.init.normal_(module.class_embedding, mean=0.0, std=module.embed_dim**-0.5 * factor)
            nn.init.normal_(module.patch_embedding.weight, std=module.config.initializer_range * factor)
            nn.init.normal_(module.position_embedding.weight, std=module.config.initializer_range * factor)
        elif isinstance(module, XCLIPAttention):
            factor = self.config.initializer_factor
            in_proj_std = (module.embed_dim**-0.5) * ((2 * module.config.num_hidden_layers) ** -0.5) * factor
            out_proj_std = (module.embed_dim**-0.5) * factor
            nn.init.normal_(module.q_proj.weight, std=in_proj_std)
            nn.init.normal_(module.k_proj.weight, std=in_proj_std)
            nn.init.normal_(module.v_proj.weight, std=in_proj_std)
            nn.init.normal_(module.out_proj.weight, std=out_proj_std)
        elif isinstance(module, XCLIPMLP):
            factor = self.config.initializer_factor
            in_proj_std = (module.config.hidden_size**-0.5) * ((2 * module.config.num_hidden_layers) ** -0.5) * factor
            fc_std = (2 * module.config.hidden_size) ** -0.5 * factor
            nn.init.normal_(module.fc1.weight, std=fc_std)
            nn.init.normal_(module.fc2.weight, std=in_proj_std)
        elif isinstance(module, XCLIPModel):
            factor = self.config.initializer_factor
            nn.init.normal_(
                module.text_projection.weight,
                std=module.text_embed_dim**-0.5 * factor,
            )
            nn.init.normal_(
                module.visual_projection.weight,
                std=module.vision_embed_dim**-0.5 * factor,
            )
            nn.init.normal_(module.prompts_visual_projection, mean=0.0, std=module.vision_embed_dim**-0.5 * factor)
        elif isinstance(module, XCLIPMultiframeIntegrationTransformer):
            nn.init.normal_(module.position_embedding, std=self.config.initializer_factor)

        if isinstance(module, nn.LayerNorm):
            module.bias.data.zero_()
            module.weight.data.fill_(1.0)
        if isinstance(module, nn.Linear):
            module.weight.data.normal_(mean=0.0, std=self.config.initializer_factor)
            if module.bias is not None:
                module.bias.data.zero_()


X_CLIP_START_DOCSTRING = r

X_CLIP_TEXT_INPUTS_DOCSTRING = r

X_CLIP_VISION_INPUTS_DOCSTRING = r

X_CLIP_INPUTS_DOCSTRING = r



class XCLIPEncoder(nn.Module):
    

    def __init__(self, config: XCLIPConfig):
        super().__init__()
        self.config = config
        self.layers = nn.ModuleList([XCLIPEncoderLayer(config) for _ in range(config.num_hidden_layers)])
        self.gradient_checkpointing = False

    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,
    ) -> Union[Tuple, BaseModelOutput]:
        r
        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,
                    causal_attention_mask,
                    output_attentions,
                )
            else:
                layer_outputs = encoder_layer(
                    hidden_states,
                    attention_mask,
                    causal_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 XCLIPTextTransformer(nn.Module):
    def __init__(self, config: XCLIPTextConfig):
        super().__init__()
        self.config = config
        embed_dim = config.hidden_size
        self.embeddings = XCLIPTextEmbeddings(config)
        self.encoder = XCLIPEncoder(config)
        self.final_layer_norm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)

    @add_start_docstrings_to_model_forward(X_CLIP_TEXT_INPUTS_DOCSTRING)
    @replace_return_docstrings(output_type=BaseModelOutputWithPooling, config_class=XCLIPTextConfig)
    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,
    ) -> Union[Tuple, BaseModelOutputWithPooling]:
        r
        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 either input_ids")

        input_shape = input_ids.size()
        input_ids = input_ids.view(-1, input_shape[-1])

        hidden_states = self.embeddings(input_ids=input_ids, position_ids=position_ids)

        
        
        causal_attention_mask = _create_4d_causal_attention_mask(
            input_shape, hidden_states.dtype, device=hidden_states.device
        )
        
        if attention_mask is not None:
            
            attention_mask = _prepare_4d_attention_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)

        
        
        pooled_output = last_hidden_state[torch.arange(last_hidden_state.shape[0]), input_ids.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,
        )


class XCLIPTextModel(XCLIPPreTrainedModel):
    config_class = XCLIPTextConfig

    def __init__(self, config: XCLIPTextConfig):
        super().__init__(config)
        self.text_model = XCLIPTextTransformer(config)
        
        self.post_init()

    def get_input_embeddings(self) -> nn.Module:
        return self.text_model.embeddings.token_embedding

    def set_input_embeddings(self, value):
        self.text_model.embeddings.token_embedding = value

    @add_start_docstrings_to_model_forward(X_CLIP_TEXT_INPUTS_DOCSTRING)
    @replace_return_docstrings(output_type=BaseModelOutputWithPooling, config_class=XCLIPTextConfig)
    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,
    ) -> Union[Tuple, BaseModelOutputWithPooling]:
        r
        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,
        )


class XCLIPVisionEncoder(nn.Module):
    

    def __init__(self, config: XCLIPConfig):
        super().__init__()
        self.config = config
        self.layers = nn.ModuleList([XCLIPVisionEncoderLayer(config) for _ in range(config.num_hidden_layers)])
        self.gradient_checkpointing = False

    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,
    ) -> Union[Tuple, BaseModelOutput]:
        r
        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,
                    causal_attention_mask,
                    output_attentions,
                )
            else:
                layer_outputs = encoder_layer(
                    hidden_states,
                    attention_mask,
                    causal_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 XCLIPVisionTransformer(nn.Module):
    

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

        self.embeddings = XCLIPVisionEmbeddings(config)
        self.pre_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
        self.encoder = XCLIPVisionEncoder(config)
        self.post_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)

    @add_start_docstrings_to_model_forward(X_CLIP_VISION_INPUTS_DOCSTRING)
    @replace_return_docstrings(output_type=BaseModelOutputWithPooling, config_class=XCLIPVisionConfig)
    def forward(
        self,
        pixel_values: torch.FloatTensor,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple, BaseModelOutputWithPooling]:
        r
        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

        hidden_states = self.embeddings(pixel_values)
        hidden_states = self.pre_layernorm(hidden_states)

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


class XCLIPVisionModel(XCLIPPreTrainedModel):
    config_class = XCLIPVisionConfig
    main_input_name = "pixel_values"

    def __init__(self, config: XCLIPVisionConfig):
        super().__init__(config)
        self.vision_model = XCLIPVisionTransformer(config)
        
        self.post_init()

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

    @add_start_docstrings_to_model_forward(X_CLIP_VISION_INPUTS_DOCSTRING)
    @replace_return_docstrings(output_type=BaseModelOutputWithPooling, config_class=XCLIPVisionConfig)
    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
        return self.vision_model(
            pixel_values=pixel_values,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )


class XCLIPMultiframeIntegrationTransformer(nn.Module):
    

    def __init__(self, config: XCLIPVisionConfig):
        super().__init__()

        self.position_embedding = nn.Parameter(torch.empty(1, config.num_frames, config.hidden_size))
        self.encoder = XCLIPEncoder(config)

    def forward(
        self,
        hidden_states,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple, BaseModelOutput]:
        residual = hidden_states

        
        hidden_states = hidden_states + self.position_embedding

        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 = last_hidden_state.type(hidden_states.dtype) + residual

        pooled_output = last_hidden_state.mean(dim=1, keepdim=False)

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


class XCLIPCrossAttention(nn.Module):
    

    def __init__(self, config):
        super().__init__()
        self.num_heads = config.prompt_num_attention_heads

        dim = config.projection_dim
        head_dim = dim // self.num_heads
        self.scale = head_dim**-0.5

        self.q_proj = nn.Linear(dim, dim, bias=False)
        self.k_proj = nn.Linear(dim, dim, bias=False)
        self.v_proj = nn.Linear(dim, dim, bias=False)

        self.attn_drop = nn.Dropout(config.prompt_attention_dropout)
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(config.prompt_projection_dropout)

    def _shape(self, tensor: torch.Tensor, seq_len: int, batch_size: int):
        return tensor.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()

    def forward(self, queries, keys, values):
        
        batch_size, query_seq_len, hidden_size = queries.shape
        batch_size, key_seq_len, hidden_size = keys.shape
        queries = (
            self.q_proj(queries)
            .reshape(batch_size, query_seq_len, self.num_heads, hidden_size // self.num_heads)
            .permute(0, 2, 1, 3)
        )
        keys = (
            self.k_proj(keys)
            .reshape(batch_size, key_seq_len, self.num_heads, hidden_size // self.num_heads)
            .permute(0, 2, 1, 3)
        )
        values = (
            self.v_proj(values)
            .reshape(batch_size, key_seq_len, self.num_heads, hidden_size // self.num_heads)
            .permute(0, 2, 1, 3)
        )

        attn = (queries @ keys.transpose(-2, -1)) * self.scale
        attn = attn.softmax(dim=-1)
        attn = self.attn_drop(attn)

        x = (attn @ values).transpose(1, 2).reshape(batch_size, query_seq_len, hidden_size)
        x = self.proj(x)
        x = self.proj_drop(x)
        return x


class PromptGeneratorLayer(nn.Module):
    def __init__(self, config):
        super().__init__()

        embed_dim = config.projection_dim
        self.cross_attn = XCLIPCrossAttention(config)
        self.norm1 = nn.LayerNorm(embed_dim, eps=config.text_config.layer_norm_eps)
        self.norm3 = nn.LayerNorm(embed_dim, eps=config.text_config.layer_norm_eps)
        self.mlp = nn.Sequential(
            nn.Linear(embed_dim, embed_dim * 4),
            ACT2FN[config.prompt_hidden_act],
            nn.Dropout(config.prompt_attention_dropout),
            nn.Linear(embed_dim * 4, embed_dim),
        )

    def forward(self, x, visual):
        x = x + self.cross_attn(self.norm1(x), visual, visual)
        x = x + self.mlp(self.norm3(x))
        return x


class XCLIPPromptGenerator(nn.Module):
    

    def __init__(self, config):
        super().__init__()
        embed_dim = config.projection_dim
        self.layernorm = nn.LayerNorm(embed_dim, eps=config.vision_config.layer_norm_eps)
        self.decoder = nn.ModuleList([PromptGeneratorLayer(config) for _ in range(config.prompt_layers)])
        self.alpha = nn.Parameter(torch.ones(embed_dim) * config.prompt_alpha)

    def forward(self, text, visual):
        visual = self.layernorm(visual)
        for layer in self.decoder:
            text = layer(text, visual)

        return self.alpha * text


@add_start_docstrings(X_CLIP_START_DOCSTRING)
class XCLIPModel(XCLIPPreTrainedModel):
    config_class = XCLIPConfig

    def __init__(self, config: XCLIPConfig):
        super().__init__(config)

        if not isinstance(config.text_config, XCLIPTextConfig):
            raise ValueError(
                "config.text_config is expected to be of type XCLIPTextConfig but is of type"
                f" {type(config.text_config)}."
            )

        if not isinstance(config.vision_config, XCLIPVisionConfig):
            raise ValueError(
                "config.vision_config is expected to be of type XCLIPVisionConfig but is of type"
                f" {type(config.vision_config)}."
            )

        text_config = config.text_config
        vision_config = config.vision_config

        self.projection_dim = config.projection_dim
        self.text_embed_dim = text_config.hidden_size
        self.vision_embed_dim = vision_config.hidden_size

        self.text_model = XCLIPTextTransformer(text_config)
        self.vision_model = XCLIPVisionTransformer(vision_config)

        self.visual_projection = nn.Linear(self.vision_embed_dim, self.projection_dim, bias=False)
        self.text_projection = nn.Linear(self.text_embed_dim, self.projection_dim, bias=False)
        self.logit_scale = nn.Parameter(torch.tensor(self.config.logit_scale_init_value))

        self.prompts_visual_layernorm = nn.LayerNorm(self.vision_embed_dim, eps=config.vision_config.layer_norm_eps)
        self.prompts_visual_projection = nn.Parameter(torch.randn(self.vision_embed_dim, self.projection_dim))

        mit_config = copy(vision_config)
        mit_config.hidden_size = vision_config.mit_hidden_size
        mit_config.intermediate_size = vision_config.mit_intermediate_size
        mit_config.num_hidden_layers = vision_config.mit_num_hidden_layers
        mit_config.num_attention_heads = vision_config.mit_num_attention_heads
        self.mit = XCLIPMultiframeIntegrationTransformer(mit_config)

        self.prompts_generator = XCLIPPromptGenerator(config)

        
        self.post_init()

    @add_start_docstrings_to_model_forward(X_CLIP_TEXT_INPUTS_DOCSTRING)
    def get_text_features(
        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,
    ) -> torch.FloatTensor:
        r
        
        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

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

        text_embeds = text_outputs[1]
        text_embeds = self.text_projection(text_embeds)

        return text_embeds

    @add_start_docstrings_to_model_forward(X_CLIP_VISION_INPUTS_DOCSTRING)
    def get_video_features(
        self,
        pixel_values: Optional[torch.FloatTensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> torch.FloatTensor:
        r
        
        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

        batch_size, num_frames, num_channels, height, width = pixel_values.shape
        pixel_values = pixel_values.reshape(-1, num_channels, height, width)

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

        video_embeds = vision_outputs[1]
        video_embeds = self.visual_projection(video_embeds)

        cls_features = video_embeds.view(batch_size, num_frames, -1)

        mit_outputs = self.mit(
            cls_features,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )
        video_embeds = mit_outputs[1]

        return video_embeds

    @add_start_docstrings_to_model_forward(X_CLIP_INPUTS_DOCSTRING)
    @replace_return_docstrings(output_type=XCLIPOutput, config_class=XCLIPConfig)
    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        pixel_values: Optional[torch.FloatTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        return_loss: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple, XCLIPOutput]:
        r
        
        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

        batch_size, num_frames, num_channels, height, width = pixel_values.shape
        pixel_values = pixel_values.reshape(-1, num_channels, height, width)

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

        video_embeds = vision_outputs[1]
        video_embeds = self.visual_projection(video_embeds)

        cls_features = video_embeds.view(batch_size, num_frames, -1)

        mit_outputs = self.mit(
            cls_features,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )
        video_embeds = mit_outputs[1]

        img_features = vision_outputs[0][:, 1:, :]
        img_features = self.prompts_visual_layernorm(img_features)
        img_features = img_features @ self.prompts_visual_projection
        img_features = img_features.view(batch_size, num_frames, -1, video_embeds.shape[-1])
        img_features = img_features.mean(dim=1, keepdim=False)

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

        text_embeds = text_outputs[1]
        text_embeds = self.text_projection(text_embeds)

        text_embeds = text_embeds.unsqueeze(0).expand(batch_size, -1, -1)
        text_embeds = text_embeds + self.prompts_generator(text_embeds, img_features)

        
        video_embeds = video_embeds / video_embeds.norm(p=2, dim=-1, keepdim=True)
        text_embeds = text_embeds / text_embeds.norm(p=2, dim=-1, keepdim=True)

        
        logit_scale = self.logit_scale.exp()
        logits_per_video = torch.einsum("bd,bkd->bk", video_embeds, logit_scale * text_embeds)
        logits_per_text = logits_per_video.T

        loss = None
        if return_loss:
            loss = x_clip_loss(logits_per_text)

        if not return_dict:
            output = (logits_per_video, logits_per_text, text_embeds, video_embeds, text_outputs, vision_outputs)
            return ((loss,) + output) if loss is not None else output

        return XCLIPOutput(
            loss=loss,
            logits_per_video=logits_per_video,
            logits_per_text=logits_per_text,
            text_embeds=text_embeds,
            video_embeds=video_embeds,
            text_model_output=text_outputs,
            vision_model_output=vision_outputs,
            mit_output=mit_outputs,
        )