khK dZddlZddlZddlmZmZddlZddlmZddl mcm Z ddl m Z ddlmZddlmZddlmZeej*j-d ddkDZ d(d eeefd eeeefd ed ej2fdZdeeefdeedffdZ d)dedefdZdeeefdeedffdZGddej>Z d*deeefdeeeffdZ!Gddej>Z" d+dedeefd Z#d!ed"ed#ej2d$ej2d ej2f d%Z$Gd&d'ej>Z%y),zf Relative position embedding modules and functions Hacked together by / Copyright 2022 Ross Wightman N)OptionalTuple)ndgrid)RegularGridInterpolator)Mlp) trunc_normal_TIMM_USE_SCIPY_INTERPq_sizek_size class_tokenreturnc |Jdtjttj|dtj|dj d}|dddddf|dddddfz }|j ddd}|dddddfxx|ddz z cc<|dddddfxx|ddz z cc<|dddddfxxd|dzdz zcc<d|dzdz d|dzdz z}|j d}|r5tj|gd}||dddf<|dz|dddf<|dz|d<|jS)Nz-Different q & k sizes not currently supportedrr)rrrrrr) torchstackrarangeflattenpermutesumFpad contiguous)r r r coordsrelative_coordsnum_relative_distancerelative_position_indexs U/var/www/teggl/fontify/venv/lib/python3.12/site-packages/timm/layers/pos_embed_rel.pygen_relative_position_indexr!sw >JJJ> [[ VAY 7fQi9PQ R Z Z[\ ]FQ4Z(6!T1*+==O%--aA6OAq!Gq A -Aq!Gq A -Aq!GF1I  11]Q.1vay=13DE0.11"5#$%%(?"N)>12&)>)BA&(=(A% " - - //new_window_sizenew_bias_shape.c l|ddzdz |ddzdz f}|jdk(r|\}}}|j\}}} ||dk(r||dk(sJ||k7s| |k7rKtjjj |j d|ddjd}|S|jdk(sJ|\} }|j\} }| |d|dzz } t| | z dz} | | f} | d|dk7s | d|dk7r| r|| dddf}|d| ddf}nd}tjjj |jddjdd | d| df|ddjd | | z jdd}|tj||fd }|S) NrrrbicubicF)sizemode align_corners?rdim) ndimshapernn functional interpolate unsqueezesqueezeint transposereshapeviewcat) rel_pos_biasr#r$dst_size_dst_hdst_wnum_attn_headssrc_hsrc_w dst_num_pos src_num_posnum_extra_tokenssrc_size extra_tokenss r resize_rel_pos_bias_table_simplerGIs#  "Q&*OA,>,BQ,FGHA(5%'3'9'9$u #!(<<< E>Ue^ 88..::&&q)# ; gaj @ 3  A%%%' Q&2&8&8# ^&(1+ *CD &663>?h' A;(1+ %! )C+-=,=,>,AB +,>.>->,>,AB #  88..::&&q!,44aXa[(ST+5VW# ; d2{%556yyA '$yy, )E1M r" interpolation antialiasc|j\}}|\}}||k(sJ||k7r|j}|j}t|dz} t|dz} t j |j ddjd|| | | | f||} | j||j dd} | j|| S|S)z Resample relative position bias table suggested in LeVit Adapted from: https://github.com/microsoft/Cream/blob/main/TinyViT/utils.py r+rr)r(r)rI) r(dtypefloatr5rr2rr8to) position_bias_tablenew_sizerHrIL1nH1L2nH2 orig_dtypeS1S2$relative_position_bias_table_resizeds r resize_rel_pos_bias_table_levitrXxs"&&(GBGB #:: Rx(.. 1779 s^ s^/0}}  ' '1 - 2 21c2r Bb 0!, 1 5 5c2 > F Fq! L -,// ;33""r"ctrddlm}|ddzdz |ddzdz f}|jdk(r1d}|\}}}||dk(r||dk(sJ|j\} } } | | f} d} nJ|jdk(sJ|\}}|j\}} ||d|dzz }t ||z dz} | | f} d} | d|dk7s | d|dk7r:|r|| d d d f}|d | d d f}nd }d fd }|| d|d}|| d|d}t j|t j|g}|dd z}|dd z}t j| |d zd}t j| |d zd}t||}g}t| D]}| r.|d d |fj| d| dj}n||d d d d fj}trgj|||jd}t j|||j!j#|j$}n;t'||}||j!j#|j$}| r|jdd}|j)|| rt j*|d}nt j*|d}|| sJt j*||fd}|S)a Resize relative position bias table using more advanced interpolation. Modified from code in Microsoft Unilm (https://github.com/microsoft/unilm) repo (BeiT, BeiT-v2, etc). https://github.com/microsoft/unilm/blob/5255d52de86dad642810f5849dd357769346c1d7/beit/run_class_finetuning.py#L351 Args: rel_pos_bias: new_window_size: new_bias_shape: Returns: r)r2rrr&Fr+TNc$|d||zz zd|z z S)N?)arns r geometric_progressionz8resize_rel_pos_bias_table..geometric_progressionsa1f %q1 1r"c(d\}}||z dkDr+||zdz } d||dz}||dzkDr|}n|}||z dkDr+g}d}t|dzD]}|j|||dzzz } t|D cgc]} | } } | dgz|zScc} w)N)g)\(?g?gư>@rrr)rangeappendreversed) srcdstleftrightqgpdiscurir<r_idsr`s r _calcz(resize_rel_pos_bias_table.._calcs#KD%$,%E\S(*1a:q=ED $,%CC3!8_ $ 3qQU|# $"*#/AaR/E/A3;$ $0s: Brbg?r[cubic)kindrr,) _USE_SCIPYscipyr2r.r/r5rtensorrrrcr8rLinterp2dnumpyTensorrrMdevicerrdr9) r:r#r$r2r;rDr<r=r>r?r@rArEhas_flat_shaperBrCrFrpyxyxtytxdydxdyxall_rel_pos_biasrnzfr^r`s @r resize_rel_pos_bias_tablersJ&%"Q&*OA,>,BQ,FGHA(5% #!(<<<'3'9'9$u5>  A%%%' Q&2&8&8# ^&(1+ *CD &663>?h'{hqk!Xa[HQK%? ')9(9(:A(=>L'(:*:):(:A(=>LL 2 %$ (1+x{ + (1+x{ +ll1ou||A /a[C  a[C  \\2#rCx - \\2#rCx -Rn~& 'A A&++HQK!EKKM Aq)//1((Aqwwyw(GLL2r+668;;L+>?FF2qM  # #A &' '*  99%52>L 99%51=L  #! !> 99lL%AqIL r"creZdZdZdfd ZdZdejfdZd de ejfdZ xZ S) RelPosBiasz_ Relative Position Bias Adapted from Swin-V1 relative position bias impl, modularized. ct||dksJ||_|d|dz|_|j|zfdz|fz|_d|dzdz d|dzdz zd|zz}t j tj|||_ |jdt|j|dkDjdd |jy) Nrrrr&r)r rF persistent)super__init__ window_size window_area bias_shaper0 Parameterrzerosrelative_position_bias_tableregister_bufferr!r8 init_weights)selfr num_heads prefix_tokensr __class__s r rzRelPosBias.__init__s !!!&&q>KN:++m;=AYLP!"[^!3a!7A A !! ! x{Q%7!EHHW x{Q%7!EHHW!KK/@BS(TU199!QBMMO v~ q !A % !!Q' */B1/E/I J * !!Q' */B1/E/I J * !!Q' *x{Q ? * !!Q' *x{Q ? *" % +@ AEJJ '++- -E/!/151!> ! !& +@ AEII '++- -E/!/ ! r"cveZdZdZ dfd Zdej fdZddeej fdZ xZ S) RelPosMlpz Log-Coordinate Relative Position MLP Based on ideas presented in Swin-V2 paper (https://arxiv.org/abs/2111.09883) This impl covers the 'swin' implementation as well as two timm specific modes ('cr', and 'rw') c:t|||_|jd|jdz|_||_||_|jfdz|fz|_|dk(r#tj|_ d|_ d}n"tj|_ d|_ d}td||tj|d |_|jd t!|j#d d |jdt%|||d y)Nrrrr)TFT)g?g)hidden_features out_features act_layerbiasdroprrFrrel_coords_log)r))rrrrrrrr0Sigmoidbias_act bias_gainIdentityrReLUmlprr!r8r) rrr hidden_dimrr)pretrained_window_sizemlp_biasrs r rzRelPosMlp.__init__Rs &++A.1A1A!1DD*"++-1YL@ 6>JJLDMDN$HKKMDM!DNH &"gg    % ' 4 9 9" =     #K1Gd S  r"rc|j|j}|jD|jd|j|j}|j|j }|j ddd}|j|}|j|j|z}|jr.tj||jd|jdg}|jdjS)Nrrrr)rrrr8rrrrrrrrr3rrs r rzRelPosMlp.get_bias~s!%$*=*=!>  ' ' 3%;%@%@T^^%TUYUqUq%r "%;%@%@%Q "!7!?!?1a!H!%/E!F >> %%)^^6L%L "   %&UU+ADDVDVXY[_[m[mopCq%r "%//2==??r"rc(||jzSrrrs r rzRelPosMlp.forwardrr")rrrrr) rrrrrrrxrrrrrs@r rrLsF#)*X @%,, @&HU\\,B&r"rlengthmax_relative_positionc||dz }d|zdz}tj|||}t|D]4}t|D]$}||z |z}t||z |kDrd||||f<&6|S)aGenerate a one_hot lookup tensor to reindex embeddings along one dimension. Args: length: the length to reindex to. max_relative_position: the maximum relative position to consider. Relative position embeddings for distances above this threshold are zeroed out. Returns: a lookup Tensor of size [length, length, vocab_size] that satisfies ret[n,m,v] = 1{m - n + max_relative_position = v}. rr)rrrcr)rr vocab_sizeretrnr|vs r generate_lookup_tensorrs$ & **Q.J ++ffj 1C 6]v AA--A1q5z11C1aL   Jr"heightwidth height_lookup width_lookupctjd||}tjd||}||z}|j|jd||S)a\Reindex 2d relative position bias with 2 independent einsum lookups. Adapted from: https://github.com/google-research/maxvit/blob/2e06a7f1f70c76e64cd3dabe5cd1b8c1a23c9fb7/maxvit/models/attention_utils.py Args: relative_position_tensor: tensor of shape [..., vocab_height, vocab_width, ...]. height: height to reindex to. width: width to reindex to. height_lookup: one-hot height lookup width_lookup: one-hot width lookup Returns: reindexed_tensor: a Tensor of shape [..., height * width, height * width, ...] z nhw,ixh->nixwznixw,jyw->nijxyr)reinsumr7r/)relative_position_tensorrrrrreindexed_tensorareas r reindex_2d_einsum_lookuprsW.||O5M}]||$57GV E>D  # #$<$B$B1$EtT RRr"creZdZdZdfd ZdZdejfdZd de ejfdZ xZ S) RelPosBiasTfz Relative Position Bias Impl (Compatible with Tensorflow MaxViT models) Adapted from: https://github.com/google-research/maxvit/blob/2e06a7f1f70c76e64cd3dabe5cd1b8c1a23c9fb7/maxvit/models/attention_utils.py ct||dksJ||_|d|dz|_||_d|dzdz }d|dzdz }|j||f|_t jtj|j |_ |jdt|dd|jdt|dd|jy)NrrrrFrr)rrrrrrr0rrrrrrr)rrrr vocab_height vocab_widthrs r rzRelPosBiasTf.__init__s !!!&&q>KN:";q>)A- +a.(1, >><E,.LLT__9U,V) _.D[QR^.Tafg ^-CKPQN-S`ef r"cZtjj|jdyr)r0initnormal_rrs r rzRelPosBiasTf.init_weightss 99sCr"rct|j|jd|jd|j|jS)Nrr)rrrrrrs r rzRelPosBiasTf.get_biassF'  - -   Q    Q          r"rc(||jzSrrrs r rzRelPosBiasTf.forwardrr"rrrrs@r rrs8 D %,, &HU\\,B&r"r)NF)r'T)rrr)&rrostypingrrrtorch.nnr0torch.nn.functionalr1rgridrr2rrr weight_initr r5environgetrsboolrxr!rGstrrXrModulerrrrrrr\r"r rs " 0&  7;