o }oi8Y@sddlZddlZddlZddlmZmZddlmZmZm Z m Z ddl Z ddl Z ddl ZddlZddlmZddlmZmZddlmZddlmZddlmZdd lmZdd lmZdd lmZdd l m!Z!dd l"m#Z#ddl$m%Z%ddl&m'Z'm(Z(dgZ)Gddde#eZ*dS)N)ABCabstractmethod)DictListOptionalUnion)Trainer) DictConfig OmegaConf)tqdm))inject_dataloader_value_from_model_config)ChannelSelectorType)audio_to_audio_dataset)LhotseAudioToTargetDataset)AudioMetricWrapper)!get_lhotse_dataloader_from_config)ModelPT)PretrainedModelInfo)logging model_utilsAudioToAudioModelcseZdZdZdGdedeffdd ZddZdHd efd d Z dHd efd dZ e dIde d efddZ fddZfddZdJde fddZdJddZdIde d efddZdJde fddZdJde fdd Zd!eefd"d#Zd$eeeeffd%d&Zd'eeeeffd(d)Zd*eeeeffd+d,Zd!ed-d.fd/d0Zed1ejd2e d-ejfd3d4Z e! 5   dKd6e"ed7ed8e d9ee d:ee#d;eed-e"efd|jvr>|j|D]\}}|}|||d|<q(d|d|iS)Ncsg|] }|dqS)_lossr).0xr'rr sz@AudioToAudioModel.multi_evaluation_epoch_end..rgr2_log)r9stackmeanr7r2r>computereset) routputsrJr' loss_meantensorboard_logsrLr5valuerrjr multi_evaluation_epoch_endsz,AudioToAudioModel.multi_evaluation_epoch_endcC|||dSrVrvrrrrJrrr multi_validation_epoch_endz,AudioToAudioModel.multi_validation_epoch_endcCrwrYrxryrrr multi_test_epoch_endr{z&AudioToAudioModel.multi_test_epoch_endconfigc Cst|j|dd|ddrt||j|jtdS|dd}|r%td|ddr/td d |vrC|d durCt d |dSt j |d }t |d rR|j }nt |jdd ra|jdj }n |jdjdj }tjjj||d||dd|d|dd|dddS)N sample_rate)key use_lhotseF) global_rank world_sizedataset is_concatzConcat not implemented is_tarredTarred datasets not supportedmanifest_filepathzJCould not load dataset as `manifest_filepath` was None. Provided config : r} collate_fnr batch_size drop_lastshuffle num_workers pin_memory)rrrrrrr)r rr8rrrrNotImplementedErrorrr$rget_audio_to_target_datasetr7rdatasetsr9rBdata DataLoader)rr}rrrrrr _setup_dataloader_from_configs8        z/AudioToAudioModel._setup_dataloader_from_configtrain_data_configcCsLd|vrd|d<|jd|d|j|d|_d|vr"|dr$tddSdS) aR Sets up the training data loader via a Dict-like object. Args: train_data_config: A config that contains the information regarding construction of a training dataset. Supported Datasets: - :class:`~nemo.collections.asr.data.audio_to_audio.AudioToTargetDataset` rTtrain dataset_namer}rrrN)_update_dataset_configr _train_dlr)rrrrr setup_training_datas z%AudioToAudioModel.setup_training_dataval_data_configcC0d|vrd|d<|jd|d|j|d|_dS)aT Sets up the validation data loader via a Dict-like object. Args: val_data_config: A config that contains the information regarding construction of a validation dataset. Supported Datasets: - :class:`~nemo.collections.asr.data.audio_to_audio.AudioToTargetDataset` rF validationrrN)rrr,)rrrrr setup_validation_data z'AudioToAudioModel.setup_validation_datatest_data_configcCr)aI Sets up the test data loader via a Dict-like object. Args: test_data_config: A config that contains the information regarding construction of a test dataset. Supported Datasets: - :class:`~nemo.collections.asr.data.audio_to_audio.AudioToTargetDataset` rFr)rrN)rrr.)rrrrr setup_test_data rz!AudioToAudioModel.setup_test_datareturnztorch.utils.data.DataLoadercCs^|d|j|d|dddd|dd|dt|dtdd d }|jt|d }|S) aIPrepare a dataloader for processing files. Args: config: A python dictionary which contains the following keys: manifest_filepath: path to a manifest file input_key: key with audio filepaths in the manifest input_channel_selector: Optional, used to select a subset of channels from input audio files batch_size: batch size for the dataloader num_workers: number of workers for the dataloader Returns: A pytorch DataLoader for the given manifest filepath. r input_keyinput_channel_selectorNrFrr(T) rr~rr target_keytarget_channel_selectorrrrrr)r~r8minos cpu_countrr )rr} dl_configtemporary_dataloaderrrr _setup_process_dataloaders  z+AudioToAudioModel._setup_process_dataloaderinput batch_lengthcCs.|d}||}d|f}tjj||ddS)zTrim or pad the output to match the batch length. Args: input: tensor with shape (B, C, T) batch_length: int Returns: Tensor with shape (B, C, T), where T matches the batch length. rconstant)sizer9r: functionalpad)rr input_length pad_lengthrrrr match_batch_length<s z$AudioToAudioModel.match_batch_lengthr(paths2audio_files output_dirrrr input_dirc Cs|dus t|dkr iS|durt|td}g}|j}t|j} z|| t } t t j t} tj| d} t| ddd } |D]}|tj|dd }| t|d qOWdn1sowY| d |t|t||d }tj|st|||}d}t|d dD]}|d}|d}|jdkr|d}|j|| || d\}}t |!dD]W}|durtjj"|||d}ntj#||}tj||}tjtj$|sttj$|||ddd||f%&}t'||j(|j)d|d7}|*|q~~qWdn 1s)wYW|j+|d|dur>|,t | |S|j+|d|durT|,t | w)a` Takes paths to audio files and returns a list of paths to processed audios. Args: paths2audio_files: paths to audio files to be processed output_dir: directory to save the processed files batch_size: (int) batch size to use during inference. num_workers: Number of workers for the dataloader input_channel_selector (int | Iterable[int] | str): select a single channel or a subset of channels from multi-channel audio. If set to `'average'`, it performs averaging across channels. Disabled if set to `None`. Defaults to `None`. input_dir: Optional, directory that contains the input files. If provided, the output directory will mirror the input directory structure. Returns: Paths to processed audio signals. Nrr(z manifest.jsonwzutf-8)encoding)path)input_filepathduration r)rrrrr Processing)desc) input_signalr)startfloat)modeT)-r-rrrtrainingnext parametersrFevalfreezer get_verbosity set_verbosityWARNINGtempfileTemporaryDirectoryrrHopenlibrosa get_durationwritejsondumpsisdirmakedirsrr ndim unsqueezeforwardrEr=rrelpathbasenamedirnamecpunumpysfTr~rDrunfreeze)rrrrrrrpaths2processed_filesrrF logging_leveltmpdirtemporary_manifest_filepathfp audio_fileentryr}rfile_idx test_batchrrprocessed_batchrl example_idxfilepath_relative output_file output_signalrrr processNs~        "  C    zAudioToAudioModel.processList[PretrainedModelInfo]cCst|}|S)z This method returns a list of pre-trained model which can be instantiated directly from NVIDIA's NGC cloud. Returns: List of available pre-trained models. )r&resolve_subclass_pretrained_model_info)clslist_of_modelsrrr list_available_modelss z'AudioToAudioModel.list_available_modelscCs2d|_d|jvr|jdr|jd|_dSdSdS)aT Utility method that must be explicitly called by the subclass in order to support optional optimization flags. This method is the only valid place to access self.cfg prior to DDP training occurs. The subclass may chose not to support this method, therefore all variables here must be checked via hasattr() F skip_nan_gradN)_skip_nan_gradr"r%rrr setup_optimization_flagssz*AudioToAudioModel.setup_optimization_flagscstt|drf|jrht|j}tjdg|tj d}| D]"\}}|j durDt |j  p:t|j  }|sD|d}nq"tjrUtjj|tjjjd|dkrjtd|jdd dSdSdSdS) zH zero-out the gradients which any of them is NAN or INF rr()rFdtypeNr)opzCdetected inf or nan values in gradients! Setting gradients to zero.F) set_to_none)ron_after_backwardr7rrrrFr9tensorfloat32named_parametersgradisnananyisinf distributedis_initialized all_reduceReduceOpMINrr$ zero_grad)rrFvalid_gradients param_nameparamis_not_nan_or_infrrr rs$  "  z#AudioToAudioModel.on_after_backwardcCs|jdd|_|js gSg}ddlm}t|jtr|j}n|jg}t|D]\}}| ||||j j |jj |jj |j|jdddq'|S)zW Create an callback to add audio/spectrogram into tensorboard & wandb. log_configNr) SpeechEnhancementLoggingCallbackmax_utts) data_loaderdata_loader_idxloggerslog_tensorboard log_wandbr~r)rr8r,nemo.collections.audio.parts.utils.callbacksrr+r,r enumeraterDrrrrr~)r log_callbacksr data_loadersrrrrr configure_callbackss*    z%AudioToAudioModel.configure_callbacksrQ)r&)rr&)r)r(NNN)rr)+__name__ __module__ __qualname____doc__r rrrstrr1rPrintrTrWrZrbrfrvrzr|rrrrrrrr staticmethodr9Tensorrno_gradrr r classmethodrrrr __classcell__rrrr r*s^ <   + z  )+rrrabcrrtypingrrrrrAr soundfilerr9lightning.pytorchr omegaconfr r r /nemo.collections.asr.data.audio_to_text_datasetr 0nemo.collections.asr.parts.preprocessing.segmentr nemo.collections.audio.datar1nemo.collections.audio.data.audio_to_audio_lhotser$nemo.collections.audio.metrics.audior#nemo.collections.common.data.lhotsernemo.core.classesrnemo.core.classes.commonr nemo.utilsrr__all__rrrrr s.