a
    d~                     @   s  d Z ddlmZ ddlmZ ddlmZ ddl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 dg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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Z(G d%d de	jZ)G d&d' d'e)Z*d(d) Z+e,e,d*d+d,d-d.e,d*d+d/d0d.e,d1d2d3d0d.e,d4d2d,d0d.e,d5d2d6d0d.e,d5d2d6d0d7d8d9e,d:d;d<d0d7d8d9e,d=d;d>d0d7d?e,d4d2d,d@d7d?e,dAd;dBd@d7d?dC
Z-d{dFdGZ.d|dIdJZ/ee/dKdLe/dKdLe/dKdLe/dKdLe/dKdLe/dKdMdNe/dKdMdNe/dKdMdNe/dKdMdNe/dKdMdNe/dOdPe/dOdPe/dOdPe/dOdPe/dOdPe/dOdPe/dOdPe/dOdPe/dOdPe/dOdPdQZ0ed}e)dRdSdTZ1ed~e)dRdUdVZ2ede)dRdWdXZ3ede)dRdYdZZ4ede)dRd[d\Z5ede)dRd]d^Z6ede)dRd_d`Z7ede)dRdadbZ8ede)dRdcddZ9ede)dRdedfZ:ede)dRdgdhZ;ede)dRdidjZ<ede)dRdkdlZ=ede)dRdmdnZ>ede)dRdodpZ?ede)dRdqdrZ@ede)dRdsdtZAede)dRdudvZBede)dRdwdxZCede)dRdydzZDdS )a   LeViT

Paper: `LeViT: a Vision Transformer in ConvNet's Clothing for Faster Inference`
    - https://arxiv.org/abs/2104.01136

@article{graham2021levit,
  title={LeViT: a Vision Transformer in ConvNet's Clothing for Faster Inference},
  author={Benjamin Graham and Alaaeldin El-Nouby and Hugo Touvron and Pierre Stock and Armand Joulin and Herv'e J'egou and Matthijs Douze},
  journal={arXiv preprint arXiv:22104.01136},
  year={2021}
}

Adapted from official impl at https://github.com/facebookresearch/LeViT, original copyright bellow.

This version combines both conv/linear models and fixes torchscript compatibility.

Modifications and additions for timm hacked together by / Copyright 2021, Ross Wightman
    )OrderedDict)partial)DictN)IMAGENET_DEFAULT_STDIMAGENET_DEFAULT_MEAN)	to_ntuple	to_2tupleget_act_layerDropPathtrunc_normal_   )build_model_with_cfg)checkpoint_seq)generate_default_cfgsregister_modelLevitc                       s6   e Zd Zd	 fdd	Ze dd Zdd Z  ZS )
ConvNormr   r   c	           	   
      sH   t    tj|||||||dd| _t|| _tj| jj	| d S NFbias)
super__init__nnConv2dlinearZBatchNorm2dbninit	constant_weight)	selfin_chsout_chskernel_sizestridepaddingdilationgroupsbn_weight_init	__class__ Z/var/www/html/stable-diffusion-webui/venv/lib/python3.9/site-packages/timm/models/levit.pyr   +   s    
zConvNorm.__init__c              	   C   s   | j | j }}|j|j|j d  }|j|d d d d d f  }|j|j|j |j|j d   }tj|	d|	d|j
dd  | j j| j j| j j| j jd}|jj| |jj| |S )N      ?r   r      )r#   r$   r%   r&   )r   r   r   running_varepsr   running_meanr   r   sizeshaper#   r$   r%   r&   datacopy_)r   cr   wbmr*   r*   r+   fuse3   s    ""zConvNorm.fusec                 C   s   |  | |S N)r   r   r   xr*   r*   r+   forward@   s    zConvNorm.forward)r   r   r   r   r   r   	__name__
__module____qualname__r   torchno_gradr9   r=   __classcell__r*   r*   r(   r+   r   *   s
    
r   c                       s6   e Zd Zd fdd	Ze dd Zdd Z  ZS )	
LinearNormr   c                    s>   t    tj||dd| _t|| _tj| jj	| d S r   )
r   r   r   Linearr   BatchNorm1dr   r   r   r   )r   in_featuresout_featuresr'   r(   r*   r+   r   E   s    
zLinearNorm.__init__c                 C   s   | j | j }}|j|j|j d  }|j|d d d f  }|j|j|j |j|j d   }t|	d|	d}|jj
| |jj
| |S )Nr,   r   r   )r   r   r   r.   r/   r   r0   r   rF   r1   r3   r4   )r   lr   r6   r7   r8   r*   r*   r+   r9   L   s    "zLinearNorm.fusec                 C   s"   |  |}| |dd|S )Nr   r   )r   r   flattenZ
reshape_asr;   r*   r*   r+   r=   W   s    
zLinearNorm.forward)r   r>   r*   r*   r(   r+   rE   D   s   

rE   c                       s6   e Zd Zd
 fdd	Ze dd Zdd	 Z  ZS )
NormLinearT{Gz?        c                    sf   t    t|| _t|| _tj|||d| _t	| jj
|d | jjd urbtj| jjd d S )Nr   )stdr   )r   r   r   rG   r   DropoutdroprF   r   r   r   r   r   r   )r   rH   rI   r   rO   rQ   r(   r*   r+   r   ]   s    
zNormLinear.__init__c                 C   s   | j | j }}|j|j|j d  }|j| j j| j j |j|j d   }|j|d d d f  }|jd u rz|| jjj }n$|j|d d d f  d| jj }t	
|d|d}|jj| |jj| |S )Nr,   r   r   )r   r   r   r.   r/   r   r0   Tviewr   rF   r1   r3   r4   )r   r   rJ   r6   r7   r8   r*   r*   r+   r9   g   s    &
$zNormLinear.fusec                 C   s   |  | | |S r:   )r   rQ   r   r;   r*   r*   r+   r=   v   s    zNormLinear.forward)TrM   rN   r>   r*   r*   r(   r+   rL   \   s   

rL   c                       s   e Zd Z fddZ  ZS )Stem8c              
      s   t    d| _| dt||d dddd | d|  | d	t|d |d dddd | d
|  | dt|d |dddd d S )N   conv1      r-   r   r#   r$   act1conv2act2conv3r   r   r#   Z
add_moduler   r   r    r!   	act_layerr(   r*   r+   r   {   s    
"zStem8.__init__r?   r@   rA   r   rD   r*   r*   r(   r+   rU   z   s   rU   c                       s   e Zd Z fddZ  ZS )Stem16c              
      s   t    d| _| dt||d dddd | d|  | d	t|d |d
 dddd | d|  | dt|d
 |d dddd | d|  | dt|d |dddd d S )N   rW   rV   rY   r-   r   rZ   r[   r\   rX   r]   r^   Zact3Zconv4r_   r`   r(   r*   r+   r      s    
""zStem16.__init__rb   r*   r*   r(   r+   rc      s   rc   c                       s&   e Zd Zd fdd	Zdd Z  ZS )
DownsampleFc                    s:   t    || _t|| _|r0tjd|dddnd | _d S )NrY   r   F)r#   r$   count_include_pad)r   r   r#   r   
resolutionr   	AvgPool2dpool)r   r#   rg   use_poolr(   r*   r+   r      s    

zDownsample.__init__c                 C   s   |j \}}}||| jd | jd |}| jd urV| |dddddddd}n"|d d d d | jd d | jf }||d|S )Nr   r   rY   r-   rR   )r2   rT   rg   ri   permuter#   reshape)r   r<   BNCr*   r*   r+   r=      s    
$"zDownsample.forward)F)r?   r@   rA   r   r=   rD   r*   r*   r(   r+   re      s   re   c                       sp   e Zd ZU eeejf ed< ddddej	f fdd	Z
e d fd	d
	ZejejdddZdd Z  ZS )	Attentionattention_bias_cacherV         @   Fc              
      sL  t    |rtnt}t|}|| _|| _|d | _|| _|| | _	t
|| | _t
|| | | _||| j| j	d  | _ttd| fd|| j|ddfg| _tt||d |d  | _ttt|d t|d d}	|	dd d d f |	dd d d f   }
|
d |d  |
d  }
| jd	|
d
d i | _d S )N      r-   actlnr   r'   r   .attention_bias_idxsF
persistent)r   r   r   rE   r   use_conv	num_headsscalekey_dimkey_attn_dimintval_dimval_attn_dimqkvr   
Sequentialr   proj	ParameterrB   zerosattention_biasesstackmeshgridarangerK   absregister_bufferrq   )r   dimr~   r|   
attn_ratiorg   r{   ra   ln_layerposrel_posr(   r*   r+   r      s*    




 ,(zAttention.__init__Tc                    s    t  | |r| jri | _d S r:   r   trainrq   r   moder(   r*   r+   r      s    
zAttention.traindevicereturnc                 C   s^   t j s| jr$| jd d | jf S t|}|| jvrP| jd d | jf | j|< | j| S d S r:   rB   jit
is_tracingtrainingr   rx   strrq   r   r   Z
device_keyr*   r*   r+   get_attention_biases   s    
zAttention.get_attention_biasesc                 C   sX  | j r|j\}}}}| ||| jd|| j| j| j| jgdd\}}}|dd| | j	 | 
|j }	|	jdd}	||	dd |d||}n|j\}}
}| |||
| jdj| j| j| jgdd\}}}|dddd}|dddd}|dddd}|| | j	 | 
|j }	|	jdd}	|	| dd||
| j}| |}|S )NrR   r-   r   rY   r   r   )r{   r2   r   rT   r|   splitr~   r   	transposer}   r   r   softmaxrk   rl   r   r   )r   r<   rm   ro   HWqkvattnrn   r*   r*   r+   r=      s0    
"


zAttention.forward)Tr?   r@   rA   r   r   rB   ZTensor__annotations__r   SiLUr   rC   r   r   r   r=   rD   r*   r*   r(   r+   rp      s   
#	rp   c                       st   e Zd ZU eeejf ed< ddddddej	f fdd	Z
e d fd
d	ZejejdddZdd Z  ZS )AttentionDownsamplerq   rV          @r-   rs   Fc              
      s  t    t|}|| _|| _|| _|| _|| | _t|| | _	| j	| j | _
|d | _|| _| jrt}ttj|	rxdnd|	rdnddd}nt}tt||	d}||| j
| j | _ttd||d	fd
||| jfg| _ttd|
 fd
|| j
|fg| _tt||d |d  | _ttt|d t|d d}tttjd|d |dtjd|d |dd}|dd d d f |dd d d f    }|d |d  |d  }| j!d|dd i | _"d S )Nrt   rY   r   r   F)r"   r$   rf   )rg   rj   Zdown)r#   rv   ru   )step.rx   ry   )#r   r   r   r#   rg   r|   r~   r   r   r   r   r}   r{   r   r   r   rh   rE   re   kvr   r   r   r   r   rB   r   r   r   r   r   rK   r   r   rq   )r   in_dimout_dimr~   r|   r   r#   rg   r{   rj   ra   r   Z	sub_layerZk_posZq_posr   r(   r*   r+   r      sP    




 ,(zAttentionDownsample.__init__Tc                    s    t  | |r| jri | _d S r:   r   r   r(   r*   r+   r   /  s    
zAttentionDownsample.trainr   c                 C   s^   t j s| jr$| jd d | jf S t|}|| jvrP| jd d | jf | j|< | j| S d S r:   r   r   r*   r*   r+   r   5  s    
z(AttentionDownsample.get_attention_biasesc                 C   s  | j r|j\}}}}|d | j d |d | j d  }}| ||| jd|| j| j| jgdd\}}	| 	||| j| jd}
|

dd| | j | |j }|jdd}|	|
dd || j||}n|j\}}}| |||| jdj| j| jgdd\}}	|dddd}|	dddd}	| 	||d| j| jdddd}
|
| | j | |j }|jdd}||	 
dd|d| j}| |}|S )Nr   rR   r-   r   r   rY   r   )r{   r2   r#   r   rT   r|   r   r~   r   r   r   r}   r   r   r   rl   r   rk   r   )r   r<   rm   ro   r   r   ZHHZWWr   r   r   r   rn   r*   r*   r+   r=   >  s$    &2" .&
zAttentionDownsample.forward)Tr   r*   r*   r(   r+   r      s   
8	r   c                       s6   e Zd ZdZdddejdf fdd	Zdd Z  ZS )	LevitMlpzL MLP for Levit w/ normalization + ability to switch btw conv and linear
    NFrN   c                    sZ   t    |p|}|p|}|r"tnt}|||| _| | _t|| _|||dd| _	d S )Nr   rw   )
r   r   r   rE   ln1ru   r   rP   rQ   ln2)r   rH   Zhidden_featuresrI   r{   ra   rQ   r   r(   r*   r+   r   [  s    	
zLevitMlp.__init__c                 C   s,   |  |}| |}| |}| |}|S r:   )r   ru   rQ   r   r;   r*   r*   r+   r=   n  s
    



zLevitMlp.forward)	r?   r@   rA   __doc__r   r   r   r=   rD   r*   r*   r(   r+   r   X  s   r   c                	       s:   e Zd Zdddejdddddf	 fdd		Zd
d Z  ZS )LevitDownsamplerV   rr   r   Nrs   FrN   c                    sf   t    |p|}t|||||||	|
|d	| _t|t|| |
|d| _|dkrXt|nt	 | _
d S )N)	r   r   r~   r|   r   ra   rg   r{   rj   r{   ra   rN   )r   r   r   attn_downsampler   r   mlpr
   r   Identity	drop_path)r   r   r   r~   r|   r   	mlp_ratiora   attn_act_layerrg   r{   rj   r   r(   r*   r+   r   w  s(    

zLevitDownsample.__init__c                 C   s"   |  |}|| | | }|S r:   )r   r   r   r;   r*   r*   r+   r=     s    
zLevitDownsample.forwardr?   r@   rA   r   r   r   r=   rD   r*   r*   r(   r+   r   v  s   &r   c                       s8   e Zd Zdddddejddf fdd		Zd
d Z  ZS )
LevitBlockrV   rr   r   rs   FNrN   c              	      s|   t    |	p|}	t|||||||	d| _|
dkr:t|
nt | _t|t	|| ||d| _
|
dkrnt|
nt | _d S )N)r   r~   r|   r   rg   r{   ra   rN   r   )r   r   rp   r   r
   r   r   
drop_path1r   r   r   
drop_path2)r   r   r~   r|   r   r   rg   r{   ra   r   r   r(   r*   r+   r     s&    
	
zLevitBlock.__init__c                 C   s,   ||  | | }|| | | }|S r:   )r   r   r   r   r;   r*   r*   r+   r=     s    zLevitBlock.forwardr   r*   r*   r(   r+   r     s   #r   c                
       s<   e Zd Zddddejdddddf
 fd	d
	Zdd Z  ZS )
LevitStagerX   rV   rr   Nrs    FrN   c                    s   t    t|
}
|rJt||||| dd||	|
||d| _dd |
D }
n||ksVJ t | _g }t|D ]&}|t|||||||	|
||d
g7 }qltj	| | _
d S )Nrr   r   )	r~   r|   r   r   ra   r   rg   r{   r   c                 S   s   g | ]}|d  d d  qS )r   r-   r*   .0rr*   r*   r+   
<listcomp>      z'LevitStage.__init__.<locals>.<listcomp>)r|   r   r   ra   r   rg   r{   r   )r   r   r   r   
downsampler   r   ranger   r   blocks)r   r   r   r~   depthr|   r   r   ra   r   rg   r   r{   r   r   _r(   r*   r+   r     sD    

zLevitStage.__init__c                 C   s   |  |}| |}|S r:   )r   r   r;   r*   r*   r+   r=     s    

zLevitStage.forwardr   r*   r*   r(   r+   r     s   6r   c                       s   e Zd ZdZd% f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   z Vision Transformer with support for patch or hybrid CNN input stage

    NOTE: distillation is defaulted to True since pretrained weights use it, will cause problems
    w/ train scripts that don't take tuple outputs,
       rY        @      rY   r   Ns16	subsample
hard_swishFavgrN   c                    s  t    t|}t|p|}|| _|| _|| _|d | _|| _|| _d| _	g | _
t|}t||ksjJ t||}t||}t||	}	|
d ur|dksJ |
| _|}nF|dv sJ |dkrt||d |d| _nt||d |d| _| jj}tdd	 tt|t|D }|d }g }t|D ]}|dkr:dnd
 |t||| ||| || || |	| |||| dkrz|nd|dg7 }| 9 }t fdd	|D }|  j
t|| |d| dg7  _
|| }q(tj| | _|dkrt|d ||dnt | _d S )NrR   Fr-   )r   s8r   r   )ra   c                 S   s   g | ]\}}|| qS r*   r*   )r   ipr*   r*   r+   r   C  r   z"Levit.__init__.<locals>.<listcomp>r   r   )
r   r|   r   r   ra   r   rg   r{   r   r   c                    s   g | ]}|d    d  qS )r   r*   r   Zstage_strider*   r+   r   Y  r   zstages.)Znum_chsZ	reductionmodulerQ   )r   r   r	   r{   num_classesglobal_poolnum_features	embed_dim	drop_rategrad_checkpointingZfeature_infolenr   stemrc   rU   r#   tuplezipr   r   r   dictr   r   stagesrL   r   head)r   Zimg_sizeZin_chansr   r   r~   r   r|   r   r   Zstem_backboneZstem_stride	stem_typeZdown_opra   r   r{   r   r   Zdrop_path_rateZ
num_stagesr#   rg   r   r   r   r(   r   r+   r     sf    

 
$zLevit.__init__c                 C   s   dd |    D S )Nc                 S   s   h | ]}d |v r|qS )r   r*   )r   r<   r*   r*   r+   	<setcomp>c  r   z(Levit.no_weight_decay.<locals>.<setcomp>)
state_dictkeysr   r*   r*   r+   no_weight_decaya  s    zLevit.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 )r   r   )r   )r   ZcoarseZmatcherr*   r*   r+   group_matchere  s
    zLevit.group_matcherTc                 C   s
   || _ d S r:   )r   r   enabler*   r*   r+   set_grad_checkpointingm  s    zLevit.set_grad_checkpointingc                 C   s   | j S r:   )r   r   r*   r*   r+   get_classifierq  s    zLevit.get_classifierc                 C   s@   || _ |d ur|| _|dkr2t| jd || jdnt | _d S )Nr   rR   r   )r   r   rL   r   r   r   r   r   r   r   r   Zdistillationr*   r*   r+   reset_classifieru  s    zLevit.reset_classifierc                 C   sN   |  |}| js"|ddd}| jr@tj s@t| j	|}n
| 	|}|S )Nr-   r   )
r   r{   rK   r   r   rB   r   is_scriptingr   r   r;   r*   r*   r+   forward_features|  s    

zLevit.forward_features
pre_logitsc                 C   s:   | j dkr(| jr|jddn
|jdd}|r0|S | |S )Nr   r   rR   r   r   )r   r{   meanr   )r   r<   r   r*   r*   r+   forward_head  s    
zLevit.forward_headc                 C   s   |  |}| |}|S r:   )r   r  r;   r*   r*   r+   r=     s    

zLevit.forward)r   rY   r   r   r   r   r   r   r   NNr   r   r   NFr   rN   rN   )F)T)NN)F)r?   r@   rA   r   r   rB   r   ignorer   r   r   r   r   r   boolr  r=   rD   r*   r*   r(   r+   r   
  sB                      P



c                       sX   e Zd Z fddZejjdd ZdddZejjdd	d
Z	de
dddZ  ZS )LevitDistilledc                    s>   t  j|i | | jdkr*t| j| jnt | _d| _d S )Nr   F)	r   r   r   rL   r   r   r   	head_distdistilled_training)r   argskwargsr(   r*   r+   r     s    "zLevitDistilled.__init__c                 C   s   | j | jfS r:   )r   r  r   r*   r*   r+   r     s    zLevitDistilled.get_classifierNc                 C   sZ   || _ |d ur|| _|dkr.t| j|| jdnt | _|dkrLt| j|nt | _d S )Nr   r   )	r   r   rL   r   r   r   r   r   r  r   r*   r*   r+   r     s    
zLevitDistilled.reset_classifierTc                 C   s
   || _ d S r:   )r  r   r*   r*   r+   set_distilled_training  s    z%LevitDistilled.set_distilled_trainingFr   c                 C   st   | j dkr(| jr|jddn
|jdd}|r0|S | || | }}| jrd| jrdtj	 sd||fS || d S d S )Nr   r  r   r   r-   )
r   r{   r  r   r  r  r   rB   r   r   )r   r<   r   Zx_distr*   r*   r+   r    s    
zLevitDistilled.forward_head)NN)T)F)r?   r@   rA   r   rB   r   r  r   r   r  r  r  rD   r*   r*   r(   r+   r    s   

r  c                 C   s   d| v r| d } dd |   D } | }i }t| |  | |  D ]`\}}}}|jdkr|jdkr|d d d d d d f }|j|jkrd|v sd|v sJ |||< qL|S )Nmodelc                 S   s   i | ]\}}d |vr||qS )rx   r*   )r   r   r   r*   r*   r+   
<dictcomp>  r   z(checkpoint_filter_fn.<locals>.<dictcomp>rX   r-   r   stem.conv1.linear)itemsr   r   r   valuesndimr2   )r   r  DZout_dictkakbvaZvbr*   r*   r+   checkpoint_filter_fn  s    *
r  )        rd   )rX      rV   )r-   rY   rX   )r   r~   r|   r   )rX   rV   r   )rX   rX   rX   )r   i   r      )rY      r  )r  r     )r  r     )r  	   r   Zsilur   )r   r~   r|   r   ra   r   )r    i  r   )rV   
   rs   )r  r  i   )rV   r   rd   )r   r~   r|   r   ra   )rX   rV   r  )r  r   r  )rV   r!  r   )

levit_128s	levit_128	levit_192	levit_256	levit_384levit_384_s8levit_512_s8	levit_512
levit_256d
levit_512dFTc           	      K   s   d| v }| dd}|dd r,|s,td|d u rR| tv rB| }n|rR| dd}tt| fi |}t|rptnt| |ft	td|dd	|}|S )
NZ_convout_indices)r   r   r-   Zfeatures_onlyzBfeatures_only not implemented for LeVit in non-convolutional mode.r   T)Zflatten_sequentialr,  )Zpretrained_filter_fnZfeature_cfg)
popgetRuntimeError
model_cfgsreplacer   r   r  r   r  )	variantZcfg_variant
pretrained	distilledr
  Zis_convr,  Z	model_cfgr  r*   r*   r+   create_levit  s,    

r5  r   c                 K   s    | ddd dddt tddd|S )	Nr   )rY   r   r   g?ZbicubicTr  )head.linearzhead_dist.linear)urlr   Z
input_size	pool_sizeZcrop_pctinterpolationZfixed_input_sizer  rO   Z
first_conv
classifier)r   r   )r7  r
  r*   r*   r+   _cfg  s    r;  ztimm/)	hf_hub_id)rX   rX   )r<  r8  r6  )r:  )zlevit_128s.fb_dist_in1kzlevit_128.fb_dist_in1kzlevit_192.fb_dist_in1kzlevit_256.fb_dist_in1kzlevit_384.fb_dist_in1kzlevit_conv_128s.fb_dist_in1kzlevit_conv_128.fb_dist_in1kzlevit_conv_192.fb_dist_in1kzlevit_conv_256.fb_dist_in1kzlevit_conv_384.fb_dist_in1kzlevit_384_s8.untrainedzlevit_512_s8.untrainedzlevit_512.untrainedzlevit_256d.untrainedzlevit_512d.untrainedzlevit_conv_384_s8.untrainedzlevit_conv_512_s8.untrainedzlevit_conv_512.untrainedzlevit_conv_256d.untrainedzlevit_conv_512d.untrained)r   c                 K   s   t dd| i|S )Nr"  r3  )r"  r5  r3  r
  r*   r*   r+   r"  B  s    r"  c                 K   s   t dd| i|S )Nr#  r3  )r#  r=  r>  r*   r*   r+   r#  G  s    r#  c                 K   s   t dd| i|S )Nr$  r3  )r$  r=  r>  r*   r*   r+   r$  L  s    r$  c                 K   s   t dd| i|S )Nr%  r3  )r%  r=  r>  r*   r*   r+   r%  Q  s    r%  c                 K   s   t dd| i|S )Nr&  r3  )r&  r=  r>  r*   r*   r+   r&  V  s    r&  c                 K   s   t dd| i|S )Nr'  r3  )r'  r=  r>  r*   r*   r+   r'  [  s    r'  c                 K   s   t d| dd|S )Nr(  Fr3  r4  )r(  r=  r>  r*   r*   r+   r(  `  s    r(  c                 K   s   t d| dd|S )Nr)  Fr?  )r)  r=  r>  r*   r*   r+   r)  e  s    r)  c                 K   s   t d| dd|S )Nr*  Fr?  )r*  r=  r>  r*   r*   r+   r*  j  s    r*  c                 K   s   t d| dd|S )Nr+  Fr?  )r+  r=  r>  r*   r*   r+   r+  o  s    r+  c                 K   s   t d| dd|S )Nlevit_conv_128sTr3  r{   )r@  r=  r>  r*   r*   r+   r@  t  s    r@  c                 K   s   t d| dd|S )Nlevit_conv_128TrA  )rB  r=  r>  r*   r*   r+   rB  y  s    rB  c                 K   s   t d| dd|S )Nlevit_conv_192TrA  )rC  r=  r>  r*   r*   r+   rC  ~  s    rC  c                 K   s   t d| dd|S )Nlevit_conv_256TrA  )rD  r=  r>  r*   r*   r+   rD    s    rD  c                 K   s   t d| dd|S )Nlevit_conv_384TrA  )rE  r=  r>  r*   r*   r+   rE    s    rE  c                 K   s   t d| dd|S )Nlevit_conv_384_s8TrA  )rF  r=  r>  r*   r*   r+   rF    s    rF  c                 K   s   t d| ddd|S )Nlevit_conv_512_s8TFr3  r{   r4  )rG  r=  r>  r*   r*   r+   rG    s    rG  c                 K   s   t d| ddd|S )Nlevit_conv_512TFrH  )rI  r=  r>  r*   r*   r+   rI    s    rI  c                 K   s   t d| ddd|S )Nlevit_conv_256dTFrH  )rJ  r=  r>  r*   r*   r+   rJ    s    rJ  c                 K   s   t d| ddd|S )Nlevit_conv_512dTFrH  )rK  r=  r>  r*   r*   r+   rK    s    rK  )NFT)r   )F)F)F)F)F)F)F)F)F)F)F)F)F)F)F)F)F)F)F)F)Er   collectionsr   	functoolsr   typingr   rB   Ztorch.nnr   Z	timm.datar   r   Ztimm.layersr   r   r	   r
   r   Z_builderr   Z_manipulater   	_registryr   r   __all__Moduler   rE   rL   r   rU   rc   re   rp   r   r   r   r   r   r   r  r  r   r0  r5  r;  Zdefault_cfgsr"  r#  r$  r%  r&  r'  r(  r)  r*  r+  r@  rB  rC  rD  rE  rF  rG  rI  rJ  rK  r*   r*   r*   r+   <module>   s,  Od-*= $


 

6