a
    dʊ                     @   sP  d Z ddlZddlmZ ddlmZ ddlmZmZ ddl	m
Z
mZmZmZ ddlZddlm  mZ ddlmZ ddlmZmZ dd	lmZmZmZmZmZ d
dlmZ d
dlmZ d
dl m!Z!m"Z"m#Z# ddgZ$eG dd dZ%dd Z&G dd dej'Z(edZee) e*eej+eej+ f dddZ,ed[e)eej+ eej+ee) f dddZ-eej+ej+e*ee) ee) ej+ej+ddd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'Z3d(d) Z4e5e%d*d+e%d,d+e%d-d+e%d.d/d0d1d2e%d,dd3e%d-dd3e%d.d/d0ddd4e%d5d6d7ddd4d8Z6d\d9d:Z7d]d<d=Z8e#e8d>d?e8d@d?e8dAd?e8dBd?e8d;d?e8dCdDdEe8dFdDdEe8dGdDdEdHZ9e!d^e3dIdJdKZ:e!d_e3dIdLdMZ;e!d`e3dIdNdOZ<e!dae3dIdPdQZ=e!dbe3dIdRdSZ>e!dce3dIdTdUZ?e!dde3dIdVdWZ@e!dee3dIdXdYZAdS )fa   Multi-Scale Vision Transformer v2

@inproceedings{li2021improved,
  title={MViTv2: Improved multiscale vision transformers for classification and detection},
  author={Li, Yanghao and Wu, Chao-Yuan and Fan, Haoqi and Mangalam, Karttikeya and Xiong, Bo and Malik, Jitendra and Feichtenhofer, Christoph},
  booktitle={CVPR},
  year={2022}
}

Code adapted from original Apache 2.0 licensed impl at https://github.com/facebookresearch/mvit
Original copyright below.

Modifications and timm support by / Copyright 2022, Ross Wightman
    N)OrderedDict)	dataclass)partialreduce)UnionListTupleOptional)nnIMAGENET_DEFAULT_MEANIMAGENET_DEFAULT_STD)MlpDropPathtrunc_normal_tf_get_norm_layer	to_2tuple   )build_model_with_cfg)register_notrace_function)register_modelregister_model_deprecationsgenerate_default_cfgsMultiScaleVitMultiScaleVitCfgc                   @   s  e Zd ZU dZeedf ed< dZeeeedf f ed< dZ	eeeedf f ed< dZ
eed	< d
Zeed< dZeed< dZeed< d
Zeed< d
Zeed< dZeed< dZeed< dZeeef ed< dZeeeeef   ed< dZeeeeef   ed< dZeeeef  ed< dZeeef ed< dZeeef ed< dZeeef ed< d Zeed!< d"Zeed#< d$Zeeeeef f ed%< d&Z eeeeef f ed'< d(Z!eed)< d*d+ Z"dS ),r            r   .depths`   	embed_dimr   	num_heads      @	mlp_ratioF
pool_firstTexpand_attnqkv_biasuse_cls_tokenuse_abs_posresidual_poolingconvmoder   r   
kernel_qkv)r   r   r   r   r0   r0   stride_qN	stride_kv   r4   stride_kv_adaptive   r7   patch_kernelpatch_stridepatch_paddingmax	pool_typespatialrel_pos_typeZgelu	act_layerZ	layernorm
norm_layergư>norm_epsc                    s   t j}tjttfs6tfddt|D _t j|ksHJ tjttfsttfddt|D _t j|ksJ jd urj	d u rj g }t|D ]Bt
j dkr fddtt  D  |t  qt|_	d S )Nc                 3   s   | ]} j d |  V  qdS r   N)r!   .0iself [/var/www/html/stable-diffusion-webui/venv/lib/python3.9/site-packages/timm/models/mvitv2.py	<genexpr>@       z1MultiScaleVitCfg.__post_init__.<locals>.<genexpr>c                 3   s   | ]} j d |  V  qdS rB   )r"   rC   rF   rH   rI   rJ   D   rK   r   c                    s(   g | ] }t  | j |  d qS r   )r;   r1   )rD   dZ
_stride_kvrE   rG   rH   rI   
<listcomp>L   s   z2MultiScaleVitCfg.__post_init__.<locals>.<listcomp>)lenr   
isinstancer!   tuplelistranger"   r5   r2   minr1   append)rG   
num_stagesZpool_kv_striderH   rN   rI   __post_init__=   s"    

zMultiScaleVitCfg.__post_init__)#__name__
__module____qualname__r   r   int__annotations__r!   r   r"   r$   floatr%   boolr&   r'   r(   r)   r*   r,   strr.   r1   r	   r2   r5   r8   r9   r:   r<   r>   r?   r@   rA   rX   rH   rH   rH   rI   r   #   s0   
c                 C   s   t tj| dS )Nr   )r   operatormul)iterablerH   rH   rI   prodT   s    rd   c                       s>   e Zd ZdZd fdd	Zeejee	 f d	d
dZ
  ZS )
PatchEmbedz
    PatchEmbed.
    r      r6   r3   r-   c                    s$   t    tj|||||d| _d S )N)Zkernel_sizestridepadding)super__init__r
   Conv2dproj)rG   dim_indim_outkernelrg   rh   	__class__rH   rI   rj   ]   s    
zPatchEmbed.__init__returnc                 C   s*   |  |}|ddd|jdd  fS )Nr   r   )rl   flatten	transposeshaperG   xrH   rH   rI   forwardo   s    
zPatchEmbed.forward)r   rf   r6   r3   r-   )rY   rZ   r[   __doc__rj   r   torchTensorr   r\   rz   __classcell__rH   rH   rp   rI   re   X   s        re   T)	feat_sizehas_cls_tokenrs   c                 C   s   |\}}|rP| d d d d d dd d f | d d d d dd d d f  }} nd }|  d||| jd dddd } | |fS )Nr   r   r   r   )reshaperw   permute
contiguous)ry   r   r   HWcls_tokrH   rH   rI   reshape_pre_poolu   s    D&r   )r"   r   rs   c                 C   sh   | j d | j d g}| j d | j d  }| d|| j d |dd} |d ur`tj|| fdd} | |fS )Nr   r   r   r   dim)rw   r   rv   r|   cat)ry   r"   r   r   ZL_pooledrH   rH   rI   reshape_post_pool   s    r   )attnqr   q_sizek_size	rel_pos_h	rel_pos_wc                 C   s  |rdnd}|\}}	|\}
}t |
| d}t ||
 d}tj||jdd| tj|
|jdd|  }||
d | 7 }t ||	 d}t |	| d}tj|	|jdd| tj||jdd|  }||d | 7 }||  }||  }|j\}}}}|dddd|df ||||	|}td||}td||}| dddd|d|df 	|d||	|
||d |d	 	|d||	 |
| | dddd|d|df< | S )
z1
    Spatial Relative Positional Embeddings.
    r   r   g      ?)devicer   Nzbyhwc,hkc->byhwkzbyhwc,wkc->byhwkrt   )
r;   r|   Zaranger   Z	unsqueezelongrw   r   Zeinsumview)r   r   r   r   r   r   r   Zsp_idxZq_hZq_wZk_hZk_wZ	q_h_ratioZ	k_h_ratioZdist_hZ	q_w_ratioZ	k_w_ratioZdist_wZrel_hZrel_wBZn_headq_Nr   Zr_qrH   rH   rI   cal_rel_pos_type   s@    (. r   c                       sH   e Zd Zddddddddddejf fdd	Zee dd	d
Z  Z	S )MultiScaleAttentionPoolFirst   Tr+   r/   r=   c              	      s  t    || _|| _|| | _| jd | _|| _tdd |D }tdd |D }tj	|||d| _
tj	|||d| _tj	|||d| _t	||| _t|dkrt|	dkrd }t|dkrt|
dkrd }|| _|dk| _d\| _| _| _d\| _| _| _|d	v rZ|d
krtjntj}|r6|||	|| _|r|||
|| _|||
|| _n|dksn|dkr|dkr|| n|}|rtj||||	||dd| _||| _|rtj||||
||dd| _||| _tj||||
||dd| _||| _ntd| || _| jdkr|d |d ks4J |d }t|	dkrV||	d  n|}t|
dkrt||
d  n|}dt|| d }tt || j| _!tt || j| _"t#| j!dd t#| j"dd || _$d S )N      c                 S   s   g | ]}t |d  qS r   r\   rD   r   rH   rH   rI   rO      rK   z9MultiScaleAttentionPoolFirst.__init__.<locals>.<listcomp>c                 S   s   g | ]}t |d  qS r   r   rD   kvrH   rH   rI   rO      rK   biasr   conv_unsharedNNNavgr;   r;   r+   Frg   rh   groupsr   Unsupported model r=   r   r   {Gz?std)%ri   rj   r"   rn   head_dimscaler   rR   r
   Linearr   kvrl   rd   r,   unsharedpool_qpool_kpool_vnorm_qnorm_knorm_v	MaxPool2d	AvgPool2drk   NotImplementedErrorr>   rP   r;   	Parameterr|   zerosr   r   r   r*   rG   r   rn   r   r"   r'   r,   kernel_q	kernel_kvr1   r2   r   r>   r*   r@   Z	padding_qZ
padding_kvZpool_opZdim_convsizer   Zkv_sizeZ
rel_sp_dimrp   rH   rI   rj      s    



	
	
	z%MultiScaleAttentionPoolFirst.__init__r   c                 C   s  |j \}}}| jrdn| j}||||ddddd}| } }}	| jd ur~t||| j\}}
| |}t|| j|
\}}n|}| j	d ur| 	|}| j
d urt||| j\}}| 
|}t|| j|\}}n|}| jd ur| |}| jd ur$t|	|| j\}	}| |	}	t|	| j|\}	}n|}| jd ur>| |	}	|d |d  t| j }|dd||d}| |||| jddd}|d |d  t| j }|dd||d}| |||| jd}|d |d  t| j }|	dd||d}	| |	||| jddd}	|| j | }| jdkrZt||| j||| j| j}|jdd}||	 }| jr~|| }|dd|d| j}| |}||fS )Nr   r   r   r   r   r=   r   )rw   r   r"   r   r   r   r   r   r   r   r   r   r   r   r\   rv   r   r   r   r   r>   r   r   r   softmaxr*   rn   rl   )rG   ry   r   r   N_Zfold_dimr   r   r   q_tokr   k_tokr   v_tokZv_sizer   Zk_NZv_Nr   rH   rH   rI   rz   *  sh    









  	
z$MultiScaleAttentionPoolFirst.forward
rY   rZ   r[   r
   	LayerNormrj   r   r\   rz   r~   rH   rH   rp   rI   r      s   dr   c                       sH   e Zd Zddddddddddejf fdd	Zee dd	d
Z  Z	S )MultiScaleAttentionr   Tr+   r/   r=   c              	      s  t    || _|| _|| | _| jd | _|| _tdd |D }tdd |D }tj	||d |d| _
t	||| _t|dkrt|	dkrd }t|dkrt|
dkrd }|| _|dk| _d	\| _| _| _d	\| _| _| _|d
v r8|dkrtjntj}|r|||	|| _|r|||
|| _|||
|| _n|dksL|dkr|dkr^|| n|}|rtj||||	||dd| _||| _|rtj||||
||dd| _||| _tj||||
||dd| _||| _ntd| || _| jdkr|d |d ksJ |d }t|	dkr4||	d  n|}t|
dkrR||
d  n|}dt|| d }tt|| j| _tt|| j| _ t!| jdd t!| j dd || _"d S )Nr   c                 S   s   g | ]}t |d  qS r   r   r   rH   rH   rI   rO     rK   z0MultiScaleAttention.__init__.<locals>.<listcomp>c                 S   s   g | ]}t |d  qS r   r   r   rH   rH   rI   rO     rK   r   r   r   r   r   r   r;   r+   Fr   r   r=   r   r   r   r   )#ri   rj   r"   rn   r   r   r   rR   r
   r   qkvrl   rd   r,   r   r   r   r   r   r   r   r   r   rk   r   r>   rP   r;   r   r|   r   r   r   r   r*   r   rp   rH   rI   rj   p  s    



	
	
	zMultiScaleAttention.__init__r   c                 C   s  |j \}}}| |||d| jdddddd}|jdd\}}}	| jd urt||| j\}}
| |}t	|| j|
\}}n|}| j
d ur| 
|}| jd urt||| j\}}| |}t	|| j|\}}n|}| jd ur| |}| jd ur$t|	|| j\}	}| |	}	t	|	| j|\}	}| jd ur:| |	}	|| j |dd }| jd	krvt||| j||| j| j}|jdd}||	 }| jr|| }|dd|d| j}| |}||fS )
Nr   r   r   r   r   r4   r   rt   r=   )rw   r   r   r"   r   Zunbindr   r   r   r   r   r   r   r   r   r   rv   r>   r   r   r   r   r*   rn   rl   )rG   ry   r   r   r   r   r   r   r   r   r   r   r   r   r   r   rH   rH   rI   rz     sR    (









	
zMultiScaleAttention.forwardr   rH   rH   rp   rI   r   o  s   br   c                       s`   e Zd Zdddejddddddddddf fdd		Zee d
ddZee d
ddZ	  Z
S )MultiScaleBlockr#   T        r/   r+   Fr=   c                    sL  t    ||k}|| _|| _|| _||| _|rB|rBt||nd | _|rt	|dkrdd |D }|}dd |D }t
|||| _nd | _|r|n|}|rtnt}|||||||	|
|||||||d| _|dkrt|nt | _||| _|}|r|st||nd | _t|t|| |d| _|dkr>t|nt | _d S )Nr   c                 S   s    g | ]}|d kr|d  n|qS rL   rH   )rD   srH   rH   rI   rO   (  rK   z,MultiScaleBlock.__init__.<locals>.<listcomp>c                 S   s   g | ]}t |d  qS r   r   )rD   skiprH   rH   rI   rO   *  rK   )r"   r   r'   r   r   r1   r2   r@   r   r,   r>   r*   r   )Zin_featuresZhidden_featuresZout_features)ri   rj   r   rn   r   norm1r
   r   shortcut_proj_attnrd   r   shortcut_pool_attnr   r   r   r   Identity
drop_path1norm2shortcut_proj_mlpr   r\   mlp
drop_path2)rG   r   rn   r"   r   r$   r'   	drop_pathr@   r   r   r1   r2   r,   r   r&   r%   r>   r*   Zproj_neededZkernel_skipZstride_skipZpadding_skipZatt_dimZ
attn_layerZmlp_dim_outrp   rH   rI   rj   	  sR    



zMultiScaleBlock.__init__r   c           	      C   s   | j d u r|S | jrL|d d d dd d f |d d dd d d f  }}nd }|j\}}}|\}}|||||dddd }|  |}|||ddd}|d urtj||fdd}|S )Nr   r   r   r   r   r   )	r   r   rw   r   r   r   rv   r|   r   )	rG   ry   r   r   r   LCr   r   rH   rH   rI   _shortcut_poolM  s    
8 
zMultiScaleBlock._shortcut_poolc                 C   s   |  |}| jd u r|n| |}| ||}| ||\}}|| | }| |}| jd u rd|n| |}|| | | }||fS N)	r   r   r   r   r   r   r   r   r   )rG   ry   r   Zx_normZ
x_shortcutZfeat_size_newrH   rH   rI   rz   ]  s    

zMultiScaleBlock.forward)rY   rZ   r[   r
   r   rj   r   r\   r   rz   r~   rH   rH   rp   rI   r     s"   Dr   c                       sN   e Zd Zddddddddddddejdf fdd		Zee d
ddZ  Z	S )MultiScaleVitStager#   Tr+   r/   Fr=   r   c                    s   t    d| _t | _|r*|f| }n|f|d  |f }t|D ]}t||| |||||	|
|dkrl|nd||||||||t|t	t
fr|| n|d}|| }| j| |dkrFt
dd t||D }qF|| _d S )NFr   r   r/   )r   rn   r"   r   r$   r'   r   r   r1   r2   r,   r   r%   r>   r*   r&   r@   r   c                 S   s   g | ]\}}|| qS rH   rH   )rD   r   rg   rH   rH   rI   rO     rK   z/MultiScaleVitStage.__init__.<locals>.<listcomp>)ri   rj   grad_checkpointingr
   
ModuleListblocksrT   r   rQ   rS   rR   rV   zipr   )rG   r   rn   depthr"   r   r$   r'   r,   r   r   r1   r2   r   r&   r%   r>   r*   r@   r   Zout_dimsrE   Zattention_blockrp   rH   rI   rj   m  s@    

zMultiScaleVitStage.__init__r   c                 C   sF   | j D ]6}| jr.tj s.t|||\}}q|||\}}q||fS r   )r   r   r|   jitZis_scripting
checkpoint)rG   ry   r   ZblkrH   rH   rI   rz     s
    
zMultiScaleVitStage.forwardr   rH   rH   rp   rI   r   k  s    	;r   c                	       s   e Zd ZdZdeeeef eee ee	e	d fdd	Z
d
d Zejjdd Zejjd ddZejjd!ddZejjdd Zd"ddZdd Zd#edddZdd Z  ZS )$r   a  
    Improved Multiscale Vision Transformers for Classification and Detection
    Yanghao Li*, Chao-Yuan Wu*, Haoqi Fan, Karttikeya Mangalam, Bo Xiong, Jitendra Malik,
        Christoph Feichtenhofer*
    https://arxiv.org/abs/2112.01526

    Multiscale Vision Transformers
    Haoqi Fan*, Bo Xiong*, Karttikeya Mangalam*, Yanghao Li*, Zhicheng Yan, Jitendra Malik,
        Christoph Feichtenhofer*
    https://arxiv.org/abs/2104.11227
       r   r   N  r   )cfgimg_sizein_chansglobal_poolnum_classesdrop_path_rate	drop_ratec                    s  t    t|}tt|j|jd}|| _|| _|d u rH|j	rDdnd}|| _
t|j| _|j| _|jd }	t||	|j|j|jd| _|d |jd  |d |jd  f}
t|
}|j	rttdd|	| _d| _|d }nd| _d | _|}|jrttd||	| _nd | _t|j}|
}dd td|t|j |jD }t! | _"t#|D ]}|jrr|j| }n|jt$|d |d  }t%|	||j| |j&| ||j'|j(|j)|j*|j|j+|j+|j,| |j-| |j	|j.|j/||| d	}|}	|j0}| j"1| qZ|	| _2||	| _3t4t5d
t6| jfd|dkr@t7| j2|nt8 fg| _9| jd urlt:| jdd | jd urt:| jdd | ;| j< d S )N)Zepstokenr   r   )rm   rn   ro   rg   rh   r   c                 S   s   g | ]}|  qS rH   )tolist)rD   ry   rH   rH   rI   rO     rK   z*MultiScaleVit.__init__.<locals>.<listcomp>)r   rn   r   r"   r   r$   r'   r,   r%   r&   r   r   r1   r2   r   r>   r*   r@   r   dropfcr   r   )=ri   rj   r   r   r   r@   rA   r   r   r(   r   rR   r   r&   r!   re   r8   r9   r:   patch_embedrd   r
   r   r|   r   	cls_tokennum_prefix_tokensr)   	pos_embedrP   Zlinspacesumsplitr   stagesrT   rU   r   r"   r$   r'   r,   r%   r.   r1   r2   r>   r*   r   rV   num_featuresnorm
Sequentialr   Dropoutr   r   headr   apply_init_weights)rG   r   r   r   r   r   r   r   r@   r!   Z
patch_dimsZnum_patchesZpos_embed_dimrW   r   ZdprrE   rn   stagerp   rH   rI   rj     s    


$

&

"
zMultiScaleVit.__init__c                 C   sD   t |tjr@t|jdd t |tjr@|jd ur@tj|jd d S )Nr   r   r   )rQ   r
   r   r   Zweightr   initZ	constant_)rG   mrH   rH   rI   r    s    zMultiScaleVit._init_weightsc                 C   s   dd |   D S )Nc                    s*   h | ]"\ }t  fd ddD r qS )c                 3   s   | ]}| v V  qd S r   rH   )rD   nr   rH   rI   rJ   $  rK   z:MultiScaleVit.no_weight_decay.<locals>.<setcomp>.<genexpr>)r   r   r   r   )any)rD   r   rH   r  rI   	<setcomp>#  s   z0MultiScaleVit.no_weight_decay.<locals>.<setcomp>)Znamed_parametersrF   rH   rH   rI   no_weight_decay!  s    zMultiScaleVit.no_weight_decayFc                 C   s   t dddgd}|S )Nz^patch_embed)z^stages\.(\d+)N)z^norm)i )stemr   )dict)rG   ZcoarseZmatcherrH   rH   rI   group_matcher&  s
    zMultiScaleVit.group_matcherTc                 C   s   | j D ]
}||_qd S r   )r   r   )rG   enabler   rH   rH   rI   set_grad_checkpointing.  s    
z$MultiScaleVit.set_grad_checkpointingc                 C   s   | j jS r   )r   r   rF   rH   rH   rI   get_classifier3  s    zMultiScaleVit.get_classifierc              	   C   sV   || _ |d ur|| _ttdt| jfd|dkr@t| j|nt	 fg| _
d S )Nr   r   r   )r   r   r
   r   r   r   r   r   r   r   r   )rG   r   r   rH   rH   rI   reset_classifier7  s     zMultiScaleVit.reset_classifierc                 C   s   |  |\}}|j\}}}| jd urF| j|dd}tj||fdd}| jd urZ|| j }| jD ]}|||\}}q`| |}|S )Nr   r   r   )	r   rw   r   expandr|   r   r   r   r   )rG   ry   r   r   r   r   Z
cls_tokensr  rH   rH   rI   forward_features@  s    




zMultiScaleVit.forward_features)
pre_logitsc                 C   sP   | j r>| j dkr.|d d | jd f d}n|d d df }|rF|S | |S )Nr   r   r   )r   r   meanr   )rG   ry   r  rH   rH   rI   forward_headQ  s
    
zMultiScaleVit.forward_headc                 C   s   |  |}| |}|S r   )r  r  rx   rH   rH   rI   rz   Y  s    

zMultiScaleVit.forward)r   r   Nr   r   r   )F)T)N)F)rY   rZ   r[   r{   r   r   r\   r	   r`   r^   rj   r  r|   r   ignorer
  r  r  r  r  r  r_   r  rz   r~   rH   rH   rp   rI   r     s:         
]


	c           	         s   d| v r| S dd l }d| v r$| d } t|dd }t|dd}|d usLJ di d t|D ]4\} fdd	t  | D   |7  q\i }|  D ]\\}}|d
fdd|}|r|dd|}n|dd|}d|v r|dd}|||< q|S )Nzstages.0.blocks.0.norm1.weightr   Zmodel_stater   r&   Tz3model requires depth attribute to remap checkpointsc                    s   i | ]}||  fqS rH   rH   rC   )	block_idx	stage_idxrH   rI   
<dictcomp>m  rK   z(checkpoint_filter_fn.<locals>.<dictcomp>zblocks\.(\d+)c                    s4   d t | d d  d t | d d  S )Nzstages.r   r   z.blocks.)r\   group)ry   )	depth_maprH   rI   <lambda>t  rK   z&checkpoint_filter_fn.<locals>.<lambda>z stages\.(\d+).blocks\.(\d+).projz&stages.\1.blocks.\2.shortcut_proj_attnz%stages.\1.blocks.\2.shortcut_proj_mlpr   zhead.projectionhead.fc)regetattr	enumerateupdaterT   itemssubreplace)	Z
state_dictmodelr  r   r&   rM   Zout_dictr   r   rH   )r  r  r  rI   checkpoint_filter_fn_  s6    $


r'  )r   r      r   )r   )r   r      r   r   )r      $   r4      r   F)r   r!   r"   r&   )r   r(   )r   r!   r"   r(   r&   )r4   r   <   r      r   )mvitv2_tinymvitv2_smallmvitv2_basemvitv2_largemvitv2_small_clsmvitv2_base_clsmvitv2_large_clsmvitv2_huge_clsc                 K   s2   t t| |f|st|  nt| ttddd|S )NT)Zflatten_sequential)Z	model_cfgZpretrained_filter_fnZfeature_cfg)r   r   
model_cfgsr'  r  )variantZcfg_variant
pretrainedkwargsrH   rH   rI   _create_mvitv2  s    r;   c                 K   s    | ddd ddt tdddd|S )	Nr   )r   r   r   g?Zbicubiczpatch_embed.projr  T)urlr   Z
input_sizeZ	pool_sizeZcrop_pctinterpolationr  r   Z
first_conv
classifierZfixed_input_sizer   )r=  r:  rH   rH   rI   _cfg  s    r@  zDhttps://dl.fbaipublicfiles.com/mvit/mvitv2_models/MViTv2_T_in1k.pyth)r=  zDhttps://dl.fbaipublicfiles.com/mvit/mvitv2_models/MViTv2_S_in1k.pythzDhttps://dl.fbaipublicfiles.com/mvit/mvitv2_models/MViTv2_B_in1k.pythzDhttps://dl.fbaipublicfiles.com/mvit/mvitv2_models/MViTv2_L_in1k.pythzEhttps://dl.fbaipublicfiles.com/mvit/mvitv2_models/MViTv2_B_in21k.pythiJ  )r=  r   zEhttps://dl.fbaipublicfiles.com/mvit/mvitv2_models/MViTv2_L_in21k.pythzEhttps://dl.fbaipublicfiles.com/mvit/mvitv2_models/MViTv2_H_in21k.pyth)zmvitv2_tiny.fb_in1kzmvitv2_small.fb_in1kzmvitv2_base.fb_in1kzmvitv2_large.fb_in1kr3  zmvitv2_base_cls.fb_inw21kzmvitv2_large_cls.fb_inw21kzmvitv2_huge_cls.fb_inw21krr   c                 K   s   t dd| i|S )Nr/  r9  )r/  r;  r9  r:  rH   rH   rI   r/    s    r/  c                 K   s   t dd| i|S )Nr0  r9  )r0  rA  rB  rH   rH   rI   r0    s    r0  c                 K   s   t dd| i|S )Nr1  r9  )r1  rA  rB  rH   rH   rI   r1    s    r1  c                 K   s   t dd| i|S )Nr2  r9  )r2  rA  rB  rH   rH   rI   r2    s    r2  c                 K   s   t dd| i|S )Nr3  r9  )r3  rA  rB  rH   rH   rI   r3    s    r3  c                 K   s   t dd| i|S )Nr4  r9  )r4  rA  rB  rH   rH   rI   r4    s    r4  c                 K   s   t dd| i|S )Nr5  r9  )r5  rA  rB  rH   rH   rI   r5    s    r5  c                 K   s   t dd| i|S )Nr6  r9  )r6  rA  rB  rH   rH   rI   r6    s    r6  )T)N)NF)r<  )F)F)F)F)F)F)F)F)Br{   ra   collectionsr   Zdataclassesr   	functoolsr   r   typingr   r   r   r	   r|   Ztorch.utils.checkpointutilsr   r
   Z	timm.datar   r   Ztimm.layersr   r   r   r   r   Z_builderr   Z_features_fxr   	_registryr   r   r   __all__r   rd   Modulere   r\   r_   r}   r   r   r   r   r   r   r   r   r'  r  r7  r;  r@  Zdefault_cfgsr/  r0  r1  r2  r3  r4  r5  r6  rH   rH   rH   rI   <module>   s   0  2 + cF /-*

