U ,-e@sdZddlZddlZddlmZddlmZmZm Z m Z ddl Z ddl Z ddlZ ddl mZddlmZddlmZmZdd lmZdd lmZmZmZmZmZd d lmZmZm Z e!e"Z#d Z$d gZ%dRe j&e j'ee(dddZ)e j&e j&dddZ*e j&e j&dddZ+e j&e(dddZ,dSe j&e-e.e(e j&dddZ/dTdd Z0d!d"Z1Gd#d$d$ej2Z3Gd%d&d&ej2Z4Gd'd(d(ej2Z5eGd)d*d*eZ6Gd+d,d,ej2Z7Gd-d.d.ej2Z8Gd/d0d0ej2Z9Gd1d2d2ej2Z:Gd3d4d4ej2Z;Gd5d6d6e;ZGd;d<dZAd?ZBd@ZCGdAdBdBej2ZDGdCdDdDej2ZEdUe jFe j'e jGe(dEdFdGZHGdHdIdIej2ZIGdJdKdKe?ZJGdLdMdMej2ZKGdNdOdOe?ZLee@GdPdQdQe?ZMdS)Vz PyTorch GroupViT model.N) dataclass)AnyOptionalTupleUnion)nn)ACT2FN)BaseModelOutputBaseModelOutputWithPooling)PreTrainedModel) ModelOutputadd_start_docstrings%add_start_docstrings_to_model_forwardloggingreplace_return_docstrings)GroupViTConfigGroupViTTextConfigGroupViTVisionConfigznvidia/groupvit-gcc-yfcc)maskdtypetgt_lencCsj|\}}|dk r|n|}|ddddddf|d|||}d|}||tjt|jS)z_ Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`. Nr?)sizeexpandtoZ masked_filltorchboolfinfomin)rrrbszsrc_lenZ expanded_maskZ inverted_maskr#o/var/www/html/Darija-Ai-Train/env/lib/python3.8/site-packages/transformers/models/groupvit/modeling_groupvit.py _expand_mask4s  *r%)logitsreturncCstj|tjt||jdS)Ndevice)r functionalZ cross_entropyrarangelenr))r&r#r#r$contrastive_lossDsr-) similarityr'cCs t|}t|}||dS)Ng@)r-t)r.Z caption_lossZ image_lossr#r#r$ groupvit_lossIs r0)r&dimcCsJ||}|j|ddd}tj|tjd||d}|||}|S)NTkeepdimrZ memory_formatr)softmaxmaxr zeros_likelegacy_contiguous_formatscatter_detach)r&r1y_softindexy_hardretr#r#r$ hard_softmaxOs  r?F)r&tauhardr1r'c Cstjjtjd|j|jdtjd|j|jd}||j}|||}| |}|r|j |ddd}tj |tj d ||d}|||} n|} | S)N)r)rrTr2rr4)r distributionsgumbelZGumbeltensorr)rsampleshaper5r6r7r8r9r:) r&rArBr1Z gumbel_distZgumbelsr;r<r=r>r#r#r$gumbel_softmaxYs   rIc Cs|||jdd}||kr@tt||}|jd|}n tt||}|jd|}|jd}|jd}|||||}tjj|||fd|d}|S)a Args: attentions (`torch.Tensor`): attention map of shape [batch_size, groups, feat_height*feat_width] height (`int`): height of the output attention map width (`int`): width of the output attention map align_corners (`bool`, *optional*): the `align_corner` argument for `nn.functional.interpolate`. Returns: `torch.Tensor`: resized attention map of shape [batch_size, groups, height, width] g?rrZbilinear)rmode align_corners)rHintnproundreshaperr* interpolate) attentionsheightwidthrLscale feat_width feat_height batch_sizegroupsr#r#r$resize_attention_mapos   rZc Csg}tbd}|D]R}|ddd}|dkr:|}n||}t|dddf|}||qW5QRX|d}|S)a1 Args: attentions (`tuple(torch.FloatTensor)`: tuple of attention maps returned by `GroupViTVisionTransformer` hw_shape (`tuple(int)`): height and width of the output attention map Returns: `torch.Tensor`: the attention map of shape [batch_size, groups, height, width] NrrJrr@)rZno_gradpermute contiguousrZappend)rRZhw_shapeZ attn_mapsZprev_attn_masksZ attn_masksZ cur_attn_mapZfinal_groupingr#r#r$get_grouping_from_attentionss  r^cs*eZdZedfdd ZddZZS)GroupViTCrossAttentionLayerconfigcsJtt||_tj|j|jd|_t ||_ tj|j|jd|_ dSNZeps) super__init__GroupViTAttentionattnr LayerNorm hidden_sizelayer_norm_epsnorm2 GroupViTMLPmlp norm_postselfra __class__r#r$res    z$GroupViTCrossAttentionLayer.__init__cCs<|}||j||dd}||||}||}|S)N)encoder_hidden_statesr)rgrmrkrn)rpquerykeyxr#r#r$forwards  z#GroupViTCrossAttentionLayer.forward)__name__ __module__ __qualname__rrerw __classcell__r#r#rqr$r_sr_cs4eZdZedfdd Zd ddZddZZS) GroupViTAssignAttentionr`csjt|jd|_t|j|j|_t|j|j|_t|j|j|_t|j|j|_ |j |_ dS)N) rdrerirUrLinearq_projk_projv_projproj assign_epsrorqr#r$res  z GroupViTAssignAttention.__init__TcCs@|r|jrt|d|d}n"|r,t|dd}ntjj|dd}|S)N)r1rBr1)trainingrIr?rr*r5)rprgrErBr#r#r$get_attns  z GroupViTAssignAttention.get_attncCs|}||}||}||}||dd|j}||}|j|ddd}||jddd|j}||}||}||fS)Nrr@F)rErBTr1r3) rrr transposerUrsumrr)rprtruvalueZraw_attnrgZ soft_attnoutr#r#r$rws     zGroupViTAssignAttention.forward)TT)rxryrzrrerrwr{r#r#rqr$r|s r|cs2eZdZedfdd ZddZddZZS)GroupViTTokenAssignr`cst||_tjjjd|_tj t j j r:j n j j f}fdd|D\}}t ||||_tjjjd|_tjjjd|_t|_t|_tjjjd|_tj|j|_dS)Nrccsg|]}t|jqSr#)rMri).0rvr`r#r$ sz0GroupViTTokenAssign.__init__..)rdrenum_output_grouprrhrirj norm_tokens isinstanceassign_mlp_ratio collectionsabcIterableGroupViTMixerMLP mlp_internorm_post_tokensnorm_xr_pre_assign_attnr|assign norm_new_xrl mlp_channels)rpranum_group_tokenrrZ tokens_dimZ channels_dimrqr`r$res    zGroupViTTokenAssign.__init__cCs||}||}|S)z Args: group_tokens (torch.Tensor): group tokens, [batch_size, num_group_tokens, channels] Returns: projected_group_tokens (torch.Tensor): [batch_size, num_output_groups, channels] )rr)rp group_tokensprojected_group_tokensr#r#r$project_group_tokens  z'GroupViTTokenAssign.project_group_tokencCs^||}||}||}|||}|||\}}||7}||||}||fS)z Args: image_tokens (`torch.Tensor`): image tokens, of shape [batch_size, input_length, channels] group_tokens (`torch.Tensor`): group tokens, [batch_size, num_group_tokens, channels] )rrrrrrr)rpZ image_tokensrrZnew_image_tokens attentionr#r#r$rw s    zGroupViTTokenAssign.forward)rxryrzrrerrwr{r#r#rqr$rs rc@seZdZUdZdZeejed<dZ ejed<dZ ejed<dZ ejed<dZ ejed<dZ ejed<dZeed <dZeed <eed d d ZdS)GroupViTModelOutputa\ Args: loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `return_loss` is `True`): Contrastive loss for image-text similarity. logits_per_image (`torch.FloatTensor` of shape `(image_batch_size, text_batch_size)`): The scaled dot product scores between `image_embeds` and `text_embeds`. This represents the image-text similarity scores. logits_per_text (`torch.FloatTensor` of shape `(text_batch_size, image_batch_size)`): The scaled dot product scores between `text_embeds` and `image_embeds`. This represents the text-image similarity scores. segmentation_logits (`torch.FloatTensor` of shape `(batch_size, config.num_labels, logits_height, logits_width)`): Classification scores for each pixel. The logits returned do not necessarily have the same size as the `pixel_values` passed as inputs. This is to avoid doing two interpolations and lose some quality when a user needs to resize the logits to the original image size as post-processing. You should always check your logits shape and resize as needed. text_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim`): The text embeddings obtained by applying the projection layer to the pooled output of [`GroupViTTextModel`]. image_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim`): The image embeddings obtained by applying the projection layer to the pooled output of [`GroupViTVisionModel`]. text_model_output (`BaseModelOutputWithPooling`): The output of the [`GroupViTTextModel`]. vision_model_output (`BaseModelOutputWithPooling`): The output of the [`GroupViTVisionModel`]. Nlosslogits_per_imagelogits_per_textsegmentation_logits text_embeds image_embedstext_model_outputvision_model_outputr'cstfddDS)Nc3s,|]$}|dkr|n t|VqdS))rrN)getattrto_tuple)rkrpr#r$ Nsz/GroupViTModelOutput.to_tuple..)tuplekeysrr#rr$rMs zGroupViTModelOutput.to_tuple)rxryrz__doc__rrr FloatTensor__annotations__rrrrrrr rrrrr#r#r#r$r!s !  rcsVeZdZdZd eeeeeeffeedfdd Zdej e ej d d d Z Z S)GroupViTPatchEmbeddingsz# Image to Patch Embedding. r image_size patch_size num_channels embed_dimcstt|tjjr|n||f}t|tjjr6|n||f}|d|d|d|d}||_||_||_t j ||||d|_ dS)Nrr)Z kernel_sizeZstride) rdrerrrrrr num_patchesrConv2d projection)rprrrrrrqr#r$reYs  z GroupViTPatchEmbeddings.__init__F pixel_valuesinterpolate_pos_encodingr'c Csx|j\}}}}|s\||jdks.||jdkr\td|d|d|jdd|jdd ||ddd}|S)NrrzInput image size (*z) doesn't match model ().rJ)rHr ValueErrorrflattenr)rprrrXrrSrTrvr#r#r$rwjs(zGroupViTPatchEmbeddings.forward)rrrr)F) rxryrzrrMrrrerTensorrrwr{r#r#rqr$rTsrcsReZdZedfdd ZejeeejdddZd eje ejdd d Z Z S) GroupViTVisionEmbeddingsr`csptt|j|j|j|jd|_|jj}t t d||j|_ t |j|_t j|j|jd|_||_dS)Nrrrc)rdrerrrrripatch_embeddingsrr Parameterrzerosposition_embeddingsZDropoutdropoutrhrj layernormra)rprarrqr#r$rews z!GroupViTVisionEmbeddings.__init__) embeddingsrSrTr'cCs|jd}||jjdkr(||kr(|jS|j}|jd}|jd}||jj}||jj} |d| d}} t|} } |dt| t| |dddd} || | | f} t j j | | ddd }|dddd dd|}|S) a# This method allows to interpolate the pre-trained position encodings, to be able to use the model on higher resolution images. Source: https://github.com/facebookresearch/dino/blob/de9ee3df6cf39fac952ab558447af1fa1365362a/vision_transformer.py#L174 rr@g?rrrJZbicubicF) scale_factorrKrL) rHrrarmathsqrtrPrMr[rr*rQview)rprrSrTZnpatchZpatch_pos_embedZnum_original_pos_embedr1rWrVZoriginal_heightZoriginal_widthZreshaped_patch_pos_embedrr#r#r$rs2     z1GroupViTVisionEmbeddings.interpolate_pos_encodingFrc Csd|j\}}}}|j||d}||}|\}}} |rL|||||}n ||j}||}|S)N)r)rHrrrrrr) rprrrXrrSrTrseq_len_r#r#r$rws   z GroupViTVisionEmbeddings.forward)F) rxryrzrrerrrMrrrwr{r#r#rqr$rvs"rcsLeZdZedfdd Zdeejeejeejej dddZ Z S) GroupViTTextEmbeddingsr`csRt|j}t|j||_t|j||_|j dt |j ddddS)N position_ids)rr@F) persistent) rdrerirZ EmbeddingZ vocab_sizetoken_embeddingZmax_position_embeddingsposition_embeddingZregister_bufferrr+rrprarrqr#r$res zGroupViTTextEmbeddings.__init__N) input_idsr inputs_embedsr'cCsb|dk r|jdn|jd}|dkr:|jddd|f}|dkrL||}||}||}|S)Nr@r)rHrrr)rprrrZ seq_lengthrrr#r#r$rws  zGroupViTTextEmbeddings.forward)NNN) rxryrzrrerr LongTensorrrrwr{r#r#rqr$rsrcseZdZdZeeeeedfdd ZeddZddZ de j e e j e j d d d Z de j e e j e eee jdddZZS) GroupViTStagezMThis corresponds to the `GroupingLayer` class in the GroupViT implementation.)radepthnum_prev_group_tokenrrcst||_||_|dkr8ttd|j|_ nd|_ d|_ t fddt |D|_ |dkr|t||d|_nd|_|dkr|dkrttjjjdt|jd||_nd|_dS) NrrFcsg|] }tqSr#GroupViTEncoderLayerrrr`r#r$rsz*GroupViTStage.__init__..)rarrrcrJ)rdrerrrrrrri group_tokengradient_checkpointing ModuleListrangelayersr downsample Sequentialrhrjrgroup_projector)rprarrrrrqr`r$res*  zGroupViTStage.__init__cCs |jdk SN)rrr#r#r$with_group_tokenszGroupViTStage.with_group_tokencCsB|jr6|ddd|j f|dd|j dffS|dfSdSr)rrrprvr#r#r$split_xs0zGroupViTStage.split_xN)rvrr'cCs|dkr |Stj||gddS)Nrr)rcat)rprvrr#r#r$concat_x szGroupViTStage.concat_xF) hidden_statesprev_group_tokenoutput_attentionsr'c Cs|jr6|j|ddd}|jdk r:|||}nd}|}|||}|jD]}||ddd}|d}qP||\}}d} |jdk r|||\}} ||f} |r| | f} | S)a Args: hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)` attention_mask (`torch.FloatTensor`): attention mask of size `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values. `(config.encoder_attention_heads,)`. output_attentions (`bool`, *optional*): Whether or not to return the grouping tensors of Grouping block. rr@N)attention_maskcausal_attention_mask) rrrrrrrrr) rprrrrrvZcat_xlayerZ layer_outroutputsr#r#r$rws$      zGroupViTStage.forward)N)NF)rxryrzrrrMrepropertyrrrrrrrrrrwr{r#r#rqr$rs&# rcsJeZdZdeeeeeeedfdd ZejejdddZ Z S) rlN)rariintermediate_size output_sizecspt||_t|j|_|dk r(|n|j}|dk r:|n|j}|dk rL|n|}t |||_ t |||_ dSr) rdrerar Z hidden_act activation_fnrirrr~fc1fc2)rprarirrrqr#r$re=s  zGroupViTMLP.__init__)rr'cCs"||}||}||}|Sr)rrr)rprr#r#r$rwMs   zGroupViTMLP.forward)NNN) rxryrzrrrMrerrrwr{r#r#rqr$rl<srlcseZdZfddZZS)rcs t|dd}|ddSNrrJ)rdrwrrrqr#r$rwUszGroupViTMixerMLP.forward)rxryrzrwr{r#r#rqr$rTsrc seZdZdZfddZejeedddZd eje eje eje ej e e e eje eje e ejfd d d Z ZS) rfz=Multi-headed attention from 'Attention Is All You Need' papercst||_|j|_|j|_|j|j|_|j|j|jkrZtd|jd|jd|jd|_ |j |_ t |j|j|_t |j|j|_t |j|j|_t |j|j|_dS)Nz;embed_dim must be divisible by num_heads (got `embed_dim`: z and `num_heads`: rr})rdrerarirZnum_attention_heads num_headshead_dimrrUZattention_dropoutrrr~rrrout_projrorqr#r$re]s  zGroupViTAttention.__init__)rFrr!cCs ||||j|jddSr)rrrrr\)rprFrr!r#r#r$_shapepszGroupViTAttention._shapeNF)rrrrsrr'cCs|\}}}|dk } |||j} | rT|||d|} |||d|} n(|||d|} |||d|} ||jd|jf} || ||j| } | j| } | j| } | d}t | | dd}|||j||fkrt d||j||fd||dk r|||d||fkrRt d|d||fd||||j|||}|||j||}|dk r||d||fkrt d|d||fd||||j|||}|||j||}t jj|dd}|r$|||j||}|||j||}nd}t jj||j|jd }t || }|||j||jfkrt d ||j||jfd||||j||j}| dd}||||}||}||fS) z#Input shape: Batch x Time x ChannelNr@rrJz$Attention weights should be of size z , but is z!Attention mask should be of size r)prz `attn_output` should be of size )rrrUrrrrrrrZbmmrrrr*r5rrrPr)rprrrrsrr!rrZis_cross_attentionZ query_statesZ key_statesZ value_statesZ proj_shaper" attn_weightsZattn_weights_reshapedZ attn_probsZ attn_outputr#r#r$rwss`         zGroupViTAttention.forward)NNNF)rxryrzrrerrrMrrrrrrwr{r#r#rqr$rfZs rfcsJeZdZedfdd Zdejejejeee ej dddZ Z S) rr`csRt|j|_t||_tj|j|jd|_ t ||_ tj|j|jd|_ dSrb) rdrerirrf self_attnrrhrj layer_norm1rlrm layer_norm2rorqr#r$res    zGroupViTEncoderLayer.__init__F)rrrrr'cCsd|}||}|j||||d\}}||}|}||}||}||}|f}|r`||f7}|S)aI 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. )rrrr)r r r rm)rprrrrZresidualr rr#r#r$rws"     zGroupViTEncoderLayer.forward)F) rxryrzrrerrrrrrrwr{r#r#rqr$rs rc@s.eZdZdZeZdZdZddZd ddZ d S) GroupViTPreTrainedModelz An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models. ZgroupvitTcCs|jj}t|tjtjfrD|jjjd|d|j dk rj|j j n&t|tj rj|j j |jj d|jj }t|tr|jjjjd|dd|jjjjd|ddnt|trB|jj }|jdd|jjd|}|jd|}tjj|jj|dtjj|jj|dtjj|jj|dtjj|jj|dnpt|tr|jj }|jjdd|jjd|}d|jjd|}tjj|jj|dtjj|jj|ddS) zInitialize the weightsrC)meanstdNrg{Gz?r}rJ)r)raZinitializer_rangerrr~rweightdataZnormal_biasZzero_rhZfill_Zinitializer_factorrrrrfrnum_hidden_layersinitrrrrrlrirr)rpmoduleZ init_rangefactorZ in_proj_stdZ out_proj_stdZfc_stdr#r#r$ _init_weightss6      z%GroupViTPreTrainedModel._init_weightsFcCst|ttfr||_dSr)rGroupViTTextEncoderGroupViTVisionEncoderr)rprrr#r#r$_set_gradient_checkpointing(sz3GroupViTPreTrainedModel._set_gradient_checkpointingN)F) rxryrzrr config_classZbase_model_prefixZsupports_gradient_checkpointingrrr#r#r#r$rs #raJ This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and behavior. Parameters: config ([`GroupViTConfig`]): Model configuration class with all the parameters of the model. Initializing with a config file does not load the weights associated with the model, only the configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights. a Args: input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide it. Indices can be obtained using [`CLIPTokenizer`]. See [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details. [What are input IDs?](../glossary#input-ids) attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: - 1 for tokens that are **not masked**, - 0 for tokens that are **masked**. [What are attention masks?](../glossary#attention-mask) position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0, config.max_position_embeddings - 1]`. [What are position IDs?](../glossary#position-ids) output_attentions (`bool`, *optional*): Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned tensors for more detail. output_hidden_states (`bool`, *optional*): Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for more detail. return_dict (`bool`, *optional*): Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. aE Args: pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`): Pixel values. Padding will be ignored by default should you provide it. Pixel values can be obtained using [`AutoImageProcessor`]. See [`CLIPImageProcessor.__call__`] for details. output_attentions (`bool`, *optional*): Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned tensors for more detail. output_hidden_states (`bool`, *optional*): Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for more detail. return_dict (`bool`, *optional*): Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. a Args: input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide it. Indices can be obtained using [`CLIPTokenizer`]. See [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details. [What are input IDs?](../glossary#input-ids) attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: - 1 for tokens that are **not masked**, - 0 for tokens that are **masked**. [What are attention masks?](../glossary#attention-mask) position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0, config.max_position_embeddings - 1]`. [What are position IDs?](../glossary#position-ids) pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`): Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See [`CLIPImageProcessor.__call__`] for details. return_loss (`bool`, *optional*): Whether or not to return the contrastive loss. output_attentions (`bool`, *optional*): Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned tensors for more detail. output_hidden_states (`bool`, *optional*): Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for more detail. return_dict (`bool`, *optional*): Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. csReZdZeddfdd Zdejeeeeeee e e fdddZ Z S) rN)rar'cs>t|_tfddttjD|_d|_ dS)Nc sFg|]>}tj|j|j||dkr:j|dnddqS)rr)rarrrr)rdepthsZnum_group_tokensZnum_output_groups)rir`r#r$rsz2GroupViTVisionEncoder.__init__..F) rdrerarrrr,rstagesrrorqr`r$res    zGroupViTVisionEncoder.__init__)routput_hidden_statesr return_dictr'c Cs|dk r |n|jj}|dk r |n|jj}|dk r4|n|jj}|rDdnd}|rPdnd}d}t|jD]P\}} |rx||f}| |||} | d}| d}|rb| ddk rb|| df}qb|r||f}|stdd|||fDSt|||dS)Nr#rrrJcss|]}|dk r|VqdSrr#rvr#r#r$rsz0GroupViTVisionEncoder.forward..last_hidden_staterrR)rarr use_return_dict enumeraterrr ) rprr rr!Zall_hidden_statesZ all_groupingsrrZstage layer_outputsr#r#r$rws0     zGroupViTVisionEncoder.forward)NNN)rxryrzrrerrrrrrr rwr{r#r#rqr$rs rc s`eZdZdZedfdd Zd eejeejee ee ee e e e fdddZ ZS) rz Transformer encoder consisting of `config.num_hidden_layers` self-attention layers. Each layer is a [`GroupViTEncoderLayer`]. Args: config: GroupViTTextConfig r`cs:t|_tfddtjD|_d|_dS)Ncsg|] }tqSr#rrr`r#r$rsz0GroupViTTextEncoder.__init__..F) rdrerarrrrrrrorqr`r$res  zGroupViTTextEncoder.__init__N)rrrr r!r'csdk r n|jj|dk r |n|jj}|dk r4|n|jj}|rDdnd}rPdnd}|} t|jD]r\} } |rx|| f}|jr|jrfdd} tj j | | | ||} n| | ||d} | d} rb|| df}qb|r|| f}|st dd | ||fDSt | ||d S) a Args: inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This is useful if you want more control over how to convert `input_ids` indices into associated vectors than the model's internal embedding lookup matrix. attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: - 1 for tokens that are **not masked**, - 0 for tokens that are **masked**. [What are attention masks?](../glossary#attention-mask) causal_attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): Causal mask for the text model. Mask values selected in `[0, 1]`: - 1 for tokens that are **not masked**, - 0 for tokens that are **masked**. [What are attention masks?](../glossary#attention-mask) output_attentions (`bool`, *optional*): Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned tensors for more detail. output_hidden_states (`bool`, *optional*): Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for more detail. return_dict (`bool`, *optional*): Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. Nr#csfdd}|S)Ncs|fSrr#)inputs)rrr#r$custom_forward szRGroupViTTextEncoder.forward..create_custom_forward..custom_forwardr#)rr*r)rr$create_custom_forward sz:GroupViTTextEncoder.forward..create_custom_forwardr+rrcss|]}|dk r|VqdSrr#r"r#r#r$r(sz.GroupViTTextEncoder.forward..r$) rarr r&r'rrrrutils checkpointrr )rprrrrr r!Zencoder_statesZall_attentionsridxZ encoder_layerr,r(r#r+r$rwsH&      zGroupViTTextEncoder.forward)NNNNN)rxryrzrrrerrrrrrr rwr{r#r#rqr$rs  r)input_ids_shaperr)past_key_values_lengthcCs|\}}tj||ft|j|d}tj|d|d}|||d|ddkd||}|dkrtj tj ||||d|gdd}|ddddddf |d|||S)zB Make causal mask used for bi-directional self-attention. r(r@rrrr)rN) rfullrr r+rZ masked_fill_rrrrr)r0rr)r1r!rrZ mask_condr#r#r$_make_causal_mask/s"  r4c sxeZdZedfdd Zeeeeedd e e j e e j e e j e e e e e e e eefdddZZS) GroupViTTextTransformerr`csHt||_|j}t||_t||_tj ||j d|_ |j |_ dSrb) rdrerarirrrencoderrrhrjfinal_layer_norm eos_token_idrrqr#r$reBs   z GroupViTTextTransformer.__init__ output_typerNrrrrr r!r'c Csn|dk r |n|jj}|dk r |n|jj}|dk r4|n|jj}|dkrLtd|}|d|d}|j||d}t||j |j d} |dk rt ||j }|j ||| |||d} | d} | | } |jdkr| tj| jd| j d|jtj| j d jdd f} n>| tj| jd| j d|jtj| j d |jkjdd f} |sZ| | f| d dSt| | | j| jd S)  Returns: NzYou have to specify input_idsr@)rrr()rrrrr r!rrJr2rrr%Z pooler_outputrrR)rarr r&rrrrr4rr)r%r6r7r8rr+rHrrMZargmaxr rrR) rprrrrr r!Z input_shaperrencoder_outputsr% pooled_outputr#r#r$rwMsV    zGroupViTTextTransformer.forward)NNNNNN)rxryrzrrerGROUPVIT_TEXT_INPUTS_DOCSTRINGrr rrrrrrrwr{r#r#rqr$r5As$   r5c seZdZeZedfdd ZejdddZddZ e e e e ed deejeejeejeeeeeeeee fd d d ZZS)GroupViTTextModelr`cs"t|t||_|dSr)rdrer5 text_model post_initrorqr#r$res  zGroupViTTextModel.__init__rcCs |jjjSrrBrrrr#r#r$get_input_embeddingssz&GroupViTTextModel.get_input_embeddingscCs||jj_dSrrD)rprr#r#r$set_input_embeddingssz&GroupViTTextModel.set_input_embeddingsr9Nr;cCs|j||||||dS)aK Returns: Examples: ```python >>> from transformers import CLIPTokenizer, GroupViTTextModel >>> tokenizer = CLIPTokenizer.from_pretrained("nvidia/groupvit-gcc-yfcc") >>> model = GroupViTTextModel.from_pretrained("nvidia/groupvit-gcc-yfcc") >>> inputs = tokenizer(["a photo of a cat", "a photo of a dog"], padding=True, return_tensors="pt") >>> outputs = model(**inputs) >>> last_hidden_state = outputs.last_hidden_state >>> pooled_output = outputs.pooler_output # pooled (EOS token) states ```rrrrr r!)rB)rprrrrr r!r#r#r$rwszGroupViTTextModel.forward)NNNNNN)rxryrzrrrerModulerErFrr@rr rrrrrrrwr{r#r#rqr$rAs*  rAc sheZdZedfdd Zeeeeedd e e j e e e e e e e eefdddZZS) GroupViTVisionTransformerr`cs@t||_|j}t||_t||_tj ||j d|_ dSrb) rdrerarirrrr6rrhrjrrrqr#r$res    z"GroupViTVisionTransformer.__init__r9N)rr rr!r'c Cs|dk r |n|jj}|dk r |n|jj}|dk r4|n|jj}|dkrLtd||}|j||||d}|d}||}|jdd}|s||f|ddSt |||j |j dS)r<Nz You have to specify pixel_values)rr rr!rrrr=) rarr r&rrr6rrr rrR) rprr rr!rr>r%r?r#r#r$rws0    z!GroupViTVisionTransformer.forward)NNNN)rxryrzrrer GROUPVIT_VISION_INPUTS_DOCSTRINGrr rrrrrrrwr{r#r#rqr$rIs   rIc s~eZdZeZdZedfdd ZedddZe e e e edd e eje ee ee eeee fd d d ZZS)GroupViTVisionModelrr`cs"t|t||_|dSr)rdrerI vision_modelrCrorqr#r$res  zGroupViTVisionModel.__init__rcCs |jjjSr)rLrrrr#r#r$rEsz(GroupViTVisionModel.get_input_embeddingsr9Nrrr r!r'cCs|j||||dS)a Returns: Examples: ```python >>> from PIL import Image >>> import requests >>> from transformers import AutoProcessor, GroupViTVisionModel >>> processor = AutoProcessor.from_pretrained("nvidia/groupvit-gcc-yfcc") >>> model = GroupViTVisionModel.from_pretrained("nvidia/groupvit-gcc-yfcc") >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" >>> image = Image.open(requests.get(url, stream=True).raw) >>> inputs = processor(images=image, return_tensors="pt") >>> outputs = model(**inputs) >>> last_hidden_state = outputs.last_hidden_state >>> pooled_output = outputs.pooler_output # pooled CLS states ```rrr r!)rL)rprrr r!r#r#r$rws zGroupViTVisionModel.forward)NNNN)rxryrzrrZmain_input_namererrErrJrr rrrrrrrwr{r#r#rqr$rK s"  rKcseZdZeZedfdd Zeedee j ee j ee j ee ee ee e j dddZ eedee j ee ee ee e j dd d Zeeeeed dee jee j ee j ee jee ee ee ee ee eeefd d dZZS) GroupViTModelr`c s6t|t|jts.tdt|jdt|jtsPtdt|jd|j}|j}|j |_ |j |_ |j |_ |j |_ t||_t||_ttj|j |j ddt|j tjddtj|j |j dd|_ttj|j |j ddt|j tjddtj|j |j dd|_tt|jj|_|dS)NzOconfig.text_config is expected to be of type GroupViTTextConfig but is of type .zSconfig.vision_config is expected to be of type GroupViTVisionConfig but is of type T)r)Zinplace) rdrer text_configrrtype vision_configrZprojection_dimZprojection_intermediate_dimriZtext_embed_dimZvision_embed_dimr5rBrIrLrrr~Z BatchNorm1dZReLUvisual_projectiontext_projectionrrrFraZlogit_scale_init_value logit_scalerC)rprarQrSrqr#r$reEs>         zGroupViTModel.__init__Nr;c Csh|dk r |n|jj}|dk r |n|jj}|dk r4|n|jj}|j||||||d}|d}||} | S)a Returns: text_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The text embeddings obtained by applying the projection layer to the pooled output of [`GroupViTTextModel`]. Examples: ```python >>> from transformers import CLIPTokenizer, GroupViTModel >>> model = GroupViTModel.from_pretrained("nvidia/groupvit-gcc-yfcc") >>> tokenizer = CLIPTokenizer.from_pretrained("nvidia/groupvit-gcc-yfcc") >>> inputs = tokenizer(["a photo of a cat", "a photo of a dog"], padding=True, return_tensors="pt") >>> text_features = model.get_text_features(**inputs) ```NrGr)rarr r&rBrU) rprrrrr r! text_outputsr?Z text_featuresr#r#r$get_text_featuresps  zGroupViTModel.get_text_featuresrMcCsd|dk r |n|jj}|dk r |n|jj}|dk r4|n|jj}|j||||d}|d}||}|S)aH Returns: image_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The image embeddings obtained by applying the projection layer to the pooled output of [`GroupViTVisionModel`]. Examples: ```python >>> from PIL import Image >>> import requests >>> from transformers import AutoProcessor, GroupViTModel >>> model = GroupViTModel.from_pretrained("nvidia/groupvit-gcc-yfcc") >>> processor = AutoProcessor.from_pretrained("nvidia/groupvit-gcc-yfcc") >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" >>> image = Image.open(requests.get(url, stream=True).raw) >>> inputs = processor(images=image, return_tensors="pt") >>> image_features = model.get_image_features(**inputs) ```NrNr)rarr r&rLrT)rprrr r!vision_outputsr?Zimage_featuresr#r#r$get_image_featuress z GroupViTModel.get_image_featuresr9) rrrr return_lossrr output_segmentationr!r'c  CsV|dk r |n|jj}|dk r |n|jj}|r0d}|dk r<|n|jj}| dk rP| n|jj} |j|||| d} |j|||||| d} | d} || } | d} || } | | j ddd} | | j ddd} |j }t | | |}|}d}|r| d}||d|jd}|r(| d }n| d }t||jd d}||j ddd}t || |}|| jdd| jddd d}||jd|jdd}t |||}||jd|jd|jd |jd }d}|rt|}| s>|dk r|||| | | | f}n||| | | | f}|dk r:|f|S|St||||| | | | d S) a Returns: Examples: ```python >>> from PIL import Image >>> import requests >>> from transformers import AutoProcessor, GroupViTModel >>> model = GroupViTModel.from_pretrained("nvidia/groupvit-gcc-yfcc") >>> processor = AutoProcessor.from_pretrained("nvidia/groupvit-gcc-yfcc") >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" >>> image = Image.open(requests.get(url, stream=True).raw) >>> inputs = processor( ... text=["a photo of a cat", "a photo of a dog"], images=image, return_tensors="pt", padding=True ... ) >>> outputs = model(**inputs) >>> logits_per_image = outputs.logits_per_image # this is the image-text similarity score >>> probs = logits_per_image.softmax(dim=1) # we can take the softmax to get the label probabilities ```NTrNrGrr@rrrrJ)rrrrrrrr)rarr\r r&rLrBrTrUZnormrVexprmatmulr/rPrHr^r[r0r)rprrrrr[rr r\r!rYrWrrrVrrZ seg_logitsZimage_group_embedsrRgroupingZlogits_per_image_groupZflatten_groupingroutputr#r#r$rws'       zGroupViTModel.forward)NNNNNN)NNNN) NNNNNNNNN)rxryrzrrrerr@rrrrrrXrJrZGROUPVIT_INPUTS_DOCSTRINGrrrrrrwr{r#r#rqr$rOAsj+.0  rO)N)rFr@)F)r)Nrcollections.abcrr dataclassesrtypingrrrrnumpyrNrZtorch.utils.checkpointrZ activationsr Zmodeling_outputsr r Zmodeling_utilsr r-r rrrrZconfiguration_groupvitrrrZ get_loggerrxloggerZ_CHECKPOINT_FOR_DOCZ&GROUPVIT_PRETRAINED_MODEL_ARCHIVE_LISTrrrMr%r-r0r?floatrrIrZr^rHr_r|rrrrrrrlrrfrrZGROUPVIT_START_DOCSTRINGr@rJrarrSizer)r4r5rArIrKrOr#r#r#r$sr       072"F!_o22  &:i ]5:4