
    ^j              	          d Z ddlZddl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ZddlmZ ddlmc mZ ddlmZmZmZmZmZmZmZmZ ddlmZ ddlmZ ddl m!Z! dd	l"m#Z#m$Z$ d
dl%m&Z&m'Z' ddgZ(dZ)dZ*de+de+dee+   fdZ,	 	 d>de+de-dej\                  fdZ/	 	 d>de+de-dej\                  fdZ0 G d dejb                        Z2 G d dejb                        Z3d?de+fdZ4 G d dejb                        Z5 G d  d!ejb                        Z6 G d" d#ejb                        Z7 G d$ d%ejb                        Z8 G d& d'ejb                        Z9 G d( d)ejb                        Z: G d* d+ejb                        Z; G d, dejb                        Z<d@d-Z= e' e=d./       e=d.d0d12       e=d.d3d14      d5      Z>d6e?d7ejb                  de?fd8Z@dAd9eAd:e-de<fd;ZBe&dAd:e-de<fd<       ZCe&dAd:e-de<fd=       ZDy)Ba  CSATv2

A frequency-domain vision model using DCT transforms with spatial attention.

Paper: TBD

This model created by members of MLPA Lab. Welcome feedback and suggestion, questions.
gusdlf93@naver.com
juno.demie.oh@gmail.com

Refined for timm by Ross Wightman
    N)partialreduce)ListOptionalTupleUnion)trunc_normal_DropPathMlpLayerNorm2d	AttentionNormMlpClassifierHead
LayerScaleLayerScale2d)GlobalResponseNorm)build_model_with_cfg)feature_take_indices)
checkpointcheckpoint_seq   )register_modelgenerate_default_cfgsCSATv2csatv2))@g;i#@g_LegmV}b?gZӼiUMu>g{g ]iUMu?gh㈵4?g	k?g)t^cwg-C6*?r   gW8g^/vgdS       g!J>h㈵ga2U0*3g	3mJ?g	.V`Z?9̗iUMugמY"gǺ6h㈵>g/nbg8*5{5?h㈵iUMug-C6*מY?r   -C6
?gGŧ h?h㈵r   gW(?h㈵?r   ga2U0*#r"   r$   g-C6?r&   h㈵?        r(   h㈵ga2U0*#?h㈵>r$   r   r'   r*   r'   r   r+   r   r   r   r#   r   )@g o_@gnqgŏ1w-?gWX?r   g'>i?gWr*   gC8
!g/$?g_LE?r   r   r   g8*5{e?r)   r!   r$   r    r   gVIk?gQ,Z?r'   r   r*   r(   r   gH}M?g9̗'?r    r*   r$   r#   r*   r   gyCnK?r   r*   r)   r'   r*   r"   r)   r!   r+   r!   r'   K8?r)   r#   r,   r"   r!   r!   r!   r!   מYr#   r!   r)   r!   r+   r!   r)   )@gQ1ߤt@gP6
ragVF摷g~tgy^?C8
!?gq@H6?gkC8Sr+   r,   g׆q&g%>?r-   r*   g-C6
gH}]?gyCn;?r+   g-C6r'   r   gcbqmhg{GzT?r"   r#   r*   r*   r*   gמYB?gyCn+?r!   r*   r+   r*   r*   r%   gǺFr"   r"   r%   r&   r   r+   r"   r!   r)   r#   r'   r$   r!   r)   r.   r   r)   r+   r   r*   r&   r#   r)   r!   r*   r!   r   r)   ))@g   AgN@@g->V@gr3܀oԙ@gߖ@g|~!\@ge@gwR~R@g[ D@gMu@gqZq@g-r@gCiw@gqu@gP6bw@gZB>g@gFg@g&:k@gx=\ri@g"[='<b@gQ_@gCVO@g\P@gqY@gK87`@gJ_{A`@ggY@geI)pX@glIFJ@gϽgK@gN],_R@gdBT@gX9vVQ@gH@gi@@gctv2B@gz1}4@g8@g\Ɏ@@@g.lIFF@g %G@gFxD@gꐛn>@gx@ٔ6@g2d:@go%;6:@gd5@g2bj1@gyt].@g!Yd%@g%zr{&@gt)@g?x-.@g.@gHm!@giR
@gMۿҤ@g~@gJF@gؙB@g%X
@gXV@g'>O@ge1?)@g[|
e@ga_Yt@g%;6qr@gQkw^S@g:dwS@g͍	CV@gHȰ@@g2_A@gNG@@g1%2<@g-(@ga7l[,@go1u2@g*Wx0@g]3f[/@g5Ry;@gqrCQ !@g 4$@g72"@g9#J{@gdz@g_5j@gt^>@gq@gіs)j@gxqZ@gEdXY@gЛT[@gN#-7?gN#-@g@g~:p@g5Ry;@g⪲@gJ5o?gsA?gT7?g	;?g]h?gTt<f?gӤt{I?g1Zd?g:ǀ?g';?g {?g2}ƅ?g<!?g^D?gnQfL2?g^D?g|wJ?grt?göE2?gQ?g"^F?gʦ\?g($?gj?gQ?g3d?g~k,	?g6[ ?gLuT5?gu7Ou?)@g[|
ٹ@g8Mr@gM;p@g^P@gmO@g+S@g7<@g*;@g,yp9@g_L7@g#@gnض(&@gqh,@gr)@g9]c)@gEV@g*Ph@gM@g㥛 P@grO@gr/@g@g@@g3.H@g3ı.n@geS@g
@gg?RD@gV	?gd`TR?gu7@gӂ}@g[%X@gs?ga)?g?gû\wb?gM֨ht?gO@a?g>??g]o%?gjt?gdu?g*<?g=~o?g~jt?gd?gmJR?gaۢ?gZM?g7qrCQ?gf?gm?g\='o|?gCVzN?g;6~?gE_A?gsFZ*o?go%;6?g~:p?g|Pk?gO|?gKqUw?goʡ?rowscolsreturnc                    t        j                  d| |z  d      j                  | |      j                         }t	        | |z   dz
        D cg c]  }g  }}g }t	        |       D ]U  }t	        |      D ]E  }||z   }|dz  dk(  r||   j                  d||   |          ,||   j                  ||   |          G W |D ]  }	|j                  |	        |S c c}w )z0Generate zigzag scan order for DCT coefficients.r   r      )nparangereshapetolistrangeinsertappendextend)
r/   r0   
idx_matrix_diazigzagijsds
             ]/var/www/ramen.bs-engineer-server.com/venv/lib/python3.12/site-packages/timm/models/csatv2.py_zigzag_permutationrE   X   s    1dTk1-55dDAHHJJTD[1_-
.!2
.C
.F4[ 0t 	0AAA1uzAaAq!12AjmA./	00  aM /s   	Ckernel_sizeorthonormalc                    t        ||      }t        j                  | fi |}|j                         j	                         j                  d|       }t        j                  ||j                  dg      gd      }t        j                  j                  |d      ddd| f   }t        j                  d|t        j                        t        j                  z  t        j                  | |t        j                        dddf   z  }t        j                  || dz  z        }||z  }|j                  }|r|dddf   t        j                   t        j                  d| d	z  z  fi |      z  |dddf<   |ddddf   t        j                   t        j                  d| dz  z  fi |      z  |ddddf<    |j	                         j
                  |j"                   }|S )
z#Generate Type-II DCT kernel matrix.devicedtyper   dimNy             r3   r      )dicttorcheyeclone
contiguousviewcatflipffttensor	complex64pir5   longexprealsqrtshape)rF   rG   rJ   rK   ddxvks           rD   _dct_kernel_type_2re   i   s    
V5	)B		+$$A		##B4A		1affaSk"+A		aR L[L1ASu?%((J
,,{6
DT1W
M	N  			!{Q'(A	AA	AAqD'EJJu||Aq4I'PR'PQQ!Q$QU8ejja;?6K)Rr)RSS!QR%QWW%AH    c                 X    t         j                  j                  t        | |||            S )z9Generate Type-III DCT kernel matrix (inverse of Type-II).)rQ   linalginvre   )rF   rG   rJ   rK   s       rD   _dct_kernel_type_3rj      s&     <<.{KQVWXXrf   c            	       v     e Zd ZdZ	 	 	 	 d
dedededdf fdZdej                  dej                  fd	Z	 xZ
S )Dct1dz#1D Discrete Cosine Transform layer.NrF   kernel_typerG   r1   c                     t        ||      }t        	| 	          t        t        d} ||    ||fi |j
                  }| j                  d|j                                | j                  dd        y )NrI   )23weightsbias)	rP   super__init__re   rj   Tregister_bufferrT   register_parameter)
selfrF   rm   rG   rJ   rK   ra   kerneldct_weights	__class__s
            rD   rt   zDct1d.__init__   sp     u-)0BC.f}.{KN2NPPY(>(>(@A-rf   rb   c                 X    t        j                  || j                  | j                        S N)Flinearrq   rr   rx   rb   s     rD   forwardzDct1d.forward   s    xx4<<33rf   r3   TNN__name__
__module____qualname____doc__intboolrt   rQ   Tensorr   __classcell__r{   s   @rD   rl   rl      sZ    -
  ! $.. . 	. 
.4 4%,, 4rf   rl   c            	       v     e Zd ZdZ	 	 	 	 d
dedededdf fdZdej                  dej                  fd	Z	 xZ
S )Dct2dz#2D Discrete Cosine Transform layer.NrF   rm   rG   r1   c                 b    t        ||      }t        | 	          t        |||fi || _        y NrI   )rP   rs   rt   rl   	transformrx   rF   rm   rG   rJ   rK   ra   r{   s          rD   rt   zDct2d.__init__   s1     u-{KKKrf   rb   c                     | j                  | j                  |      j                  dd            j                  dd      S )NrL   )r   	transposer   s     rD   r   zDct2d.forward   s5    ~~dnnQ/99"bABLLRQSTTrf   r   r   r   s   @rD   r   r      sb    -
  ! $
L
L 
L 	
L 

LU U%,, Urf   r   out_chsc           
      &   t        t        j                  |      t        fd|D              }t	        |      }| |z  dk(  r| |k\  sJ d|  d| d| d| d	       | |z  fd|D        \  }}}||z   |z   | k(  rt        |||      dkD  sJ |||fS )	Nc              3   (   K   | ]	  }|z    y wr}    ).0rb   gs     rD   	<genexpr>z!_split_out_chs.<locals>.<genexpr>   s     $a1f$   r   zout_chs=z( can't be split into Y/Cb/Cr with ratio z
 (reduced z!); out_chs must be a multiple of .c              3   (   K   | ]	  }|z    y wr}   r   )r   riunits     rD   r   z!_split_out_chs.<locals>.<genexpr>   s     'rd'r   )r   mathgcdtuplesummin)	r   ratiordenomycbcrr   r   s	          @@rD   _split_out_chsr      s    txxA$e$$AFEU?aGu$4 
7)CE7 K37wa	A4
 eD'Q'IAr2r6B;'!c!Rnq&888b"9rf   c                   <    e Zd ZdZ	 	 	 	 	 ddededededdf
 fdZdd	Zdd
ZddZ	de
j                  de
j                  fdZde
j                  de
j                  fdZde
j                  de
j                  fdZde
j                  de
j                  fdZ xZS )LearnableDct2dzKLearnable 2D DCT stem with RGB to YCbCr conversion and frequency selection.NrF   rm   rG   r   r1   c           
         t        ||      }t        | 	          || _        t	        |||fi || _        t        ||      | _        t        |d      \  }}	}
t        j                  |dz  |fddd|| _        t        j                  |dz  |	fddd|| _        t        j                  |dz  |
fddd|| _        | j                  dt        j                   d	d
||      d       | j                  dt        j                   d	d
||      d       | j                  dt        j                   d	dd||      d       | j                  dt        j                   d	dd||      d       | j#                          y )NrI      rO   rO   )r   r3   r   r   rF   paddingmean   @   F)
persistentvarimagenet_meanimagenet_std)rP   rs   rt   rd   r   r   rE   permutationr   nnConv2dconv_yconv_cbconv_crrv   rQ   emptyreset_parameters)rx   rF   rm   rG   r   rJ   rK   ra   y_chcb_chcr_chr{   s              rD   rt   zLearnableDct2d.__init__   ss    u-{KKK.{KH+G:FeUyy!14YaQRYVXYyy!15YaQRYVXYyy!15YaQRYVXY 	VU[[BvU%S`efUEKK2fE$R_de_ekk!Q&X].^kpq^U[[AqW\-]jop 	rf   c                 $    | j                          y)zInitialize buffers.N_init_buffersrx   s    rD   r   zLearnableDct2d.reset_parameters       rf   c                    | j                   j                  t        j                  t                     | j
                  j                  t        j                  t                     | j                  j                  t        j                  g d      j                  ddd             | j                  j                  t        j                  g d      j                  ddd             y)z.Compute and fill non-persistent buffer values.g
ףp=
?gv/?gCl?r   r   gZd;O?gy&1?g?N)
r   copy_rQ   rY   	_DCT_MEANr   _DCT_VARr   rU   r   r   s    rD   r   zLearnableDct2d._init_buffers   s    		Y/0u||H-.  .C!D!I!I!QPQ!RS-B C H HAq QRrf   c                 $    | j                          y)z"Initialize non-persistent buffers.Nr   r   s    rD   init_non_persistent_buffersz*LearnableDct2d.init_non_persistent_buffers   r   rf   rb   c                 p    |j                  | j                        j                  | j                        dz  S )z3Convert from ImageNet normalized to [0, 255] range.   )mulr   add_r   r   s     rD   _denormalizezLearnableDct2d._denormalize   s-    uuT&&',,T-?-?@3FFrf   c                     |dddf   |dddf   |dddf   }}}|dz  |dz  z   |dz  z   }d||z
  z  d	z   }d
||z
  z  d	z   }t        j                  |||gd      S )z5Convert RGB to YCbCr color space (BCHW input/output).Nr   r   r3   gA`"?gbX9?gv/?g?5^I?   g7A`?rM   )rQ   stack)rx   rb   r   r   br   r   r   s           rD   _rgb_to_ycbcrzLearnableDct2d._rgb_to_ycbcr   s}    AqD'1QT7AadGa1IE	!AI-a!e_s"a!e_s"{{Ar2;A..rf   c                 J    | j                   dz  dz   }|| j                  z
  |z  S )z8Normalize DCT coefficients using precomputed statistics.g      ?g:0yE>)r   r   )rx   rb   stds      rD   _frequency_normalizez#LearnableDct2d._frequency_normalize  s'    hh#o$DII$$rf   c                 H   |j                   \  }}}}| j                  |      }| j                  |      }|j                  |||| j                  z  | j                  || j                  z  | j                        }|j                  dddddd      }| j                  |      }|j                  d|| j                  | j                  z        }|d d d d | j                  f   }| j                  |      }|j                  ||| j                  z  || j                  z  |d      }|j                  ddddd      j                         }| j                  |d d df         }| j                  |d d df         }| j                  |d d df         }t        j                  |||gd      S )	Nr   r3   rO   r   r      rL   rM   )r`   r   r   r6   rd   permuter   r   r   rT   r   r   r   rQ   rV   )	rx   rb   r   chwx_yx_cbx_crs	            rD   r   zLearnableDct2d.forward  sm   WW
1aa q!IIaAKdffdffEIIaAq!Q'NN1IIb!TVVdff_-aD$$$%%%a(IIadffa466k1b9IIaAq!$//1kk!AqD'"||AadG$||AadG$yy#tT*22rf   )r3   T    NNr1   N)r   r   r   r   r   r   rt   r   r   r   rQ   r   r   r   r   r   r   r   s   @rD   r   r      s    U
  ! $     	 
   
 <SGell Gu|| G/u|| / /%ell %u|| %
3 3%,, 3rf   r   c            	            e Zd ZdZ	 	 	 	 d
dedededdf fdZdej                  de	ej                  ej                  f   fd	Z
 xZS )
Dct2dStatsz5Utility module to compute DCT coefficient statistics.NrF   rm   rG   r1   c                     t        ||      }t        | 	          || _        t	        |||fi || _        t        ||      | _        y r   )rP   rs   rt   rd   r   r   rE   r   r   s          rD   rt   zDct2dStats.__init__  sG     u-{KKK.{KHrf   rb   c                    |j                   \  }}}}|j                  |||| j                  z  | j                  || j                  z  | j                        }|j                  dddddd      }| j	                  |      }|j                  d|| j                  | j                  z        }|d d d d | j
                  f   }|j                  ||| j                  z  z  || j                  z  z  |d      }t        j                  ddg      }t        j                  ddg      }t        d      D ]D  }t        j                  |d d |f   d	      ||<   t        j                  |d d |f   d	      ||<   F ||fS )
Nr   r3   rO   r   r   r   rL   r   rM   )r`   r6   rd   r   r   r   rQ   zerosr8   r   r   )	rx   rb   r   r   r   r   	mean_listvar_listr@   s	            rD   r   zDct2dStats.forward,  sF   WW
1aIIaAKdffdffEIIaAq!Q'NN1IIb!TVVdff_-aD$$$%IIa1;'1;7B?KKB(	;;2w'q 	4A ::a1g15IaL))AadG3HQK	4 (""rf   r   )r   r   r   r   r   r   rt   rQ   r   r   r   r   r   s   @rD   r   r     sm    ?
  ! $II I 	I 
I# #%ell0J*K #rf   r   c            	       |     e Zd ZdZ	 	 	 	 d
dededee   ddf fdZdej                  dej                  fd	Z
 xZS )Blockz,ConvNeXt-style block with spatial attention.NrN   	drop_pathls_init_valuer1   c                 T   t        ||      }t        | 	          t        j                  ||fdd|d|| _        t        j                  |fddi|| _        t        j                  |d|z  fi || _	        t        j                         | _        t        d|z  fdd	i|| _        t        j                  d|z  |fi || _        |rt        |fd
|i|nt        j                          | _        |dkD  rt%        |      nt        j                          | _        t)        di || _        y )NrI      r   )rF   r   groupsepsư>rO   channels_lastTinit_valuesr)   r   )rP   rs   rt   r   r   dwconv	LayerNormnormLinearpwconv1GELUactr   grnpwconv2r   Identitylsr
   r   SpatialAttentionattn)rx   rN   r   r   rJ   rK   ra   r{   s          rD   rt   zBlock.__init__A  s    u-iiSUa3URTULL5$5"5	yya#g44779%a#gHTHRHyyS#44HU,sDDD[][f[f[h09B),BKKM$*r*	rf   rb   c                    |}| j                  |      }|j                  dddd      }| j                  |      }| j                  |      }| j	                  |      }| j                  |      }| j                  |      }|j                  dddd      }| j                  |      }t        j                  ||j                  dd  dd      }||z  }| j                  |      }|| j                  |      z   S )Nr   r3   r   r   bilinearT)sizemodealign_corners)r   r   r   r  r  r  r  r
  r~   interpolater`   r  r   )rx   rb   shortcutr
  s       rD   r   zBlock.forwardU  s    KKNIIaAq!IIaLLLOHHQKHHQKLLOIIaAq!yy|}}T*TXYHGGAJ$..+++rf   )r)   NNN)r   r   r   r   r   floatr   rt   rQ   r   r   r   r   s   @rD   r   r   >  s^    6
  "-1++ + $E?	+ 
+(, ,%,, ,rf   r   c                   d     e Zd ZdZ	 	 d	 d fdZdej                  dej                  fdZ xZS )SpatialTransformerBlockzLightweight transformer block for spatial attention (1-channel, 7x7 grid).

    This is a simplified transformer with single-head, 1-dim attention over spatial
    positions. Used inside SpatialAttention where input is 1 channel at 7x7 resolution.
    r1   c                 @   t        ||      }t        | 	          t        dddi|| _        t        j                  d	i || _        t        j                  d
ddi|| _	        t        j                  d	i || _
        t        ddt
        j                  i|| _        y )NrI   in_chansr   rr   F)r   rO   r   	act_layerr   )r   )r   r   )rP   rs   rt   PosConv	pos_embedr   r   norm1r  qkvnorm2r   r  mlprx   rJ   rK   ra   r{   s       rD   rt   z SpatialTransformerBlock.__init__o  s    
 u- 2!2r2\\*r*
99444 \\*r*
8"''8R8rf   rb   c                    |j                   \  }}}}|}|j                  d      j                  dd      }| j                  |      }| j	                  |||f      }| j                  |      }|j                  d      \  }	}
}|	|
j                  dd      z  j                  d      }||z  j                  d      }|j                  dd      j                  ||||      }||z   }|}|j                  d      j                  dd      }| j                  | j                  |            }|j                  dd      j                  ||||      }||z   }|S )Nr3   r   rL   r   rM   )r`   flattenr   r  r  r  unbindsoftmax	unsqueezer6   r  r  )rx   rb   BCHWr  x_tr  qrd   rc   r
  s                rD   r   zSpatialTransformerBlock.forward  s>   WW
1a iil$$Q*jjonnS1a&) hhsm**R.1aAKKB''00R08ax""2&mmAq!))!Q15sN iil$$Q*hhtzz#'mmAq!))!Q15sNrf   NNr   	r   r   r   r   rt   rQ   r   r   r   r   s   @rD   r  r  h  s:     9 
	9  %,, rf   r  c                   d     e Zd ZdZ	 	 d	 d fdZdej                  dej                  fdZ xZS )r	  zBSpatial attention module using channel statistics and transformer.r1   c                     t        ||      }t        | 	          t        j                  d      | _        t        j                  dddd|| _        t        di || _	        y )NrI   )r   r   r   r   r   )r3   r   r   )
rP   rs   rt   r   AdaptiveAvgPool2davgpoolr   convr  r
  r  s       rD   rt   zSpatialAttention.__init__  sW    
 u-++F3IIC1CC	+1b1	rf   rb   c                     |j                  dd      }|j                  dd      }t        j                  ||gd      }| j	                  |      }| j                  |      }| j                  |      }|S )Nr   T)rN   keepdimrM   )r   amaxrQ   rV   r/  r0  r
  )rx   rb   x_avgx_maxs       rD   r   zSpatialAttention.forward  sf    1d+1d+IIuen!,LLOIIaLIIaLrf   r*  r   r+  r   s   @rD   r	  r	    s8    L 	2 
		2 %,, rf   r	  c                        e Zd ZdZ	 	 	 	 	 	 	 	 	 ddedededededed	ed
edee   ddf fdZde	j                  de	j                  fdZ xZS )TransformerBlockzQTransformer block with optional downsampling and convolutional position encoding.Ninpoup	num_headsattn_head_dim
downsample	attn_drop	proj_dropr   r   r1   c           
      "   t        |
|      }t        | 	          t        |dz        }|| _        | j                  rZt        j                  ddd      | _        t        j                  ddd      | _        t        j                  ||dddfddi|| _
        nKt        j                         | _        t        j                         | _        t        j                         | _
        t        dd	|i|| _        t        j                  |fi || _        t!        d||||||d
|| _        |	rt%        |fd|	i|nt        j                         | _        |dkD  rt)        |      nt        j                         | _        t        j                  |fi || _        t/        |||ft
        j0                  |d|| _        |	rt%        |fd|	i|nt        j                         | _        |dkD  rt)        |      | _        y t        j                         | _        y )NrI   rO   r   r3   r   r   rr   Fr  )rN   r:  r;  dim_outr=  r>  r   r)   )r  dropr   )rP   rs   rt   r   r<  r   	MaxPool2dpool1pool2r   projr  r  r  r   r  r   r
  r   ls1r
   
drop_path1r  r   r  r  ls2
drop_path2)rx   r8  r9  r:  r;  r<  r=  r>  r   r   rJ   rK   ra   
hidden_dimr{   s                 rD   rt   zTransformerBlock.__init__  s    u-q\
$??aA.DJaA.DJ		#sAq!F%F2FDIDJDJDI 4#44\\#,,
 
'
 
	 HU:cC}CCZ\ZeZeZg1:R(9-R[[]\\#,,
sJUrwwYURTUGT:cC}CCZ\ZeZeZg1:R(9-R[[]rf   rb   c                 x   | j                   r| j                  | j                  |            }| j                  |      }|j                  \  }}}}|j                  d      j                  dd      }| j                  |      }| j                  |||f      }| j                  | j                  |            }|j                  dd      j                  |d||      }|| j                  |      z   }n|j                  \  }}}}|}|j                  d      j                  dd      }| j                  |      }| j                  |||f      }| j                  | j                  |            }|j                  dd      j                  |d||      }|| j                  |      z   }|j                  \  }}}}|}|j                  d      j                  dd      }| j                  | j                  | j                  |                  }|j                  dd      j                  ||||      }|| j!                  |      z   }|S )Nr3   r   rL   )r<  rE  rC  rD  r`   r   r   r  r  rF  r
  r6   rG  rH  r  r  rI  )rx   rb   r  r(  r$  r%  r&  r'  s           rD   r   zTransformerBlock.forward  s   ??yyA/H**Q-CJAq!Q++a.**1a0C**S/C..q!f-C((499S>*C--1%--aQ:C4??3//AJAq!QH))A,((A.C**S/C..q!f-C((499S>*C--1%--aQ:C4??3//A WW
1aiil$$Q*hhtxx

301mmAq!))!Q15ts++rf   )	   r   Fr)   r)   r)   NNN)r   r   r   r   r   r   r  r   rt   rQ   r   r   r   r   s   @rD   r7  r7    s    [ !#$!!!-1-S-S -S 	-S
 -S -S -S -S -S $E?-S 
-S^ %,, rf   r7  c                   x     e Zd ZdZ	 	 d	deddf fdZdej                  deeef   dej                  fdZ	 xZ
S )
r  z Convolutional position encoding.Nr  r1   c           	          t        ||      }t        | 	          t        j                  ||fdddd|d|| _        y )NrI   r   r   T)rF   strider   rr   r   )rP   rs   rt   r   r   rE  )rx   r  rJ   rK   ra   r{   s        rD   rt   zPosConv.__init__  sD     u-IIhwaST[_hpwtvw	rf   rb   r  c                     |j                   \  }}}|\  }}|j                  dd      j                  ||||      }| j                  |      |z   }|j	                  d      j                  dd      S )Nr   r3   )r`   r   rU   rE  r   )	rx   rb   r  r$  Nr%  r&  r'  cnn_feats	            rD   r   zPosConv.forward  si    ''1a1;;q!$))!Q15IIh(*yy|%%a++rf   r*  )r   r   r   r   r   rt   rQ   r   r   r   r   r   s   @rD   r  r    sR    *
 	xx
 
x, ,U38_ , ,rf   r  c                       e Zd ZdZ	 	 	 	 	 	 	 	 	 	 	 d%dededeedf   deedf   deedf   d	ed
edee   de	ddf fdZ
d&defdZd&dej                  deddfdZej                   j"                  dej                  fd       Zd'dedee	   ddfdZej                   j"                  d&deddfd       Zdej*                  dej*                  fdZ	 	 	 	 	 d(dej*                  deeeee   f      dedede	dedeeej*                     eej*                  eej*                     f   f   fdZ	 	 	 d)deeee   f   ded edee   fd!Zd*dej*                  d"edej*                  fd#Zdej*                  dej*                  fd$Z xZS )+r   zCSATv2: Frequency-domain vision model with spatial attention.

    A hybrid architecture that processes images in the DCT frequency domain
    with ConvNeXt-style blocks and transformer attention.
    Nnum_classesr  dims.depthstransformer_depthsdrop_path_ratetransformer_drop_pathr   global_poolr1   c                    t        |
|      }t        | 	          |dk7  rt        j                  d| d       || _        || _        |	| _        d| _        |d   | _	        | j                  | _
        t        |d   dd	
      g| _        d}t        |      D ]9  \  }}|dkD  r|dz  }| j                  j                  t        ||d| 
             ; |rt        |      nt        d t        ||      D              }t!        t#        j$                  d||      j'                               }g }t        ||      D ]V  \  }}|t)        ||z
        D cg c]  }t+        |       c}z  }|t)        |      D cg c]  }|rt+        |      nd c}z  }X t-        dd|d   i|| _        t!        |      }g }t        t        |||            D ]  \  }\  }}}|dkD  r!t1        j2                  ||dz
     |fddd|gng t)        ||z
        D cg c]  }t5        d|t+        |      |d| c}z   t)        |      D cg c]  }t7        d||t+        |      |d| c}z   |t9        |      dz
  k  rt;        |fddi|gng z   }|j                  t1        j<                  |         t1        j<                  | | _        tA        |d   |fd|	i|| _!        | jE                  d       y c c}w c c}w c c}w c c}w )NrI   r   z5CSATv2 is designed for 3-channel RGB input. in_chans=z* may not work correctly with the DCT stem.FrL   r   rL  stem_dct)num_chs	reductionmoduler3   stages.c              3   ,   K   | ]  \  }}||z
    y wr}   r   )r   rC   ts      rD   r   z"CSATv2.__init__.<locals>.<genexpr>H  s     DwtqRSQUDws   r)   r   r   )rF   rO  )rN   r   r   )r8  r9  r   r   r   r   	pool_typeneeds_reset)rL  r   )#rP   rs   rt   warningswarnrT  r  rZ  grad_checkpointingnum_featureshead_hidden_sizefeature_info	enumerater:   r   zipiterrQ   linspacer7   r8   nextr   r\  r   r   r   r7  lenr   
Sequentialstagesr   headinit_weights)rx   rT  r  rU  rV  rW  rX  rY  r   rZ  rJ   rK   kwargsra   r^  r@   rN   total_blocksdp_iterdp_ratesdeptht_depthr=   rs  layersr{   s                            rD   rt   zCSATv2.__init__!  sD    u-q=MM$:%OQ ' &"' H $ 1 1 "$q'QzRS	o 	cFAs1uQ	$$T#U\]^\_S`%ab	c '<s6{DwWZ[acuWvDwAwu~~aFMMOP!&*<= 	aNE7ego0FG1gGGHQVW^Q_`A*?gRG``H	a '@$q'@R@ x.(1#dFDV2W(X 	2$A$UG RSUVQV"))DQKL!ALLM\^ejkpszkze{|`a[3$w-}[XZ[|} z  @G  zH  Itu!ocsd7m[holno  I	I 893v;?7J+c2t2r23PRT  MM"--01	2 mmV,)$r(K];]Z\]	 	e,3 H` } Is   K%3K*	K/9 K4re  c                 P    | j                  t        | j                  |             y )Nrd  )applyr   _init_weights)rx   re  s     rD   ru  zCSATv2.init_weightsg  s    

74--;GHrf   mc                 6   t        |t        j                  t        j                  f      rOt	        |j
                  d       |j                  +t        j                  j                  |j                  d       y y |rt        |d      r|j                          y y y )Ng{Gz?)r   r   r   )
isinstancer   r   r  r	   weightrr   init	constant_hasattrr   )rx   r  re  s      rD   r  zCSATv2._init_weightsj  sn    a"))RYY/0!((-vv!!!!&&!, "WQ(:;  <[rf   c                 .    | j                   j                  S r}   )rt  fcr   s    rD   get_classifierzCSATv2.get_classifierr  s    yy||rf   c                 ^    || _         ||| _        | j                  j                  ||       y )N)rc  )rT  rZ  rt  reset)rx   rT  rZ  s      rD   reset_classifierzCSATv2.reset_classifierv  s,    &"*D		{;rf   enablec                     || _         y r}   )rh  )rx   r  s     rD   set_grad_checkpointingzCSATv2.set_grad_checkpointing|  s
    "(rf   rb   c                     | j                  |      }| j                  r6t        j                  j	                         st        | j                  |      }|S | j                  |      }|S r}   )r\  rh  rQ   jitis_scriptingr   rs  r   s     rD   forward_featureszCSATv2.forward_features  sS    MM!""599+A+A+Ct{{A.A  AArf   indicesr   
stop_early
output_fmtintermediates_onlyc                     |dk(  sJ d       g }t        t        | j                        dz   |      \  }}	| j                  |      }d|v r|j	                  |       t
        j                  j                         s|s| j                  }
n|	dkD  r| j                  d|	 ng }
t        |
      D ]]  \  }}| j                  r+t
        j                  j                         st        ||      }n ||      }|dz   |v sM|j	                  |       _ |r|S ||fS )aO  Forward pass returning intermediate features.

        Args:
            x: Input image tensor.
            indices: Indices of features to return (0=stem_dct, 1-4=stages). None returns all.
            norm: Apply norm layer to final intermediate (unused, for API compat).
            stop_early: Stop iterating when last desired intermediate is reached.
            output_fmt: Output format, must be 'NCHW'.
            intermediates_only: Only return intermediate features.

        Returns:
            List of intermediate features or tuple of (final features, intermediates).
        NCHWzOutput format must be NCHW.r   r   N)r   rq  rs  r\  r:   rQ   r  r  rl  rh  r   )rx   rb   r  r   r  r  r  intermediatestake_indices	max_indexrs  feat_idxstages                rD   forward_intermediateszCSATv2.forward_intermediates  s
   , V#B%BB#"6s4;;7G!7KW"UiMM!  #99!!#:[[F 1:AT[[),2F(0 	(OHe&&uyy/E/E/Gua(!H!||+$$Q'	(   -rf   
prune_norm
prune_headc                    t        t        | j                        dz   |      \  }}|dkD  r| j                  d| nt        j                         | _        |r#t        j
                         | j                  _        |r| j                  dd       |S )a_  Prune layers not required for specified intermediates.

        Args:
            indices: Indices of intermediate layers to keep (0=stem_dct, 1-4=stages).
            prune_norm: Whether to prune the final norm layer.
            prune_head: Whether to prune the classifier head.

        Returns:
            List of indices that were kept.
        r   r   N )	r   rq  rs  r   rr  r  rt  r   r  )rx   r  r  r  r  r  s         rD   prune_intermediate_layersz CSATv2.prune_intermediate_layers  sp    " #7s4;;7G!7KW"Ui1:Qdkk*9-BMMO[[]DIIN!!!R(rf   
pre_logitsc                 (    | j                  ||      S )N)r  )rt  )rx   rb   r  s      rD   forward_headzCSATv2.forward_head  s    yyzy22rf   c                 F    | j                  |      }| j                  |      S r}   )r  r  r   s     rD   r   zCSATv2.forward  s#    !!!$  ##rf   )  r   )r   H      i  )r3   r3   rL     )r   r   r3   r3   r)   FNavgNN)Tr}   )NFFr  F)r   FTF)r   r   r   r   r   r   r  r   r   strrt   ru  r   Moduler  rQ   r  ignorer  r  r  r   r  r   r   r  r  r  r   r   r   s   @rD   r   r     s     $$6&22>$'*/-1$D-D- D- S/	D-
 #s(OD- !&c3hD- "D- $(D- $E?D- D- 
D-LI I!ryy !t !t ! YY		  <C <hsm <W[ < YY)T )T ) )%,, 5<<  8<$$',0 ||0  eCcN340  	0 
 0  0  !%0  
tELL!5tELL7I)I#JJ	K0 h ./$#	3S	>*  	
 
c83ell 3 3 3$ $%,, $rf   c                      | ddddddddg d	
|S )
Nr  )r      r  )rL  rL  r   r   r  g      ?zhead.fc)
urlrT  
input_size	pool_sizer   r   interpolationcrop_pct
classifier
first_convr   )r  rv  s     rD   _cfgr    s0    =v%.C#r  rf   ztimm/)	hf_hub_id)r     r  bicubic)r  r  r  )
   r  )r  r  r  )zcsatv2.r512_in1kzcsatv2_21m.sw_r640_in1kzcsatv2_21m.sw_r512_in1k
state_dictmodelc                    d| v r| S ddl }dddd|j                  d      }|j                  d      }|j                  d	      }|j                  d
      }d|j                  dt        ffd}i }| j	                         D ]  \  }	}
|j                  d|	      }	|	j                  dd      j                  dd      j                  dd      }	|j                  ||	      }	d|	v r$|	j                  dd      |
j                  d      }
}	n'd|	v r#|	j                  dd      |
j                  d      }
}	d|	v r|	j                  dd      }	nDd|	v r|	j                  dd      }	n-d|	v r|	j                  dd      }	nd |	v r|	j                  d d!      }	d"|	v r3|	j                  d#d$      j                  d%d&      j                  d"d&      }	nd'|	v r|	j                  d'd(      }	d)|	v r|	j                  d)d*      }	nd+|	v r|	j                  d+d,      }	d-|	v rd&|	vr|	j                  d-d.      }	|j                  d/|	      }	|j                  d0|	      }	|
||	<    |S )1zRemap original CSATv2 checkpoint to timm format.

    Handles two key structural changes:
    1) Stage naming: stages1/2/3/4 -> stages.0/1/2/3
    2) Downsample position: moved from end of stage N to start of stage N+1
    zstages.0.0.grn.weightr   Nr   	   )r   r3   r   z^dct\.z^stages([1-4])\.(\d+)\.(.*)$z^head\.z^norm\.r  r1   c                     t        | j                  d            t        | j                  d            | j                  d      }}}|v r||   k(  rd| d| S |dk(  rd| d| S d|dz
   d|dz    d| S )Nr   r3   r   r`  z.0.z	stages.0.r   )r   group)r  r  idxrestdownsample_idxs       rD   remap_stagez)checkpoint_filter_fn.<locals>.remap_stage  s    qwwqz?C
OQWWQZDsN"snU.C'CUG3tf--A:se1TF++1S1WIQtf55rf   z	stem_dct.z.Y_Conv.z.conv_y.z	.Cb_Conv.z	.conv_cb.z	.Cr_Conv.z	.conv_cr.z	grn.gammaz
grn.weightrL   zgrn.betazgrn.biasz
.ff.net.0.z	.mlp.fc1.z
.ff.net.3.z	.mlp.fc2.z	.ff_norm.z.norm2.z.attn_norm.z.norm1.z.attention.attention.z!.attention.attention.attn.to_qkv.z.attn.attn.qkv.z.attention.attention.attn.z.attn.attn.z.attention.z.attn.z.attn.to_qkv.z
.attn.qkv.z.attn.to_out.0.z.attn.proj.z.attn.pos_embed.z.pos_embed.zhead.fc.z
head.norm.)recompileMatchr  itemssubreplacer6   )r  r  r  dct_restage_rehead_renorm_rer  outrd   rc   r  s              @rD   checkpoint_filter_fnr    sk    *, qQ'Nzz)$Fzz9:Hzz*%Gzz*%G6rxx 6C 6 C  " 01JJ{A&YYzJ/[1[1 	

 LLa( !99[,72qA1_99Z4aiimqA 1		,4AQ		,4AA		+y1Aa		-3A #a'>@QRG8OG3O  a		-2A a		/<8A!#		+];A "}A'=		,m<A KK
A&KKa(Aa0d Jrf   variant
pretrainedc                 ~    |j                  dd      }t        t        | |ft        t	        |d      t
        |    d|S )Nout_indices)r   r3   r   rO   T)r  flatten_sequential)pretrained_filter_fnfeature_cfgdefault_cfg)popr   r   r  rP   default_cfgs)r  r  rv  r  s       rD   _create_csatv2r  L  sN    **]L9K 2[TJ )  rf   c                     t        d| fi |S )Nr   )r  )r  rv  s     rD   r   r   Y  s    (J9&99rf   c           	      L    t        ddd      }t        d| fi t        |fi |S )N)0   `      i  )r   r   r  rL  )r   r   rO   r   )rU  rV  rW  
csatv2_21m)rP   r  )r  rv  
model_argss      rD   r  r  ^  s5     !)J ,
Qd:6P6PQQrf   r*  )r   )r  r  )Er   r   rf  	functoolsr   r   typingr   r   r   r   numpyr4   rQ   torch.nnr   torch.nn.functional
functionalr~   timm.layersr	   r
   r   r   r   r   r   r   timm.layers.grnr   timm.models._builderr   timm.models._featuresr   timm.models._manipulater   r   	_registryr   r   __all__r   r   r   rE   r   r   re   rj   r  rl   r   r   r   r   r   r  r	  r7  r  r   r  r  rP   r  r  r  r   r  r   rf   rD   <module>r     sX     % / /      } } } . 5 6 > <X
	68c  c ( 	
 \\: 	YYY
 \\Y4BII 4,UBII U&C "Q3RYY Q3h #  #F',BII ',T0bii 0fryy 0Oryy Od,bii ,*A$RYY A$H %  $  
  $ & "OT O")) O Od
C 
T 
 
 :t :& : : R4 Rf R Rrf   