a
    dfd                    @   s  d Z ddlZddlZddlmZ ddlmZ ddlmZm	Z	m
Z
mZmZmZ ddlZddlmZ ddlm  mZ ddlZddlmZ ddlmZmZmZmZmZmZ ddlmZm Z m!Z!m"Z"m#Z#m$Z$m%Z%m&Z&m'Z'm(Z(m)Z) d	d
l*m+Z+ d	dl,m-Z-m.Z.m/Z/ d	dl0m1Z1m2Z2m3Z3 dgZ4e5e6Z7G dd dej8Z9G dd dej8Z:G dd dej8Z;G dd dej8Z<G dd dej8Z=G dd dej8Z>G dd dej8Z?d<ej8e@dddZAd=ej8e@eBd d!d"ZCd>ej8e@dd#d$ZDd?eBd&d'd(ZEd@d,d-ZFeG dAe?e@e@d.d/d0ZHd1d2 ZId3d4 ZJdBd6d7ZKdCd8d9ZLe1eLd:d;eL eLd:d;eLd<d:d5d=eLd>d:d5d?d@dAeLdBd:d5d=eLdCd:d5d?d@dAeLdDd:d5d=eLdEd:d5d?d@dAeLdFd:d5d=eLdGd:d5d?d@dAeLdHd:d5d=eLdId:d5d?d@dAeLdJd:d5d=eLdKd:d5d=eLdLd:d5d?d@dAeLdMd:dNeLdOd:d?d@dPeLdQd:d?d@dPeLdRd:d5d=eLdSd:d5d?d@dAeLdTd:d5d=eLdUd:d5d?d@dAeLdVd:d5d=eLdWd:d5d?d@dAeLddXeLddXeLddXeLddXeLdYd:dZd[eLd\d:d5dZd]eLd^d:d5dZd]eLd_d:d5dZd]eLd`d:d5dZd]eLdad:d5dZd]eLdbd:d5dZd]eLdcd:d5dZd]eLddd:d5dZd]eLded5d:dfeLdgd5d:dfeLdhd:eeddieLdjd:eeddieLdkd:eeddieLdld:eeddieLdmd:dneeddod@dpeLdqd:dneeddod@dpeLdrd:dneeddod@dpeLdsd:dneeddod@dpeLdtd:dudvdwdxdydzeLd{d:dudvdwdxd|eLd}d:dNeLd:d~dddeLd:dddeLd:d?dddeL eLd:eedeLd:eed@d?deLd:eed@ddeLd:eeddeLd:eed@d?ddeLd:eed@deLd:eed@dddeLd:eed@deLd:eed@dddeLeedeLd:eedd?ddeLd:eeddeLd:eedd?ddeLd:eed@deLd:eed@dddeLd:eedeLd:eed@deLd:eed@d?ddeLd:eed@deLd:eed@dddeLd:eed@deLdeed@dddeLd:eedeLd:eedeLd:eed@d?ddeLd:eed@deLeeddeLd:eeddeLd:eed@ddeLd:eed@ddeLeeddeLd:eeddeLd:eed@ddeLddeeddeLddeed@ddeLddeed@ddeLddeed@ddeLddeed@ddeLddeed@ddeLddeed@ddeLd:eeddeLd:eeddeLd:eed@ddeLd:deed@dddeLddddeLdd~ddeLddXeLddXeLddXeLd:deedd@deLd:deedd@ddeLd:deedd@deLd:deedd@ddeLdd5d:d~ddeLdd5d:d~ddeLdd5d:d~ddeLdd5d:d~ddeLdd5d:d~ddeLdd5d:d~ddeLdd5d:d~ddZdeLdd5d:d~ddZdeLdd5d:d~ddeLdd5d:d~ddeLdd5d:d~ddeLdd5d:d~ddZdeLdd5d:d~ddZdeLddXeLddXeLddXeLdd:dneeddeLdd:dneeddeLdd:dneedddZMdDddZNe2dEe?dddZOe2dFe?dddZPe2dGe?dddZQe2dHe?dddZRe2dIe?dddĄZSe2dJe?dddƄZTe2dKe?dddȄZUe2dLe?dddʄZVe2dMe?ddd̄ZWe2dNe?ddd΄ZXe2dOe?dddЄZYe2dPe?ddd҄ZZe2dQe?dddԄZ[e2dRe?dddքZ\e2dSe?ddd؄Z]e2dTe?dddڄZ^e2dUe?ddd܄Z_e2dVe?dddބZ`e2dWe?dddZae2dXe?dddZbe2dYe?dddZce2dZe?dddZde2d[e?dddZee2d\e?dddZfe2d]e?dddZge2d^e?dddZhe2d_e?dddZie2d`e?dddZje2dae?dddZke2dbe?dddZle2dce?dddZme2dde?dddZne2dee?dddZoe2dfe?dddZpe2dge?ddd Zqe2dhe?dddZre2die?dddZse2dje?dddZte2dke?dddZue2dle?dd	d
Zve2dme?dddZwe2dne?dddZxe2doe?dddZye2dpe?dddZze2dqe?dddZ{e2dre?dddZ|e2dse?dddZ}e2dte?dddZ~e2due?dddZe2dve?dddZe2dwe?ddd Ze2dxe?dd!d"Ze2dye?dd#d$Ze2dze?dd%d&Ze3e6d'd(d)d*d+d,d-d.d/d0d1d2d3d4d5d6d7d8d9d:d; dS ({  a   Vision Transformer (ViT) in PyTorch

A PyTorch implement of Vision Transformers as described in:

'An Image Is Worth 16 x 16 Words: Transformers for Image Recognition at Scale'
    - https://arxiv.org/abs/2010.11929

`How to train your ViT? Data, Augmentation, and Regularization in Vision Transformers`
    - https://arxiv.org/abs/2106.10270

`FlexiViT: One Model for All Patch Sizes`
    - https://arxiv.org/abs/2212.08013

The official jax code is released and available at
  * https://github.com/google-research/vision_transformer
  * https://github.com/google-research/big_vision

Acknowledgments:
  * The paper authors for releasing code and weights, thanks!
  * I fixed my class token impl based on Phil Wang's https://github.com/lucidrains/vit-pytorch
  * Simple transformer style inspired by Andrej Karpathy's https://github.com/karpathy/minGPT
  * Bert reference code checks against Huggingface Transformers and Tensorflow Bert

Hacked together by / Copyright 2020, Ross Wightman
    N)OrderedDict)partial)CallableListOptionalSequenceTupleUnion)Final)IMAGENET_DEFAULT_MEANIMAGENET_DEFAULT_STDIMAGENET_INCEPTION_MEANIMAGENET_INCEPTION_STDOPENAI_CLIP_MEANOPENAI_CLIP_STD)
PatchEmbedMlpDropPathtrunc_normal_lecun_normal_resample_patch_embedresample_abs_pos_embedRmsNormPatchDropoutuse_fused_attnSwiGLUPacked   )build_model_with_cfg)named_applycheckpoint_seqadapt_input_conv)generate_default_cfgsregister_modelregister_model_deprecationsVisionTransformerc                       sB   e Zd ZU ee ed< dddddejf fdd	Zdd Z	  Z
S )		Attention
fused_attn   F        c                    s   t    || dksJ d|| _|| | _| jd | _t | _tj||d |d| _	|rf|| jnt
 | _|r~|| jnt
 | _t|| _t||| _t|| _d S )Nr   $dim should be divisible by num_heads         bias)super__init__	num_headshead_dimscaler   r&   nnLinearqkvIdentityq_normk_normDropout	attn_dropproj	proj_drop)selfdimr0   qkv_biasqk_normr:   r<   
norm_layer	__class__ g/var/www/html/stable-diffusion-webui/venv/lib/python3.9/site-packages/timm/models/vision_transformer.pyr/   7   s    


zAttention.__init__c           
      C   s   |j \}}}| |||d| j| jddddd}|d\}}}| || | }}| j	rzt
j|||| jjd}n8|| j }||dd }	|	jdd	}	| |	}	|	| }|dd|||}| |}| |}|S )
Nr+      r   r      Z	dropout_pr>   )shaper5   reshaper0   r1   permuteZunbindr7   r8   r&   Fscaled_dot_product_attentionr:   pr2   	transposesoftmaxr;   r<   )
r=   xBNCr5   qkvattnrD   rD   rE   forwardO   s$    *



zAttention.forward)__name__
__module____qualname__r
   bool__annotations__r3   	LayerNormr/   r\   __classcell__rD   rD   rB   rE   r%   4   s   
r%   c                       s&   e Zd Zd fdd	Zdd Z  ZS )
LayerScaleh㈵>Fc                    s*   t    || _t|t| | _d S N)r.   r/   inplacer3   	ParametertorchZonesgamma)r=   r>   init_valuesrg   rB   rD   rE   r/   h   s    
zLayerScale.__init__c                 C   s   | j r|| jS || j S rf   )rg   Zmul_rj   r=   rT   rD   rD   rE   r\   m   s    zLayerScale.forward)re   F)r]   r^   r_   r/   r\   rc   rD   rD   rB   rE   rd   g   s   rd   c                
       s>   e Zd Zdddddddejejef
 fdd	Zdd Z  Z	S )	Block      @Fr(   Nc              	      s   t    ||| _t|||||||d| _|r<t||dnt | _|	dkrVt	|	nt | _
||| _||t|| |
|d| _|rt||dnt | _|	dkrt	|	nt | _d S )Nr0   r?   r@   r:   r<   rA   rk   r(   Zin_featureshidden_features	act_layerdrop)r.   r/   norm1r%   r[   rd   r3   r6   ls1r   
drop_path1norm2intmlpls2
drop_path2r=   r>   r0   	mlp_ratior?   r@   r<   r:   rk   	drop_pathrs   rA   	mlp_layerrB   rD   rE   r/   s   s,    

	

zBlock.__init__c              
   C   sD   ||  | | | | }|| | | | | }|S rf   )rw   rv   r[   ru   r|   r{   rz   rx   rl   rD   rD   rE   r\      s      zBlock.forward)
r]   r^   r_   r3   GELUrb   r   r/   r\   rc   rD   rD   rB   rE   rm   q   s   'rm   c                
       sF   e Zd Zdddddddejejef
 fdd	Zdd Zd	d
 Z	  Z
S )ResPostBlockrn   Fr(   Nc              	      s   t    || _t|||||||d| _||| _|	dkrBt|	nt | _	||t
|| |
|d| _||| _|	dkrt|	nt | _|   d S )Nro   r(   rq   )r.   r/   rk   r%   r[   ru   r   r3   r6   rw   ry   rz   rx   r|   init_weightsr}   rB   rD   rE   r/      s,    
	


zResPostBlock.__init__c                 C   s6   | j d ur2tj| jj| j  tj| jj| j  d S rf   )rk   r3   init	constant_ru   weightrx   r=   rD   rD   rE   r      s    
zResPostBlock.init_weightsc                 C   s8   ||  | | | }|| | | | }|S rf   )rw   ru   r[   r|   rx   rz   rl   rD   rD   rE   r\      s    zResPostBlock.forward)r]   r^   r_   r3   r   rb   r   r/   r   r\   rc   rD   rD   rB   rE   r      s   )r   c                
       sP   e Zd ZU dZee ed< dddddddejej	df
 fdd	Z
d	d
 Z  ZS )ParallelScalingBlockz Parallel ViT block (MLP & Attention in parallel)
    Based on:
      'Scaling Vision Transformers to 22 Billion Parameters` - https://arxiv.org/abs/2302.05442
    r&   rn   Fr(   Nc                    s~  t    || dksJ d|| _|| | _| jd | _t | _t|| }|d|  }||| _t	j
|||d| _|g|gd  | _|r| dd  | dd  n,| jdtd| dd	 t	t|| _|r|| jnt	 | _|r|| jnt	 | _t	|| _t	
||| _t	|| _|
 | _t	
||| _|d urTt||d
nt	 | _|	dkrpt|	nt	 | _d S )Nr   r)   r*   r+   r,   r?   mlp_biasF)
persistentrp   r(   ) r.   r/   r0   r1   r2   r   r&   ry   in_normr3   r4   in_projin_splitZregister_bufferZregister_parameterri   zerosrh   r   r6   r7   r8   r9   r:   attn_out_projmlp_dropmlp_actmlp_out_projrd   lsr   r   )r=   r>   r0   r~   r?   r@   r<   r:   rk   r   rs   rA   r   Zmlp_hidden_dimZin_proj_out_dimrB   rD   rE   r/      s2    


 zParallelScalingBlock.__init__c                 C   sz  |j \}}}| |}| jd urBt|| jjt| j	| jf}n
| |}tj
|| jdd\}}}}	| |||| j| jdd}| |||| j| jdd}|	||| j| jdd}	| jrtj|||	| jjd}
n8|| j }||dd }|jdd}| |}||	 }
|
dd|||}
| |
}
| |}| |}| |}| | |
| }|| }|S )NrJ   rK   r   rF   rH   rI   )rL   r   r   rO   Zlinearr   r   ri   catr?   splitr   r7   viewr0   r1   rR   r8   r&   rP   r:   rQ   r2   rS   rM   r   r   r   r   r   r   )r=   rT   rU   rV   rW   yZx_mlprX   rY   rZ   Zx_attnr[   rD   rD   rE   r\     s6    

"
""





zParallelScalingBlock.forward)r]   r^   r_   __doc__r
   r`   ra   r3   r   rb   r/   r\   rc   rD   rD   rB   rE   r      s   
.r   c                       s\   e Zd ZdZddddddddejejef fdd	Zd	d
 Z	e
jjdd Zdd Z  ZS )ParallelThingsBlockz Parallel ViT block (N parallel attention followed by N parallel MLP)
    Based on:
      `Three things everyone should know about Vision Transformers` - https://arxiv.org/abs/2203.09795
    rF   rn   FNr(   c                    s  t    || _t | _t | _t|D ]}| jt	t
d||fdt|||||	||dfd|rpt||dnt fd|
dkrt|
nt fg | jt	t
d||fd||t|| ||d	fd|rt||dnt fd|
dkrt|
nt fg q,d S )
Nnormr[   ro   r   rp   r   r(   rz   )rr   rs   rt   )r.   r/   num_parallelr3   Z
ModuleListattnsffnsrangeappend
Sequentialr   r%   rd   r6   r   ry   )r=   r>   r0   r   r~   r?   r@   rk   r<   r:   r   rs   rA   r   _rB   rD   rE   r/   ;  s<    



	


zParallelThingsBlock.__init__c                    sP    t  fdd| jD jdd   t  fdd| jD jdd   S )Nc                    s   g | ]}| qS rD   rD   .0r[   rT   rD   rE   
<listcomp>k      z4ParallelThingsBlock._forward_jit.<locals>.<listcomp>r   rK   c                    s   g | ]}| qS rD   rD   r   Zffnr   rD   rE   r   l  r   )ri   stackr   sumr   rl   rD   r   rE   _forward_jitj  s    &&z ParallelThingsBlock._forward_jitc                    s<    t  fdd| jD    t  fdd| jD    S )Nc                 3   s   | ]}| V  qd S rf   rD   r   r   rD   rE   	<genexpr>q  r   z/ParallelThingsBlock._forward.<locals>.<genexpr>c                 3   s   | ]}| V  qd S rf   rD   r   r   rD   rE   r   r  r   )r   r   r   rl   rD   r   rE   _forwardo  s    zParallelThingsBlock._forwardc                 C   s,   t j st j r| |S | |S d S rf   )ri   jitis_scripting
is_tracingr   r   rl   rD   rD   rE   r\   u  s    
zParallelThingsBlock.forward)r]   r^   r_   r   r3   r   rb   r   r/   r   ri   r   ignorer   r\   rc   rD   rD   rB   rE   r   6  s"   /
r   c                       s  e Zd ZdZddddddddd	d
ddd
ddddddddddeddeefeee	eef f eee	eef f eee
eeeeeeee eeeee eeeeeee
eee ee eed fddZd3ddZdd Zej d4ddZejjdd Zejjd5ddZejjd6ddZejjdd Zd7ed d!d"Zd#d$ Zd8ejeeef d&d'd(Zd9ejeeef eeee	eeje	ej f  d)d*d+Zd,d- Z d:ed.d/d0Z!d1d2 Z"  Z#S );r$   z Vision Transformer

    A PyTorch impl of : `An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale`
        - https://arxiv.org/abs/2010.11929
          r+     token      rn   TFNr(    )img_size
patch_sizein_chansnum_classesglobal_pool	embed_dimdepthr0   r~   r?   r@   rk   class_tokenno_embed_classpre_normfc_norm	drop_ratepos_drop_ratepatch_drop_rateproj_drop_rateattn_drop_ratedrop_path_rateweight_initembed_layerrA   rs   block_fnr   c                     s   t    |dv sJ |s&|dks&J |du r6|dkn|}pJttjdd pTtj || _|| _ | _| _	|rvdnd| _
|| _d	| _||||| d
| _| jj}|rttddnd| _|r|n|| j
 }ttd|d | _tj|d| _|dkrt|| j
d| _n
t | _|r2nt | _dd td||D tj 	
fddt|D  | _|snt | _|rnt | _ t|| _!|dkrt"| j	|nt | _#|dkr| $| dS )a  
        Args:
            img_size: Input image size.
            patch_size: Patch size.
            in_chans: Number of image input channels.
            num_classes: Mumber of classes for classification head.
            global_pool: Type of global pooling for final sequence (default: 'token').
            embed_dim: Transformer embedding dimension.
            depth: Depth of transformer.
            num_heads: Number of attention heads.
            mlp_ratio: Ratio of mlp hidden dim to embedding dim.
            qkv_bias: Enable bias for qkv projections if True.
            init_values: Layer-scale init values (layer-scale enabled if not None).
            class_token: Use class token.
            fc_norm: Pre head norm after pool (instead of before), if None, enabled when global_pool == 'avg'.
            drop_rate: Head dropout rate.
            pos_drop_rate: Position embedding dropout rate.
            attn_drop_rate: Attention dropout rate.
            drop_path_rate: Stochastic depth rate.
            weight_init: Weight initialization scheme.
            embed_layer: Patch embedding layer.
            norm_layer: Normalization layer.
            act_layer: MLP activation layer.
            block_fn: Transformer block layer.
        r   avgr   r   Nr   ư>)Zepsr   r   F)r   r   r   r   r-   {Gz?)rQ   num_prefix_tokensc                 S   s   g | ]}|  qS rD   )item)r   rT   rD   rD   rE   r     r   z.VisionTransformer.__init__.<locals>.<listcomp>c                    s0   g | ](}	
|  d qS ))r>   r0   r~   r?   r@   rk   r<   r:   r   rA   rs   r   rD   )r   irs   r   r   Zdprr   rk   r   r~   rA   r0   r   r@   r?   rD   rE   r     s   skip)%r.   r/   r   r3   rb   r   r   r   Znum_featuresr   r   r   grad_checkpointingpatch_embednum_patchesrh   ri   r   	cls_tokenZrandn	pos_embedr9   pos_dropr   
patch_dropr6   norm_preZlinspacer   r   blocksr   r   	head_dropr4   headr   ) r=   r   r   r   r   r   r   r   r0   r~   r?   r@   rk   r   r   r   r   r   r   r   r   r   r   r   r   rA   rs   r   r   Zuse_fc_normr   Z	embed_lenrB   r   rE   r/     sR    8




&
"
zVisionTransformer.__init__c                 C   sd   |dv sJ d|v r"t | j nd}t| jdd | jd urPtjj| jdd t	t
|||  d S )N)jaxZjax_nlhbmocor   Znlhbr(   r   stdr   )mathlogr   r   r   r   r3   r   normal_r   get_init_weights_vit)r=   mode	head_biasrD   rD   rE   r     s    
zVisionTransformer.init_weightsc                 C   s   t | d S rf   )init_weights_vit_timm)r=   mrD   rD   rE   _init_weights  s    zVisionTransformer._init_weightsc                 C   s   t | || d S rf   )_load_weights)r=   checkpoint_pathprefixrD   rD   rE   load_pretrained  s    z!VisionTransformer.load_pretrainedc                 C   s   h dS )N>   r   Z
dist_tokenr   rD   r   rD   rD   rE   no_weight_decay
  s    z!VisionTransformer.no_weight_decayc                 C   s   t dddgdS )Nz ^cls_token|pos_embed|patch_embed)z^blocks\.(\d+)N)z^norm)i )stemr   )dict)r=   ZcoarserD   rD   rE   group_matcher  s    zVisionTransformer.group_matcherc                 C   s
   || _ d S rf   )r   )r=   enablerD   rD   rE   set_grad_checkpointing  s    z(VisionTransformer.set_grad_checkpointingc                 C   s   | j S rf   )r   r   rD   rD   rE   get_classifier  s    z VisionTransformer.get_classifier)r   c                 C   sD   || _ |d ur |dv sJ || _|dkr6t| j|nt | _d S )Nr   r   )r   r   r3   r4   r   r6   r   )r=   r   r   rD   rD   rE   reset_classifier  s
    z"VisionTransformer.reset_classifierc                 C   s   | j r@|| j }| jd urxtj| j|jd dd|fdd}n8| jd urntj| j|jd dd|fdd}|| j }| |S )Nr   rJ   r   rK   )r   r   r   ri   r   expandrL   r   rl   rD   rD   rE   
_pos_embed$  s    

&
$
zVisionTransformer._pos_embedr   )rT   nc                 C   s   g t | j }}tt|tr*t|| |n|}| |}| |}| |}| 	|}t
| jD ]"\}}||}||v rb|| qb|S rf   )lenr   set
isinstancery   r   r   r   r   r   	enumerater   )r=   rT   r   outputsZ
num_blocksZtake_indicesr   ZblkrD   rD   rE   _intermediate_layers3  s     



z&VisionTransformer._intermediate_layers)rT   r   rM   return_class_tokenr   returnc                    s    |}|r"fdd|D }fdd|D }fdd|D }|rfjj  fdd|D }|rxtt||S t|S )zs Intermediate layer accessor (NOTE: This is a WIP experiment).
        Inspired by DINO / DINOv2 interface
        c                    s   g | ]}  |qS rD   )r   r   outr   rD   rE   r   U  r   z=VisionTransformer.get_intermediate_layers.<locals>.<listcomp>c                    s"   g | ]}|d d d j f qS )Nr   r   r  r   rD   rE   r   V  r   c                    s"   g | ]}|d d  j d f qS rf   r   r  r   rD   rE   r   W  r   c                    s:   g | ]2}| jd   d   d dd ddd qS )r   r   rJ   r+   rF   )rM   rL   rN   
contiguousr  )	grid_sizerT   rD   rE   r   [  s   )r  r   r  tuplezip)r=   rT   r   rM   r  r   r  Zclass_tokensrD   )r  r=   rT   rE   get_intermediate_layersG  s    z)VisionTransformer.get_intermediate_layersc                 C   s^   |  |}| |}| |}| |}| jrFtj sFt| j	|}n
| 	|}| 
|}|S rf   )r   r   r   r   r   ri   r   r   r   r   r   rl   rD   rD   rE   forward_featuresd  s    





z"VisionTransformer.forward_features)
pre_logitsc                 C   sd   | j r>| j dkr.|d d | jd f jddn|d d df }| |}| |}|rZ|S | |S )Nr   r   rK   r   )r   r   meanr   r   r   )r=   rT   r  rD   rD   rE   forward_headp  s
    8

zVisionTransformer.forward_headc                 C   s   |  |}| |}|S rf   )r  r  rl   rD   rD   rE   r\   w  s    

zVisionTransformer.forward)r   )r   )F)T)N)r   )r   FFF)F)$r]   r^   r_   r   r   rm   r   r	   ry   r   strfloatr`   r   r   r/   r   r   ri   r   r   r   r   r   r   r   r   r   ZTensorr   r  r  r  r  r\   rc   rD   rD   rB   rE   r$   |  s   w


 
    
r   modulenamec                 C   sJ   t | tjr4t| jdd | jdurFtj| j nt| drF| 	  dS )zE ViT weight initialization, original timm impl (for reproducibility) r   r   Nr   )
r   r3   r4   r   r   r-   r   zeros_hasattrr   r  rD   rD   rE   r   }  s    

r   r(   r  r  r   c                 C   s   t | tjrx|dr6tj| j tj| j| qtj	| j | jdurd|v rhtjj
| jddntj| j nBt | tjrt| j | jdurtj| j nt| dr|   dS )z5 ViT weight initialization, matching JAX (Flax) impl r   Nrz   r   r   r   )r   r3   r4   
startswithr   r  r   r   r-   xavier_uniform_r   ZConv2dr   r  r   r  rD   rD   rE   init_weights_vit_jax  s    

*


r  c                 C   s   t | tjr|d|v rTtdt| jjd d | jjd   }tj	| j| | ntj
| j | jdurtj| j nt| dr|   dS )zI ViT weight initialization, matching moco-v3 impl minus fixed PatchEmbed r5   g      @r   r+   r   Nr   )r   r3   r4   r   sqrtr  r   rL   r   Zuniform_r  r-   r  r  r   )r  r  valrD   rD   rE   init_weights_vit_moco  s    *

r  r   r   c                 C   s(   d| v rt t|dS d| v r tS tS d S )Nr   r  r   )r   r  r  r   )r   r   rD   rD   rE   r     s
    r   rD   bicubicFc           
      C   s2  |j d }|r>| ddd|f | d|df  }}||8 }n| ddddf | d  }}ttt|}	t|stt|gd }t|dksJ td| j  d|	|	g d|j  d| d	 |d|	|	d	dd
dd}t	j
||||dd}|ddd
dd|d |d  d	}tj||gdd} | S )a=   Rescale the grid of position embeddings when loading from state_dict.

    *DEPRECATED* This function is being deprecated in favour of resample_abs_pos_embed

    Adapted from:
        https://github.com/google-research/vision_transformer/blob/00883dd691c63a6830751563748663526e811cee/vit_jax/checkpoint.py#L224
    r   Nr   rF   zResized position embedding: z (z) to z).rJ   r+   F)sizer   	antialiasZalign_cornersrK   )rL   ry   r   r  r   _loggerinforM   rN   rO   Zinterpolateri   r   )
ZposembZ
posemb_newr   Zgs_newinterpolationr!  Zntok_newZposemb_prefixZposemb_gridZgs_oldrD   rD   rE   resize_pos_embed  s    
&
,&r%  )modelr   r   c                    s*  ddl }d@dd ||d}d}d}|sJdv r:d	}nd
v rJd}d}t| jdrV| jj}t|d }|rt|n|j}	|	jjt	|	jjj
d  | d  |	jj | d  |	jj | d  |sBt|jD ]P\}
}t|jD ]:\}}| d|
d  d|d  d}tdD ]}t|d|d  j | d|d  d  t|d|d  j | d|d  d  t|d|d  j | d|d  d  q2|jdur|jjj | d  |jjj | d  |jjj | d  qq | d }n$t	| jjjj
d  | d }|j
d d | jjjj
d d krt|| jjjj
d d ||dd!}| jjj| | jjj | d"  | jdur| j | d# dd$ |r6 | d% dd$}n | d& dd$}|j
| jj
kr|j
}t| d'drtdn
t| d(d}t|| jj|||dd)}| j| | jj | d*  | jj | d+  t| jtjrH| jjj
d | d, j
d- krH| jj | d.  | jj | d,  |rRd/nd0\}}}t| j D ]\}
}| d1|
 d}|d2| d |jj | d3  |jj | d4  |jj jt!" fd5d6d7D  |jj jt!" fd8d6d7D  |jjj  d9 #d |jjj  d:  td;D ]v}t|j$d<|d  j | d=| d>| d  t|j$d<|d  j | d=| d>| d  qf|j%j | d?| d  |j%j | d?| d  qjdS )AzV Load weights from .npz checkpoints for official Google Brain Flax implementation
    r   NTc                 S   s   | j dkrF| jd | jd   kr:| jd   kr:dkrFn n|  } |r| j dkrd| g d} n2| j dkr~| g d} n| j dkr| ddg} t| S )NrG   r   r   rF   )r+   rF   r   r   r+   )rF   r   r   )ndimrL   flattenrR   ri   Z
from_numpy)wtrD   rD   rE   _n2p  s    >


z_load_weights.<locals>._n2pbilinearFzopt/target/embedding/kernelzopt/target/zparams/embedding/kernelzparams/backboner   r   zconv_root/kernelzgn_root/scalezgn_root/biasblockz/unit/r+   conv/kernelr   Zgnz/scale/biaszconv_proj/kernelzgn_proj/scalezgn_proj/biaszembedding/kernelrI   r$  r!  verbosezembedding/biasclsr*  Zpos_embeddingz(Transformer/posembed_input/pos_embeddingr   r   Znew_sizer   r$  r!  r4  zTransformer/encoder_norm/scalezTransformer/encoder_norm/biasz	head/biasrJ   zhead/kernel)r   r   r   )r   r+   rF   zTransformer/encoderblock_ZMultiHeadDotProductAttention_zLayerNorm_0/scalezLayerNorm_0/biasc                    s.   g | ]&}  | d  dd djqS )r1  Fr6  r   )r(  Tr   r   r+  Z
mha_prefixr)  rD   rE   r   3  s   z!_load_weights.<locals>.<listcomp>)querykeyvaluec                    s,   g | ]$}  | d  dd dqS )r2  Fr6  rJ   )rM   r9  r:  rD   rE   r   5  s   z
out/kernelzout/biasrF   ZfcZ	MlpBlock_z/Dense_Z
LayerNorm_)T)&numpyloadr  r   r-  r   r0  r   Zcopy_r    rL   r   r-   r   Zstagesr   r   getattrZ
downsampler;   r   r   r   r   r  r   r   r3   r4   childrenru   r[   r5   ri   r   r(  rz   rx   )r&  r   r   npr$  r!  Z
big_visionr-  Z	stem_onlyr   r   Zstagejr.  ZbprZembed_conv_wZpos_embed_wZ	old_shaper   Zmha_subZb_subZln1_subZblock_prefixrD   r:  rE   r     s    

,448$$4

$8<"r   c                 C   s   i }g d}|   D ]\}}|ds(q|D ]}||d |d }q,|dkrtd}|dd}t|jd |d< np|dkrd	}|dd}nR|d
kr|d}|jd |jjd krt	||jt
|drdn
t
|dd|jj}|||< q|S )N))visual.r   )Zconv1patch_embed.proj)Zpositional_embeddingr   )ztransformer.resblocks.zblocks.)Zln_prer   )Zln_postr   )Zln_r   )Zin_proj_zqkv.)Zout_projr;   )zmlp.c_fczmlp.fc1)z
mlp.c_projzmlp.fc2rE  r   r   r;   zhead.weightz	head.biasZclass_embeddingr   r   r   r   )itemsr  replacerR   ri   r   rL   Z	unsqueezer   r%  r@  r   r  )
state_dictr&  out_dictZswapsrY   rZ   sprD   rD   rE   _convert_openai_clip@  s2    


rL  c                 C   st   dd l }i }|  D ]Z\}}|dkr(qn>|d|rH|||dd< qn|d|rf|||dd< q|||< q|S )	Nr   
mask_tokenz(blocks\.(\d+)\.mlp\.w12\.(?:weight|bias)Zw12Zfc1z'blocks\.(\d+)\.mlp\.w3\.(?:weight|bias)Zw3Zfc2)rerG  matchrH  )rI  r&  rN  rJ  rY   rZ   rD   rD   rE   _convert_dinov2b  s    
rP  Tc              	   C   sp  ddl }i }| d| } | d| } d| v r6t| |S d| v rHt| |S |  D ]\}}d|v r|jjjj\}	}
}}t	|jdk r|jjjj\}	}
}}|
|	d	||}|jd	 |ks|jd
 |krt|||f||dd}n|dkr6|jd |jjd kr6t|ddrdn
t|dd}t||jj|||dd}n,|rVd|v rV|dd|}nd|v rbqP|||< qP|S )zJ convert patch embedding weight from manual patchify + linear proj to convr   Nr&  rI  zvisual.class_embeddingrM  zpatch_embed.proj.weightrG   rJ   rI   Tr3  r   r   r   Fr   r7  Zgamma_zgamma_([0-9])z
ls\1.gammar  )rN  getrL  rP  rG  r   r;   r   rL   r   rM   r   r   r@  r   r  sub)rI  r&  Zadapt_layer_scaler$  r!  rN  rJ  rY   rZ   OIHWr   rD   rD   rE   checkpoint_filter_fnr  sL    

"

rW  c                 K   s    | ddd dddt tddd|S )	Nr   )r+   r   r   g?r  TrF  r   )urlr   
input_sizeZ	pool_sizecrop_pctr$  Zfixed_input_sizer  r   Z
first_conv
classifier)r   r   )rX  kwargsrD   rD   rE   _cfg  s    r]  ztimm/)	hf_hub_idzhttps://storage.googleapis.com/vit_models/augreg/Ti_16-i21k-300ep-lr_0.001-aug_none-wd_0.03-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.03-res_224.npz)rX  r^  custom_loadzhttps://storage.googleapis.com/vit_models/augreg/Ti_16-i21k-300ep-lr_0.001-aug_none-wd_0.03-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.03-res_384.npz)r+     r`        ?)rX  r^  r_  rY  rZ  zhttps://storage.googleapis.com/vit_models/augreg/S_32-i21k-300ep-lr_0.001-aug_light1-wd_0.03-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.03-res_224.npzzhttps://storage.googleapis.com/vit_models/augreg/S_32-i21k-300ep-lr_0.001-aug_light1-wd_0.03-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.03-res_384.npzzhttps://storage.googleapis.com/vit_models/augreg/S_16-i21k-300ep-lr_0.001-aug_light1-wd_0.03-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.03-res_224.npzzhttps://storage.googleapis.com/vit_models/augreg/S_16-i21k-300ep-lr_0.001-aug_light1-wd_0.03-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.03-res_384.npzzhttps://storage.googleapis.com/vit_models/augreg/B_32-i21k-300ep-lr_0.001-aug_medium1-wd_0.03-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.03-res_224.npzzhttps://storage.googleapis.com/vit_models/augreg/B_32-i21k-300ep-lr_0.001-aug_light1-wd_0.1-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.03-res_384.npzzhttps://storage.googleapis.com/vit_models/augreg/B_16-i21k-300ep-lr_0.001-aug_medium1-wd_0.1-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.01-res_224.npzzhttps://storage.googleapis.com/vit_models/augreg/B_16-i21k-300ep-lr_0.001-aug_medium1-wd_0.1-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.01-res_384.npzzhttps://storage.googleapis.com/vit_models/augreg/B_8-i21k-300ep-lr_0.001-aug_medium1-wd_0.1-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.01-res_224.npzzhttps://storage.googleapis.com/vit_models/augreg/L_16-i21k-300ep-lr_0.001-aug_medium1-wd_0.1-do_0.1-sd_0.1--imagenet2012-steps_20k-lr_0.01-res_224.npzzhttps://storage.googleapis.com/vit_models/augreg/L_16-i21k-300ep-lr_0.001-aug_medium1-wd_0.1-do_0.1-sd_0.1--imagenet2012-steps_20k-lr_0.01-res_384.npzzohttps://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_base_p16_224-80ecf9dd.pth)rX  r^  zohttps://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_base_p16_384-83fb41ba.pth)rX  r^  rY  rZ  zphttps://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_large_p32_384-9b920ba8.pthzhttps://storage.googleapis.com/vit_models/augreg/S_16-i1k-300ep-lr_0.001-aug_medium2-wd_0.1-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.01-res_224.npzzhttps://storage.googleapis.com/vit_models/augreg/S_16-i1k-300ep-lr_0.001-aug_medium2-wd_0.1-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.01-res_384.npzzhttps://storage.googleapis.com/vit_models/augreg/B_32-i1k-300ep-lr_0.001-aug_medium2-wd_0.1-do_0.1-sd_0.1--imagenet2012-steps_20k-lr_0.01-res_224.npzzhttps://storage.googleapis.com/vit_models/augreg/B_32-i1k-300ep-lr_0.001-aug_medium2-wd_0.1-do_0.1-sd_0.1--imagenet2012-steps_20k-lr_0.01-res_384.npzzhttps://storage.googleapis.com/vit_models/augreg/B_16-i1k-300ep-lr_0.001-aug_strong2-wd_0.1-do_0.1-sd_0.1--imagenet2012-steps_20k-lr_0.01-res_224.npzzhttps://storage.googleapis.com/vit_models/augreg/B_16-i1k-300ep-lr_0.001-aug_strong2-wd_0.1-do_0.1-sd_0.1--imagenet2012-steps_20k-lr_0.01-res_384.npz)rX  zzhttps://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_large_patch32_224_in21k-9046d2e7.pthiSU  )rX  r^  r   zBhttps://storage.googleapis.com/vit_models/imagenet21k/ViT-H_14.npz)rX  r^  r_  r   zmhttps://storage.googleapis.com/vit_models/augreg/Ti_16-i21k-300ep-lr_0.001-aug_none-wd_0.03-do_0.0-sd_0.0.npzznhttps://storage.googleapis.com/vit_models/augreg/S_32-i21k-300ep-lr_0.001-aug_light1-wd_0.03-do_0.0-sd_0.0.npzznhttps://storage.googleapis.com/vit_models/augreg/S_16-i21k-300ep-lr_0.001-aug_light1-wd_0.03-do_0.0-sd_0.0.npzzohttps://storage.googleapis.com/vit_models/augreg/B_32-i21k-300ep-lr_0.001-aug_medium1-wd_0.03-do_0.0-sd_0.0.npzznhttps://storage.googleapis.com/vit_models/augreg/B_16-i21k-300ep-lr_0.001-aug_medium1-wd_0.1-do_0.0-sd_0.0.npzzmhttps://storage.googleapis.com/vit_models/augreg/B_8-i21k-300ep-lr_0.001-aug_medium1-wd_0.1-do_0.0-sd_0.0.npzznhttps://storage.googleapis.com/vit_models/augreg/L_16-i21k-300ep-lr_0.001-aug_medium1-wd_0.1-do_0.1-sd_0.1.npzz:https://storage.googleapis.com/vit_models/sam/ViT-B_32.npz)rX  r_  r^  z:https://storage.googleapis.com/vit_models/sam/ViT-B_16.npzz[https://dl.fbaipublicfiles.com/dino/dino_deitsmall16_pretrain/dino_deitsmall16_pretrain.pth)rX  r^  r  r   r   zYhttps://dl.fbaipublicfiles.com/dino/dino_deitsmall8_pretrain/dino_deitsmall8_pretrain.pthzWhttps://dl.fbaipublicfiles.com/dino/dino_vitbase16_pretrain/dino_vitbase16_pretrain.pthzUhttps://dl.fbaipublicfiles.com/dino/dino_vitbase8_pretrain/dino_vitbase8_pretrain.pthzNhttps://dl.fbaipublicfiles.com/dinov2/dinov2_vits14/dinov2_vits14_pretrain.pthzcc-by-nc-4.0)r+     rb  )rX  r^  licenser  r   r   rY  rZ  zNhttps://dl.fbaipublicfiles.com/dinov2/dinov2_vitb14/dinov2_vitb14_pretrain.pthzNhttps://dl.fbaipublicfiles.com/dinov2/dinov2_vitl14/dinov2_vitl14_pretrain.pthzNhttps://dl.fbaipublicfiles.com/dinov2/dinov2_vitg14/dinov2_vitg14_pretrain.pthz}https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-tresnet/vit_base_patch16_224_in21k_miil-887286df.pth)r(   r(   r(   )ra  ra  ra  g      ?r,  i+  )rX  r^  r  r   rZ  r$  r   zhttps://github.com/rwightman/pytorch-image-models/releases/download/v0.1-tresnet/vit_base_patch16_224_1k_miil_84_4-2deb18e3.pth)rX  r^  r  r   rZ  r$  z}https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-tpu-weights/vit_base_patch16_rpn_224-sw-3b07e89d.pth)r+      rd  gffffff?i-.  )r^  rY  rZ  r   )r+      re  )r^  rY  rZ  Zsquash)r^  rY  rZ  	crop_mode)r^  r  r   )r^  r  r   rZ  rY  )r+     rg  )r^  r  r   rZ  )r^  r  r   rZ  rY  rf  )r+   P  rh  )r  r   )r  r   r   )r^  r  r   r   )r^  r  r   rZ  r   z%laion/CLIP-ViT-B-32-laion2B-s34B-b79Kzopen_clip_pytorch_model.bin   )r^  hf_hub_filenamer  r   r   z%laion/CLIP-ViT-B-16-laion2B-s34B-b88K)r^  rj  r  r   rZ  r   z%laion/CLIP-ViT-L-14-laion2B-s32B-b82Kr   z)laion/CLIP-ViT-L-14-DataComp.XL-s13B-b90Kz%laion/CLIP-ViT-H-14-laion2B-s32B-b79K   z%laion/CLIP-ViT-g-14-laion2B-s12B-b42Kz(laion/CLIP-ViT-bigG-14-laion2B-39B-b160k   )r^  rj  r  r   rZ  rY  r   )rX  rY  rZ  Zmit)r+      rm  )r^  rc  r  r   rY  rZ  )r^  rc  r  r   rY  rZ  rf  zEhttps://storage.googleapis.com/big_vision/flexivit/flexivit_s_i1k.npz)rX  r_  r^  rY  rZ  zKhttps://storage.googleapis.com/big_vision/flexivit/flexivit_s_i1k_600ep.npzzKhttps://storage.googleapis.com/big_vision/flexivit/flexivit_s_i1k_300ep.npzzEhttps://storage.googleapis.com/big_vision/flexivit/flexivit_b_i1k.npzzKhttps://storage.googleapis.com/big_vision/flexivit/flexivit_b_i1k_600ep.npzzKhttps://storage.googleapis.com/big_vision/flexivit/flexivit_b_i1k_300ep.npzzMhttps://storage.googleapis.com/big_vision/flexivit/flexivit_b_i21k_1000ep.npz)rX  r_  r^  rY  rZ  r   zLhttps://storage.googleapis.com/big_vision/flexivit/flexivit_b_i21k_300ep.npzzEhttps://storage.googleapis.com/big_vision/flexivit/flexivit_l_i1k.npzzKhttps://storage.googleapis.com/big_vision/flexivit/flexivit_l_i1k_600ep.npzzKhttps://storage.googleapis.com/big_vision/flexivit/flexivit_l_i1k_300ep.npzzIhttps://storage.googleapis.com/big_vision/flexivit/vit_b16_i21k_300ep.npzzIhttps://storage.googleapis.com/big_vision/flexivit/vit_b30_i21k_300ep.npzzEhttps://dl.fbaipublicfiles.com/mae/pretrain/mae_pretrain_vit_base.pth)rX  r^  rc  r  r   r   zFhttps://dl.fbaipublicfiles.com/mae/pretrain/mae_pretrain_vit_large.pthzEhttps://dl.fbaipublicfiles.com/mae/pretrain/mae_pretrain_vit_huge.pth)z*vit_base_patch16_224.augreg2_in21k_ft_in1kz*vit_base_patch16_384.augreg2_in21k_ft_in1kz)vit_base_patch8_224.augreg2_in21k_ft_in1kz)vit_tiny_patch16_224.augreg_in21k_ft_in1kz)vit_tiny_patch16_384.augreg_in21k_ft_in1kz*vit_small_patch32_224.augreg_in21k_ft_in1kz*vit_small_patch32_384.augreg_in21k_ft_in1kz*vit_small_patch16_224.augreg_in21k_ft_in1kz*vit_small_patch16_384.augreg_in21k_ft_in1kz)vit_base_patch32_224.augreg_in21k_ft_in1kz)vit_base_patch32_384.augreg_in21k_ft_in1kz)vit_base_patch16_224.augreg_in21k_ft_in1kz)vit_base_patch16_384.augreg_in21k_ft_in1kz(vit_base_patch8_224.augreg_in21k_ft_in1kz*vit_large_patch16_224.augreg_in21k_ft_in1kz*vit_large_patch16_384.augreg_in21k_ft_in1kz'vit_base_patch16_224.orig_in21k_ft_in1kz'vit_base_patch16_384.orig_in21k_ft_in1kz(vit_large_patch32_384.orig_in21k_ft_in1kz!vit_small_patch16_224.augreg_in1kz!vit_small_patch16_384.augreg_in1kz vit_base_patch32_224.augreg_in1kz vit_base_patch32_384.augreg_in1kz vit_base_patch16_224.augreg_in1kz vit_base_patch16_384.augreg_in1kzvit_large_patch14_224.untrainedzvit_huge_patch14_224.untrainedzvit_giant_patch14_224.untrainedz"vit_gigantic_patch14_224.untrained vit_large_patch32_224.orig_in21kvit_huge_patch14_224.orig_in21k!vit_tiny_patch16_224.augreg_in21k"vit_small_patch32_224.augreg_in21k"vit_small_patch16_224.augreg_in21k!vit_base_patch32_224.augreg_in21k!vit_base_patch16_224.augreg_in21k vit_base_patch8_224.augreg_in21k"vit_large_patch16_224.augreg_in21kzvit_base_patch32_224.sam_in1kzvit_base_patch16_224.sam_in1kvit_small_patch16_224.dinovit_small_patch8_224.dinovit_base_patch16_224.dinovit_base_patch8_224.dinoz vit_small_patch14_dinov2.lvd142mzvit_base_patch14_dinov2.lvd142mz vit_large_patch14_dinov2.lvd142mz vit_giant_patch14_dinov2.lvd142mvit_base_patch16_224_miil.in21kz'vit_base_patch16_224_miil.in21k_ft_in1kz vit_base_patch16_rpn_224.sw_in1kz#vit_medium_patch16_gap_240.sw_in12kz+vit_medium_patch16_gap_256.sw_in12k_ft_in1kz+vit_medium_patch16_gap_384.sw_in12k_ft_in1kvit_base_patch16_gap_224z/vit_base_patch32_clip_224.laion2b_ft_in12k_in1kz/vit_base_patch32_clip_384.laion2b_ft_in12k_in1kz/vit_base_patch32_clip_448.laion2b_ft_in12k_in1kz/vit_base_patch16_clip_224.laion2b_ft_in12k_in1kz/vit_base_patch16_clip_384.laion2b_ft_in12k_in1kz0vit_large_patch14_clip_224.laion2b_ft_in12k_in1kz0vit_large_patch14_clip_336.laion2b_ft_in12k_in1kz/vit_huge_patch14_clip_224.laion2b_ft_in12k_in1kz/vit_huge_patch14_clip_336.laion2b_ft_in12k_in1kz.vit_base_patch32_clip_224.openai_ft_in12k_in1kz.vit_base_patch32_clip_384.openai_ft_in12k_in1kz.vit_base_patch16_clip_224.openai_ft_in12k_in1kz.vit_base_patch16_clip_384.openai_ft_in12k_in1kz/vit_large_patch14_clip_224.openai_ft_in12k_in1kz/vit_large_patch14_clip_336.openai_ft_in12k_in1kz)vit_base_patch32_clip_224.laion2b_ft_in1kz)vit_base_patch16_clip_224.laion2b_ft_in1kz)vit_base_patch16_clip_384.laion2b_ft_in1kz*vit_large_patch14_clip_224.laion2b_ft_in1kz*vit_large_patch14_clip_336.laion2b_ft_in1kz)vit_huge_patch14_clip_224.laion2b_ft_in1kz)vit_huge_patch14_clip_336.laion2b_ft_in1kz(vit_base_patch32_clip_224.openai_ft_in1kz(vit_base_patch16_clip_224.openai_ft_in1kz(vit_base_patch16_clip_384.openai_ft_in1kz)vit_large_patch14_clip_224.openai_ft_in1kz*vit_base_patch32_clip_224.laion2b_ft_in12kz*vit_base_patch16_clip_224.laion2b_ft_in12kz+vit_large_patch14_clip_224.laion2b_ft_in12kz*vit_huge_patch14_clip_224.laion2b_ft_in12kz)vit_base_patch32_clip_224.openai_ft_in12kz)vit_base_patch16_clip_224.openai_ft_in12kz*vit_large_patch14_clip_224.openai_ft_in12k!vit_base_patch32_clip_224.laion2bz!vit_base_patch16_clip_224.laion2b"vit_large_patch14_clip_224.laion2bz%vit_large_patch14_clip_224.datacompxl!vit_huge_patch14_clip_224.laion2b"vit_giant_patch14_clip_224.laion2bz%vit_gigantic_patch14_clip_224.laion2bz vit_base_patch32_clip_224.openaiz vit_base_patch16_clip_224.openaiz!vit_large_patch14_clip_224.openaiz!vit_large_patch14_clip_336.openaiz#vit_base_patch32_plus_256.untrainedz#vit_base_patch16_plus_240.untrainedz$vit_small_patch16_36x1_224.untrainedz$vit_small_patch16_18x2_224.untrainedz#vit_base_patch16_18x2_224.untrainedz)eva_large_patch14_196.in22k_ft_in22k_in1kz)eva_large_patch14_336.in22k_ft_in22k_in1kz#eva_large_patch14_196.in22k_ft_in1kz#eva_large_patch14_336.in22k_ft_in1kzflexivit_small.1200ep_in1kzflexivit_small.600ep_in1kzflexivit_small.300ep_in1kzflexivit_base.1200ep_in1kzflexivit_base.600ep_in1kzflexivit_base.300ep_in1kzflexivit_base.1000ep_in21kzflexivit_base.300ep_in21kzflexivit_large.1200ep_in1kzflexivit_large.600ep_in1kzflexivit_large.300ep_in1kzflexivit_base.patch16_in21kzflexivit_base.patch30_in21kz!vit_base_patch16_xp_224.untrainedz"vit_large_patch14_xp_224.untrainedz!vit_huge_patch14_xp_224.untrainedzvit_base_patch16_224.maezvit_large_patch16_224.maezvit_huge_patch14_224.maec                 K   sH   | dd rtdd| v r,ttddd}nt}tt| |fd|i|S )NZfeatures_onlyz<features_only not implemented for Vision Transformer models.Zflexir,  F)r$  r!  Zpretrained_filter_fn)rQ  RuntimeErrorr   rW  r   r$   )variant
pretrainedr\  Z
_filter_fnrD   rD   rE   _create_vision_transformer  s    r  )r  c                 K   s2   t ddddd}tdd| it |fi |}|S )	z ViT-Tiny (Vit-Ti/16)
    r      r   r+   r   r   r   r0   vit_tiny_patch16_224r  )r  r   r  r  r\  Z
model_argsr&  rD   rD   rE   r    s    r  c                 K   s2   t ddddd}tdd| it |fi |}|S )	z% ViT-Tiny (Vit-Ti/16) @ 384x384.
    r   r  r   r+   r  vit_tiny_patch16_384r  )r  r  r  rD   rD   rE   r    s    r  c                 K   s2   t ddddd}tdd| it |fi |}|S )	z ViT-Small (ViT-S/32)
        r`  r      r  vit_small_patch32_224r  )r  r  r  rD   rD   rE   r    s    r  c                 K   s2   t ddddd}tdd| it |fi |}|S )	z& ViT-Small (ViT-S/32) at 384x384.
    r  r`  r   r  r  vit_small_patch32_384r  )r  r  r  rD   rD   rE   r    s    r  c                 K   s2   t ddddd}tdd| it |fi |}|S )	 ViT-Small (ViT-S/16)
    r   r`  r   r  r  vit_small_patch16_224r  )r  r  r  rD   rD   rE   r    s    r  c                 K   s2   t ddddd}tdd| it |fi |}|S )	r  r   r`  r   r  r  vit_small_patch16_384r  )r  r  r  rD   rD   rE   r    s    r  c                 K   s2   t ddddd}tdd| it |fi |}|S )	z ViT-Small (ViT-S/8)
    r'   r`  r   r  r  vit_small_patch8_224r  )r  r  r  rD   rD   rE   r    s    r  c                 K   s2   t ddddd}tdd| it |fi |}|S )z ViT-Base (ViT-B/32) from original paper (https://arxiv.org/abs/2010.11929).
    ImageNet-1k weights fine-tuned from in21k, source https://github.com/google-research/vision_transformer.
    r  r   r   r  vit_base_patch32_224r  )r  r  r  rD   rD   rE   r    s    r  c                 K   s2   t ddddd}tdd| it |fi |}|S )z ViT-Base model (ViT-B/32) from original paper (https://arxiv.org/abs/2010.11929).
    ImageNet-1k weights fine-tuned from in21k @ 384x384, source https://github.com/google-research/vision_transformer.
    r  r   r   r  vit_base_patch32_384r  )r  r  r  rD   rD   rE   r    s    r  c                 K   s2   t ddddd}tdd| it |fi |}|S )z ViT-Base (ViT-B/16) from original paper (https://arxiv.org/abs/2010.11929).
    ImageNet-1k weights fine-tuned from in21k @ 224x224, source https://github.com/google-research/vision_transformer.
    r   r   r   r  vit_base_patch16_224r  )r  r  r  rD   rD   rE   r    s    r  c                 K   s2   t ddddd}tdd| it |fi |}|S )z ViT-Base model (ViT-B/16) from original paper (https://arxiv.org/abs/2010.11929).
    ImageNet-1k weights fine-tuned from in21k @ 384x384, source https://github.com/google-research/vision_transformer.
    r   r   r   r  vit_base_patch16_384r  )r  r  r  rD   rD   rE   r    s    r  c                 K   s2   t ddddd}tdd| it |fi |}|S )z ViT-Base (ViT-B/8) from original paper (https://arxiv.org/abs/2010.11929).
    ImageNet-1k weights fine-tuned from in21k @ 224x224, source https://github.com/google-research/vision_transformer.
    r'   r   r   r  vit_base_patch8_224r  )r  r  r  rD   rD   rE   r    s    r  c                 K   s2   t ddddd}tdd| it |fi |}|S )	zo ViT-Large model (ViT-L/32) from original paper (https://arxiv.org/abs/2010.11929). No pretrained weights.
    r  rk     r   r  vit_large_patch32_224r  )r  r  r  rD   rD   rE   r  $  s    r  c                 K   s2   t ddddd}tdd| it |fi |}|S )	z ViT-Large model (ViT-L/32) from original paper (https://arxiv.org/abs/2010.11929).
    ImageNet-1k weights fine-tuned from in21k @ 384x384, source https://github.com/google-research/vision_transformer.
    r  rk  r  r   r  vit_large_patch32_384r  )r  r  r  rD   rD   rE   r  -  s    r  c                 K   s2   t ddddd}tdd| it |fi |}|S )z ViT-Large model (ViT-L/16) from original paper (https://arxiv.org/abs/2010.11929).
    ImageNet-1k weights fine-tuned from in21k @ 224x224, source https://github.com/google-research/vision_transformer.
    r   rk  r  r  vit_large_patch16_224r  )r  r  r  rD   rD   rE   r  7  s    r  c                 K   s2   t ddddd}tdd| it |fi |}|S )z ViT-Large model (ViT-L/16) from original paper (https://arxiv.org/abs/2010.11929).
    ImageNet-1k weights fine-tuned from in21k @ 384x384, source https://github.com/google-research/vision_transformer.
    r   rk  r  r  vit_large_patch16_384r  )r  r  r  rD   rD   rE   r  A  s    r  c                 K   s2   t ddddd}tdd| it |fi |}|S )	z  ViT-Large model (ViT-L/14)
       rk  r  r   r  vit_large_patch14_224r  )r  r  r  rD   rD   rE   r  K  s    r  c                 K   s2   t ddddd}tdd| it |fi |}|S )	zW ViT-Huge model (ViT-H/14) from original paper (https://arxiv.org/abs/2010.11929).
    r  rl  r  r   r  vit_huge_patch14_224r  )r  r  r  rD   rD   rE   r  T  s    r  c                 K   s4   t dddddd}td	d| it |fi |}|S )
zq ViT-Giant (little-g) model (ViT-g/14) from `Scaling Vision Transformers` - https://arxiv.org/abs/2106.04560
    r    tE]t@(   r   r   r   r~   r   r0   vit_giant_patch14_224r  )r  r  r  rD   rD   rE   r  ]  s    r  c                 K   s4   t dddddd}td	d| it |fi |}|S )
zq ViT-Gigantic (big-G) model (ViT-G/14) from `Scaling Vision Transformers` - https://arxiv.org/abs/2106.04560
    r    ;;@0   r   r  vit_gigantic_patch14_224r  )r  r  r  rD   rD   rE   r  f  s     r  c                 K   s4   t dddddd}tdd| it |fi |}|S )	z ViT-Base (ViT-B/16) from original paper (https://arxiv.org/abs/2010.11929).
    Weights taken from: https://github.com/Alibaba-MIIL/ImageNet21K
    r   r   r   F)r   r   r   r0   r?   vit_base_patch16_224_miilr  )r  r  r  rD   rD   rE   r  p  s     r  c                 K   s<   t dddddddddd	}tdd
| it |fi |}|S )zB ViT-Medium (ViT-M/16) w/o class token, w/ avg-pool @ 240x240
    r   ri  r   r'   Fr   r   	r   r   r   r0   r   r   r?   rk   r   vit_medium_patch16_gap_240r  )r  r  r  rD   rD   rE   r  {  s    
 r  c                 K   s<   t dddddddddd	}tdd
| it |fi |}|S )zB ViT-Medium (ViT-M/16) w/o class token, w/ avg-pool @ 256x256
    r   ri  r   r'   Fr   r   r  vit_medium_patch16_gap_256r  )r  r  r  rD   rD   rE   r    s    
 r  c                 K   s<   t dddddddddd	}tdd
| it |fi |}|S )zB ViT-Medium (ViT-M/16) w/o class token, w/ avg-pool @ 384x384
    r   ri  r   r'   Fr   r   r  vit_medium_patch16_gap_384r  )r  r  r  rD   rD   rE   r    s    
 r  c              	   K   s8   t dddddddd}td	d| it |fi |}|S )
z@ ViT-Base (ViT-B/16) w/o class token, w/ avg-pool @ 256x256
    r   r   r   Fr   )r   r   r   r0   r   r   r   r|  r  )r|  r  r  rD   rD   rE   r|    s     r|  c                 K   s8   t dddddtjd}tdd| it |fi |}|S )	z) ViT-B/32 CLIP image tower @ 224x224
    r  r   r   Tr   r   r   r0   r   rA   vit_base_patch32_clip_224r  )r  r   r3   rb   r  r  rD   rD   rE   r    s     r  c                 K   s8   t dddddtjd}tdd| it |fi |}|S )	z) ViT-B/32 CLIP image tower @ 384x384
    r  r   r   Tr  vit_base_patch32_clip_384r  )r  r  r  rD   rD   rE   r    s     r  c                 K   s8   t dddddtjd}tdd| it |fi |}|S )	z) ViT-B/32 CLIP image tower @ 448x448
    r  r   r   Tr  vit_base_patch32_clip_448r  )r  r  r  rD   rD   rE   r    s     r  c                 K   s8   t dddddtjd}tdd| it |fi |}|S )	z ViT-B/16 CLIP image tower
    r   r   r   Tr  vit_base_patch16_clip_224r  )r  r  r  rD   rD   rE   r    s     r  c                 K   s8   t dddddtjd}tdd| it |fi |}|S )	z) ViT-B/16 CLIP image tower @ 384x384
    r   r   r   Tr  vit_base_patch16_clip_384r  )r  r  r  rD   rD   rE   r    s     r  c                 K   s8   t dddddtjd}td	d| it |fi |}|S )
z1 ViT-Large model (ViT-L/14) CLIP image tower
    r  rk  r  r   Tr  vit_large_patch14_clip_224r  )r  r  r  rD   rD   rE   r    s     r  c                 K   s8   t dddddtjd}td	d| it |fi |}|S )
z; ViT-Large model (ViT-L/14) CLIP image tower @ 336x336
    r  rk  r  r   Tr  vit_large_patch14_clip_336r  )r  r  r  rD   rD   rE   r    s     r  c                 K   s8   t dddddtjd}td	d| it |fi |}|S )
z1 ViT-Huge model (ViT-H/14) CLIP image tower.
    r  rl  r  r   Tr  vit_huge_patch14_clip_224r  )r  r  r  rD   rD   rE   r    s     r  c                 K   s8   t dddddtjd}td	d| it |fi |}|S )
z: ViT-Huge model (ViT-H/14) CLIP image tower @ 336x336
    r  rl  r  r   Tr  vit_huge_patch14_clip_336r  )r  r  r  rD   rD   rE   r    s     r  c              	   K   s:   t ddddddtjd}td
d	| it |fi |}|S )z ViT-Giant (little-g) model (ViT-g/14) from `Scaling Vision Transformers` - https://arxiv.org/abs/2106.04560
    Pretrained weights from CLIP image tower.
    r  r  r  r  r   Tr   r   r~   r   r0   r   rA   vit_giant_patch14_clip_224r  )r  r  r  rD   rD   rE   r    s     r  c              	   K   s:   t ddddddtjd}td
d	| it |fi |}|S )z ViT-bigG model (ViT-G/14) from `Scaling Vision Transformers` - https://arxiv.org/abs/2106.04560
    Pretrained weights from CLIP image tower.
    r  r  r  r  r   Tr  vit_gigantic_patch14_clip_224r  )r  r  r  rD   rD   rE   r    s     r  c                 K   s4   t dddddd}td	d| it |fi |}|S )
z ViT-Base (ViT-B/32+)
    r    r   r  re   r   r   r   r0   rk   vit_base_patch32_plus_256r  )r  r  r  rD   rD   rE   r     s     r  c                 K   s4   t dddddd}td	d| it |fi |}|S )
z ViT-Base (ViT-B/16+)
    r   r  r   r  re   r  vit_base_patch16_plus_240r  )r  r  r  rD   rD   rE   r  *  s     r  c                 K   s<   t dddddddtdd	}td
d	| it |fi |}|S )z/ ViT-Base (ViT-B/16) w/ residual post-norm
    r   r   r   Fre   r   )	r   r   r   r0   r?   rk   r   r   r   vit_base_patch16_rpn_224r  )r  )r   r   r  r  rD   rD   rE   r  4  s     r  c                 K   s4   t dddddd}td	d| it |fi |}|S )
a   ViT-Base w/ LayerScale + 36 x 1 (36 block serial) config. Experimental, may remove.
    Based on `Three things everyone should know about Vision Transformers` - https://arxiv.org/abs/2203.09795
    Paper focuses on 24x2 + 48x1 for 'Small' width but those are extremely slow.
    r   r`  $   r  re   r  vit_small_patch16_36x1_224r  )r  r  r  rD   rD   rE   r  @  s     r  c                 K   s6   t dddddtd}td	d| it |fi |}|S )
a   ViT-Small w/ LayerScale + 18 x 2 (36 block parallel) config. Experimental, may remove.
    Based on `Three things everyone should know about Vision Transformers` - https://arxiv.org/abs/2203.09795
    Paper focuses on 24x2 + 48x1 for 'Small' width but those are extremely slow.
    r   r`     r  re   r   r   r   r0   rk   r   vit_small_patch16_18x2_224r  )r  r   r   r  r  rD   rD   rE   r  L  s     r  c                 K   s6   t dddddtd}td	d| it |fi |}|S )
z ViT-Base w/ LayerScale + 18 x 2 (36 block parallel) config. Experimental, may remove.
    Based on `Three things everyone should know about Vision Transformers` - https://arxiv.org/abs/2203.09795
    r   r   r  r   re   r  vit_base_patch16_18x2_224r  )r  r  r  rD   rD   rE   r  Y  s     r  c                 K   s4   t dddddd}td	d| it |fi |}|S )
zG EVA-large model https://arxiv.org/abs/2211.07636 /via MAE MIM pretrainr  rk  r  r   r   r   r   r   r0   r   eva_large_patch14_196r  )r  r  r  rD   rD   rE   r  e  s     r  c                 K   s4   t dddddd}td	d| it |fi |}|S )
zF EVA-large model https://arxiv.org/abs/2211.07636 via MAE MIM pretrainr  rk  r  r   r   r  eva_large_patch14_336r  )r  r  r  rD   rD   rE   r  n  s    r  c                 K   s4   t dddddd}td	d| it |fi |}|S )
z FlexiViT-Small
    r   r`  r   r  Tr   r   r   r0   r   flexivit_smallr  )r  r  r  rD   rD   rE   r  v  s    r  c                 K   s4   t dddddd}tdd| it |fi |}|S )	z FlexiViT-Base
    r   r   r   Tr  flexivit_baser  )r  r  r  rD   rD   rE   r    s    r  c                 K   s4   t dddddd}tdd| it |fi |}|S )	z FlexiViT-Large
    r   rk  r  Tr  flexivit_larger  )r  r  r  rD   rD   rE   r    s    r  c                 K   s>   t ddddddttddd
}td	d| it |fi |}|S )
H ViT-Large model (ViT-L/14) w/ parallel blocks and qk norm enabled.
    r   r   r   TF
r   r   r   r0   r   r   rA   r   r?   r@   vit_base_patch16_xp_224r  )r  r   r   r   r  r  rD   rD   rE   r    s     r  c                 K   s>   t ddddddttddd
}td
d	| it |fi |}|S )r  r  rk  r  r   TFr  vit_large_patch14_xp_224r  )r  r  r  rD   rD   rE   r    s     r  c                 K   s>   t ddddddttddd
}td
d	| it |fi |}|S )zG ViT-Huge model (ViT-H/14) w/ parallel blocks and qk norm enabled.
    r  rl  r  r   TFr  vit_huge_patch14_xp_224r  )r  r  r  rD   rD   rE   r    s     r  c                 K   s6   t ddddddd}td
d	| it |fi |}|S )z ViT-S/14 for DINOv2
    r  r`  r   r  ra  rb  r   r   r   r0   rk   r   vit_small_patch14_dinov2r  )r  r  r  rD   rD   rE   r    s     r  c                 K   s6   t ddddddd}td	d| it |fi |}|S )
z ViT-B/14 for DINOv2
    r  r   r   ra  rb  r  vit_base_patch14_dinov2r  )r  r  r  rD   rD   rE   r    s     r  c                 K   s6   t ddddddd}td
d	| it |fi |}|S )z ViT-L/14 for DINOv2
    r  rk  r  r   ra  rb  r  vit_large_patch14_dinov2r  )r  r  r  rD   rD   rE   r    s     r  c                 K   s>   t ddddddtdtjd	}tdd
| it |fi |}|S )z ViT-G/14 for DINOv2
    r  i   r  r  ra  gh˹WU@rb  )	r   r   r   r0   rk   r~   r   r   rs   vit_giant_patch14_dinov2r  )r  )r   r   r3   ZSiLUr  r  rD   rD   rE   r    s    


 r  rp  rq  rr  rs  rt  ru  rn  rv  ro  zvit_base_patch32_224.samzvit_base_patch16_224.samrw  rx  ry  rz  r{  r}  r~  r  r  )Zvit_tiny_patch16_224_in21kZvit_small_patch32_224_in21kZvit_small_patch16_224_in21kZvit_base_patch32_224_in21kZvit_base_patch16_224_in21kZvit_base_patch8_224_in21kZvit_large_patch32_224_in21kZvit_large_patch16_224_in21kZvit_huge_patch14_224_in21kZvit_base_patch32_224_samZvit_base_patch16_224_samZvit_small_patch16_224_dinoZvit_small_patch8_224_dinoZvit_base_patch16_224_dinoZvit_base_patch8_224_dinoZvit_base_patch16_224_miil_in21kZ!vit_base_patch32_224_clip_laion2bZ"vit_large_patch14_224_clip_laion2bZ!vit_huge_patch14_224_clip_laion2bZ"vit_giant_patch14_224_clip_laion2b)r   )r   r(   )r   )r   r(   )r   rD   r  F)r   )Fr  T)r   )F)F)F)F)F)F)F)F)F)F)F)F)F)F)F)F)F)F)F)F)F)F)F)F)F)F)F)F)F)F)F)F)F)F)F)F)F)F)F)F)F)F)F)F)F)F)F)F)F)F)F)F)F)F)F)r   loggingr   collectionsr   	functoolsr   typingr   r   r   r   r   r	   ri   Ztorch.nnr3   Ztorch.nn.functionalZ
functionalrO   Ztorch.utils.checkpointZ	torch.jitr
   Z	timm.datar   r   r   r   r   r   Ztimm.layersr   r   r   r   r   r   r   r   r   r   r   Z_builderr   Z_manipulater   r   r    	_registryr!   r"   r#   __all__	getLoggerr]   r"  Moduler%   rd   rm   r   r   r   r$   r  r   r  r  r  r   r%  Zno_gradr   rL  rP  rW  r]  Zdefault_cfgsr  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r|  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  rD   rD   rD   rE   <module>   s    4
3
/7_F  
   
!m"  
7
      n									




								