o ϯi6@sddlZddlmZddlmmZddlmZddlZ ddl Z ddl Z ddl m Z ddlmmZddlZddlmZmZddlmZddlmZddlmZmZddlmZmZmZd d ZedZed Zed Z ed Z!eZ"d8de#de$fddZ%Gdddej&Z'Gdddej&Z(Gdddej&Z)ddZ*d9dd Z+d:d#d$Z,d%d&Z-d'd(Z.d)d*Z/Gd+d,d,ej&Z0Gd-d.d.ej&Z1Gd/d0d0ej&Z2Gd1d2d2ej&Z3Gd3d4d4ej&Z4d;d6d7Z5dS)<N)repeat)_calculate_fan_in_and_fan_out) SpectrogramLogmelFilterBank)SpecAugmentation)do_mixup interpolate)iAFFAFFDAFcsfdd}|S)Ncs t|tjjr |Stt|SN) isinstance collectionsabcIterabletuplerxnP/home/ubuntu/.local/lib/python3.10/site-packages/laion_clap/clap_module/htsat.pyparsesz_ntuple..parser)rrrrr_ntuples rF drop_probtrainingcCsd|dks|s|Sd|}|jdfd|jd}|tj||j|jd}||||}|S)a&Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks). This is the same as the DropConnect impl I created for EfficientNet, etc networks, however, the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper... See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for changing the layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use 'survival rate' as the argument. rrrr)dtypedevice)shapendimtorchrandr"r#floor_div)rrr keep_probr$ random_tensoroutputrrr drop_path+s r-cs*eZdZdZdfdd ZddZZS)DropPathz^Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks). Ncstt|||_dSr )superr.__init__r)selfr __class__rrr0@s zDropPath.__init__cCst||j|jSr )r-rr r1rrrrforwardDszDropPath.forwardr __name__ __module__ __qualname____doc__r0r5 __classcell__rrr2rr.=sr.cs0eZdZdZ  dfd d Zdd d ZZS) PatchEmbedz! 2D Image to Patch Embedding rNTFNonec stt|}t|}t|}||_||_||_|d|d|d|df|_|jd|jd|_||_||_ ||_ ||_ | |_ |d|dd|d|ddf} |j ro|j dkrot j|d|||| d|_n t j||||| d|_|r||nt |_|j r|j dvrt j|||d|ddf|d|ddf| d|_|j d krt|_dS|j d krt|d d |_dS|j d krt|d d |_dSdSdSdS)Nrrr channel_mapr kernel_sizestridepaddingdaf_2daff_2diaff_2drrGrH2DchannelstyperI)r/r0 to_2tupleimg_size patch_size patch_stride grid_size num_patchesflattenin_chans embed_dim enable_fusion fusion_typennConv2dprojIdentitynorm mel_conv2dr fusion_modelr r ) r1rOrPrUrV norm_layerrTrQrWrXrEr2rrr0Js: "(6    zPatchEmbed.__init__c CsB|jr|jdvr|ddddddddf}|j\}}}}||jdkr-||jdksDJd|d|d|jdd|jdd ||}|d}t|dkr||ddddddf} | j\}}}}| ||d||} | | } | ||| d| d | d } | d  d } | \} } } } | d|krt j | t j| | | || df|jd gdd } n| ddddddd|f} |||| ||<|}n2|j\}}}}||jdkr||jdksJd|d|d|jdd|jdd ||}|j r| d dd }||}|S)NrFrrzInput image size (*z) doesn't match model (z).rr)rrrrrr#dim)rWrXr$rOr[sizelen contiguousviewr^permuterTr&catzerosr#r_ transposer])r1r longer_idxglobal_xBCHWTWlocal_xTBTCTH_rrrr5ls< (     $2  (  zPatchEmbed.forward) r=r>rr?NTr>Fr@r r6rrr2rr<Gs "r<cs4eZdZdZddejdffdd ZddZZS)MlpzG MLP as used in Vision Transformer, MLP-Mixer and related networks NrcsNt|p|}|p |}t|||_||_t|||_t||_dSr ) r/r0rYLinearfc1actfc2Dropoutdrop)r1 in_featureshidden_features out_features act_layerrr2rrr0s z Mlp.__init__cCs6||}||}||}||}||}|Sr )r|r}rr~r4rrrr5s     z Mlp.forward) r7r8r9r:rYGELUr0r5r;rrr2rrzs rzcCsdd}||d|ks||d|krtjdddtB||||}||||}|d|dd|d|||td| ||j ||d|WdS1sdwYdS) NcSsdt|tddS)N?@)matherfsqrtrrrrnorm_cdfsz(_no_grad_trunc_normal_..norm_cdfrzjmean is more than 2 std from [a, b] in nn.init.trunc_normal_. The distribution of values may be incorrect.) stacklevelrr)minmax) warningswarnr&no_graduniform_erfinv_mul_rradd_clamp_)tensormeanstdabrlurrr_no_grad_trunc_normal_s   $rrrcCst|||||S)aFills the input Tensor with values drawn from a truncated normal distribution. The values are effectively drawn from the normal distribution :math:`\mathcal{N}(\text{mean}, \text{std}^2)` with values outside :math:`[a, b]` redrawn until they are within the bounds. The method used for generating the random values works best when :math:`a \leq \text{mean} \leq b`. Args: tensor: an n-dimensional `torch.Tensor` mean: the mean of the normal distribution std: the standard deviation of the normal distribution a: the minimum cutoff value b: the maximum cutoff value Examples: >>> w = torch.empty(3, 5) >>> nn.init.trunc_normal_(w) )r)rrrrrrrr trunc_normal_srfan_innormalc Cst|\}}|dkr |}n|dkr|}n |dkr||d}||}|dkr3t|t|dddS|dkrB|jt|ddS|d krVtd |}|| |dStd |) Nrfan_outfan_avgrtruncated_normalg۶%?rruniformrzinvalid distribution )rrrrnormal_r ValueError) rscalemode distributionrrdenomvarianceboundrrrvariance_scaling_s  rcCst|ddddS)Nrr)rr)r)rrrr lecun_normal_srcCsR|j\}}}}|||||||||}|ddddddd|||}|S)z Args: x: (B, H, W, C) window_size (int): window size Returns: windows: (num_windows*B, window_size, window_size, C) rrrrrrb)r$rirjrh)r window_sizerprrrsrqwindowsrrrwindow_partitions$rcCsbt|jd||||}||||||||d}|dddddd|||d}|S)z Args: windows: (num_windows*B, window_size, window_size, C) window_size (int): Window size H (int): Height of image W (int): Width of image Returns: x: (B, H, W, C) rrbrrrrr)intr$rirjrh)rrrrrsrprrrrwindow_reverses $rcs4eZdZdZd fdd Zd ddZd d ZZS) WindowAttentiona Window based multi-head self attention (W-MSA) module with relative position bias. It supports both of shifted and non-shifted window. Args: dim (int): Number of input channels. window_size (tuple[int]): The height and width of the window. num_heads (int): Number of attention heads. qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0 proj_drop (float, optional): Dropout ratio of output. Default: 0.0 TNrcst||_||_||_||}|p|d|_tt d|ddd|dd||_ t |jd} t |jd} t t | | g} t| d} | dddddf| dddddf} | ddd} | dddddf|jdd7<| dddddf|jdd7<| dddddfd|jdd9<| d}|d|tj||d|d|_t||_t|||_t||_t|j d d tjdd |_dS) Ngrrrrbrelative_position_indexrbias{Gz?rrd)r/r0rer num_headsrrY Parameterr&rlrelative_position_bias_tablearangestackmeshgridrTrjrhsumregister_bufferr{qkvr attn_dropr[ proj_droprSoftmaxsoftmax)r1rerrqkv_biasqk_scalerrhead_dimcoords_hcoords_wcoordscoords_flattenrelative_coordsrr2rrr0$s4 & ,((,    zWindowAttention.__init__c Csn|j\}}}||||d|j||jddddd}|d|d|d}}} ||j}||dd} |j|j d |j d|j d|j d|j dd} | ddd } | | d} |dur|jd} | || | |j||| d d} | d|j||} | | } n| | } || } | | dd|||}||}||}|| fS) z Args: x: input features with shape of (num_windows*B, N, C) mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None rrrrrrbN)r$rreshaperrjrrmrrrirrh unsqueezerrr[r) r1rmaskB_Nrqrqkvattnrelative_position_biasnWrrrr5Fs* . & (     zWindowAttention.forwardcCd|jd|jd|jS)Ndim=, window_size= , num_heads=)rerrr1rrr extra_reprgzWindowAttention.extra_repr)TNrrr )r7r8r9r:r0r5rr;rrr2rrs  "!rc sLeZdZdZddddddddejejdf fd d Zd d Zd dZ Z S)SwinTransformerBlocka Swin Transformer Block. Args: dim (int): Number of input channels. input_resolution (tuple[int]): Input resulotion. num_heads (int): Number of attention heads. window_size (int): Window size. shift_size (int): Shift size for SW-MSA. mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set. drop (float, optional): Dropout rate. Default: 0.0 attn_drop (float, optional): Attention dropout rate. Default: 0.0 drop_path (float, optional): Stochastic depth rate. Default: 0.0 act_layer (nn.Module, optional): Activation layer. Default: nn.GELU norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm r@TNrlnc s:t|_||_||_||_||_||_||_t |j|jkr+d|_t |j|_d|jkr:|jks?JdJd| |_ t t |j|||| | d|_ | dkr[t| nt|_|jdkrlt|_n|jdkryfdd|_ntt|}t|| | d |_|jdkr|j\}}td ||d f}td|j t|j |j t|j df}td|j t|j |j t|j df}d}|D]}|D]}||dd||ddf<|d 7}qqt||j}|d |j|j}|d |d }||dktd |dktd}nd}| d|dS)Nrz shift_size must in 0-window_size)rrrrrrrrbncst|ddddS)Nrr)rY BatchNorm1drmrrdrrsz/SwinTransformerBlock.__init__..)rrrrrrbrgY attn_mask)!r/r0reinput_resolutionrr shift_size mlp_rationorm_before_mlprnorm1rrNrr.rYr\r- LayerNormnorm2NotImplementedErrorrrzmlpr&rlslicerrir masked_fillfloatr)r1rerrrrrrrrrr-rr`rmlp_hidden_dimrrrsimg_maskh_slicesw_slicescnthw mask_windowsrr2rdrr0~s`  (             &zSwinTransformerBlock.__init__c Cs$|j\}}|j\}}}|}||}|||||}|jdkr.tj||j |j fdd}n|}t||j} | d|j|j|} |j | |j d\} } | d|j|j|} t | |j||}|jdkrotj||j|jfdd}n|}|||||}|| |}|| | ||}|| fS)Nr)rr)shiftsdimsrb)r)rr$rrirr&rollrrrrrr-rr) r1rrrrsrpLrqshortcut shifted_x x_windows attn_windowsrrrrr5s(      zSwinTransformerBlock.forwardc Cs4d|jd|jd|jd|jd|jd|j S)Nr, input_resolution=rrz , shift_size=z , mlp_ratio=)rerrrrrrrrrrszSwinTransformerBlock.extra_repr) r7r8r9r:rYrrr0r5rr;rrr2rrls  9+rcs6eZdZdZejffdd ZddZddZZ S) PatchMergingz Patch Merging Layer. Args: input_resolution (tuple[int]): Resolution of input feature. dim (int): Number of input channels. norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm csBt||_||_tjd|d|dd|_|d||_dS)NrrFr)r/r0rrerYr{ reductionr])r1rrer`r2rrr0s zPatchMerging.__init__c Cs6|j\}}|j\}}}|||ksJd|ddkr!|ddks,Jd|d|d|||||}|ddddddddddf}|ddddddddddf}|ddddddddddf} |ddddddddddf} t||| | gd }||d d |}||}||}|S) z x: B, H*W, C zinput feature has wrong sizerrzx size (raz) are not even.Nrrbr)rr$rir&rkr]r ) r1rrrrsrprrqx0x1x2x3rrrr5s  .$$$$  zPatchMerging.forwardcCsd|jd|jS)Nzinput_resolution=z, dim=)rrerrrrrszPatchMerging.extra_repr r7r8r9r:rYrr0r5rr;rrr2rr s r c sHeZdZdZddddddejdddf fdd Zd d Zd d ZZ S) BasicLayera. A basic Swin Transformer layer for one stage. Args: dim (int): Number of input channels. input_resolution (tuple[int]): Input resolution. depth (int): Number of blocks. num_heads (int): Number of attention heads. window_size (int): Local window size. mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set. drop (float, optional): Dropout rate. Default: 0.0 attn_drop (float, optional): Attention dropout rate. Default: 0.0 drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0 norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False. rTNrFrcs|t|_|_||_||_t f ddt|D|_ | dur9| d|_ dSd|_ dS)NcsTg|]&}t |ddkrdn d ttr"|nd qS)rr) rerrrrrrrrrr-r`r)rrlist).0i rrerr-rrrr`rrrrrr 1s z'BasicLayer.__init__..)rer`) r/r0rerdepthuse_checkpointrY ModuleListrangeblocks downsample)r1rerrrrrrrrrr-r`rrrr2rrr0%s $   zBasicLayer.__init__cCsg}|jD]}|jrt||}q||\}}|js"||dq|jdur-||}|js>tj|dd}tj |dd}||fS)Nrrd) rr checkpointr appendrrr&rkr)r1rattnsblkrrrrr5Bs    zBasicLayer.forwardcCr)Nrr z, depth=)rerrrrrrrRrzBasicLayer.extra_reprrrrr2rrs  rcseZdZdZddddddgdgd d d d d dddejdd ddd ddffdd ZddZej j ddZ ej j ddZ d&ddZ d&ddZdd Zd!d"Zd'd#ejfd$d%ZZS)(HTSAT_Swin_Transformera*HTSAT based on the Swin Transformer Args: spec_size (int | tuple(int)): Input Spectrogram size. Default 256 patch_size (int | tuple(int)): Patch size. Default: 4 path_stride (iot | tuple(int)): Patch Stride for Frequency and Time Axis. Default: 4 in_chans (int): Number of input image channels. Default: 1 (mono) num_classes (int): Number of classes for classification head. Default: 527 embed_dim (int): Patch embedding dimension. Default: 96 depths (tuple(int)): Depth of each HTSAT-Swin Transformer layer. num_heads (tuple(int)): Number of attention heads in different layers. window_size (int): Window size. Default: 8 mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4 qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True qk_scale (float): Override default qk scale of head_dim ** -0.5 if set. Default: None drop_rate (float): Dropout rate. Default: 0 attn_drop_rate (float): Attention dropout rate. Default: 0 drop_path_rate (float): Stochastic depth rate. Default: 0.1 norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm. ape (bool): If True, add absolute position embedding to the patch embedding. Default: False patch_norm (bool): If True, add normalization after patch embedding. Default: True use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False config (module): The configuration Module from config.py rrrri`rrrrr> r(rTNrg?Frr@c% stt|||_||_||_||_| |_||_||_ ||_ ||_ ||_ ||_ t|j |_t|jd|jd|_| |_||_||_| |_d|_||_|jrR|nd|_||_| |_||_||_||_|j|jj|_d}d}d}d}d}d}d|_ t!|j|j"|j|||dd |_#t$|j%|j|j|j&|j'|||dd |_(t)d dd dd |_*t+,|jj|_-t.|j|j|j |j|j||j|jd|_/|j/j0}|j/j1} | |_2|j rt+3t45d||j|_6t7|j6ddt+j8|jd|_9ddt4:d|jt;|j D}!t+<|_=t>|jD]]}"t?t|jd|"| dd|"| dd|"f|j |"|j |"|j|j|j|j|j|j|!t;|j d|"t;|j d|"d|j|"|jdkrTt@nd||jd}#|j=A|#q||j|_Bt+Cd|_Dt+Ed|_F|jdt|j d|jd|j}$t+jG|j|j |$dfdd|_Ht+I|||_J|jr|jdvrt+Kt+jLd d ddddt+Md |_N|jdkrtO|_Pn|jdkrtQd dd|_Pn |jd krtRd dd|_P|S|jTdS)!NrrhannTreflectrg|=r))n_fft hop_length win_lengthwindowcenterpad_modefreeze_parameters) srr,n_melsfminfmaxrefamintop_dbr2@r()time_drop_widthtime_stripes_numfreq_drop_widthfreq_stripes_num)rOrPrUrVr`rQrWrXrr)pcSsg|]}|qSr)item)rrrrrrsz3HTSAT_Swin_Transformer.__init__..r)rerrrrrrrrrr-r`rrrr)rr) in_channels out_channelsrCrEdaf_1daff_1diaff_1drrBrDrE1DrKrF)Ur/r!r0config spec_sizerQrPrrVdepthsaperU num_classesrrg num_layersr num_features drop_rateattn_drop_ratedrop_path_raterr patch_normr`rrrrWrXmel_bins freq_ratiointerpolate_ratiorhop_sizespectrogram_extractorr sample_rater5r6logmel_extractorrspec_augmenterrY BatchNorm2dbn0r< patch_embedrSrRpatches_resolutionrr&rlabsolute_pos_embedrrpos_droplinspacerrlayersrrr rr]AdaptiveAvgPool1davgpoolAdaptiveMaxPool1dmaxpoolrZ tscam_convr{head SequentialConv1dr mel_conv1dr r_r r apply _init_weights)%r1rIrPrQrUrLrVrJrrrrrrOrPrQr`rKrRrrrHrWrXkwargsr/r0r1r7r8r9rSr^dpri_layerlayerSFr2rrr0ps         *  (    zHTSAT_Swin_Transformer.__init__cCst|tjr&t|jddt|tjr"|jdur$tj|jddSdSdSt|tjr>tj|jdtj|jddSdS)Nrrrr) rrYr{rweightrinit constant_r)r1mrrrrms  z$HTSAT_Swin_Transformer._init_weightscCdhS)Nr_rrrrrno_weight_decayz&HTSAT_Swin_Transformer.no_weight_decaycCrw)Nrrrrrrno_weight_decay_keywordsryz/HTSAT_Swin_Transformer.no_weight_decay_keywordscCs|jd}|j||d}|jr||j}||}t|jD] \}}||\}}q||}|j\}}} |dt|j d|j d} |dt|j d|j d} | ddd || | | }|j\}} } } | |j}| || | ||| }| ddddd || |d}tj|dd}t| ddd d |j d}|t|d}t|d}||}t|d}tt| ddd d |j d}||}t|d}|t|||d }|S) Nrrnrrrrrbrdr()framewise_outputclipwise_outputfine_grained_embedding embedding)r$r]rKr_r` enumeraterbr]rgrJrQrjrhrrTr&rr rdrTrgsigmoid)r1rrn frames_numrrqrrprrqrrSTFT c_freq_binfine_grained_latent_output latent_outputfpx output_dictrrrforward_featuress>        ""   (  z'HTSAT_Swin_Transformer.forward_featurescCs|jd}t|jd|jd||jd|j}tt|D]%}|dur1td||d}n|}||d|||ddf||d<q |S)Nrrrr) r$r&rltor#rrgrandomrandint)r1r crop_sizespe_pos time_stepstxrcrop_posrrrcrop_wav5s *&zHTSAT_Swin_Transformer.crop_wavcCs |j\}}}}t|j|j}|j|j}||kr||ks!Jd||kr4tjj|||jdfddd}||krGtjj||jd|fddd}|dddd}| |jd|jd|jd|j|jd|j}|ddddd }| |jd|jd|jd|jd|jd }|S) N=the wav size should less than or equal to the swin input sizerbicubicTr align_cornersrrrr) r$rrIrTrY functionalr rjrhr)r1rrprqrrtarget_Ttarget_Frrrreshape_wav2imgAs 22z&HTSAT_Swin_Transformer.reshape_wav2imgc Cs|j\}}}}t|j|j}|j|j}||kr||ks!Jd||kr4tjj|||jdfddd}||krGtjj||jd|fddd}|dddd}|dddddd|||jf}|j d d }|S) NrrrTrrrr)rrrr)repeats) r$rrIrTrYrr rjrhr) r1rcur_posrprqrrrrrrrrepeat_wat2imgSs & z%HTSAT_Swin_Transformer.repeat_wat2imgrc CsF|jrG|ddkrG|jrd|dtd|djdd<n(|dj|dd}|dd}||}|dd}| |}|j |gd }|S|js|d j|dd}| |}| |}|dd}||}|dd}|jrv| |}|jr|durt||}| |}| |}|S|dj|dd}|dj|dd}|dd}||}|dd}t|d}|jd vrv|ddddddddf}t|dkrs||ddddddf} | \} } } } | | | | | } t| d } || } | | | | | d } t| dd} | d | kr>tj| tj| | | | d f|dgd d} n | ddddd| f} |dd }|||| ||<|d dddddddf}n |}n|jdvr~|}|jr| |}|jr|durt||}| |}|j ||d }|S)NlongerrTr! mel_fusion)r# non_blockingrrr{waveformrC)rrrrb)rrrrrrcrd)rGrHrIrA)rWrr r&rr$rrmr\rrrWrYrZrwhererXclonerhrgrfrirjrkrTrkrlsqueezer_)r1r mixup_lambda infer_moder#r longer_listlonger_list_idxnew_xfusion_x_localFBFCFTFFrrrr5bsp"             T    ($ .(    -zHTSAT_Swin_Transformer.forwardr )NFN)r7r8r9r:rYrr0rmr&jitignorerxrzrrrrTensorr5r;rrr2rr!Ws.    / r!r@c Csz]|jdvs Jd|jdkr&tddd|jdgdgd d |||d }|WS|jd krBtddd|jd gdgd d |||d }|WS|jdkr[tddd|jdgdgd d |||d }|WStd|jd)N)tinybaselargezmodel name for HTS-AT is wrong!rr"rr#r$r%r'r() rIrPrQrLrVrJrrrHrWrXr)rr rrzImport Model for z7 not found, or the audio cfg parameters are not enough.) model_namer! class_num RuntimeError) audio_cfgrWrXmodelrrrcreate_htsat_modelsb *  r)rF)rrrr)rrr)Fr@)6r&torch.nnrYtorch.nn.functionalrr itertoolsrcollections.abcrrr torch.nn.initrtorch.utils.checkpointutilsrrtorchlibrosa.stftrrtorchlibrosa.augmentationrrr feature_fusionr r r r to_1tuplerN to_3tuple to_4tuple to_ntuplerboolr-Moduler.r<rzrrrrrrrrr rr!rrrrrsP      K $ U|*E