
    khF                    *   d Z ddlZddlmZmZmZ ddlmZ ddlm	Z	m
Z
mZmZmZmZmZmZ ddlZddlmZ ddlmZmZmZmZ ddlmZmZmZmZmZmZm Z 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.m/Z/ ddl0m1Z1m2Z2 g dZ3e G d d             Z4e G d d             Z5dde	e6df   de	e7df   de6de	e4df   fdZ8	 	 	 	 dde	e6df   de	e7df   de	e6df   de6de
e
e4      f
dZ9	 	 dde	e:e:f   de6dee6e
e6   f   de;de	e4df   f
dZ<d ee4ee4   f   de
e4   fd!Z=d"ee6   d#e6de6fd$Z>e G d% d&             Z? G d' d(ej                        ZAd)e:d*e6d+e6d,e6d-e	e6e6f   d.e?deej                     fd/ZB G d0 d1ej                        ZC G d2 d3ej                        ZD G d4 d5ej                        ZE G d6 d7ej                        ZF G d8 d9ej                        ZG G d: d;ej                        ZH G d< d=ej                        ZI eJeCeDeEeFeGeHeI>      ZKd?e:d@ej                  fdAZLdBee:ej                  f   fdCZM G dD dEej                        ZO	 	 	 	 dd*e6d+e6dGe:dHe:dIe:d.e?fdJZPddLZQdM ZRdNee:ef   dOe4dPe5fdQZSddeSfdRe5dSe7dTe6dUee:ef   dVee6   d.ee?   dWee   fdXZTddRe5dZe;fd[ZU G d\ d]ej                        ZVdd^ej                  d_e:d`e;ddfdaZW eJdi db e5 e4dcddddKddef       e4dcdKdgdKddef       e4dhdidjdKddkf       e4dhdldjdKddmf       e4dhdndjdddmf      fdoddpq      dr e5 e4dcddddKddef       e4dcdKdgdKddef       e4dhdidjdKddkf       e4dhdndjdKddmf       e4dhddjdddmf      fdoddpq      ds e5 e4dcddtdKddef       e4dcdudtdKddef       e4dhdvdwdKddkf       e4dhdKdxdKddmf       e4dhddydddmf      fdzdd{q      d| e5 e8d}d~      ddt      d e5 e8d}d      dd      d e5 e8d}d      dd      d e5 e8d      dd      d e5 e8d      dd      d e5 e8ddn      dd      d e5 e8d      dd      d e5 e8ddn      dd      d e5 e8d      dd      d e5 e8ddn      dd      d e5 e8dd      ddd eJdd            d e5 e4dhdKdyddodkf       e4dhdnddKdodkf       e4dhdiddKdodkf       e4dhdnddKddef      fdddddd      d e5 e4dddyddde eJ              e4dhdnddKdodkf       e4dhdiddKdodkf       e4dhdnddKddef      fdddddd eJdY            d e5 e4dhdKdyddodkf       e4dhdKddKdodkf       e4dhdKddKdodkf       e4dhdKddKdodkf      fdddd      d e5 e4dhdKdyddodkf       e4dhdKddKdodkf       e4dhdKddKdodkf       e4dhdKddKdodkf      fddddd      d e5 e4dhdKdyddodkf       e4dhdKddKdodkf       e4dhdKddKdodkf       e4dhdKddKdodkf      fddddd      d e5 e4dhdKdyddodkf       e4dhdKddKdodkf       e4dhdKddKdodkf       e4dhdKddKdodkf      fddddd      d e5 e4dhdKdyddodkf       e4dhdKddKdodkf       e4dhdKddKdodkf       e4dhdKddKdodkf      fddddd eJd            d e5 e4dhdKdydddkf       e4dhduddKddkf       e4dhduddKddkf       e4dhdKddKddkf      fdddFdd      d e5 e4dhdKdydddkf       e4dhduddKddkf       e4dhduddKddkf       e4dhdKddKddkf      fdddFdd      d e5 e4dhdKdydddkf       e4dhduddKddkf       e4dhduddKddkf       e4dhdKddKddkf      fdddFddd      d e5 e4dhdKdydddkf       e4dhduddKddkf       e4dhduddKddkf       e4dhdKddKddkf      fdddFddd      d e5 e4dhdKdydddkf       e4dhduddKddkf       e4dhduddKddkf       e4dhdKddKddkf      fdddFddd      d e5 e4dhdudyddk       e4dhdnddKdk       e4dhdiddKdk       e4dhduddKdk      fdddFd      d e5 e4dhdudyddodkf       e4dhdnddKdodkf       e4dhdiddKdodkf       e4dhduddKdodkf      fddddd      d e5 e4dhdKdtdKdduf       e4dhdiddKdduf       e4dhddgdKdduf       e4dhdKddKdduf      fdodFdFddd eJdkǫ       eJdYdYȫ      ɫ	      d e5 e4dhdKdtdKddnf       e4dhdiddKddnf       e4dhddgdKddnf       e4dhdKddKddnf      fdodFdFddd eJdkǫ       eJdYdYȫ      ɫ	      d e5 e4dhdudddodnf       e4dhdidddKdodnf       e4dhddydKdodnf       e4dhdudwdKdodnf      fdddFdFddd eJdkǫ       eJdYdYȫ      ͫ
      d e5 e4dhduddddnf       e4dhdidddKddnf       e4dhddydKddnf       e4dhdudwdKddnf      fdddFdFddd eJdkǫ       eJdYdYȫ      ͫ
      d e5 e4dhduddddnf       e4dhddgdKddnf       e4dhddwdKddnf       e4dhduddKddnf      fdddFdFddd eJdkǫ       eJdYdYȫ      ͫ
      d e5 e4dhdKdtdKdduf       e4dhdiddKdduf       e4dhddgdKdduf       e4dhdKddKdduf      fdodFdFdd eedìѫ      d eJdkǫ       eJdYdYȫ      ҫ
      d e5 e4dhdKdtdKddnf       e4dhdiddKddnf       e4dhddgdKddnf       e4dhdKddKddnf      fdodFdFdd eedìѫ      d eJdkǫ       eJdYdYȫ      ҫ
      d e5 e4dhduddddnf       e4dhdidddKddnf       e4dhddydKddnf       e4dhdudwdKddnf      fdddFdFdd eedìѫ      d eJdkǫ       eJdYdYȫ      ֫      d e5 e9ddn٫      ddt      d e5 e9dܬ      dd      d e5 e9dެ      dd      d e5 e9dଉ      dd      d e5 e9dd      dd      d e5 e4dhdudyddk       e4dhdnddKdk       e4dhdiddKdk       e4dhduddKdk      fddFdddd      d e5 e4dhdudyddk       e4dhdnddKdk       e4dhdddKdk       e4dhduddKdk      fddFdddd      d e5 e4dhdndyddk       e4dhdiddKdk       e4dhdddKdk       e4dhdiddKdk      fdddFdddd      d e5 e4dhdidyddk       e4dhdddKdk       e4dhdddKdk       e4dhdddKdk      fdddFdddd      d e5 e4dhdudyddk       e4dhdddKdk       e4dhdddKdk       e4dhdddKdk      fdddFdddd      d e5 e4dhdudyddk       e4dhdnddKdk       e4dhdiddKdk       e4dhduddKdk      fddFddddd      d e5 e4dddodKddf       e4ddddKddf       e4dcddddKdodkf       e4dhddydKddkf      fdddFdd eJdkǫ             ZXdD ]  ZY eeXeY   d      eXeYdz   <    	 ddee:ej                  f   deVde:dee:ej                  f   fdZ[dee:ej                  f   deVfd	Z\dÐd
e:de;deVfdZ]dĐde:dee:ef   fdZ^dĐde:dee:ef   fdZ_ e1i d e^d      d e^d      d e^ddd      d e^ddd      d e^ddd      d e^ddd      d e^ddd      d e^ddd      d  e^ddd      d! e^ddd      d" e^ddd      d# e^ddd      d$ e^ddd      d% e^dddd&d'de(      d) e^dd*d+ddd,de-      d. e_dd/d,de0      d1 e_dd2d,de0      i d3 e_dd4d,de0      d5 e_dd6d,de0      d7 e_dd8d,de0      d9 e_dd:d;      d< e_dd=d,de0      d> e_dd?d,de0      d@ e_ddAd,de0      dB e_ddCd,de0      dD e_ddEd,de0      dF e_ddGd,de0      dH e_ddId,de0      dJ e_ddKdLdMdMdNdOdPd,deQ
      dR e_ddSdLdMdMdPd&deT      dU e_ddVdMdMdWd&X      dY e_ddZdMdMdPd&de[      d\ e_dd]dMdMdPd&de[      d^ e_dLdMdMdNdOdWd,_      i d` e_ddadLdMdMdWd&b      dc e_ddddMdMdWd&de[      de e^ddfdgh      di e^ddjdgh      dk e^ddjdgh      dl e^ddjdgh      dm e^ddjdgh      dn e_ddeedYdNdOdop      dq e_ddeedYdNdOdop      dr e_ddjeedYd,dsdop      dt e_ddueedYdvdwdop      dx e_ddeedYdydzdop      d{ e_ddeedYdNdOdop      d| e_ddeedYdNdOdop      d} e_ddeedYdNdOdop      d~ e_ddeedNdO      d e_ddeedNdO       e_ddeed,ds       e_ddeedvdw       e_ddeedydz       e_ddeedNdO       e_ddeedNdO       e_ddeedNdO       e_dd       e_ddLddWd      d      Z`e2ddeVfd       Zae2ddeVfd       Zbe2ddeVfd       Zce2ddeVfd       Zde2ddeVfd       Zee2ddeVfd       Zfe2ddeVfd       Zge2ddeVfd       Zhe2ddeVfd       Zie2ddeVfd       Zje2ddeVfd       Zke2ddeVfd       Zle2ddeVfd       Zme2ddeVfd       Zne2ddeVfd       Zoe2ddeVfd       Zpe2ddeVfd       Zqe2ddeVfd       Zre2ddeVfd       Zse2ddeVfd       Zte2ddeVfd       Zue2ddeVfd       Zve2ddeVfd       Zwe2ddeVfd       Zxe2ddeVfd       Zye2ddeVfd       Zze2ddeVfd       Z{e2ddeVfd       Z|e2ddeVfd       Z}e2ddeVfd       Z~e2ddeVfd       Ze2ddeVfd       Ze2ddeVfd       Ze2ddeVfd       Ze2ddeVfd       Ze2ddeVfd       Ze2ddeVfd       Ze2ddeVfd       Ze2ddeVfd       Ze2ddeVfd       Ze2ddeVfd       Ze2ddeVfd       Ze2ddeVfd       Ze2ddeVfd       Ze2ddeVfd       Ze2ddeVfd       Ze2ddeVfd       Ze2ddeVfd       Ze2ddeVfd       Ze2ddeVfd       Ze2ddeVfd       Ze2ddeVfd       Ze2ddeVfd       Zy(  a   Bring-Your-Own-Blocks Network

A flexible network w/ dataclass based config for stacking those NN blocks.

This model is currently used to implement the following networks:

GPU Efficient (ResNets) - gernet_l/m/s (original versions called genet, but this was already used (by SENet author)).
Paper: `Neural Architecture Design for GPU-Efficient Networks` - https://arxiv.org/abs/2006.14090
Code and weights: https://github.com/idstcv/GPU-Efficient-Networks, licensed Apache 2.0

RepVGG - repvgg_*
Paper: `Making VGG-style ConvNets Great Again` - https://arxiv.org/abs/2101.03697
Code and weights: https://github.com/DingXiaoH/RepVGG, licensed MIT

MobileOne - mobileone_*
Paper: `MobileOne: An Improved One millisecond Mobile Backbone` - https://arxiv.org/abs/2206.04040
Code and weights: https://github.com/apple/ml-mobileone, licensed MIT

In all cases the models have been modified to fit within the design of ByobNet. I've remapped
the original weights and verified accuracies.

For GPU Efficient nets, I used the original names for the blocks since they were for the most part
the same as original residual blocks in ResNe(X)t, DarkNet, and other existing models. Note also some
changes introduced in RegNet were also present in the stem and bottleneck blocks for this model.

A significant number of different network archs can be implemented here, including variants of the
above nets that include attention.

Hacked together by / copyright Ross Wightman, 2021.
    N)	dataclassfieldreplace)partial)TupleListDictOptionalUnionAnyCallableSequence)IMAGENET_DEFAULT_MEANIMAGENET_DEFAULT_STDOPENAI_CLIP_MEANOPENAI_CLIP_STD)ClassifierHeadNormMlpClassifierHeadConvNormActBatchNormAct2dEvoNorm2dS0aAttentionPool2dRotAttentionPool2dDropPathAvgPool2dSamecreate_conv2dget_act_layerget_norm_act_layerget_attnmake_divisible	to_2tuple   )build_model_with_cfg)feature_take_indices)named_applycheckpoint_seq)generate_default_cfgsregister_model)ByobNetByoModelCfgByoBlockCfgcreate_byob_stemcreate_blockc                   
   e Zd ZU dZeeej                  f   ed<   e	ed<   e	ed<   dZ
e	ed<   dZeee	ef      ed<   d	Zeed
<   dZee   ed<   dZeeeef      ed<   dZee   ed<   dZeeeef      ed<   dZeeeef      ed<   y)r+   zuBlock configuration for Bring-Your-Own-Blocks.

    Defines configuration for a single block or stage of blocks.
    typedc   sNgs      ?br
attn_layerattn_kwargsself_attn_layerself_attn_kwargsblock_kwargs)__name__
__module____qualname____doc__r   strnnModule__annotations__intr3   r4   r
   r   r6   floatr7   r8   r	   r   r9   r:   r;        O/var/www/teggl/fontify/venv/lib/python3.12/site-packages/timm/models/byobnet.pyr+   r+   5   s     RYY

F
FAsJ)-BsH}%&-BN !%J$,0K$sCx.)0%)OXc])15htCH~.5-1L(4S>*1rG   r+   c                      e Zd ZU dZeeeeedf   f   df   ed<   dZe	ed<   dZ
e	ed<   dZee	   ed	<   d
Zeeee   eedf   f   ed<   dZeed<   dZeed<   dZeed<   dZeed<   dZe	ed<   dZe	ed<   dZe	ed<   dZee   ed<   dZe	ed<   dZee	   ed<    ed        Zeed!<   dZee	   ed"<    ed#        Z eed$<    ed%        Z!e"e	e#f   ed&<   y)'r*   zlModel configuration for Bring-Your-Own-Blocks network.

    Defines overall architecture configuration.
    .blocksconv1x1
downsample3x3	stem_typemaxpool	stem_pool    stem_chsr5   width_factorr   num_featuresTzero_init_lastFfixed_input_sizerelu	act_layer	batchnorm
norm_layer aa_layerNhead_hidden_size
classifier	head_typer7   c                      t               S NdictrF   rG   rH   <lambda>zByoModelCfg.<lambda>f   s    df rG   )default_factoryr8   r9   c                      t               S ra   rb   rF   rG   rH   rd   zByoModelCfg.<lambda>h   s    46 rG   r:   c                      t               S ra   rb   rF   rG   rH   rd   zByoModelCfg.<lambda>i   s     rG   r;   )$r<   r=   r>   r?   r   r   r+   rC   rL   r@   rN   rP   r
   rR   rD   r   rS   rE   rT   rU   boolrV   rX   rZ   r\   r]   r_   r7   r   r8   rc   r9   r:   r;   r	   r   rF   rG   rH   r*   r*   J   s)    %U;+;%<<=sBCCJIs(Ix}(79HeCcE#s(O349L%L#ND"d" Is!J!Hc '+hsm*!Is! !%J$n=K=%)OXc])">BdB#(#HL$sCx.HrG   r*   r0   .wfgroupsreturnc                     d}d}dkD  rfd}t        t        | ||      D  cg c]  \  } }}t        d| ||z  |       c}}}       }|S c c}}} w )zCreate RepVGG block configuration.

    Args:
        d: Depth (number of blocks) per stage.
        wf: Width factor per stage.
        groups: Number of groups for grouped convolution.

    Returns:
        Tuple of block configurations.
    @            r   c                 (    |dz   dz  dk(  r| z  S dS )Nr"   r2   r   rF   )chsidxrj   s     rH   rd   z_rep_vgg_bcfg.<locals>.<lambda>{   s    a1}7IcVm q rG   rep)r/   r0   r1   r4   )tuplezipr+   )r0   ri   rj   r1   
group_sizebcfgs     `   rH   _rep_vgg_bcfgrz   l   sa     	AJzP
X[\]_`bdXeffHAqRT+5ARJGfgDK gs   ArF   	se_blocksnum_conv_branchesc                 j   d}t        d|d   |d   z        }|xs dt        |       z  }g }t        | |||      D ]v  \  } }}}g }	t        |       D ]Y  }
||z  }t	        |      }i }|
| |z
  k\  rd|d<   |	t        ddd	|d	|d
|gz  }	|	t        ddd	|dt	        ddd	i|d
|gz  }	|}[ ||	gz  }x |S )a,  Create MobileOne block configuration.

    Args:
        d: Depth (number of blocks) per stage.
        wf: Width factor per stage.
        se_blocks: Number of SE blocks per stage.
        num_conv_branches: Number of conv branches.

    Returns:
        List of block configurations per stage.
    rm   rn   r   )r   )r|   ser7   oner"   )r/   r0   r1   r4   r;   kernel_sizerF   )minlenrw   rangerc   r+   )r0   ri   r{   r|   r1   prev_cry   wr~   scfgiout_cbkaks                 rH   _mobileone_bcfgr      s   " 	AQqTBqE\"F*TCF]ID1aY/ 1aq 		AEE(9:BBAF{#'< [XeqFqrXUWXYYD[ ^a5QT=Va=VSU=V^Z\^ _ _DF		 	 KrG   typeseveryfirstc                    t        |       dk(  sJ t        |t              r%t        t	        |rdn|||dz               }|s|dz
  g}t        |       g }t	        |      D ]"  }||v r| d   n| d   }|t        d|dd|gz  }$ t        |      S )a\  Interleave 2 block types in stack.

    Args:
        types: Two block type names to interleave.
        d: Total depth of blocks.
        every: Interval for alternating blocks.
        first: Whether to start with alternate block.
        **kwargs: Additional block arguments.

    Returns:
        Tuple of interleaved block configurations.
    r2   r   r"   )r/   r0   rF   )r   
isinstancerD   listr   setr+   rv   )r   r0   r   r   kwargsrJ   r   
block_types           rH   interleave_blocksr      s    & u:??%U15!UQY?@UGEJF1X @!"eU1Xq
;>J!>v>??@ =rG   stage_blocks_cfgc           	          t        | t              s| f} g }t        |       D ]6  \  }}|t        |j                        D cg c]  }t        |d       c}z  }8 |S c c}w )zExpand block config into individual block instances.

    Args:
        stage_blocks_cfg: Block configuration(s) for a stage.

    Returns:
        List of individual block configurations.
    r"   )r0   )r   r   	enumerater   r0   r   )r   
block_cfgsr   cfg_s        rH   expand_blocks_cfgr      sh     &1,.J,- @3%,?Qwsa(??
@ @s   Arx   channelsc                 &    | sy|| z  dk(  sJ || z  S )zCalculate number of groups for grouped convolution.

    Args:
        group_size: Size of each group (1 for depthwise).
        channels: Number of channels.

    Returns:
        Number of groups.
    r"   r   rF   )rx   r   s     rH   
num_groupsr      s*      *$))):%%rG   c                   x    e Zd ZU dZeZeed<   eZ	eed<   e
j                  Zeed<   dZee   ed<   dZee   ed<   y)LayerFnz&Container for layer factory functions.conv_norm_actnorm_actactNattn	self_attn)r<   r=   r>   r?   r   r   r   rC   r   r   rA   ReLUr   r   r
   r   rF   rG   rH   r   r      sE    0)M8)'Hh'GGC#D(8
#$(Ix!(rG   r   c                        e Zd ZdZ	 	 	 	 ddedededededee   f fdZd	e	j                  d
e	j                  fdZ xZS )DownsampleAvgz[Average pool downsampling module.

    AvgPool Downsampling as in 'D' ResNet variants.
    in_chsout_chsstridedilation	apply_actlayersc                 6   t         t        |           |xs
 t               }|dk(  r|nd}|dkD  s|dkD  r2|dk(  r|dkD  rt        nt
        j                  } |d|dd      | _        nt        j                         | _        |j                  ||d|      | _
        y)aB  Initialize DownsampleAvg.

        Args:
            in_chs: Number of input channels.
            out_chs: Number of output channels.
            stride: Stride for downsampling.
            dilation: Dilation rate.
            apply_act: Whether to apply activation.
            layers: Layer factory functions.
        r"   r2   TF)	ceil_modecount_include_padr   N)superr   __init__r   r   rA   	AvgPool2dpoolIdentityr   conv)
selfr   r   r   r   r   r   
avg_strideavg_pool_fn	__class__s
            rH   r   zDownsampleAvg.__init__   s    & 	mT+-$79'1}V!
A:A+5?x!|-QSQ]Q]K#AzTUZ[DIDI((!y(Q	rG   xrk   c                 B    | j                  | j                  |            S )zoForward pass.

        Args:
            x: Input tensor.

        Returns:
            Output tensor.
        )r   r   r   r   s     rH   forwardzDownsampleAvg.forward  s     yy1&&rG   )r"   r"   FN)r<   r=   r>   r?   rD   rh   r
   r   r   torchTensorr   __classcell__r   s   @rH   r   r      s|     #(,RR R 	R
 R R W%R:	' 	'%,, 	'rG   r   downsample_typer   r   r   r   r   c                     | dv sJ ||k7  s|dk7  s|d   |d   k7  r6| sy| dk(  rt        ||f||d   d|S  |j                  ||fd||d   d|S t        j                         S )a  Create shortcut connection for residual blocks.

    Args:
        downsample_type: Type of downsampling ('avg', 'conv1x1', or '').
        in_chs: Input channels.
        out_chs: Output channels.
        stride: Stride for downsampling.
        dilation: Dilation rates.
        layers: Layer factory functions.
        **kwargs: Additional arguments.

    Returns:
        Shortcut module or None.
    )avgrK   r[   r"   r   Nr   r   r   )r   r   r   )r   r   rA   r   )r   r   r   r   r   r   r   s          rH   create_shortcutr   !  s    . 4444FaK8A;(1++E% `(ST+`Y_``'6''vQv`hij`kvouvv{{}rG   c                        e Zd ZdZ	 	 	 	 	 	 	 	 	 	 	 ddededededeeef   dee   ded	ed
e	de	de
dedef fdZdde	fdZd Z xZS )
BasicBlockz$ ResNet Basic Block - kxk + kxk
    r   r   r   r   r   rx   bottle_ratiorL   	attn_last
linear_outr   
drop_blockdrop_path_ratec           	         t         t        |           |xs
 t               }t	        ||z        }t        ||      }t        |||||d|      | _        |j                  |||||d         | _	        |	s|j                  t        j                         n|j                  |      | _
        |j                  ||||d   ||d      | _        |	r|j                  t        j                         n|j                  |      | _        |dkD  rt        |      nt        j                         | _        |
rt        j                         | _        y |j#                  d	      | _        y )
NFr   r   r   r   r   r   r"   )r   rj   
drop_layerr           Tinplace)r   r   r   r   r    r   r   shortcutr   	conv1_kxkr   rA   r   	conv2_kxkr   r   	drop_pathr   )r   r   r   r   r   r   rx   r   rL   r   r   r   r   r   mid_chsrj   r   s                   rH   r   zBasicBlock.__init__H  s0     	j$(*$79 <!78J0'Hf

  --fg{SYdlmndo-p%.&++2EBKKM6;;W^K_	--Wka[JRW . 
 /86;;;NTZT_T_`gTh5Cb5H.1bkkm$.2;;=FJJtJ4LrG   rU   c                 P   |rj| j                   ^t        | j                  j                  dd       =t        j
                  j                  | j                  j                  j                         | j                  | j                  fD ]  }t        |d      s|j                          ! y Nweightreset_parametersr   getattrr   bnrA   initzeros_r   r   r   hasattrr   r   rU   r   s      rH   init_weightszBasicBlock.init_weightsl  {    dmm7GDNNDUDUW_ae<f<rGGNN4>>,,334YY/ 	(Dt/0%%'	(rG   c                    |}| j                  |      }| j                  |      }| j                  |      }| j                  |      }| j	                  |      }| j
                  || j                  |      z   }| j                  |      S ra   )r   r   r   r   r   r   r   r   r   r   s      rH   r   zBasicBlock.forwards  w    NN1IIaLNN1NN1NN1==$DMM(++Axx{rG   )   r"   r"   r"   Nr5   r   TFNNr   F)r<   r=   r>   r?   rD   r   r
   rE   r@   rh   r   r   r   r   r   r   r   s   @rH   r   r   D  s      !(.(,"%#"$"#'$&"M"M "M 	"M
 "M CHo"M !"M  "M "M "M "M "M !"M ""MH(4 (	rG   r   c                        e Zd ZdZ	 	 	 	 	 	 	 	 	 	 	 	 	 ddededededeeef   dedee   d	ed
e	de	de	de	de
dedef fdZdde	fdZd Z xZS )BottleneckBlockz4 ResNet-like Bottleneck Block - 1x1 - kxk - 1x1
    r   r   r   r   r   r   rx   rL   r   r   
extra_conv	bottle_inr   r   r   c           	      `   t         t        |           |xs
 t               }t	        |r|n||z        }t        ||      }t        |||||d|      | _        |j                  ||d      | _	        |j                  |||||d   ||      | _
        |r|j                  ||||d   |      | _        nt        j                         | _        |	s|j                  t        j                         n|j                  |      | _        |j                  ||dd      | _        |	r|j                  t        j                         n|j                  |      | _        |dkD  rt#        |      nt        j                         | _        |
rt        j                         | _        y |j'                  d	
      | _        y )NFr   r"   r   r   r   rj   r   )r   rj   r   r   Tr   )r   r   r   r   r    r   r   r   r   	conv1_1x1r   
conv2b_kxkrA   r   r   	conv3_1x1r   r   r   r   )r   r   r   r   r   r   r   rx   rL   r   r   r   r   r   r   r   r   rj   r   s                     rH   r   zBottleneckBlock.__init__  sx   $ 	ot-/$79 I&7l!RSJ0'Hf

  --fgqA--WkHQK: . 
 $22+F 3 TDO !kkmDO%.&++2EBKKM6;;W^K_	--gwU-S.76;;;NTZT_T_`gTh5Cb5H.1bkkm$.2;;=FJJtJ4LrG   rU   c                 P   |rj| j                   ^t        | j                  j                  dd       =t        j
                  j                  | j                  j                  j                         | j                  | j                  fD ]  }t        |d      s|j                          ! y r   )r   r   r   r   rA   r   r   r   r   r   r   r   r   s      rH   r   zBottleneckBlock.init_weights  r   rG   c                 V   |}| j                  |      }| j                  |      }| j                  |      }| j                  |      }| j	                  |      }| j                  |      }| j                  |      }| j                  || j                  |      z   }| j                  |      S ra   )	r   r   r   r   r   r   r   r   r   r   s      rH   r   zBottleneckBlock.forward  s    NN1NN1OOAIIaLNN1NN1NN1==$DMM(++Axx{rG   )r   r"   r   r5   Nr   FFFFNNr   r   r<   r=   r>   r?   rD   r   rE   r
   r@   rh   r   r   r   r   r   r   r   s   @rH   r   r     s      !(."$(,##$$#"#'$&!*M*M *M 	*M
 *M CHo*M  *M !*M *M *M *M *M *M *M !*M  "!*MX(4 (rG   r   c                        e Zd ZdZ	 	 	 	 	 	 	 	 	 	 	 ddededededeeef   dedee   d	ed
e	de	de
dedef fdZdde	fdZd Z xZS )	DarkBlocka
   DarkNet-like (1x1 + 3x3 w/ stride) block

    The GE-Net impl included a 1x1 + 3x3 block in their search space. It was not used in the feature models.
    This block is pretty much a DarkNet block (also DenseNet) hence the name. Neither DarkNet or DenseNet
    uses strides within the block (external 3x3 or maxpool downsampling is done in front of the block repeats).

    If one does want to use a lot of these blocks w/ stride, I'd recommend using the EdgeBlock (3x3 /w stride + 1x1)
    for more optimal compute.
    r   r   r   r   r   r   rx   rL   r   r   r   r   r   c           
         t         t        |           |xs
 t               }t	        ||z        }t        ||      }t        |||||d|      | _        |j                  ||d      | _	        |	s|j                  t        j                         n|j                  |      | _
        |j                  |||||d   ||d      | _        |	r|j                  t        j                         n|j                  |      | _        |dkD  rt        |      nt        j                         | _        |
rt        j                         | _        y |j#                  d      | _        y )	NFr   r"   r   r   r   rj   r   r   r   Tr   )r   r   r   r   r    r   r   r   r   r   r   rA   r   r   r   r   r   r   r   r   r   r   r   r   r   rx   rL   r   r   r   r   r   r   rj   r   s                   rH   r   zDarkBlock.__init__  s#     	i')$79 <!78J0'Hf

  --fgqA%.&++2EBKKM6;;W^K_	--WkHQK:af . 
 /86;;;NTZT_T_`gTh5Cb5H.1bkkm$.2;;=FJJtJ4LrG   rU   c                 P   |rj| j                   ^t        | j                  j                  dd       =t        j
                  j                  | j                  j                  j                         | j                  | j                  fD ]  }t        |d      s|j                          ! y r   r   r   s      rH   r   zDarkBlock.init_weights  r   rG   c                    |}| j                  |      }| j                  |      }| j                  |      }| j                  |      }| j	                  |      }| j
                  || j                  |      z   }| j                  |      S ra   )r   r   r   r   r   r   r   r   s      rH   r   zDarkBlock.forward  r   rG   )r   r"   r   r5   Nr   TFNNr   r   r   r   s   @rH   r   r     s      !(."%(,#"$"#'$&"M"M "M 	"M
 "M CHo"M  "M !"M "M "M "M "M !"M ""MH(4 (	rG   r   c                        e Zd ZdZ	 	 	 	 	 	 	 	 	 	 	 ddededededeeef   dedee   d	ed
e	de	de
dedef fdZdde	fdZd Z xZS )	EdgeBlocka   EdgeResidual-like (3x3 + 1x1) block

    A two layer block like DarkBlock, but with the order of the 3x3 and 1x1 convs reversed.
    Very similar to the EfficientNet Edge-Residual block but this block it ends with activations, is
    intended to be used with either expansion or bottleneck contraction, and can use DW/group/non-grouped convs.

    FIXME is there a more common 3x3 + 1x1 conv block to name this after?
    r   r   r   r   r   r   rx   rL   r   r   r   r   r   c           	         t         t        |           |xs
 t               }t	        ||z        }t        ||      }t        |||||d|      | _        |j                  |||||d   ||      | _	        |	s|j                  t        j                         n|j                  |      | _
        |j                  ||dd      | _        |	r|j                  t        j                         n|j                  |      | _        |dkD  rt        |      nt        j                         | _        |
rt        j                         | _        y |j#                  d	      | _        y )
NFr   r   r   r"   r   r   Tr   )r   r  r   r   r    r   r   r   r   r   r   rA   r   	conv2_1x1r   r   r   r   r   s                   rH   r   zEdgeBlock.__init__  s%     	i')$79 <!78J0'Hf
  --G[HQK: . 
 &/&++2EBKKM6;;W^K_	--gwU-S.76;;;NTZT_T_`gTh5Cb5H.1bkkm$.2;;=FJJtJ4LrG   rU   c                 P   |rj| j                   ^t        | j                  j                  dd       =t        j
                  j                  | j                  j                  j                         | j                  | j                  fD ]  }t        |d      s|j                          ! y r   )r   r   r  r   rA   r   r   r   r   r   r   r   r   s      rH   r   zEdgeBlock.init_weights3  r   rG   c                    |}| j                  |      }| j                  |      }| j                  |      }| j                  |      }| j	                  |      }| j
                  || j                  |      z   }| j                  |      S ra   )r   r   r  r   r   r   r   r   s      rH   r   zEdgeBlock.forward:  r   rG   )r   r"   r   r5   Nr   FFNNr   r   r   r   s   @rH   r  r    s      !(."%(,##$"#'$&!M!M !M 	!M
 !M CHo!M  !M !!M !M !M !M !M !!M "!MF(4 (	rG   r  c                       e Zd ZdZ	 	 	 	 	 	 	 	 	 	 ddededededeeef   dedee   d	ed
e	de
dedef fdZddefdZd Zd Zdeej"                  ej"                  f   fdZdeej"                  ej"                  f   fdZ xZS )RepVggBlockzQ RepVGG Block.

    Adapted from impl at https://github.com/DingXiaoH/RepVGG
    r   r   r   r   r   r   rx   rL   r   r   r   inference_modec           
         t         t        |           t        ||      x| _        }|	xs
 t               }	|r"t        j                  ||||||d      | _        nd | _        ||k(  xr |dk(  xr |d   |d   k(  }|r|	j                  |d      nd | _
        |	j                  |||||d   ||
d      | _        |	j                  ||d||d      | _        |d	kD  r|rt        |      nt        j                         | _        |	j"                  t        j                         n|	j#                  |      | _        |	j%                  d
      | _        y )NTin_channelsout_channelsr   r   r   rj   biasr"   r   Fr   r   )r   rj   r   r   r   )r   r
  r   r   rj   r   rA   Conv2dreparam_convr   identityr   conv_kxkconv_1x1r   r   r   r   r   )r   r   r   r   r   r   r   rx   rL   r   r   r   r  rj   	use_identr   s                  rH   r   zRepVggBlock.__init__L  sI    	k4)+)*f==f$79 "		"$'!!D !%D')XfkXhqkXVW[>XIIRFOOGuOEX\DM"00Fzej 1 DM #00!F[amr0sDM9G"9LQZXn5`b`k`k`mDN%+[[%8BKKMfkk'>R	::d:+rG   rU   c                 r   | j                         D ]s  }t        |t        j                        st        j                  j                  |j                  dd       t        j                  j                  |j                  dd       u t        | j                  d      r| j                  j                          y y )Ng?r   r   )modulesr   rA   BatchNorm2dr   normal_r   r  r   r   r   )r   rU   ms      rH   r   zRepVggBlock.init_weightsw  s{     	/A!R^^,"b12.	/ 49901II&&( 2rG   c                    | j                   /| j                  | j                  | j                  |                  S | j                  $| j	                  |      | j                  |      z   }nJ| j                  |      }| j	                  |      | j                  |      z   }| j                  |      }||z  }| j                  |      }| j                  |      S ra   )r  r   r   r  r  r  r   )r   r   r  s      rH   r   zRepVggBlock.forward  s    (88DIId&7&7&:;<<== a 4==#33A}}Q'Ha 4==#33Aq!AMAIIaLxx{rG   c           
      z   | j                   y| j                         \  }}t        j                  | j                  j
                  j                  | j                  j
                  j                  | j                  j
                  j                  | j                  j
                  j                  | j                  j
                  j                  | j                  j
                  j                  | j                  j
                  j                  d      | _         || j                   j                  _        || j                   j                  _        | j!                         D ]  \  }}d|v r|j#                           | j%                  d       | j%                  d       | j%                  d       | j%                  d       y)	   Following works like `RepVGG: Making VGG-style ConvNets Great Again` -
        https://arxiv.org/pdf/2101.03697.pdf. We re-parameterize multi-branched
        architecture used at training time to obtain a plain CNN-like structure
        for inference.
        NTr  r  r   r   paddingr   rj   r  r  r  r  r  r   r  _get_kernel_biasrA   r  r  r   r  r  r   r   r   r   rj   r   datar  named_parametersdetach___delattr__r   kernelr  nameparas        rH   reparameterizezRepVggBlock.reparameterize  sT    (,,.II**66++88**66==%%,,MM&&..]]''00==%%,,	
 )/  %&*# //1 	JD$%LLN	 	$$$%rG   rk   c                    d}d}| j                   r| j                  | j                         \  }}| j                  j                  j                  d   dz  }t
        j                  j                  j                  |||||g      }d}d}| j                  | j                  | j                        \  }}| j                  | j                        \  }}||z   |z   }||z   |z   }	||	fS z Method to obtain re-parameterized kernel and bias.
        Reference: https://github.com/DingXiaoH/RepVGG/blob/main/repvgg.py#L83
        r   r2   )
r  _fuse_bn_tensorr  r   r   r   rA   
functionalpadr  )
r   
kernel_1x1bias_1x1r0  kernel_identitybias_identitykernel_conv	bias_convkernel_final
bias_finals
             rH   r"  zRepVggBlock._get_kernel_bias  s    
 
==$#'#7#7#F J--$$003q8C,,00c3S=QRJ ==$-1-A-A$---P*O] "&!5!5dmm!DY"Z//A)M9
Z''rG   c                    t        |t              r|j                  j                  }|j                  j
                  }|j                  j                  }|j                  j                  }|j                  j                  }|j                  j                  }n!t        |t        j                        sJ t        | d      s| j                  j                  j                  }|| j                  z  }	| j                  j                  j                  }
t!        j"                  | j                  j                  j                        }t%        |      D ]  }d||||	z  |
d   dz  |
d   dz  f<    || _        | j&                  }|j
                  }|j                  }|j                  }|j                  }|j                  }||z   j)                         }||z  j+                  dddd      }||z  |||z  |z  z
  fS ) Method to fuse batchnorm layer with preceding conv layer.
        Reference: https://github.com/DingXiaoH/RepVGG/blob/main/repvgg.py#L95
        	id_tensorr"   r   r2   r   r   r   r   r   running_meanrunning_varr  epsrA   r  r   r  r  rj   r   r   
zeros_liker   r;  sqrtreshaper   branchr(  r>  r?  gammabetar@  r   	input_dimr   kernel_valuer   stdts                  rH   r.  zRepVggBlock._fuse_bn_tensor  s    fk*[[''F!9911L ))//KII$$E99>>D))--Cfbnn5554-++77"dkk1	"mm00<<$//0B0B0I0IJv aA_`LA	M;q>Q3FTUZ[H[![\a!-^^F!..L ,,KMME;;D**CS &&(S[!!"aA.z4,"6"<<<<rG   )
r   r"   r   r5   Nr[   NNr   Fr   )r<   r=   r>   r?   rD   r   rE   r
   r@   r   r   rh   r   r   r   r+  r   r   r"  r.  r   r   s   @rH   r
  r
  F  s     !(."%(, "#'$&#(),), ), 	),
 ), CHo),  ), !), ), ), !), "), !),V)4 )&B(%ell(B"C (4=u||U\\/I)J =rG   r
  c                   >    e Zd ZdZ	 	 	 	 	 	 	 	 	 	 	 ddededededeeef   ded	ee   d
ede	dede
dededdf fdZdej                  dej                  fdZd Zdeej                  ej                  f   fdZdeej                  ej                  f   fdZ xZS )MobileOneBlocka8   MobileOne building block.

        This block has a multi-branched architecture at train-time
        and plain-CNN style architecture at inference time
        For more details, please refer to our paper:
        `An Improved One millisecond Mobile Backbone` -
        https://arxiv.org/pdf/2206.04040.pdf
    Nr   r   r   r   r   r   rx   rL   r  r|   r   r   r   rk   c                 2   t         t        |           |
| _        t	        ||      x| _        }|xs
 t               }|	r"t        j                  ||||||d      | _	        nd| _	        ||k(  xr |dk(  xr |d   |d   k(  }|r|j                  |d      nd| _        g }t        | j                        D ](  }|j                  |j                  |||||d             * t        j                  |      | _        d| _        |dkD  r|j                  ||d||d      | _        |d	kD  r|rt%        |      nt        j&                         | _        |j*                  t        j&                         n|j+                  |      | _        |j-                  d
      | _        y)z, Construct a MobileOneBlock module.
        Tr  Nr"   r   Fr   )r   r   rj   r   r   r   )r   rM  r   r|   r   rj   r   rA   r  r  r   r  r   appendr   
ModuleListr  
conv_scaler   r   r   r   r   )r   r   r   r   r   r   r   rx   rL   r  r|   r   r   r   rj   r  convsr   r   s                     rH   r   zMobileOneBlock.__init__  s   $ 	nd,.!2)*f==f$79 "		"$'!!D !%D ')XfkXhqkXVW[>XIIRFOOGuOEX\DM E4112 DV11G!&E 2 C DD MM%0DM #DOQ"("6"6G!&E #7 #C :H"9LQZXn5`b`k`k`mDN%+[[%8BKKMfkk'>R	::d:+rG   r   c                    | j                   /| j                  | j                  | j                  |                  S d}| j                  | j                  |      }d}| j                  | j	                  |      }|}| j
                  D ]  }| ||      z  } | j                  |      }||z  }| j                  | j                  |            S )z Apply forward pass. r   )r  r   r   r  rQ  r  r   )r   r   identity_out	scale_outoutcks         rH   r   zMobileOneBlock.forward-  s     (88DIId&7&7&:;<< ==$==+L 	??&*I -- 	B2a5LC	nnS!|xx		#''rG   c           
         | j                   y| j                         \  }}t        j                  | j                  d   j
                  j                  | j                  d   j
                  j                  | j                  d   j
                  j                  | j                  d   j
                  j                  | j                  d   j
                  j                  | j                  d   j
                  j                  | j                  d   j
                  j                  d      | _         || j                   j                  _        || j                   j                  _        | j!                         D ]  \  }}d|v r|j#                           | j%                  d       | j%                  d       | j%                  d       | j%                  d	       y)
r  Nr   Tr  r  r  rQ  r  r   r!  r'  s        rH   r+  zMobileOneBlock.reparameterizeG  sw    (,,.IIa(--99q)..;;a(--99==#((//MM!$))11]]1%**33==#((// )/  %&*# //1 	JD$%LLN	 	$&$%rG   c                 "   d}d}| j                   u| j                  | j                         \  }}| j                  d   j                  j                  d   dz  }t
        j                  j                  j                  |||||g      }d}d}| j                  | j                  | j                        \  }}d}d}t        | j                        D ]-  }| j                  | j                  |         \  }	}
||	z  }||
z  }/ ||z   |z   }||z   |z   }||fS r-  )rQ  r.  r  r   r   r   rA   r/  r0  r  r   r|   )r   kernel_scale
bias_scaler0  r3  r4  r5  r6  ix_kernel_biasr7  r8  s                rH   r"  zMobileOneBlock._get_kernel_biasg  s,   
 
??&'+';';DOO'L$L*--"''33A6!;C 88..22<#sCQTAUVL ==$-1-A-A$---P*O] 	../ 	B!11$--2CDNGU7"KI	
 #\1OC+m;
Z''rG   c                    t        |t              r|j                  j                  }|j                  j
                  }|j                  j                  }|j                  j                  }|j                  j                  }|j                  j                  }n*t        |t        j                        sJ t        | d      s| j                  d   j                  j                  }|| j                  z  }	| j                  d   j                  j                  }
t!        j"                  | j                  d   j                  j                        }t%        |      D ]  }d||||	z  |
d   dz  |
d   dz  f<    || _        | j&                  }|j
                  }|j                  }|j                  }|j                  }|j                  }||z   j)                         }||z  j+                  dddd      }||z  |||z  |z  z
  fS )r:  r;  r   r"   r2   r<  r=  rD  s                  rH   r.  zMobileOneBlock._fuse_bn_tensor  s    fk*[[''F!9911L ))//KII$$E99>>D))--Cfbnn5554-q)..::"dkk1	"mmA.33??$//a0@0E0E0L0LMv aA_`LA	M;q>Q3FTUZ[H[![\a!-^^F!..L ,,KMME;;D**CS &&(S[!!"aA.z4,"6"<<<<rG   )r   r"   r   r5   Nr[   Fr"   NNr   )r<   r=   r>   r?   rD   r   rE   r
   r@   rh   r   r   r   r   r   r   r+  r"  r.  r   r   s   @rH   rM  rM    s$     !(."%(, #(%&"#'$&8,8, 8, 	8,
 8, CHo8,  8, !8, 8, !8,  #8, 8, !8, "8, 
8,t( (%,, (4&@(%ell(B"C (>=u||U\\/I)J =rG   rM  c            !            e Zd ZdZ	 	 	 	 	 	 	 	 	 	 	 	 	 	 ddededededeeef   dedee   d	ed
e	de	de	de	deeeef      de
dedef  fdZdde	fdZd Z xZS )SelfAttnBlockzI ResNet-like Bottleneck Block - 1x1 - optional kxk - self attn - 1x1
    r   r   r   r   r   r   rx   rL   r   r   r   post_attn_na	feat_sizer   r   r   c           	         t         t        |           |J t        |r|n||z        }t	        ||      }t        |||||d|      | _        |j                  ||d      | _        |	r#|j                  |||||d   ||      | _	        d}nt        j                         | _	        |i nt        |      } |j                  |fd|i|| _        |r|j                  |      nt        j                         | _        |j                  ||dd      | _        |d	kD  rt#        |      nt        j                         | _        |
rt        j                         | _        y |j'                  d
      | _        y )NFr   r"   r   r   )rc  r   r   r   Tr   )r   ra  r   r    r   r   r   r   r   r   rA   r   rc   r   r   	post_attnr   r   r   r   )r   r   r   r   r   r   r   rx   rL   r   r   r   rb  rc  r   r   r   r   rj   
opt_kwargsr   s                       rH   r   zSelfAttnBlock.__init__  sR   & 	mT+-!!! I&7l!RSJ0'Hf

  --fgqA#11+Fz 2 DN F[[]DN$,R$2K
)))'O&OJO5A1r{{}--gwU-S5Cb5H.1bkkm$.2;;=FJJtJ4LrG   rU   c                 >   |rj| j                   ^t        | j                  j                  dd       =t        j
                  j                  | j                  j                  j                         t        | j                  d      r| j                  j                          y y r   )r   r   r   r   rA   r   r   r   r   r   r   )r   rU   s     rH   r   zSelfAttnBlock.init_weights  sl    dmm7GDNNDUDUW_ae<f<rGGNN4>>,,3344>>#56NN++- 7rG   c                 4   |}| j                  |      }| j                  |      }| j                  |      }| j                  |      }| j	                  |      }| j                  |      }| j                  || j                  |      z   }| j                  |      S ra   )r   r   r   re  r   r   r   r   r   s      rH   r   zSelfAttnBlock.forward  s    NN1NN1NN1NN1NN1NN1==$DMM(++Axx{rG   )r   r"   r   r5   Nr   FFFTNNNr   r   r   r   s   @rH   ra  ra    s     !(."$(,#$$#!%37"#'$&#,M,M ,M 	,M
 ,M CHo,M  ,M !,M ,M ,M ,M ,M ,M  c3h0,M ,M  !!,M" "#,M\.4 .
rG   ra  )basicbottledarkedgeru   r   r   r   block_fnc                     |t         | <   y ra   )_block_registry)r   rm  s     rH   register_blockrp    s    "*OJrG   blockc                     t        | t        j                  t        f      r | di |S | t        v s
J d|         t	        |    di |S )NzUnknown block type (rF   )r   rA   rB   r   ro  )rq  r   s     rH   r-   r-     sO    %"))W-.vO#C';E7%CC#5!+F++rG   c                        e Zd Z	 	 	 	 	 	 	 ddedeeee   eedf   f   dededededee   d	e	d
e
f fdZdeej                  eej                     f   fdZ xZS )Stemr   r   .r   r   r   num_repnum_act	chs_decayr   c
           
      "   t         |           |dv sJ |	xs
 t               }	t        |t        t
        f      rt        |      }|}
n.t        |      D cg c]  }t        |||z  z         c}d d d   }
|| _	        g | _
        d}dgdg|dz
  z  z   }|dk(  r|sd|d<   ||n|}dg||z
  z  dg|z  z   }|}d}d}t        t        |
||            D ]  \  }\  }}}|r|	j                  nt        }d	|dz    }|d
kD  r2|dkD  r-|dz
  }| j                  j                  t!        |||d
             | j#                  | |||||             |}||z  }|} |r|j%                         }|dv sJ }| j                  j                  t!        |||d
             |dk(  r&| j#                  dt'        j(                  d             n|dk(  r&| j#                  dt'        j*                  d             nZd|v r)| j#                  dt'        j(                  ddd             n-d|v r)| j#                  dt'        j*                  dddd             |dz  }d}|d
k\  r|nd | _        | j                  j                  t!        |||d
             ||k(  sJ y c c}w )N)r2      r<  r[   r2   r"   ry  FTr   r   num_chs	reductionmodulestage)r   r   )maxrO   r   avgpoolmax2avg2r  r   r  r  r   )r   r   r   r   )r   r   r   r   )r   r   r   r   r   rv   r   r   roundr   feature_infor   rw   r   r   rO  rc   
add_modulelowerrA   	MaxPool2dr   last_feat_idx)r   r   r   r   r   r   ru  rv  rw  r   rR   r   	prev_featstem_stridesstem_norm_actsprev_chscurr_strider  chr3   nalayer_fn	conv_namer   s                          rH   r   zStem.__init__  s    	$79ge}-'lGHAFwPAg	Q67PQUSUQUVH	saSGaK00Q;t L$_''Gg$56$'9II'HlN(ST 		"NA{Ar/1v++}Hq1ugI1uQ !A!!((h+^gop)qrOOIx"+^_'`aH1K!I		" ::<DOOOOM$$T(kZckl%mnv~Q8Q8$1VW(XY$1VWkp(qr1KI.;q.@]d  h+V_gh!ijf$$$[ Qs   Jrk   c                     d }t        |       D ],  \  }} ||      }| j                  || j                  k(  s+|}. ||fS ra   )r   r  )r   r   intermediater   r  s        rH   forward_intermediateszStem.forward_intermediatesF  sQ    /3dO 	!DAq!A!!-!t7I7I2I 	! ,rG   )r   ry  rO   r   N      ?N)r<   r=   r>   rD   r   r   r   r@   r
   rE   r   r   r   r   r  r   r   s   @rH   rt  rt    s      !!%)""A%A% 3S	5c?:;A% 	A%
 A% A% A% c]A% A% A%F%hu||>T0T*U rG   rt  r[   rN   	pool_typefeat_prefixc                    |xs
 t               }|dv sJ d|v rd|v rdnd }t        | |d|||      }nd|v rt        | d|z  d	z  |dz  |f||
      }nd|v rt        | |dd||      }nd|v rt        | |d|      }nd|v rt        | |dd|      }n~d|v r*|rt        | |dd||      }nf|j	                  | |dd      }nPt        |t        t        f      rt        | |d||
      }n)|rt        | |dd||      }n|j	                  | |dd      }t        |t              r<|j                  D cg c]#  }t        |dj                  ||d   g            % }	}||	fS t        |d|d      g}	||	fS c c}w )N)	r[   quadquad2tiereddeepru   r   7x7rM   r  r  r2   ry  )ru  rv  r   r   r  r      )r   r   r  r5   )ru  rw  r   r   ru   )r   r   r   )r   r   r   r     r"   )ru  r   r   r   .r}  )r}  r   rz  )r   rt  r
  rM  r   r   rv   r   r  rc   join)
r   r   rN   r  r  r   rv  stemfr  s
             rH   r,   r,   O  s     wyF[[[[)+!FGQiX^_	Y	FQ[A-w!|WEI^de	9	FGQ#IV\]	)	671VD	)	fg1QvV	)	AIfUD''1'EDgt}-6JD FGQ	RXY++FGQq+I$VZVgVghQRQsxxak0J'KLhh  W+UVWX is   .(E0r2   c                 L    | d S t        | D cg c]  }||z  	 c}      S c c}w ra   )rv   )rc  r   r3   s      rH   reduce_feat_sizer  ~  s)    $4Q%i0Pf0P*QQ0Ps   !c                     | | n|}|xs i S )a2   Override model level attn/self-attn/block kwargs w/ block level

    NOTE: kwargs are NOT merged across levels, block_kwargs will fully replace model_kwargs
    for the block if set to anything that isn't None.

    i.e. an empty block_kwargs dict will remove kwargs set at model level for that block
    rF   )r;   model_kwargs
out_kwargss      rH   override_kwargsr    s     ".!9|JrG   r;   	block_cfg	model_cfgc                    | d   }|j                   d u}|s|j                  q|r|j                   sd }nSt        |j                  |j                        }|j                   xs |j                   }|t        t	        |      fi |nd }t        ||      }|j                  d u}|s|j                  q|r|j                  sd }nSt        |j                  |j                        }	|j                  xs |j                  }|t        t	        |      fi |	nd }t        ||      }|| d<   | j                  t        |j                  |j                               y )Nr   )r   r   )
r7   r8   r  r   r   r   r9   r:   updater;   )
r;   r  r  	layer_fnsattn_setr7   r8   self_attn_setr9   r:   s
             rH   update_block_kwargsr    sM   X&I ##4/H9((4I00J))*?*?AVAVWK"--E1E1EJISI_*!5EEeiJIJ7	 --T9M	22>!:!:"O.y/I/I9KeKef'77T9;T;TO". &h&?TCST48 IA	&L 	(>(>	@V@VWXrG   r   r   output_stride	stem_featrc  block_kwargs_fnc                    |xs
 t               }g }| j                  D cg c]  }t        |       }	}|	D 
cg c]%  }
t        |
D cg c]  }|j                   c}      ' }}
}t        j                  d|t        |            j                  |      D cg c]  }|j                          }}d}|d   }|d   }|}g }t        |	      D ]]  \  }}|d   j                  }|dk7  r|r|j                  |       ||k\  r|dkD  r||z  }d}||z  }|dv rdnd}g }t        |      D ]  \  }}t        |j                  | j                  z        }|j                  }t!        |t"              r	 |||      }t%        |||dk(  r|nd||f||j&                  | j(                  ||   |   |	      }|j*                  dv r||d	<    |||| 
       |t-        |j*                  fi |gz  }|}|}|dkD  s|dk(  st/        ||      } |t1        j2                  | gz  }t%        ||d| |dz         }` |j                  |       t1        j2                  | ||fS c c}w c c}w c c}}
w c c}w )Nr   r"   r|  r{  )r"   r2   r2   )	r   r   r   r   rx   r   rL   r   r   r  rc  )r  r  stages.rz  )r   rJ   r   sumr0   r   linspacesplittolistr   r3   rO  r    r1   rS   r4   r   r   rc   r6   rL   r/   r-   r  rA   
Sequential)r   r   r  r  rc  r   r  r  r3   r   	stage_bcsbcdepthsr   dprr   
net_strider  r  stages	stage_idxstage_block_cfgsr   first_dilationrJ   	block_idxr  r   rx   r;   s                                 rH   create_byob_stagesr    s     wyFL03

;1#A&;J;?IJ)c),B244,-JFJ$~~aVMSSTZ[
\!188:
\C
\H;'J#HIF'0'< &t#	#!!$&&Q;9	*&6A:HFf
&&0a$-.>$? 	@ Iy$Y[[33C3C%CDG"J*h/';
!*avQ((3%&\\>>"9~i8
L ~~/,5[)LIM|INNClCDDF%NHzi1n,Y?	1	@4 	2==&)**Z'R[Q\H]enqrers	M&tP 	"==&!<::c <,J
\s"   III	II!ITallow_aac                    t        | j                        }t        | j                  |      }| j                  r*|r(t        t        | j                  || j                        }nt        t        | j                  |      }| j                  r)t        t        | j                        fi | j                  nd }| j                  r)t        t        | j                        fi | j                  nd }t        |||||      }|S )N)rZ   rX   )rZ   rX   r\   )r   r   r   r   r   )r   rX   r   rZ   r\   r   r   r7   r   r8   r9   r:   r   )r   r  r   r   r   r   r   r  s           rH   get_layer_fnsr    s    

&C!S^^sKH
||RU`c`l`lmRUVCF>>78CNN+?s?W[DRUReRe!4!45N9M9MNkoI]X3UYenoHOrG   c                       e Zd ZdZ	 	 	 	 	 	 	 	 d#dedededee   dedeeee	eef   f      d	e
d
e
def fdZej                  j                  d$dedeeef   fd       Zej                  j                  d%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	 	 	 	 	 	 d'dej0                  deeeee   f      dedededededeeej0                     e	ej0                  eej0                     f   f   fdZ	 	 	 d(deeee   f   dededee   fdZdej0                  dej0                  fdZd$dej0                  d edej0                  fd!Zdej0                  dej0                  fd"Z xZS ))r)   a%  Bring-your-own-blocks Network.

    A flexible network backbone that allows building model stem + blocks via
    dataclass cfg definition w/ factory functions for module instantiation.

    Current assumption is that both stem and blocks are in conv-bn-act order (w/ block ending in act).
    Nr   num_classesin_chansglobal_poolr  img_size	drop_rater   rU   c
           
         t         |           || _        || _        d| _        t        |fi |
}t        |d      }t        |      }|j                  r	|J d       |t        |      nd}g | _	        t        |j                  t        t        f      r9|j                  D cg c]#  }t        t        ||j                   z              % }}nFt        t        |j                  xs |j"                  d   j$                  |j                   z              }t'        |||j(                  |j*                  |      \  | _        }| j                  j/                  |dd        t1        ||d   d   	      }t3        ||||d   ||
      \  | _        }}| j                  j/                  |dd        |d   d   }|d   d   }|j6                  rSt        t        |j                   |j6                  z              | _        |j9                  || j6                  d      | _        n || _        t=        j>                         | _        | xj                  tA        | j6                  |dtC        | j4                              gz  c_	        | j                  D cg c]  }|d   	 c}| _"        | j6                  | _#        |jH                  dv sJ |jH                  dk(  rj|d}tK        | j6                  ||jF                  ||jL                  |jN                  | j                        | _(        | jP                  jR                  | _#        n|jH                  dk(  r[|d}|dv sJ tU        | j6                  |jF                  |||| j                  d      | _(        | jP                  jV                  | _#        n|jH                  dk(  r[|d}|dv sJ tY        | j6                  |jF                  |||| j                  d      | _(        | jP                  jV                  | _#        n:|d}|jF                  J t[        | j6                  ||| j                        | _(        || _.        t_        ta        tb        |	      |        yc c}w c c}w )a_  
        Args:
            cfg: Model architecture configuration.
            num_classes: Number of classifier classes.
            in_chans: Number of input channels.
            global_pool: Global pooling type.
            output_stride: Output stride of network, one of (8, 16, 32).
            img_size: Image size for fixed image size models (i.e. self-attn).
            drop_rate: Classifier dropout rate.
            drop_path_rate: Stochastic depth drop-path rate.
            zero_init_last: Zero-init last weight of residual path.
            **kwargs: Extra kwargs overlayed onto cfg.
        F)r  Nz8img_size argument is required for fixed input size modelr   )r   r   rN   r  r   r<  r|  r  )r   rc  r{  r"   
final_convrz  r~  )r[   r^   mlpattn_absattn_rotr  r   )hidden_sizer  rZ   rX   r  r  token)r[   r  T)	embed_dimout_featuresrc  r  r  qkv_separater  )r  r  ref_feat_sizer  r  r  )r  r  rU   )2r   r   r  r  grad_checkpointingr   r  rV   r!   r  r   rR   r   rv   rD   r  rS   rJ   r1   r,   rN   rP   r  extendr  r  r  rT   r   r  rA   r   rc   r   
stage_endsr]   r_   r   rZ   rX   headr  r   r  r   r   r  r%   r   _init_weights)r   r   r  r  r  r  r  r  r   rU   r   stem_layersstage_layersrc  r1   rR   r  
stage_featr|  r  r  r   s                        rH   r   zByobNet.__init__  s   4 	&""'c$V$#C%8$S)'c)cc'+3+?Ih'T	cllT5M2BE,,OQE!c&6&6"678OHO5#,,"A#**Q-//SEUEU!UVWH/mmmm 
	9 	  3B0$Yy}[7QR	-?bM.
*Z 	  CR1rN;/	b>), #E#*:*:S=M=M*M$N OD*884CTCTVWXDO (D kkmDO**i\_`d`k`k\lmo 	o/3/@/@A!1W:A $ 1 1}} QQQQ==E!"#-!!00%>>--..DI %)II$9$9D!]]j("%-///'!!..(#%..!DI %)II$7$7D!]]j("%-///*!!..('%..!DI %)II$7$7D!"#''///&!!%..	DI ' 	GM.I4Pw PB Bs   "(Q'Q,coarserk   c                 0    t        d|rdnddfdg      }|S )zGroup matcher for parameter groups.

        Args:
            coarse: Whether to use coarse grouping.

        Returns:
            Dictionary mapping group names to patterns.
        z^stemz^stages\.(\d+)z^stages\.(\d+)\.(\d+)N)z^final_conv)i )r  rJ   rb   )r   r  matchers      rH   group_matcherzByobNet.group_matcher  s.     &,"2JDQ*
 rG   enablec                     || _         y)zEnable or disable gradient checkpointing.

        Args:
            enable: Whether to enable gradient checkpointing.
        N)r  )r   r  s     rH   set_grad_checkpointingzByobNet.set_grad_checkpointing  s     #)rG   c                 .    | j                   j                  S )zPGet classifier module.

        Returns:
            Classifier module.
        )r  fc)r   s    rH   get_classifierzByobNet.get_classifier  s     yy||rG   c                 J    || _         | j                  j                  ||       y)zReset classifier.

        Args:
            num_classes: Number of classes for new classifier.
            global_pool: Global pooling type.
        N)r  r  reset)r   r  r  s      rH   reset_classifierzByobNet.reset_classifier  s     '		[1rG   r   indicesnorm
stop_early
output_fmtintermediates_onlyexclude_final_convc                 N   |dv sJ d       g }t        t        | j                        |      \  }	}
|	D cg c]  }| j                  |    }	}| j                  |
   }
d}t        | j                  d      r| j                  j                  |      \  }}n| j	                  |      d}}||	v r|j                  ||n|       | j                  d   }t        j                  j                         s|s| j                  }n| j                  d|
 }|D ]t  }|dz  }| j                  r+t        j                  j                         st        ||      }n ||      }|s||k(  r| j                  |      }||	v sd|j                  |       v |r|S |r||k(  r| j                  |      }||fS c c}w )a9   Forward features that returns intermediates.

        Args:
            x: Input image tensor
            indices: Take last n blocks if int, all if None, select matching indices if sequence
            norm: Apply norm layer to compatible intermediates
            stop_early: Stop iterating over blocks when last desired intermediate hit
            output_fmt: Shape of intermediate feature outputs
            intermediates_only: Only return intermediate features
            exclude_final_conv: Exclude final_conv from last intermediate
        Returns:

        )NCHWzOutput shape must be NCHW.r   r  Nr<  r"   )r$   r   r  r   r  r  rO  r   jitis_scriptingr  r  r&   r  )r   r   r  r  r  r  r  r  intermediatestake_indices	max_indexr   feat_idxx_interlast_idxr  r~  s                    rH   r  zByobNet.forward_intermediates  s   . Y&D(DD&"6s4??7KW"Ui4@Aq*AAOOI.	4995688;JAw1twA|#  go7C??2&99!!#:[[F[[),F 
	(EMH&&uyy/E/E/G"5!,!H%(h*>OOA&<'$$Q'
	(   (h"6"A-E Bs   F"
prune_norm
prune_headc                    t        t        | j                        |      \  }}| j                  |   }| j                  d| | _        || j                  d   k  rt	        j
                         | _        |r| j                  dd       |S )aE  Prune layers not required for specified intermediates.

        Args:
            indices: Indices of intermediate layers to keep.
            prune_norm: Whether to prune normalization layer.
            prune_head: Whether to prune the classifier head.

        Returns:
            List of indices that were kept.
        Nr<  r   r[   )r$   r   r  r  rA   r   r  r  )r   r  r  r  r  r  s         rH   prune_intermediate_layersz!ByobNet.prune_intermediate_layers  su      #7s4??7KW"UiOOI.	kk*9-tr** kkmDO!!!R(rG   c                     | j                  |      }| j                  r5t        j                  j	                         st        | j                  |      }n| j                  |      }| j                  |      }|S )zForward pass through feature extraction.

        Args:
            x: Input tensor.

        Returns:
            Feature tensor.
        )r  r  r   r  r  r&   r  r  r   s     rH   forward_featureszByobNet.forward_features  sY     IIaL""599+A+A+Ct{{A.AAAOOArG   
pre_logitsc                 N    |r| j                  ||      S | j                  |      S )zForward pass through head.

        Args:
            x: Input features.
            pre_logits: Return features before final linear layer.

        Returns:
            Classification logits or features.
        )r	  )r  )r   r   r	  s      rH   forward_headzByobNet.forward_head#  s(     7Atyyzy2RdiiPQlRrG   c                 J    | j                  |      }| j                  |      }|S )zoForward pass.

        Args:
            x: Input tensor.

        Returns:
            Output logits.
        )r  r  r   s     rH   r   zByobNet.forward/  s)     !!!$a rG   )  r   NrQ   Nr   r   Tr   Tra   )NFFr  FF)r"   FT) r<   r=   r>   r?   r*   rD   r
   r@   r   r   rE   rh   r   r   r  ignorer	   r   r  r  rA   rB   r  r  r   r   r  r  r  r  r   r   r   s   @rH   r)   r)     s     $)-!#>B!$&#'CQCQ CQ 	CQ
 "#CQ CQ uS%S/%9:;CQ CQ "CQ !CQJ YYD T#s(^  $ YY)T )T ) ) YY		  2C 2hsm 2W[ 2 8<$$',',< ||<  eCcN34<  	< 
 <  <  !%<  !%<  
tELL!5tELL7I)I#JJ	K< @ ./$#	3S	>*  	
 
c2%,, 5<< "
Sell 
S 
S 
S %,, rG   r)   r}  r)  rU   c                    t        | t        j                        r| j                  d   | j                  d   z  | j                  z  }|| j
                  z  }| j                  j                  j                  dt        j                  d|z               | j                  %| j                  j                  j                          yyt        | t        j                        rct        j                  j                  | j                  dd       | j                  *t        j                  j                  | j                         yyt        | t        j                         rSt        j                  j#                  | j                         t        j                  j                  | j                         yt%        | d      r| j'                  |	       yy)
zInitialize weights.

    Args:
        module: Module to initialize.
        name: Module name.
        zero_init_last: Zero-initialize last layer.
    r   r"          @Nr   g{Gz?)meanrJ  r   r  )r   rA   r  r   r  rj   r   r#  r  mathrB  r  zero_Linearr   r   r  ones_r   r   )r}  r)  rU   fan_outs       rH   r  r  =  s8    &"))$$$Q'&*<*<Q*??&BUBUUFMM!""1diig&>?;;"KK""$ #	FBII	&
CT:;;"GGNN6;;' #	FBNN	+
fmm$
v{{#		(>: 
)rG   gernet_lri  ro   r5   )r/   r0   r1   r3   r4   r6      rj     i  g      ?         @ry  rQ   i 
  )rJ   rR   rP   rT   gernet_mgernet_s0   r   r    i0  rp      i  	repvgg_a0)r2   ry     r"   )      ?r$  r$        @)r0   ri   ru   )rJ   rN   rR   	repvgg_a1)r"   r"   r"   r%  rn   	repvgg_a2)      ?r(  r(  g      @	repvgg_b0)r5   r5   r5   r%  )ri   	repvgg_b1)r  r  r        @repvgg_b1g4)ri   rj   	repvgg_b2)r%  r%  r%        @repvgg_b2g4	repvgg_b3)r  r  r  r.  repvgg_b3g4repvgg_d2se)r  r#     r"   r~   g      ?)rd_ratio
rd_divisor)rJ   rN   rR   r7   r8   	resnet51qrq   i   r  i   silu)rJ   rR   rN   rP   rT   rX   	resnet61qrl  )r/   r0   r1   r3   r4   r6   r;   r  )r   )rJ   rR   rN   rP   rT   rX   r;   resnext26tsi   r  rO   )rJ   rR   rN   rP   rX   gcresnext26tsgca)rJ   rR   rN   rP   rX   r7   seresnext26tseca_resnext26tsecabat_resnext26tsbatr  )
block_size)rJ   rR   rN   rP   rX   r7   r8   
resnet32ts
resnet33tsi   gcresnet33ts)rJ   rR   rN   rP   rT   rX   r7   seresnet33tseca_resnet33tsgcresnet50t)r/   r0   r1   r3   r6   )rJ   rR   rN   rP   r7   gcresnext50tsregnetz_b16   `         )r4  )r   r   )	rJ   rR   rP   rL   rT   rX   r7   r8   r;   regnetz_c16regnetz_d32i   )
rJ   rR   rN   rP   rL   rT   rX   r7   r8   r;   
regnetz_d8
regnetz_e8regnetz_b16_evos)rx   )
rJ   rR   rP   rL   rT   rX   rZ   r7   r8   r;   regnetz_c16_evosregnetz_d8_evosr  )rJ   rR   rN   rP   rL   rT   rX   rZ   r7   r8   r;   mobileone_s0)r$  r5   r5   r  )ri   r|   r   mobileone_s1)r(  r(  r  r%  mobileone_s2)r(  r  r%  r+  mobileone_s3)r  r%  r  r+  mobileone_s4)r        @rZ  r+  )r   r   r  r"   )ri   r{   resnet50_clip)rQ   rQ   rn   r  r   r  )rJ   rR   rN   rP   rL   r\   r_   resnet101_clip   resnet50x4_clip
   g      ?)rJ   rS   rR   rN   rP   rL   r\   r_   resnet50x16_clip   r(  resnet50x64_clip   $   r  resnet50_mlpr  )rJ   rR   rN   rP   rL   r\   r]   r_   test_byobnetr  rk  r3  rW   )rJ   rR   rL   rP   rX   r7   r8   )r[  r\  r^  r`  rb  r^   )r_   _gap
state_dictmodelprefixc                 x   t        |j                  t        t        f      }dd l}d }d }i }| j                         D ]  \  }}	|j                  |      s|j                  | dd|      }|j                  | dd|      }|j                  | d||      }|j                  | d	||      }|j                  | d
      rr|s|j                  |d
z   d      }|j                  dd      }|j                  dd      }|j                  dd      }|j                  dd      }|j                  dd      }|	||<    |S )Nr   c                     t        | j                  d            dz
  }t        | j                  d            | j                  d      t        | j                  d            }}}d| d| d}ddd	d
}||   |z   }||z   S )Nr"   r2   r   ry  r  r  z
conv1_1x1.z
conv2_kxk.z
conv3_1x1.)r"   r2   r   rD   group)r  r  	layer_idx
layer_typelayer_id
prefix_strid_map
suffix_strs           rH   
_stage_subz(_convert_openai_clip.<locals>._stage_sub  s    
Oa'	*-aggaj/1771:s177ST:x:	yk9+Q7
!l|DH%
2
J&&rG   c                     t        | j                  d            dz
  }t        | j                  d            t        | j                  d            }}d| d| d|dk(  rdz   S d	z   S )
Nr"   r2   r   r  r  z
.shortcut.r   z	conv.convzconv.bnrm  )r  r  ro  rq  s       rH   	_down_subz'_convert_openai_clip.<locals>._down_sub  sc    
Oa'	!!''!*os1771:8	1YKz:XYZ]kjj`ijjrG   zconv([0-9])zstem.conv\1.convz	bn([0-9])zstem.conv\1.bnz'layer([0-9])\.([0-9]+)\.([a-z]+)([0-9])z+layer([0-9])\.([0-9]+)\.downsample\.([0-9])attnpoolr  positional_embedding	pos_embedq_projqk_projkv_projvc_projproj)	r   r  r   r   reitems
startswithsubr   )
rh  ri  rj  model_has_attn_poolr  ru  rw  out_dictr~  r  s
             rH   _convert_openai_clipr    sT   
 %UZZ2Do1VW'k
 H  " 1||F#FFvhk*,?CFFvhi(*;Q?FFvhEF
TUVFFvhIJIWXY<<6((+,&		&:-v6A		0+>A		(C(A		(C(A		(C(A		(F+A!$ OrG   c                 &    d| v rt        | |      } | S )Nzvisual.conv1.weight)r  )rh  ri  s     rH   checkpoint_filter_fnr    s     
*)*e<
rG   variant
pretrainedc                 X    t        t        | |ft        |    t        t	        d      d|S )zCreate a ByobNet model.

    Args:
        variant: Model variant name.
        pretrained: Load pretrained weights.
        **kwargs: Additional model arguments.

    Returns:
        ByobNet model instance.
    T)flatten_sequential)r  pretrained_filter_fnfeature_cfg)r#   r)   
model_cfgsr  rc   )r  r  r   s      rH   _create_byobnetr    s:      *W%1D1	
  rG   urlc                 0    | dddddt         t        ddd
|S )	zCreate default configuration dictionary.

    Args:
        url: Model weight URL.
        **kwargs: Additional configuration options.

    Returns:
        Configuration dictionary.
    r  r      r  r  r        ?bilinear	stem.convhead.fc
r  r  
input_size	pool_sizecrop_pctinterpolationr  rJ  
first_convr^   r   r   r  r   s     rH   _cfgr    s2     4}SYJ%.B!	
  rG   c                 0    | dddddt         t        ddd
|S )	zCreate RepVGG configuration dictionary.

    Args:
        url: Model weight URL.
        **kwargs: Additional configuration options.

    Returns:
        Configuration dictionary.
    r  r   rp   rp   r  r  ?bicubiczstem.conv1.convr  r  r  r  s     rH   _cfgrr    s2     4}SY)%.B'y	
  rG   zgernet_s.idstcv_in1kztimm/)	hf_hub_idzgernet_m.idstcv_in1kzgernet_l.idstcv_in1kr  r  )r  r  r  zrepvgg_a0.rvgg_in1k)zstem.conv_kxk.convzstem.conv_1x1.convmit)r  r  licensezrepvgg_a1.rvgg_in1kzrepvgg_a2.rvgg_in1kzrepvgg_b0.rvgg_in1kzrepvgg_b1.rvgg_in1kzrepvgg_b1g4.rvgg_in1kzrepvgg_b2.rvgg_in1kzrepvgg_b2g4.rvgg_in1kzrepvgg_b3.rvgg_in1kzrepvgg_b3g4.rvgg_in1kzrepvgg_d2se.rvgg_in1k)r   @  r  )r_  r_  )r  r  r  r  r  r  zresnet51q.ra2_in1kzkhttps://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/resnet51q_ra2-d47dcc76.pthz
stem.conv1)r   rM  rM  )r  r  r  r  r  test_input_sizetest_crop_pctzresnet61q.ra2_in1kzkhttps://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/resnet61q_ra2-6afc536c.pth)r  r  r  r  zresnext26ts.ra2_in1kzvhttps://github.com/rwightman/pytorch-image-models/releases/download/v0.1-attn-weights/resnext26ts_256_ra2-8bbd9106.pthzseresnext26ts.ch_in1kzthttps://github.com/rwightman/pytorch-image-models/releases/download/v0.1-attn-weights/seresnext26ts_256-6f0d74a3.pthzgcresnext26ts.ch_in1kzthttps://github.com/rwightman/pytorch-image-models/releases/download/v0.1-attn-weights/gcresnext26ts_256-e414378b.pthzeca_resnext26ts.ch_in1kzvhttps://github.com/rwightman/pytorch-image-models/releases/download/v0.1-attn-weights/eca_resnext26ts_256-5a1d030f.pthzbat_resnext26ts.ch_in1kzvhttps://github.com/rwightman/pytorch-image-models/releases/download/v0.1-attn-weights/bat_resnext26ts_256-fa6fd595.pth)r  r  min_input_sizezresnet32ts.ra2_in1kzqhttps://github.com/rwightman/pytorch-image-models/releases/download/v0.1-attn-weights/resnet32ts_256-aacf5250.pthzresnet33ts.ra2_in1kzqhttps://github.com/rwightman/pytorch-image-models/releases/download/v0.1-attn-weights/resnet33ts_256-e91b09a4.pthzgcresnet33ts.ra2_in1kzshttps://github.com/rwightman/pytorch-image-models/releases/download/v0.1-attn-weights/gcresnet33ts_256-0e0cd345.pthzseresnet33ts.ra2_in1kzshttps://github.com/rwightman/pytorch-image-models/releases/download/v0.1-attn-weights/seresnet33ts_256-f8ad44d9.pthzeca_resnet33ts.ra2_in1kzuhttps://github.com/rwightman/pytorch-image-models/releases/download/v0.1-attn-weights/eca_resnet33ts_256-8f98face.pthzgcresnet50t.ra2_in1kzrhttps://github.com/rwightman/pytorch-image-models/releases/download/v0.1-attn-weights/gcresnet50t_256-96374d1c.pthzgcresnext50ts.ch_in1kzthttps://github.com/rwightman/pytorch-image-models/releases/download/v0.1-attn-weights/gcresnext50ts_256-3e0f515e.pthzregnetz_b16.ra3_in1kzphttps://github.com/rwightman/pytorch-image-models/releases/download/v0.1-attn-weights/regnetz_b_raa-677d9606.pthr  )r  r  r  r  r  gGz?)
r  r  r  r  rJ  r  r  r  r  r  zregnetz_c16.ra3_in1kzuhttps://github.com/rwightman/pytorch-image-models/releases/download/v0.1-attn-weights/regnetz_c_rab2_256-a54bf36a.pth)r  r  r  r  rJ  r  r  r  zregnetz_d32.ra3_in1kzthttps://github.com/rwightman/pytorch-image-models/releases/download/v0.1-attn-weights/regnetz_d_rab_256-b8073a89.pthgffffff?)r  r  r  rJ  r  r  zregnetz_d8.ra3_in1kzphttps://github.com/rwightman/pytorch-image-models/releases/download/v0.1-attn-weights/regnetz_d8_bh-afc03c55.pth)r  r  r  rJ  r  r  r  zregnetz_e8.ra3_in1kzphttps://github.com/rwightman/pytorch-image-models/releases/download/v0.1-attn-weights/regnetz_e8_bh-aace8e6e.pthzregnetz_b16_evos.untrained)r  r  rJ  r  r  r  r  zregnetz_c16_evos.ch_in1kzuhttps://github.com/rwightman/pytorch-image-models/releases/download/v0.1-tpu-weights/regnetz_c16_evos_ch-d8311942.pth)r  r  r  r  rJ  r  r  zregnetz_d8_evos.ch_in1kzthttps://github.com/rwightman/pytorch-image-models/releases/download/v0.1-tpu-weights/regnetz_d8_evos_ch-2bc12646.pthzmobileone_s0.apple_in1kr  )zstem.conv_kxk.0.convzstem.conv_scale.conv)r  r  r  zmobileone_s1.apple_in1kr  zmobileone_s2.apple_in1kzmobileone_s3.apple_in1kzmobileone_s4.apple_in1kzresnet50_clip.openaiz	head.proj)r  r  r  rJ  rV   r  r  r^   zresnet101_clip.openaizresnet50x4_clip.openai)	   r  zresnet50x16_clip.openaii   )r   r   r   )rL  rL  zresnet50x64_clip.openai)r     r  )r#  r#  zresnet50_clip.cc12mzresnet50_clip.yfcc15mzresnet101_clip.yfcc15mzresnet50_clip_gap.openai)r  r  r  rJ  r  r  zresnet101_clip_gap.openai)r  r  )r      r  )r  r  )r  r  r  r  r  )zresnet50x4_clip_gap.openaizresnet50x16_clip_gap.openaizresnet50x64_clip_gap.openaizresnet50_clip_gap.cc12mzresnet50_clip_gap.yfcc15mzresnet101_clip_gap.yfcc15mzresnet50_mlp.untrainedztest_byobnet.r160_in1kc                     t        dd| i|S )z GEResNet-Large (GENet-Large from official impl)
    `Neural Architecture Design for GPU-Efficient Networks` - https://arxiv.org/abs/2006.14090
    r  )r  r  r  r   s     rH   r  r  
      
 G*GGGrG   c                     t        dd| i|S )z GEResNet-Medium (GENet-Normal from official impl)
    `Neural Architecture Design for GPU-Efficient Networks` - https://arxiv.org/abs/2006.14090
    r  )r  r  r  s     rH   r  r  
  r  rG   c                     t        dd| i|S )z EResNet-Small (GENet-Small from official impl)
    `Neural Architecture Design for GPU-Efficient Networks` - https://arxiv.org/abs/2006.14090
    r  )r  r  r  s     rH   r  r  $
  r  rG   c                     t        dd| i|S )z^ RepVGG-A0
    `Making VGG-style ConvNets Great Again` - https://arxiv.org/abs/2101.03697
    r  )r"  r  r  s     rH   r"  r"  ,
      
 H:HHHrG   c                     t        dd| i|S )z^ RepVGG-A1
    `Making VGG-style ConvNets Great Again` - https://arxiv.org/abs/2101.03697
    r  )r&  r  r  s     rH   r&  r&  4
  r  rG   c                     t        dd| i|S )z^ RepVGG-A2
    `Making VGG-style ConvNets Great Again` - https://arxiv.org/abs/2101.03697
    r  )r'  r  r  s     rH   r'  r'  <
  r  rG   c                     t        dd| i|S )z^ RepVGG-B0
    `Making VGG-style ConvNets Great Again` - https://arxiv.org/abs/2101.03697
    r  )r)  r  r  s     rH   r)  r)  D
  r  rG   c                     t        dd| i|S )z^ RepVGG-B1
    `Making VGG-style ConvNets Great Again` - https://arxiv.org/abs/2101.03697
    r  )r*  r  r  s     rH   r*  r*  L
  r  rG   c                     t        dd| i|S )z` RepVGG-B1g4
    `Making VGG-style ConvNets Great Again` - https://arxiv.org/abs/2101.03697
    r  )r,  r  r  s     rH   r,  r,  T
      
 JZJ6JJrG   c                     t        dd| i|S )z^ RepVGG-B2
    `Making VGG-style ConvNets Great Again` - https://arxiv.org/abs/2101.03697
    r  )r-  r  r  s     rH   r-  r-  \
  r  rG   c                     t        dd| i|S )z` RepVGG-B2g4
    `Making VGG-style ConvNets Great Again` - https://arxiv.org/abs/2101.03697
    r  )r/  r  r  s     rH   r/  r/  d
  r  rG   c                     t        dd| i|S )z^ RepVGG-B3
    `Making VGG-style ConvNets Great Again` - https://arxiv.org/abs/2101.03697
    r  )r0  r  r  s     rH   r0  r0  l
  r  rG   c                     t        dd| i|S )z` RepVGG-B3g4
    `Making VGG-style ConvNets Great Again` - https://arxiv.org/abs/2101.03697
    r  )r1  r  r  s     rH   r1  r1  t
  r  rG   c                     t        dd| i|S )z` RepVGG-D2se
    `Making VGG-style ConvNets Great Again` - https://arxiv.org/abs/2101.03697
    r  )r2  r  r  s     rH   r2  r2  |
  r  rG   c                     t        dd| i|S )
    r  )r6  r  r  s     rH   r6  r6  
       H:HHHrG   c                     t        dd| i|S )r  r  )r8  r  r  s     rH   r8  r8  
  r  rG   c                     t        dd| i|S )r  r  )r9  r  r  s     rH   r9  r9  
       JZJ6JJrG   c                     t        dd| i|S )r  r  )r:  r  r  s     rH   r:  r:  
       LzLVLLrG   c                     t        dd| i|S )r  r  )r<  r  r  s     rH   r<  r<  
  r  rG   c                     t        dd| i|S )r  r  )r=  r  r  s     rH   r=  r=  
       NNvNNrG   c                     t        dd| i|S )r  r  )r?  r  r  s     rH   r?  r?  
  r  rG   c                     t        dd| i|S )r  r  )rB  r  r  s     rH   rB  rB  
       IJI&IIrG   c                     t        dd| i|S )r  r  )rC  r  r  s     rH   rC  rC  
  r  rG   c                     t        dd| i|S )r  r  )rD  r  r  s     rH   rD  rD  
       KjKFKKrG   c                     t        dd| i|S )r  r  )rE  r  r  s     rH   rE  rE  
  r  rG   c                     t        dd| i|S )r  r  )rF  r  r  s     rH   rF  rF  
       M
MfMMrG   c                     t        dd| i|S )r  r  )rG  r  r  s     rH   rG  rG  
  r  rG   c                     t        dd| i|S )r  r  )rH  r  r  s     rH   rH  rH  
  r  rG   c                     t        dd| i|S )r  r  )rI  r  r  s     rH   rI  rI  
  r  rG   c                     t        dd| i|S )r  r  )rN  r  r  s     rH   rN  rN  
  r  rG   c                     t        dd| i|S )r  r  )rO  r  r  s     rH   rO  rO  
  r  rG   c                     t        dd| i|S )r  r  )rP  r  r  s     rH   rP  rP  
  r  rG   c                     t        dd| i|S )r  r  )rQ  r  r  s     rH   rQ  rQ    r  rG   c                     t        dd| i|S )r  r  )rR  r  r  s     rH   rR  rR  	       O*OOOrG   c                     t        dd| i|S )r  r  )rS  r  r  s     rH   rS  rS    r  rG   c                     t        dd| i|S )r  r  )rT  r  r  s     rH   rT  rT    r  rG   c                     t        dd| i|S )r  r  )rU  r  r  s     rH   rU  rU    r  rG   c                     t        dd| i|S )r  r  )rV  r  r  s     rH   rV  rV  %  r  rG   c                     t        dd| i|S )r  r  )rW  r  r  s     rH   rW  rW  ,  r  rG   c                     t        dd| i|S )r  r  )rX  r  r  s     rH   rX  rX  3  r  rG   c                     t        dd| i|S )r  r  )rY  r  r  s     rH   rY  rY  :  r  rG   c                     t        dd| i|S )z0 OpenAI Modified ResNet-50 CLIP image tower
    r  )r[  r  r  s     rH   r[  r[  A  r  rG   c                     t        dd| i|S )z1 OpenAI Modified ResNet-101 CLIP image tower
    r  )r\  r  r  s     rH   r\  r\  H  r  rG   c                     t        dd| i|S )z2 OpenAI Modified ResNet-50x4 CLIP image tower
    r  )r^  r  r  s     rH   r^  r^  O  r  rG   c                     t        dd| i|S )z3 OpenAI Modified ResNet-50x16 CLIP image tower
    r  )r`  r  r  s     rH   r`  r`  V  r  rG   c                     t        dd| i|S )z3 OpenAI Modified ResNet-50x64 CLIP image tower
    r  )rb  r  r  s     rH   rb  rb  ]  r  rG   c                     t        dd| i|S )zP OpenAI Modified ResNet-50 CLIP image tower w/ avg pool (no attention pool)
    r  )resnet50_clip_gapr  r  s     rH   r  r  d  s     P:PPPrG   c                     t        dd| i|S )zQ OpenAI Modified ResNet-101 CLIP image tower w/ avg pool (no attention pool)
    r  )resnet101_clip_gapr  r  s     rH   r  r  k  s     QJQ&QQrG   c                     t        dd| i|S )zR OpenAI Modified ResNet-50x4 CLIP image tower w/ avg pool (no attention pool)
    r  )resnet50x4_clip_gapr  r  s     rH   r  r  r  s     RZR6RRrG   c                     t        dd| i|S )zS OpenAI Modified ResNet-50x16 CLIP image tower w/ avg pool (no attention pool)
    r  )resnet50x16_clip_gapr  r  s     rH   r  r  y       SjSFSSrG   c                     t        dd| i|S )zS OpenAI Modified ResNet-50x64 CLIP image tower w/ avg pool (no attention pool)
    r  )resnet50x64_clip_gapr  r  s     rH   r  r    r  rG   c                     t        dd| i|S )r  r  )re  r  r  s     rH   re  re    r  rG   c                     t        dd| i|S )z- Minimal test ResNet (BYOB based) model.
    r  )rf  r  r  s     rH   rf  rf    r  rG   ))ry  r  rJ  r"   r5   r5   r5   r5   r   ))r2   r  r_  r"   r  rF   r"   )r"   F)r[   r[   r  N)r2   r  )r[   F)zvisual.r   )r[   )r?   r  dataclassesr   r   r   	functoolsr   typingr   r   r	   r
   r   r   r   r   r   torch.nnrA   	timm.datar   r   r   r   timm.layersr   r   r   r   r   r   r   r   r   r   r   r   r   r    r!   _builderr#   	_featuresr$   _manipulater%   r&   	_registryr'   r(   __all__r+   r*   rD   rE   rz   r   r@   rh   r   r   r   r   rB   r   r   r   r   r   r  r
  rM  ra  rc   ro  rp  r-   r  rt  r,   r  r  r  r  r  r)   r  r  r~  r   r  r  r  r  r  default_cfgsr  r  r  r"  r&  r'  r)  r*  r,  r-  r/  r0  r1  r2  r6  r8  r9  r:  r<  r=  r?  rB  rC  rD  rE  rF  rG  rH  rI  rN  rO  rP  rQ  rR  rS  rT  rU  rV  rW  rX  rY  r[  r\  r^  r`  rb  r  r  r  r  r  re  rf  rF   rG   rH   <module>r     sM$  <  1 1  N N N   d d    
 + + 4 <
W 2 2 2( I I IBU38_ %s
:K hk  k3 * + 0%'!"	"c?"%*" c?" 	"
 
${
"P ()	S#X S$s)^$ 	 ;@k8K;P.P(Q VZ[fVg "&8C= &C &C &$ ) ) ),'BII ,'^      	 
 S/    bii F8 8vBbii BJ?		 ?D=		 =@`=")) `=Fz=RYY z=zBBII BJ 
		+s +bii +,c299n- ,K2== Kb !,,, , 	,
 , ,^R	!Yd38n !Y !YYd !YR $($(.A<;<;<; <; S>	<;
 C=<; !<; "(+<;~
{ 
d 
|bii |~	;")) ;3 ;T ;VZ ;2  xWSA!CWSA!CXcQ1GXcQ1DXcQ1D
 x WSA!CWSA!CXcQ1GXcQ1DXcQ1D
 x2 WR1rBWR1rBXcQ1GXcQ1DXcQ1D
 3xL }1HIMxV }@Wx` }1FGaxj  12kxt  01ux~  0;xH  34IxR  3A>Sx\  01]xf  0;gxp ~2EF&Q7qxD XcQ2$GXcQ2$GXdaB4HXdaA#F	
 Exd VqC1sQUQWXXcQ2$GXdaB4HXdaA#F	
 T*exF XcQ2$GXcQ2$GXdaB4HXdaB4H	
 Gx^ XcQ2$GXcQ2$GXdaB4HXdaB4H	
 _xx XcQ2$GXcQ2$GXdaB4HXdaB4H	
 yxR  XcQ2$GXcQ2$GXdaB4HXdaB4H	
 Sxl  XcQ2$GXcQ2$GXdaB4HXdaB4H	
 A&mxL XcQ1FXcQ1FXdaA$GXdaA$G	
 Mxj XcQ1FXcQ1FXdaA$GXdaA$G	
 kxJ XcQ1FXcQ1FXdaA$GXdaA$G	
 Kxf XcQ1FXcQ1FXdaA$GXdaA$G	
 gxB XcQ1FXcQ1FXdaA$GXdaA$G	
 Cx` XcQ4@XcQ4@XdaDAXdaDA	
 axz XcQ2$GXcQ2$GXdaB4HXdaB4H	
 {xX	 XbA"CXbA"CXsaB1EXcQ2!D	
 $'DT:Y	xx	 XbA"CXbA"CXsaB1EXcQ2!D	
 $'DT:y	xX
 XbA"CXcQ2!DXsaB1EXcQ2!D	
 $'DT:Y
xz
 XbA!BXcQ1CXsaA!DXcQ1C	
 $'DT:{
x\ XbA!BXcQ1CXsaA!DXcQ1C	
 $'DT:]xB !XbA"CXbA"CXsaB1EXcQ2!D	
 <B7$'DT:Cxd !XbA"CXbA"CXsaB1EXcQ2!D	
 <B7$'DT:exF  XbA!BXcQ1CXsaA!DXcQ1C	
 <B7$'DT:!Gxl "6!Lmxv "67wx@ "67AxJ "67KxT "6,OUx` XcQ4@XcQ4@XdaDAXdaDA	
 ax| XcQ4@XcQ4@XtqTBXdaDA	
 }xX  XcQ4@XcQ4@XtqTBXdaDA	
 Yxv !XcQ4@XcQ4@XtqTBXdaDA	
 wxT !XcQ4@XsaDAXtqTBXtqTB	
 Uxt XcQ4@XcQ4@XdaDAXdaDA	
 uxT VqB!cBVqB!cBWSA"FXcQ2$G	
 $'Ux
r 
h LA$Z]lKJq6zL  (ell*+(( ( 
#u||
	(Vell*+S d  (c T#s(^ &s d38n & % P&D73P& D73P& D7}X^_	P& 4?PP& 4?PP& 4?PP&  4?P!P&& 4?P'P&, T?P-P&2 4?P3P&8 T?P9P&> 4?P?P&D T?PEP&J T? HsKP&X $yMV%S	:YP&b %y%S:cP&n E E%S:oP&v U C%S:wP&~ U C%S:P&F u E%S :GP&N u E$ &OP&Z 5%S:[P&b 5%S:cP&j U B%S:kP&r U B%S:sP&z u D%S :{P&D E A%S:EP&N U C%S:OP&Z E~_/ FTS`ps	u[P&d E D_/}C	IeP&n E C/DR_aoP&v 5~/DR_ortwP&~ 5~/DR_ortP&H !%_/ FTS`#bIP&N  D_/}	!6OP&X u C/DR_or tYP&b tC cP&l tC mP&v tC wP&@ tC AP&J tC KP&X E/_-6	YP&d U.O-6	eP&p e.O-6	qP&| u.O-8	 }P&H u/_-8	 IP&T 5/_-6	UP&` U/_-6	aP&l e.O-6	mP&| ,/ F!}P&F  ,/ F"GP&P #(,/ F#
 $),/ H$
 $),/ H$
  %,/ F 
 "',/ F"
 #(,/ F# $ F $ 46WP& Pf HG H H HG H H HG H H IW I I IW I I IW I I IW I I IW I I Kw K K IW I I Kw K K IW I I Kw K K Kw K K IW I I IW I I Kw K K M M M M M M O7 O O O7 O O Jg J J Jg J J L L L L L L N' N N Kw K K M M M Kw K K Kw K K Kw K K Jg J J Jg J J PG P P PG P P O7 O O L L L L L L L L L L L L L L L M M M N' N N O7 O O PG P P PG P P QW Q Q Rg R R Sw S S T T T T T T L L L L L LrG   