o ib @sddlZddlZddlmZmZmZmZddlZddlmZddl m Z Gdddejj Z Gdddejj ZGd d d ejj ZGd d d ejj ZGd ddejj ZGdddejj ZGdddej ZGdddej ZdejdededejfddZddZd/dejded ed!edejf d"d#Zd0d$ejd ed%ed!edejf d&d'Zd(eedefd)d*Zd(eedefd+d,Zd(eedefd-d.ZdS)1N)AnyDictListOptional)nn) functionalc s^eZdZdZddedededeffdd Zed e j fd d Z d e j d e j fddZ Z S)_ScaledEmbeddingaFMake continuous embeddings and boost learning rate Args: num_embeddings (int): number of embeddings embedding_dim (int): embedding dimensions scale (float, optional): amount to scale learning rate (Default: 10.0) smooth (bool, optional): choose to apply smoothing (Default: ``False``) $@Fnum_embeddings embedding_dimscalesmoothcstt|||_|r3tj|jjjdd}|t d|d dddf}||jjjdd<|jjj|_||_ dS)Nrdim) super__init__r Embedding embeddingtorchcumsumweightdataarangesqrtr )selfr r r r r __class__N/home/ubuntu/.local/lib/python3.10/site-packages/torchaudio/models/_hdemucs.pyr-s $ z_ScaledEmbedding.__init__returncCs|jj|jSN)rrr )rrrrr8sz_ScaledEmbedding.weightxcCs|||j}|S)zForward pass for embedding with scale. Args: x (torch.Tensor): input tensor of shape `(num_embeddings)` Returns: (Tensor): Embedding output of shape `(num_embeddings, embedding_dim)` )rr )rr"outrrrforward<s z_ScaledEmbedding.forward)r F)__name__ __module__ __qualname____doc__intfloatboolrpropertyrTensorrr$ __classcell__rrrrr#s   rcseZdZdZ         dd ed ed ed ed edededededeeee fdeffdd Z dde j dee j de j fddZ ZS) _HEncLayeratEncoder layer. This used both by the time and the frequency branch. Args: chin (int): number of input channels. chout (int): number of output channels. kernel_size (int, optional): Kernel size for encoder (Default: 8) stride (int, optional): Stride for encoder layer (Default: 4) norm_groups (int, optional): number of groups for group norm. (Default: 4) empty (bool, optional): used to make a layer with just the first conv. this is used before merging the time and freq. branches. (Default: ``False``) freq (bool, optional): boolean for whether conv layer is for frequency domain (Default: ``True``) norm_type (string, optional): Norm type, either ``group_norm `` or ``none`` (Default: ``group_norm``) context (int, optional): context size for the 1x1 conv. (Default: 0) dconv_kw (Dict[str, Any] or None, optional): dictionary of kwargs for the DConv class. (Default: ``None``) pad (bool, optional): true to pad the input. Padding is done so that the output size is always the input size / stride. (Default: ``True``) FT group_normrNchinchout kernel_sizestride norm_groupsemptyfreq norm_typecontextdconv_kwpadc st| dur i} dd} |dkrfdd} | r|dnd} tj}||_||_||_||_| |_|rD|dg}|dg}| dg} tj }|||||| |_ | ||_ |jrft |_ t |_t |_dS||d|dd| d| |_ | d||_t|fi| |_dS) NcStSr!rIdentitydrrrmz%_HEncLayer.__init__..r2c t|Sr!r GroupNormrAr7rrrCo r1rr)rrrConv1dr9r5r6r8r=Conv2dconvnorm1r@rewritenorm2dconv_DConv)rr3r4r5r6r7r8r9r:r;r<r=norm_fnpad_valklassrrHrr\s6     z_HEncLayer.__init__r"injectr c Csh|js|dkr|j\}}}}||d|}|js4|jd}||jdks4t|d|j||jf}||}|jr>|S|durk|jd|jdkrPt d|dkrg|dkrg|dddddf}||}t | |}|jr|j\}}}}| dddd d||}||}||||| dddd}n||}|||} tj| dd } | S) a]Forward pass for encoding layer. Size depends on whether frequency or time Args: x (torch.Tensor): tensor input of shape `(B, C, F, T)` for frequency and shape `(B, C, T)` for time inject (torch.Tensor, optional): on last layer, combine frequency and time branches through inject param, same shape as x (default: ``None``) Returns: Tensor output tensor after encoder layer of shape `(B, C, F / stride, T)` for frequency and shape `(B, C, ceil(T / stride))` for time r1rNzInjection shapes do not alignrJrr)r9rshapeviewr6Fr=rMr8 ValueErrorgelurNpermutereshaperQrPrOglu) rr"rVBCFrTleyzrrrr$s4    z_HEncLayer.forward) r0r1r1FTr2rNTr!r%r&r'r(r)r+strrrrrrr-r$r.rrrrr/IsF    *,r/cseZdZdZ          dd ed ed ed ed ededededededeeee fdeffdd Z de j dee j fddZ ZS) _HDecLayeraDecoder layer. This used both by the time and the frequency branches. Args: chin (int): number of input channels. chout (int): number of output channels. last (bool, optional): whether current layer is final layer (Default: ``False``) kernel_size (int, optional): Kernel size for encoder (Default: 8) stride (int): Stride for encoder layer (Default: 4) norm_groups (int, optional): number of groups for group norm. (Default: 1) empty (bool, optional): used to make a layer with just the first conv. this is used before merging the time and freq. branches. (Default: ``False``) freq (bool, optional): boolean for whether conv layer is for frequency (Default: ``True``) norm_type (str, optional): Norm type, either ``group_norm `` or ``none`` (Default: ``group_norm``) context (int, optional): context size for the 1x1 conv. (Default: 1) dconv_kw (Dict[str, Any] or None, optional): dictionary of kwargs for the DConv class. (Default: ``None``) pad (bool, optional): true to pad the input. Padding is done so that the output size is always the input size / stride. (Default: ``True``) Fr0r1rTr2Nr3r4lastr5r6r7r8r9r:r;r<r=c st| dur i} dd} | dkrfdd} | r.||ddkr'td||d} nd} | |_||_||_||_||_||_||_ t j }t j }|r[|dg}|dg}t j }t j}||||||_| ||_|jrwt |_t |_dS||d|dd| d| |_| d||_dS) NcSr>r!r?rArrrrCrDz%_HDecLayer.__init__..r2crEr!rFrArHrrrCrIrJrz#Kernel size and stride do not alignr)rrr\r=rkr9r3r8r6r5rrKConvTranspose1drLConvTranspose2dconv_trrPr@rOrN)rr3r4rkr5r6r7r8r9r:r;r<r=rSrUklass_trrrHrrs@    z_HDecLayer.__init__r"skipc Cs|jr|dkr|j\}}}|||jd|}|js-||}tj|| |dd}n |}|dur7t d| | |}|jrT|j rS|d|j |j ddf}n|d|j |j |f}|jd|krkt d|jsst|}||fS) a,Forward pass for decoding layer. Size depends on whether frequency or time Args: x (torch.Tensor): tensor input of shape `(B, C, F, T)` for frequency and shape `(B, C, T)` for time skip (torch.Tensor, optional): on first layer, separate frequency and time branches using param (default: ``None``) length (int): Size of tensor for output Returns: (Tensor, Tensor): Tensor output tensor after decoder layer of shape `(B, C, F * stride, T)` for frequency domain except last frequency layer shape is `(B, C, kernel_size, T)`. Shape is `(B, C, stride * T)` for time domain. Tensor contains the output just before final transposed convolution, which is used when the freq. and time branch separate. Otherwise, does not matter. Shape is `(B, C, F, T)` for frequency and `(B, C, T)` for time. rXrWrrNz%Skip must be none when empty is true..z'Last index of z must be equal to length)r9rrYrZr3r8r[r`rNrOr\rPrnr=rkr]) rr"rplengthrarbrdrfrgrrrr$s(  z_HDecLayer.forward) Fr0r1rFTr2rNTrhrrrrrjsL     "2rjc+seZdZdZ           d:deedededededededededededededededededed ed!ed"ef*fd#d$ Z d%d&Z d;d(d)Z d