o jm8@sLddlmZddlmZddlmZmZddlZddlm Z ddl m m Z ddl m mmZdZGdddej jZGdd d ej jZGd d d e jZeGd d d ZGdddeZedkreZedddZedddZeejeeeZeejejegdksJe dde!DZ"ee"dSdS)) dataclass)Enum)CallableOptionalNg@csBeZdZdZdddddddiffdd Zd d Zd d ZZS)KernelPredictorz7Kernel predictor for the location-variable convolutions@ LeakyReLUnegative_slopeg?c sNt||_||_||_||_||||} ||} ttjj tj ||ddddt t| di| |_ t|_|dd} tdD]<}|jtt|tjj tj |||| ddt t| di| tjj tj |||| ddt t| di| qHtjj tj || || dd|_tjj tj || || dd|_dS) a7 Args: cond_channels (int): number of channel for the conditioning sequence, conv_in_channels (int): number of channel for the input sequence, conv_out_channels (int): number of channel for the output sequence, conv_layers (int): number of layers T)paddingbiasrN)super__init__conv_in_channelsconv_out_channelsconv_kernel_size conv_layersnn Sequentialutilsparametrizations weight_normConv1dgetattr input_conv ModuleListresidual_convsrangeappendDropout kernel_conv bias_conv)self cond_channelsrrrrkpnet_hidden_channelskpnet_conv_size kpnet_dropoutkpnet_nonlinear_activation!kpnet_nonlinear_activation_paramskpnet_kernel_channelskpnet_bias_channelsr_ __class__rR/home/kuhnn/.local/lib/python3.10/site-packages/TTS/tts/layers/tortoise/vocoder.pyrsv        zKernelPredictor.__init__c Cs|j\}}}||}|jD]}||j|||}q||}||}|||j |j |j |j |}|||j |j |} || fS)zm Args: c (Tensor): the conditioning sequence (batch, cond_channels, cond_length) ) shaperr!todevicer%r& contiguousviewrrrr) r'cbatchr0 cond_length residual_convkbkernelsrrrr3forwardds,      zKernelPredictor.forwardcCs\t|jddt|jdt|j|jD]}t|ddt|ddqdS)Nrweightrr) parametrizeremove_parametrizationsrr%r&r!r'blockrrr3remove_weight_norms  z"KernelPredictor.remove_weight_norm)__name__ __module__ __qualname____doc__rr@rF __classcell__rrr1r3r sTrcsNeZdZdZgdddddddffdd Zd d Zdd dZddZZS)LVCBlockz"the location-variable convolutionsrr 皙?rrr c st||_t||_||_t||d|t|||| | d|id |_t t |tj j tj||d|||d|d|dd|_t|_|D]&} |jt t |tj j tj|||| |dd| dt |qMdS)Nr r ) r(rrrrr)r*r+r-)strideroutput_paddingr)rdilation)rrcond_hop_lengthlenrrrkernel_predictorrrr rrrConvTranspose1d convt_prer conv_blocksr#r) r' in_channelsr(rR dilations lReLU_sloperrUr)r*r+rTr1rr3rs\     zLVCBlock.__init__c Cs|j\}}}||}||\}}t|jD]V\}}||} |dd|ddddddddf} |dd|ddddf} |j| | | |jd} |t| ddd|ddft | dd|dddf}q|S)aLforward propagation of the location-variable convolutions. Args: x (Tensor): the input sequence (batch, in_channels, in_length) c (Tensor): the conditioning sequence (batch, cond_channels, cond_length) Returns: Tensor: the output sequence (batch, in_channels, in_length) N)hop_size) r4rYrW enumeraterZlocation_variable_convolutionrUtorchsigmoidtanh) r'xr9r0r[r?riconvoutputr=r>rrr3r@s ( $ zLVCBlock.forwardrcCs,|j\}}}|j\}}} } } || |ksJd|t| dd} t|| | fdd}|d|d| |}||krEt|d|fdd}|d||}|ddddddddd|f}|dd}|d| d}td ||} | jtj d } | d  d jtj d }| |} | || d } | S) uperform location-variable convolution operation on the input sequence (x) using the local convolution kernl. Time: 414 μs ± 309 ns per loop (mean ± std. dev. of 7 runs, 1000 loops each), test on NVIDIA V100. Args: x (Tensor): the input sequence (batch, in_channels, in_length). kernel (Tensor): the local convolution kernel (batch, in_channel, out_channels, kernel_size, kernel_length) bias (Tensor): the bias for the local convolution (batch, out_channels, kernel_length) dilation (int): the dilation of convolution. hop_size (int): the hop_size of the conditioning sequence. Returns: (Tensor): the output sequence after performing local convolution. (batch, out_channels, in_length). z$length of (x, kernel) is not matchedrr constantrrNzbildsk,biokl->bolsd) memory_format) r4intFpadunfold transposeraeinsumr5channels_last_3d unsqueezer7r8)r'rdkernelrrTr^r:r0 in_length out_channels kernel_size kernel_lengthrorrr3r`s( & z&LVCBlock.location_variable_convolutioncCs<|jt|jdd|jD] }t|ddqdS)NrrA)rWrFrBrCrYrZrDrrr3rFs  zLVCBlock.remove_weight_norm)rrQ) rGrHrIrJrr@r`rFrKrrr1r3rLs? %rLcsbeZdZdZddgdgddddd ffd d Zd d Zdfdd ZddZdddZZ S)UnivNetGeneratorzw UnivNet Generator Originally from https://github.com/mindslab-ai/univnet/blob/master/model/generator.py. r rM)r|rirPrrQdc stt|||_||_||_|}|}t|_d}|D]} | |}|j t ||| ||||dqtj j tj||dddd|_tt|tj j tj|dddddt|_dS)Nr)rRr\r]rUr*rreflect)r padding_mode)rrzr mel_channel noise_dim hop_lengthrr res_stackr#rLrrrrconv_prerr Tanh conv_post) r'r channel_sizer\stridesr]r*rn_mel_channelsrRr1rr3rs:    zUnivNetGenerator.__init__cCs:||}|jD] }||j|||}q||}|S)z Args: c (Tensor): the conditioning sequence of mel-spectrogram (batch, mel_channels, in_length) z (Tensor): the noise sequence (batch, noise_dim, in_length) )rrr5r6r)r'r9z res_blockrrr3r@Is     zUnivNetGenerator.forwardFcs"tt||r|dSdSN)rrzevalrF)r' inferencer1rr3rZs zUnivNetGenerator.evalcCsNt|jd|jD]}t|dkrt|dq |jD]}|qdS)NrAr)rBrCrrrV state_dictrrF)r'layerrrrr3rF`s    z#UnivNetGenerator.remove_weight_normNcCst|jd|jdfd|j}tj||fdd}|dur0t|jd|j| d|j}| ||}|ddddd|j d f}|j ddd}|S) Nr g<,Ԛ'r )dimrkr)minmax) rafullr4rr5r6catrandnrsizer@rclamp)r'r9rzeromelaudiorrr3rjs"$ "zUnivNetGenerator.inference)Fr) rGrHrIrJrr@rrFrrKrrr1r3rzs- rzc@s@eZdZUegejfed<eed<dZe eed<ddZ dS)VocType constructor model_pathNsubkeycCs|jdur ||jS|Sr)r)r' model_dictrrr3optionally_indexs  zVocType.optionally_index) rGrHrIrrModule__annotations__strrrrrrrr3rys  rc@seZdZeeddZdS)VocConfz vocoder.pthmodel_gN)rGrHrIrrzUnivnetrrrr3rsr__main__rr}rr)rri ccs|] }|jr|VqdSr) requires_gradnumel).0prrr3 sr)# dataclassesrenumrtypingrrratorch.nnrtorch.nn.functional functionalrmtorch.nn.utils.parametrizerrB MAX_WAV_VALUErrrLrzrrrGmodelrr9rprintr4ySizesum parameterspytorch_total_paramsrrrr3s4  } d