
    Ng[              
          d Z dgZddlZddlmZ ddlmZ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 dd	lmZ dd
lmZ ddlmZmZ  G d dej
        j                  Z G d dej
        j                  Z G d dej
        j                  Z G d dej
        j                  Z G d dej
        j                  Z  G d dej
        j                  Z! G d dej
        j                  Z" G d dej
        j                  Z# G d dej
        j                  Z$ G d dej
        j                  Z% G d  de
j                  Z&d.d"Z' e e'd#$           e'd#$           e'd#$           e'd#$           e'd#$           e'd#$          d%          Z(d/d'Z)ed/d(            Z*ed/d)            Z+ed/d*            Z,ed/d+            Z-ed/d,            Z.ed/d-            Z/dS )0z EfficientViT (by MSRA)

Paper: `EfficientViT: Memory Efficient Vision Transformer with Cascaded Group Attention`
    - https://arxiv.org/abs/2305.07027

Adapted from official impl at https://github.com/microsoft/Cream/tree/main/EfficientViT
EfficientVitMsra    N)OrderedDict)DictOptionalIMAGENET_DEFAULT_MEANIMAGENET_DEFAULT_STD)SqueezeExciteSelectAdaptivePool2dtrunc_normal__assert   )build_model_with_cfg)checkpoint_seq)register_modelgenerate_default_cfgsc                   P     e Zd Zd fd	Z ej                    d             Z xZS )ConvNormr   r   c	           
      x   t                                                       t          j        |||||||d          | _        t          j        |          | _        t          j        j        	                    | j        j
        |           t          j        j        	                    | j        j        d           d S )NFbiasr   )super__init__nnConv2dconvBatchNorm2dbntorchinit	constant_weightr   )
selfin_chsout_chsksstridepaddilationgroupsbn_weight_init	__class__s
            Y/var/www/html/ai-engine/env/lib/python3.11/site-packages/timm/models/efficientvit_msra.pyr   zConvNorm.__init__   s    Ifgr63&W\]]]	.))???a00000    c           	      n   | j         | j        }}|j        |j        |j        z   dz  z  }|j        |d d d d d f         z  }|j        |j        |j        z  |j        |j        z   dz  z  z
  }t          j        	                    |
                    d          | j         j        z  |
                    d          |j        dd          | j         j        | j         j        | j         j        | j         j                  }|j        j                            |           |j        j                            |           |S )N      ?r   r      )r'   paddingr)   r*   )r   r   r"   running_varepsr   running_meanr   r   r   sizer*   shaper'   r2   r)   datacopy_)r#   cr   wbms         r-   fusezConvNorm.fuse    s   	472I"&0366HqD$,--Gbo	1^bf$s*+ +HOOFF1II	((!&&))QWQRR[9#TY->I[dhdmdt  v v 	
A	!r.   )r   r   r   r   r   r   __name__
__module____qualname__r   r   no_gradr>   __classcell__r,   s   @r-   r   r      s`        1 1 1 1 1 1 U]__  _    r.   r   c                   P     e Zd Zd fd	Z ej                    d             Z xZS )
NormLinearT{Gz?        c                    t                                                       t          j        |          | _        t          j        |          | _        t          j        |||          | _        t          | j        j
        |           | j        j        ,t          j                            | j        j        d           d S d S )Nr   )stdr   )r   r   r   BatchNorm1dr   DropoutdropLinearlinearr   r"   r   r    r!   )r#   in_featuresout_featuresr   rK   rN   r,   s         r-   r   zNormLinear.__init__0   s    .--Jt$$	i\EEEdk(c2222;'Gdk.22222 ('r.   c                    | j         | j        }}|j        |j        |j        z   dz  z  }|j        | j         j        | j         j        z  |j        |j        z   dz  z  z
  }|j        |d d d f         z  }|j        || j        j        j        z  }n4|j        |d d d f         z                      d          | j        j        z   }t          j
                            |                    d          |                    d                    }|j        j                            |           |j        j                            |           |S )Nr0   r   r   )r   rP   r"   r3   r4   r   r5   Tviewr   r   rO   r6   r8   r9   )r#   r   rP   r;   r<   r=   s         r-   r>   zNormLinear.fuse:   s    WdkFI"&0366Gdg*GN nrv5;< <MAdAAAgJ&;DK&((AA111d7+11"558HHAHOOAFF1IIqvvayy11	A	!r.   )TrH   rI   r?   rE   s   @r-   rG   rG   /   s`        3 3 3 3 3 3 U]__  _    r.   rG   c                   $     e Zd Z fdZd Z xZS )PatchMergingc                 p   t                                                       t          |dz            }t          ||ddd          | _        t
          j                                        | _        t          ||ddd|          | _	        t          |d          | _        t          ||ddd          | _        d S )N   r   r      r1   r*   g      ?)r   r   intr   conv1r   r   ReLUactconv2r
   seconv3)r#   dimout_dimhid_dimr,   s       r-   r   zPatchMerging.__init__L   s    cAg,,c7Aq!44
8==??gw1aHHH
--gw1a88


r.   c                     |                      |                     |                     |                     |                     |                     |                                                            }|S N)rc   rb   r`   ra   r^   r#   xs     r-   forwardzPatchMerging.forwardU   sR    JJtwwtxx

488DJJqMM3J3J(K(KLLMMNNr.   r@   rA   rB   r   rk   rD   rE   s   @r-   rX   rX   K   sG        9 9 9 9 9      r.   rX   c                   &     e Zd Zd fd	Zd Z xZS )ResidualDroprI   c                 d    t                                                       || _        || _        d S rh   )r   r   r=   rN   )r#   r=   rN   r,   s      r-   r   zResidualDrop.__init__[   s+    			r.   c           	      t   | j         r| j        dk    r||                     |          t          j        |                    d          ddd|j                                      | j                                      d| j        z
            	                                z  z   S ||                     |          z   S )Nr   r   )device)
trainingrN   r=   r   randr6   rq   ge_divdetachri   s     r-   rk   zResidualDrop.forward`   s    = 	!TY]]tvvayy5:q		1a18$5 $5 $558S^^CCDIDVDVW]W]W_W_` ` ` tvvayy= r.   )rI   rl   rE   s   @r-   rn   rn   Z   sL             
! ! ! ! ! ! !r.   rn   c                   $     e Zd Z fdZd Z xZS )ConvMlpc                     t                                                       t          ||          | _        t          j                                        | _        t          ||d          | _        d S )Nr   r+   )	r   r   r   pw1r   r   r_   r`   pw2)r#   edhr,   s      r-   r   zConvMlp.__init__i   sS    B??8==??Ar!444r.   c                 |    |                      |                     |                     |                              }|S rh   )r|   r`   r{   ri   s     r-   rk   zConvMlp.forwardo   s.    HHTXXdhhqkk**++r.   rl   rE   s   @r-   rx   rx   h   sG        5 5 5 5 5      r.   rx   c                        e Zd ZU eeej        f         ed<   	 	 	 	 	 d fd	Z ej	                    d fd	            Z
d	ej        d
ej        fdZd Z xZS )CascadedGroupAttentionattention_bias_cache   rZ         r   r   r   c                    t                                                       || _        |dz  | _        || _        t          ||z            | _        || _        g }g }t          |          D ]~}	|	                    t          ||z  | j        dz  | j        z                        |	                    t          | j        | j        ||	         d||	         dz  | j                             t          j                            |          | _        t          j                            |          | _        t          j                            t          j                                        t          | j        |z  |d                    | _        t'          t)          j        t          |          t          |                              }
t-          |
          }i }g }|
D ]t}|
D ]o}t/          |d         |d         z
            t/          |d         |d         z
            f}||vrt-          |          ||<   |	                    ||                    put          j                            t          j        |t-          |                              | _        |                     dt          j        |                              ||          d	           i | _        d S )
Ng      r1   r   r\   r   rz   attention_bias_idxsF)
persistent)r   r   	num_headsscalekey_dimr]   val_dim
attn_ratiorangeappendr   r   r   
ModuleListqkvsdws
Sequentialr_   projlist	itertoolsproductlenabs	Parameterzerosattention_biasesregister_buffer
LongTensorrV   r   )r#   rd   r   r   r   
resolutionkernelsr   r   ipointsNattention_offsetsidxsp1p2offsetr,   s                    r-   r   zCascadedGroupAttention.__init__   s    	"_
:/00$y!! 	r 	rAKK!3T\A5E5TUUVVVJJxdlGAJ7ST:YZ?cgcopppqqqqH''--	8&&s++H''HMMOOT\I-s1EEE
 
	
 i'j(9(95;L;LMMNNKK 	7 	7B 7 7bebem,,c"Q%"Q%-.@.@A!222034E0F0F%f--f56666	7
 !& 2 25;y#N_J`J`3a3a b b2E4DT4J4J4O4OPQST4U4Ubghhh$&!!!r.   Tc                 r    t                                          |           |r| j        ri | _        d S d S d S rh   )r   trainr   )r#   moder,   s     r-   r   zCascadedGroupAttention.train   sM    d 	+D- 	+(*D%%%	+ 	+ 	+ 	+r.   rq   returnc                     t           j                                        s| j        r| j        d d | j        f         S t          |          }|| j        vr| j        d d | j        f         | j        |<   | j        |         S rh   )r   jit
is_tracingrr   r   r   strr   )r#   rq   
device_keys      r-   get_attention_biasesz+CascadedGroupAttention.get_attention_biases   s    9!! 	9T] 	9(D,D)DEEVJ!:::8<8MaaaQUQiNi8j)*5,Z88r.   c                    |j         \  }}}}|                    t          | j                  d          }g }|d         }|                     |j                  }	t          t          | j        | j                            D ]A\  }
\  }}|
dk    r|||
         z   } ||          }|	                    |d||          
                    | j        | j        | j        gd          \  }}} ||          }|                    d          |                    d          |                    d          }}}|| j        z  }|                    dd          |z  }||	|
         z   }|                    d          }||                    dd          z  }|	                    || j        ||          }|                    |           C|                     t'          j        |d                    }|S )Nr   )rd   r   rT   r1   )r7   chunkr   r   r   rq   	enumeratezipr   rV   splitr   r   flattenr   	transposesoftmaxr   r   r   cat)r#   rj   BCHWfeats_in	feats_outfeat	attn_biashead_idxqkvr   qkvattns                    r-   rk   zCascadedGroupAttention.forward   s   W
1a773ty>>q711	{--ah77	$-c$)TX.F.F$G$G 	# 	# HjsC!||hx003t99Dii2q!,,22DL$,PTP\3]cd2eeGAq!AAiillAIIaLL!))A,,!qADJA;;r2&&*D)H--D<<B<''Dt~~b"---D99Qa33DT""""IIei	1--..r.   )r   rZ   r   r   T)r@   rA   rB   r   r   r   Tensor__annotations__r   rC   r   rq   r   rk   rD   rE   s   @r-   r   r   t   s         sEL01111	  (' (' (' (' (' ('T U]__+ + + + + _+
95< 9EL 9 9 9 9      r.   r   c                   4     e Zd ZdZ	 	 	 	 	 d	 fd	Zd Z xZS )
LocalWindowAttentiona   Local Window Attention.

    Args:
        dim (int): Number of input channels.
        key_dim (int): The dimension for query and key.
        num_heads (int): Number of attention heads.
        attn_ratio (int): Multiplier for the query dim for value dimension.
        resolution (int): Input resolution.
        window_resolution (int): Local window resolution.
        kernels (List[int]): The kernel size of the dw conv on query.
    r   rZ   r      r   c                     t                                                       || _        || _        || _        |dk    s
J d            || _        t          ||          }t          ||||||          | _        d S )Nr   z"window_size must be greater than 0)r   r   r   )	r   r   rd   r   r   window_resolutionminr   r   	r#   rd   r   r   r   r   r   r   r,   s	           r-   r   zLocalWindowAttention.__init__   s     	"$ 1$$$&J$$$!2 1:>>*)!(	
 
 
			r.   c           	         | j         x}}|j        \  }}}}t          ||k    d||f d||f            t          ||k    d||f d||f            || j        k    r"|| j        k    r|                     |          }n|                    dddd          }| j        || j        z  z
  | j        z  }| j        || j        z  z
  | j        z  }	t          j        j        	                    |ddd|	d|f          }||z   ||	z   }}
|
| j        z  }|| j        z  }|
                    ||| j        || j        |                              dd          }|                    ||z  |z  | j        | j        |                              dddd          }|                     |          }|                    dddd          
                    |||| j        | j        |          }|                    dd                              ||
||          }|d d d |d |f                                         }|                    dddd          }|S )Nz%input feature has wrong size, expect z, got r   r1   r[   r   )r   r7   r   r   r   permuter   r   
functionalr(   rV   r   reshape
contiguous)r#   rj   r   r   r   r   H_W_pad_bpad_rpHpWnHnWs                 r-   rk   zLocalWindowAttention.forward   sk   Aw1b"RY!QYYPRTVxYYZZZRY!QYYPRTVxYYZZZ&&&10F+F+F		!AA		!Q1%%A+a$2H.HHDLbbE+a$2H.HHDLbbE#''Aq!UAu+EFFAYE	Bt--Bt--Bq"d4b$:PRSTT^^_`bcddA		!b&2+t'=t?UWXYYaabcefhiklmmA		!A		!Q1%%**1b"d6LdNdfghhAAq!!))!RQ77A!!!RaR!)''))A		!Q1%%Ar.   )r   rZ   r   r   r   r@   rA   rB   __doc__r   rk   rD   rE   s   @r-   r   r      sg        
 
  
 
 
 
 
 
0      r.   r   c                   8     e Zd ZdZddddg df fd	Zd Z xZS )	EfficientVitBlocka   A basic EfficientVit building block.

    Args:
        dim (int): Number of input channels.
        key_dim (int): Dimension for query and key in the token mixer.
        num_heads (int): Number of attention heads.
        attn_ratio (int): Multiplier for the query dim for value dimension.
        resolution (int): Input resolution.
        window_resolution (int): Local window resolution.
        kernels (List[int]): The kernel size of the dw conv on query.
    r   rZ   r   r   r   c                     t                                                       t          t          ||ddd|d                    | _        t          t          |t          |dz                                | _        t          t          |||||||                    | _	        t          t          ||ddd|d                    | _
        t          t          |t          |dz                                | _        d S )Nr[   r   rI   )r*   r+   r1   )r   r   r   r   )r   r   rn   r   dw0rx   r]   ffn0r   mixerdw1ffn1r   s	           r-   r   zEfficientVitBlock.__init__  s     	c1a3WY Z Z Z[[ c#'ll!;!;<<	! Wi%%"3  
 

  c1a3WY Z Z Z[[ c#'ll!;!;<<			r.   c                     |                      |                     |                     |                     |                     |                                                  S rh   )r   r   r   r   r   ri   s     r-   rk   zEfficientVitBlock.forward;  sB    yy$**TYYtxx{{-C-C"D"DEEFFFr.   r   rE   s   @r-   r   r     sr        
 
  LL= = = = = =8G G G G G G Gr.   r   c                   8     e Zd Zdddddg ddf fd	Zd	 Z xZS )
EfficientVitStage r   r   rZ   r   r   r   r   c                     t                                                       |d         dk    rd|dz
  |d         z  dz   | _        g }|                    dt          j                            t          t          ||ddd|                    t          t          |t          |dz                                          f           |                    dt          ||          f           |                    d	t          j                            t          t          ||ddd|                    t          t          |t          |dz                                          f           t          j        t          |                    | _        n'||k    sJ t          j                    | _        || _        g }t          |
          D ]/}|                    t!          ||||| j        ||	                     0t          j        | | _        d S )
Nr   	subsampler   res1r[   r\   r1   
patchmergeres2)r   r   r   r   r   r   r   rn   r   rx   r]   rX   r   
downsampleIdentityr   r   blocks)r#   in_dimre   r   r   r   r   r   r   r   depthdown_blocksr   dr,   s                 r-   r   zEfficientVitStage.__init__@  s    	a=K'')A~*Q-?!CDOK## &&!Q&!Q!Q!QRR VaZ!A!ABB      l67.K.KLMMM## '7Aq!G!T!T!TUU #gk2B2B!C!CDD      !mK,D,DEEDOOW$$$$ kmmDO(DOu 	C 	CAMM+GWiUYUdfw  zA  B  B  C  C  C  CmV,r.   c                 Z    |                      |          }|                     |          }|S rh   )r   r   ri   s     r-   rk   zEfficientVitStage.forwardk  s'    OOAKKNNr.   rl   rE   s   @r-   r   r   ?  sh          LL)- )- )- )- )- )-V      r.   r   c                        e Zd Z fdZ xZS )PatchEmbeddingc           
         t                                                       |                     dt          ||dz  ddd                     |                     dt          j                                                   |                     dt          |dz  |dz  ddd                     |                     d	t          j                                                   |                     d
t          |dz  |dz  ddd                     |                     dt          j                                                   |                     dt          |dz  |ddd                     d| _        d S )Nr^   r   r[   r1   r   relu1ra   rZ   relu2rc   relu3conv4   )r   r   
add_moduler   r   r   r_   
patch_size)r#   in_chansrd   r,   s      r-   r   zPatchEmbedding.__init__r  s&   (C1HaA!F!FGGG111#(C1HaA!F!FGGG111#(C1HaA!F!FGGG111#(CAq!A!ABBBr.   )r@   rA   rB   r   rD   rE   s   @r-   r   r   q  s8        	 	 	 	 	 	 	 	 	r.   r   c                   ,    e Zd Z	 	 	 	 	 	 	 	 	 	 	 	 d fd	Zej        j        d             Zej        j        dd            Zej        j        dd            Z	ej        j        de
j        fd            Zd dedee         fdZd ZddefdZd Z xZS )!r      r[     @         r  r  r  r   r1   r[   rZ   rZ   rZ   r   r   r   r   r   r   r1   r  avgrI   c                 
   t          t          |                                            d| _        || _        || _        t          |d                   | _        | j        j        }|| j        j        z  }fdt          t                              D             }g | _        g }d         }t          t          ||||
                    D ]\  }\  }}}}}}}t          |||||||||	|
  
        }|}|d         dk    r|dk    r||d         z  }|j        }|                    |           | xj        t#          ||d|           gz  c_        t%          j        | | _        |d	k    rt+          |d
          | _        n |dk    sJ t%          j                    | _        d         x| _        | _        |dk    rt5          | j        || j                  nt6          j                                        | _        d S )NFr   c                 D    g | ]}|         |         |         z  z  S  r  ).0r   	embed_dimr   r   s     r-   
<listcomp>z-EfficientVitMsra.__init__.<locals>.<listcomp>  s/    ```Qilgaj9Q<&?@```r.   )
r   re   r   r   r   r   r   r   r   r   r   r   zstages.)num_chs	reductionmoduler  T	pool_typer   rT   rN   )r   r   r   grad_checkpointingnum_classes	drop_rater   patch_embedr  r   r   feature_infor   r   r   r   r   dictr   r   stagesr   global_poolr   num_featureshead_hidden_sizerG   r   head)r#   img_sizer	  r%  r  r   r   r   window_sizer   down_opsr+  r&  r'   r   r   r*  pre_edr   r}   kddpthnharwddostager,   s       `` `                   r-   r   zEfficientVitMsra.__init__  s]    	%%..000"'&" *(IaLAA!,!1!<<
``````%PST]P^P^J_J_```
 11:Iwy*kS[\\2^ 2^ 	\ 	\-A-Bb"b"%%"$  E F!u##Q"Q%)JMM%   $rVMVWMM"Z"Z"Z![[mV,%3kSWXXXD!####!{}}D4=bMAD1DORSOO {A A A AY^YaYjYjYlYl 				r.   c                 b    d |                                                                  D             S )Nc                     h | ]}d |v |	S )r   r  )r  rj   s     r-   	<setcomp>z3EfficientVitMsra.no_weight_decay.<locals>.<setcomp>  s#    OOOa7IQ7N7N7N7N7Nr.   )
state_dictkeysr#   s    r-   no_weight_decayz EfficientVitMsra.no_weight_decay  s,    OO4??,,1133OOOOr.   Fc                 4    t          d|rdnddg          }|S )Nz^patch_embedz^stages\.(\d+))z^stages\.(\d+).downsample)r   )z^stages\.(\d+)\.\w+\.(\d+)N)stemr   )r)  )r#   coarsematchers      r-   group_matcherzEfficientVitMsra.group_matcher  s9     (. $$455
 
 
 r.   Tc                     || _         d S rh   )r$  )r#   enables     r-   set_grad_checkpointingz'EfficientVitMsra.set_grad_checkpointing  s    "(r.   r   c                     | j         j        S rh   )r.  rP   r?  s    r-   get_classifierzEfficientVitMsra.get_classifier  s    yr.   Nr%  r+  c                    || _         |=|dk    rt          |d          | _        n |dk    sJ t          j                    | _        |dk    rt          | j        || j                  nt          j                                        | _	        d S )Nr  Tr!  r   r#  )
r%  r   r+  r   r   rG   r,  r&  r   r.  )r#   r%  r+  s      r-   reset_classifierz!EfficientVitMsra.reset_classifier  s    &"e###7+W[#\#\#\  "a''''#%;== DORSOO {A A A AY^YaYjYjYlYl 				r.   c                     |                      |          }| j        r4t          j                                        st          | j        |          }n|                     |          }|S rh   )r'  r$  r   r   is_scriptingr   r*  ri   s     r-   forward_featuresz!EfficientVitMsra.forward_features  s\    Q" 	59+A+A+C+C 	t{A..AAAAr.   
pre_logitsc                 ^    |                      |          }|r|n|                     |          S rh   )r+  r.  )r#   rj   rP  s      r-   forward_headzEfficientVitMsra.forward_head  s.    Q0qqDIIaLL0r.   c                 Z    |                      |          }|                     |          }|S rh   )rO  rR  ri   s     r-   rk   zEfficientVitMsra.forward  s-    !!!$$a  r.   )r  r[   r  r  r  r  r  r  r   r  r  rI   Fr   rh   )r@   rA   rB   r   r   r   ignorer@  rE  rH  r   ModulerJ  r]   r   r   rL  rO  boolrR  rk   rD   rE   s   @r-   r   r   ~  sy        $ ! B;m ;m ;m ;m ;m ;mz YP P P Y    Y) ) ) ) Y 	        	m 	mC 	mhsm 	m 	m 	m 	m  1 1$ 1 1 1 1      r.   r   c           	      0    | dt           t          ddddd|S )Nr  zpatch_embed.conv1.convzhead.linearT)rZ   rZ   )urlr%  meanrK   
first_conv
classifierfixed_input_size	pool_sizer   )rY  kwargss     r-   _cfgr`    s3    %#.# 
 
 
 
r.   ztimm/)	hf_hub_id)zefficientvit_m0.r224_in1kzefficientvit_m1.r224_in1kzefficientvit_m2.r224_in1kzefficientvit_m3.r224_in1kzefficientvit_m4.r224_in1kzefficientvit_m5.r224_in1kFc                 |    |                     dd          }t          t          | |fdt          d|          i|}|S )Nout_indices)r   r   r1   feature_cfgT)flatten_sequentialrc  )popr   r   r)  )variant
pretrainedr_  rc  models        r-   _create_efficientvit_msrarj  <  sZ    **]I66K   DkJJJ	
  E Lr.   c           	      r    t          dg dg dg dg dg d          }t          d
d	| it          |fi |S )Nr  r  r  r  r  r   r/  r  r   r   r0  r   efficientvit_m0rh  )rm  r)  rj  rh  r_  
model_argss      r-   rm  rm  H  sj     ..ii))II  J %ll:lQUV`QkQkdjQkQklllr.   c           	      r    t          dg dg dg dg dg d          }t          d
d	| it          |fi |S )Nr  )r     r  r  )r1   r[   r[   r  r   r   r[   r[   rl  efficientvit_m1rh  )rt  rn  ro  s      r-   rt  rt  U  j    !//ii))II  J %ll:lQUV`QkQkdjQkQklllr.   c           	      r    t          dg dg dg dg dg d          }t          d
d	| it          |fi |S )Nr  )r  r  r  r  )rZ   r[   r1   r  rs  rl  efficientvit_m2rh  )rw  rn  ro  s      r-   rw  rw  b  ru  r.   c           	      r    t          dg dg dg dg dg d          }t          d
d	| it          |fi |S )Nr  )r     i@  r  )rZ   r[   rZ   r  r   rl  efficientvit_m3rh  )rz  rn  ro  s      r-   rz  rz  o  ru  r.   c           	      r    t          dg dg dg dg dg d          }t          d
d	| it          |fi |S )Nr  )r       r  r  r  rs  rl  efficientvit_m4rh  )r~  rn  ro  s      r-   r~  r~  |  ru  r.   c           	      r    t          dg dg dg dg dg d          }t          d
d	| it          |fi |S )Nr  )r  i   r}  )r   r[   rZ   )r[   r[   rZ   r  rs  rl  efficientvit_m5rh  )r  rn  ro  s      r-   r  r    ru  r.   )r   rT  )0r   __all__r   collectionsr   typingr   r   r   torch.nnr   	timm.datar   r	   timm.layersr
   r   r   r   _builderr   _manipulater   	_registryr   r   r   r   rG   rV  rX   rn   rx   r   r   r   r   r   r   r`  default_cfgsrj  rm  rt  rw  rz  r~  r  r  r.   r-   <module>r     s    
     # # # # # # ! ! ! ! ! ! ! !        A A A A A A A A S S S S S S S S S S S S * * * * * * ' ' ' ' ' ' < < < < < < < <    ux"   .    $   8    58?   ! ! ! ! !58? ! ! !	 	 	 	 	eho 	 	 	[ [ [ [ [UX_ [ [ [|> > > > >58? > > >B)G )G )G )G )G )G )G )GX/ / / / / / / /d
 
 
 
 
UX( 
 
 
o o o o ory o o oh    %$!%" " " "&" " " "&" " " "&" " " "&" " " "&" " "+& &  8	 	 	 	 	m 	m 	m 	m 	m 	m 	m 	m 	m 	m 	m 	m 	m 	m 	m 	m 	m 	m 	m 	m 	m 	m 	m 	m 	m 	mr.   