o id@sdZddlmZddlmZmZddlZddlmZddlm Z ddl m Z dd l m Z mZd d lmZeeZee d d Gddde ZGdddejZGdddejZGdddejZGdddejZe Gddde ZdgZdS)zPyTorch UnivNetModel model.) dataclass)OptionalUnionN)nn) ModelOutput)PreTrainedModel)auto_docstringlogging) UnivNetConfigz 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`]). ) custom_introc@s6eZdZUdZdZeejed<dZ eejed<dS)UnivNetModelOutputa" 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__rrri/home/ubuntu/veenaModal/venv/lib/python3.10/site-packages/transformers/models/univnet/modeling_univnet.pyrs 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_channelschannelskernel_predictor_conv_size kernel_sizekernel_predictor_dropout dropout_probleaky_relu_sloperDropoutdropoutConv1dconv1conv2)selfrr __class__rrr!<s  z,UnivNetKernelPredictorResidualBlock.__init__ hidden_statescCsJ|}||}||}tj||j}||}tj||j}||SN)r*r,r functional leaky_relur(r-)r.r1residualrrrforwardLs   z+UnivNetKernelPredictorResidualBlock.forwardcCs8tjj}ttjjdrtjjj}||j||jdSN weight_norm)rutilsr8hasattrparametrizationsr,r-r.r8rrrapply_weight_normVs   z5UnivNetKernelPredictorResidualBlock.apply_weight_normcCs tj|jtj|jdSr2)rr9remove_weight_normr,r-r.rrrr>^sz6UnivNetKernelPredictorResidualBlock.remove_weight_norm) rrrrr r!rrr6r=r> __classcell__rrr/rr2s  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_channelsrCrDkernel_channels bias_channels num_mel_binsresnet_in_channels kernel_predictor_hidden_channelsresnet_hidden_channelsr$resnet_kernel_sizekernel_predictor_num_blocks num_blocksr(rr+ input_conv ModuleListrange resblocks kernel_conv bias_conv)r.rrCrDrr/rHrr!ts,    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)`. )shaperVrr3r4r(rYrZr[viewrDrKrLrC contiguous) r.r\ batch_sizerG seq_lengthr1resblockkernel_hidden_statesbias_hidden_stateskernelsbiasesrrrr6s4      zUnivNetKernelPredictor.forwardcCsVtjj}ttjjdrtjjj}||j|jD]}|q||j||j dSr7) rr9r8r:r;rVrYr=rZr[r.r8layerrrrr=s     z(UnivNetKernelPredictor.apply_weight_normcCsBtj|j|jD]}|q tj|jtj|jdSr2)rr9r>rVrYrZr[r.rhrrrr>s   z)UnivNetKernelPredictor.remove_weight_norm)rrB rrrrr intr!rrr6r=r>r@rrr/rrAcs&. rAc 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)rrm) r r!rJhidden_channelsr%rmr(rr+conv)r.rr%rmrr/rrr!s  z UnivNetLvcResidualBlock.__init__cCs|}tj||j}||}tj||j}|j||||d}t|ddd|jddft |dd|jdddf}||}|SN)hop_size) rr3r4r(rolocation_variable_convolutionrsigmoidrntanh)r.r1kernelrrrr5rrrr6s $zUnivNetLvcResidualBlock.forwardr r1rvrrrcCsB|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://huggingface.co/papers/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 rconstantrrNrBzbildsk,biokl->bolsd) memory_format)r] ValueErrorrkrr3padunfold transposereinsumtochannels_last_3d unsqueezer_r^)r.r1rvrrmrrbatchrG in_length out_channelsr% kernel_lengthroutput_hidden_statesrrrrss*  & z5UnivNetLvcResidualBlock.location_variable_convolutioncCs.tjj}ttjjdrtjjj}||jdSr7)rr9r8r:r;ror<rrrr=Os z)UnivNetLvcResidualBlock.apply_weight_normcCstj|jdSr2)rr9r>ror?rrrr>Vsz*UnivNetLvcResidualBlock.remove_weight_normrp)r rp)rrrrr rkr!r6rrrsr=r>r@rrr/rrls2   ArlcsXeZdZdZ 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. rprlayer_id lvc_hop_sizecstj_j|_j|_j|_ |_ j _ t j _ tjjjdjjjdjdjdd_tjj _tfddtj D_dS)Nr)strideroutput_paddingcs g|] }tjj|qSr)rlr% dilationsrFirr.rrrIs z,UnivNetLvcBlock.__init__..)r r!rJrnresblock_kernel_sizesr%resblock_stride_sizesrresblock_dilation_sizesrcond_hop_lengthr(lenrUrConvTranspose1d convt_prerAkernel_predictorrWrXrY)r.rrrr/rrr!ls(       zUnivNetLvcBlock.__init__r1r\c Cstj||j}||}||\}}t|jD]/\}}|dd|ddddddddf}|dd|ddddf}|||||jd}q|Srq) rr3r4r(rr enumeraterYr) r.r1r\rerfrrbrvrrrrr6s (zUnivNetLvcBlock.forwardcCsLtjj}ttjjdrtjjj}||j|j|jD]}|qdSr7) rr9r8r:r;rrr=rYrgrrrr=s     z!UnivNetLvcBlock.apply_weight_normcCs0tj|j|j|jD]}|qdSr2)rr9r>rrrYrirrrr>s    z"UnivNetLvcBlock.remove_weight_normrrjrrr/rrZs rcseZdZUeed<dZdeffdd Ze    ddej de ej de ej de ej d e e d e eej eff d d Zd dZddZddZZS) UnivNetModelrinput_featurescsttj|_j|_tjjj ddddd|_ tj }d}gj D] }||} |q*t fddt|D|_tjj ddddd|_|dS) Nr rreflect)r%rr padding_modecsg|] }t||dqS))rr)rrr hop_lengthsrrrIsz)UnivNetModel.__init__..)rr)r r!rr num_kernelsr(rr+r"rJconv_prerappendrWrXrY conv_post post_init)r.r num_layers hop_lengthrr/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 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)ruse_return_dictrrr]r"rrandnrrrepeatrzr}rrYrr3r4r(rrusqueezesumr)r.rrrrrspectrogram_batchedspectrogram_batch_sizespectrogram_lengthrGnoise_sequence_batchednoise_sequence_shapenoise_sequence_batch_sizepadding_mask_batch_sizer1rbwaveformroutputsrrrr6sl0                  zUnivNetModel.forwardcCsNt|tjtjtjfr#|jjjd|jj d|j dur%|j j dSdSdS)zInitialize the weights.g)meanstdN) isinstancerLinearr+rweightdatanormal_rinitializer_rangerzero_)r.modulerrr _init_weightsQs  zUnivNetModel._init_weightscCsLtjj}ttjjdrtjjj}||j|jD]}|q||jdSr7) rr9r8r:r;rrYr=rrgrrrr=Xs    zUnivNetModel.apply_weight_normcCs4tj|j|jD]}|q tj|jdSr2)rr9r>rrYrrirrrr>bs  zUnivNetModel.remove_weight_norm)NNNN)rrrr rmain_input_namer!r rrr Generatorboolrtuplerr6rr=r>r@rrr/rrs2 ' z r)r dataclassesrtypingrrrrmodeling_outputsrmodeling_utilsrr9r r configuration_univnetr get_loggerrloggerrModulerrArlrr__all__rrrrs.       1xP ?