o )jD@sdZddlZddlZddlmZddlmmZddlm Z ddl m Z m Z ddl mZmZddlmZmZdZGd d d ejZGd d d ejZd dZGdddejZGdddejZGdddejZGdddejZGdddejZe jej e j!dGddde Z"dS)zLibrary to support dual-path speech separation. Authors * Cem Subakan 2020 * Mirco Ravanelli 2020 * Samuele Cornell 2020 * Mirko Bronzi 2020 * Jianyuan Zhong 2020 N)Models)MODELS TorchModel) ModelFileTasks)MossformerBlockGFSMNScaledSinuEmbedding:0yE>cs*eZdZdZdfdd ZddZZS) GlobalLayerNorma3Calculate Global Layer Normalization. Args: dim : (int or list or torch.Size) Input shape from an expected input of size. eps : float A value added to the denominator for numerical stability. elementwise_affine : bool A boolean value that when set to True, this module has learnable per-element affine parameters initialized to ones (for weights) and zeros (for biases). Example: >>> x = torch.randn(5, 10, 20) >>> GLN = GlobalLayerNorm(10, 3) >>> x_norm = GLN(x) r Tcstt|||_||_||_|jrM|dkr-tt |jd|_ tt |jd|_ |dkrKtt |jdd|_ tt |jdd|_ dSdS| dd| dddS)Nrweightbias)superr __init__dimepselementwise_affinenn ParametertorchZonesrZzerosrZregister_parameter)selfrshaperr __class__q/var/www/html/Deteccion_Ine/venv/lib/python3.10/site-packages/modelscope/models/audio/separation/m2/mossformer.pyr/s zGlobalLayerNorm.__init__cCs|dkr>> x = torch.randn(5, 10, 20) >>> CLN = CumulativeLayerNorm(10) >>> x_norm = CLN(x) Tcstt|j||dddS)Nr )rr)rr*r)rrrrrrrss  zCumulativeLayerNorm.__init__csx|dkr |dddd}t|}|dddd}|dkr:t|dd}t|}t|dd}|S)zReturns the normalized tensor. Arguments --------- x : torch.Tensor Tensor size [N, C, K, S] or [N, C, L] r rrr r)rpermute contiguousrr#r transposerr!rrrr#ws    zCumulativeLayerNorm.forward)Tr$rrrrr*csr*cCsL|dkr t||ddS|dkrt|ddS|dkr!tjd|ddSt|S) z5Just a wrapper to select the normalization type. ZglnT)rclnlnrr r)r r*rZ GroupNormZ BatchNorm1d)normrrrrr select_norms  r3cs*eZdZdZd fdd ZddZZS) EncoderaConvolutional Encoder Layer. Args: kernel_size : int Length of filters. in_channels : int Number of input channels. out_channels : int Number of output channels. Example: >>> x = torch.randn(2, 1000) >>> encoder = Encoder(kernel_size=4, out_channels=64) >>> h = encoder(x) >>> h.shape torch.Size([2, 64, 499]) r@rcs4tt|tj||||dddd|_||_dS)NrrF) in_channels out_channels kernel_sizestridegroupsr)rr4rrConv1dconv1dr6)rr8r7r6rrrrs zEncoder.__init__cCs0|jdkr tj|dd}||}t|}|S)aReturn the encoded output. Args: x : torch.Tensor Input tensor with dimensionality [B, L]. Returns: x : torch.Tensor Encoded tensor with dimensionality [B, N, T_out]. where B = Batchsize L = Number of timepoints N = Number of filters T_out = Number of timepoints at the output of the encoder rr)r6r unsqueezer<FZrelur.rrrr#s   zEncoder.forward)rr5rr$rrrrr4s r4cs,eZdZdZfddZfddZZS)DecoderaA decoder layer that consists of ConvTranspose1d. Args: kernel_size : int Length of filters. in_channels : int Number of input channels. out_channels : int Number of output channels. Example: --------- >>> x = torch.randn(2, 100, 1000) >>> decoder = Decoder(kernel_size=4, in_channels=100, out_channels=1) >>> h = decoder(x) >>> h.shape torch.Size([2, 1003]) cstt|j|i|dS)N)rr@r)rargskwargsrrrrszDecoder.__init__csr|dvrtd|jt|dkr|nt|d}t|dkr2tj|dd}|St|}|S)aReturn the decoded output. Args: x : torch.Tensor Input tensor with dimensionality [B, N, L]. where, B = Batchsize, N = number of filters L = time points )rr z{} accept 3/4D tensor as inputr rr=) r RuntimeErrorformatr%rr#rr>Zsqueezer.rrrr#s $ zDecoder.forwardr$rrrrr@s r@cs6eZdZdZ      d fdd Zd d ZZS) MossFormerMa9This class implements the transformer encoder. Args: num_blocks : int Number of mossformer blocks to include. d_model : int The dimension of the input embedding. attn_dropout : float Dropout for the self-attention (Optional). group_size: int the chunk size query_key_dim: int the attention vector dimension expansion_factor: int the expansion factor for the linear projection in conv module causal: bool true for causal / false for non causal Example: ------- >>> import torch >>> x = torch.rand((8, 60, 512)) >>> net = TransformerEncoder_MossFormerM(num_blocks=8, d_model=512) >>> output, _ = net(x) >>> output.shape torch.Size([8, 60, 512]) NF@皙?c s6tt|||||||d|_tj|dd|_dS)N)rdepth group_size query_key_dimexpansion_factorcausal attn_dropoutgư>r1)rrr mossformerMr LayerNormr2)r num_blocksd_modelrNrKrLrMrOrrrr"s zMossFormerM.__init__cCs||}||}|S)a Args: src : torch.Tensor Tensor shape [B, L, N], where, B = Batchsize, L = time points N = number of filters The sequence to the encoder layer (required). )rPr2)rsrcoutputrrrr#6s zMossFormerM.forward)NFrFrGrHrIr$rrrrrEsrEcs.eZdZdZ  dfdd ZddZZS) ComputationBlocka&Computation block for dual-path processing. Args: num_blocks : int Number of mossformer blocks to include. out_channels : int Dimensionality of inter/intra model. norm : str Normalization type. skip_around_intra : bool Skip connection around the intra layer. Example: --------- >>> comp_block = ComputationBlock(64) >>> x = torch.randn(10, 64, 100) >>> x = comp_block(x) >>> x.shape torch.Size([10, 64, 100]) r0TcsFtt|t||d|_||_||_|dur!t||d|_dSdS)N)rRrSr ) rrVrrE intra_mdlskip_around_intrar2r3 intra_norm)rrRr7r2rXrrrr\szComputationBlock.__init__cCsd|j\}}}|ddd}||}|ddd}|jdur'||}|jr.||}|}|S)adReturns the output tensor. Args: x : torch.Tensor Input tensor of dimension [B, N, S]. Returns: out: torch.Tensor Output tensor of dimension [B, N, S]. where, B = Batchsize, N = number of filters S = sequence time index rrrN)rr+r,rWr2rYrX)rr!BNSZintraoutrrrr#os    zComputationBlock.forward)r0Tr$rrrrrVFs rVcs6eZdZdZ      d fdd Zd d ZZS) MossFormerMaskNetaThe dual path model which is the basis for dualpathrnn, sepformer, dptnet. Args: in_channels : int Number of channels at the output of the encoder. out_channels : int Number of channels that would be inputted to the intra and inter blocks. norm : str Normalization type. num_spks : int Number of sources (speakers). skip_around_intra : bool Skip connection around intra. use_global_pos_enc : bool Global positional encodings. max_length : int Maximum sequence length. Example: --------- >>> mossformer_block = MossFormerM(1, 64, 8) >>> mossformer_masknet = MossFormerMaskNet(64, 64, intra_block, num_spks=2) >>> x = torch.randn(10, 64, 2000) >>> x = mossformer_masknet(x) >>> x.shape torch.Size([2, 10, 64, 2000]) r0rT Nc stt|||_||_t||d|_tj||ddd|_ ||_ |j r)t ||_ t ||||d|_tj|||dd|_tj||ddd|_t|_t|_tt||dt|_tt||dt|_dS)Nr rF)r)rX)r8)rr^rnum_spksrRr3r2rr;conv1d_encoderuse_global_pos_encr pos_encrVmdl conv1d_out conv1_decoderZPReLUpreluZReLU activationZ SequentialZTanhrUZSigmoid output_gate) rr6r7rRr2rarXrc max_lengthrrrrs<      zMossFormerMaskNet.__init__c Cs||}||}|jr$|}|dd}||}|dd}||}||}||}||}|j\}}}| ||j d|}| || |}| |}|j\}}}| ||j ||}||}|dd}|S)aReturns the output tensor. Args: x : torch.Tensor Input tensor of dimension [B, N, S]. Returns: out : torch.Tensor Output tensor of dimension [spks, B, N, S] where, spks = Number of speakers B = Batchsize, N = number of filters S = the number of time frames rr)r2rbrcr-rdrerhrfrviewrarUrjrgri) rr!baseZembrZ_r\r[Lrrrr#s(             zMossFormerMaskNet.forward)r_r0rTTr`r$rrrrr^s )r^) module_namecsLeZdZdZ         dd effd d Zd d ZdddZZS) MossFormer2ziLibrary to support MossFormer speech separation. Args: model_dir (str): the model path. r_r0rTr` model_dirc sftj|g| Ri| ||_t||dd|_t||||||| | d|_t|d||ddd|_dS)Nr)r8r7r6)r6r7rRr2rarXrcrkrF)r6r7r8r9r) rrrar4encr^mask_netr@dec) rrur6r7rRr8r2rarXrcrkrArBrrrrs,   zMossFormer2.__init__cs|}|}t|gj}||tjfddtjDdd}|d}|d}||krEt |ddd||f}|S|ddd|ddf}|S)Ncs g|] }|dqS)rl)rxr>).0irZsep_xrr Cs z'MossFormer2.forward..rlr=rr) rvrwrstackracatrangesizer?pad)rinputr!maskZ est_sourceZT_originZT_estrr{rr#;s    zMossFormer2.forwardNcCs@|s|j}|s td}|jtjtj|tj |ddddS)Ncpu)Z map_locationF)strict) rurdeviceZload_state_dictloadospathjoinrZTORCH_MODEL_FILE)rZ load_pathrrrrload_check_pointNs  zMossFormer2.load_check_point) rsrsr_rtr0rTTr`)NN) r%r&r'r(strrr#rr)rrrrrrs "rr)#r(rrZtorch.nnrZtorch.nn.functionalZ functionalr?Zmodelscope.metainforZmodelscope.modelsrrZmodelscope.utils.constantrrZmossformer_blockrr ZEPSModuler rQr*r3r4ZConvTranspose1dr@rErVr^Zregister_moduleZspeech_separationZ)speech_mossformer2_separation_temporal_8krrrrrrs.   G-8/AK~