a
    dDT                     @   s  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 dd	lmZmZmZ dd
lmZ ddlmZmZmZmZ ddlmZ eeZ dZ!dZ"g dZ#dZ$dZ%dgZ&d5e'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*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Z5d-Z6d.Z7ed/e6G d0d1 d1e5Z8ed2e6G d3d4 d4e5Z9dS )6z. PyTorch Visual Attention Network (VAN) model.    N)OrderedDict)OptionalTupleUnion)nn)BCEWithLogitsLossCrossEntropyLossMSELoss   )ACT2FN)BaseModelOutputWithNoAttention(BaseModelOutputWithPoolingAndNoAttention$ImageClassifierOutputWithNoAttention)PreTrainedModel)add_code_sample_docstringsadd_start_docstrings%add_start_docstrings_to_model_forwardlogging   )	VanConfigr   z!Visual-Attention-Network/van-base)r   i      r   ztabby, tabby cat        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ndimtorchZrandr   r   Zfloor_div)inputr   r   Z	keep_probr   Zrandom_tensoroutput r"   m/var/www/html/stable-diffusion-webui/venv/lib/python3.9/site-packages/transformers/models/van/modeling_van.py	drop_path9   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 )VanDropPathzXDrop 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#   r)   Q   s    
zVanDropPath.__init__)hidden_statesr&   c                 C   s   t || j| jS r'   )r$   r   r   )r*   r-   r"   r"   r#   forwardU   s    zVanDropPath.forward)r&   c                 C   s   d | jS )Nzp={})formatr   )r*   r"   r"   r#   
extra_reprX   s    zVanDropPath.extra_repr)N)__name__
__module____qualname____doc__r   floatr)   r   Tensorr.   strr0   __classcell__r"   r"   r+   r#   r%   N   s   r%   c                       sB   e Zd ZdZd
eeeed fddZejejddd	Z  Z	S )VanOverlappingPatchEmbedderz
    Downsamples the input using a patchify operation with a `stride` of 4 by default making adjacent windows overlap by
    half of the area. From [PVTv2: Improved Baselines with Pyramid Vision
    Transformer](https://arxiv.org/abs/2106.13797).
    r      )in_channelshidden_size
patch_sizestridec                    s4   t    tj|||||d d| _t|| _d S )N   )kernel_sizer>   padding)r(   r)   r   Conv2dconvolutionBatchNorm2dnormalization)r*   r;   r<   r=   r>   r+   r"   r#   r)   c   s
    
z$VanOverlappingPatchEmbedder.__init__)r    r&   c                 C   s   |  |}| |}|S r'   )rC   rE   )r*   r    hidden_stater"   r"   r#   r.   j   s    

z#VanOverlappingPatchEmbedder.forward)r   r:   
r1   r2   r3   r4   intr)   r   r6   r.   r8   r"   r"   r+   r#   r9   \   s   r9   c                       sD   e Zd ZdZd
eeeeed fddZej	ej	ddd	Z
  ZS )VanMlpLayerz
    MLP with depth-wise convolution, from [PVTv2: Improved Baselines with Pyramid Vision
    Transformer](https://arxiv.org/abs/2106.13797).
    gelu      ?)r;   r<   out_channels
hidden_actdropout_ratec                    sj   t    tj||dd| _tj||dd|d| _t| | _t|| _	tj||dd| _
t|| _d S )Nr   r@   r
   r@   rA   groups)r(   r)   r   rB   in_dense
depth_wiser   
activationZDropoutdropout1	out_densedropout2)r*   r;   r<   rL   rM   rN   r+   r"   r#   r)   v   s    

zVanMlpLayer.__init__rF   r&   c                 C   s@   |  |}| |}| |}| |}| |}| |}|S r'   )rR   rS   rT   rU   rV   rW   r*   rF   r"   r"   r#   r.      s    





zVanMlpLayer.forward)rJ   rK   )r1   r2   r3   r4   rH   r7   r5   r)   r   r6   r.   r8   r"   r"   r+   r#   rI   p   s   
  rI   c                       s:   e Zd ZdZed fddZejejdddZ  Z	S )VanLargeKernelAttentionz-
    Basic Large Kernel Attention (LKA).
    r<   c                    sN   t    tj||dd|d| _tj||ddd|d| _tj||dd	| _d S )
N   r?   rP   r   r
   	   )r@   ZdilationrA   rQ   r   rO   )r(   r)   r   rB   rS   depth_wise_dilated
point_wiser*   r<   r+   r"   r#   r)      s    
z VanLargeKernelAttention.__init__rX   c                 C   s"   |  |}| |}| |}|S r'   )rS   r^   r_   rY   r"   r"   r#   r.      s    


zVanLargeKernelAttention.forwardrG   r"   r"   r+   r#   rZ      s   rZ   c                       s:   e Zd ZdZed fddZejejdddZ  Z	S )VanLargeKernelAttentionLayerzV
    Computes attention using Large Kernel Attention (LKA) and attends the input.
    r[   c                    s   t    t|| _d S r'   )r(   r)   rZ   	attentionr`   r+   r"   r#   r)      s    
z%VanLargeKernelAttentionLayer.__init__rX   c                 C   s   |  |}|| }|S r'   )rb   )r*   rF   rb   Zattendedr"   r"   r#   r.      s    
z$VanLargeKernelAttentionLayer.forwardrG   r"   r"   r+   r#   ra      s   ra   c                       s>   e Zd ZdZd	eed fddZejejdddZ	  Z
S )
VanSpatialAttentionLayerz
    Van spatial attention layer composed by projection (via conv) -> act -> Large Kernel Attention (LKA) attention ->
    projection (via conv) + residual connection.
    rJ   )r<   rM   c              	      sV   t    ttdtj||ddfdt| fg| _t|| _	tj||dd| _
d S )Nconvr   rO   Zact)r(   r)   r   
Sequentialr   rB   r   pre_projectionra   attention_layerpost_projection)r*   r<   rM   r+   r"   r#   r)      s    


z!VanSpatialAttentionLayer.__init__rX   c                 C   s.   |}|  |}| |}| |}|| }|S r'   )rf   rg   rh   r*   rF   Zresidualr"   r"   r#   r.      s    


z VanSpatialAttentionLayer.forward)rJ   )r1   r2   r3   r4   rH   r7   r)   r   r6   r.   r8   r"   r"   r+   r#   rc      s   rc   c                       s>   e Zd ZdZd	eed fddZejejdddZ	  Z
S )
VanLayerScalingzT
    Scales the inputs by a learnable parameter initialized by `initial_value`.
    {Gz?)r<   initial_valuec                    s(   t    tj|t| dd| _d S )NT)Zrequires_grad)r(   r)   r   	Parameterr   Zonesweight)r*   r<   rl   r+   r"   r#   r)      s    
zVanLayerScaling.__init__rX   c                 C   s   | j dd| }|S )N)rn   Z	unsqueezerY   r"   r"   r#   r.      s    zVanLayerScaling.forward)rk   )r1   r2   r3   r4   rH   r5   r)   r   r6   r.   r8   r"   r"   r+   r#   rj      s   rj   c                       sB   e Zd ZdZd
eeeed fddZej	ej	ddd	Z
  ZS )VanLayerzv
    Van layer composed by normalization layers, large kernel attention (LKA) and a multi layer perceptron (MLP).
    r:   rK   )configr<   	mlp_ratiodrop_path_ratec                    s   t    |dkrt|nt | _t|| _t||j	| _
t||j| _t|| _t||| ||j	|j| _t||j| _d S )Nr   )r(   r)   r%   r   Identityr$   rD   pre_normomalizationrc   rM   rb   rj   Zlayer_scale_init_valueattention_scalingpost_normalizationrI   rN   mlpmlp_scaling)r*   rq   r<   rr   rs   r+   r"   r#   r)      s    
zVanLayer.__init__rX   c                 C   sl   |}|  |}| |}| |}| |}|| }|}| |}| |}| |}| |}|| }|S r'   )ru   rb   rv   r$   rw   rx   ry   ri   r"   r"   r#   r.      s    







zVanLayer.forward)r:   rK   r1   r2   r3   r4   r   rH   r5   r)   r   r6   r.   r8   r"   r"   r+   r#   rp      s     rp   c                
       sJ   e Zd ZdZd
eeeeeeeed fddZej	ej	ddd	Z
  ZS )VanStagez2
    VanStage, consisting of multiple layers.
    r:   r   )rq   r;   r<   r=   r>   depthrr   rs   c	           	         sT   t    t|||| _tj fddt|D  | _tj j	d| _
d S )Nc                    s   g | ]}t  d qS ))rr   rs   )rp   ).0_rq   rs   r<   rr   r"   r#   
<listcomp>  s   z%VanStage.__init__.<locals>.<listcomp>Zeps)r(   r)   r9   
embeddingsr   re   rangelayers	LayerNormlayer_norm_epsrE   )	r*   rq   r;   r<   r=   r>   r|   rr   rs   r+   r   r#   r)     s    
zVanStage.__init__rX   c                 C   s^   |  |}| |}|j\}}}}|ddd}| |}|||||dddd}|S )Nr?   r   r   r
   )r   r   r   flattenZ	transposerE   viewZpermute)r*   rF   Z
batch_sizer<   heightwidthr"   r"   r#   r.   (  s    


zVanStage.forward)r:   r   rz   r"   r"   r+   r#   r{   	  s     r{   c                       sN   e Zd ZdZed fddZd
ejee	 ee	 e
eef ddd	Z  ZS )
VanEncoderz4
    VanEncoder, consisting of multiple stages.
    )rq   c                    s   t    tg | _|j}|j}|j}|j}|j	}dd t
d|jt|jD }tt||||||D ]R\}\}	}
}}}}|dk}||d  }|r|j}| jt||||	|
|||d qjd S )Nc                 S   s   g | ]}|  qS r"   )item)r}   xr"   r"   r#   r   A      z'VanEncoder.__init__.<locals>.<listcomp>r   r   )r=   r>   r|   rr   rs   )r(   r)   r   Z
ModuleListstagespatch_sizesstrideshidden_sizesdepths
mlp_ratiosr   Zlinspacers   sum	enumeratezipZnum_channelsappendr{   )r*   rq   r   r   r   r   r   Zdrop_path_ratesZ	num_stager=   r>   r<   r|   Zmlp_expantionrs   Zis_first_stager;   r+   r"   r#   r)   9  s6    
 zVanEncoder.__init__FT)rF   output_hidden_statesreturn_dictr&   c                 C   s\   |rdnd }t | jD ]\}}||}|r||f }q|sPtdd ||fD S t||dS )Nr"   c                 s   s   | ]}|d ur|V  qd S r'   r"   )r}   vr"   r"   r#   	<genexpr>f  r   z%VanEncoder.forward.<locals>.<genexpr>)last_hidden_stater-   )r   r   tupler   )r*   rF   r   r   Zall_hidden_statesr~   Zstage_moduler"   r"   r#   r.   W  s    zVanEncoder.forward)FT)r1   r2   r3   r4   r   r)   r   r6   r   boolr   r   r   r.   r8   r"   r"   r+   r#   r   4  s   !  
r   c                   @   s2   e Zd ZdZeZdZdZdZdd Z	ddd	Z
d
S )VanPreTrainedModelz
    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
    models.
    vanpixel_valuesTc                 C   s   t |tjrJtjj|j| jjd t |tjr|jdurtj	|jd nt |tj
rxtj	|jd tj	|jd n`t |tjr|jd |jd  |j }||j }|jjdtd|  |jdur|jj  dS )zInitialize the weights)ZstdNr   g      ?r   g       @)
isinstancer   LinearinitZtrunc_normal_rn   rq   Zinitializer_rangeZbiasZ	constant_r   rB   r@   rL   rQ   dataZnormal_mathsqrtZzero_)r*   moduleZfan_outr"   r"   r#   _init_weightsv  s    

z VanPreTrainedModel._init_weightsFc                 C   s   t |tr||_d S r'   )r   VanModelZgradient_checkpointing)r*   r   valuer"   r"   r#   _set_gradient_checkpointing  s    
z.VanPreTrainedModel._set_gradient_checkpointingN)F)r1   r2   r3   r4   r   config_classZbase_model_prefixZmain_input_nameZsupports_gradient_checkpointingr   r   r"   r"   r"   r#   r   k  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 ([`VanConfig`]): 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.
aF  
    Args:
        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See
            [`ConvNextImageProcessor.__call__`] for details.

        output_hidden_states (`bool`, *optional*):
            Whether or not to return the hidden states of all stages. 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.
zxThe bare VAN model outputting raw features without any specific head on top. Note, VAN does not have an embedding layer.c                	       sb   e Zd Z f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 )
r   c                    s@   t  | || _t|| _tj|jd |jd| _	| 
  d S )Nro   r   )r(   r)   rq   r   encoderr   r   r   r   Z	layernorm	post_initr*   rq   r+   r"   r#   r)     s
    
zVanModel.__init__Zvision)
checkpointoutput_typer   Zmodalityexpected_outputN)r   r   r   r&   c                 C   sx   |d ur|n| j j}|d ur |n| j j}| j|||d}|d }|jddgd}|sh||f|dd   S t|||jdS )Nr   r   r   ro   )Zdimr   )r   pooler_outputr-   )rq   r   use_return_dictr   meanr   r-   )r*   r   r   r   Zencoder_outputsr   pooled_outputr"   r"   r#   r.     s"    zVanModel.forward)NN)r1   r2   r3   r)   r   VAN_INPUTS_DOCSTRINGr   _CHECKPOINT_FOR_DOCr   _CONFIG_FOR_DOC_EXPECTED_OUTPUT_SHAPEr   r   FloatTensorr   r   r   r.   r8   r"   r"   r+   r#   r     s"   	
  
r   z
    VAN Model with an image classification head on top (a linear layer on top of the pooled features), 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 )	VanForImageClassificationc                    sJ   t  | t|| _|jdkr4t|jd |jnt | _	| 
  d S )Nr   ro   )r(   r)   r   r   
num_labelsr   r   r   rt   
classifierr   r   r+   r"   r#   r)     s
    
$z"VanForImageClassification.__init__)r   r   r   r   N)r   labelsr   r   r&   c                 C   st  |dur|n| j j}| j|||d}|r.|jn|d }| |}d}|dur4| j jdu r| j jdkrpd| 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rt }	| j jdkr|	| | }n
|	||}nP| j jdkrt }	|	|d| j j|d}n| j jdkr4t }	|	||}|sd|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   Z
regressionZsingle_label_classificationZmulti_label_classificationro   r?   )losslogitsr-   )rq   r   r   r   r   Zproblem_typer   r   r   longrH   r	   Zsqueezer   r   r   r   r-   )r*   r   r   r   r   outputsr   r   r   Zloss_fctr!   r"   r"   r#   r.     s6    


$

z!VanForImageClassification.forward)NNNN)r1   r2   r3   r)   r   r   r   _IMAGE_CLASS_CHECKPOINTr   r   _IMAGE_CLASS_EXPECTED_OUTPUTr   r   r   Z
LongTensorr   r   r   r.   r8   r"   r"   r+   r#   r     s&       
r   )r   F):r4   r   collectionsr   typingr   r   r   r   Ztorch.utils.checkpointr   Ztorch.nnr   r   r	   Zactivationsr   Zmodeling_outputsr   r   r   Zmodeling_utilsr   utilsr   r   r   r   Zconfiguration_vanr   Z
get_loggerr1   loggerr   r   r   r   r   Z!VAN_PRETRAINED_MODEL_ARCHIVE_LISTr5   r   r$   Moduler%   r9   rI   rZ   ra   rc   rj   rp   r{   r   r   ZVAN_START_DOCSTRINGr   r   r   r"   r"   r"   r#   <module>   sX   
 ++7 0