codingutf8 2022 google ai 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 owlvit model import warnings from dataclasses import dataclass from typing import any dict optional tuple union import numpy as np import torch import torch utils checkpoint from torch import tensor 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 isvisionavailable logging replacereturndocstrings from configurationowlvit import owlvitconfig owlvittextconfig owlvitvisionconfig if isvisionavailable from transformers imagetransforms import centertocornersformat logger logging getloggername checkpointfordoc googleowlvitbasepatch32 see all owlvit models at https huggingface comodels filterowlvit owlvitpretrainedmodelarchivelist googleowlvitbasepatch32 googleowlvitbasepatch16 googleowlvitlargepatch14 copied from transformers models clip modelingclip contrastiveloss with clipowlvit def contrastivelosslogits torch tensor torch tensor return nn functional crossentropylogits torch arangelenlogits devicelogits device copied from transformers models clip modelingclip cliploss with clipowlvit def owlvitlosssimilarity torch tensor torch tensor captionloss contrastivelosssimilarity imageloss contrastivelosssimilarity t return captionloss imageloss 2 0 dataclass class owlvitoutputmodeloutput loss optionaltorch floattensor none logitsperimage torch floattensor none logitspertext 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 copied from transformers models detr modelingdetr upcast def upcastt tensor tensor protects from numerical overflows in multiplications by upcasting to the equivalent higher type if t isfloatingpoint return t if t dtype in torch float32 torch float64 else t float else return t if t dtype in torch int32 torch int64 else t int copied from transformers models detr modelingdetr boxarea def boxareaboxes tensor tensor boxes upcastboxes return boxes 2 boxes 0 boxes 3 boxes 1 copied from transformers models detr modelingdetr boxiou def boxiouboxes1 boxes2 area1 boxareaboxes1 area2 boxareaboxes2 lefttop torch maxboxes1 none 2 boxes2 2 n m 2 rightbottom torch minboxes1 none 2 boxes2 2 n m 2 widthheight rightbottom lefttop clampmin0 n m 2 inter widthheight 0 widthheight 1 n m union area1 none area2 inter iou inter union return iou union copied from transformers models detr modelingdetr generalizedboxiou def generalizedboxiouboxes1 boxes2 degenerate boxes gives inf nan results so do an early check if not boxes1 2 boxes1 2 all raise valueerrorfboxes1 must be in x0 y0 x1 y1 corner format but got boxes1 if not boxes2 2 boxes2 2 all raise valueerrorfboxes2 must be in x0 y0 x1 y1 corner format but got boxes2 iou union boxiouboxes1 boxes2 topleft torch minboxes1 none 2 boxes2 2 bottomright torch maxboxes1 none 2 boxes2 2 widthheight bottomright topleft clampmin0 n m 2 area widthheight 0 widthheight 1 return iou area union area dataclass class owlvitobjectdetectionoutputmodeloutput loss optionaltorch floattensor none lossdict optionaldict none logits torch floattensor none predboxes torch floattensor none textembeds torch floattensor none imageembeds torch floattensor none classembeds 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 dataclass class owlvitimageguidedobjectdetectionoutputmodeloutput logits torch floattensor none imageembeds torch floattensor none queryimageembeds torch floattensor none targetpredboxes torch floattensor none querypredboxes torch floattensor none classembeds 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 owlvitvisionembeddingsnn module def initself config owlvitvisionconfig super init self config config self embeddim config hiddensize self classembedding nn parametertorch randnconfig hiddensize self patchembedding nn conv2d inchannelsconfig numchannels outchannelsself embeddim kernelsizeconfig patchsize strideconfig patchsize biasfalse self numpatches config imagesize config patchsize 2 self numpositions self numpatches 1 self positionembedding nn embeddingself numpositions self embeddim self registerbufferpositionids torch arangeself numpositions expand1 1 persistentfalse def forwardself pixelvalues torch floattensor torch tensor batchsize pixelvalues shape0 patchembeds self patchembeddingpixelvalues shape batchsize numchannels height width patchembeds patchembeds flatten2 transpose1 2 classembeds self classembedding expandbatchsize 1 1 embeddings torch catclassembeds patchembeds dim1 embeddings embeddings self positionembeddingself positionids return embeddings class owlvittextembeddingsnn module def initself config owlvittextconfig super init self tokenembedding nn embeddingconfig vocabsize config hiddensize self positionembedding nn embeddingconfig maxpositionembeddings config hiddensize 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 owlvitattentionnn module input shape batch x time x channel bsz tgtlen embeddim hiddenstates size get query proj querystates self qprojhiddenstates self scale keystates self shapeself kprojhiddenstates 1 bsz valuestates self shapeself vprojhiddenstates 1 bsz projshape bsz self numheads 1 self headdim querystates self shapequerystates tgtlen bsz viewprojshape keystates keystates viewprojshape valuestates valuestates viewprojshape srclen keystates size1 attnweights torch bmmquerystates keystates transpose1 2 if attnweights size bsz self numheads tgtlen srclen raise valueerror fattention weights should be of size bsz self numheads tgtlen srclen but is f attnweights size apply the causalattentionmask first if causalattentionmask is not none if causalattentionmask size bsz 1 tgtlen srclen raise valueerror fattention mask should be of size bsz 1 tgtlen srclen but is f causalattentionmask size attnweights attnweights viewbsz self numheads tgtlen srclen causalattentionmask attnweights attnweights viewbsz self numheads tgtlen srclen if attentionmask is not none if attentionmask size bsz 1 tgtlen srclen raise valueerror fattention mask should be of size bsz 1 tgtlen srclen but is attentionmask size attnweights attnweights viewbsz self numheads tgtlen srclen attentionmask attnweights attnweights viewbsz self numheads tgtlen srclen attnweights nn functional softmaxattnweights dim1 if outputattentions this operation is a bit akward but it s required to make sure that attnweights keeps its gradient in order to do so attnweights have to reshaped twice and have to be reused in the following attnweightsreshaped attnweights viewbsz self numheads tgtlen srclen attnweights attnweightsreshaped viewbsz self numheads tgtlen srclen else attnweightsreshaped none attnprobs nn functional dropoutattnweights pself dropout trainingself training for int8 compatibility sometimes the attnprobs are in fp32 attnprobs attnprobs tovaluestates dtype attnoutput torch bmmattnprobs valuestates if attnoutput size bsz self numheads tgtlen self headdim raise valueerror fattnoutput should be of size bsz self numheads tgtlen self headdim but is f attnoutput size attnoutput attnoutput viewbsz self numheads tgtlen self headdim attnoutput attnoutput transpose1 2 attnoutput attnoutput reshapebsz tgtlen embeddim attnoutput self outprojattnoutput return attnoutput attnweightsreshaped copied from transformers models clip modelingclip clipmlp with clipowlvit class owlvitmlpnn module def initself config super init self config config self activationfn act2fnconfig hiddenact self fc1 nn linearconfig hiddensize config intermediatesize self fc2 nn linearconfig intermediatesize config hiddensize def forwardself hiddenstates torch tensor torch tensor hiddenstates self fc1hiddenstates hiddenstates self activationfnhiddenstates hiddenstates self fc2hiddenstates return hiddenstates copied from transformers models clip modelingclip clipencoderlayer with clipowlvit class owlvitencoderlayernn module def initself config owlvitconfig super init self embeddim config hiddensize self selfattn owlvitattentionconfig self layernorm1 nn layernormself embeddim epsconfig layernormeps self mlp owlvitmlpconfig self layernorm2 nn layernormself embeddim epsconfig layernormeps def forward self hiddenstates torch tensor attentionmask torch tensor causalattentionmask torch tensor outputattentions optionalbool false tupletorch floattensor residual hiddenstates hiddenstates self layernorm1hiddenstates hiddenstates attnweights self selfattn hiddenstateshiddenstates attentionmaskattentionmask causalattentionmaskcausalattentionmask outputattentionsoutputattentions hiddenstates residual hiddenstates residual hiddenstates hiddenstates self layernorm2hiddenstates hiddenstates self mlphiddenstates hiddenstates residual hiddenstates outputs hiddenstates if outputattentions outputs attnweights return outputs class owlvitpretrainedmodelpretrainedmodel configclass owlvitconfig basemodelprefix owlvit supportsgradientcheckpointing true nosplitmodules owlvitencoderlayer def initweightsself module owlvittextinputsdocstring r args inputids torch longtensor of shape batchsize nummaxtextqueries sequencelength indices of input sequence tokens in the vocabulary indices can be obtained using autotokenizer see pretrainedtokenizer encode and pretrainedtokenizer call for details what are input ids glossaryinputids attentionmask torch tensor of shape batchsize nummaxtextqueries 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 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 owlvitinputsdocstring r args inputids torch longtensor of shape batchsize sequencelength indices of input sequence tokens in the vocabulary indices can be obtained using autotokenizer see pretrainedtokenizer encode and pretrainedtokenizer call for details what are input ids glossaryinputids 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 pixelvalues torch floattensor of shape batchsize numchannels height width pixel values 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 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 owlvitimageguidedobjectdetectioninputsdocstring r args pixelvalues torch floattensor of shape batchsize numchannels height width pixel values querypixelvalues torch floattensor of shape batchsize numchannels height width pixel values of query images to be detected pass in one query image per target image 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 transformer encoder consisting of config numhiddenlayers self attention layers each layer is a owlvitencoderlayer args config owlvitconfig 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 encoderstates if outputhiddenstates else none allattentions if outputattentions else none hiddenstates inputsembeds for encoderlayer in self layers if outputhiddenstates encoderstates encoderstates hiddenstates if self gradientcheckpointing and self training layeroutputs self gradientcheckpointingfunc encoderlayer call hiddenstates attentionmask causalattentionmask outputattentions else layeroutputs encoderlayer hiddenstates attentionmask causalattentionmask outputattentionsoutputattentions hiddenstates layeroutputs0 if outputattentions allattentions allattentions layeroutputs1 if outputhiddenstates encoderstates encoderstates hiddenstates if not returndict return tuplev for v in hiddenstates encoderstates allattentions if v is not none return basemodeloutput lasthiddenstatehiddenstates hiddenstatesencoderstates attentionsallattentions class owlvittexttransformernn module def initself config owlvittextconfig super init self config config embeddim config hiddensize self embeddings owlvittextembeddingsconfig self encoder owlvitencoderconfig self finallayernorm nn layernormembeddim epsconfig layernormeps addstartdocstringstomodelforwardowlvittextinputsdocstring replacereturndocstringsoutputtypebasemodeloutputwithpooling configclassowlvittextconfig def forward self inputids torch tensor attentionmask optionaltorch tensor none positionids optionaltorch tensor none outputattentions optionalbool none outputhiddenstates optionalbool none returndict optionalbool none uniontuple basemodeloutputwithpooling r returns numsamples seqlen inputshape where numsamples batchsize nummaxtextqueries owlvit s text model uses causal mask prepare it here https github comopenaiclipblobcfcffb90e69f37bf2ff1e988237a0fbe41f33c04clipmodel pyl324 expand attentionmask numsamples seqlen numsamples 1 tgtseqlen srcseqlen take features from the end of tokens embedding end of token is the highest number in each sequence casting to torch int for onnx compatibility argmax doesn t support int64 inputs with opset 14 initialize weights and apply final processing get embeddings for all text queries in all batch samples 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 cast the input to the expected dtype expectedinputdtype self embeddings patchembedding weight dtype pixelvalues pixelvalues toexpectedinputdtype hiddenstates self embeddingspixelvalues hiddenstates self prelayernormhiddenstates encoderoutputs self encoder inputsembedshiddenstates outputattentionsoutputattentions outputhiddenstatesoutputhiddenstates returndictreturndict lasthiddenstate encoderoutputs0 pooledoutput lasthiddenstate 0 pooledoutput self postlayernormpooledoutput if not returndict return lasthiddenstate pooledoutput encoderoutputs1 return basemodeloutputwithpooling lasthiddenstatelasthiddenstate pooleroutputpooledoutput hiddenstatesencoderoutputs hiddenstates attentionsencoderoutputs attentions class owlvitvisionmodelowlvitpretrainedmodel configclass owlvitvisionconfig maininputname pixelvalues def initself config owlvitvisionconfig super initconfig self visionmodel owlvitvisiontransformerconfig initialize weights and apply final processing self postinit def getinputembeddingsself nn module return self visionmodel embeddings patchembedding addstartdocstringstomodelforwardowlvitvisioninputsdocstring replacereturndocstringsoutputtypebasemodeloutputwithpooling configclassowlvitvisionconfig def forward self pixelvalues optionaltorch floattensor none outputattentions optionalbool none outputhiddenstates optionalbool none returndict optionalbool none uniontuple basemodeloutputwithpooling r returns examples python from pil import image import requests from transformers import autoprocessor owlvitvisionmodel model owlvitvisionmodel frompretrainedgoogleowlvitbasepatch32 processor autoprocessor frompretrainedgoogleowlvitbasepatch32 url http images cocodataset orgval2017000000039769 jpg image image openrequests geturl streamtrue raw inputs processorimagesimage returntensorspt outputs modelinputs lasthiddenstate outputs lasthiddenstate pooledoutput outputs pooleroutput pooled cls states return self visionmodel pixelvaluespixelvalues outputattentionsoutputattentions outputhiddenstatesoutputhiddenstates returndictreturndict addstartdocstringsowlvitstartdocstring class owlvitmodelowlvitpretrainedmodel configclass owlvitconfig def initself config owlvitconfig super initconfig if not isinstanceconfig textconfig owlvittextconfig raise valueerror config textconfig is expected to be of type owlvittextconfig but is of type f typeconfig textconfig if not isinstanceconfig visionconfig owlvitvisionconfig raise valueerror config visionconfig is expected to be of type owlvitvisionconfig but is of type f typeconfig visionconfig textconfig config textconfig visionconfig config visionconfig self projectiondim config projectiondim self textembeddim textconfig hiddensize self visionembeddim visionconfig hiddensize self textmodel owlvittexttransformertextconfig self visionmodel owlvitvisiontransformervisionconfig self visualprojection nn linearself visionembeddim self projectiondim biasfalse self textprojection nn linearself textembeddim self projectiondim biasfalse self logitscale nn parametertorch tensorconfig logitscaleinitvalue initialize weights and apply final processing self postinit addstartdocstringstomodelforwardowlvittextinputsdocstring def gettextfeatures self inputids optionaltorch tensor none attentionmask optionaltorch tensor none outputattentions optionalbool none outputhiddenstates optionalbool none returndict optionalbool none torch floattensor r returns textfeatures torch floattensor of shape batchsize outputdim the text embeddings obtained by applying the projection layer to the pooled output of owlvittextmodel examples python from transformers import autoprocessor owlvitmodel model owlvitmodel frompretrainedgoogleowlvitbasepatch32 processor autoprocessor frompretrainedgoogleowlvitbasepatch32 inputs processor texta photo of a cat a photo of a dog photo of a astranaut returntensorspt textfeatures model gettextfeaturesinputs use owlvit model s config for some fields if specified instead of those of vision text components returndict returndict if returndict is not none else self config usereturndict get embeddings for all text queries in all batch samples textoutput self textmodelinputidsinputids attentionmaskattentionmask returndictreturndict pooledoutput textoutput1 textfeatures self textprojectionpooledoutput return textfeatures addstartdocstringstomodelforwardowlvitvisioninputsdocstring def getimagefeatures self pixelvalues optionaltorch floattensor none outputattentions optionalbool none outputhiddenstates optionalbool none returndict optionalbool none torch floattensor r returns imagefeatures torch floattensor of shape batchsize outputdim the image embeddings obtained by applying the projection layer to the pooled output of owlvitvisionmodel examples python from pil import image import requests from transformers import autoprocessor owlvitmodel model owlvitmodel frompretrainedgoogleowlvitbasepatch32 processor autoprocessor frompretrainedgoogleowlvitbasepatch32 url http images cocodataset orgval2017000000039769 jpg image image openrequests geturl streamtrue raw inputs processorimagesimage returntensorspt imagefeatures model getimagefeaturesinputs use owlvit model s config for some fields if specified instead of those of vision text components 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 visionoutputs self visionmodel pixelvaluespixelvalues outputattentionsoutputattentions outputhiddenstatesoutputhiddenstates returndictreturndict pooledoutput visionoutputs1 imagefeatures self visualprojectionpooledoutput return imagefeatures addstartdocstringstomodelforwardowlvitinputsdocstring replacereturndocstringsoutputtypeowlvitoutput configclassowlvitconfig def forward self inputids optionaltorch longtensor none pixelvalues optionaltorch floattensor none attentionmask optionaltorch tensor none returnloss optionalbool none outputattentions optionalbool none outputhiddenstates optionalbool none returnbaseimageembeds optionalbool none returndict optionalbool none uniontuple owlvitoutput r returns examples python from pil import image import requests from transformers import autoprocessor owlvitmodel model owlvitmodel frompretrainedgoogleowlvitbasepatch32 processor autoprocessor frompretrainedgoogleowlvitbasepatch32 url http images cocodataset orgval2017000000039769 jpg image image openrequests geturl streamtrue raw inputs processortexta photo of a cat a photo of a dog imagesimage returntensorspt outputs modelinputs logitsperimage outputs logitsperimage this is the imagetext similarity score probs logitsperimage softmaxdim1 we can take the softmax to get the label probabilities use owlvit model s config for some fields if specified instead of those of vision text components 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 visionoutputs self visionmodel pixelvaluespixelvalues outputattentionsoutputattentions outputhiddenstatesoutputhiddenstates returndictreturndict get embeddings for all text queries in all batch samples textoutputs self textmodel inputidsinputids attentionmaskattentionmask outputattentionsoutputattentions outputhiddenstatesoutputhiddenstates returndictreturndict textembeds textoutputs1 textembeds self textprojectiontextembeds imageembeds visionoutputs1 imageembeds self visualprojectionimageembeds normalized features imageembeds imageembeds torch linalg normimageembeds ord2 dim1 keepdimtrue textembedsnorm textembeds torch linalg normtextembeds ord2 dim1 keepdimtrue cosine similarity as logits and set it on the correct device logitscale self logitscale exp toimageembeds device logitspertext torch matmultextembedsnorm imageembeds t logitscale logitsperimage logitspertext t loss none if returnloss loss owlvitlosslogitspertext if returnbaseimageembeds warnings warn returnbaseimageembeds is deprecated and will be removed in v4 27 of transformers one can obtain the base unprojected image embeddings from outputs visionmodeloutput futurewarning lasthiddenstate visionoutputs0 imageembeds self visionmodel postlayernormlasthiddenstate else textembeds textembedsnorm if not returndict output logitsperimage logitspertext textembeds imageembeds textoutputs visionoutputs return loss output if loss is not none else output return owlvitoutput lossloss logitsperimagelogitsperimage logitspertextlogitspertext textembedstextembeds imageembedsimageembeds textmodeloutputtextoutputs visionmodeloutputvisionoutputs class owlvitboxpredictionheadnn module def initself config owlvitconfig outdim int 4 super init width config visionconfig hiddensize self dense0 nn linearwidth width self dense1 nn linearwidth width self gelu nn gelu self dense2 nn linearwidth outdim def forwardself imagefeatures torch tensor torch floattensor output self dense0imagefeatures output self geluoutput output self dense1output output self geluoutput output self dense2output return output class owlvitclasspredictionheadnn module def initself config owlvitconfig super init outdim config textconfig hiddensize self querydim config visionconfig hiddensize self dense0 nn linearself querydim outdim self logitshift nn linearself querydim 1 self logitscale nn linearself querydim 1 self elu nn elu def forward self imageembeds torch floattensor queryembeds optionaltorch floattensor querymask optionaltorch tensor tupletorch floattensor imageclassembeds self dense0imageembeds if queryembeds is none device imageclassembeds device batchsize numpatches imageclassembeds shape 2 predlogits torch zerosbatchsize numpatches self querydim todevice return predlogits imageclassembeds normalize image and text features imageclassembeds imageclassembeds torch linalg normimageclassembeds dim1 keepdimtrue 1e6 queryembeds queryembeds torch linalg normqueryembeds dim1 keepdimtrue 1e6 get class predictions predlogits torch einsum pd qd pq imageclassembeds queryembeds apply a learnable shift and scale to logits logitshift self logitshiftimageembeds logitscale self logitscaleimageembeds logitscale self elulogitscale 1 predlogits predlogits logitshift logitscale if querymask is not none if querymask ndim 1 querymask torch unsqueezequerymask dim2 predlogits predlogits totorch float64 predlogits torch wherequerymask 0 1e6 predlogits predlogits predlogits totorch float32 return predlogits imageclassembeds class owlvitforobjectdetectionowlvitpretrainedmodel configclass owlvitconfig def initself config owlvitconfig super initconfig self owlvit owlvitmodelconfig self classhead owlvitclasspredictionheadconfig self boxhead owlvitboxpredictionheadconfig self layernorm nn layernormconfig visionconfig hiddensize epsconfig visionconfig layernormeps self sigmoid nn sigmoid def normalizegridcornercoordinatesself featuremap torch floattensor computes normalized xy corner coordinates from featuremap if not featuremap ndim 4 raise valueerrorexpected input shape is batchsize numpatches numpatches hiddendim device featuremap device numpatches featuremap shape1 boxcoordinates np stack np meshgridnp arange1 numpatches 1 np arange1 numpatches 1 axis1 astypenp float32 boxcoordinates np arraynumpatches numpatches np float32 flatten h w 2 hw 2 boxcoordinates boxcoordinates reshape boxcoordinates shape0 boxcoordinates shape1 boxcoordinates shape2 boxcoordinates torch fromnumpyboxcoordinates todevice return boxcoordinates def computeboxbiasself featuremap torch floattensor torch floattensor the box center is biased to its position on the feature grid boxcoordinates self normalizegridcornercoordinatesfeaturemap boxcoordinates torch clipboxcoordinates 0 0 1 0 unnormalize xy boxcoordbias torch logboxcoordinates 1e4 torch log1pboxcoordinates 1e4 the box size is biased to the patch size boxsize torch fulllikeboxcoordbias 1 0 featuremap shape2 boxsizebias torch logboxsize 1e4 torch log1pboxsize 1e4 compute box bias boxbias torch catboxcoordbias boxsizebias dim1 return boxbias def boxpredictor self imagefeats torch floattensor featuremap torch floattensor torch floattensor bounding box detection head batchsize numboxes 4 predboxes self boxheadimagefeats compute the location of each token on the grid and use it to compute a bias for the bbox prediction predboxes self computeboxbiasfeaturemap predboxes self sigmoidpredboxes return predboxes def classpredictor self imagefeats torch floattensor queryembeds optionaltorch floattensor none querymask optionaltorch tensor none tupletorch floattensor predlogits imageclassembeds self classheadimagefeats queryembeds querymask return predlogits imageclassembeds def imagetextembedder self inputids torch tensor pixelvalues torch floattensor attentionmask torch tensor outputattentions optionalbool none outputhiddenstates optionalbool none tupletorch floattensor encode text and image outputs self owlvit pixelvaluespixelvalues inputidsinputids attentionmaskattentionmask outputattentionsoutputattentions outputhiddenstatesoutputhiddenstates returndicttrue get image embeddings lasthiddenstate outputs visionmodeloutput0 imageembeds self owlvit visionmodel postlayernormlasthiddenstate resize class token newsize tuplenp arrayimageembeds shape np array0 1 0 classtokenout torch broadcasttoimageembeds 1 newsize merge image embedding with class tokens imageembeds imageembeds 1 classtokenout imageembeds self layernormimageembeds resize to batchsize numpatches numpatches hiddensize newsize imageembeds shape0 intnp sqrtimageembeds shape1 intnp sqrtimageembeds shape1 imageembeds shape1 imageembeds imageembeds reshapenewsize textembeds outputs4 return textembeds imageembeds outputs def imageembedder self pixelvalues torch floattensor outputattentions optionalbool none outputhiddenstates optionalbool none tupletorch floattensor get owlvitmodel vision embeddings same as clip visionoutputs self owlvit visionmodelpixelvaluespixelvalues returndicttrue apply postlayernorm to lasthiddenstate return nonprojected output lasthiddenstate visionoutputs0 imageembeds self owlvit visionmodel postlayernormlasthiddenstate resize class token newsize tuplenp arrayimageembeds shape np array0 1 0 classtokenout torch broadcasttoimageembeds 1 newsize merge image embedding with class tokens imageembeds imageembeds 1 classtokenout imageembeds self layernormimageembeds resize to batchsize numpatches numpatches hiddensize newsize imageembeds shape0 intnp sqrtimageembeds shape1 intnp sqrtimageembeds shape1 imageembeds shape1 imageembeds imageembeds reshapenewsize return imageembeds visionoutputs def embedimagequery self queryimagefeatures torch floattensor queryfeaturemap torch floattensor torch floattensor classembeds self classpredictorqueryimagefeatures predboxes self boxpredictorqueryimagefeatures queryfeaturemap predboxesascorners centertocornersformatpredboxes loop over query images bestclassembeds bestboxindices predboxesdevice predboxesascorners device for i in rangequeryimagefeatures shape0 eachquerybox torch tensor0 0 1 1 devicepredboxesdevice eachquerypredboxes predboxesascornersi ious boxioueachquerybox eachquerypredboxes if there are no overlapping boxes fall back to generalized iou if torch allious0 0 0 ious generalizedboxioueachquerybox eachquerypredboxes use an adaptive threshold to include all boxes within 80 of the best iou iouthreshold torch maxious 0 8 selectedinds ious0 iouthreshold nonzero if selectedinds numel selectedembeddings classembedsiselectedinds squeeze1 meanembeds torch meanclassembedsi axis0 meansim torch einsumd idi meanembeds selectedembeddings bestboxind selectedindstorch argminmeansim bestclassembeds appendclassembedsibestboxind bestboxindices appendbestboxind if bestclassembeds queryembeds torch stackbestclassembeds boxindices torch stackbestboxindices else queryembeds boxindices none none return queryembeds boxindices predboxes addstartdocstringstomodelforwardowlvitimageguidedobjectdetectioninputsdocstring replacereturndocstringsoutputtypeowlvitimageguidedobjectdetectionoutput configclassowlvitconfig def imageguideddetection self pixelvalues torch floattensor querypixelvalues optionaltorch floattensor none outputattentions optionalbool none outputhiddenstates optionalbool none returndict optionalbool none owlvitimageguidedobjectdetectionoutput r returns examples python import requests from pil import image import torch from transformers import autoprocessor owlvitforobjectdetection processor autoprocessor frompretrainedgoogleowlvitbasepatch16 model owlvitforobjectdetection frompretrainedgoogleowlvitbasepatch16 url http images cocodataset orgval2017000000039769 jpg image image openrequests geturl streamtrue raw queryurl http images cocodataset orgval2017000000001675 jpg queryimage image openrequests getqueryurl streamtrue raw inputs processorimagesimage queryimagesqueryimage returntensorspt with torch nograd outputs model imageguideddetectioninputs target image sizes height width to rescale box predictions batchsize 2 targetsizes torch tensorimage size 1 convert outputs bounding boxes and class logits to coco api results processor postprocessimageguideddetection outputsoutputs threshold0 6 nmsthreshold0 3 targetsizestargetsizes i 0 retrieve predictions for the first image boxes scores resultsiboxes resultsiscores for box score in zipboxes scores box roundi 2 for i in box tolist printfdetected similar object with confidence roundscore item 3 at location box detected similar object with confidence 0 856 at location 10 94 50 4 315 8 471 39 detected similar object with confidence 1 0 at location 334 84 25 33 636 16 374 71 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 returndict compute feature maps for the input and query images queryfeaturemap self imageembedderpixelvaluesquerypixelvalues0 featuremap visionoutputs self imageembedder pixelvaluespixelvalues outputattentionsoutputattentions outputhiddenstatesoutputhiddenstates batchsize numpatches numpatches hiddendim featuremap shape imagefeats torch reshapefeaturemap batchsize numpatches numpatches hiddendim batchsize numpatches numpatches hiddendim queryfeaturemap shape queryimagefeats torch reshapequeryfeaturemap batchsize numpatches numpatches hiddendim get top class embedding and best box index for each query image in batch queryembeds bestboxindices querypredboxes self embedimagequeryqueryimagefeats queryfeaturemap predict object classes batchsize numpatches numqueries1 predlogits classembeds self classpredictorimagefeatsimagefeats queryembedsqueryembeds predict object boxes targetpredboxes self boxpredictorimagefeats featuremap if not returndict output featuremap queryfeaturemap targetpredboxes querypredboxes predlogits classembeds visionoutputs totuple output tuplex for x in output if x is not none return output return owlvitimageguidedobjectdetectionoutput imageembedsfeaturemap queryimageembedsqueryfeaturemap targetpredboxestargetpredboxes querypredboxesquerypredboxes logitspredlogits classembedsclassembeds textmodeloutputnone visionmodeloutputvisionoutputs addstartdocstringstomodelforwardowlvitobjectdetectioninputsdocstring replacereturndocstringsoutputtypeowlvitobjectdetectionoutput configclassowlvitconfig def forward self inputids torch tensor pixelvalues torch floattensor attentionmask optionaltorch tensor none outputattentions optionalbool none outputhiddenstates optionalbool none returndict optionalbool none owlvitobjectdetectionoutput r returns examples python import requests from pil import image import torch from transformers import autoprocessor owlvitforobjectdetection processor autoprocessor frompretrainedgoogleowlvitbasepatch32 model owlvitforobjectdetection frompretrainedgoogleowlvitbasepatch32 url http images cocodataset orgval2017000000039769 jpg image image openrequests geturl streamtrue raw texts a photo of a cat a photo of a dog inputs processortexttexts imagesimage returntensorspt outputs modelinputs target image sizes height width to rescale box predictions batchsize 2 targetsizes torch tensorimage size 1 convert outputs bounding boxes and class logits to final bounding boxes and scores results processor postprocessobjectdetection outputsoutputs threshold0 1 targetsizestargetsizes i 0 retrieve predictions for the first image for the corresponding text queries text textsi boxes scores labels resultsiboxes resultsiscores resultsilabels for box score label in zipboxes scores labels box roundi 2 for i in box tolist printfdetected textlabel with confidence roundscore item 3 at location box detected a photo of a cat with confidence 0 707 at location 324 97 20 44 640 58 373 29 detected a photo of a cat with confidence 0 717 at location 1 46 55 26 315 55 472 17 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 returndict embed images and text queries queryembeds featuremap outputs self imagetextembedder inputidsinputids pixelvaluespixelvalues attentionmaskattentionmask outputattentionsoutputattentions outputhiddenstatesoutputhiddenstates text and vision model outputs textoutputs outputs textmodeloutput visionoutputs outputs visionmodeloutput batchsize numpatches numpatches hiddendim featuremap shape imagefeats torch reshapefeaturemap batchsize numpatches numpatches hiddendim reshape from batchsize maxtextqueries hiddendim batchsize maxtextqueries hiddendim maxtextqueries inputids shape0 batchsize queryembeds queryembeds reshapebatchsize maxtextqueries queryembeds shape1 if first token is 0 then this is a padded query batchsize numqueries inputids inputids reshapebatchsize maxtextqueries inputids shape1 querymask inputids 0 0 predict object classes batchsize numpatches numqueries1 predlogits classembeds self classpredictorimagefeats queryembeds querymask predict object boxes predboxes self boxpredictorimagefeats featuremap if not returndict output predlogits predboxes queryembeds featuremap classembeds textoutputs totuple visionoutputs totuple output tuplex for x in output if x is not none return output return owlvitobjectdetectionoutput imageembedsfeaturemap textembedsqueryembeds predboxespredboxes logitspredlogits classembedsclassembeds textmodeloutputtextoutputs visionmodeloutputvisionoutputs coding utf 8 2022 google ai 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 owl vit model see all owlvit models at https huggingface co models filter owlvit copied from transformers models clip modeling_clip contrastive_loss with clip owlvit copied from transformers models clip modeling_clip clip_loss with clip owlvit 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 text_embeds torch floattensor of shape batch_size num_max_text_queries output_dim the text embeddings obtained by applying the projection layer to the pooled output of owlvittextmodel image_embeds torch floattensor of shape batch_size output_dim the image embeddings obtained by applying the projection layer to the pooled output of owlvitvisionmodel text_model_output tuple basemodeloutputwithpooling the output of the owlvittextmodel vision_model_output basemodeloutputwithpooling the output of the owlvitvisionmodel copied from transformers models detr modeling_detr _upcast protects from numerical overflows in multiplications by upcasting to the equivalent higher type copied from transformers models detr modeling_detr box_area computes the area of a set of bounding boxes which are specified by its x1 y1 x2 y2 coordinates args boxes torch floattensor of shape number_of_boxes 4 boxes for which the area will be computed they are expected to be in x1 y1 x2 y2 format with 0 x1 x2 and 0 y1 y2 returns torch floattensor a tensor containing the area for each box copied from transformers models detr modeling_detr box_iou n m 2 n m 2 n m 2 n m copied from transformers models detr modeling_detr generalized_box_iou generalized iou from https giou stanford edu the boxes should be in x0 y0 x1 y1 corner format returns torch floattensor a n m pairwise matrix where n len boxes1 and m len boxes2 degenerate boxes gives inf nan results so do an early check n m 2 output type of owlvitforobjectdetection args loss torch floattensor of shape 1 optional returned when labels are provided total loss as a linear combination of a negative log likehood cross entropy for class prediction and a bounding box loss the latter is defined as a linear combination of the l1 loss and the generalized scale invariant iou loss loss_dict dict optional a dictionary containing the individual losses useful for logging logits torch floattensor of shape batch_size num_patches num_queries classification logits including no object for all queries pred_boxes torch floattensor of shape batch_size num_patches 4 normalized boxes coordinates for all queries represented as center_x center_y width height these values are normalized in 0 1 relative to the size of each individual image in the batch disregarding possible padding you can use owlvitimageprocessor post_process_object_detection to retrieve the unnormalized bounding boxes text_embeds torch floattensor of shape batch_size num_max_text_queries output_dim the text embeddings obtained by applying the projection layer to the pooled output of owlvittextmodel image_embeds torch floattensor of shape batch_size patch_size patch_size output_dim pooled output of owlvitvisionmodel owl vit represents images as a set of image patches and computes image embeddings for each patch class_embeds torch floattensor of shape batch_size num_patches hidden_size class embeddings of all image patches owl vit represents images as a set of image patches where the total number of patches is image_size patch_size 2 text_model_output tuple basemodeloutputwithpooling the output of the owlvittextmodel vision_model_output basemodeloutputwithpooling the output of the owlvitvisionmodel output type of owlvitforobjectdetection image_guided_detection args logits torch floattensor of shape batch_size num_patches num_queries classification logits including no object for all queries target_pred_boxes torch floattensor of shape batch_size num_patches 4 normalized boxes coordinates for all queries represented as center_x center_y width height these values are normalized in 0 1 relative to the size of each individual target image in the batch disregarding possible padding you can use owlvitimageprocessor post_process_object_detection to retrieve the unnormalized bounding boxes query_pred_boxes torch floattensor of shape batch_size num_patches 4 normalized boxes coordinates for all queries represented as center_x center_y width height these values are normalized in 0 1 relative to the size of each individual query image in the batch disregarding possible padding you can use owlvitimageprocessor post_process_object_detection to retrieve the unnormalized bounding boxes image_embeds torch floattensor of shape batch_size patch_size patch_size output_dim pooled output of owlvitvisionmodel owl vit represents images as a set of image patches and computes image embeddings for each patch query_image_embeds torch floattensor of shape batch_size patch_size patch_size output_dim pooled output of owlvitvisionmodel owl vit represents images as a set of image patches and computes image embeddings for each patch class_embeds torch floattensor of shape batch_size num_patches hidden_size class embeddings of all image patches owl vit represents images as a set of image patches where the total number of patches is image_size patch_size 2 text_model_output tuple basemodeloutputwithpooling the output of the owlvittextmodel vision_model_output basemodeloutputwithpooling the output of the owlvitvisionmodel shape batch_size num_channels height width position_ids 1 len position emb is contiguous in memory and exported when serialized 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 for int8 compatibility sometimes the attn_probs are in fp32 copied from transformers models clip modeling_clip clipmlp with clip owlvit copied from transformers models clip modeling_clip clipencoderlayer with clip owlvit 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 this model inherits from pretrainedmodel 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 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 owlvitconfig 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 num_max_text_queries sequence_length indices of input sequence tokens in the vocabulary indices can be obtained using autotokenizer see pretrainedtokenizer encode and pretrainedtokenizer __call__ for details what are input ids glossary input ids attention_mask torch tensor of shape batch_size num_max_text_queries 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 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 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 indices can be obtained using autotokenizer see pretrainedtokenizer encode and pretrainedtokenizer __call__ for details what are input ids glossary input ids attention_mask torch tensor of shape batch_size sequence_length optional mask to avoid performing attention on padding token indices mask values selected in 0 1 1 for tokens that are not masked 0 for tokens that are masked what are attention masks glossary attention mask pixel_values torch floattensor of shape batch_size num_channels height width pixel values 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 args pixel_values torch floattensor of shape batch_size num_channels height width pixel values input_ids torch longtensor of shape batch_size num_max_text_queries sequence_length optional indices of input sequence tokens in the vocabulary indices can be obtained using autotokenizer see pretrainedtokenizer encode and pretrainedtokenizer __call__ for details what are input ids glossary input ids attention_mask torch tensor of shape batch_size num_max_text_queries 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 output_hidden_states bool optional whether or not to return the last hidden state see text_model_last_hidden_state and vision_model_last_hidden_state 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 query_pixel_values torch floattensor of shape batch_size num_channels height width pixel values of query image s to be detected pass in one query image per target image 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 owlvitencoderlayer args config owlvitconfig args inputs_embeds torch floattensor of shape batch_size sequence_length hidden_size attention_mask torch tensor of shape batch_size sequence_length optional mask to avoid performing attention on padding token indices mask values selected in 0 1 1 for tokens that are not masked 0 for tokens that are masked what are attention masks glossary attention mask causal_attention_mask torch tensor of shape batch_size sequence_length optional causal mask for the text model mask values selected in 0 1 1 for tokens that are not masked 0 for tokens that are masked what are attention masks glossary attention mask output_attentions bool optional whether or not to return the attentions tensors of all attention layers see attentions under returned tensors for more detail output_hidden_states bool optional whether or not to return the hidden states of all layers see hidden_states under returned tensors for more detail return_dict bool optional whether or not to return a utils modeloutput instead of a plain tuple returns num_samples seq_len input_shape where num_samples batch_size num_max_text_queries owlvit s text model uses causal mask prepare it here https github com openai clip blob cfcffb90e69f37bf2ff1e988237a0fbe41f33c04 clip model py l324 expand attention_mask num_samples seq_len num_samples 1 tgt_seq_len src_seq_len take features from the end of tokens embedding end of token is the highest number in each sequence casting to torch int for onnx compatibility argmax doesn t support int64 inputs with opset 14 initialize weights and apply final processing returns examples python from transformers import autoprocessor owlvittextmodel model owlvittextmodel from_pretrained google owlvit base patch32 processor autoprocessor from_pretrained google owlvit base patch32 inputs processor text a photo of a cat a photo of a dog photo of a astranaut return_tensors pt outputs model inputs last_hidden_state outputs last_hidden_state pooled_output outputs pooler_output pooled eos token states get embeddings for all text queries in all batch samples returns cast the input to the expected dtype initialize weights and apply final processing returns examples python from pil import image import requests from transformers import autoprocessor owlvitvisionmodel model owlvitvisionmodel from_pretrained google owlvit base patch32 processor autoprocessor from_pretrained google owlvit base patch32 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 owlvittextmodel examples python from transformers import autoprocessor owlvitmodel model owlvitmodel from_pretrained google owlvit base patch32 processor autoprocessor from_pretrained google owlvit base patch32 inputs processor text a photo of a cat a photo of a dog photo of a astranaut return_tensors pt text_features model get_text_features inputs use owl vit model s config for some fields if specified instead of those of vision text components get embeddings for all text queries in all batch samples 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 owlvitvisionmodel examples python from pil import image import requests from transformers import autoprocessor owlvitmodel model owlvitmodel from_pretrained google owlvit base patch32 processor autoprocessor from_pretrained google owlvit base patch32 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 owl vit model s config for some fields if specified instead of those of vision text components returns examples python from pil import image import requests from transformers import autoprocessor owlvitmodel model owlvitmodel from_pretrained google owlvit base patch32 processor autoprocessor from_pretrained google owlvit base patch32 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 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 owl vit model s config for some fields if specified instead of those of vision text components get embeddings for all text queries in all batch samples normalized features cosine similarity as logits and set it on the correct device normalize image and text features get class predictions apply a learnable shift and scale to logits computes normalized xy corner coordinates from feature_map flatten h w 2 h w 2 the box center is biased to its position on the feature grid unnormalize xy the box size is biased to the patch size compute box bias args image_feats features extracted from the image returned by the image_text_embedder method feature_map a spatial re arrangement of image_features also returned by the image_text_embedder method returns pred_boxes list of predicted boxes cxcywh normalized to 0 1 nested within a dictionary bounding box detection head batch_size num_boxes 4 compute the location of each token on the grid and use it to compute a bias for the bbox prediction args image_feats features extracted from the image_text_embedder query_embeds text query embeddings query_mask must be provided with query_embeddings a mask indicating which query embeddings are valid encode text and image get image embeddings resize class token merge image embedding with class tokens resize to batch_size num_patches num_patches hidden_size get owlvitmodel vision embeddings same as clip apply post_layernorm to last_hidden_state return non projected output resize class token merge image embedding with class tokens resize to batch_size num_patches num_patches hidden_size loop over query images if there are no overlapping boxes fall back to generalized iou use an adaptive threshold to include all boxes within 80 of the best iou returns examples python import requests from pil import image import torch from transformers import autoprocessor owlvitforobjectdetection processor autoprocessor from_pretrained google owlvit base patch16 model owlvitforobjectdetection from_pretrained google owlvit base patch16 url http images cocodataset org val2017 000000039769 jpg image image open requests get url stream true raw query_url http images cocodataset org val2017 000000001675 jpg query_image image open requests get query_url stream true raw inputs processor images image query_images query_image return_tensors pt with torch no_grad outputs model image_guided_detection inputs target image sizes height width to rescale box predictions batch_size 2 target_sizes torch tensor image size 1 convert outputs bounding boxes and class logits to coco api results processor post_process_image_guided_detection outputs outputs threshold 0 6 nms_threshold 0 3 target_sizes target_sizes i 0 retrieve predictions for the first image boxes scores results i boxes results i scores for box score in zip boxes scores box round i 2 for i in box tolist print f detected similar object with confidence round score item 3 at location box detected similar object with confidence 0 856 at location 10 94 50 4 315 8 471 39 detected similar object with confidence 1 0 at location 334 84 25 33 636 16 374 71 compute feature maps for the input and query images get top class embedding and best box index for each query image in batch predict object classes batch_size num_patches num_queries 1 predict object boxes returns examples python import requests from pil import image import torch from transformers import autoprocessor owlvitforobjectdetection processor autoprocessor from_pretrained google owlvit base patch32 model owlvitforobjectdetection from_pretrained google owlvit base patch32 url http images cocodataset org val2017 000000039769 jpg image image open requests get url stream true raw texts a photo of a cat a photo of a dog inputs processor text texts images image return_tensors pt outputs model inputs target image sizes height width to rescale box predictions batch_size 2 target_sizes torch tensor image size 1 convert outputs bounding boxes and class logits to final bounding boxes and scores results processor post_process_object_detection outputs outputs threshold 0 1 target_sizes target_sizes i 0 retrieve predictions for the first image for the corresponding text queries text texts i boxes scores labels results i boxes results i scores results i labels for box score label in zip boxes scores labels box round i 2 for i in box tolist print f detected text label with confidence round score item 3 at location box detected a photo of a cat with confidence 0 707 at location 324 97 20 44 640 58 373 29 detected a photo of a cat with confidence 0 717 at location 1 46 55 26 315 55 472 17 embed images and text queries text and vision model outputs reshape from batch_size max_text_queries hidden_dim batch_size max_text_queries hidden_dim if first token is 0 then this is a padded query batch_size num_queries predict object classes batch_size num_patches num_queries 1 predict object boxes
import warnings
from dataclasses import dataclass
from typing import Any, Dict, Optional, Tuple, Union

import numpy as np
import torch
import torch.utils.checkpoint
from torch import Tensor, 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,
    is_vision_available,
    logging,
    replace_return_docstrings,
)
from .configuration_owlvit import OwlViTConfig, OwlViTTextConfig, OwlViTVisionConfig


if is_vision_available():
    from transformers.image_transforms import center_to_corners_format


logger = logging.get_logger(__name__)

_CHECKPOINT_FOR_DOC = "google/owlvit-base-patch32"


OWLVIT_PRETRAINED_MODEL_ARCHIVE_LIST = [
    "google/owlvit-base-patch32",
    "google/owlvit-base-patch16",
    "google/owlvit-large-patch14",
]



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



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


@dataclass
class OwlViTOutput(ModelOutput):
    

    loss: Optional[torch.FloatTensor] = None
    logits_per_image: torch.FloatTensor = None
    logits_per_text: 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()
        )



def _upcast(t: Tensor) -> Tensor:
    
    if t.is_floating_point():
        return t if t.dtype in (torch.float32, torch.float64) else t.float()
    else:
        return t if t.dtype in (torch.int32, torch.int64) else t.int()



def box_area(boxes: Tensor) -> Tensor:
    
    boxes = _upcast(boxes)
    return (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])



def box_iou(boxes1, boxes2):
    area1 = box_area(boxes1)
    area2 = box_area(boxes2)

    left_top = torch.max(boxes1[:, None, :2], boxes2[:, :2])  
    right_bottom = torch.min(boxes1[:, None, 2:], boxes2[:, 2:])  

    width_height = (right_bottom - left_top).clamp(min=0)  
    inter = width_height[:, :, 0] * width_height[:, :, 1]  

    union = area1[:, None] + area2 - inter

    iou = inter / union
    return iou, union



def generalized_box_iou(boxes1, boxes2):
    
    
    
    if not (boxes1[:, 2:] >= boxes1[:, :2]).all():
        raise ValueError(f"boxes1 must be in [x0, y0, x1, y1] (corner) format, but got {boxes1}")
    if not (boxes2[:, 2:] >= boxes2[:, :2]).all():
        raise ValueError(f"boxes2 must be in [x0, y0, x1, y1] (corner) format, but got {boxes2}")
    iou, union = box_iou(boxes1, boxes2)

    top_left = torch.min(boxes1[:, None, :2], boxes2[:, :2])
    bottom_right = torch.max(boxes1[:, None, 2:], boxes2[:, 2:])

    width_height = (bottom_right - top_left).clamp(min=0)  
    area = width_height[:, :, 0] * width_height[:, :, 1]

    return iou - (area - union) / area


@dataclass
class OwlViTObjectDetectionOutput(ModelOutput):
    

    loss: Optional[torch.FloatTensor] = None
    loss_dict: Optional[Dict] = None
    logits: torch.FloatTensor = None
    pred_boxes: torch.FloatTensor = None
    text_embeds: torch.FloatTensor = None
    image_embeds: torch.FloatTensor = None
    class_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()
        )


@dataclass
class OwlViTImageGuidedObjectDetectionOutput(ModelOutput):
    

    logits: torch.FloatTensor = None
    image_embeds: torch.FloatTensor = None
    query_image_embeds: torch.FloatTensor = None
    target_pred_boxes: torch.FloatTensor = None
    query_pred_boxes: torch.FloatTensor = None
    class_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 OwlViTVisionEmbeddings(nn.Module):
    def __init__(self, config: OwlViTVisionConfig):
        super().__init__()
        self.config = config
        self.embed_dim = config.hidden_size
        self.class_embedding = nn.Parameter(torch.randn(config.hidden_size))

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

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

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

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

        return embeddings


class OwlViTTextEmbeddings(nn.Module):
    def __init__(self, config: OwlViTTextConfig):
        super().__init__()
        self.token_embedding = nn.Embedding(config.vocab_size, config.hidden_size)
        self.position_embedding = nn.Embedding(config.max_position_embeddings, config.hidden_size)

        
        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 OwlViTAttention(nn.Module):
    

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

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

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

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

        bsz, tgt_len, embed_dim = hidden_states.size()

        
        query_states = self.q_proj(hidden_states) * self.scale
        key_states = self._shape(self.k_proj(hidden_states), -1, bsz)
        value_states = self._shape(self.v_proj(hidden_states), -1, bsz)

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

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

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

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

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

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

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

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

        
        attn_probs = attn_probs.to(value_states.dtype)

        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 OwlViTMLP(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.activation_fn = ACT2FN[config.hidden_act]
        self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size)
        self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size)

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



class OwlViTEncoderLayer(nn.Module):
    def __init__(self, config: OwlViTConfig):
        super().__init__()
        self.embed_dim = config.hidden_size
        self.self_attn = OwlViTAttention(config)
        self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
        self.mlp = OwlViTMLP(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 OwlViTPreTrainedModel(PreTrainedModel):
    

    config_class = OwlViTConfig
    base_model_prefix = "owlvit"
    supports_gradient_checkpointing = True
    _no_split_modules = ["OwlViTEncoderLayer"]

    def _init_weights(self, module):
        
        factor = self.config.initializer_factor
        if isinstance(module, OwlViTTextEmbeddings):
            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, OwlViTVisionEmbeddings):
            factor = self.config.initializer_factor
            nn.init.normal_(module.class_embedding, mean=0.0, std=module.embed_dim**-0.5 * factor)
            nn.init.normal_(module.patch_embedding.weight, std=module.config.initializer_range * factor)
            nn.init.normal_(module.position_embedding.weight, std=module.config.initializer_range * factor)
        elif isinstance(module, OwlViTAttention):
            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, OwlViTMLP):
            factor = self.config.initializer_factor
            in_proj_std = (module.config.hidden_size**-0.5) * ((2 * module.config.num_hidden_layers) ** -0.5) * factor
            fc_std = (2 * module.config.hidden_size) ** -0.5 * factor
            nn.init.normal_(module.fc1.weight, std=fc_std)
            nn.init.normal_(module.fc2.weight, std=in_proj_std)
        elif isinstance(module, OwlViTModel):
            nn.init.normal_(
                module.text_projection.weight,
                std=module.text_embed_dim**-0.5 * self.config.initializer_factor,
            )
            nn.init.normal_(
                module.visual_projection.weight,
                std=module.vision_embed_dim**-0.5 * self.config.initializer_factor,
            )
        if isinstance(module, nn.LayerNorm):
            module.bias.data.zero_()
            module.weight.data.fill_(1.0)
        if isinstance(module, nn.Linear) and module.bias is not None:
            module.bias.data.zero_()


OWLVIT_START_DOCSTRING = r

OWLVIT_TEXT_INPUTS_DOCSTRING = r

OWLVIT_VISION_INPUTS_DOCSTRING = r

OWLVIT_INPUTS_DOCSTRING = r

OWLVIT_OBJECT_DETECTION_INPUTS_DOCSTRING = r

OWLVIT_IMAGE_GUIDED_OBJECT_DETECTION_INPUTS_DOCSTRING = r


class OwlViTEncoder(nn.Module):
    

    def __init__(self, config: OwlViTConfig):
        super().__init__()
        self.layers = nn.ModuleList([OwlViTEncoderLayer(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 encoder_layer in 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 OwlViTTextTransformer(nn.Module):
    def __init__(self, config: OwlViTTextConfig):
        super().__init__()
        self.config = config
        embed_dim = config.hidden_size
        self.embeddings = OwlViTTextEmbeddings(config)
        self.encoder = OwlViTEncoder(config)
        self.final_layer_norm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)

    @add_start_docstrings_to_model_forward(OWLVIT_TEXT_INPUTS_DOCSTRING)
    @replace_return_docstrings(output_type=BaseModelOutputWithPooling, config_class=OwlViTTextConfig)
    def forward(
        self,
        input_ids: torch.Tensor,
        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

        input_shape = input_ids.size()
        input_ids = input_ids.view(-1, input_shape[-1])
        hidden_states = self.embeddings(input_ids=input_ids, position_ids=position_ids)

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

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

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

        
        
        pooled_output = last_hidden_state[
            torch.arange(last_hidden_state.shape[0], device=last_hidden_state.device),
            input_ids.to(torch.int).argmax(dim=-1).to(last_hidden_state.device),
        ]

        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 OwlViTTextModel(OwlViTPreTrainedModel):
    config_class = OwlViTTextConfig

    def __init__(self, config: OwlViTTextConfig):
        super().__init__(config)
        self.text_model = OwlViTTextTransformer(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(OWLVIT_TEXT_INPUTS_DOCSTRING)
    @replace_return_docstrings(output_type=BaseModelOutputWithPooling, config_class=OwlViTTextConfig)
    def forward(
        self,
        input_ids: torch.Tensor,
        attention_mask: 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,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )


class OwlViTVisionTransformer(nn.Module):
    def __init__(self, config: OwlViTVisionConfig):
        super().__init__()
        self.config = config

        self.embeddings = OwlViTVisionEmbeddings(config)
        self.pre_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.encoder = OwlViTEncoder(config)
        self.post_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)

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

        
        expected_input_dtype = self.embeddings.patch_embedding.weight.dtype
        pixel_values = pixel_values.to(expected_input_dtype)

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

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

        last_hidden_state = encoder_outputs[0]
        pooled_output = last_hidden_state[:, 0, :]

        pooled_output = self.post_layernorm(pooled_output)

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

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


class OwlViTVisionModel(OwlViTPreTrainedModel):
    config_class = OwlViTVisionConfig
    main_input_name = "pixel_values"

    def __init__(self, config: OwlViTVisionConfig):
        super().__init__(config)
        self.vision_model = OwlViTVisionTransformer(config)
        
        self.post_init()

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

    @add_start_docstrings_to_model_forward(OWLVIT_VISION_INPUTS_DOCSTRING)
    @replace_return_docstrings(output_type=BaseModelOutputWithPooling, config_class=OwlViTVisionConfig)
    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(OWLVIT_START_DOCSTRING)
class OwlViTModel(OwlViTPreTrainedModel):
    config_class = OwlViTConfig

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

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

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

        text_config = config.text_config
        vision_config = config.vision_config

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

        self.text_model = OwlViTTextTransformer(text_config)
        self.vision_model = OwlViTVisionTransformer(vision_config)

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

        
        self.post_init()

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

        
        text_output = self.text_model(input_ids=input_ids, attention_mask=attention_mask, return_dict=return_dict)
        pooled_output = text_output[1]
        text_features = self.text_projection(pooled_output)

        return text_features

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

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

        text_embeds = text_outputs[1]
        text_embeds = self.text_projection(text_embeds)
        image_embeds = vision_outputs[1]
        image_embeds = self.visual_projection(image_embeds)

        
        image_embeds = image_embeds / torch.linalg.norm(image_embeds, ord=2, dim=-1, keepdim=True)
        text_embeds_norm = text_embeds / torch.linalg.norm(text_embeds, ord=2, dim=-1, keepdim=True)

        
        logit_scale = self.logit_scale.exp().to(image_embeds.device)

        logits_per_text = torch.matmul(text_embeds_norm, image_embeds.t()) * logit_scale
        logits_per_image = logits_per_text.t()

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

        if return_base_image_embeds:
            warnings.warn(
                "`return_base_image_embeds` is deprecated and will be removed in v4.27 of Transformers, one can"
                " obtain the base (unprojected) image embeddings from outputs.vision_model_output.",
                FutureWarning,
            )
            last_hidden_state = vision_outputs[0]
            image_embeds = self.vision_model.post_layernorm(last_hidden_state)
        else:
            text_embeds = text_embeds_norm

        if not return_dict:
            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 OwlViTOutput(
            loss=loss,
            logits_per_image=logits_per_image,
            logits_per_text=logits_per_text,
            text_embeds=text_embeds,
            image_embeds=image_embeds,
            text_model_output=text_outputs,
            vision_model_output=vision_outputs,
        )


class OwlViTBoxPredictionHead(nn.Module):
    def __init__(self, config: OwlViTConfig, out_dim: int = 4):
        super().__init__()

        width = config.vision_config.hidden_size
        self.dense0 = nn.Linear(width, width)
        self.dense1 = nn.Linear(width, width)
        self.gelu = nn.GELU()
        self.dense2 = nn.Linear(width, out_dim)

    def forward(self, image_features: torch.Tensor) -> torch.FloatTensor:
        output = self.dense0(image_features)
        output = self.gelu(output)
        output = self.dense1(output)
        output = self.gelu(output)
        output = self.dense2(output)
        return output


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

        out_dim = config.text_config.hidden_size
        self.query_dim = config.vision_config.hidden_size

        self.dense0 = nn.Linear(self.query_dim, out_dim)
        self.logit_shift = nn.Linear(self.query_dim, 1)
        self.logit_scale = nn.Linear(self.query_dim, 1)
        self.elu = nn.ELU()

    def forward(
        self,
        image_embeds: torch.FloatTensor,
        query_embeds: Optional[torch.FloatTensor],
        query_mask: Optional[torch.Tensor],
    ) -> Tuple[torch.FloatTensor]:
        image_class_embeds = self.dense0(image_embeds)
        if query_embeds is None:
            device = image_class_embeds.device
            batch_size, num_patches = image_class_embeds.shape[:2]
            pred_logits = torch.zeros((batch_size, num_patches, self.query_dim)).to(device)
            return (pred_logits, image_class_embeds)

        
        image_class_embeds = image_class_embeds / (torch.linalg.norm(image_class_embeds, dim=-1, keepdim=True) + 1e-6)
        query_embeds = query_embeds / (torch.linalg.norm(query_embeds, dim=-1, keepdim=True) + 1e-6)

        
        pred_logits = torch.einsum("...pd,...qd->...pq", image_class_embeds, query_embeds)

        
        logit_shift = self.logit_shift(image_embeds)
        logit_scale = self.logit_scale(image_embeds)
        logit_scale = self.elu(logit_scale) + 1
        pred_logits = (pred_logits + logit_shift) * logit_scale

        if query_mask is not None:
            if query_mask.ndim > 1:
                query_mask = torch.unsqueeze(query_mask, dim=-2)

            pred_logits = pred_logits.to(torch.float64)
            pred_logits = torch.where(query_mask == 0, -1e6, pred_logits)
            pred_logits = pred_logits.to(torch.float32)

        return (pred_logits, image_class_embeds)


class OwlViTForObjectDetection(OwlViTPreTrainedModel):
    config_class = OwlViTConfig

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

        self.owlvit = OwlViTModel(config)
        self.class_head = OwlViTClassPredictionHead(config)
        self.box_head = OwlViTBoxPredictionHead(config)

        self.layer_norm = nn.LayerNorm(config.vision_config.hidden_size, eps=config.vision_config.layer_norm_eps)
        self.sigmoid = nn.Sigmoid()

    def normalize_grid_corner_coordinates(self, feature_map: torch.FloatTensor):
        
        if not feature_map.ndim == 4:
            raise ValueError("Expected input shape is [batch_size, num_patches, num_patches, hidden_dim]")

        device = feature_map.device
        num_patches = feature_map.shape[1]

        box_coordinates = np.stack(
            np.meshgrid(np.arange(1, num_patches + 1), np.arange(1, num_patches + 1)), axis=-1
        ).astype(np.float32)
        box_coordinates /= np.array([num_patches, num_patches], np.float32)

        
        box_coordinates = box_coordinates.reshape(
            box_coordinates.shape[0] * box_coordinates.shape[1], box_coordinates.shape[2]
        )
        box_coordinates = torch.from_numpy(box_coordinates).to(device)

        return box_coordinates

    def compute_box_bias(self, feature_map: torch.FloatTensor) -> torch.FloatTensor:
        
        box_coordinates = self.normalize_grid_corner_coordinates(feature_map)
        box_coordinates = torch.clip(box_coordinates, 0.0, 1.0)

        
        box_coord_bias = torch.log(box_coordinates + 1e-4) - torch.log1p(-box_coordinates + 1e-4)

        
        box_size = torch.full_like(box_coord_bias, 1.0 / feature_map.shape[-2])
        box_size_bias = torch.log(box_size + 1e-4) - torch.log1p(-box_size + 1e-4)

        
        box_bias = torch.cat([box_coord_bias, box_size_bias], dim=-1)
        return box_bias

    def box_predictor(
        self,
        image_feats: torch.FloatTensor,
        feature_map: torch.FloatTensor,
    ) -> torch.FloatTensor:
        
        
        pred_boxes = self.box_head(image_feats)

        
        pred_boxes += self.compute_box_bias(feature_map)
        pred_boxes = self.sigmoid(pred_boxes)
        return pred_boxes

    def class_predictor(
        self,
        image_feats: torch.FloatTensor,
        query_embeds: Optional[torch.FloatTensor] = None,
        query_mask: Optional[torch.Tensor] = None,
    ) -> Tuple[torch.FloatTensor]:
        
        (pred_logits, image_class_embeds) = self.class_head(image_feats, query_embeds, query_mask)

        return (pred_logits, image_class_embeds)

    def image_text_embedder(
        self,
        input_ids: torch.Tensor,
        pixel_values: torch.FloatTensor,
        attention_mask: torch.Tensor,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
    ) -> Tuple[torch.FloatTensor]:
        
        outputs = self.owlvit(
            pixel_values=pixel_values,
            input_ids=input_ids,
            attention_mask=attention_mask,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=True,
        )

        
        last_hidden_state = outputs.vision_model_output[0]
        image_embeds = self.owlvit.vision_model.post_layernorm(last_hidden_state)

        
        new_size = tuple(np.array(image_embeds.shape) - np.array((0, 1, 0)))
        class_token_out = torch.broadcast_to(image_embeds[:, :1, :], new_size)

        
        image_embeds = image_embeds[:, 1:, :] * class_token_out
        image_embeds = self.layer_norm(image_embeds)

        
        new_size = (
            image_embeds.shape[0],
            int(np.sqrt(image_embeds.shape[1])),
            int(np.sqrt(image_embeds.shape[1])),
            image_embeds.shape[-1],
        )
        image_embeds = image_embeds.reshape(new_size)
        text_embeds = outputs[-4]

        return (text_embeds, image_embeds, outputs)

    def image_embedder(
        self,
        pixel_values: torch.FloatTensor,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
    ) -> Tuple[torch.FloatTensor]:
        
        vision_outputs = self.owlvit.vision_model(pixel_values=pixel_values, return_dict=True)

        
        last_hidden_state = vision_outputs[0]
        image_embeds = self.owlvit.vision_model.post_layernorm(last_hidden_state)

        
        new_size = tuple(np.array(image_embeds.shape) - np.array((0, 1, 0)))
        class_token_out = torch.broadcast_to(image_embeds[:, :1, :], new_size)

        
        image_embeds = image_embeds[:, 1:, :] * class_token_out
        image_embeds = self.layer_norm(image_embeds)

        
        new_size = (
            image_embeds.shape[0],
            int(np.sqrt(image_embeds.shape[1])),
            int(np.sqrt(image_embeds.shape[1])),
            image_embeds.shape[-1],
        )
        image_embeds = image_embeds.reshape(new_size)

        return (image_embeds, vision_outputs)

    def embed_image_query(
        self, query_image_features: torch.FloatTensor, query_feature_map: torch.FloatTensor
    ) -> torch.FloatTensor:
        _, class_embeds = self.class_predictor(query_image_features)
        pred_boxes = self.box_predictor(query_image_features, query_feature_map)
        pred_boxes_as_corners = center_to_corners_format(pred_boxes)

        
        best_class_embeds = []
        best_box_indices = []
        pred_boxes_device = pred_boxes_as_corners.device

        for i in range(query_image_features.shape[0]):
            each_query_box = torch.tensor([[0, 0, 1, 1]], device=pred_boxes_device)
            each_query_pred_boxes = pred_boxes_as_corners[i]
            ious, _ = box_iou(each_query_box, each_query_pred_boxes)

            
            if torch.all(ious[0] == 0.0):
                ious = generalized_box_iou(each_query_box, each_query_pred_boxes)

            
            iou_threshold = torch.max(ious) * 0.8

            selected_inds = (ious[0] >= iou_threshold).nonzero()
            if selected_inds.numel():
                selected_embeddings = class_embeds[i][selected_inds.squeeze(1)]
                mean_embeds = torch.mean(class_embeds[i], axis=0)
                mean_sim = torch.einsum("d,id->i", mean_embeds, selected_embeddings)
                best_box_ind = selected_inds[torch.argmin(mean_sim)]
                best_class_embeds.append(class_embeds[i][best_box_ind])
                best_box_indices.append(best_box_ind)

        if best_class_embeds:
            query_embeds = torch.stack(best_class_embeds)
            box_indices = torch.stack(best_box_indices)
        else:
            query_embeds, box_indices = None, None

        return query_embeds, box_indices, pred_boxes

    @add_start_docstrings_to_model_forward(OWLVIT_IMAGE_GUIDED_OBJECT_DETECTION_INPUTS_DOCSTRING)
    @replace_return_docstrings(output_type=OwlViTImageGuidedObjectDetectionOutput, config_class=OwlViTConfig)
    def image_guided_detection(
        self,
        pixel_values: torch.FloatTensor,
        query_pixel_values: Optional[torch.FloatTensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> OwlViTImageGuidedObjectDetectionOutput:
        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.return_dict

        
        query_feature_map = self.image_embedder(pixel_values=query_pixel_values)[0]
        feature_map, vision_outputs = self.image_embedder(
            pixel_values=pixel_values,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
        )

        batch_size, num_patches, num_patches, hidden_dim = feature_map.shape
        image_feats = torch.reshape(feature_map, (batch_size, num_patches * num_patches, hidden_dim))

        batch_size, num_patches, num_patches, hidden_dim = query_feature_map.shape
        query_image_feats = torch.reshape(query_feature_map, (batch_size, num_patches * num_patches, hidden_dim))
        
        query_embeds, best_box_indices, query_pred_boxes = self.embed_image_query(query_image_feats, query_feature_map)

        
        (pred_logits, class_embeds) = self.class_predictor(image_feats=image_feats, query_embeds=query_embeds)

        
        target_pred_boxes = self.box_predictor(image_feats, feature_map)

        if not return_dict:
            output = (
                feature_map,
                query_feature_map,
                target_pred_boxes,
                query_pred_boxes,
                pred_logits,
                class_embeds,
                vision_outputs.to_tuple(),
            )
            output = tuple(x for x in output if x is not None)
            return output

        return OwlViTImageGuidedObjectDetectionOutput(
            image_embeds=feature_map,
            query_image_embeds=query_feature_map,
            target_pred_boxes=target_pred_boxes,
            query_pred_boxes=query_pred_boxes,
            logits=pred_logits,
            class_embeds=class_embeds,
            text_model_output=None,
            vision_model_output=vision_outputs,
        )

    @add_start_docstrings_to_model_forward(OWLVIT_OBJECT_DETECTION_INPUTS_DOCSTRING)
    @replace_return_docstrings(output_type=OwlViTObjectDetectionOutput, config_class=OwlViTConfig)
    def forward(
        self,
        input_ids: torch.Tensor,
        pixel_values: torch.FloatTensor,
        attention_mask: Optional[torch.Tensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> OwlViTObjectDetectionOutput:
        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.return_dict

        
        query_embeds, feature_map, outputs = self.image_text_embedder(
            input_ids=input_ids,
            pixel_values=pixel_values,
            attention_mask=attention_mask,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
        )

        
        text_outputs = outputs.text_model_output
        vision_outputs = outputs.vision_model_output

        batch_size, num_patches, num_patches, hidden_dim = feature_map.shape
        image_feats = torch.reshape(feature_map, (batch_size, num_patches * num_patches, hidden_dim))

        
        max_text_queries = input_ids.shape[0] // batch_size
        query_embeds = query_embeds.reshape(batch_size, max_text_queries, query_embeds.shape[-1])

        
        input_ids = input_ids.reshape(batch_size, max_text_queries, input_ids.shape[-1])
        query_mask = input_ids[..., 0] > 0

        
        (pred_logits, class_embeds) = self.class_predictor(image_feats, query_embeds, query_mask)

        
        pred_boxes = self.box_predictor(image_feats, feature_map)

        if not return_dict:
            output = (
                pred_logits,
                pred_boxes,
                query_embeds,
                feature_map,
                class_embeds,
                text_outputs.to_tuple(),
                vision_outputs.to_tuple(),
            )
            output = tuple(x for x in output if x is not None)
            return output

        return OwlViTObjectDetectionOutput(
            image_embeds=feature_map,
            text_embeds=query_embeds,
            pred_boxes=pred_boxes,
            logits=pred_logits,
            class_embeds=class_embeds,
            text_model_output=text_outputs,
            vision_model_output=vision_outputs,
        )