o Zhe@sdZddlmZddlmZmZmZddlZddlZddlm Z ddl m Z ddl m Z dd lmZmZd d lmZeeZeGd d d e ZGddde jZGddde jZGddde jZGddde jZeGddde ZdgZdS)zPyTorch UnivNetModel model.) dataclass)OptionalTupleUnionN)nn) ModelOutput)PreTrainedModel)auto_docstringlogging) UnivNetConfigc@s6eZdZUdZdZeejed<dZ eejed<dS)UnivNetModelOutputa Output class for the [`UnivNetModel`], which includes the generated audio waveforms and the original unpadded lengths of those waveforms (so that the padding can be removed by [`UnivNetModel.batch_decode`]). Args: waveforms (`torch.FloatTensor` of shape `(batch_size, sequence_length)`): Batched 1D (mono-channel) output audio waveforms. waveform_lengths (`torch.FloatTensor` of shape `(batch_size,)`): The batched length in samples of each unpadded waveform in `waveforms`. N waveformswaveform_lengths) __name__ __module__ __qualname____doc__rrtorch FloatTensor__annotations__rrr[/var/www/auris/lib/python3.10/site-packages/transformers/models/univnet/modeling_univnet.pyr s  rcsFeZdZdZdeffdd ZdejfddZdd Z d d Z Z S) #UnivNetKernelPredictorResidualBlockz Implementation of the residual block for the kernel predictor network inside each location variable convolution block (LVCBlock). Parameters: config: (`UnivNetConfig`): Config for the `UnivNetModel` model. configcst|j|_|j|_|j|_|j|_|jdd}t |j|_ t j |j|j|j|dd|_ t j |j|j|j|dd|_dS)Nr Tpaddingbias)super__init__model_in_channelsZchannelskernel_predictor_conv_size kernel_sizeZkernel_predictor_dropoutZ dropout_probleaky_relu_sloperZDropoutdropoutConv1dconv1conv2)selfrr __class__rrr!;s  z,UnivNetKernelPredictorResidualBlock.__init__ hidden_statescCsJ|}||}||}tj||j}||}tj||j}||SN)r&r(r functional leaky_relur%r))r*r-residualrrrforwardKs   z+UnivNetKernelPredictorResidualBlock.forwardcCs8tjj}ttjjdrtjjj}||j||jdSN weight_norm)rutilsr4hasattrparametrizationsr(r)r*r4rrrapply_weight_normUs   z5UnivNetKernelPredictorResidualBlock.apply_weight_normcCs tj|jtj|jdSr.)rr5remove_weight_normr(r)r*rrrr:]sz6UnivNetKernelPredictorResidualBlock.remove_weight_norm) rrrrr r!rrr2r9r: __classcell__rrr+rr1s  rcsTeZdZdZ  ddededeffdd Zd ejfd d Z d d Z ddZ Z S)UnivNetKernelPredictora  Implementation of the kernel predictor network which supplies the kernel and bias for the location variable convolutional layers (LVCs) in each UnivNet LVCBlock. Based on the KernelPredictor implementation in [maum-ai/univnet](https://github.com/maum-ai/univnet/blob/9bb2b54838bb6d7ce767131cc7b8b61198bc7558/model/lvcnet.py#L7). Parameters: config: (`UnivNetConfig`): Config for the `UnivNetModel` model. conv_kernel_size (`int`, *optional*, defaults to 3): The kernel size for the location variable convolutional layer kernels (convolutional weight tensor). conv_layers (`int`, *optional*, defaults to 4): The number of location variable convolutional layers to output kernels and biases for. rrconv_kernel_size conv_layerscstj|_dj|_||_||_|j|j|j|j|_|j|j|_j |_ j |_ j |_j|_j|_|jdd}tj|j |j dddd|_tfddt|jD|_tj|j |j|j|dd|_tj|j |j|j|dd|_dS)Nrr Trcsg|]}tqSr)r).0_rrr sz3UnivNetKernelPredictor.__init__..)r r!model_hidden_channelsconv_in_channelsconv_out_channelsr?r@Zkernel_channelsZ bias_channelsZ num_mel_binsZresnet_in_channelsZ kernel_predictor_hidden_channelsZresnet_hidden_channelsr#Zresnet_kernel_sizeZkernel_predictor_num_blocks num_blocksr%rr' input_conv ModuleListrange resblocks kernel_conv bias_conv)r*rr?r@rr+rDrr!ss,    zUnivNetKernelPredictor.__init__ spectrogramc Cs|j\}}}||}tj||j}|jD]}||}q||}||}| ||j |j |j |j |} | ||j |j |} | | fS)a Maps a conditioning log-mel spectrogram to a tensor of convolutional kernels and biases, for use in location variable convolutional layers. Note that the input spectrogram should have shape (batch_size, input_channels, seq_length). Args: spectrogram (`torch.FloatTensor` of shape `(batch_size, input_channels, seq_length)`): Tensor containing the log-mel spectrograms. Returns: Tuple[`torch.FloatTensor, `torch.FloatTensor`]: tuple of tensors where the first element is the tensor of location variable convolution kernels of shape `(batch_size, self.conv_layers, self.conv_in_channels, self.conv_out_channels, self.conv_kernel_size, seq_length)` and the second element is the tensor of location variable convolution biases of shape `(batch_size, self.conv_layers. self.conv_out_channels, seq_length)`. )shaperJrr/r0r%rMrNrOviewr@rGrHr? contiguous) r*rPZ batch_sizerCZ seq_lengthr-resblockZkernel_hidden_statesZbias_hidden_stateskernelsbiasesrrrr2s4      zUnivNetKernelPredictor.forwardcCsVtjj}ttjjdrtjjj}||j|jD]}|q||j||j dSr3) rr5r4r6r7rJrMr9rNrOr*r4layerrrrr9s     z(UnivNetKernelPredictor.apply_weight_normcCsBtj|j|jD]}|q tj|jtj|jdSr.)rr5r:rJrMrNrOr*rXrrrr:s   z)UnivNetKernelPredictor.remove_weight_norm)rr> rrrrr intr!rrr2r9r:r<rrr+rr=bs&. r=c sreZdZdZdededeffdd Zddd Z dd ej d ej d ej dedef ddZ ddZ ddZ Z S)UnivNetLvcResidualBlocka Implementation of the location variable convolution (LVC) residual block for the UnivNet residual network. Parameters: config: (`UnivNetConfig`): Config for the `UnivNetModel` model. kernel_size (`int`): The kernel size for the dilated 1D convolutional layer. dilation (`int`): The dilation for the dilated 1D convolutional layer. rr$dilationcs\t|j|_||_||_|j|_|j|jdd}tj|j|j|j||jd|_ dS)Nr r)rr]) r r!rFhidden_channelsr$r]r%rr'conv)r*rr$r]rr+rrr!s  z UnivNetLvcResidualBlock.__init__cCs|}tj||j}||}tj||j}|j||||d}t|ddd|jddft |dd|jdddf}||}|SN)hop_size) rr/r0r%r_location_variable_convolutionrZsigmoidr^tanh)r*r-kernelrrbr1rrrr2s $zUnivNetLvcResidualBlock.forwardr r-rerrbcCsB|j\}}}|j\}}} } } || |kr!td| |d|d|t| dd} tj|| | fdd}|d|d| |}||krPtj|d|fdd}|d||}|d d d d d d d d d |f}|dd }|d | d}t d ||} | j tj d } | d  d j tj d }| |} | || d } | S)u Performs location-variable convolution operation on the input sequence (hidden_states) using the local convolution kernel. This was introduced in [LVCNet: Efficient Condition-Dependent Modeling Network for Waveform Generation](https://arxiv.org/abs/2102.10815) by Zhen Zheng, Jianzong Wang, Ning Cheng, and Jing Xiao. Time: 414 μs ± 309 ns per loop (mean ± std. dev. of 7 runs, 1000 loops each), test on NVIDIA V100. Args: hidden_states (`torch.FloatTensor` of shape `(batch_size, in_channels, in_length)`): The input sequence of shape (batch, in_channels, in_length). kernel (`torch.FloatTensor` of shape `(batch_size, in_channels, out_channels, kernel_size, kernel_length)`): The local convolution kernel of shape (batch, in_channels, out_channels, kernel_size, kernel_length). bias (`torch.FloatTensor` of shape `(batch_size, out_channels, kernel_length)`): The bias for the local convolution of shape (batch, out_channels, kernel_length). dilation (`int`, *optional*, defaults to 1): The dilation of convolution. hop_size (`int`, *optional*, defaults to 256): The hop_size of the conditioning sequence. Returns: `torch.FloatTensor`: the output sequence after performing local convolution with shape (batch_size, out_channels, in_length). z#Dim 2 of `hidden_states` should be z ) but got zX. Please check `hidden_states` or `kernel` and `hop_size` to make sure they are correct.r rZconstantrrNr>zbildsk,biokl->bolsd)Z memory_format)rQ ValueErrorr[rr/padZunfold transposerZeinsumtoZchannels_last_3d unsqueezerSrR)r*r-rerr]rbbatchrCZ in_lengthZ out_channelsr$Z kernel_lengthrZoutput_hidden_statesrrrrc s*  & z5UnivNetLvcResidualBlock.location_variable_convolutioncCs.tjj}ttjjdrtjjj}||jdSr3)rr5r4r6r7r_r8rrrr9Ns z)UnivNetLvcResidualBlock.apply_weight_normcCstj|jdSr.)rr5r:r_r;rrrr:Usz*UnivNetLvcResidualBlock.remove_weight_normr`)r r`)rrrrr r[r!r2rrrcr9r:r<rrr+rr\s2   Ar\csXeZdZdZ ddededeffdd Zdejd ejfd d Z d d Z ddZ Z S)UnivNetLvcBlocka# Implementation of the location variable convolution (LVC) residual block of the UnivNet residual block. Includes a `UnivNetKernelPredictor` inside to predict the kernels and biases of the LVC layers. Based on LVCBlock in [maum-ai/univnet](https://github.com/maum-ai/univnet/blob/9bb2b54838bb6d7ce767131cc7b8b61198bc7558/model/lvcnet.py#L98) Parameters: config (`UnivNetConfig`): Config for the `UnivNetModel` model. layer_id (`int`): An integer corresponding to the index of the current LVC resnet block layer. This should be between 0 and `len(config.resblock_stride_sizes) - 1)` inclusive. lvc_hop_size (`int`, *optional*, defaults to 256): The hop size for the location variable convolutional layers. r`rlayer_id lvc_hop_sizecstj_j|_j|_j|_ |_ j _ t j _ tjjjdjjjdjdjdd_tjj _tfddtj D_dS)Nr)striderZoutput_paddingcs g|] }tjj|qSr)r\r$ dilationsrBirr*rrrEs z,UnivNetLvcBlock.__init__..)r r!rFr^resblock_kernel_sizesr$resblock_stride_sizesrqZresblock_dilation_sizesrrcond_hop_lengthr%lenrIrConvTranspose1d convt_prer=kernel_predictorrKrLrM)r*rrorpr+rurr!ks(       zUnivNetLvcBlock.__init__r-rPc Cstj||j}||}||\}}t|jD]/\}}|dd|ddddddddf}|dd|ddddf}|||||jd}q|Sra) rr/r0r%r{r| enumeraterMrx) r*r-rPrUrVrtrTrerrrrr2s (zUnivNetLvcBlock.forwardcCsLtjj}ttjjdrtjjj}||j|j|jD]}|qdSr3) rr5r4r6r7r{r|r9rMrWrrrr9s     z!UnivNetLvcBlock.apply_weight_normcCs0tj|j|j|jD]}|qdSr.)rr5r:r{r|rMrYrrrr:s    z"UnivNetLvcBlock.remove_weight_normrmrZrrr+rrnYs rncseZdZeZdZdeffdd Ze    ddej de ej de ej de ej d e e d e eej eff d d Zd dZddZddZZS) UnivNetModelinput_featuresrcsttj|_j|_tjjj ddddd|_ tj }d}gj D] }||} |q*t fddt|D|_tjj ddddd|_|dS) Nr rZreflect)r$rqr padding_modecsg|] }t||dqS))rorp)rnrsrZ hop_lengthsrrrEsz)UnivNetModel.__init__..)rr)r r!ryrvZ num_kernelsr%rr'r"rFconv_prerwappendrKrLrM conv_postZ post_init)r*rZ num_layersZ hop_lengthrqr+rrr!s0       zUnivNetModel.__init__Nnoise_sequence padding_mask generator return_dictreturncCs|dur|n|jj}|dk}|s|d}|j\}}} |dur/|dk} | s.|d}n|||jjf} tj| ||j|j d}|jd} |dkrV| dkrV| |dd}n| dkre|dkre| | dd}| |krtt d| d|d|dur|dkr|d}|jd} | |krt d | d|d| d d}| d d}| |}|jD]}|||}qtj||j}||}t|}|d}d}|durtj|dd }|s||f}|St||d S) a input_features (`torch.FloatTensor`): Tensor containing the log-mel spectrograms. Can be batched and of shape `(batch_size, sequence_length, config.num_mel_channels)`, or un-batched and of shape `(sequence_length, config.num_mel_channels)`. noise_sequence (`torch.FloatTensor`, *optional*): Tensor containing a noise sequence of standard Gaussian noise. Can be batched and of shape `(batch_size, sequence_length, config.model_in_channels)`, or un-batched and of shape (sequence_length, config.model_in_channels)`. If not supplied, will be randomly generated. padding_mask (`torch.BoolTensor`, *optional*): Mask indicating which parts of each sequence are padded. Mask values are selected in `[0, 1]`: - 1 for tokens that are **not masked** - 0 for tokens that are **masked** The mask can be batched and of shape `(batch_size, sequence_length)` or un-batched and of shape `(sequence_length,)`. generator (`torch.Generator`, *optional*): A [torch generator](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. return_dict: Whether to return a [`~utils.ModelOutput`] subclass instead of a plain tuple. Example: ```python >>> from transformers import UnivNetFeatureExtractor, UnivNetModel >>> from datasets import load_dataset, Audio >>> model = UnivNetModel.from_pretrained("dg845/univnet-dev") >>> feature_extractor = UnivNetFeatureExtractor.from_pretrained("dg845/univnet-dev") >>> ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation") >>> # Resample the audio to the feature extractor's sampling rate. >>> ds = ds.cast_column("audio", Audio(sampling_rate=feature_extractor.sampling_rate)) >>> inputs = feature_extractor( ... ds[0]["audio"]["array"], sampling_rate=ds[0]["audio"]["sampling_rate"], return_tensors="pt" ... ) >>> audio = model(**inputs).waveforms >>> list(audio.shape) [1, 140288] ``` Nrr)rdtypedevicer z&The batch size of `noise_sequence` is z+ and the batch size of `input_features` is z', but the two are expected to be equal.z$The batch size of `padding_mask` is r)dim)rr)rZuse_return_dictrrkrQr"rZrandnrrrepeatrgrirrMrr/r0r%rrdZsqueezesumr)r*rrrrrZspectrogram_batchedZspectrogram_batch_sizeZspectrogram_lengthrCZnoise_sequence_batchedZnoise_sequence_shapeZnoise_sequence_batch_sizeZpadding_mask_batch_sizer-rTZwaveformrZoutputsrrrr2sl3                  zUnivNetModel.forwardcCsNt|tjtjtjfr#|jjjd|jj d|j dur%|j j dSdSdS)zInitialize the weights.g)meanZstdN) isinstancerZLinearr'rzweightdataZnormal_rZinitializer_rangerZzero_)r*modulerrr _init_weightsSs  zUnivNetModel._init_weightscCsLtjj}ttjjdrtjjj}||j|jD]}|q||jdSr3) rr5r4r6r7rrMr9rrWrrrr9Zs    zUnivNetModel.apply_weight_normcCs4tj|j|jD]}|q tj|jdSr.)rr5r:rrMrrYrrrr:ds  zUnivNetModel.remove_weight_norm)NNNN)rrrr Z config_classZmain_input_namer!r rrr Generatorboolrrrr2rr9r:r<rrr+rr~s2' } r~)r dataclassesrtypingrrrrZtorch.utils.checkpointrZmodeling_outputsrZmodeling_utilsr r5r r Zconfiguration_univnetr Z get_loggerrloggerrModulerr=r\rnr~__all__rrrrs*      1xP B