a
    dp                     @   sH  d Z ddlZddlmZ ddlmZmZmZ ddl	Z	ddl
Z	ddl	mZ ddlmZmZmZ ddlmZmZmZ dd	lmZmZ dd
lmZmZmZ ddlmZ ddlmZ ee Z!dZ"dZ#g dZ$dZ%dZ&g dZ'eG dd deZ(d@e)e*dddZ+G dd dej,Z-G dd dej,Z.G dd dej,Z/G d d! d!ej,Z0G d"d# d#ej,Z1G d$d% d%ej,Z2G d&d' d'ej,Z3G d(d) d)ej,Z4G d*d+ d+ej,Z5G d,d- d-ej,Z6G d.d/ d/ej,Z7G d0d1 d1ej,Z8G d2d3 d3ej,Z9G d4d5 d5ej,Z:G d6d7 d7eZ;d8Z<d9Z=ed:e<G d;d< d<e;Z>ed=e<G d>d? d?e;Z?dS )Az PyTorch CvT model.    N)	dataclass)OptionalTupleUnion)nn)BCEWithLogitsLossCrossEntropyLossMSELoss   )add_code_sample_docstringsadd_start_docstrings%add_start_docstrings_to_model_forward)$ImageClassifierOutputWithNoAttentionModelOutput)PreTrainedModel find_pruneable_heads_and_indicesprune_linear_layer)logging   )	CvtConfigr   microsoft/cvt-13)r   i     r   ztabby, tabby cat)r   zmicrosoft/cvt-13-384zmicrosoft/cvt-13-384-22kzmicrosoft/cvt-21zmicrosoft/cvt-21-384zmicrosoft/cvt-21-384-22kc                   @   sD   e Zd ZU dZdZejed< dZejed< dZ	e
eej  ed< dS )BaseModelOutputWithCLSTokena  
    Base class for model's outputs, with potential hidden states and attentions.

    Args:
        last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
            Sequence of hidden-states at the output of the last layer of the model.
        cls_token_value (`torch.FloatTensor` of shape `(batch_size, 1, hidden_size)`):
            Classification token at the output of the last layer of the model.
        hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
            Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of
            shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of the model at the output of each layer
            plus the initial embedding outputs.
    Nlast_hidden_statecls_token_valuehidden_states)__name__
__module____qualname____doc__r   torchZFloatTensor__annotations__r   r   r   r    r"   r"   m/var/www/html/stable-diffusion-webui/venv/lib/python3.9/site-packages/transformers/models/cvt/modeling_cvt.pyr   ;   s   
r           F)	drop_probtrainingc                 C   sd   |dks|s| S d| }| j d fd| jd   }|tj|| j| jd }|  | || }|S )aF  
    Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).

    Comment by Ross Wightman: This is the same as the DropConnect impl I created for EfficientNet, etc networks,
    however, the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...
    See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for changing the
    layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use 'survival rate' as the
    argument.
    r$   r   r   )r   )dtypedevice)shapendimr    Zrandr'   r(   Zfloor_div)inputr%   r&   Z	keep_probr)   Zrandom_tensoroutputr"   r"   r#   	drop_pathQ   s    
r.   c                       sP   e Zd ZdZdee dd fddZejejdddZ	e
d	d
dZ  ZS )CvtDropPathzXDrop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).N)r%   returnc                    s   t    || _d S N)super__init__r%   )selfr%   	__class__r"   r#   r3   i   s    
zCvtDropPath.__init__)r   r0   c                 C   s   t || j| jS r1   )r.   r%   r&   )r4   r   r"   r"   r#   forwardm   s    zCvtDropPath.forward)r0   c                 C   s   d | jS )Nzp={})formatr%   )r4   r"   r"   r#   
extra_reprp   s    zCvtDropPath.extra_repr)N)r   r   r   r   r   floatr3   r    Tensorr7   strr9   __classcell__r"   r"   r5   r#   r/   f   s   r/   c                       s(   e Zd ZdZ fddZdd Z  ZS )CvtEmbeddingsz'
    Construct the CvT embeddings.
    c                    s.   t    t|||||d| _t|| _d S )N)
patch_sizenum_channels	embed_dimstridepadding)r2   r3   CvtConvEmbeddingsconvolution_embeddingsr   Dropoutdropout)r4   r?   r@   rA   rB   rC   dropout_rater5   r"   r#   r3   y   s
    

zCvtEmbeddings.__init__c                 C   s   |  |}| |}|S r1   )rE   rG   )r4   pixel_valueshidden_stater"   r"   r#   r7      s    

zCvtEmbeddings.forwardr   r   r   r   r3   r7   r=   r"   r"   r5   r#   r>   t   s   r>   c                       s(   e Zd ZdZ fddZdd Z  ZS )rD   z"
    Image to Conv Embedding.
    c                    sP   t    t|tjjr|n||f}|| _tj|||||d| _	t
|| _d S )N)kernel_sizerB   rC   )r2   r3   
isinstancecollectionsabcIterabler?   r   Conv2d
projection	LayerNormnormalization)r4   r?   r@   rA   rB   rC   r5   r"   r#   r3      s
    
zCvtConvEmbeddings.__init__c                 C   sf   |  |}|j\}}}}|| }||||ddd}| jrH| |}|ddd||||}|S Nr      r   )rR   r)   viewpermuterT   )r4   rI   
batch_sizer@   heightwidthhidden_sizer"   r"   r#   r7      s    

zCvtConvEmbeddings.forwardrK   r"   r"   r5   r#   rD      s   rD   c                       s$   e Zd Z fddZdd Z  ZS )CvtSelfAttentionConvProjectionc              	      s4   t    tj|||||d|d| _t|| _d S )NF)rL   rC   rB   biasgroups)r2   r3   r   rQ   convolutionZBatchNorm2drT   )r4   rA   rL   rC   rB   r5   r"   r#   r3      s    
	z'CvtSelfAttentionConvProjection.__init__c                 C   s   |  |}| |}|S r1   )r`   rT   r4   rJ   r"   r"   r#   r7      s    

z&CvtSelfAttentionConvProjection.forwardr   r   r   r3   r7   r=   r"   r"   r5   r#   r]      s   r]   c                   @   s   e Zd Zdd ZdS ) CvtSelfAttentionLinearProjectionc                 C   s2   |j \}}}}|| }||||ddd}|S rU   )r)   rW   rX   )r4   rJ   rY   r@   rZ   r[   r\   r"   r"   r#   r7      s    z(CvtSelfAttentionLinearProjection.forwardN)r   r   r   r7   r"   r"   r"   r#   rc      s   rc   c                       s&   e Zd Zd fdd	Zdd Z  ZS )CvtSelfAttentionProjectiondw_bnc                    s.   t    |dkr"t||||| _t | _d S )Nre   )r2   r3   r]   convolution_projectionrc   linear_projection)r4   rA   rL   rC   rB   projection_methodr5   r"   r#   r3      s    
z#CvtSelfAttentionProjection.__init__c                 C   s   |  |}| |}|S r1   )rf   rg   ra   r"   r"   r#   r7      s    

z"CvtSelfAttentionProjection.forward)re   rb   r"   r"   r5   r#   rd      s   rd   c                       s.   e Zd Zd fdd	Zdd Zdd Z  ZS )	CvtSelfAttentionTc                    s   t    |d | _|| _|| _|| _t|||||dkr<dn|d| _t|||||d| _t|||||d| _	t
j|||	d| _t
j|||	d| _t
j|||	d| _t
|
| _d S )Ng      avgZlinear)rh   )r^   )r2   r3   scalewith_cls_tokenrA   	num_headsrd   convolution_projection_queryconvolution_projection_keyconvolution_projection_valuer   Linearprojection_queryprojection_keyprojection_valuerF   rG   )r4   rm   rA   rL   	padding_q
padding_kvstride_q	stride_kvqkv_projection_methodqkv_biasattention_drop_raterl   kwargsr5   r"   r#   r3      s,    



zCvtSelfAttention.__init__c                 C   s6   |j \}}}| j| j }|||| j|ddddS )Nr   rV   r   r
   )r)   rA   rm   rW   rX   )r4   rJ   rY   r\   _head_dimr"   r"   r#   "rearrange_for_multi_head_attention   s    z3CvtSelfAttention.rearrange_for_multi_head_attentionc                 C   sT  | j r t|d|| gd\}}|j\}}}|ddd||||}| |}| |}	| |}
| j rtj	||	fdd}	tj	||fdd}tj	||
fdd}
| j
| j }| | |	}	| | |}| | |
}
td|	|g| j }tjjj|dd}| |}td||
g}|j\}}}}|dddd ||| j| }|S )	Nr   r   rV   Zdimzbhlk,bhtk->bhltzbhlt,bhtv->bhlvr
   )rl   r    splitr)   rX   rW   ro   rn   rp   catrA   rm   r   rr   rs   rt   Zeinsumrk   r   Z
functionalZsoftmaxrG   
contiguous)r4   rJ   rZ   r[   	cls_tokenrY   r\   r@   keyqueryvaluer~   Zattention_scoreZattention_probscontextr}   r"   r"   r#   r7      s,    



$zCvtSelfAttention.forward)T)r   r   r   r3   r   r7   r=   r"   r"   r5   r#   ri      s    )ri   c                       s(   e Zd ZdZ fddZdd Z  ZS )CvtSelfOutputz
    The residual connection is defined in CvtLayer instead of here (as is the case with other models), due to the
    layernorm applied before each block.
    c                    s(   t    t||| _t|| _d S r1   )r2   r3   r   rq   denserF   rG   )r4   rA   	drop_rater5   r"   r#   r3      s    
zCvtSelfOutput.__init__c                 C   s   |  |}| |}|S r1   r   rG   r4   rJ   Zinput_tensorr"   r"   r#   r7   %  s    

zCvtSelfOutput.forwardrK   r"   r"   r5   r#   r     s   r   c                       s.   e Zd Zd fdd	Zdd Zdd Z  ZS )	CvtAttentionTc                    s@   t    t|||||||||	|
|| _t||| _t | _d S r1   )r2   r3   ri   	attentionr   r-   setpruned_heads)r4   rm   rA   rL   ru   rv   rw   rx   ry   rz   r{   r   rl   r5   r"   r#   r3   ,  s     
zCvtAttention.__init__c                 C   s   t |dkrd S t|| jj| jj| j\}}t| jj|| j_t| jj|| j_t| jj	|| j_	t| j
j|dd| j
_| jjt | | j_| jj| jj | j_| j|| _d S )Nr   r   r   )lenr   r   Znum_attention_headsZattention_head_sizer   r   r   r   r   r-   r   Zall_head_sizeunion)r4   headsindexr"   r"   r#   prune_headsL  s    zCvtAttention.prune_headsc                 C   s   |  |||}| ||}|S r1   )r   r-   )r4   rJ   rZ   r[   Zself_outputattention_outputr"   r"   r#   r7   ^  s    zCvtAttention.forward)T)r   r   r   r3   r   r7   r=   r"   r"   r5   r#   r   +  s     r   c                       s$   e Zd Z fddZdd Z  ZS )CvtIntermediatec                    s.   t    t|t|| | _t | _d S r1   )r2   r3   r   rq   intr   ZGELU
activation)r4   rA   	mlp_ratior5   r"   r#   r3   e  s    
zCvtIntermediate.__init__c                 C   s   |  |}| |}|S r1   )r   r   ra   r"   r"   r#   r7   j  s    

zCvtIntermediate.forwardrb   r"   r"   r5   r#   r   d  s   r   c                       s$   e Zd Z fddZdd Z  ZS )	CvtOutputc                    s0   t    tt|| || _t|| _d S r1   )r2   r3   r   rq   r   r   rF   rG   )r4   rA   r   r   r5   r"   r#   r3   q  s    
zCvtOutput.__init__c                 C   s    |  |}| |}|| }|S r1   r   r   r"   r"   r#   r7   v  s    

zCvtOutput.forwardrb   r"   r"   r5   r#   r   p  s   r   c                       s*   e Zd ZdZd fdd	Zdd Z  ZS )CvtLayerzb
    CvtLayer composed by attention layers, normalization and multi-layer perceptrons (mlps).
    Tc                    s|   t    t|||||||||	|
||| _t||| _t|||| _|dkrVt|dnt	
 | _t	|| _t	|| _d S )Nr$   )r%   )r2   r3   r   r   r   intermediater   r-   r/   r   Identityr.   rS   layernorm_beforelayernorm_after)r4   rm   rA   rL   ru   rv   rw   rx   ry   rz   r{   r   r   drop_path_raterl   r5   r"   r#   r3     s(    
zCvtLayer.__init__c                 C   sX   |  | |||}|}| |}|| }| |}| |}| ||}| |}|S r1   )r   r   r.   r   r   r-   )r4   rJ   rZ   r[   Zself_attention_outputr   Zlayer_outputr"   r"   r#   r7     s    



zCvtLayer.forward)TrK   r"   r"   r5   r#   r   }  s    'r   c                       s$   e Zd Z fddZdd Z  ZS )CvtStagec                    s   t     _|_jjj rBttddjj	d _t
 jj  jj jdkrh jn j	jd   j	j  jj  jj d_dd td jj  j| D tj fddt jj D  _d S )Nr   r   r   )r?   rB   r@   rA   rC   rH   c                 S   s   g | ]}|  qS r"   )item).0xr"   r"   r#   
<listcomp>      z%CvtStage.__init__.<locals>.<listcomp>c                    s   g | ]}t  jj  jj  jj  jj  jj  jj  jj  j	j  j
j  jj  jj j  jj  jj d qS ))rm   rA   rL   ru   rv   rx   rw   ry   rz   r{   r   r   r   rl   )r   rm   stagerA   Z
kernel_qkvru   rv   rx   rw   ry   rz   r{   r   r   r   )r   r}   configZdrop_path_ratesr4   r"   r#   r     s"   












)r2   r3   r   r   r   r   	Parameterr    ZrandnrA   r>   Zpatch_sizesZpatch_strider@   Zpatch_paddingr   	embeddingZlinspacer   depthZ
Sequentialrangelayers)r4   r   r   r5   r   r#   r3     s&    





	&zCvtStage.__init__c           	      C   s   d }|  |}|j\}}}}||||| ddd}| jj| j rh| j|dd}tj	||fdd}| j
D ]}||||}|}qn| jj| j rt|d|| gd\}}|ddd||||}||fS )Nr   rV   r   r   r   )r   r)   rW   rX   r   r   r   expandr    r   r   r   )	r4   rJ   r   rY   r@   rZ   r[   layerZlayer_outputsr"   r"   r#   r7     s    

zCvtStage.forwardrb   r"   r"   r5   r#   r     s   (r   c                       s&   e Zd Z fddZdddZ  ZS )
CvtEncoderc                    sF   t    || _tg | _tt|jD ]}| j	t
|| q*d S r1   )r2   r3   r   r   Z
ModuleListstagesr   r   r   appendr   )r4   r   Z	stage_idxr5   r"   r#   r3     s
    
zCvtEncoder.__init__FTc           	      C   sl   |rdnd }|}d }t | jD ]"\}}||\}}|r||f }q|s^tdd |||fD S t|||dS )Nr"   c                 s   s   | ]}|d ur|V  qd S r1   r"   )r   vr"   r"   r#   	<genexpr>  r   z%CvtEncoder.forward.<locals>.<genexpr>r   r   r   )	enumerater   tupler   )	r4   rI   output_hidden_statesreturn_dictZall_hidden_statesrJ   r   r}   Zstage_moduler"   r"   r#   r7     s    zCvtEncoder.forward)FTrb   r"   r"   r5   r#   r     s   r   c                   @   s$   e Zd ZdZeZdZdZdd ZdS )CvtPreTrainedModelz
    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
    models.
    cvtrI   c                 C   s   t |tjtjfrHtjj|jjd| jj	d|j_|j
dur|j
j  nlt |tjrp|j
j  |jjd nDt |tr| jj|j rtjjtdd| jjd d| jj	d|j_dS )zInitialize the weightsr$   )meanZstdNg      ?r   r   )rM   r   rq   rQ   initZtrunc_normal_Zweightdatar   Zinitializer_ranger^   Zzero_rS   Zfill_r   r   r   r    zerosrA   )r4   moduler"   r"   r#   _init_weights"  s    

z CvtPreTrainedModel._init_weightsN)	r   r   r   r   r   config_classZbase_model_prefixZmain_input_namer   r"   r"   r"   r#   r     s
   r   aE  
    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 ([`CvtConfig`]): 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.
aE  
    Args:
        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See [`CvtImageProcessor.__call__`]
            for details.
        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 [`~file_utils.ModelOutput`] instead of a plain tuple.
z]The bare Cvt Model transformer outputting raw hidden-states without any specific head on top.c                	       sl   e Zd Zd fdd	Zdd Zeeeee	e
deddeej ee ee eee	f d	d
dZ  ZS )CvtModelTc                    s(   t  | || _t|| _|   d S r1   )r2   r3   r   r   encoder	post_init)r4   r   add_pooling_layerr5   r"   r#   r3   O  s    
zCvtModel.__init__c                 C   s*   |  D ]\}}| jj| j| qdS )z
        Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
        class PreTrainedModel
        N)itemsr   r   r   r   )r4   Zheads_to_pruner   r   r"   r"   r#   _prune_headsU  s    zCvtModel._prune_headsZvision)
checkpointoutput_typer   Zmodalityexpected_outputN)rI   r   r   r0   c                 C   sx   |d ur|n| j j}|d ur |n| j j}|d u r8td| j|||d}|d }|sf|f|dd   S t||j|jdS )Nz You have to specify pixel_valuesr   r   r   r   r   )r   r   use_return_dict
ValueErrorr   r   r   r   )r4   rI   r   r   Zencoder_outputssequence_outputr"   r"   r#   r7   ]  s$    zCvtModel.forward)T)NNN)r   r   r   r3   r   r   CVT_INPUTS_DOCSTRINGr   _CHECKPOINT_FOR_DOCr   _CONFIG_FOR_DOC_EXPECTED_OUTPUT_SHAPEr   r    r;   boolr   r   r7   r=   r"   r"   r5   r#   r   J  s&   	   
r   z
    Cvt Model transformer with an image classification head on top (a linear layer on top of the final hidden state of
    the [CLS] token) e.g. for ImageNet.
    c                
       sh   e Zd Z fddZeeeeee	e
ddeej eej ee ee eeef dddZ  ZS )	CvtForImageClassificationc                    sh   t  | |j| _t|dd| _t|jd | _|jdkrRt	|jd |jnt
 | _|   d S )NF)r   r   r   )r2   r3   
num_labelsr   r   r   rS   rA   	layernormrq   r   
classifierr   )r4   r   r5   r"   r#   r3     s    $z"CvtForImageClassification.__init__)r   r   r   r   N)rI   labelsr   r   r0   c                 C   s  |dur|n| j j}| j|||d}|d }|d }| j jd rL| |}n4|j\}}	}
}|||	|
| ddd}| |}|jdd}| 	|}d}|dur| j j
du r| j jdkrd| j _
n6| j jdkr|jtjks|jtjkrd	| j _
nd
| j _
| j j
dkr>t }| j jdkr2|| | }n
|||}nP| j j
d	krpt }||d| j j|d}n| j j
d
krt }|||}|s|f|dd  }|dur|f| S |S t|||jdS )a  
        labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
            Labels for computing the image classification/regression loss. Indices should be in `[0, ...,
            config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
            `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
        Nr   r   r   r   rV   r   Z
regressionZsingle_label_classificationZmulti_label_classification)losslogitsr   )r   r   r   r   r   r)   rW   rX   r   r   Zproblem_typer   r'   r    longr   r	   Zsqueezer   r   r   r   )r4   rI   r   r   r   outputsr   r   rY   r@   rZ   r[   Zsequence_output_meanr   r   Zloss_fctr-   r"   r"   r#   r7     sL    



$

z!CvtForImageClassification.forward)NNNN)r   r   r   r3   r   r   r   _IMAGE_CLASS_CHECKPOINTr   r   _IMAGE_CLASS_EXPECTED_OUTPUTr   r    r;   r   r   r   r7   r=   r"   r"   r5   r#   r     s&       
r   )r$   F)@r   collections.abcrN   Zdataclassesr   typingr   r   r   r    Ztorch.utils.checkpointr   Ztorch.nnr   r   r	   Z
file_utilsr   r   r   Zmodeling_outputsr   r   Zmodeling_utilsr   r   r   utilsr   Zconfiguration_cvtr   Z
get_loggerr   loggerr   r   r   r   r   Z!CVT_PRETRAINED_MODEL_ARCHIVE_LISTr   r:   r   r.   Moduler/   r>   rD   r]   rc   rd   ri   r   r   r   r   r   r   r   r   ZCVT_START_DOCSTRINGr   r   r   r"   r"   r"   r#   <module>   sb   
	Q9B=6