
    RЦiI                        S r SSKrSSKJr  SSKJr  SSKJrJrJ	r	J
r
Jr  SSKrSSKJr  SSKJrJrJr  SSKJrJrJrJrJr  SS	KJr  SS
KJr  SSKJrJr  S/r S*S\RB                  S\"S\#4S jjr$ " S S\RB                  5      r% " S S\RB                  5      r& " S S\&5      r' " S S\RB                  5      r( " S S\RB                  5      r) " S S\RB                  5      r* " S S\RB                  5      r+ " S S\RB                  5      r,S  r-S+S! jr.S,S" jr/\" \/" S#S$9\/" S#S$9\/" S#S$9S%.5      r0\S+S&\,4S' jj5       r1\S+S&\,4S( jj5       r2\S+S&\,4S) jj5       r3g)-zjSequencer

Paper: `Sequencer: Deep LSTM for Image Classification` - https://arxiv.org/pdf/2205.01972.pdf

    N)partial)
accumulate)ListOptionalTupleTypeUnion)IMAGENET_DEFAULT_MEANIMAGENET_DEFAULT_STDDEFAULT_CROP_PCT)lecun_normal_DropPathMlp
PatchEmbedClassifierHead   )build_model_with_cfg)named_apply)register_modelgenerate_default_cfgsSequencer2dmodulename	head_biasc                 R   [        U [        R                  5      (       GaN  UR                  S5      (       aT  [        R                  R                  U R                  5        [        R                  R                  U R                  U5        g U(       aM  [        U R                  5        U R                  b*  [        R                  R                  U R                  5        g g [        R                  R                  U R                  5        U R                  bY  SU;   a)  [        R                  R                  U R                  SS9  g [        R                  R                  U R                  5        g g [        U [        R                  5      (       aM  [        U R                  5        U R                  b*  [        R                  R                  U R                  5        g g [        U [        R                  [        R                  [        R                  45      (       aS  [        R                  R!                  U R                  5        [        R                  R                  U R                  5        g [        U [        R"                  [        R$                  [        R&                  45      (       a\  S[(        R*                  " U R,                  5      -  nU R/                  5        H$  n[        R                  R1                  XT* U5        M&     g [3        U S5      (       a  U R5                  5         g g )Nheadmlpư>)stdg      ?init_weights)
isinstancennLinear
startswithinitzeros_weight	constant_biasr   xavier_uniform_normal_Conv2d	LayerNormBatchNorm2d	GroupNormones_RNNGRULSTMmathsqrthidden_size
parametersuniform_hasattrr    )r   r   r   flaxstdvr'   s         T/var/www/html/ai-image-ml/venv/lib/python3.13/site-packages/timm/models/sequencer.py_init_weightsr=      s   &"))$$??6""GGNN6==)GGfkk95fmm,;;*GGNN6;;/ + ''6;;*}>v{{3	 +
 
FBII	&	&fmm$;;"GGNN6;;' #	FR\\2>>2<<H	I	I
fmm$
v{{#	FRVVRVVRWW5	6	6TYYv1122'')FGGVUD1 *		(	( 
)    c                   l   ^  \ rS rSrU 4S jrS\R                  S\\R                  S4   4S jrSr	U =r
$ )RNNIdentity;   c                 "   > [         TU ]  5         g Nsuper__init__)selfargskwargs	__class__s      r<   rF   RNNIdentity.__init__<       r>   xreturnNc                 
    US 4$ rC    rG   rM   s     r<   forwardRNNIdentity.forward?   s    $wr>   rP   )__name__
__module____qualname____firstlineno__rF   torchTensorr   rR   __static_attributes____classcell__rJ   s   @r<   r@   r@   ;   s2     %d0B*C  r>   r@   c                   `   ^  \ rS rSr       SS\S\S\S\S\S\S\4U 4S	 jjjrS
 rSr	U =r
$ )	RNN2dBaseC   
input_sizer6   
num_layersr)   bidirectionalunionwith_fcc
                   > XS.n
[         TU ]  5         Xl        X l        U(       a  SU-  OUU l        X`l        SU l        SU l        Xpl        S U l	        U(       a  US:X  a,  [        R                  " SU R                  -  U40 U
D6U l	        GOUS:X  a)  [        R                  " U R                  U40 U
D6U l	        GOiUS:X  a0  [        R                  " U R                  U40 U
D6U l	        SU l        GO3US:X  a/  [        R                  " U R                  U40 U
D6U l	        SU l        O[        S	U-   5      eUS:X  a4   SU R                  -  U:w  a  [        S
SU R                  -   SU S35      eOUS:X  a.   U R                  U:w  a  [        S
U R                   SU S35      eOUS:X  a4  U R                  U:w  a  [        S
U R                   SU S35      eSU l        OHUS:X  a4  U R                  U:w  a  [        S
U R                   SU S35      eSU l        O[        S	U-   5      e[        5       U l        [        5       U l        g )Ndevicedtype   TcataddverticalF
horizontalzUnrecognized union: zThe output channel z% is different from the input channel .)rE   rF   r`   r6   output_sizerc   with_verticalwith_horizontalrd   fcr"   r#   
ValueErrorr@   rnn_vrnn_hrG   r`   r6   ra   r)   rb   rc   rd   rg   rh   ddrJ   s              r<   rF   RNN2dBase.__init__E   sP    /$&.;1{?
!#~))A(8(8$8*KK%))D$4$4jGBG*$))D$4$4jGBG',$,&))D$4$4jGBG%*" !7%!?@@e^4###z1 #6q4;K;K7K6LLqr|q}}~!  A  A 2e^:- #6t7G7G6HHmnxmyyz!{|| .j :- #6t7G7G6HHmnxmyyz!{||#(D l":- #6t7G7G6HHmnxmyyz!{||!&D3e;<< ]
 ]
r>   c                 Z   UR                   u  p#pEU R                  (       aa  UR                  SSSS5      nUR                  SX55      nU R	                  U5      u  pgUR                  X$US5      nUR                  SSSS5      nOS nU R
                  (       a9  UR                  SXE5      nU R                  U5      u  pUR                  X#US5      nOS nUb/  Ub,  U R                  S:X  a  [        R                  " Xh/SS9nOXh-   nOUb  UnOUb  UnU R                  b  U R                  U5      nU$ )Nr   ri   r      rj   )dim)shaperp   permutereshapert   rq   ru   rc   rX   rj   rr   )	rG   rM   BHWCv_hs	            r<   rR   RNN2dBase.forward   s   WW
a		!Q1%A		"a#A::a=DA		!2&A		!Q1%AA		"a#A::a=DA		!2&AA=Q]zzU"IIqf"-E]A]A77
Ar>   )
rr   r6   r`   ro   ru   rt   rc   rd   rq   rp   r   TTrj   TNN)rT   rU   rV   rW   intboolstrrF   rR   rZ   r[   r\   s   @r<   r^   r^   C   ss      "& :#:# :# 	:#
 :#  :# :# :# :#x   r>   r^   c                   Z   ^  \ rS rSr       SS\S\S\S\S\S\S\4U 4S	 jjjrS
rU =r	$ )LSTM2d   r`   r6   ra   r)   rb   rc   rd   c
                   > XS.n
[         TU ]  XX4XVXxU	5	        U R                  (       a"  [        R                  " UUU4SUUS.U
D6U l        U R                  (       a#  [        R                  " UUU4SUUS.U
D6U l        g g )Nrf   T)batch_firstr)   rb   )rE   rF   rp   r"   r3   rt   rq   ru   rv   s              r<   rF   LSTM2d.__init__   s     /*MZakpq !+ DJ  !+ DJ  r>   )ru   rt   r   )
rT   rU   rV   rW   r   r   r   rF   rZ   r[   r\   s   @r<   r   r      sh      "& !! ! 	!
 !  ! ! ! !r>   r   c                     ^  \ rS rSrS\\\" \R                  SS9\R                  SSSSSSS	S	4S
\
S\
S\S\\R                     S\\R                     S\\R                     S\\R                     S\
S\S\S\S\S\4U 4S jjjrS rSrU =r$ )Sequencer2dBlock         @r   epsr   Trj           Nr|   r6   	mlp_ratio	rnn_layer	mlp_layer
norm_layer	act_layerra   rb   rc   rd   drop	drop_pathc                   > XS.n[         TU ]  5         [        X1-  5      nU" U40 UD6U l        U" UU4UU	U
US.UD6U l        US:  a  [        U5      O[        R                  " 5       U l        U" U40 UD6U l	        U" UU4X|S.UD6U l
        g )Nrf   )ra   rb   rc   rd   r   )r   r   )rE   rF   r   norm1
rnn_tokensr   r"   Identityr   norm2mlp_channels)rG   r|   r6   r   r   r   r   r   ra   rb   rc   rd   r   r   rg   rh   rw   channels_dimrJ   s                     r<   rF   Sequencer2dBlock.__init__   s    $ /9?+*r*
#
 "'
 
 1:B),BKKM*r*
%c<^9^[]^r>   c                     XR                  U R                  U R                  U5      5      5      -   nXR                  U R                  U R	                  U5      5      5      -   nU$ rC   )r   r   r   r   r   rQ   s     r<   rR   Sequencer2dBlock.forward   sM    ttzz!}=>>t00A?@@r>   )r   r   r   r   r   )rT   rU   rV   rW   r   r   r   r"   r-   GELUr   floatr   Moduler   r   rF   rR   rZ   r[   r\   s   @r<   r   r      s    
  #)/),*1",,D*I)+"& !!!_!_ !_ 	!_
 BII!_ BII!_ RYY!_ BII!_ !_  !_ !_ !_ !_ !_ !_F r>   r   c                   .   ^  \ rS rSrU 4S jrS rSrU =r$ )Shuffle   c                 "   > [         TU ]  5         g rC   rD   )rG   rJ   s    r<   rF   Shuffle.__init__   rL   r>   c                     U R                   (       aW  UR                  u  p#pE[        R                  " X4-  5      nUR	                  USU5      nUS S 2US S 24   R	                  X#US5      nU$ )Nr{   )trainingr}   rX   randpermr   )rG   rM   r   r   r   r   rs          r<   rR   Shuffle.forward   s[    ==JA!qu%A		!R#A!Q'
""1B/Ar>   rP   )rT   rU   rV   rW   rF   rR   rZ   r[   r\   s   @r<   r   r      s     r>   r   c                   F   ^  \ rS rSr  SS\S\S\4U 4S jjjrS rSrU =r$ )	Downsample2di  	input_dim
output_dim
patch_sizec                 f   > XES.n[         TU ]  5         [        R                  " X4X3S.UD6U l        g )Nrf   )kernel_sizestride)rE   rF   r"   r,   down)rG   r   r   r   rg   rh   rw   rJ   s          r<   rF   Downsample2d.__init__  s3     /IIieebde	r>   c                 x    UR                  SSSS5      nU R                  U5      nUR                  SSSS5      nU$ )Nr   rz   r   ri   )r~   r   rQ   s     r<   rR   Downsample2d.forward  s<    IIaAq!IIaLIIaAq!r>   )r   )NN)	rT   rU   rV   rW   r   rF   rR   rZ   r[   r\   s   @r<   r   r     s@     
f
f 
f 	
f 
f r>   r   c            (       ^  ^  \ rS rSrS\\\\" \R                  SS9\R                  SSSSSSS	S	4S
\S\S\S\S\S\S\S\\R                     S\\R                     S\\R                     S\\R                     S\\R                     S\S\S\S\S\S\\\\   4   4$U 4S jjjrS rSrU =r$ )Sequencer2dStagei  Fr   r   r   Trj   r   Nr|   dim_outdepthr   r6   r   
downsampleblock_layerr   r   r   r   ra   rb   rc   rd   r   r   c                   > [         TU ]  5         UUS.nU(       a  [        XU40 UD6U l        O!X:X  d   e[        R
                  " 5       U l        / n[        U5       HH  nUR                  U" UU4UU	U
UUUUUUU[        U[        [        45      (       a  UU   OUS.UD65        MJ     [        R                  " U6 U l        g )Nrf   )r   r   r   r   r   ra   rb   rc   rd   r   r   )rE   rF   r   r   r"   r   rangeappendr!   listtuple
Sequentialblocks)rG   r|   r   r   r   r6   r   r   r   r   r   r   r   ra   rb   rc   rd   r   r   rg   rh   rw   r   	block_idxrJ   s                           r<   rF   Sequencer2dStage.__init__  s    . 	/*3JrJDO>!> kkmDOuIMM+ $##%#%+2<Yu2V2V)I.\e   &" mmV,r>   c                 J    U R                  U5      nU R                  U5      nU$ rC   )r   r   rQ   s     r<   rR   Sequencer2dStage.forwardI  s"    OOAKKNr>   )r   r   )rT   rU   rV   rW   r   r   r   r   r"   r-   r   r   r   r   r   r   r   r	   r   rF   rR   rZ   r[   r\   s   @r<   r   r     s;     %+;)/),*1",,D*I)+"& 35+1-1- 1- 	1-
 1- 1- 1- 1- bii1- BII1- BII1- RYY1- BII1- 1-  1-  !1-" #1-$ %1-& UDK/0'1- 1-f r>   r   c            -         ^  \ rS rSrSSSSSSSS	S
\\\\" \R                  SS9\R                  SSSSSSSSSS4S\S\S\S\S\\S4   S\\S4   S\\S4   S\\S4   S\\S4   S\\R                      S\\R                      S\\R                      S \\R                      S!\\R                      S"\S#\S$\S%\S&\S'\S(\S)\4,U 4S* jjjrS6S+ jr\R*                  R,                  S6S, j5       r\R*                  R,                  S7S- j5       r\R*                  R,                  S.\R                   4S/ j5       rS8S\S\\   4S0 jjrS1 rS6S2\4S3 jjrS4 rS5rU =r $ )9r   iO       rz   avg   rz      rz   )   ri   ri   r        r   r   0   `   r   r   r   r   r   r   r   r   r   Trj   r   FNnum_classesimg_sizein_chansglobal_poollayers.patch_sizes
embed_dimshidden_sizes
mlp_ratiosr   r   r   r   r   num_rnn_layersrb   rc   rd   	drop_ratedrop_path_ratenlhb	stem_normc                   > [         TU ]  5         UUS.nUS;   d   eXl        X0l        X@l        US   =U l        U l        SU l        SU l        / U l	        [        SS US   UUS   U(       a  UOS SSS.UD6U l        [        U5      [        U5      s=:X  a0  [        U5      s=:X  a   [        U5      s=:X  a  [        U	5      :X  d   e   e[        [        US 5      5      n/ nUS   n[        U5       H  u  nnU[!        UUU   40 S	UU   _S
US:  _SUU   _SUU   _SU	U   _SU
_SU_SU_SU_SU_SU_SU_SU_SU_SU_SU_UD6/-  nUU   nU =R                  [#        UUU   SU 3S9/-  sl	        M     [$        R&                  " U6 U l        U" US   40 UD6U l        [-        U R
                  U4UUU R                  S.UD6U l        U R1                  US9  g )Nrf   ) r   r{   NHWCr   F)r   r   r   	embed_dimr   flatten
output_fmtc                 
    X-  $ rC   rP   )rM   ys     r<   <lambda>&Sequencer2d.__init__.<locals>.<lambda>  s    qur>   r   r   r   r6   r   r   r   r   r   r   ra   rb   rc   rd   r   r   stages.)num_chs	reductionr   )	pool_typer   	input_fmt)r   rP   )rE   rF   r   r   r   num_featureshead_hidden_sizefeature_dimr   feature_infor   stemlenr   r   	enumerater   dictr"   r   stagesnormr   r   r    ) rG   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   rb   rc   rd   r   r   r   r   rg   rh   rw   
reductionsr  prev_dimir   rJ   s                                   r<   rF   Sequencer2d.__init__P  s   6 	/k)))& &4>rNBD1  	
"1~ m%.zD	
 	
	 6{c+.i#j/iSEViZ]^hZiiiiii*[2DEF
a=j)DAq'1 Qi q5	
 'q> )O %Q- ( $ $ & $ * ,    !" #$ )'  F* "!}H$x:a=[bcdbeYf"g!hh/ *2 mmV,z"~44	"
 "oo
 
	 	t$r>   c                     U(       a!  [         R                  " U R                  5      * OSn[        [	        [
        US9U S9  g )Nr   )r   )r   )r4   logr   r   r   r=   )rG   r   r   s      r<   r    Sequencer2d.init_weights  s.    37TXXd..//R	GMY?Mr>   c                 4    [        SU(       a  SS/S9$ / SQS9$ )Nz^stem)z^stages\.(\d+)Nz^norm)i ))z^stages\.(\d+)\.blocks\.(\d+)N)z^stages\.(\d+)\.downsample)r   r  )r  r   )r  )rG   coarses     r<   group_matcherSequencer2d.group_matcher  s3      *$

 
	



 
	
r>   c                      U(       a   S5       eg )Nz$gradient checkpointing not supportedrP   )rG   enables     r<   set_grad_checkpointing"Sequencer2d.set_grad_checkpointing  s    AAAz6r>   rN   c                     U R                   $ rC   r   )rG   s    r<   get_classifierSequencer2d.get_classifier  s    yyr>   c                 B    Xl         U R                  R                  XS9  g )N)r   )r   r   reset)rG   r   r   s      r<   reset_classifierSequencer2d.reset_classifier  s    &		;r>   c                 l    U R                  U5      nU R                  U5      nU R                  U5      nU$ rC   )r  r  r	  rQ   s     r<   forward_featuresSequencer2d.forward_features  s.    IIaLKKNIIaLr>   
pre_logitsc                 R    U(       a  U R                  USS9$ U R                  U5      $ )NT)r%  r  )rG   rM   r%  s      r<   forward_headSequencer2d.forward_head  s$    0:tyyty,L		!Lr>   c                 J    U R                  U5      nU R                  U5      nU$ rC   )r#  r'  rQ   s     r<   rR   Sequencer2d.forward  s'    !!!$a r>   )r  r  r   r   r  r   r	  r   r   r   r  r  F)TrC   )!rT   rU   rV   rW   r   r   r   r   r"   r-   r   r   r   r   r   r   r   r   rF   r    rX   jitignorer  r  r  r   r   r#  r'  rR   rZ   r[   r\   s   @r<   r   r   O  sK     $$&2+7*>,<,@+;)/),*1",,D*I)+"#"& !$&#3Y%Y% Y% 	Y%
 Y% #s(OY% sCxY% c3hY%  S/Y% eSj)Y% biiY% BIIY% BIIY% RYYY% BIIY%   !Y%"  #Y%$ %Y%& 'Y%( )Y%* "+Y%, -Y%. /Y% Y%vN YY
 
 YYB B YY		  <C <hsm <M$ M r>   c                     SU ;   a  U $ SU ;   a  U S   n SSK n0 nU R                  5        HB  u  pEUR                  SS U5      nUR                  SSU5      nUR                  S	S
5      nXSU'   MD     U$ )z#Remap original checkpoints -> timm zstages.0.blocks.0.norm1.weightmodelr   Nzblocks.([0-9]+).([0-9]+).downc                 D    S[        U R                  S5      5      S-    S3$ )Nr   r   z.downsample.down)r   group)rM   s    r<   r   &checkpoint_filter_fn.<locals>.<lambda>  s"    QWWUVZ[\I\H]]m?nr>   zblocks.([0-9]+).([0-9]+)zstages.\1.blocks.\2zhead.zhead.fc.)reitemssubreplace)
state_dictr/  r3  out_dictkr   s         r<   checkpoint_filter_fnr:    s    ':5*(
H  "FF35npqrFF.0FJIIgz*	 # Or>   c           	          [        [        S5      5      nUR                  SU5      n[        [        U U4[
        [        SUS9S.UD6nU$ )Nrz   out_indicesT)flatten_sequentialr<  )pretrained_filter_fnfeature_cfg)r   r   popr   r   r:  r  )variant
pretrainedrI   default_out_indicesr<  r/  s         r<   _create_sequencer2drD    sY    a/**],?@K  2DkJ E Lr>   c                 <    U SSS [         SS[        [        SSSS.UE$ )	Nr   )rz   r   r   bicubicTz	stem.projzhead.fcz
apache-2.0)urlr   r`   	pool_sizecrop_pctinterpolationfixed_input_sizemeanr   
first_conv
classifierlicense)r   r
   r   )rG  rI   s     r<   _cfgrP    s6    =t$yVZ%.B!  r>   ztimm/)	hf_hub_id)zsequencer2d_s.in1kzsequencer2d_m.in1kzsequencer2d_l.in1krN   c                 r    [        / SQ/ SQ/ SQ/ SQ/ SQ[        SSSS9	n[        S
S	U 0[        U40 UD6D6nU$ )Nr   r   ri   r   r   r   r   r   Trj   	r   r   r   r   r   r   rb   rc   rd   rB  )sequencer2d_sr  r   rD  rB  rI   
model_argsr/  s       r<   rU  rU    sK     '%'
J  eJe$zJd]cJdeELr>   c                 z    [        S
/ SQ/ SQ/ SQ/ SQ/ SQ[        SSSS.	UD6n[        SS	U 0[        U40 UD6D6nU$ )N)r   rz      rz   rS  r   r   r   Trj   rT  rB  rP   )sequencer2d_mrV  rW  s       r<   r[  r[    Z     
 '%'
 
J  eJe$zJd]cJdeELr>   c                 z    [        S
/ SQ/ SQ/ SQ/ SQ/ SQ[        SSSS.	UD6n[        SS	U 0[        U40 UD6D6nU$ )N)r   r      r   rS  r   r   r   Trj   rT  rB  rP   )sequencer2d_lrV  rW  s       r<   r_  r_  .  r\  r>   )r   Fr+  )r   )4__doc__r4   	functoolsr   	itertoolsr   typingr   r   r   r   r	   rX   torch.nnr"   	timm.datar
   r   r   timm.layersr   r   r   r   r   _builderr   _manipulater   	_registryr   r   __all__r   r   r   r=   r@   r^   r   r   r   r   r   r   r:  rD  rP  default_cfgsrU  r[  r_  rP   r>   r<   <module>rl     sj       5 5   S S P P * $ </")) 3 5 B")) ^		 ^B#Y #L'ryy 'T
bii 
299 (7ryy 7tF")) FR$	 %111&            r>   