a
    da                 0   @   s  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 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# dd	l$m%Z% dd
l&m'Z'm(Z( ddl)m*Z*m+Z+ g dZ,eG dd dZ-eG dd dZ.dddZ/de	e0e0f ee1e
e1 f e2e	e- dddZ3ee-ee- f e
e- dddZ4dd Z5eG dd dZ6G d d! d!ej7Z8e0e1e1e1e	e1e1f e6d"d#d$Z9G d%d& d&ej7Z:G d'd( d(ej7Z;G d)d* d*ej7Z<G d+d, d,ej7Z=G d-d. d.ej7Z>G d/d0 d0ej7Z?e@e:e;e<e=e>e?d1ZAe0ej7d2d3d4ZBee0ej7f d5d6d7ZCG d8d9 d9ejDZEde1e1e0e0e0e6d<d=d>ZFdd@dAZGdBdC ZHee0ef e-e.dDdEdFZIddeIfe.eJe1ee0ef ee1 ee6 ee dGdHdIZKe.dJdKdLZLG dMdN dNej7ZMddOdPZNe@e.e-dQddRd?ddSdTe-dQd?dUd?ddSdTe-dVdWdXd?ddYdTe-dVdZdXd?dd[dTe-dVd\dXddd[dTfd]dd^d_e.e-dQddRd?ddSdTe-dQd?dUd?ddSdTe-dVdWdXd?ddYdTe-dVd\dXd?dd[dTe-dVddXddd[dTfd]dd^d_e.e-dQdd`d?ddSdTe-dQdad`d?ddSdTe-dVdbdcd?ddYdTe-dVd?ddd?dd[dTe-dVddeddd[dTfdfddgd_e.e/dhdidjdkdldme.e/dndodkdldme.e/dpdodkdldme.e/dpd\dqdkdldme.e/drdodkdldme.e/drd\dqdkdldme.e/dsdodkdldme.e/dsd\dqdkdldme.e-dVd?dedd]dYdTe-dVd\dtd?d]dYdTe-dVdWdud?d]dYdTe-dVd\dud?ddSdTfdRdvddwdxdye.e-dzddedddSe@ d{e-dVd\dtd?d]dYdTe-dVdWdud?d]dYdTe-dVd\dud?ddSdTfdRd|ddwdxe@d}d~de.e-dVd?dedd]dYdTe-dVd?dtd?d]dYdTe-dVd?dd?d]dYdTe-dVd?dwd?d]dYdTfdldddxde.e-dVd?dedd]dYdTe-dVd?dtd?d]dYdTe-dVd?dd?d]dYdTe-dVd?dwd?d]dYdTfdldddxdde.e-dVd?dedd]dYdTe-dVd?dtd?d]dYdTe-dVd?dd?d]dYdTe-dVd?dwd?d]dYdTfdldddxdde.e-dVd?dedd]dYdTe-dVd?dtd?d]dYdTe-dVd?dd?d]dYdTe-dVd?dwd?d]dYdTfdldddxdde.e-dVd?dedd]dYdTe-dVd?dtd?d]dYdTe-dVd?dd?d]dYdTe-dVd?dwd?d]dYdTfdldddxde@ddde.e-dVd?dedddYdTe-dVdadtd?ddYdTe-dVdadud?ddYdTe-dVd?dud?ddYdTfdldd:ddxdye.e-dVd?dedddYdTe-dVdadtd?ddYdTe-dVdadud?ddYdTe-dVd?dud?ddYdTfdldd:ddxdye.e-dVd?dedddYdTe-dVdadtd?ddYdTe-dVdadud?ddYdTe-dVd?dud?ddYdTfdldd:ddxdde.e-dVd?dedddYdTe-dVdadtd?ddYdTe-dVdadud?ddYdTe-dVd?dud?ddYdTfdldd:ddxdde.e-dVd?dedddYdTe-dVdadtd?ddYdTe-dVdadud?ddYdTe-dVd?dud?ddYdTfdldd:ddxdde.e-dVdadeddYde-dVd\dtd?dYde-dVdWdd?dYde-dVdadwd?dYdfdldd:dde.e-dVdadedd]dYdTe-dVd\dtd?d]dYdTe-dVdWdd?d]dYdTe-dVdadwd?d]dYdTfdldddxdde.e-dVd?d`d?ddadTe-dVdWdd?ddadTe-dVddUd?ddadTe-dVd?dd?ddadTfd]d:d:dudxde@dYde@d}d}dd	e.e-dVd?d`d?dd\dTe-dVdWdd?dd\dTe-dVddUd?dd\dTe-dVd?dd?dd\dTfd]d:d:dudxde@dYde@d}d}dd	e.e-dVdadldd]d\dTe-dVdWdRd?d]d\dTe-dVdded?d]d\dTe-dVdadcd?d]d\dTfdldd:d:ddxde@dYde@d}d}dd
e.e-dVdadlddd\dTe-dVdWdRd?dd\dTe-dVdded?dd\dTe-dVdadcd?dd\dTfdldd:d:ddxde@dYde@d}d}dd
e.e-dVdadddd\dTe-dVddUd?dd\dTe-dVddcd?dd\dTe-dVdadtd?dd\dTfdldd:d:dwdxde@dYde@d}d}dd
e.e-dVd?d`d?ddadTe-dVdWdd?ddadTe-dVddUd?ddadTe-dVd?dd?ddadTfd]d:d:dudxee#ddde@dYde@d}d}dd
e.e-dVd?d`d?dd\dTe-dVdWdd?dd\dTe-dVddUd?dd\dTe-dVd?dd?dd\dTfd]d:d:dudxee#ddde@dYde@d}d}dd
e.e-dVdadlddd\dTe-dVdWdRd?dd\dTe-dVdded?dd\dTe-dVdadcd?dd\dTfdldd:d:ddxee#ddde@dYde@d}d}ddd!ZOdddZPdddZQdddZRe*eQddeQddeQddddeQddddeQddddeQddddeQddddeQddddeQddddeQddddeQddddeQdddddddSdeRddddSdeRddddSdeRddddSdeRddddSdeRddddSdeRddddeRddddSdeRddddSdeRddddSdeRddddSdeRddddSdeRddddSdeRddddSdeRddddddddddSdō
eRddddddddSdȍeRdddddddˍeRdddddddSd͍eRdddddddSd͍eRddddddddύeRddddddddэeRdddddddSd͍dӜ!ZSe+deMdԜddքZTe+d eMdԜdd؄ZUe+d!eMdԜddڄZVe+d"eMdԜdd܄ZWe+d#eMdԜddބZXe+d$eMdԜddZYe+d%eMdԜddZZe+d&eMdԜddZ[e+d'eMdԜddZ\e+d(eMdԜddZ]e+d)eMdԜddZ^e+d*eMdԜddZ_e+d+eMdԜddZ`e+d,eMdԜddZae+d-eMdԜddZbe+d.eMdԜddZce+d/eMdԜddZde+d0eMdԜddZee+d1eMdԜddZfe+d2eMdԜddZge+d3eMdԜddZhe+d4eMdԜdd Zie+d5eMdԜddZje+d6eMdԜddZke+d7eMdԜddZle+d8eMdԜddZme+d9eMdԜd	d
Zne+d:eMdԜddZoe+d;eMdԜddZpe+d<eMdԜddZqe+d=eMdԜddZre+d>eMdԜddZse+d?eMdԜddZtdS (@  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

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)ClassifierHeadConvNormActBatchNormAct2dDropPathAvgPool2dSamecreate_conv2dget_act_layerget_norm_act_layerget_attnmake_divisible	to_2tupleEvoNorm2dS0a   )build_model_with_cfg)named_applycheckpoint_seq)generate_default_cfgsregister_model)ByobNetByoModelCfgByoBlockCfgcreate_byob_stemcreate_blockc                   @   s   e Zd ZU 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< dS )r%   typedc   sNgs      ?br
attn_layerattn_kwargsself_attn_layerself_attn_kwargsblock_kwargs)__name__
__module____qualname__r
   strnnModule__annotations__intr,   r-   r	   r   r/   floatr0   r1   r   r   r2   r3   r4    r>   r>   \/var/www/html/stable-diffusion-webui/venv/lib/python3.9/site-packages/timm/models/byobnet.pyr%   -   s   
r%   c                   @   s  e Zd ZU 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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< edd dZeed< dZee ed< edd dZeed< ed d dZeeef ed!< dS )"r$   .blocksconv1x1
downsample3x3	stem_typemaxpool	stem_pool    stem_chsr.   width_factorr   num_featuresTzero_init_lastFfixed_input_sizeZrelu	act_layerZ	batchnorm
norm_layerNr0   c                   C   s   t  S Ndictr>   r>   r>   r?   <lambda>O       zByoModelCfg.<lambda>)default_factoryr1   r2   c                   C   s   t  S rO   rP   r>   r>   r>   r?   rR   Q   rS   r3   c                   C   s   t  S rO   rP   r>   r>   r>   r?   rR   R   rS   r4   )r5   r6   r7   r   r
   r%   r;   rB   r8   rD   rF   r	   rH   r<   rI   r=   rJ   rK   boolrL   rM   rN   r0   r   r1   rQ   r2   r3   r4   r   r   r>   r>   r>   r?   r$   >   s    
 r$            r   r.   r.   r.   r.   c                    s>   d}d dkrfdd t  fddt| ||D }|S )N)@            r   c                    s   |d d dkr|   S dS )Nr   r+   r   r>   )Zchsidx)groupsr>   r?   rR   Y   rS   z_rep_vgg_bcfg.<locals>.<lambda>c                    s&   g | ]\}}}t d |||  dqS )rep)r(   r)   r*   r-   )r%   ).0r)   r*   wf
group_sizer>   r?   
<listcomp>Z   rS   z!_rep_vgg_bcfg.<locals>.<listcomp>)tuplezip)r)   rc   r`   r*   Zbcfgr>   )re   r`   r?   _rep_vgg_bcfgU   s    ri   F)typeseveryfirstreturnc                 K   s   t | dksJ t|trDtt|r&dn|||d }|sD|d g}t| g }t|D ]6}||v rl| d n| d }|tf |dd|g7 }qXt|S )z' interleave 2 block types in stack
    r+   r   r   )r(   r)   )len
isinstancer<   listrangesetr%   rg   )rj   r)   rk   rl   kwargsr@   i
block_typer>   r>   r?   interleave_blocks^   s    

rv   )stage_blocks_cfgrm   c                    sF   t | ts| f} g }t| D ]$\} | fddt jD 7 }q|S )Nc                    s   g | ]}t  d dqS )r   r)   )r   )rb   _cfgr>   r?   rf   x   rS   z%expand_blocks_cfg.<locals>.<listcomp>)ro   r   	enumeraterq   r)   )rw   
block_cfgsrt   r>   rz   r?   expand_blocks_cfgs   s    
r~   c                 C   s$   | sdS ||  dksJ ||  S d S )Nr   r   r>   )re   Zchannelsr>   r>   r?   
num_groups|   s    r   c                   @   sT   e Zd ZU 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< dS )LayerFnconv_norm_actnorm_actactNattn	self_attn)r5   r6   r7   r   r   r   r;   r   r   r9   ZReLUr   r   r	   r   r>   r>   r>   r?   r      s
   
r   c                       s6   e Zd Zd	eeeeeed fddZdd Z  ZS )
DownsampleAvgr   FN)in_chsout_chsstridedilation	apply_actlayersc           	         s   t t|   |pt }|dkr$|nd}|dks8|dkrf|dkrL|dkrLtntj}|d|ddd| _n
t | _|j	||d|d| _
dS )z0 AvgPool Downsampling as in 'D' ResNet variants.r   r+   TF)Z	ceil_modeZcount_include_padr   N)superr   __init__r   r   r9   Z	AvgPool2dpoolIdentityr   conv)	selfr   r   r   r   r   r   Z
avg_strideZavg_pool_fn	__class__r>   r?   r      s    


zDownsampleAvg.__init__c                 C   s   |  | |S rO   )r   r   r   xr>   r>   r?   forward   s    zDownsampleAvg.forward)r   r   FN)	r5   r6   r7   r<   rU   r   r   r   __classcell__r>   r>   r   r?   r      s       r   )downsample_typer   r   r   r   r   c                 K   s   | dv sJ ||ks,|dks,|d |d krz| s4d S | dkrXt ||f||d d|S |j||fd||d d|S nt S d S )N)avgrA    r   r   r   r   r   )kernel_sizer   r   )r   r   r9   r   )r   r   r   r   r   r   rs   r>   r>   r?   create_shortcut   s    	 "r   c                       sd   e Zd ZdZdeeeeeeef ee eee	e	e
eed fddZde	dddZdd Z  ZS )
BasicBlockz$ ResNet Basic Block - kxk + kxk
       r   r   r   Nr.   r   TF        )r   r   r   r   r   re   bottle_ratiorB   	attn_last
linear_outr   
drop_blockdrop_path_ratec              	      s   t t|   |pt }t|| }t||}t|||||d|d| _|j|||||d d| _	|	sn|j
d u rvt n|
|| _
|j||||d ||dd| _|	r|j
d u rt n|
|| _|dkrt|nt | _|
rt n
|jdd	| _d S )
NFr   r   r   r   r   r   r   )r   r`   
drop_layerr   r   TZinplace)r   r   r   r   r   r   r   shortcutr   	conv1_kxkr   r9   r   	conv2_kxkr   r   	drop_pathr   )r   r   r   r   r   r   re   r   rB   r   r   r   r   r   mid_chsr`   r   r>   r?   r      s"    

""zBasicBlock.__init__rK   c                 C   s\   |r4| j d ur4t| jjdd d ur4tj| jjj | j| j	fD ]}t
|dr@|  q@d S Nweightreset_parametersr   getattrr   bnr9   initzeros_r   r   r   hasattrr   r   rK   r   r>   r>   r?   init_weights   s
    "
zBasicBlock.init_weightsc                 C   sN   |}|  |}| |}| |}| |}| jd urD|| | }| |S rO   )r   r   r   r   r   r   r   r   r   r>   r>   r?   r      s    




zBasicBlock.forward)r   r   r   Nr.   r   TFNNr   )F)r5   r6   r7   __doc__r<   r   r	   r=   r8   rU   r   r   r   r   r   r   r>   r>   r   r?   r      s:              
$r   c                       sh   e Zd ZdZdeeeeeeef eee ee	e	e	e	e
eed
 fddZde	dddZdd Z  ZS )BottleneckBlockz4 ResNet-like Bottleneck Block - 1x1 - kxk - 1x1
    r   r   r   r.   Nr   Fr   )r   r   r   r   r   r   re   rB   r   r   
extra_conv	bottle_inr   r   r   c              	      s:  t t|   |pt }t|r"|n|| }t||}t|||||d|d| _|||d| _	|j|||||d ||d| _
|r|j||||d |d| _n
t | _|	s|jd u rt n||| _|j||ddd| _|	r|jd u rt n||| _|dkrt|nt | _|
r(t n
|jd	d
| _d S )NFr   r   r   r   r   r`   r   )r   r`   r   r   Tr   )r   r   r   r   r   r   r   r   r   	conv1_1x1r   
conv2b_kxkr9   r   r   	conv3_1x1r   r   r   r   )r   r   r   r   r   r   r   re   rB   r   r   r   r   r   r   r   r   r`   r   r>   r?   r      s.    



""zBottleneckBlock.__init__r   c                 C   s\   |r4| j d ur4t| jjdd d ur4tj| jjj | j| j	fD ]}t
|dr@|  q@d S r   )r   r   r   r   r9   r   r   r   r   r   r   r   r   r>   r>   r?   r   &  s
    "
zBottleneckBlock.init_weightsc                 C   sl   |}|  |}| |}| |}| |}| |}| |}| |}| jd urb|| | }| |S rO   )	r   r   r   r   r   r   r   r   r   r   r>   r>   r?   r   -  s    







zBottleneckBlock.forward)r   r   r   r.   Nr   FFFFNNr   )Fr5   r6   r7   r   r<   r   r=   r	   r8   rU   r   r   r   r   r   r   r>   r>   r   r?   r      sB                
,r   c                       sd   e Zd ZdZdeeeeeeef eee ee	e	e
eed fddZde	dddZdd Z  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.   Nr   TFr   r   r   r   r   r   r   re   rB   r   r   r   r   r   c              
      s   t t|   |pt }t|| }t||}t|||||d|d| _|||d| _	|	sd|j
d u rlt n|
|| _
|j|||||d ||dd| _|	r|j
d u rt n|
|| _|dkrt|nt | _|
rt n
|jdd| _d S )	NFr   r   r   r   r   r`   r   r   r   Tr   )r   r   r   r   r   r   r   r   r   r   r   r9   r   r   r   r   r   r   r   r   r   r   r   r   r   re   rB   r   r   r   r   r   r   r`   r   r>   r?   r   F  s"    

""zDarkBlock.__init__r   c                 C   s\   |r4| j d ur4t| jjdd d ur4tj| jjj | j| j	fD ]}t
|dr@|  q@d S r   r   r   r>   r>   r?   r   j  s
    "
zDarkBlock.init_weightsc                 C   sX   |}|  |}| |}| |}| |}| |}| jd urN|| | }| |S rO   )r   r   r   r   r   r   r   r   r>   r>   r?   r   q  s    





zDarkBlock.forward)r   r   r   r.   Nr   TFNNr   )Fr   r>   r>   r   r?   r   ;  s:              
$r   c                       sd   e Zd ZdZdeeeeeeef eee ee	e	e
eed
 fddZde	dddZdd Z  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.   Nr   Fr   r   c              	      s   t t|   |pt }t|| }t||}t|||||d|d| _|j|||||d ||d| _	|	sr|j
d u rzt n|
|| _
|j||ddd| _|	r|j
d u rt n|
|| _|dkrt|nt | _|
rt n
|jdd	| _d S )
NFr   r   r   r   r   r   Tr   )r   r   r   r   r   r   r   r   r   r   r   r9   r   	conv2_1x1r   r   r   r   r   r   r>   r?   r     s"    

""zEdgeBlock.__init__r   c                 C   s\   |r4| j d ur4t| jjdd d ur4tj| jjj | j| j	fD ]}t
|dr@|  q@d S r   )r   r   r   r   r9   r   r   r   r   r   r   r   r   r>   r>   r?   r     s
    "
zEdgeBlock.init_weightsc                 C   sX   |}|  |}| |}| |}| |}| |}| jd urN|| | }| |S rO   )r   r   r   r   r   r   r   r   r>   r>   r?   r     s    





zEdgeBlock.forward)r   r   r   r.   Nr   FFNNr   )Fr   r>   r>   r   r?   r   }  s:              
$r   c                       s`   e Zd ZdZdeeeeeeef eee ee	e
ed	 fd
dZdedddZdd Z  ZS )RepVggBlockz RepVGG Block.

    Adapted from impl at https://github.com/DingXiaoH/RepVGG

    This version does not currently support the deploy optimization. It is currently fixed in 'train' mode.
    r   r   r   r.   Nr   r   )r   r   r   r   r   r   re   rB   r   r   r   c              
      s   t t|   |	pt }	t||}||ko@|dko@|d |d k}|rT|	j|ddnd | _|	j|||||d ||
dd| _|	j||d||dd| _	|	j
d u rt n|	
|| _
|dkr|rt|nt | _|	jdd	| _d S )
Nr   r   Fr   r   )r   r`   r   r   Tr   )r   r   r   r   r   r   identityr   conv_kxkconv_1x1r   r9   r   r   r   r   )r   r   r   r   r   r   r   re   rB   r   r   r   r`   Z	use_identr   r>   r?   r     s    

 zRepVggBlock.__init__Fr   c                 C   sX   |   D ]4}t|tjrtj|jdd tj|jdd qt| j	drT| j	
  d S )Ng?r   r   )modulesro   r9   BatchNorm2dr   normal_r   biasr   r   r   )r   rK   mr>   r>   r?   r     s    zRepVggBlock.init_weightsc                 C   sd   | j d u r | || | }n0|  |}| || | }| |}|| }| |}| |S rO   )r   r   r   r   r   r   )r   r   r   r>   r>   r?   r     s    



zRepVggBlock.forward)	r   r   r   r.   Nr   NNr   )F)r5   r6   r7   r   r<   r   r=   r	   r8   r   r   r   rU   r   r   r   r>   r>   r   r?   r     s2            
	r   c                       sv   e Zd ZdZdeeeeeeef eee ee	e	e	e	eeeef  e
eed fddZde	dddZdd Z  ZS )SelfAttnBlockzI ResNet-like Bottleneck Block - 1x1 - optional kxk - self attn - 1x1
    r   r   r   r.   Nr   FTr   )r   r   r   r   r   r   re   rB   r   r   r   post_attn_na	feat_sizer   r   r   c              	      s&  t t|   |d usJ t|r$|n|| }t||}t|||||d|d| _|||d| _|	r|j|||||d ||d| _	d}n
t
 | _	|d u ri nt|d}|j|fd|i|| _|r||nt
 | _|j||ddd| _|d	krt|nt
 | _|
rt
 n
|jd
d| _d S )NFr   r   r   r   )r   r   r   r   Tr   )r   r   r   r   r   r   r   r   r   r   r9   r   rQ   r   r   	post_attnr   r   r   r   )r   r   r   r   r   r   r   re   rB   r   r   r   r   r   r   r   r   r   r`   Z
opt_kwargsr   r>   r?   r     s,    

zSelfAttnBlock.__init__r   c                 C   sN   |r4| j d ur4t| jjdd d ur4tj| jjj t| j	drJ| j	
  d S r   )r   r   r   r   r9   r   r   r   r   r   r   )r   rK   r>   r>   r?   r   *  s    "zSelfAttnBlock.init_weightsc                 C   sb   |}|  |}| |}| |}| |}| |}| |}| jd urX|| | }| |S rO   )r   r   r   r   r   r   r   r   r   r>   r>   r?   r   0  s    






zSelfAttnBlock.forward)r   r   r   r.   Nr   FFFTNNNr   )Fr   r>   r>   r   r?   r     sF                 
.r   )basicbottleZdarkedgera   r   ru   Zblock_fnc                 C   s   |t | < d S rO   )_block_registryr   r>   r>   r?   register_blockG  s    r   )blockc                 K   sF   t | tjtfr| f i |S | tv s4J d|  t|  f i |S )NzUnknown block type ()ro   r9   r:   r   r   )r   rs   r>   r>   r?   r'   K  s    r'   c                       s8   e Zd Zd	eeeeeeee eed	 fddZ  Z	S )
Stemr   rW   rE   N      ?)	r   r   r   r   r   num_repnum_act	chs_decayr   c
              
      s  t    |dv sJ |	pt }	tttfr<t}}
n" fddt|D d d d }
|| _g | _	d}dgdg|d   }|dkr|sd|d< |d u r|n|}d	g||  d
g|  }|}d}t
t|
||D ]|\}\}}}|r|	jnt}d|d  }|dkr*|dkr*| j	t|||d | ||||||d |}||9 }|}q|rd| v r| j	t|||d | dtddd |d9 }d}| j	t|||d ||ksJ d S )N)r+   rW   c                    s   g | ]}t  |  qS r>   )round)rb   rt   r   r   r>   r?   rf   h  rS   z!Stem.__init__.<locals>.<listcomp>r   r+   r   rW   FTr   r   num_chs	reductionmodule)r   r   maxr   r   )r   r   r   ro   rp   rg   rn   rq   r   feature_infor|   rh   r   r   appendrQ   Z
add_modulelowerr9   Z	MaxPool2d)r   r   r   r   r   r   r   r   r   r   rH   	prev_featZstem_stridesZstem_norm_actsprev_chsZcurr_stridert   chr,   nalayer_fnZ	conv_namer   r   r?   r   T  sB    

"zStem.__init__)r   rW   rE   r   Nr   N)
r5   r6   r7   r<   r8   r	   r=   r   r   r   r>   r>   r   r?   r   R  s$          r   r   stem)r   r   rD   	pool_typefeat_prefixr   c           	         sD  |pt  }|dv sJ d|v rDd|v r*dnd }t| |d|||d}nd|v rpt| d| d	 |d |f||d
}nd|v rt| |dd||d}n|d|v rt| |d|d}nbd|v r|rt| |dd||d}n|j| |ddd}n,|rt| |dd||d}n|j| |ddd}t|tr, fdd|jD }nt|d dg}||fS )N)r   quadquad2tiereddeepra   7x7rC   r   r   r+   rW   )r   r   r   r   r   r      )r   r   r   r.   )r   r   r   r   ra   )r   r   r      r   )r   r   r   r   c              	      s&   g | ]}t |d  |d gdqS ).r   )r   )rQ   join)rb   fr   r>   r?   rf     rS   z$create_byob_stem.<locals>.<listcomp>r   )r   r   r   r   ro   r   rQ   )	r   r   rD   r   r   r   r   r   r   r>   r  r?   r&     s,    
$r&   r+   c                    s"   | d u rd S t  fdd| D S )Nc                    s   g | ]}|  qS r>   r>   rb   r,   r   r>   r?   rf     rS   z$reduce_feat_size.<locals>.<listcomp>)rg   )r   r   r>   r   r?   reduce_feat_size  s    r  c                 C   s   | dur| n|}|p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
    Nr>   )r4   Zmodel_kwargsZ
out_kwargsr>   r>   r?   override_kwargs  s    r  )r4   	block_cfg	model_cfgc           
      C   s  | d }|j d u}|s |jd urv|r0|j s0d }n:t|j|j}|j pH|j }|d urftt|fi |nd }t||d}|jd u}|s|jd ur|r|jsd }n:t|j|j}	|jp|j}|d urtt|fi |	nd }t||d}|| d< | t|j	|j	 d S )Nr   )r   r   )
r0   r1   r  r   r   r   r2   r3   updater4   )
r4   r  r	  Z	layer_fnsZattn_setr0   r1   Zself_attn_setr2   r3   r>   r>   r?   update_block_kwargs  s,    

 

r  )r{   r   output_stride	stem_featr   r   block_kwargs_fnc                 C   s  |pt  }g }dd | jD }dd |D }	dd td|t|	|	D }
d}|d }|d }|}g }t|D ]V\}}|d j}|dkr|r|| ||kr|dkr||9 }d}||9 }|d	v rdnd
}g }t|D ]\}}t	|j
| j }|j}t|tr|||}t|||dkr$|nd||f||j| j|
| | |d	}|jdv r\||d< |||| d |t|jfi |g7 }|}|}|dkr|dkrt||}q|tj| g7 }t||d| d}qr|| tj| |fS )Nc                 S   s   g | ]}t |qS r>   )r~   r  r>   r>   r?   rf     rS   z&create_byob_stages.<locals>.<listcomp>c                 S   s   g | ]}t d d |D qS )c                 S   s   g | ]
}|j qS r>   rx   )rb   Zbcr>   r>   r?   rf     rS   z1create_byob_stages.<locals>.<listcomp>.<listcomp>)sum)rb   Z	stage_bcsr>   r>   r?   rf     rS   c                 S   s   g | ]}|  qS r>   )tolist)rb   r   r>   r>   r?   rf     rS   r   r   r   r   )r   r+   r+   )	r   r   r   r   re   r   rB   r   r   r
  r   )r  r	  zstages.r   )r   r@   torchZlinspacer  splitr|   r,   r   r   r*   rI   r-   ro   r   rQ   r/   rB   r(   r'   r  r9   
Sequential)r{   r   r  r  r   r   r  r   r}   ZdepthsZdprr   Z
net_strider   r   stagesZ	stage_idxZstage_block_cfgsr   Zfirst_dilationr@   Z	block_idxr  r   re   r4   r>   r>   r?   create_byob_stages  s`    

"




r  rz   c                 C   s   t | j}t| j|d}tt| j|d}| jrFtt| jfi | jnd }| j	rhtt| j	fi | j
nd }t|||||d}|S )N)rN   rM   )r   r   r   r   r   )r   rM   r   rN   r   r   r0   r   r1   r2   r3   r   )r{   r   r   r   r   r   r   r>   r>   r?   get_layer_fns)  s    
""r  c                       s   e Zd ZdZdeeeeeeeee	eef f  e
e
ed		 fd
dZejjdddZejjdddZejjdd ZdddZdd Zd edddZdd Z  ZS )!r#   a#   'Bring-your-own-blocks' Net

    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).
      r   r   rG   Nr   T)	r{   num_classesin_chansglobal_poolr  img_size	drop_rater   rK   c
                    s  t    || _|| _d| _t|fi |
}t|}|jrJ|dusJJ d|durZt|nd}g | _	t
t|jpx|jd j|j }t|||j|j|d\| _}| j	|dd  t||d d d}t||||d ||d	\| _}| j	|dd  |d d
 }|jr8t
t|j|j | _||| jd| _n|| _t | _|  j	t| j|d d ddg7  _	t| j||| jd| _t t!t"|	d|  dS )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.
        FNz8img_size argument is required for fixed input size modelr   )r   r   r   r   )r   r   r   r   
final_convr   )r   r  r   )#r   r   r  r  grad_checkpointingr   r  rL   r   r   r<   r   rH   r@   r*   rI   r&   rD   rF   r   extendr  r  r  rJ   r   r  r9   r   rQ   r   headr   r   _init_weights)r   r{   r  r  r  r  r  r  r   rK   rs   r   r   rH   r  Z
stage_featr   r   r>   r?   r   ;  sN    
 
zByobNet.__init__Fc                 C   s    t d|rdndd fdgd}|S )Nz^stemz^stages\.(\d+)z^stages\.(\d+)\.(\d+))z^final_conv)i )r   r@   rP   )r   ZcoarseZmatcherr>   r>   r?   group_matcher  s    zByobNet.group_matcherc                 C   s
   || _ d S rO   )r  )r   enabler>   r>   r?   set_grad_checkpointing  s    zByobNet.set_grad_checkpointingc                 C   s   | j jS rO   )r!  Zfc)r   r>   r>   r?   get_classifier  s    zByobNet.get_classifierc                 C   s   | j || d S rO   )r!  reset)r   r  r  r>   r>   r?   reset_classifier  s    zByobNet.reset_classifierc                 C   s@   |  |}| jr(tj s(t| j|}n
| |}| |}|S rO   )r   r  r  jitZis_scriptingr    r  r  r   r>   r>   r?   forward_features  s    


zByobNet.forward_features
pre_logitsc                 C   s   | j ||dS )Nr+  )r!  )r   r   r,  r>   r>   r?   forward_head  s    zByobNet.forward_headc                 C   s   |  |}| |}|S rO   )r*  r-  r   r>   r>   r?   r     s    

zByobNet.forward)r  r   r   rG   Nr   r   T)F)T)r   )F)r5   r6   r7   r   r$   r<   r8   r	   r
   r   r=   rU   r   r  r)  ignorer#  r%  r&  r(  r*  r-  r   r   r>   r>   r   r?   r#   3  s<   
        I


	r#   c                 C   s   t | tjrb| jd | jd  | j }|| j }| jjdt	
d|  | jd ur| jj  nzt | tjrtjj| jddd | jd urtj| j n@t | tjrtj| j tj| j nt| dr| j|d d S )	Nr   r          @r   g{Gz?)meanstdr   r   )ro   r9   ZConv2dr   Zout_channelsr`   r   datar   mathsqrtr   Zzero_ZLinearr   r   r   Zones_r   r   )r   namerK   Zfan_outr>   r>   r?   r"    s    



r"  r   r\   r.   )r(   r)   r*   r,   r-   r/      r   rX   i  g      ?         @rW   rG   i 
  )r@   rH   rF   rJ   0   r   r   i  i0  r]      i  )r+   rW      r   )      ?r<  r<  g      @)r)   rc   ra   r[   )r@   rD   rH   )r.   r.   r.         @)rc   )r/  r/  r/  g      @)rc   r`   )r=  r=  r=        @)r8  r8  r8  r>  r^   i   r   i   Zsilu)r@   rH   rD   rF   rJ   rM   r   )r(   r)   r*   r,   r-   r/   r4   r   T)r   )r@   rH   rD   rF   rJ   rM   r4   i   r   rE   )r@   rH   rD   rF   rM   Zgca)r@   rH   rD   rF   rM   r0   seZecaZbatr   )
block_size)r@   rH   rD   rF   rM   r0   r1   i   )r@   rH   rD   rF   rJ   rM   r0   )r(   r)   r*   r,   r/   )r@   rH   rD   rF   r0   rY   `         )Zrd_ratio)r   r   )	r@   rH   rF   rB   rJ   rM   r0   r1   r4   i   )
r@   rH   rD   rF   rB   rJ   rM   r0   r1   r4   rd   )
r@   rH   rF   rB   rJ   rM   rN   r0   r1   r4   r   )r@   rH   rD   rF   rB   rJ   rM   rN   r0   r1   r4   )!gernet_lgernet_mgernet_s	repvgg_a2	repvgg_b0	repvgg_b1repvgg_b1g4	repvgg_b2repvgg_b2g4	repvgg_b3repvgg_b3g4	resnet51q	resnet61qresnext26tsgcresnext26tsseresnext26tseca_resnext26tsbat_resnext26ts
resnet32ts
resnet33tsgcresnet33tsseresnet33tseca_resnet33tsgcresnet50tgcresnext50tsregnetz_b16regnetz_c16regnetz_d32
regnetz_d8
regnetz_e8regnetz_b16_evosregnetz_c16_evosregnetz_d8_evosc                 K   s$   t t| |ft|  tddd|S )NT)Zflatten_sequential)r	  Zfeature_cfg)r   r#   
model_cfgsrQ   )variant
pretrainedrs   r>   r>   r?   _create_byobnete  s    rh  c                 K   s   | dddddt tddd
|S )	Nr  r      rj  r   r   g      ?Zbilinear	stem.convhead.fc
urlr  
input_size	pool_sizecrop_pctinterpolationr0  r1  
first_conv
classifierr   ro  rs   r>   r>   r?   _cfgm  s    rw  c                 K   s   | dddddt tddd
|S )	Nr  r   r]   r]   r   r   g?Zbicubiczstem.conv1.convrm  rn  r   rv  r>   r>   r?   _cfgrw  s    rz  ztimm/)	hf_hub_idrx  ry  )r{  rp  rq  )zstem.conv_kxk.convzstem.conv_1x1.convZmit)r{  rt  licensezkhttps://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/resnet51q_ra2-d47dcc76.pthz
stem.conv1)r   rC  rC  )r{  ro  rt  rp  rq  test_input_sizetest_crop_pctzkhttps://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/resnet61q_ra2-6afc536c.pth)r{  ro  r}  r~  zvhttps://github.com/rwightman/pytorch-image-models/releases/download/v0.1-attn-weights/resnext26ts_256_ra2-8bbd9106.pthzthttps://github.com/rwightman/pytorch-image-models/releases/download/v0.1-attn-weights/seresnext26ts_256-6f0d74a3.pthzthttps://github.com/rwightman/pytorch-image-models/releases/download/v0.1-attn-weights/gcresnext26ts_256-e414378b.pthzvhttps://github.com/rwightman/pytorch-image-models/releases/download/v0.1-attn-weights/eca_resnext26ts_256-5a1d030f.pthzvhttps://github.com/rwightman/pytorch-image-models/releases/download/v0.1-attn-weights/bat_resnext26ts_256-fa6fd595.pth)r{  ro  Zmin_input_sizezqhttps://github.com/rwightman/pytorch-image-models/releases/download/v0.1-attn-weights/resnet32ts_256-aacf5250.pthzqhttps://github.com/rwightman/pytorch-image-models/releases/download/v0.1-attn-weights/resnet33ts_256-e91b09a4.pthzshttps://github.com/rwightman/pytorch-image-models/releases/download/v0.1-attn-weights/gcresnet33ts_256-0e0cd345.pthzshttps://github.com/rwightman/pytorch-image-models/releases/download/v0.1-attn-weights/seresnet33ts_256-f8ad44d9.pthzuhttps://github.com/rwightman/pytorch-image-models/releases/download/v0.1-attn-weights/eca_resnet33ts_256-8f98face.pthzrhttps://github.com/rwightman/pytorch-image-models/releases/download/v0.1-attn-weights/gcresnet50t_256-96374d1c.pthzthttps://github.com/rwightman/pytorch-image-models/releases/download/v0.1-attn-weights/gcresnext50ts_256-3e0f515e.pthzphttps://github.com/rwightman/pytorch-image-models/releases/download/v0.1-attn-weights/regnetz_b_raa-677d9606.pthrl  )r   r   r   ri  rk  gGz?)
r{  ro  rt  r0  r1  rp  rq  rr  r}  r~  zuhttps://github.com/rwightman/pytorch-image-models/releases/download/v0.1-attn-weights/regnetz_c_rab2_256-a54bf36a.pth)r   @  r  )r{  ro  rt  r0  r1  rr  r}  r~  zthttps://github.com/rwightman/pytorch-image-models/releases/download/v0.1-attn-weights/regnetz_d_rab_256-b8073a89.pthgffffff?)r{  ro  r0  r1  rr  r}  zphttps://github.com/rwightman/pytorch-image-models/releases/download/v0.1-attn-weights/regnetz_d8_bh-afc03c55.pth)r{  ro  r0  r1  rr  r}  r~  zphttps://github.com/rwightman/pytorch-image-models/releases/download/v0.1-attn-weights/regnetz_e8_bh-aace8e6e.pth)rt  r0  r1  rp  rq  rr  r}  zuhttps://github.com/rwightman/pytorch-image-models/releases/download/v0.1-tpu-weights/regnetz_c16_evos_ch-d8311942.pth)r{  ro  rt  r0  r1  rr  r}  zthttps://github.com/rwightman/pytorch-image-models/releases/download/v0.1-tpu-weights/regnetz_d8_evos_ch-2bc12646.pth)!zgernet_s.idstcv_in1kzgernet_m.idstcv_in1kzgernet_l.idstcv_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resnet51q.ra2_in1kzresnet61q.ra2_in1kzresnext26ts.ra2_in1kzseresnext26ts.ch_in1kzgcresnext26ts.ch_in1kzeca_resnext26ts.ch_in1kzbat_resnext26ts.ch_in1kzresnet32ts.ra2_in1kzresnet33ts.ra2_in1kzgcresnet33ts.ra2_in1kzseresnet33ts.ra2_in1kzeca_resnet33ts.ra2_in1kzgcresnet50t.ra2_in1kzgcresnext50ts.ch_in1kzregnetz_b16.ra3_in1kzregnetz_c16.ra3_in1kzregnetz_d32.ra3_in1kzregnetz_d8.ra3_in1kzregnetz_e8.ra3_in1kzregnetz_b16_evos.untrainedzregnetz_c16_evos.ch_in1kzregnetz_d8_evos.ch_in1k)rm   c                 K   s   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
    rD  rg  )rD  rh  rg  rs   r>   r>   r?   rD  	  s    rD  c                 K   s   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
    rE  rg  )rE  r  r  r>   r>   r?   rE    s    rE  c                 K   s   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
    rF  rg  )rF  r  r  r>   r>   r?   rF    s    rF  c                 K   s   t dd| i|S )z^ RepVGG-A2
    `Making VGG-style ConvNets Great Again` - https://arxiv.org/abs/2101.03697
    rG  rg  )rG  r  r  r>   r>   r?   rG  !  s    rG  c                 K   s   t dd| i|S )z^ RepVGG-B0
    `Making VGG-style ConvNets Great Again` - https://arxiv.org/abs/2101.03697
    rH  rg  )rH  r  r  r>   r>   r?   rH  )  s    rH  c                 K   s   t dd| i|S )z^ RepVGG-B1
    `Making VGG-style ConvNets Great Again` - https://arxiv.org/abs/2101.03697
    rI  rg  )rI  r  r  r>   r>   r?   rI  1  s    rI  c                 K   s   t dd| i|S )z` RepVGG-B1g4
    `Making VGG-style ConvNets Great Again` - https://arxiv.org/abs/2101.03697
    rJ  rg  )rJ  r  r  r>   r>   r?   rJ  9  s    rJ  c                 K   s   t dd| i|S )z^ RepVGG-B2
    `Making VGG-style ConvNets Great Again` - https://arxiv.org/abs/2101.03697
    rK  rg  )rK  r  r  r>   r>   r?   rK  A  s    rK  c                 K   s   t dd| i|S )z` RepVGG-B2g4
    `Making VGG-style ConvNets Great Again` - https://arxiv.org/abs/2101.03697
    rL  rg  )rL  r  r  r>   r>   r?   rL  I  s    rL  c                 K   s   t dd| i|S )z^ RepVGG-B3
    `Making VGG-style ConvNets Great Again` - https://arxiv.org/abs/2101.03697
    rM  rg  )rM  r  r  r>   r>   r?   rM  Q  s    rM  c                 K   s   t dd| i|S )z` RepVGG-B3g4
    `Making VGG-style ConvNets Great Again` - https://arxiv.org/abs/2101.03697
    rN  rg  )rN  r  r  r>   r>   r?   rN  Y  s    rN  c                 K   s   t dd| i|S )
    rO  rg  )rO  r  r  r>   r>   r?   rO  a  s    rO  c                 K   s   t dd| i|S )r  rP  rg  )rP  r  r  r>   r>   r?   rP  h  s    rP  c                 K   s   t dd| i|S )r  rQ  rg  )rQ  r  r  r>   r>   r?   rQ  o  s    rQ  c                 K   s   t dd| i|S )r  rR  rg  )rR  r  r  r>   r>   r?   rR  v  s    rR  c                 K   s   t dd| i|S )r  rS  rg  )rS  r  r  r>   r>   r?   rS  }  s    rS  c                 K   s   t dd| i|S )r  rT  rg  )rT  r  r  r>   r>   r?   rT    s    rT  c                 K   s   t dd| i|S )r  rU  rg  )rU  r  r  r>   r>   r?   rU    s    rU  c                 K   s   t dd| i|S )r  rV  rg  )rV  r  r  r>   r>   r?   rV    s    rV  c                 K   s   t dd| i|S )r  rW  rg  )rW  r  r  r>   r>   r?   rW    s    rW  c                 K   s   t dd| i|S )r  rX  rg  )rX  r  r  r>   r>   r?   rX    s    rX  c                 K   s   t dd| i|S )r  rY  rg  )rY  r  r  r>   r>   r?   rY    s    rY  c                 K   s   t dd| i|S )r  rZ  rg  )rZ  r  r  r>   r>   r?   rZ    s    rZ  c                 K   s   t dd| i|S )r  r[  rg  )r[  r  r  r>   r>   r?   r[    s    r[  c                 K   s   t dd| i|S )r  r\  rg  )r\  r  r  r>   r>   r?   r\    s    r\  c                 K   s   t dd| i|S )r  r]  rg  )r]  r  r  r>   r>   r?   r]    s    r]  c                 K   s   t dd| i|S )r  r^  rg  )r^  r  r  r>   r>   r?   r^    s    r^  c                 K   s   t dd| i|S )r  r_  rg  )r_  r  r  r>   r>   r?   r_    s    r_  c                 K   s   t dd| i|S )r  r`  rg  )r`  r  r  r>   r>   r?   r`    s    r`  c                 K   s   t dd| i|S )r  ra  rg  )ra  r  r  r>   r>   r?   ra    s    ra  c                 K   s   t dd| i|S )r  rb  rg  )rb  r  r  r>   r>   r?   rb    s    rb  c                 K   s   t dd| i|S )r  rc  rg  )rc  r  r  r>   r>   r?   rc    s    rc  c                 K   s   t dd| i|S )r  rd  rg  )rd  r  r  r>   r>   r?   rd    s    rd  )rV   rZ   r   )r   F)r   r   r   N)r+   )r   F)F)r   )r   )F)F)F)F)F)F)F)F)F)F)F)F)F)F)F)F)F)F)F)F)F)F)F)F)F)F)F)F)F)F)F)F)F)ur   r3  Zdataclassesr   r   r   	functoolsr   typingr   r   r   r	   r
   r   r   r   r  Ztorch.nnr9   Z	timm.datar   r   Ztimm.layersr   r   r   r   r   r   r   r   r   r   r   r   Z_builderr   Z_manipulater   r    	_registryr!   r"   __all__r%   r$   ri   r8   r<   rU   rv   r~   r   r   r:   r   r   r   r   r   r   r   r   rQ   r   r   r'   r  r   r&   r  r  r  r=   r  r  r#   r"  re  rh  rw  rz  Zdefault_cfgsrD  rE  rF  rG  rH  rI  rJ  rK  rL  rM  rN  rO  rP  rQ  rR  rS  rT  rU  rV  rW  rX  rY  rZ  r[  r\  r]  r^  r_  r`  ra  rb  rc  rd  r>   r>   r>   r?   <module>   sX  (8 
 		
:EBA:E
<   *)
@
y














      *





  	