codingutf8 2022 nvidia 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 tf 2 0 groupvit model from future import annotations import collections abc import math from dataclasses import dataclass from typing import any optional tuple union import numpy as np import tensorflow as tf from activationstf import gettfactivation from modelingtfoutputs import tfbasemodeloutput tfbasemodeloutputwithpooling from modelingtfutils import tfmodelinputtype tfpretrainedmodel getinitializer kerasserializable unpackinputs from tfutils import checkembeddingswithinbounds shapelist stablesoftmax from utils import modeloutput addstartdocstrings addstartdocstringstomodelforward istensorflowprobabilityavailable logging replacereturndocstrings from configurationgroupvit import groupvitconfig groupvittextconfig groupvitvisionconfig logger logging getloggername soft dependency if istensorflowprobabilityavailable try import tensorflowprobability as tfp on the first call check whether a compatible version of tensorflow is installed tensorflow probability depends on a recent stable release of tensorflow tfp distributions normalloc0 0 scale1 0 except importerror logger error groupvit models are not usable since tensorflowprobability can t be loaded it seems you have tensorflowprobability installed with the wrong tensorflow version please try to reinstall it following the instructions here https github comtensorflowprobability checkpointfordoc nvidiagroupvitgccyfcc tfgroupvitpretrainedmodelarchivelist nvidiagroupvitgccyfcc see all groupvit models at https huggingface comodels filtergroupvit largenegative 1e8 copied from transformers models bart modelingtfbart expandmask def expandmaskmask tf tensor tgtlen optionalint none srclen shapelistmask1 tgtlen tgtlen if tgtlen is not none else srclen onecst tf constant1 0 mask tf castmask dtypeonecst dtype expandedmask tf tilemask none none 1 1 tgtlen 1 return onecst expandedmask largenegative contrastive loss function adapted from https sachinruk github ioblogpytorchpytorch20lightningloss20functiongpu20210307clip html def contrastivelosslogits tf tensor tf tensor return tf math reducemean tf keras metrics sparsecategoricalcrossentropy ytruetf rangeshapelistlogits0 ypredlogits fromlogitstrue copied from transformers models clip modelingtfclip cliploss with clipgroupvit def groupvitlosssimilarity tf tensor tf tensor captionloss contrastivelosssimilarity imageloss contrastivelosstf transposesimilarity return captionloss imageloss 2 0 def hardsoftmaxlogits tf tensor dim int tf tensor ysoft stablesoftmaxlogits dim straight through index tf argmaxysoft dim yhard tf onehot index depthshapelistlogitsdim tensorflow expects axis to be 1 or between 0 3 but received 2 this is why the following code snippet is used axisrangelenshapelistlogitsdim dtypeysoft dtype ret yhard tf stopgradientysoft ysoft return ret def gumbelsoftmaxlogits tf tensor tau float 1 hard bool false dim int 1 tf tensor gumbeldist tfp distributions gumbel0 0 1 0 gumbels gumbeldist sampletf shapelogits dtypelogits dtype gumbels logits gumbels tau gumbellogits tau ysoft stablesoftmaxgumbels dim if hard straight through index tf argmaxysoft dim yhard tf onehot index depthshapelistlogitsdim tensorflow expects axis to be 1 or between 0 3 but received 2 this is why the following code snippet is used axisrangelenshapelistlogitsdim dtypeysoft dtype ret yhard tf stopgradientysoft ysoft else reparametrization trick ret ysoft return ret def resizeattentionmapattentions tf tensor height int width int aligncorners bool false tf tensor scale height width attentions shape2 0 5 if height width featwidth intnp roundwidth scale featheight shapelistattentions2 featwidth else featheight intnp roundheight scale featwidth shapelistattentions2 featheight batchsize shapelistattentions0 groups shapelistattentions1 number of group token batchsize groups height x width groups batchsize groups height width attentions tf reshapeattentions batchsize groups featheight featwidth attentions tf transposeattentions perm0 2 3 1 if aligncorners attentions tf compat v1 image resize attentions sizeheight width methodbilinear aligncornersaligncorners else attentions tf image resizeattentions sizeheight width methodbilinear attentions tf transposeattentions perm0 3 1 2 return attentions def getgroupingfromattentionsattentions tupletf tensor hwshape tupleint tf tensor attnmaps prevattnmasks none for attnmasks in attentions batchsize numgroups height x width batchsize height x width numgroups attnmasks tf transposeattnmasks perm0 2 1 if prevattnmasks is none prevattnmasks attnmasks else prevattnmasks tf matmulprevattnmasks attnmasks batchsize height x width numgroups batchsize numgroups height x width batchsize numgroups height width curattnmap resizeattentionmaptf transposeprevattnmasks perm0 2 1 hwshape attnmaps appendcurattnmap batchsize numgroups height width finalgrouping attnmaps1 return tf stopgradientfinalgrouping dataclass class tfgroupvitmodeloutputmodeloutput loss tf tensor none none logitsperimage tf tensor none logitspertext tf tensor none segmentationlogits tf tensor none textembeds tf tensor none imageembeds tf tensor none textmodeloutput tfbasemodeloutputwithpooling none visionmodeloutput tfbasemodeloutputwithpooling none def totupleself tupleany return tuple selfk if k not in textmodeloutput visionmodeloutput else getattrself k totuple for k in self keys class tfgroupvitcrossattentionlayertf keras layers layer def initself config groupvitvisionconfig kwargs super initkwargs self attn tfgroupvitattentionconfig nameattn self norm2 tf keras layers layernormalizationepsilonconfig layernormeps namenorm2 self mlp tfgroupvitmlpconfig namemlp self normpost tf keras layers layernormalizationepsilonconfig layernormeps namenormpost def callself query tf tensor key tf tensor training bool false tf tensor x query x x self attnquery encoderhiddenstateskey0 x x self mlpself norm2x x self normpostx return x class tfgroupvitassignattentiontf keras layers layer def initself config groupvitvisionconfig kwargs super initkwargs self scale config hiddensize0 5 self qproj tf keras layers denseconfig hiddensize nameqproj self kproj tf keras layers denseconfig hiddensize namekproj self vproj tf keras layers denseconfig hiddensize namevproj self proj tf keras layers denseconfig hiddensize nameproj self assigneps config assigneps def getattnself attn tf tensor gumbel bool true hard bool true training bool false tf tensor if gumbel and training attn gumbelsoftmaxattn dim2 hardhard else if hard attn hardsoftmaxattn dim2 else attn stablesoftmaxattn axis2 return attn def callself query tf tensor key tf tensor training bool false value key batchsize querylength channels query self qprojquery batchsize keylength channels key self kprojkey batchsize keylength channels value self vprojvalue batchsize querylength keylength rawattn tf matmulquery key transposebtrue self scale attn self getattnrawattn trainingtraining softattn self getattnrawattn trainingtraining gumbelfalse hardfalse attn attn tf math reducesumattn axis1 keepdimstrue self assigneps out tf matmulattn value out self projout return out softattn class tfgroupvittokenassigntf keras layers layer def initself config groupvitvisionconfig numgrouptoken int numoutputgroup int kwargs super initkwargs self numoutputgroup numoutputgroup norm on grouptokens self normtokens tf keras layers layernormalizationepsilonconfig layernormeps namenormtokens assignmlpratio config assignmlpratio if isinstanceconfig assignmlpratio collections abc iterable else config assignmlpratio config assignmlpratio tokensdim channelsdim intx config hiddensize for x in assignmlpratio self mlpinter tfgroupvitmixermlpconfig numgrouptoken tokensdim numoutputgroup namemlpinter self normposttokens tf keras layers layernormalization epsilonconfig layernormeps namenormposttokens norm on x self normx tf keras layers layernormalizationepsilonconfig layernormeps namenormx self preassignattn tfgroupvitcrossattentionlayerconfig namepreassignattn self assign tfgroupvitassignattentionconfig nameassign self normnewx tf keras layers layernormalizationepsilonconfig layernormeps namenormnewx self mlpchannels tfgroupvitmlp config config hiddensize channelsdim config hiddensize namemlpchannels def projectgrouptokenself grouptokens tf tensor tf tensor b numoutputgroups c b numgrouptokens c projectedgrouptokens self mlpintergrouptokens projectedgrouptokens self normposttokensprojectedgrouptokens return projectedgrouptokens def callself imagetokens tf tensor grouptokens tf tensor training bool false grouptokens self normtokensgrouptokens imagetokens self normximagetokens batchsize numoutputgroups channels projectedgrouptokens self projectgrouptokengrouptokens projectedgrouptokens self preassignattnprojectedgrouptokens imagetokens newimagetokens attention self assignprojectedgrouptokens imagetokens newimagetokens projectedgrouptokens newimagetokens newimagetokens self mlpchannelsself normnewxnewimagetokens return newimagetokens attention adapted from transformers models vit modelingtfvit tfvitpatchembeddings with vitgroupvit class tfgroupvitpatchembeddingstf keras layers layer def initself config groupvitconfig kwargs super initkwargs imagesize patchsize config imagesize config patchsize numchannels config numchannels hiddensize is a member as it will be required in the call method self hiddensize config hiddensize imagesize imagesize if isinstanceimagesize collections abc iterable else imagesize imagesize patchsize patchsize if isinstancepatchsize collections abc iterable else patchsize patchsize numpatches imagesize1 patchsize1 imagesize0 patchsize0 self imagesize imagesize self patchsize patchsize self numpatches numpatches self numchannels numchannels self config config self projection tf keras layers conv2d filtersself hiddensize kernelsizepatchsize stridespatchsize paddingvalid dataformatchannelslast usebiastrue kernelinitializergetinitializerself config initializerrange biasinitializerzeros nameprojection def call self pixelvalues tf tensor interpolateposencoding bool false training bool false tf tensor batchsize numchannels height width shapelistpixelvalues if tf executingeagerly and numchannels self numchannels raise valueerror make sure that the channel dimension of the pixel values match with the one set in the configuration if not interpolateposencoding and tf executingeagerly and height self imagesize0 or width self imagesize1 raise valueerror finput image size heightwidth doesn t match model self imagesize0self imagesize1 when running on cpu tf keras layers conv2d doesn t support nchw format so change the input format from nchw to nhwc shape batchsize inheight inwidth inchannelsnumchannels pixelvalues tf transposepixelvalues perm0 2 3 1 projection self projectionpixelvalues change the 2d spatial dimensions to a single temporal dimension shape batchsize numpatches outchannelsembeddim numpatches width self patchsize1 height self patchsize0 in the tfgroupvitvisionembeddings the embeddings from this layer will be layer normalized layernormalization layer needs to have static last dimension otherwise the testkerassaveload fails with symbolic tensors this is why we have used the hiddensize in the reshape method embeddings tf reshapetensorprojection shapebatchsize numpatches self hiddensize return embeddings adapted from transformers vit modelingtfvit tfvitembeddings class tfgroupvitvisionembeddingstf keras layers layer def initself config groupvitvisionconfig kwargs super initkwargs self patchembeddings tfgroupvitpatchembeddingsconfig namepatchembeddings self dropout tf keras layers dropoutrateconfig dropout namedropout self layernorm tf keras layers layernormalizationepsilonconfig layernormeps namelayernorm self config config def buildself inputshape tf tensorshape numpatches self patchembeddings numpatches self positionembeddings self addweight shape1 numpatches self config hiddensize initializerzeros trainabletrue namepositionembeddings super buildinputshape def interpolateposencodingself embeddings height width tf tensor batchsize numpatches dim shapelistembeddings numpositions shapelistself positionembeddings1 if numpatches numpositions and height width return self positionembeddings patchposembed self positionembeddings h0 height self config patchsize w0 width self config patchsize patchposembed tf image resize imagestf reshape patchposembed shape1 intmath sqrtnumpositions intmath sqrtnumpositions dim sizeh0 w0 methodbicubic patchposembed tf reshapetensorpatchposembed shape1 1 dim return patchposembed def call self pixelvalues tf tensor interpolateposencoding bool false training bool false tf tensor height width shapelistpixelvalues embeddings self patchembeddingspixelvalues interpolateposencodinginterpolateposencoding embeddings self layernormembeddings add positional encoding to each token if interpolateposencoding embeddings embeddings self interpolateposencodingembeddings height width else embeddings embeddings self positionembeddings embeddings self dropoutembeddings return embeddings copied from transformers models clip modelingtfclip tfcliptextembeddings with clipgroupvit class tfgroupvittextembeddingstf keras layers layer def initself config groupvittextconfig kwargs super initkwargs self embeddim config hiddensize self config config def buildself inputshape tf tensorshape none with tf namescopetokenembedding self weight self addweight shapeself config vocabsize self embeddim initializergetinitializerself config initializerfactor self config initializerrange trainabletrue nameweight with tf namescopepositionembedding self positionembedding self addweight shapeself config maxpositionembeddings self embeddim initializergetinitializerself config initializerfactor self config initializerrange trainabletrue nameembeddings super buildinputshape def call self inputids tf tensor none positionids tf tensor none inputsembeds tf tensor none tf tensor if inputids is none and inputsembeds is none raise valueerroryou have to specify either inputids or inputsembeds if inputsembeds is none checkembeddingswithinboundsinputids self config vocabsize inputsembeds tf gatherparamsself weight indicesinputids inputshape shapelistinputsembeds 1 if positionids is none positionids tf expanddimstf rangestart0 limitinputshape1 axis0 positionembeds tf gatherparamsself positionembedding indicespositionids positionembeds tf tileinputpositionembeds multiplesinputshape0 1 1 finalembeddings inputsembeds positionembeds return finalembeddings class tfgroupvitstagetf keras layers layer args hiddenstates tf tensor input to the layer of shape batch seqlen embeddim attentionmask tf tensor attention mask of size batch 1 tgtlen srclen where padding elements are indicated by very large negative values config encoderattentionheads outputattentions bool optional whether or not to return the grouping tensors of grouping block adapted from transformers models clip modelingtfclip tfclipattention multiheaded attention from attention is all you need paper def initself config groupvitconfig kwargs super initkwargs self embeddim config hiddensize self numattentionheads config numattentionheads self attentionheadsize self embeddim self numattentionheads if self attentionheadsize self numattentionheads self embeddim raise valueerror fembeddim must be divisible by numheads got embeddim self embeddim and numheads f self numattentionheads factor config initializerfactor inprojstd self embeddim0 5 2 config numhiddenlayers 0 5 factor outprojstd self embeddim0 5 factor self sqrtattheadsize math sqrtself attentionheadsize self qproj tf keras layers dense unitsself embeddim kernelinitializergetinitializerinprojstd nameqproj self kproj tf keras layers dense unitsself embeddim kernelinitializergetinitializerinprojstd namekproj self vproj tf keras layers dense unitsself embeddim kernelinitializergetinitializerinprojstd namevproj self dropout tf keras layers dropoutrateconfig attentiondropout self outproj tf keras layers dense unitsself embeddim kernelinitializergetinitializeroutprojstd nameoutproj copied from transformers models bert modelingtfbert tfbertselfattention transposeforscores def transposeforscoresself tensor tf tensor batchsize int tf tensor reshape from batchsize seqlength allheadsize to batchsize seqlength numattentionheads attentionheadsize tensor tf reshapetensortensor shapebatchsize 1 self numattentionheads self attentionheadsize transpose the tensor from batchsize seqlength numattentionheads attentionheadsize to batchsize numattentionheads seqlength attentionheadsize return tf transposetensor perm0 2 1 3 def call self hiddenstates tf tensor attentionmask tf tensor none causalattentionmask tf tensor none outputattentions bool none encoderhiddenstates tf tensor none training bool false tupletf tensor take the dot product between query and key to get the raw attention scores batch size numheads seqlenq seqlenk apply the causalattentionmask first apply the causal attention mask precomputed for all layers in tfclipmodel call function apply the attention mask precomputed for all layers in tfclipmodel call function normalize the attention scores to probabilities this is actually dropping out entire tokens to attend to which might seem a bit unusual but is taken from the original transformer paper batchsize seqlenq embeddim in tfbert attention weights are returned after dropout however in clip they are returned before dropout copied from transformers models clip modelingtfclip tfclipencoderlayer with clipgroupvit args hiddenstates tf tensor input to the layer of shape batch seqlen embeddim attentionmask tf tensor attention mask of size batch 1 tgtlen srclen where padding elements are indicated by very large negative values causalattentionmask tf tensor causal attention mask of size batch 1 tgtlen srclen where padding elements are indicated by very large negative values outputattentions bool whether or not to return the attentions tensors of all attention layers see outputs under returned tensors for more detail adapted from transformers models clip modelingtfclip tfgroupvittextencoder copied from transformers models clip modelingtfclip tfcliptexttransformer with cliptextgroupvittext clipencodergroupvittextencoder for pooledoutput computation clip s text model uses causal mask prepare it here https github comopenaiclipblobcfcffb90e69f37bf2ff1e988237a0fbe41f33c04clipmodel pyl324 check attention mask and invert bsz seqlen bsz 1 tgtseqlen srcseqlen the eostokenid was incorrect before pr 24773 let s keep what have been done here a clip model with such eostokenid in the config can t work correctly with extra new tokens added textembeds shape batchsize nctx transformer width take features from the eot embedding eottoken is the highest number in each sequence the config gets updated eostokenid from pr 24773 so the use of exta new tokens is possible it is possible with an unspecified sequence length for seqlength to be a runtime value which is unsupported by tf constant per the tensorflow docs tf fill can handle runtime dynamic shapes https www tensorflow orgapidocspythontffill set an additive 2d attention mask with all places being masked set diagonal lower triangular parts to 0 i e the places not to be masked tip think the 2d matrix as the space of queryseq keyseq tomask tf linalg bandparttomask 1 0 adapted from transformers models clip modelingtfclip tfclipvisiontransformer normalize the last hidden state copied from transformers models clip modelingtfclip tfcliptextmainlayer with clipgroupvit copied from transformers models clip modelingtfclip tfclipvisionmainlayer with clipgroupvit adapted from transformers models clip modelingtfclip tfclipmainlayer normalized features cosine similarity as logits grouped features batchsizeimage numgroup hiddensize batchsizeimagenumgroup hiddensize batchsizeimage numgroup height width normalized features batchsizeimage x numgroup batchsizetext batchsizeimage batchsizetext numgroup batchsizeimage batchsizetext height x width batchsizeimage batchsizetext height width an abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models groupvittextinputsdocstring r args inputids np ndarray tf tensor listtf tensor dictstr tf tensor or dictstr np ndarray and each example must have the shape 0 indices of input sequence tokens in the vocabulary indices can be obtained using autotokenizer see pretrainedtokenizer call and pretrainedtokenizer encode for details what are input ids glossaryinputids attentionmask np ndarray or tf tensor of shape 0 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 positionids np ndarray or tf tensor of shape 0 optional indices of positions of each input sequence tokens in the position embeddings selected in the range 0 config maxpositionembeddings 1 what are position ids glossarypositionids outputattentions bool optional whether or not to return the attentions tensors of all attention layers see attentions under returned tensors for more detail this argument can be used only in eager mode in graph mode the value in the config will be used instead outputhiddenstates bool optional whether or not to return the hidden states of all layers see hiddenstates under returned tensors for more detail this argument can be used only in eager mode in graph mode the value in the config will be used instead returndict bool optional whether or not to return a utils modeloutput instead of a plain tuple this argument can be used in eager mode in graph mode the value will always be set to true training bool optional defaults to false whether or not to use the model in training mode some modules like dropout modules have different behaviors between training and evaluation groupvitinputsdocstring r args inputids np ndarray tf tensor listtf tensor dictstr tf tensor or dictstr np ndarray and each example must have the shape 0 indices of input sequence tokens in the vocabulary indices can be obtained using autotokenizer see pretrainedtokenizer call and pretrainedtokenizer encode for details what are input ids glossaryinputids pixelvalues np ndarray tf tensor listtf tensor dictstr tf tensor or dictstr np ndarray and each example must have the shape batchsize numchannels height width pixel values pixel values can be obtained using autoimageprocessor see clipimageprocessor call for details attentionmask np ndarray or tf tensor of shape 0 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 positionids np ndarray or tf tensor of shape 0 optional indices of positions of each input sequence tokens in the position embeddings selected in the range 0 config maxpositionembeddings 1 what are position ids glossarypositionids returnloss bool optional whether or not to return the contrastive loss outputattentions bool optional whether or not to return the attentions tensors of all attention layers see attentions under returned tensors for more detail this argument can be used only in eager mode in graph mode the value in the config will be used instead outputhiddenstates bool optional whether or not to return the hidden states of all layers see hiddenstates under returned tensors for more detail this argument can be used only in eager mode in graph mode the value in the config will be used instead returndict bool optional whether or not to return a utils modeloutput instead of a plain tuple this argument can be used in eager mode in graph mode the value will always be set to true training bool optional defaults to false whether or not to use the model in training mode some modules like dropout modules have different behaviors between training and evaluation todo as is this currently fails with savedmodeltrue because tensorflow cannot trace through nested dataclasses reference https github comhuggingfacetransformerspull16886 coding utf 8 2022 nvidia 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 tf 2 0 groupvit model soft dependency on the first call check whether a compatible version of tensorflow is installed tensorflow probability depends on a recent stable release of tensorflow see all groupvit models at https huggingface co models filter groupvit copied from transformers models bart modeling_tf_bart _expand_mask expands attention_mask from bsz seq_len to bsz 1 tgt_seq_len src_seq_len 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_tf_clip clip_loss with clip groupvit straight through tensorflow expects axis to be 1 or between 0 3 but received 2 this is why the following code snippet is used gumbel logits tau straight through tensorflow expects axis to be 1 or between 0 3 but received 2 this is why the following code snippet is used reparametrization trick args attentions tf tensor attention map of shape batch_size groups feat_height feat_width height int height of the output attention map width int width of the output attention map align_corners bool optional the align_corner argument for nn functional interpolate returns tf tensor resized attention map of shape batch_size groups height width number of group token batch_size groups height x width groups batch_size groups height width args attentions tuple tf tensor tuple of attention maps returned by tfgroupvitvisiontransformer hw_shape tuple int height and width of the output attention map returns tf tensor the attention map of shape batch_size groups height width batch_size num_groups height x width batch_size height x width num_groups batch_size height x width num_groups batch_size num_groups height x width batch_size num_groups height width batch_size num_groups height width args loss tf tensor of shape 1 optional returned when return_loss is true contrastive loss for image text similarity logits_per_image tf tensor of shape image_batch_size text_batch_size the scaled dot product scores between image_embeds and text_embeds this represents the image text similarity scores logits_per_text tf tensor of shape text_batch_size image_batch_size the scaled dot product scores between text_embeds and image_embeds this represents the text image similarity scores segmentation_logits tf tensor of shape batch_size config num_labels logits_height logits_width classification scores for each pixel tip warning true the logits returned do not necessarily have the same size as the pixel_values passed as inputs this is to avoid doing two interpolations and lose some quality when a user needs to resize the logits to the original image size as post processing you should always check your logits shape and resize as needed tip text_embeds tf tensor of shape batch_size output_dim the text embeddings obtained by applying the projection layer to the pooled output of tfgroupvittextmodel image_embeds tf tensor of shape batch_size output_dim the image embeddings obtained by applying the projection layer to the pooled output of tfgroupvitvisionmodel text_model_output tfbasemodeloutputwithpooling the output of the tfgroupvittextmodel vision_model_output tfbasemodeloutputwithpooling the output of the tfgroupvitvisionmodel batch_size query_length channels batch_size key_length channels batch_size key_length channels batch_size query_length key_length norm on group_tokens norm on x args group_tokens tf tensor group tokens batch_size num_group_tokens channels returns projected_group_tokens tf tensor batch_size num_output_groups channels b num_output_groups c b num_group_tokens c args image_tokens tf tensor image tokens of shape batch_size input_length channels group_tokens tf tensor group tokens batch_size num_group_tokens channels batch_size num_output_groups channels adapted from transformers models vit modeling_tf_vit tfvitpatchembeddings with vit groupvit this class turns pixel_values of shape batch_size num_channels height width into the initial hidden_states patch embeddings of shape batch_size seq_length hidden_size to be consumed by a transformer hidden_size is a member as it will be required in the call method when running on cpu tf keras layers conv2d doesn t support nchw format so change the input format from nchw to nhwc shape batch_size in_height in_width in_channels num_channels change the 2d spatial dimensions to a single temporal dimension shape batch_size num_patches out_channels embed_dim in the tfgroupvitvisionembeddings the embeddings from this layer will be layer normalized layernormalization layer needs to have static last dimension otherwise the test_keras_save_load fails with symbolic tensors this is why we have used the hidden_size in the reshape method adapted from transformers vit modeling_tf_vit tfvitembeddings construct the position and patch embeddings this method allows to interpolate the pre trained position encodings to be able to use the model on higher resolution images source https github com facebookresearch dino blob de9ee3df6cf39fac952ab558447af1fa1365362a vision_transformer py l174 add positional encoding to each token copied from transformers models clip modeling_tf_clip tfcliptextembeddings with clip groupvit applies embedding based on inputs tensor returns final_embeddings tf tensor output embedding tensor this corresponds to the groupinglayer class in the groupvit implementation args hidden_states tf tensor input to the layer of shape batch seq_len embed_dim attention_mask tf tensor 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 grouping tensors of grouping block adapted from transformers models clip modeling_tf_clip tfclipattention multi headed attention from attention is all you need paper copied from transformers models bert modeling_tf_bert tfbertselfattention transpose_for_scores reshape from batch_size seq_length all_head_size to batch_size seq_length num_attention_heads attention_head_size transpose the tensor from batch_size seq_length num_attention_heads attention_head_size to batch_size num_attention_heads seq_length attention_head_size input shape batch x time x channel take the dot product between query and key to get the raw attention scores batch size num_heads seq_len_q seq_len_k apply the causal_attention_mask first apply the causal attention mask precomputed for all layers in tfclipmodel call function apply the attention mask precomputed for all layers in tfclipmodel call function normalize the attention scores to probabilities this is actually dropping out entire tokens to attend to which might seem a bit unusual but is taken from the original transformer paper batch_size seq_len_q embed_dim in tfbert attention weights are returned after dropout however in clip they are returned before dropout copied from transformers models clip modeling_tf_clip tfclipencoderlayer with clip groupvit args hidden_states tf tensor input to the layer of shape batch seq_len embed_dim attention_mask tf tensor attention mask of size batch 1 tgt_len src_len where padding elements are indicated by very large negative values causal_attention_mask tf tensor causal attention mask of size batch 1 tgt_len src_len where padding elements are indicated by very large negative values output_attentions bool whether or not to return the attentions tensors of all attention layers see outputs under returned tensors for more detail add attentions if we output them adapted from transformers models clip modeling_tf_clip tfgroupvittextencoder copied from transformers models clip modeling_tf_clip tfcliptexttransformer with cliptext groupvittext clipencoder groupvittextencoder for pooled_output computation clip s text model uses causal mask prepare it here https github com openai clip blob cfcffb90e69f37bf2ff1e988237a0fbe41f33c04 clip model py l324 check attention mask and invert bsz seq_len bsz 1 tgt_seq_len src_seq_len the eos_token_id was incorrect before pr 24773 let s keep what have been done here a clip model with such eos_token_id in the config can t work correctly with extra new tokens added text_embeds shape batch_size n_ctx transformer width take features from the eot embedding eot_token is the highest number in each sequence the config gets updated eos_token_id from pr 24773 so the use of exta new tokens is possible it is possible with an unspecified sequence length for seq_length to be a runtime value which is unsupported by tf constant per the tensorflow docs tf fill can handle runtime dynamic shapes https www tensorflow org api_docs python tf fill set an additive 2d attention mask with all places being masked set diagonal lower triangular parts to 0 i e the places not to be masked tip think the 2d matrix as the space of query_seq key_seq to_mask tf linalg band_part to_mask 1 0 adapted from transformers models clip modeling_tf_clip tfclipvisiontransformer normalize the last hidden state copied from transformers models clip modeling_tf_clip tfcliptextmainlayer with clip groupvit copied from transformers models clip modeling_tf_clip tfclipvisionmainlayer with clip groupvit adapted from transformers models clip modeling_tf_clip tfclipmainlayer normalized features cosine similarity as logits grouped features batch_size_image num_group hidden_size batch_size_image num_group hidden_size batch_size_image num_group height width normalized features batch_size_image x num_group batch_size_text batch_size_image batch_size_text num_group batch_size_image batch_size_text height x width batch_size_image batch_size_text height width an abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models this model inherits from tfpretrainedmodel check the superclass documentation for the generic methods the library implements for all its model such as downloading or saving resizing the input embeddings pruning heads etc this model is also a tf keras model https www tensorflow org api_docs python tf keras model subclass use it as a regular tf 2 0 keras model and refer to the tf 2 0 documentation for all matter related to general usage and behavior tip tf 2 0 models accepts two formats as inputs having all inputs as keyword arguments like pytorch models or having all inputs as a list tuple or dict in the first positional arguments this second option is useful when using tf keras model fit method which currently requires having all the tensors in the first argument of the model call function model inputs if you choose this second option there are three possibilities you can use to gather all the input tensors in the first positional argument a single tensor with input_ids only and nothing else model input_ids a list of varying length with one or several input tensors in the order given in the docstring model input_ids attention_mask or model input_ids attention_mask token_type_ids a dictionary with one or several input tensors associated to the input names given in the docstring model input_ids input_ids token_type_ids token_type_ids tip args config groupvitconfig 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 np ndarray tf tensor list tf tensor dict str tf tensor or dict str np ndarray and each example must have the shape 0 indices of input sequence tokens in the vocabulary indices can be obtained using autotokenizer see pretrainedtokenizer __call__ and pretrainedtokenizer encode for details what are input ids glossary input ids attention_mask np ndarray or tf tensor of shape 0 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 np ndarray or tf tensor of shape 0 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 this argument can be used only in eager mode in graph mode the value in the config will be used instead 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 this argument can be used only in eager mode in graph mode the value in the config will be used instead return_dict bool optional whether or not to return a utils modeloutput instead of a plain tuple this argument can be used in eager mode in graph mode the value will always be set to true training bool optional defaults to false whether or not to use the model in training mode some modules like dropout modules have different behaviors between training and evaluation args pixel_values np ndarray tf tensor list tf tensor dict str tf tensor or dict str np ndarray and each example must have the shape batch_size num_channels height width pixel values 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 this argument can be used only in eager mode in graph mode the value in the config will be used instead 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 this argument can be used only in eager mode in graph mode the value in the config will be used instead return_dict bool optional whether or not to return a utils modeloutput instead of a plain tuple this argument can be used in eager mode in graph mode the value will always be set to true training bool optional defaults to false whether or not to use the model in training mode some modules like dropout modules have different behaviors between training and evaluation args input_ids np ndarray tf tensor list tf tensor dict str tf tensor or dict str np ndarray and each example must have the shape 0 indices of input sequence tokens in the vocabulary indices can be obtained using autotokenizer see pretrainedtokenizer __call__ and pretrainedtokenizer encode for details what are input ids glossary input ids pixel_values np ndarray tf tensor list tf tensor dict str tf tensor or dict str np ndarray and each example must have the shape batch_size num_channels height width pixel values pixel values can be obtained using autoimageprocessor see clipimageprocessor __call__ for details attention_mask np ndarray or tf tensor of shape 0 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 np ndarray or tf tensor of shape 0 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 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 this argument can be used only in eager mode in graph mode the value in the config will be used instead 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 this argument can be used only in eager mode in graph mode the value in the config will be used instead return_dict bool optional whether or not to return a utils modeloutput instead of a plain tuple this argument can be used in eager mode in graph mode the value will always be set to true training bool optional defaults to false whether or not to use the model in training mode some modules like dropout modules have different behaviors between training and evaluation returns examples python from transformers import cliptokenizer tfgroupvittextmodel tokenizer cliptokenizer from_pretrained nvidia groupvit gcc yfcc model tfgroupvittextmodel from_pretrained nvidia groupvit gcc yfcc inputs tokenizer a photo of a cat a photo of a dog padding true return_tensors tf outputs model inputs last_hidden_state outputs last_hidden_state pooled_output outputs pooler_output pooled eos token states returns examples python from pil import image import requests from transformers import autoprocessor tfgroupvitvisionmodel processor autoprocessor from_pretrained nvidia groupvit gcc yfcc model tfgroupvitvisionmodel from_pretrained nvidia groupvit gcc yfcc url http images cocodataset org val2017 000000039769 jpg image image open requests get url stream true raw inputs processor images image return_tensors tf outputs model inputs last_hidden_state outputs last_hidden_state pooled_output outputs pooler_output pooled cls states returns text_features tf tensor of shape batch_size output_dim the text embeddings obtained by applying the projection layer to the pooled output of tfgroupvittextmodel examples python from transformers import cliptokenizer tfgroupvitmodel model tfgroupvitmodel from_pretrained nvidia groupvit gcc yfcc tokenizer cliptokenizer from_pretrained nvidia groupvit gcc yfcc inputs tokenizer a photo of a cat a photo of a dog padding true return_tensors tf text_features model get_text_features inputs returns image_features tf tensor of shape batch_size output_dim the image embeddings obtained by applying the projection layer to the pooled output of tfgroupvitvisionmodel examples python from pil import image import requests from transformers import autoprocessor tfgroupvitmodel model tfgroupvitmodel from_pretrained nvidia groupvit gcc yfcc processor autoprocessor from_pretrained nvidia groupvit gcc yfcc url http images cocodataset org val2017 000000039769 jpg image image open requests get url stream true raw inputs processor images image return_tensors tf image_features model get_image_features inputs returns examples python from pil import image import requests from transformers import autoprocessor tfgroupvitmodel import tensorflow as tf model tfgroupvitmodel from_pretrained nvidia groupvit gcc yfcc processor autoprocessor from_pretrained nvidia groupvit gcc yfcc url http images cocodataset org val2017 000000039769 jpg image image open requests get url stream true raw inputs processor text a photo of a cat a photo of a dog images image return_tensors tf padding true outputs model inputs logits_per_image outputs logits_per_image this is the image text similarity score probs tf math softmax logits_per_image axis 1 we can take the softmax to get the label probabilities todo as is this currently fails with saved_model true because tensorflow cannot trace through nested dataclasses reference https github com huggingface transformers pull 16886
from __future__ import annotations

import collections.abc
import math
from dataclasses import dataclass
from typing import Any, Optional, Tuple, Union

import numpy as np
import tensorflow as tf

from ...activations_tf import get_tf_activation
from ...modeling_tf_outputs import TFBaseModelOutput, TFBaseModelOutputWithPooling
from ...modeling_tf_utils import (
    TFModelInputType,
    TFPreTrainedModel,
    get_initializer,
    keras_serializable,
    unpack_inputs,
)
from ...tf_utils import check_embeddings_within_bounds, shape_list, stable_softmax
from ...utils import (
    ModelOutput,
    add_start_docstrings,
    add_start_docstrings_to_model_forward,
    is_tensorflow_probability_available,
    logging,
    replace_return_docstrings,
)
from .configuration_groupvit import GroupViTConfig, GroupViTTextConfig, GroupViTVisionConfig


logger = logging.get_logger(__name__)


if is_tensorflow_probability_available():
    try:
        import tensorflow_probability as tfp

        
        
        _ = tfp.distributions.Normal(loc=0.0, scale=1.0)
    except ImportError:
        logger.error(
            "GroupViT models are not usable since `tensorflow_probability` can't be loaded. "
            "It seems you have `tensorflow_probability` installed with the wrong tensorflow version."
            "Please try to reinstall it following the instructions here: https://github.com/tensorflow/probability."
        )

_CHECKPOINT_FOR_DOC = "nvidia/groupvit-gcc-yfcc"

TF_GROUPVIT_PRETRAINED_MODEL_ARCHIVE_LIST = [
    "nvidia/groupvit-gcc-yfcc",
    
]


LARGE_NEGATIVE = -1e8



def _expand_mask(mask: tf.Tensor, tgt_len: Optional[int] = None):
    
    src_len = shape_list(mask)[1]
    tgt_len = tgt_len if tgt_len is not None else src_len
    one_cst = tf.constant(1.0)
    mask = tf.cast(mask, dtype=one_cst.dtype)
    expanded_mask = tf.tile(mask[:, None, None, :], (1, 1, tgt_len, 1))

    return (one_cst - expanded_mask) * LARGE_NEGATIVE




def contrastive_loss(logits: tf.Tensor) -> tf.Tensor:
    return tf.math.reduce_mean(
        tf.keras.metrics.sparse_categorical_crossentropy(
            y_true=tf.range(shape_list(logits)[0]), y_pred=logits, from_logits=True
        )
    )



def groupvit_loss(similarity: tf.Tensor) -> tf.Tensor:
    caption_loss = contrastive_loss(similarity)
    image_loss = contrastive_loss(tf.transpose(similarity))
    return (caption_loss + image_loss) / 2.0


def hard_softmax(logits: tf.Tensor, dim: int) -> tf.Tensor:
    y_soft = stable_softmax(logits, dim)
    
    index = tf.argmax(y_soft, dim)
    y_hard = tf.one_hot(
        index,
        depth=shape_list(logits)[dim],
        
        
        axis=range(len(shape_list(logits)))[dim],
        dtype=y_soft.dtype,
    )
    ret = y_hard - tf.stop_gradient(y_soft) + y_soft

    return ret


def gumbel_softmax(logits: tf.Tensor, tau: float = 1, hard: bool = False, dim: int = -1) -> tf.Tensor:
    gumbel_dist = tfp.distributions.Gumbel(0.0, 1.0)
    gumbels = gumbel_dist.sample(tf.shape(logits), dtype=logits.dtype)

    gumbels = (logits + gumbels) / tau  
    y_soft = stable_softmax(gumbels, dim)

    if hard:
        
        index = tf.argmax(y_soft, dim)
        y_hard = tf.one_hot(
            index,
            depth=shape_list(logits)[dim],
            
            
            axis=range(len(shape_list(logits)))[dim],
            dtype=y_soft.dtype,
        )
        ret = y_hard - tf.stop_gradient(y_soft) + y_soft
    else:
        
        ret = y_soft
    return ret


def resize_attention_map(attentions: tf.Tensor, height: int, width: int, align_corners: bool = False) -> tf.Tensor:
    

    scale = (height * width // attentions.shape[2]) ** 0.5
    if height > width:
        feat_width = int(np.round(width / scale))
        feat_height = shape_list(attentions)[2] // feat_width
    else:
        feat_height = int(np.round(height / scale))
        feat_width = shape_list(attentions)[2] // feat_height

    batch_size = shape_list(attentions)[0]
    groups = shape_list(attentions)[1]  
    
    attentions = tf.reshape(attentions, (batch_size, groups, feat_height, feat_width))
    attentions = tf.transpose(attentions, perm=(0, 2, 3, 1))
    if align_corners:
        attentions = tf.compat.v1.image.resize(
            attentions,
            size=(height, width),
            method="bilinear",
            align_corners=align_corners,
        )
    else:
        attentions = tf.image.resize(attentions, size=(height, width), method="bilinear")
    attentions = tf.transpose(attentions, perm=(0, 3, 1, 2))
    return attentions


def get_grouping_from_attentions(attentions: Tuple[tf.Tensor], hw_shape: Tuple[int]) -> tf.Tensor:
    

    attn_maps = []
    prev_attn_masks = None
    for attn_masks in attentions:
        
        attn_masks = tf.transpose(attn_masks, perm=(0, 2, 1))
        if prev_attn_masks is None:
            prev_attn_masks = attn_masks
        else:
            prev_attn_masks = tf.matmul(prev_attn_masks, attn_masks)
        
        cur_attn_map = resize_attention_map(tf.transpose(prev_attn_masks, perm=(0, 2, 1)), *hw_shape)
        attn_maps.append(cur_attn_map)

    
    final_grouping = attn_maps[-1]

    return tf.stop_gradient(final_grouping)


@dataclass
class TFGroupViTModelOutput(ModelOutput):
    

    loss: tf.Tensor | None = None
    logits_per_image: tf.Tensor = None
    logits_per_text: tf.Tensor = None
    segmentation_logits: tf.Tensor = None
    text_embeds: tf.Tensor = None
    image_embeds: tf.Tensor = None
    text_model_output: TFBaseModelOutputWithPooling = None
    vision_model_output: TFBaseModelOutputWithPooling = None

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


class TFGroupViTCrossAttentionLayer(tf.keras.layers.Layer):
    def __init__(self, config: GroupViTVisionConfig, **kwargs):
        super().__init__(**kwargs)
        self.attn = TFGroupViTAttention(config, name="attn")
        self.norm2 = tf.keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="norm2")
        self.mlp = TFGroupViTMLP(config, name="mlp")
        self.norm_post = tf.keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="norm_post")

    def call(self, query: tf.Tensor, key: tf.Tensor, training: bool = False) -> tf.Tensor:
        x = query
        x = x + self.attn(query, encoder_hidden_states=key)[0]
        x = x + self.mlp(self.norm2(x))
        x = self.norm_post(x)
        return x


class TFGroupViTAssignAttention(tf.keras.layers.Layer):
    def __init__(self, config: GroupViTVisionConfig, **kwargs):
        super().__init__(**kwargs)
        self.scale = config.hidden_size**-0.5

        self.q_proj = tf.keras.layers.Dense(config.hidden_size, name="q_proj")
        self.k_proj = tf.keras.layers.Dense(config.hidden_size, name="k_proj")
        self.v_proj = tf.keras.layers.Dense(config.hidden_size, name="v_proj")
        self.proj = tf.keras.layers.Dense(config.hidden_size, name="proj")
        self.assign_eps = config.assign_eps

    def get_attn(self, attn: tf.Tensor, gumbel: bool = True, hard: bool = True, training: bool = False) -> tf.Tensor:
        if gumbel and training:
            attn = gumbel_softmax(attn, dim=-2, hard=hard)
        else:
            if hard:
                attn = hard_softmax(attn, dim=-2)
            else:
                attn = stable_softmax(attn, axis=-2)

        return attn

    def call(self, query: tf.Tensor, key: tf.Tensor, training: bool = False):
        value = key
        
        query = self.q_proj(query)

        
        key = self.k_proj(key)

        
        value = self.v_proj(value)

        
        raw_attn = tf.matmul(query, key, transpose_b=True) * self.scale

        attn = self.get_attn(raw_attn, training=training)
        soft_attn = self.get_attn(raw_attn, training=training, gumbel=False, hard=False)

        attn = attn / (tf.math.reduce_sum(attn, axis=-1, keepdims=True) + self.assign_eps)

        out = tf.matmul(attn, value)

        out = self.proj(out)

        return out, soft_attn


class TFGroupViTTokenAssign(tf.keras.layers.Layer):
    def __init__(self, config: GroupViTVisionConfig, num_group_token: int, num_output_group: int, **kwargs):
        super().__init__(**kwargs)
        self.num_output_group = num_output_group
        
        self.norm_tokens = tf.keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="norm_tokens")
        assign_mlp_ratio = (
            config.assign_mlp_ratio
            if isinstance(config.assign_mlp_ratio, collections.abc.Iterable)
            else (config.assign_mlp_ratio, config.assign_mlp_ratio)
        )
        tokens_dim, channels_dim = [int(x * config.hidden_size) for x in assign_mlp_ratio]
        self.mlp_inter = TFGroupViTMixerMLP(config, num_group_token, tokens_dim, num_output_group, name="mlp_inter")
        self.norm_post_tokens = tf.keras.layers.LayerNormalization(
            epsilon=config.layer_norm_eps, name="norm_post_tokens"
        )
        
        self.norm_x = tf.keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="norm_x")
        self.pre_assign_attn = TFGroupViTCrossAttentionLayer(config, name="pre_assign_attn")

        self.assign = TFGroupViTAssignAttention(config, name="assign")
        self.norm_new_x = tf.keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="norm_new_x")
        self.mlp_channels = TFGroupViTMLP(
            config, config.hidden_size, channels_dim, config.hidden_size, name="mlp_channels"
        )

    def project_group_token(self, group_tokens: tf.Tensor) -> tf.Tensor:
        
        
        projected_group_tokens = self.mlp_inter(group_tokens)
        projected_group_tokens = self.norm_post_tokens(projected_group_tokens)
        return projected_group_tokens

    def call(self, image_tokens: tf.Tensor, group_tokens: tf.Tensor, training: bool = False):
        

        group_tokens = self.norm_tokens(group_tokens)
        image_tokens = self.norm_x(image_tokens)
        
        projected_group_tokens = self.project_group_token(group_tokens)
        projected_group_tokens = self.pre_assign_attn(projected_group_tokens, image_tokens)
        new_image_tokens, attention = self.assign(projected_group_tokens, image_tokens)
        new_image_tokens += projected_group_tokens

        new_image_tokens = new_image_tokens + self.mlp_channels(self.norm_new_x(new_image_tokens))

        return new_image_tokens, attention



class TFGroupViTPatchEmbeddings(tf.keras.layers.Layer):
    

    def __init__(self, config: GroupViTConfig, **kwargs):
        super().__init__(**kwargs)
        image_size, patch_size = config.image_size, config.patch_size
        num_channels = config.num_channels
        
        self.hidden_size = config.hidden_size

        image_size = image_size if isinstance(image_size, collections.abc.Iterable) else (image_size, image_size)
        patch_size = patch_size if isinstance(patch_size, collections.abc.Iterable) else (patch_size, patch_size)
        num_patches = (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
        self.image_size = image_size
        self.patch_size = patch_size
        self.num_patches = num_patches
        self.num_channels = num_channels
        self.config = config

        self.projection = tf.keras.layers.Conv2D(
            filters=self.hidden_size,
            kernel_size=patch_size,
            strides=patch_size,
            padding="valid",
            data_format="channels_last",
            use_bias=True,
            kernel_initializer=get_initializer(self.config.initializer_range),
            bias_initializer="zeros",
            name="projection",
        )

    def call(
        self, pixel_values: tf.Tensor, interpolate_pos_encoding: bool = False, training: bool = False
    ) -> tf.Tensor:
        batch_size, num_channels, height, width = shape_list(pixel_values)
        if tf.executing_eagerly() and num_channels != self.num_channels:
            raise ValueError(
                "Make sure that the channel dimension of the pixel values match with the one set in the configuration."
            )
        if (
            not interpolate_pos_encoding
            and tf.executing_eagerly()
            and (height != self.image_size[0] or width != self.image_size[1])
        ):
            raise ValueError(
                f"Input image size ({height}*{width}) doesn't match model ({self.image_size[0]}*{self.image_size[1]})."
            )

        
        
        
        pixel_values = tf.transpose(pixel_values, perm=(0, 2, 3, 1))

        projection = self.projection(pixel_values)

        
        
        num_patches = (width // self.patch_size[1]) * (height // self.patch_size[0])
        
        
        
        embeddings = tf.reshape(tensor=projection, shape=(batch_size, num_patches, self.hidden_size))

        return embeddings



class TFGroupViTVisionEmbeddings(tf.keras.layers.Layer):
    

    def __init__(self, config: GroupViTVisionConfig, **kwargs):
        super().__init__(**kwargs)

        self.patch_embeddings = TFGroupViTPatchEmbeddings(config, name="patch_embeddings")
        self.dropout = tf.keras.layers.Dropout(rate=config.dropout, name="dropout")
        self.layernorm = tf.keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="layernorm")
        self.config = config

    def build(self, input_shape: tf.TensorShape):
        num_patches = self.patch_embeddings.num_patches
        self.position_embeddings = self.add_weight(
            shape=(1, num_patches, self.config.hidden_size),
            initializer="zeros",
            trainable=True,
            name="position_embeddings",
        )

        super().build(input_shape)

    def interpolate_pos_encoding(self, embeddings, height, width) -> tf.Tensor:
        

        batch_size, num_patches, dim = shape_list(embeddings)
        num_positions = shape_list(self.position_embeddings)[1]

        if num_patches == num_positions and height == width:
            return self.position_embeddings
        patch_pos_embed = self.position_embeddings
        h0 = height // self.config.patch_size
        w0 = width // self.config.patch_size
        patch_pos_embed = tf.image.resize(
            images=tf.reshape(
                patch_pos_embed, shape=(1, int(math.sqrt(num_positions)), int(math.sqrt(num_positions)), dim)
            ),
            size=(h0, w0),
            method="bicubic",
        )
        patch_pos_embed = tf.reshape(tensor=patch_pos_embed, shape=(1, -1, dim))
        return patch_pos_embed

    def call(
        self, pixel_values: tf.Tensor, interpolate_pos_encoding: bool = False, training: bool = False
    ) -> tf.Tensor:
        _, _, height, width = shape_list(pixel_values)
        embeddings = self.patch_embeddings(pixel_values, interpolate_pos_encoding=interpolate_pos_encoding)
        embeddings = self.layernorm(embeddings)

        
        if interpolate_pos_encoding:
            embeddings = embeddings + self.interpolate_pos_encoding(embeddings, height, width)
        else:
            embeddings = embeddings + self.position_embeddings

        embeddings = self.dropout(embeddings)

        return embeddings



class TFGroupViTTextEmbeddings(tf.keras.layers.Layer):
    def __init__(self, config: GroupViTTextConfig, **kwargs):
        super().__init__(**kwargs)

        self.embed_dim = config.hidden_size

        self.config = config

    def build(self, input_shape: tf.TensorShape = None):
        with tf.name_scope("token_embedding"):
            self.weight = self.add_weight(
                shape=(self.config.vocab_size, self.embed_dim),
                initializer=get_initializer(self.config.initializer_factor * self.config.initializer_range),
                trainable=True,
                name="weight",
            )

        with tf.name_scope("position_embedding"):
            self.position_embedding = self.add_weight(
                shape=(self.config.max_position_embeddings, self.embed_dim),
                initializer=get_initializer(self.config.initializer_factor * self.config.initializer_range),
                trainable=True,
                name="embeddings",
            )

        super().build(input_shape)

    def call(
        self,
        input_ids: tf.Tensor = None,
        position_ids: tf.Tensor = None,
        inputs_embeds: tf.Tensor = None,
    ) -> tf.Tensor:
        
        if input_ids is None and inputs_embeds is None:
            raise ValueError("You have to specify either input_ids or inputs_embeds")

        if inputs_embeds is None:
            check_embeddings_within_bounds(input_ids, self.config.vocab_size)
            inputs_embeds = tf.gather(params=self.weight, indices=input_ids)

        input_shape = shape_list(inputs_embeds)[:-1]

        if position_ids is None:
            position_ids = tf.expand_dims(tf.range(start=0, limit=input_shape[-1]), axis=0)

        position_embeds = tf.gather(params=self.position_embedding, indices=position_ids)
        position_embeds = tf.tile(input=position_embeds, multiples=(input_shape[0], 1, 1))
        final_embeddings = inputs_embeds + position_embeds

        return final_embeddings


class TFGroupViTStage(tf.keras.layers.Layer):
    

    def __init__(
        self,
        config: GroupViTVisionConfig,
        depth: int,
        num_prev_group_token: int,
        num_group_token: int,
        num_output_group: int,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.config = config
        self.depth = depth
        self.num_group_token = num_group_token
        self.layers = [TFGroupViTEncoderLayer(config, name=f"layers_._{i}") for i in range(depth)]

        if num_group_token > 0:
            self.downsample = TFGroupViTTokenAssign(
                config=config,
                num_group_token=num_group_token,
                num_output_group=num_output_group,
                name="downsample",
            )
        else:
            self.downsample = None

        if num_prev_group_token > 0 and num_group_token > 0:
            self.group_projector = [
                tf.keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="group_projector.0"),
                TFGroupViTMixerMLP(
                    config, num_prev_group_token, config.hidden_size // 2, num_group_token, name="group_projector.1"
                ),
            ]
        else:
            self.group_projector = None

    def build(self, input_shape: tf.TensorShape):
        if self.num_group_token > 0:
            self.group_token = self.add_weight(
                shape=(1, self.num_group_token, self.config.hidden_size),
                initializer="zeros",
                trainable=True,
                name="group_token",
            )
        else:
            self.group_token = None
        super().build(input_shape)

    @property
    def with_group_token(self):
        return self.group_token is not None

    def split_x(self, x: tf.Tensor) -> tf.Tensor:
        if self.with_group_token:
            return x[:, : -self.num_group_token], x[:, -self.num_group_token :]
        else:
            return x, None

    def concat_x(self, x: tf.Tensor, group_token: tf.Tensor | None = None) -> tf.Tensor:
        if group_token is None:
            return x
        return tf.concat([x, group_token], axis=1)

    def call(
        self,
        hidden_states: tf.Tensor,
        prev_group_token: tf.Tensor | None = None,
        output_attentions: bool = False,
        training: bool = False,
    ) -> Tuple[tf.Tensor]:
        
        if self.with_group_token:
            group_token = tf.tile(self.group_token, multiples=(shape_list(hidden_states)[0], 1, 1))
            if self.group_projector is not None:
                for layer in self.group_projector:
                    prev_group_token = layer(prev_group_token)
                group_token = group_token + prev_group_token
        else:
            group_token = None

        x = hidden_states

        cat_x = self.concat_x(x, group_token)
        for layer in self.layers:
            layer_out = layer(
                cat_x,
                attention_mask=None,
                causal_attention_mask=None,
                output_attentions=None,
            )
            cat_x = layer_out[0]

        x, group_token = self.split_x(cat_x)

        attention = None
        if self.downsample is not None:
            x, attention = self.downsample(x, group_token)

        outputs = (x, group_token)
        if output_attentions:
            outputs = outputs + (attention,)

        return outputs


class TFGroupViTMLP(tf.keras.layers.Layer):
    def __init__(
        self,
        config: GroupViTVisionConfig,
        hidden_size: Optional[int] = None,
        intermediate_size: Optional[int] = None,
        output_size: Optional[int] = None,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.config = config
        self.activation_fn = get_tf_activation(config.hidden_act)
        hidden_size = hidden_size if hidden_size is not None else config.hidden_size
        intermediate_size = intermediate_size if intermediate_size is not None else config.intermediate_size
        output_size = output_size if output_size is not None else hidden_size
        self.fc1 = tf.keras.layers.Dense(intermediate_size, name="fc1")
        self.fc2 = tf.keras.layers.Dense(output_size, name="fc2")

    def call(self, hidden_states: tf.Tensor, training: bool = False) -> tf.Tensor:
        hidden_states = self.fc1(hidden_states)
        hidden_states = self.activation_fn(hidden_states)
        hidden_states = self.fc2(hidden_states)
        return hidden_states


class TFGroupViTMixerMLP(TFGroupViTMLP):
    def call(self, x, training: bool = False):
        x = super().call(hidden_states=tf.transpose(x, perm=(0, 2, 1)))
        return tf.transpose(x, perm=(0, 2, 1))



class TFGroupViTAttention(tf.keras.layers.Layer):
    

    def __init__(self, config: GroupViTConfig, **kwargs):
        super().__init__(**kwargs)

        self.embed_dim = config.hidden_size
        self.num_attention_heads = config.num_attention_heads
        self.attention_head_size = self.embed_dim // self.num_attention_heads
        if self.attention_head_size * self.num_attention_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_attention_heads})."
            )

        factor = config.initializer_factor
        in_proj_std = (self.embed_dim**-0.5) * ((2 * config.num_hidden_layers) ** -0.5) * factor
        out_proj_std = (self.embed_dim**-0.5) * factor

        self.sqrt_att_head_size = math.sqrt(self.attention_head_size)

        self.q_proj = tf.keras.layers.Dense(
            units=self.embed_dim, kernel_initializer=get_initializer(in_proj_std), name="q_proj"
        )
        self.k_proj = tf.keras.layers.Dense(
            units=self.embed_dim, kernel_initializer=get_initializer(in_proj_std), name="k_proj"
        )
        self.v_proj = tf.keras.layers.Dense(
            units=self.embed_dim, kernel_initializer=get_initializer(in_proj_std), name="v_proj"
        )

        self.dropout = tf.keras.layers.Dropout(rate=config.attention_dropout)

        self.out_proj = tf.keras.layers.Dense(
            units=self.embed_dim, kernel_initializer=get_initializer(out_proj_std), name="out_proj"
        )

    
    def transpose_for_scores(self, tensor: tf.Tensor, batch_size: int) -> tf.Tensor:
        
        tensor = tf.reshape(tensor=tensor, shape=(batch_size, -1, self.num_attention_heads, self.attention_head_size))

        
        return tf.transpose(tensor, perm=[0, 2, 1, 3])

    def call(
        self,
        hidden_states: tf.Tensor,
        attention_mask: tf.Tensor = None,
        causal_attention_mask: tf.Tensor = None,
        output_attentions: bool = None,
        encoder_hidden_states: tf.Tensor = None,
        training: bool = False,
    ) -> Tuple[tf.Tensor]:
        

        batch_size = shape_list(hidden_states)[0]
        is_cross_attention = encoder_hidden_states is not None

        mixed_query_layer = self.q_proj(inputs=hidden_states)
        if is_cross_attention:
            mixed_key_layer = self.k_proj(inputs=encoder_hidden_states)
            mixed_value_layer = self.v_proj(inputs=encoder_hidden_states)
        else:
            mixed_key_layer = self.k_proj(inputs=hidden_states)
            mixed_value_layer = self.v_proj(inputs=hidden_states)

        query_layer = self.transpose_for_scores(mixed_query_layer, batch_size)
        key_layer = self.transpose_for_scores(mixed_key_layer, batch_size)
        value_layer = self.transpose_for_scores(mixed_value_layer, batch_size)

        
        
        attention_scores = tf.matmul(query_layer, key_layer, transpose_b=True)
        dk = tf.cast(self.sqrt_att_head_size, dtype=attention_scores.dtype)
        attention_scores = tf.divide(attention_scores, dk)

        
        if causal_attention_mask is not None:
            
            attention_scores = tf.add(attention_scores, causal_attention_mask)

        if attention_mask is not None:
            
            attention_scores = tf.add(attention_scores, attention_mask)

        
        _attention_probs = stable_softmax(logits=attention_scores, axis=-1)

        
        
        attention_probs = self.dropout(inputs=_attention_probs)

        attention_output = tf.matmul(attention_probs, value_layer)
        attention_output = tf.transpose(attention_output, perm=[0, 2, 1, 3])

        
        attention_output = tf.reshape(tensor=attention_output, shape=(batch_size, -1, self.embed_dim))

        attention_output = self.out_proj(attention_output)
        
        
        outputs = (attention_output, _attention_probs) if output_attentions else (attention_output,)

        return outputs



class TFGroupViTEncoderLayer(tf.keras.layers.Layer):
    def __init__(self, config: GroupViTConfig, **kwargs):
        super().__init__(**kwargs)

        self.embed_dim = config.hidden_size
        self.self_attn = TFGroupViTAttention(config, name="self_attn")
        self.layer_norm1 = tf.keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="layer_norm1")
        self.mlp = TFGroupViTMLP(config, name="mlp")
        self.layer_norm2 = tf.keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="layer_norm2")

    def call(
        self,
        hidden_states: tf.Tensor,
        attention_mask: tf.Tensor,
        causal_attention_mask: tf.Tensor,
        output_attentions: bool,
        training: bool = False,
    ) -> Tuple[tf.Tensor]:
        
        residual = hidden_states

        hidden_states = self.layer_norm1(inputs=hidden_states)
        attention_outputs = self.self_attn(
            hidden_states=hidden_states,
            attention_mask=attention_mask,
            causal_attention_mask=causal_attention_mask,
            output_attentions=output_attentions,
            training=training,
        )
        hidden_states = attention_outputs[0]
        hidden_states = residual + hidden_states

        residual = hidden_states
        hidden_states = self.layer_norm2(inputs=hidden_states)
        hidden_states = self.mlp(hidden_states=hidden_states)
        hidden_states = residual + hidden_states

        outputs = (hidden_states,) + attention_outputs[1:]  

        return outputs



class TFGroupViTTextEncoder(tf.keras.layers.Layer):
    def __init__(self, config: GroupViTTextConfig, **kwargs):
        super().__init__(**kwargs)

        self.layers = [TFGroupViTEncoderLayer(config, name=f"layers_._{i}") for i in range(config.num_hidden_layers)]

    def call(
        self,
        hidden_states,
        attention_mask: tf.Tensor,
        causal_attention_mask: tf.Tensor,
        output_attentions: bool,
        output_hidden_states: bool,
        return_dict: bool,
        training: bool = False,
    ) -> Union[Tuple, TFBaseModelOutput]:
        encoder_states = () if output_hidden_states else None
        all_attentions = () if output_attentions else None

        for idx, encoder_layer in enumerate(self.layers):
            if output_hidden_states:
                encoder_states = encoder_states + (hidden_states,)

            layer_outputs = encoder_layer(
                hidden_states,
                attention_mask,
                causal_attention_mask,
                output_attentions=output_attentions,
            )
            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 TFBaseModelOutput(
            last_hidden_state=hidden_states, hidden_states=encoder_states, attentions=all_attentions
        )


class TFGroupViTVisionEncoder(tf.keras.layers.Layer):
    def __init__(self, config: GroupViTVisionConfig, **kwargs) -> None:
        super().__init__(**kwargs)

        self.stages = [
            TFGroupViTStage(
                config=config,
                depth=config.depths[i],
                num_group_token=config.num_group_tokens[i],
                num_output_group=config.num_output_groups[i],
                num_prev_group_token=config.num_output_groups[i - 1] if i > 0 else 0,
                name=f"stages_._{i}",
            )
            for i in range(len(config.depths))
        ]

    def call(
        self,
        hidden_states: tf.Tensor,
        output_hidden_states: bool,
        output_attentions: bool,
        return_dict: bool,
        training: bool = False,
    ) -> Union[tuple, TFBaseModelOutput]:
        all_hidden_states = () if output_hidden_states else None
        all_groupings = () if output_attentions else None

        group_tokens = None

        for stage in self.stages:
            if output_hidden_states:
                all_hidden_states = all_hidden_states + (hidden_states,)

            layer_outputs = stage(hidden_states, group_tokens, output_attentions)

            hidden_states = layer_outputs[0]
            group_tokens = layer_outputs[1]

            if output_attentions and layer_outputs[2] is not None:
                all_groupings = all_groupings + (layer_outputs[2],)

        if output_hidden_states:
            all_hidden_states = all_hidden_states + (hidden_states,)

        if not return_dict:
            return tuple(v for v in [hidden_states, all_hidden_states, all_groupings] if v is not None)
        return TFBaseModelOutput(
            last_hidden_state=hidden_states, hidden_states=all_hidden_states, attentions=all_groupings
        )



class TFGroupViTTextTransformer(tf.keras.layers.Layer):
    def __init__(self, config: GroupViTTextConfig, **kwargs):
        super().__init__(**kwargs)

        self.embeddings = TFGroupViTTextEmbeddings(config, name="embeddings")
        self.encoder = TFGroupViTTextEncoder(config, name="encoder")
        self.final_layer_norm = tf.keras.layers.LayerNormalization(
            epsilon=config.layer_norm_eps, name="final_layer_norm"
        )

        
        self.eos_token_id = config.eos_token_id

    def call(
        self,
        input_ids: TFModelInputType,
        attention_mask: tf.Tensor,
        position_ids: tf.Tensor,
        output_attentions: bool,
        output_hidden_states: bool,
        return_dict: bool,
        training: bool = False,
    ) -> Union[TFBaseModelOutputWithPooling, Tuple[tf.Tensor]]:
        input_shape = shape_list(input_ids)

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

        batch_size, seq_length = input_shape
        
        
        causal_attention_mask = self._build_causal_attention_mask(batch_size, seq_length, dtype=embedding_output.dtype)

        
        
        attention_mask = _expand_mask(attention_mask)

        encoder_outputs = self.encoder(
            hidden_states=embedding_output,
            attention_mask=attention_mask,
            causal_attention_mask=causal_attention_mask,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            training=training,
        )

        sequence_output = encoder_outputs[0]
        sequence_output = self.final_layer_norm(inputs=sequence_output)

        if self.eos_token_id == 2:
            
            
            
            
            
            pooled_output = tf.gather_nd(
                params=sequence_output,
                indices=tf.stack(
                    values=(tf.range(input_shape[0], dtype=tf.int64), tf.math.argmax(input_ids, axis=-1)), axis=1
                ),
            )
        else:
            
            pooled_output = tf.gather_nd(
                params=sequence_output,
                indices=tf.stack(
                    values=(
                        tf.range(input_shape[0], dtype=tf.int64),
                        tf.math.argmax(tf.cast(input_ids == self.eos_token_id, dtype=tf.int8), axis=-1),
                    ),
                    axis=1,
                ),
            )

        if not return_dict:
            return (sequence_output, pooled_output) + encoder_outputs[1:]

        return TFBaseModelOutputWithPooling(
            last_hidden_state=sequence_output,
            pooler_output=pooled_output,
            hidden_states=encoder_outputs.hidden_states,
            attentions=encoder_outputs.attentions,
        )

    def _build_causal_attention_mask(self, batch_size, seq_length, dtype=tf.float32):
        
        
        
        
        diag = tf.cast(tf.fill((seq_length,), 0.0), dtype)

        
        to_mask = tf.cast(tf.fill((seq_length, seq_length), -10000.0), dtype)

        
        
        to_mask = tf.linalg.band_part(to_mask, 0, -1)
        
        to_mask = tf.linalg.set_diag(to_mask, diagonal=diag)

        return tf.broadcast_to(input=to_mask, shape=(batch_size, 1, seq_length, seq_length))



class TFGroupViTVisionTransformer(tf.keras.layers.Layer):
    def __init__(self, config: GroupViTVisionConfig, **kwargs):
        super().__init__(**kwargs)

        self.embeddings = TFGroupViTVisionEmbeddings(config, name="embeddings")
        self.encoder = TFGroupViTVisionEncoder(config, name="encoder")
        self.layernorm = tf.keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="layernorm")

    def call(
        self,
        pixel_values: TFModelInputType,
        output_attentions: bool,
        output_hidden_states: bool,
        return_dict: bool,
        training: bool = False,
    ) -> Union[Tuple, TFBaseModelOutputWithPooling]:
        embedding_output = self.embeddings(pixel_values)

        encoder_outputs = self.encoder(
            hidden_states=embedding_output,
            output_hidden_states=output_hidden_states,
            output_attentions=output_attentions,
            return_dict=return_dict,
        )

        last_hidden_state = encoder_outputs[0]

        
        last_hidden_state = self.layernorm(last_hidden_state)
        pooled_output = tf.math.reduce_mean(last_hidden_state, axis=1)

        if not return_dict:
            return (last_hidden_state, pooled_output) + encoder_outputs[1:]

        return TFBaseModelOutputWithPooling(
            last_hidden_state=last_hidden_state,
            pooler_output=pooled_output,
            hidden_states=encoder_outputs.hidden_states,
            attentions=encoder_outputs.attentions,
        )


@keras_serializable

class TFGroupViTTextMainLayer(tf.keras.layers.Layer):
    config_class = GroupViTTextConfig

    def __init__(self, config: GroupViTTextConfig, **kwargs):
        super().__init__(**kwargs)
        self.config = config
        self.text_model = TFGroupViTTextTransformer(config, name="text_model")

    def get_input_embeddings(self) -> tf.keras.layers.Layer:
        return self.text_model.embeddings

    def set_input_embeddings(self, value: tf.Variable):
        self.text_model.embeddings.weight = value
        self.text_model.embeddings.vocab_size = shape_list(value)[0]

    @unpack_inputs
    def call(
        self,
        input_ids: TFModelInputType | None = None,
        attention_mask: np.ndarray | tf.Tensor | None = None,
        position_ids: np.ndarray | tf.Tensor | None = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        training: bool = False,
    ) -> Union[TFBaseModelOutputWithPooling, Tuple[tf.Tensor]]:
        if input_ids is None:
            raise ValueError("You have to specify input_ids")

        input_shape = shape_list(input_ids)

        if attention_mask is None:
            attention_mask = tf.fill(dims=input_shape, value=1)

        text_model_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,
            training=training,
        )

        return text_model_outputs


@keras_serializable

class TFGroupViTVisionMainLayer(tf.keras.layers.Layer):
    config_class = GroupViTVisionConfig

    def __init__(self, config: GroupViTVisionConfig, **kwargs):
        super().__init__(**kwargs)
        self.config = config
        self.vision_model = TFGroupViTVisionTransformer(config, name="vision_model")

    def get_input_embeddings(self) -> tf.keras.layers.Layer:
        return self.vision_model.embeddings

    @unpack_inputs
    def call(
        self,
        pixel_values: TFModelInputType | None = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        training: bool = False,
    ) -> Union[TFBaseModelOutputWithPooling, Tuple[tf.Tensor]]:
        if pixel_values is None:
            raise ValueError("You have to specify pixel_values")

        vision_model_outputs = self.vision_model(
            pixel_values=pixel_values,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            training=training,
        )

        return vision_model_outputs


@keras_serializable

class TFGroupViTMainLayer(tf.keras.layers.Layer):
    config_class = GroupViTConfig

    def __init__(self, config: GroupViTConfig, **kwargs):
        super().__init__(**kwargs)

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

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

        self.config = config

        text_config = config.text_config
        vision_config = config.vision_config

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

        self.text_model = TFGroupViTTextTransformer(text_config, name="text_model")
        self.vision_model = TFGroupViTVisionTransformer(vision_config, name="vision_model")

        self.visual_projection = [
            tf.keras.layers.Dense(self.projection_intermediate_dim, name="visual_projection.0"),
            tf.keras.layers.BatchNormalization(name="visual_projection.1", momentum=0.9, epsilon=1e-5),
            tf.keras.layers.ReLU(name="visual_projection.2"),
            tf.keras.layers.Dense(self.projection_dim, name="visual_projection.3"),
        ]
        self.text_projection = [
            tf.keras.layers.Dense(self.projection_intermediate_dim, name="text_projection.0"),
            tf.keras.layers.BatchNormalization(name="text_projection.1", momentum=0.9, epsilon=1e-5),
            tf.keras.layers.ReLU(name="text_projection.2"),
            tf.keras.layers.Dense(self.projection_dim, name="text_projection.3"),
        ]

    def build(self, input_shape: tf.TensorShape):
        self.logit_scale = self.add_weight(
            shape=(1,),
            initializer=tf.keras.initializers.Constant(self.config.logit_scale_init_value),
            trainable=True,
            name="logit_scale",
        )

        super().build(input_shape)

    @unpack_inputs
    def get_text_features(
        self,
        input_ids: TFModelInputType | None = None,
        attention_mask: np.ndarray | tf.Tensor | None = None,
        position_ids: np.ndarray | tf.Tensor | None = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        training: bool = False,
    ) -> tf.Tensor:
        if input_ids is None:
            raise ValueError("You have to specify either input_ids")

        input_shape = shape_list(input_ids)

        if attention_mask is None:
            attention_mask = tf.fill(dims=input_shape, value=1)

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

        pooled_output = text_outputs[1]
        for layer in self.text_projection:
            pooled_output = layer(pooled_output)

        text_features = pooled_output
        return text_features

    @unpack_inputs
    def get_image_features(
        self,
        pixel_values: TFModelInputType | None = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        training: bool = False,
    ) -> tf.Tensor:
        if pixel_values is None:
            raise ValueError("You have to specify pixel_values")

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

        pooled_output = vision_outputs[1]
        for layer in self.visual_projection:
            pooled_output = layer(pooled_output)

        image_features = pooled_output
        return image_features

    @unpack_inputs
    def call(
        self,
        input_ids: TFModelInputType | None = None,
        pixel_values: TFModelInputType | None = None,
        attention_mask: np.ndarray | tf.Tensor | None = None,
        position_ids: np.ndarray | tf.Tensor | None = None,
        return_loss: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        output_segmentation: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        training: bool = False,
    ) -> Union[TFGroupViTModelOutput, Tuple[tf.Tensor]]:
        if input_ids is None:
            raise ValueError("You have to specify either input_ids")
        if pixel_values is None:
            raise ValueError("You have to specify pixel_values")

        input_shape = shape_list(input_ids)

        if attention_mask is None:
            attention_mask = tf.fill(dims=input_shape, value=1)
        if output_segmentation:
            output_attentions = True
        vision_outputs = self.vision_model(
            pixel_values=pixel_values,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            training=training,
        )

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

        image_embeds = vision_outputs[1]
        for layer in self.visual_projection:
            image_embeds = layer(image_embeds)

        text_embeds = text_outputs[1]
        for layer in self.text_projection:
            text_embeds = layer(text_embeds)

        
        image_embeds = image_embeds / tf.norm(image_embeds, axis=-1, keepdims=True)
        text_embeds = text_embeds / tf.norm(text_embeds, axis=-1, keepdims=True)

        
        logit_scale = tf.math.exp(self.logit_scale)
        logits_per_text = tf.matmul(text_embeds, image_embeds, transpose_b=True) * logit_scale
        logits_per_image = tf.transpose(logits_per_text)

        seg_logits = None
        if output_segmentation:
            
            
            image_group_embeds = vision_outputs[0]
            
            image_group_embeds = tf.reshape(image_group_embeds, shape=(-1, shape_list(image_group_embeds)[-1]))
            for layer in self.visual_projection:
                image_group_embeds = layer(image_group_embeds)
            if output_hidden_states:
                attentions = vision_outputs[3]
            else:
                attentions = vision_outputs[2]
            
            grouping = get_grouping_from_attentions(attentions, pixel_values.shape[2:])

            
            image_group_embeds = image_group_embeds / tf.norm(
                tensor=image_group_embeds, ord="euclidean", axis=-1, keepdims=True
            )
            
            logits_per_image_group = tf.matmul(image_group_embeds, text_embeds, transpose_b=True) * logit_scale
            
            logits_per_image_group = tf.reshape(
                logits_per_image_group, shape=(image_embeds.shape[0], -1, text_embeds.shape[0])
            )
            logits_per_image_group = tf.transpose(logits_per_image_group, perm=(0, 2, 1))

            
            flatten_grouping = tf.reshape(grouping, shape=(shape_list(grouping)[0], shape_list(grouping)[1], -1))

            
            seg_logits = tf.matmul(logits_per_image_group, flatten_grouping) * logit_scale
            seg_logits = tf.reshape(
                seg_logits, shape=(seg_logits.shape[0], seg_logits.shape[1], grouping.shape[2], grouping.shape[3])
            )

        loss = None
        if return_loss:
            loss = groupvit_loss(logits_per_text)[None, ...]

        if not return_dict:
            if seg_logits is not None:
                output = (
                    logits_per_image,
                    logits_per_text,
                    seg_logits,
                    text_embeds,
                    image_embeds,
                    text_outputs,
                    vision_outputs,
                )
            else:
                output = (logits_per_image, logits_per_text, text_embeds, image_embeds, text_outputs, vision_outputs)
            return ((loss,) + output) if loss is not None else output

        return TFGroupViTModelOutput(
            loss=loss,
            logits_per_image=logits_per_image,
            logits_per_text=logits_per_text,
            segmentation_logits=seg_logits,
            text_embeds=text_embeds,
            image_embeds=image_embeds,
            text_model_output=text_outputs,
            vision_model_output=vision_outputs,
        )


class TFGroupViTPreTrainedModel(TFPreTrainedModel):
    

    config_class = GroupViTConfig
    base_model_prefix = "groupvit"


GROUPVIT_START_DOCSTRING = r

GROUPVIT_TEXT_INPUTS_DOCSTRING = r

GROUPVIT_VISION_INPUTS_DOCSTRING = r

GROUPVIT_INPUTS_DOCSTRING = r


class TFGroupViTTextModel(TFGroupViTPreTrainedModel):
    config_class = GroupViTTextConfig
    main_input_name = "input_ids"

    def __init__(self, config: GroupViTTextConfig, *inputs, **kwargs):
        super().__init__(config, *inputs, **kwargs)

        self.groupvit = TFGroupViTTextMainLayer(config, name="groupvit")

    @unpack_inputs
    @add_start_docstrings_to_model_forward(GROUPVIT_TEXT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
    @replace_return_docstrings(output_type=TFBaseModelOutputWithPooling, config_class=GroupViTTextConfig)
    def call(
        self,
        input_ids: TFModelInputType | None = None,
        attention_mask: np.ndarray | tf.Tensor | None = None,
        position_ids: np.ndarray | tf.Tensor | None = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        training: bool = False,
    ) -> Union[TFBaseModelOutputWithPooling, Tuple[tf.Tensor]]:
        r

        outputs = self.groupvit(
            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,
            training=training,
        )

        return outputs


class TFGroupViTVisionModel(TFGroupViTPreTrainedModel):
    config_class = GroupViTVisionConfig
    main_input_name = "pixel_values"

    def __init__(self, config: GroupViTVisionConfig, *inputs, **kwargs):
        super().__init__(config, *inputs, **kwargs)

        self.groupvit = TFGroupViTVisionMainLayer(config, name="groupvit")

    @unpack_inputs
    @add_start_docstrings_to_model_forward(GROUPVIT_VISION_INPUTS_DOCSTRING)
    @replace_return_docstrings(output_type=TFBaseModelOutputWithPooling, config_class=GroupViTVisionConfig)
    def call(
        self,
        pixel_values: TFModelInputType | None = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        training: bool = False,
    ) -> Union[TFBaseModelOutputWithPooling, Tuple[tf.Tensor]]:
        r

        outputs = self.groupvit(
            pixel_values=pixel_values,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            training=training,
        )

        return outputs


@add_start_docstrings(GROUPVIT_START_DOCSTRING)
class TFGroupViTModel(TFGroupViTPreTrainedModel):
    config_class = GroupViTConfig

    def __init__(self, config: GroupViTConfig, *inputs, **kwargs):
        super().__init__(config, *inputs, **kwargs)

        self.groupvit = TFGroupViTMainLayer(config, name="groupvit")

    @unpack_inputs
    @add_start_docstrings_to_model_forward(GROUPVIT_TEXT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
    def get_text_features(
        self,
        input_ids: TFModelInputType | None = None,
        attention_mask: np.ndarray | tf.Tensor | None = None,
        position_ids: np.ndarray | tf.Tensor | None = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        training: bool = False,
    ) -> tf.Tensor:
        r

        text_features = self.groupvit.get_text_features(
            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,
            training=training,
        )

        return text_features

    @unpack_inputs
    @add_start_docstrings_to_model_forward(GROUPVIT_VISION_INPUTS_DOCSTRING)
    def get_image_features(
        self,
        pixel_values: TFModelInputType | None = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        training: bool = False,
    ) -> tf.Tensor:
        r

        image_features = self.groupvit.get_image_features(
            pixel_values=pixel_values,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            training=training,
        )

        return image_features

    @unpack_inputs
    @add_start_docstrings_to_model_forward(GROUPVIT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
    @replace_return_docstrings(output_type=TFGroupViTModelOutput, config_class=GroupViTConfig)
    def call(
        self,
        input_ids: TFModelInputType | None = None,
        pixel_values: TFModelInputType | None = None,
        attention_mask: np.ndarray | tf.Tensor | None = None,
        position_ids: np.ndarray | tf.Tensor | None = None,
        return_loss: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        output_segmentation: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        training: bool = False,
    ) -> Union[TFGroupViTModelOutput, Tuple[tf.Tensor]]:
        r

        outputs = self.groupvit(
            input_ids=input_ids,
            pixel_values=pixel_values,
            attention_mask=attention_mask,
            position_ids=position_ids,
            return_loss=return_loss,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            output_segmentation=output_segmentation,
            return_dict=return_dict,
            training=training,
        )

        return outputs

    def serving_output(self, output: TFGroupViTModelOutput) -> TFGroupViTModelOutput:
        
        
        
        return output