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 pytorch groupvit model import collections abc import math from dataclasses import dataclass from typing import any optional tuple union import numpy as np import torch import torch utils checkpoint from torch import nn from activations import act2fn from modelingattnmaskutils import create4dcausalattentionmask prepare4dattentionmask from modelingoutputs import basemodeloutput basemodeloutputwithpooling from modelingutils import pretrainedmodel from utils import modeloutput addstartdocstrings addstartdocstringstomodelforward logging replacereturndocstrings from configurationgroupvit import groupvitconfig groupvittextconfig groupvitvisionconfig logger logging getloggername checkpointfordoc nvidiagroupvitgccyfcc groupvitpretrainedmodelarchivelist nvidiagroupvitgccyfcc see all groupvit models at https huggingface comodels filtergroupvit contrastive loss function adapted from https sachinruk github ioblogpytorchpytorch20lightningloss20functiongpu20210307clip html def contrastivelosslogits torch tensor torch tensor return nn functional crossentropylogits torch arangelenlogits devicelogits device copied from transformers models clip modelingclip cliploss with clipgroupvit def groupvitlosssimilarity torch tensor torch tensor captionloss contrastivelosssimilarity imageloss contrastivelosssimilarity t return captionloss imageloss 2 0 def hardsoftmaxlogits torch tensor dim int ysoft logits softmaxdim straight through index ysoft maxdim keepdimtrue1 yhard torch zeroslikelogits memoryformattorch legacycontiguousformat scatterdim index 1 0 ret yhard ysoft detach ysoft return ret def gumbelsoftmaxlogits torch tensor tau float 1 hard bool false dim int 1 torch tensor more stable https github compytorchpytorchissues41663 gumbeldist torch distributions gumbel gumbel torch tensor0 0 devicelogits device dtypelogits dtype torch tensor1 0 devicelogits device dtypelogits dtype gumbels gumbeldist samplelogits shape gumbels logits gumbels tau gumbellogits tau ysoft gumbels softmaxdim if hard straight through index ysoft maxdim keepdimtrue1 yhard torch zeroslikelogits memoryformattorch legacycontiguousformat scatterdim index 1 0 ret yhard ysoft detach ysoft else reparametrization trick ret ysoft return ret def resizeattentionmapattentions height width aligncornersfalse scale height width attentions shape2 0 5 if height width featwidth intnp roundwidth scale featheight attentions shape2 featwidth else featheight intnp roundheight scale featwidth attentions shape2 featheight batchsize attentions shape0 groups attentions shape1 number of group token batchsize groups heightwidth groups batchsize groups height width attentions attentions reshapebatchsize groups featheight featwidth attentions nn functional interpolate attentions sizeheight width modebilinear aligncornersaligncorners return attentions def getgroupingfromattentionsattentions hwshape attnmaps with torch nograd prevattnmasks none for attnmasks in attentions batchsize numgroups height x width batchsize height x width numgroups attnmasks attnmasks permute0 2 1 contiguous if prevattnmasks is none prevattnmasks attnmasks else prevattnmasks prevattnmasks attnmasks batchsize heightxwidth numgroups batchsize numgroups heightxwidth batchsize numgroups height width curattnmap resizeattentionmapprevattnmasks permute0 2 1 contiguous hwshape attnmaps appendcurattnmap batchsize numgroups height width finalgrouping attnmaps1 return finalgrouping class groupvitcrossattentionlayernn module def initself config groupvitvisionconfig super init self attn groupvitattentionconfig self norm2 nn layernormconfig hiddensize epsconfig layernormeps self mlp groupvitmlpconfig self normpost nn layernormconfig hiddensize epsconfig layernormeps def forwardself query key x query x x self attnquery encoderhiddenstateskey0 x x self mlpself norm2x x self normpostx return x class groupvitassignattentionnn module def initself config groupvitvisionconfig super init self scale config hiddensize0 5 self qproj nn linearconfig hiddensize config hiddensize self kproj nn linearconfig hiddensize config hiddensize self vproj nn linearconfig hiddensize config hiddensize self proj nn linearconfig hiddensize config hiddensize self assigneps config assigneps def getattnself attn gumbeltrue hardtrue if gumbel and self training attn gumbelsoftmaxattn dim2 hardhard else if hard attn hardsoftmaxattn dim2 else attn nn functional softmaxattn dim2 return attn def forwardself query key value key batchsize querylength channels query self qprojquery batchsize keylength channels key self kprojkey batchsize keylength channels value self vprojvalue batchsize querylength keylength rawattn query key transpose2 1 self scale attn self getattnrawattn softattn self getattnrawattn gumbelfalse hardfalse attn attn attn sumdim1 keepdimtrue self assigneps out attn value out self projout return out softattn class groupvittokenassignnn module def initself config groupvitvisionconfig numgrouptoken numoutputgroup super init self numoutputgroup numoutputgroup norm on grouptokens self normtokens nn layernormconfig hiddensize epsconfig layernormeps 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 groupvitmixermlpconfig numgrouptoken tokensdim numoutputgroup self normposttokens nn layernormconfig hiddensize epsconfig layernormeps norm on x self normx nn layernormconfig hiddensize epsconfig layernormeps self preassignattn groupvitcrossattentionlayerconfig self assign groupvitassignattentionconfig self normnewx nn layernormconfig hiddensize epsconfig layernormeps self mlpchannels groupvitmlpconfig config hiddensize channelsdim config hiddensize def projectgrouptokenself grouptokens b numoutputgroups c b numgrouptokens c projectedgrouptokens self mlpintergrouptokens projectedgrouptokens self normposttokensprojectedgrouptokens return projectedgrouptokens def forwardself imagetokens grouptokens 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 dataclass class groupvitmodeloutputmodeloutput loss optionaltorch floattensor none logitsperimage torch floattensor none logitspertext torch floattensor none segmentationlogits torch floattensor none textembeds torch floattensor none imageembeds torch floattensor none textmodeloutput basemodeloutputwithpooling none visionmodeloutput basemodeloutputwithpooling none def totupleself tupleany return tuple selfk if k not in textmodeloutput visionmodeloutput else getattrself k totuple for k in self keys class groupvitpatchembeddingsnn module def init self imagesize int 224 patchsize unionint tupleint int 16 numchannels int 3 embeddim int 768 super init 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 projection nn conv2dnumchannels embeddim kernelsizepatchsize stridepatchsize def forwardself pixelvalues torch tensor interpolateposencoding bool false torch tensor batchsize numchannels height width pixelvalues shape if not interpolateposencoding if height self imagesize0 or width self imagesize1 raise valueerror finput image size heightwidth doesn t match model f self imagesize0self imagesize1 x self projectionpixelvalues flatten2 transpose1 2 return x class groupvitvisionembeddingsnn module def initself config groupvitvisionconfig super init self patchembeddings groupvitpatchembeddings imagesizeconfig imagesize patchsizeconfig patchsize numchannelsconfig numchannels embeddimconfig hiddensize numpatches self patchembeddings numpatches self positionembeddings nn parametertorch zeros1 numpatches config hiddensize self dropout nn dropoutconfig dropout self layernorm nn layernormconfig hiddensize epsconfig layernormeps self config config def interpolateposencodingself embeddings torch tensor height int width int torch tensor npatch embeddings shape1 if npatch self positionembeddings shape1 and height width return self positionembeddings patchposembed self positionembeddings numoriginalposembed patchposembed shape1 dim embeddings shape1 featheight height self config patchsize featwidth width self config patchsize we add a small number to avoid floating point error in the interpolation see discussion at https github comfacebookresearchdinoissues8 featheight featwidth featheight 0 1 featwidth 0 1 originalheight originalwidth math sqrtnumoriginalposembed reshapedpatchposembed patchposembed reshape1 intoriginalheight intoriginalwidth dim permute 0 3 1 2 scalefactor featheight originalheight featwidth originalwidth patchposembed nn functional interpolate reshapedpatchposembed scalefactorscalefactor modebicubic aligncornersfalse patchposembed patchposembed permute0 2 3 1 view1 1 dim return patchposembed def forwardself pixelvalues torch tensor interpolateposencoding bool false torch tensor batchsize numchannels height width pixelvalues shape embeddings self patchembeddingspixelvalues interpolateposencodinginterpolateposencoding embeddings self layernormembeddings batchsize seqlen embeddings size 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 modelingclip cliptextembeddings with clipgroupvit class groupvittextembeddingsnn module def initself config groupvittextconfig super init embeddim config hiddensize self tokenembedding nn embeddingconfig vocabsize embeddim self positionembedding nn embeddingconfig maxpositionembeddings embeddim positionids 1 len position emb is contiguous in memory and exported when serialized self registerbuffer positionids torch arangeconfig maxpositionembeddings expand1 1 persistentfalse def forward self inputids optionaltorch longtensor none positionids optionaltorch longtensor none inputsembeds optionaltorch floattensor none torch tensor seqlength inputids shape1 if inputids is not none else inputsembeds shape2 if positionids is none positionids self positionids seqlength if inputsembeds is none inputsembeds self tokenembeddinginputids positionembeddings self positionembeddingpositionids embeddings inputsembeds positionembeddings return embeddings class groupvitstagenn module args hiddenstates torch floattensor input to the layer of shape batch seqlen embeddim attentionmask torch floattensor attention mask of size batch 1 tgtlen srclen where padding elements are indicated by very large negative values config encoderattentionheads outputattentions bool optional whether or not to return the grouping tensors of grouping block multiheaded attention from attention is all you need paper def initself config super init self config config self embeddim config hiddensize self numheads config numattentionheads self headdim self embeddim self numheads if self headdim self numheads self embeddim raise valueerror fembeddim must be divisible by numheads got embeddim self embeddim and numheads f self numheads self scale self headdim0 5 self dropout config attentiondropout self kproj nn linearself embeddim self embeddim self vproj nn linearself embeddim self embeddim self qproj nn linearself embeddim self embeddim self outproj nn linearself embeddim self embeddim def shapeself tensor torch tensor seqlen int bsz int return tensor viewbsz seqlen self numheads self headdim transpose1 2 contiguous def forward self hiddenstates torch tensor attentionmask optionaltorch tensor none causalattentionmask optionaltorch tensor none encoderhiddenstates optionaltorch floattensor none outputattentions optionalbool false tupletorch tensor optionaltorch tensor optionaltupletorch tensor get query proj apply the causalattentionmask first this operation is a bit akward but it s required to make sure that attnweights keeps its gradient in order to do so attnweights have to reshaped twice and have to be reused in the following copied from transformers models clip modelingclip clipencoderlayer with clipgroupvit args hiddenstates torch floattensor input to the layer of shape batch seqlen embeddim attentionmask torch floattensor attention mask of size batch 1 tgtlen srclen where padding elements are indicated by very large negative values config encoderattentionheads outputattentions bool optional whether or not to return the attentions tensors of all attention layers see attentions under returned tensors for more detail an abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models initialize the weights initrange self config initializerrange if isinstancemodule nn linear nn conv2d slightly different from the tf version which uses truncatednormal for initialization cf https github compytorchpytorchpull5617 module weight data normalmean0 0 stdinitrange if module bias is not none module bias data zero elif isinstancemodule nn layernorm module bias data zero module weight data fill1 0 factor self config initializerfactor if isinstancemodule groupvittextembeddings module tokenembedding weight data normalmean0 0 stdfactor 0 02 module positionembedding weight data normalmean0 0 stdfactor 0 02 elif isinstancemodule groupvitattention factor self config initializerfactor inprojstd module embeddim0 5 2 module config numhiddenlayers 0 5 factor outprojstd module embeddim0 5 factor nn init normalmodule qproj weight stdinprojstd nn init normalmodule kproj weight stdinprojstd nn init normalmodule vproj weight stdinprojstd nn init normalmodule outproj weight stdoutprojstd elif isinstancemodule groupvitmlp factor self config initializerfactor inprojstd module config hiddensize0 5 2 module config numhiddenlayers 0 5 factor fcstd 2 module config hiddensize 0 5 factor nn init normalmodule fc1 weight stdfcstd nn init normalmodule fc2 weight stdinprojstd groupvitstartdocstring r this model is a pytorch torch nn modulehttps pytorch orgdocsstablenn htmltorch nn module subclass use it as a regular pytorch module and refer to the pytorch documentation for all matter related to general usage and behavior parameters config 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 frompretrained method to load the model weights groupvitvisioninputsdocstring r args pixelvalues torch floattensor of shape batchsize numchannels height width pixel values padding will be ignored by default should you provide it pixel values can be obtained using autoimageprocessor see clipimageprocessor call for details outputattentions bool optional whether or not to return the attentions tensors of all attention layers see attentions under returned tensors for more detail outputhiddenstates bool optional whether or not to return the hidden states of all layers see hiddenstates under returned tensors for more detail returndict bool optional whether or not to return a utils modeloutput instead of a plain tuple class groupvitvisionencodernn module def initself config groupvitvisionconfig none super init self config config self stages nn modulelist groupvitstage configconfig depthconfig depthsi numgrouptokenconfig numgrouptokensi numoutputgroupconfig numoutputgroupsi numprevgrouptokenconfig numoutputgroupsi 1 if i 0 else 0 for i in rangelenconfig depths self gradientcheckpointing false def forward self hiddenstates torch tensor outputhiddenstates optionalbool none outputattentions optionalbool none returndict optionalbool none uniontuple basemodeloutput outputattentions outputattentions if outputattentions is not none else self config outputattentions outputhiddenstates outputhiddenstates if outputhiddenstates is not none else self config outputhiddenstates returndict returndict if returndict is not none else self config usereturndict allhiddenstates if outputhiddenstates else none allgroupings if outputattentions else none grouptokens none for i stage in enumerateself stages if outputhiddenstates allhiddenstates allhiddenstates hiddenstates layeroutputs stagehiddenstates grouptokens outputattentions hiddenstates layeroutputs0 grouptokens layeroutputs1 if outputattentions and layeroutputs2 is not none allgroupings allgroupings layeroutputs2 if outputhiddenstates allhiddenstates allhiddenstates hiddenstates if not returndict return tuplev for v in hiddenstates allhiddenstates allgroupings if v is not none return basemodeloutput lasthiddenstatehiddenstates hiddenstatesallhiddenstates attentionsallgroupings class groupvittextencodernn module def initself config groupvittextconfig super init self config config self layers nn modulelistgroupvitencoderlayerconfig for in rangeconfig numhiddenlayers self gradientcheckpointing false def forward self inputsembeds attentionmask optionaltorch tensor none causalattentionmask optionaltorch tensor none outputattentions optionalbool none outputhiddenstates optionalbool none returndict optionalbool none uniontuple basemodeloutput r args inputsembeds torch floattensor of shape batchsize sequencelength hiddensize optionally instead of passing inputids you can choose to directly pass an embedded representation this is useful if you want more control over how to convert inputids indices into associated vectors than the model s internal embedding lookup matrix attentionmask torch tensor of shape batchsize sequencelength optional mask to avoid performing attention on padding token indices mask values selected in 0 1 1 for tokens that are not masked 0 for tokens that are masked what are attention masks glossaryattentionmask causalattentionmask torch tensor of shape batchsize sequencelength optional causal mask for the text model mask values selected in 0 1 1 for tokens that are not masked 0 for tokens that are masked what are attention masks glossaryattentionmask outputattentions bool optional whether or not to return the attentions tensors of all attention layers see attentions under returned tensors for more detail outputhiddenstates bool optional whether or not to return the hidden states of all layers see hiddenstates under returned tensors for more detail returndict bool optional whether or not to return a utils modeloutput instead of a plain tuple copied from transformers models clip modelingclip cliptexttransformer with cliptextgroupvittext clipencodergroupvittextencoder cliptextgroupvittext for pooledoutput computation outputattentions outputattentions if outputattentions is not none else self config outputattentions outputhiddenstates outputhiddenstates if outputhiddenstates is not none else self config outputhiddenstates returndict returndict if returndict is not none else self config usereturndict if inputids is none raise valueerroryou have to specify inputids inputshape inputids size inputids inputids view1 inputshape1 hiddenstates self embeddingsinputidsinputids positionidspositionids clip s text model uses causal mask prepare it here https github comopenaiclipblobcfcffb90e69f37bf2ff1e988237a0fbe41f33c04clipmodel pyl324 causalattentionmask create4dcausalattentionmask inputshape hiddenstates dtype devicehiddenstates device expand attentionmask if attentionmask is not none bsz seqlen bsz 1 tgtseqlen srcseqlen attentionmask prepare4dattentionmaskattentionmask hiddenstates dtype encoderoutputs self encoder inputsembedshiddenstates attentionmaskattentionmask causalattentionmaskcausalattentionmask outputattentionsoutputattentions outputhiddenstatesoutputhiddenstates returndictreturndict lasthiddenstate encoderoutputs0 lasthiddenstate self finallayernormlasthiddenstate if self eostokenid 2 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 sequencelength transformer width take features from the eot embedding eottoken is the highest number in each sequence casting to torch int for onnx compatibility argmax doesn t support int64 inputs with opset 14 pooledoutput lasthiddenstate torch arangelasthiddenstate shape0 devicelasthiddenstate device inputids todtypetorch int devicelasthiddenstate device argmaxdim1 else the config gets updated eostokenid from pr 24773 so the use of exta new tokens is possible pooledoutput lasthiddenstate torch arangelasthiddenstate shape0 devicelasthiddenstate device we need to get the first position of eostokenid value padtokenids might equal to eostokenid inputids todtypetorch int devicelasthiddenstate device self eostokenid int argmaxdim1 if not returndict return lasthiddenstate pooledoutput encoderoutputs1 return basemodeloutputwithpooling lasthiddenstatelasthiddenstate pooleroutputpooledoutput hiddenstatesencoderoutputs hiddenstates attentionsencoderoutputs attentions class groupvittextmodelgroupvitpretrainedmodel configclass groupvittextconfig def initself config groupvittextconfig super initconfig self textmodel groupvittexttransformerconfig initialize weights and apply final processing self postinit def getinputembeddingsself nn module return self textmodel embeddings tokenembedding def setinputembeddingsself value self textmodel embeddings tokenembedding value addstartdocstringstomodelforwardgroupvittextinputsdocstring replacereturndocstringsoutputtypebasemodeloutputwithpooling configclassgroupvittextconfig def forward self inputids optionaltorch tensor none attentionmask optionaltorch tensor none positionids optionaltorch tensor none outputattentions optionalbool none outputhiddenstates optionalbool none returndict optionalbool none uniontuple basemodeloutputwithpooling r returns examples python from transformers import cliptokenizer groupvittextmodel tokenizer cliptokenizer frompretrainednvidiagroupvitgccyfcc model groupvittextmodel frompretrainednvidiagroupvitgccyfcc inputs tokenizera photo of a cat a photo of a dog paddingtrue returntensorspt outputs modelinputs lasthiddenstate outputs lasthiddenstate pooledoutput outputs pooleroutput pooled eos token states return self textmodel inputidsinputids attentionmaskattentionmask positionidspositionids outputattentionsoutputattentions outputhiddenstatesoutputhiddenstates returndictreturndict class groupvitvisiontransformernn module def initself config groupvitvisionconfig super init self config config embeddim config hiddensize self embeddings groupvitvisionembeddingsconfig self encoder groupvitvisionencoderconfig self layernorm nn layernormembeddim epsconfig layernormeps addstartdocstringstomodelforwardgroupvitvisioninputsdocstring replacereturndocstringsoutputtypebasemodeloutputwithpooling configclassgroupvitvisionconfig def forward self pixelvalues optionaltorch floattensor none outputhiddenstates optionalbool none outputattentions optionalbool none returndict optionalbool none uniontuple basemodeloutputwithpooling r returns normalize the last hidden state initialize weights and apply final processing initialize weights and apply final processing use groupvit model s config for some fields if specified instead of those of vision text components use groupvit model s config for some fields if specified instead of those of vision text components use groupvit model s config for some fields if specified instead of those of vision text components 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 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 pytorch groupvit model see all groupvit models at https huggingface co models filter groupvit contrastive loss function adapted from https sachinruk github io blog pytorch pytorch 20lightning loss 20function gpu 2021 03 07 clip html copied from transformers models clip modeling_clip clip_loss with clip groupvit straight through more stable https github com pytorch pytorch issues 41663 gumbel logits tau straight through reparametrization trick args attentions torch 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 torch tensor resized attention map of shape batch_size groups height width number of group token batch_size groups height width groups batch_size groups height width args attentions tuple torch floattensor tuple of attention maps returned by groupvitvisiontransformer hw_shape tuple int height and width of the output attention map returns torch 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 heightxwidth num_groups batch_size num_groups heightxwidth batch_size num_groups height width batch_size num_groups height width 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 torch tensor group tokens batch_size num_group_tokens channels returns projected_group_tokens torch tensor batch_size num_output_groups channels b num_output_groups c b num_group_tokens c args image_tokens torch tensor image tokens of shape batch_size input_length channels group_tokens torch tensor group tokens batch_size num_group_tokens channels batch_size num_output_groups channels args loss torch floattensor of shape 1 optional returned when return_loss is true contrastive loss for image text similarity logits_per_image torch floattensor 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 torch floattensor 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 torch floattensor 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 torch floattensor of shape batch_size output_dim the text embeddings obtained by applying the projection layer to the pooled output of groupvittextmodel image_embeds torch floattensor of shape batch_size output_dim the image embeddings obtained by applying the projection layer to the pooled output of groupvitvisionmodel text_model_output basemodeloutputwithpooling the output of the groupvittextmodel vision_model_output basemodeloutputwithpooling the output of the groupvitvisionmodel image to patch embedding 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 we add a small number to avoid floating point error in the interpolation see discussion at https github com facebookresearch dino issues 8 add positional encoding to each token copied from transformers models clip modeling_clip cliptextembeddings with clip groupvit position_ids 1 len position emb is contiguous in memory and exported when serialized this corresponds to the groupinglayer class in the groupvit implementation args hidden_states torch floattensor input to the layer of shape batch seq_len embed_dim attention_mask torch floattensor attention mask of size batch 1 tgt_len src_len where padding elements are indicated by very large negative values config encoder_attention_heads output_attentions bool optional whether or not to return the grouping tensors of grouping block multi headed attention from attention is all you need paper input shape batch x time x channel get query proj apply the causal_attention_mask first this operation is a bit akward but it s required to make sure that attn_weights keeps its gradient in order to do so attn_weights have to reshaped twice and have to be reused in the following copied from transformers models clip modeling_clip clipencoderlayer with clip groupvit args hidden_states torch floattensor input to the layer of shape batch seq_len embed_dim attention_mask torch floattensor attention mask of size batch 1 tgt_len src_len where padding elements are indicated by very large negative values config encoder_attention_heads output_attentions bool optional whether or not to return the attentions tensors of all attention layers see attentions under returned tensors for more detail an abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models initialize the weights slightly different from the tf version which uses truncated_normal for initialization cf https github com pytorch pytorch pull 5617 this model is a pytorch torch nn module https pytorch org docs stable nn html torch nn module subclass use it as a regular pytorch module and refer to the pytorch documentation for all matter related to general usage and behavior parameters config 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 torch longtensor of shape batch_size sequence_length indices of input sequence tokens in the vocabulary padding will be ignored by default should you provide it indices can be obtained using cliptokenizer see pretrainedtokenizer encode and pretrainedtokenizer __call__ for details what are input ids glossary input ids attention_mask torch tensor of shape batch_size sequence_length optional mask to avoid performing attention on padding token indices mask values selected in 0 1 1 for tokens that are not masked 0 for tokens that are masked what are attention masks glossary attention mask position_ids torch longtensor of shape batch_size sequence_length optional indices of positions of each input sequence tokens in the position embeddings selected in the range 0 config max_position_embeddings 1 what are position ids glossary position ids output_attentions bool optional whether or not to return the attentions tensors of all attention layers see attentions under returned tensors for more detail output_hidden_states bool optional whether or not to return the hidden states of all layers see hidden_states under returned tensors for more detail return_dict bool optional whether or not to return a utils modeloutput instead of a plain tuple args pixel_values torch floattensor of shape batch_size num_channels height width pixel values padding will be ignored by default should you provide it pixel values can be obtained using autoimageprocessor see clipimageprocessor __call__ for details output_attentions bool optional whether or not to return the attentions tensors of all attention layers see attentions under returned tensors for more detail output_hidden_states bool optional whether or not to return the hidden states of all layers see hidden_states under returned tensors for more detail return_dict bool optional whether or not to return a utils modeloutput instead of a plain tuple args input_ids torch longtensor of shape batch_size sequence_length indices of input sequence tokens in the vocabulary padding will be ignored by default should you provide it indices can be obtained using cliptokenizer see pretrainedtokenizer encode and pretrainedtokenizer __call__ for details what are input ids glossary input ids attention_mask torch tensor of shape batch_size sequence_length optional mask to avoid performing attention on padding token indices mask values selected in 0 1 1 for tokens that are not masked 0 for tokens that are masked what are attention masks glossary attention mask position_ids torch longtensor of shape batch_size sequence_length optional indices of positions of each input sequence tokens in the position embeddings selected in the range 0 config max_position_embeddings 1 what are position ids glossary position ids pixel_values torch floattensor of shape batch_size num_channels height width pixel values pixel values can be obtained using autoimageprocessor see clipimageprocessor __call__ for details return_loss bool optional whether or not to return the contrastive loss output_attentions bool optional whether or not to return the attentions tensors of all attention layers see attentions under returned tensors for more detail output_hidden_states bool optional whether or not to return the hidden states of all layers see hidden_states under returned tensors for more detail return_dict bool optional whether or not to return a utils modeloutput instead of a plain tuple transformer encoder consisting of config num_hidden_layers self attention layers each layer is a groupvitencoderlayer args config groupvittextconfig args inputs_embeds torch floattensor of shape batch_size sequence_length hidden_size optionally instead of passing input_ids you can choose to directly pass an embedded representation this is useful if you want more control over how to convert input_ids indices into associated vectors than the model s internal embedding lookup matrix attention_mask torch tensor of shape batch_size sequence_length optional mask to avoid performing attention on padding token indices mask values selected in 0 1 1 for tokens that are not masked 0 for tokens that are masked what are attention masks glossary attention mask causal_attention_mask torch tensor of shape batch_size sequence_length optional causal mask for the text model mask values selected in 0 1 1 for tokens that are not masked 0 for tokens that are masked what are attention masks glossary attention mask output_attentions bool optional whether or not to return the attentions tensors of all attention layers see attentions under returned tensors for more detail output_hidden_states bool optional whether or not to return the hidden states of all layers see hidden_states under returned tensors for more detail return_dict bool optional whether or not to return a utils modeloutput instead of a plain tuple copied from transformers models clip modeling_clip cliptexttransformer with cliptext groupvittext clipencoder groupvittextencoder clip_text groupvit_text for pooled_output computation returns clip s text model uses causal mask prepare it here https github com openai clip blob cfcffb90e69f37bf2ff1e988237a0fbe41f33c04 clip model py l324 expand attention_mask 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 sequence_length transformer width take features from the eot embedding eot_token is the highest number in each sequence casting to torch int for onnx compatibility argmax doesn t support int64 inputs with opset 14 the config gets updated eos_token_id from pr 24773 so the use of exta new tokens is possible we need to get the first position of eos_token_id value pad_token_ids might equal to eos_token_id initialize weights and apply final processing returns examples python from transformers import cliptokenizer groupvittextmodel tokenizer cliptokenizer from_pretrained nvidia groupvit gcc yfcc model groupvittextmodel from_pretrained nvidia groupvit gcc yfcc inputs tokenizer a photo of a cat a photo of a dog padding true return_tensors pt outputs model inputs last_hidden_state outputs last_hidden_state pooled_output outputs pooler_output pooled eos token states returns normalize the last hidden state initialize weights and apply final processing returns examples python from pil import image import requests from transformers import autoprocessor groupvitvisionmodel processor autoprocessor from_pretrained nvidia groupvit gcc yfcc model groupvitvisionmodel 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 pt outputs model inputs last_hidden_state outputs last_hidden_state pooled_output outputs pooler_output pooled cls states initialize weights and apply final processing returns text_features torch floattensor of shape batch_size output_dim the text embeddings obtained by applying the projection layer to the pooled output of groupvittextmodel examples python from transformers import cliptokenizer groupvitmodel model groupvitmodel 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 pt text_features model get_text_features inputs use groupvit model s config for some fields if specified instead of those of vision text components returns image_features torch floattensor of shape batch_size output_dim the image embeddings obtained by applying the projection layer to the pooled output of groupvitvisionmodel examples python from pil import image import requests from transformers import autoprocessor groupvitmodel model groupvitmodel 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 pt image_features model get_image_features inputs use groupvit model s config for some fields if specified instead of those of vision text components pooled_output returns examples python from pil import image import requests from transformers import autoprocessor groupvitmodel model groupvitmodel 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 pt padding true outputs model inputs logits_per_image outputs logits_per_image this is the image text similarity score probs logits_per_image softmax dim 1 we can take the softmax to get the label probabilities use groupvit model s config for some fields if specified instead of those of vision text components 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
import collections.abc
import math
from dataclasses import dataclass
from typing import Any, Optional, Tuple, Union

import numpy as np
import torch
import torch.utils.checkpoint
from torch import nn

from ...activations import ACT2FN
from ...modeling_attn_mask_utils import _create_4d_causal_attention_mask, _prepare_4d_attention_mask
from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling
from ...modeling_utils import PreTrainedModel
from ...utils import (
    ModelOutput,
    add_start_docstrings,
    add_start_docstrings_to_model_forward,
    logging,
    replace_return_docstrings,
)
from .configuration_groupvit import GroupViTConfig, GroupViTTextConfig, GroupViTVisionConfig


logger = logging.get_logger(__name__)

_CHECKPOINT_FOR_DOC = "nvidia/groupvit-gcc-yfcc"

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




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



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


def hard_softmax(logits: torch.Tensor, dim: int):
    y_soft = logits.softmax(dim)
    
    index = y_soft.max(dim, keepdim=True)[1]
    y_hard = torch.zeros_like(logits, memory_format=torch.legacy_contiguous_format).scatter_(dim, index, 1.0)
    ret = y_hard - y_soft.detach() + y_soft

    return ret


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

    gumbels = (logits + gumbels) / tau  
    y_soft = gumbels.softmax(dim)

    if hard:
        
        index = y_soft.max(dim, keepdim=True)[1]
        y_hard = torch.zeros_like(logits, memory_format=torch.legacy_contiguous_format).scatter_(dim, index, 1.0)
        ret = y_hard - y_soft.detach() + y_soft
    else:
        
        ret = y_soft
    return ret


def resize_attention_map(attentions, height, width, align_corners=False):
    

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

    batch_size = attentions.shape[0]
    groups = attentions.shape[1]  
    
    attentions = attentions.reshape(batch_size, groups, feat_height, feat_width)
    attentions = nn.functional.interpolate(
        attentions, size=(height, width), mode="bilinear", align_corners=align_corners
    )
    return attentions


def get_grouping_from_attentions(attentions, hw_shape):
    

    attn_maps = []
    with torch.no_grad():
        prev_attn_masks = None
        for attn_masks in attentions:
            
            attn_masks = attn_masks.permute(0, 2, 1).contiguous()
            if prev_attn_masks is None:
                prev_attn_masks = attn_masks
            else:
                prev_attn_masks = prev_attn_masks @ attn_masks
            
            cur_attn_map = resize_attention_map(prev_attn_masks.permute(0, 2, 1).contiguous(), *hw_shape)
            attn_maps.append(cur_attn_map)

    
    final_grouping = attn_maps[-1]

    return final_grouping


class GroupViTCrossAttentionLayer(nn.Module):
    def __init__(self, config: GroupViTVisionConfig):
        super().__init__()
        self.attn = GroupViTAttention(config)
        self.norm2 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.mlp = GroupViTMLP(config)
        self.norm_post = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)

    def forward(self, query, key):
        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 GroupViTAssignAttention(nn.Module):
    def __init__(self, config: GroupViTVisionConfig):
        super().__init__()
        self.scale = config.hidden_size**-0.5

        self.q_proj = nn.Linear(config.hidden_size, config.hidden_size)
        self.k_proj = nn.Linear(config.hidden_size, config.hidden_size)
        self.v_proj = nn.Linear(config.hidden_size, config.hidden_size)
        self.proj = nn.Linear(config.hidden_size, config.hidden_size)
        self.assign_eps = config.assign_eps

    def get_attn(self, attn, gumbel=True, hard=True):
        if gumbel and self.training:
            attn = gumbel_softmax(attn, dim=-2, hard=hard)
        else:
            if hard:
                attn = hard_softmax(attn, dim=-2)
            else:
                attn = nn.functional.softmax(attn, dim=-2)

        return attn

    def forward(self, query, key):
        value = key
        
        query = self.q_proj(query)

        
        key = self.k_proj(key)

        
        value = self.v_proj(value)

        
        raw_attn = (query @ key.transpose(-2, -1)) * self.scale

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

        attn = attn / (attn.sum(dim=-1, keepdim=True) + self.assign_eps)

        out = attn @ value

        out = self.proj(out)

        return out, soft_attn


class GroupViTTokenAssign(nn.Module):
    def __init__(self, config: GroupViTVisionConfig, num_group_token, num_output_group):
        super().__init__()
        self.num_output_group = num_output_group
        
        self.norm_tokens = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        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 = GroupViTMixerMLP(config, num_group_token, tokens_dim, num_output_group)
        self.norm_post_tokens = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        
        self.norm_x = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.pre_assign_attn = GroupViTCrossAttentionLayer(config)

        self.assign = GroupViTAssignAttention(config)
        self.norm_new_x = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.mlp_channels = GroupViTMLP(config, config.hidden_size, channels_dim, config.hidden_size)

    def project_group_token(self, group_tokens):
        
        
        projected_group_tokens = self.mlp_inter(group_tokens)
        projected_group_tokens = self.norm_post_tokens(projected_group_tokens)
        return projected_group_tokens

    def forward(self, image_tokens, group_tokens):
        

        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


@dataclass
class GroupViTModelOutput(ModelOutput):
    

    loss: Optional[torch.FloatTensor] = None
    logits_per_image: torch.FloatTensor = None
    logits_per_text: torch.FloatTensor = None
    segmentation_logits: torch.FloatTensor = None
    text_embeds: torch.FloatTensor = None
    image_embeds: torch.FloatTensor = None
    text_model_output: BaseModelOutputWithPooling = None
    vision_model_output: BaseModelOutputWithPooling = 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 GroupViTPatchEmbeddings(nn.Module):
    

    def __init__(
        self,
        image_size: int = 224,
        patch_size: Union[int, Tuple[int, int]] = 16,
        num_channels: int = 3,
        embed_dim: int = 768,
    ):
        super().__init__()
        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.projection = nn.Conv2d(num_channels, embed_dim, kernel_size=patch_size, stride=patch_size)

    def forward(self, pixel_values: torch.Tensor, interpolate_pos_encoding: bool = False) -> torch.Tensor:
        batch_size, num_channels, height, width = pixel_values.shape
        if not interpolate_pos_encoding:
            if height != self.image_size[0] or width != self.image_size[1]:
                raise ValueError(
                    f"Input image size ({height}*{width}) doesn't match model"
                    f" ({self.image_size[0]}*{self.image_size[1]})."
                )
        x = self.projection(pixel_values).flatten(2).transpose(1, 2)
        return x


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

        self.patch_embeddings = GroupViTPatchEmbeddings(
            image_size=config.image_size,
            patch_size=config.patch_size,
            num_channels=config.num_channels,
            embed_dim=config.hidden_size,
        )
        num_patches = self.patch_embeddings.num_patches
        self.position_embeddings = nn.Parameter(torch.zeros(1, num_patches, config.hidden_size))
        self.dropout = nn.Dropout(config.dropout)
        self.layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.config = config

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

        npatch = embeddings.shape[1]
        if npatch == self.position_embeddings.shape[1] and height == width:
            return self.position_embeddings
        patch_pos_embed = self.position_embeddings
        num_original_pos_embed = patch_pos_embed.shape[1]
        dim = embeddings.shape[-1]
        feat_height = height // self.config.patch_size
        feat_width = width // self.config.patch_size
        
        
        feat_height, feat_width = feat_height + 0.1, feat_width + 0.1
        original_height = original_width = math.sqrt(num_original_pos_embed)
        reshaped_patch_pos_embed = patch_pos_embed.reshape(1, int(original_height), int(original_width), dim).permute(
            0, 3, 1, 2
        )
        scale_factor = (feat_height / original_height, feat_width / original_width)
        patch_pos_embed = nn.functional.interpolate(
            reshaped_patch_pos_embed,
            scale_factor=scale_factor,
            mode="bicubic",
            align_corners=False,
        )
        patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
        return patch_pos_embed

    def forward(self, pixel_values: torch.Tensor, interpolate_pos_encoding: bool = False) -> torch.Tensor:
        batch_size, num_channels, height, width = pixel_values.shape
        embeddings = self.patch_embeddings(pixel_values, interpolate_pos_encoding=interpolate_pos_encoding)

        embeddings = self.layernorm(embeddings)

        batch_size, seq_len, _ = embeddings.size()

        
        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 GroupViTTextEmbeddings(nn.Module):
    def __init__(self, config: GroupViTTextConfig):
        super().__init__()
        embed_dim = config.hidden_size

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

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

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

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

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

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

        return embeddings


class GroupViTStage(nn.Module):
    

    def __init__(
        self,
        config: GroupViTVisionConfig,
        depth: int,
        num_prev_group_token: int,
        num_group_token: int,
        num_output_group: int,
    ):
        super().__init__()
        self.depth = depth
        self.num_group_token = num_group_token
        if num_group_token > 0:
            self.group_token = nn.Parameter(torch.zeros(1, num_group_token, config.hidden_size))
        else:
            self.group_token = None
        self.layers = nn.ModuleList([GroupViTEncoderLayer(config) for _ in range(depth)])

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

        if num_prev_group_token > 0 and num_group_token > 0:
            self.group_projector = nn.Sequential(
                nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps),
                GroupViTMixerMLP(config, num_prev_group_token, config.hidden_size // 2, num_group_token),
            )
        else:
            self.group_projector = None

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

    def split_x(self, x):
        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: torch.Tensor, group_token: Optional[torch.Tensor] = None) -> torch.Tensor:
        if group_token is None:
            return x
        return torch.cat([x, group_token], dim=1)

    def forward(
        self,
        hidden_states: torch.Tensor,
        prev_group_token: Optional[torch.Tensor] = None,
        output_attentions: Optional[bool] = False,
    ) -> Tuple[torch.FloatTensor]:
        
        if self.with_group_token:
            group_token = self.group_token.expand(hidden_states.size(0), -1, -1)
            if self.group_projector is not None:
                group_token = group_token + self.group_projector(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)
            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 GroupViTMLP(nn.Module):
    def __init__(
        self,
        config: GroupViTVisionConfig,
        hidden_size: Optional[int] = None,
        intermediate_size: Optional[int] = None,
        output_size: Optional[int] = None,
    ):
        super().__init__()
        self.config = config
        self.activation_fn = ACT2FN[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 = nn.Linear(hidden_size, intermediate_size)
        self.fc2 = nn.Linear(intermediate_size, output_size)

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


class GroupViTMixerMLP(GroupViTMLP):
    def forward(self, x):
        x = super().forward(x.transpose(1, 2))
        return x.transpose(1, 2)


class GroupViTAttention(nn.Module):
    

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

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

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

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

        bsz, tgt_len, embed_dim = hidden_states.size()
        is_cross_attention = encoder_hidden_states is not None

        
        query_states = self.q_proj(hidden_states) * self.scale
        if is_cross_attention:
            key_states = self._shape(self.k_proj(encoder_hidden_states), -1, bsz)
            value_states = self._shape(self.v_proj(encoder_hidden_states), -1, bsz)
        else:
            key_states = self._shape(self.k_proj(hidden_states), -1, bsz)
            value_states = self._shape(self.v_proj(hidden_states), -1, bsz)

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

        src_len = key_states.size(1)
        attn_weights = torch.bmm(query_states, key_states.transpose(1, 2))

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

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

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

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

        if output_attentions:
            
            
            
            
            attn_weights_reshaped = attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
            attn_weights = attn_weights_reshaped.view(bsz * self.num_heads, tgt_len, src_len)
        else:
            attn_weights_reshaped = None

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

        attn_output = torch.bmm(attn_probs, value_states)

        if attn_output.size() != (bsz * self.num_heads, tgt_len, self.head_dim):
            raise ValueError(
                f"`attn_output` should be of size {(bsz, self.num_heads, tgt_len, self.head_dim)}, but is"
                f" {attn_output.size()}"
            )

        attn_output = attn_output.view(bsz, self.num_heads, tgt_len, self.head_dim)
        attn_output = attn_output.transpose(1, 2)
        attn_output = attn_output.reshape(bsz, tgt_len, embed_dim)

        attn_output = self.out_proj(attn_output)

        return attn_output, attn_weights_reshaped



class GroupViTEncoderLayer(nn.Module):
    def __init__(self, config: GroupViTConfig):
        super().__init__()
        self.embed_dim = config.hidden_size
        self.self_attn = GroupViTAttention(config)
        self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
        self.mlp = GroupViTMLP(config)
        self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)

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

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

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

        outputs = (hidden_states,)

        if output_attentions:
            outputs += (attn_weights,)

        return outputs


class GroupViTPreTrainedModel(PreTrainedModel):
    

    config_class = GroupViTConfig
    base_model_prefix = "groupvit"
    supports_gradient_checkpointing = True

    def _init_weights(self, module):
        

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

        factor = self.config.initializer_factor
        if isinstance(module, GroupViTTextEmbeddings):
            module.token_embedding.weight.data.normal_(mean=0.0, std=factor * 0.02)
            module.position_embedding.weight.data.normal_(mean=0.0, std=factor * 0.02)
        elif isinstance(module, GroupViTAttention):
            factor = self.config.initializer_factor
            in_proj_std = (module.embed_dim**-0.5) * ((2 * module.config.num_hidden_layers) ** -0.5) * factor
            out_proj_std = (module.embed_dim**-0.5) * factor
            nn.init.normal_(module.q_proj.weight, std=in_proj_std)
            nn.init.normal_(module.k_proj.weight, std=in_proj_std)
            nn.init.normal_(module.v_proj.weight, std=in_proj_std)
            nn.init.normal_(module.out_proj.weight, std=out_proj_std)
        elif isinstance(module, GroupViTMLP):
            factor = self.config.initializer_factor
            in_proj_std = (module.config.hidden_size**-0.5) * ((2 * module.config.num_hidden_layers) ** -0.5) * factor
            fc_std = (2 * module.config.hidden_size) ** -0.5 * factor
            nn.init.normal_(module.fc1.weight, std=fc_std)
            nn.init.normal_(module.fc2.weight, std=in_proj_std)


GROUPVIT_START_DOCSTRING = r

GROUPVIT_TEXT_INPUTS_DOCSTRING = r

GROUPVIT_VISION_INPUTS_DOCSTRING = r

GROUPVIT_INPUTS_DOCSTRING = r


class GroupViTVisionEncoder(nn.Module):
    def __init__(self, config: GroupViTVisionConfig) -> None:
        super().__init__()
        self.config = config
        self.stages = nn.ModuleList(
            [
                GroupViTStage(
                    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,
                )
                for i in range(len(config.depths))
            ]
        )
        self.gradient_checkpointing = False

    def forward(
        self,
        hidden_states: torch.Tensor,
        output_hidden_states: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[tuple, BaseModelOutput]:
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        all_hidden_states = () if output_hidden_states else None
        all_groupings = () if output_attentions else None

        group_tokens = None

        for i, stage in enumerate(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 BaseModelOutput(
            last_hidden_state=hidden_states, hidden_states=all_hidden_states, attentions=all_groupings
        )


class GroupViTTextEncoder(nn.Module):
    

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

    def forward(
        self,
        inputs_embeds,
        attention_mask: Optional[torch.Tensor] = None,
        causal_attention_mask: Optional[torch.Tensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple, BaseModelOutput]:
        r
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        encoder_states = () if output_hidden_states else None
        all_attentions = () if output_attentions else None

        hidden_states = inputs_embeds
        for idx, encoder_layer in enumerate(self.layers):
            if output_hidden_states:
                encoder_states = encoder_states + (hidden_states,)
            if self.gradient_checkpointing and self.training:
                layer_outputs = self._gradient_checkpointing_func(
                    encoder_layer.__call__,
                    hidden_states,
                    attention_mask,
                    causal_attention_mask,
                    output_attentions,
                )
            else:
                layer_outputs = encoder_layer(
                    hidden_states,
                    attention_mask,
                    causal_attention_mask,
                    output_attentions=output_attentions,
                )

            hidden_states = layer_outputs[0]

            if output_attentions:
                all_attentions = all_attentions + (layer_outputs[1],)

        if output_hidden_states:
            encoder_states = encoder_states + (hidden_states,)

        if not return_dict:
            return tuple(v for v in [hidden_states, encoder_states, all_attentions] if v is not None)
        return BaseModelOutput(
            last_hidden_state=hidden_states, hidden_states=encoder_states, attentions=all_attentions
        )



class GroupViTTextTransformer(nn.Module):
    def __init__(self, config: GroupViTTextConfig):
        super().__init__()
        self.config = config
        embed_dim = config.hidden_size
        self.embeddings = GroupViTTextEmbeddings(config)
        self.encoder = GroupViTTextEncoder(config)
        self.final_layer_norm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)

        
        self.eos_token_id = config.eos_token_id

    @add_start_docstrings_to_model_forward(GROUPVIT_TEXT_INPUTS_DOCSTRING)
    @replace_return_docstrings(output_type=BaseModelOutputWithPooling, config_class=GroupViTTextConfig)
    def forward(
        self,
        input_ids: Optional[torch.Tensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.Tensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple, BaseModelOutputWithPooling]:
        r
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        if input_ids is None:
            raise ValueError("You have to specify input_ids")

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

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

        
        
        causal_attention_mask = _create_4d_causal_attention_mask(
            input_shape, hidden_states.dtype, device=hidden_states.device
        )
        
        if attention_mask is not None:
            
            attention_mask = _prepare_4d_attention_mask(attention_mask, hidden_states.dtype)

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

        last_hidden_state = encoder_outputs[0]
        last_hidden_state = self.final_layer_norm(last_hidden_state)

        if self.eos_token_id == 2:
            
            
            
            
            
            
            pooled_output = last_hidden_state[
                torch.arange(last_hidden_state.shape[0], device=last_hidden_state.device),
                input_ids.to(dtype=torch.int, device=last_hidden_state.device).argmax(dim=-1),
            ]
        else:
            
            pooled_output = last_hidden_state[
                torch.arange(last_hidden_state.shape[0], device=last_hidden_state.device),
                
                (input_ids.to(dtype=torch.int, device=last_hidden_state.device) == self.eos_token_id)
                .int()
                .argmax(dim=-1),
            ]

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

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


class GroupViTTextModel(GroupViTPreTrainedModel):
    config_class = GroupViTTextConfig

    def __init__(self, config: GroupViTTextConfig):
        super().__init__(config)
        self.text_model = GroupViTTextTransformer(config)
        
        self.post_init()

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

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

    @add_start_docstrings_to_model_forward(GROUPVIT_TEXT_INPUTS_DOCSTRING)
    @replace_return_docstrings(output_type=BaseModelOutputWithPooling, config_class=GroupViTTextConfig)
    def forward(
        self,
        input_ids: Optional[torch.Tensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.Tensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple, BaseModelOutputWithPooling]:
        r
        return self.text_model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            position_ids=position_ids,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )


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

        self.embeddings = GroupViTVisionEmbeddings(config)
        self.encoder = GroupViTVisionEncoder(config)
        self.layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)

    @add_start_docstrings_to_model_forward(GROUPVIT_VISION_INPUTS_DOCSTRING)
    @replace_return_docstrings(output_type=BaseModelOutputWithPooling, config_class=GroupViTVisionConfig)
    def forward(
        self,
        pixel_values: Optional[torch.FloatTensor] = None,
        output_hidden_states: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple, BaseModelOutputWithPooling]:
        r
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        if pixel_values is None:
            raise ValueError("You have to specify pixel_values")

        hidden_states = self.embeddings(pixel_values)

        encoder_outputs = self.encoder(
            hidden_states=hidden_states,
            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 = last_hidden_state.mean(dim=1)

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

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


class GroupViTVisionModel(GroupViTPreTrainedModel):
    config_class = GroupViTVisionConfig
    main_input_name = "pixel_values"

    def __init__(self, config: GroupViTVisionConfig):
        super().__init__(config)
        self.vision_model = GroupViTVisionTransformer(config)
        
        self.post_init()

    def get_input_embeddings(self) -> GroupViTPatchEmbeddings:
        return self.vision_model.embeddings.patch_embeddings

    @add_start_docstrings_to_model_forward(GROUPVIT_VISION_INPUTS_DOCSTRING)
    @replace_return_docstrings(output_type=BaseModelOutputWithPooling, config_class=GroupViTVisionConfig)
    def forward(
        self,
        pixel_values: Optional[torch.FloatTensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple, BaseModelOutputWithPooling]:
        r
        return self.vision_model(
            pixel_values=pixel_values,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )


@add_start_docstrings(GROUPVIT_START_DOCSTRING)
class GroupViTModel(GroupViTPreTrainedModel):
    config_class = GroupViTConfig

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

        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)}."
            )

        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 = GroupViTTextTransformer(text_config)
        self.vision_model = GroupViTVisionTransformer(vision_config)

        self.visual_projection = nn.Sequential(
            nn.Linear(self.vision_embed_dim, self.projection_intermediate_dim, bias=True),
            nn.BatchNorm1d(self.projection_intermediate_dim),
            nn.ReLU(inplace=True),
            nn.Linear(self.projection_intermediate_dim, self.projection_dim, bias=True),
        )
        self.text_projection = nn.Sequential(
            nn.Linear(self.text_embed_dim, self.projection_intermediate_dim, bias=True),
            nn.BatchNorm1d(self.projection_intermediate_dim),
            nn.ReLU(inplace=True),
            nn.Linear(self.projection_intermediate_dim, self.projection_dim, bias=True),
        )
        self.logit_scale = nn.Parameter(torch.tensor(self.config.logit_scale_init_value))

        
        self.post_init()

    @add_start_docstrings_to_model_forward(GROUPVIT_TEXT_INPUTS_DOCSTRING)
    def get_text_features(
        self,
        input_ids: Optional[torch.Tensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.Tensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> torch.FloatTensor:
        r
        
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        text_outputs = self.text_model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            position_ids=position_ids,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

        pooled_output = text_outputs[1]
        text_features = self.text_projection(pooled_output)

        return text_features

    @add_start_docstrings_to_model_forward(GROUPVIT_VISION_INPUTS_DOCSTRING)
    def get_image_features(
        self,
        pixel_values: Optional[torch.FloatTensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> torch.FloatTensor:
        r
        
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

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

        pooled_output = vision_outputs[1]  
        image_features = self.visual_projection(pooled_output)

        return image_features

    @add_start_docstrings_to_model_forward(GROUPVIT_INPUTS_DOCSTRING)
    @replace_return_docstrings(output_type=GroupViTModelOutput, config_class=GroupViTConfig)
    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        pixel_values: Optional[torch.FloatTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        return_loss: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        output_segmentation: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple, GroupViTModelOutput]:
        r
        
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_segmentation = (
            output_segmentation if output_segmentation is not None else self.config.output_segmentation
        )
        if output_segmentation:
            output_attentions = True
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

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

        text_outputs = self.text_model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            position_ids=position_ids,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

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

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

        
        image_embeds = image_embeds / image_embeds.norm(dim=-1, keepdim=True)
        text_embeds = text_embeds / text_embeds.norm(dim=-1, keepdim=True)

        
        logit_scale = self.logit_scale.exp()
        logits_per_text = torch.matmul(text_embeds, image_embeds.t()) * logit_scale
        logits_per_image = logits_per_text.t()

        seg_logits = None
        if output_segmentation:
            
            
            image_group_embeds = vision_outputs[0]
            
            image_group_embeds = self.visual_projection(image_group_embeds.reshape(-1, image_group_embeds.shape[-1]))
            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 / image_group_embeds.norm(dim=-1, keepdim=True)
            
            logits_per_image_group = torch.matmul(image_group_embeds, text_embeds.t()) * logit_scale
            
            logits_per_image_group = logits_per_image_group.reshape(
                image_embeds.shape[0], -1, text_embeds.shape[0]
            ).permute(0, 2, 1)

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

            
            seg_logits = torch.matmul(logits_per_image_group, flatten_grouping) * logit_scale
            seg_logits = seg_logits.reshape(
                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)

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