o j~@sddlZddlZddlZddlZddlmZmZmZddlZ ddl Z ddl Z ddl m Z ddlmZmZmZddlmZddlmZe jddd Zd d Zd d ZGddde ZGddde ZGdddZGdddZdS)N)DictListUnion)Dataset) prepare_dataprepare_stop_targetprepare_tensor)AudioProcessor)compute_energy file_systemcCsld}d}t|dkr|\}}}}}nt|dkr|\}}}}nt|dkr+|\}}}ntd|||||fS)Nz% [!] Dataset cannot parse the sample.)len ValueError)item language_name attn_filetextwav_file speaker_namerK/home/kuhnn/.local/lib/python3.10/site-packages/TTS/tts/datasets/dataset.py _parse_samples    rcCs|dtjj|jS)Ng?)nprandomrandshape)wavrrrnoise_augment_audio#srcCst|ddd}|S)Nzutf-8ignore)base64urlsafe_b64encodeencodedecode)stringfilenamerrrstring2filename'sr'c/seZdZddddddddddddeddedddddddddfdededed eed d d ed ede de dedededededede dededededededef.fdd Z e d d!Z e d"d#Z e jd$d#Z d%d&Zd'd(ZdNd)ed*dfd+d,Zd-d.Zd/d0Zd1d2Zd3d4Zed5d6Zd7d8Zd9d:Zed;d<Zed=eed>ed?efd@dAZed eefdBdCZedefdDdEZedFdGZdHdIZedJdKZ dLdMZ!Z"S)O TTSDatasetFNrinfoutputs_per_stepcompute_linear_specapsamples tokenizer TTSTokenizer compute_f0r f0_cache_pathenergy_cache_path return_wavbatch_group_size min_text_len max_text_len min_audio_len max_audio_lenphoneme_cache_pathprecompute_num_workersspeaker_id_mappingd_vector_mappinglanguage_id_mappinguse_noise_augmentstart_by_longestverbosecst| |_||_||_||_| |_||_||_||_ | |_ ||_ ||_ | |_ | |_||_||_||_||_||_||_||_||_d|_d|_||_|jjr\t|j|j||d|_|rit|j|j||d|_|rvt |j|j| |d|_!|jr|"dSdS)u8 Generic 📂 data loader for `tts` models. It is configurable for different outputs and needs. If you need something different, you can subclass and override. Args: outputs_per_step (int): Number of time frames predicted per step. compute_linear_spec (bool): compute linear spectrogram if True. ap (TTS.tts.utils.AudioProcessor): Audio processor object. samples (list): List of dataset samples. tokenizer (TTSTokenizer): tokenizer to convert text to sequence IDs. If None init internally else use the given. Defaults to None. compute_f0 (bool): compute f0 if True. Defaults to False. compute_energy (bool): compute energy if True. Defaults to False. f0_cache_path (str): Path to store f0 cache. Defaults to None. energy_cache_path (str): Path to store energy cache. Defaults to None. return_wav (bool): Return the waveform of the sample. Defaults to False. batch_group_size (int): Range of batch randomization after sorting sequences by length. It shuffles each batch with bucketing to gather similar lenght sequences in a batch. Set 0 to disable. Defaults to 0. min_text_len (int): Minimum length of input text to be used. All shorter samples will be ignored. Defaults to 0. max_text_len (int): Maximum length of input text to be used. All longer samples will be ignored. Defaults to float("inf"). min_audio_len (int): Minimum length of input audio to be used. All shorter samples will be ignored. Defaults to 0. max_audio_len (int): Maximum length of input audio to be used. All longer samples will be ignored. The maximum length in the dataset defines the VRAM used in the training. Hence, pay attention to this value if you encounter an OOM error in training. Defaults to float("inf"). phoneme_cache_path (str): Path to cache computed phonemes. It writes phonemes of each sample to a separate file. Defaults to None. precompute_num_workers (int): Number of workers to precompute features. Defaults to 0. speaker_id_mapping (dict): Mapping of speaker names to IDs used to compute embedding vectors by the embedding layer. Defaults to None. d_vector_mapping (dict): Mapping of wav files to computed d-vectors. Defaults to None. use_noise_augment (bool): Enable adding random noise to wav for augmentation. Defaults to False. start_by_longest (bool): Start by longest sequence. It is especially useful to check OOM. Defaults to False. verbose (bool): Print diagnostic information. Defaults to false. r)F)r;) cache_pathr;N)#super__init__r5_samplesr+r,r4r1r r2r3r8r9r6r7r-r:r<r=r>r?r@rArescue_item_idxpitch_computedr/ use_phonemesPhonemeDatasetr.phoneme_dataset F0Dataset f0_dataset EnergyDatasetenergy_dataset print_logs)selfr+r,r-r.r/r1r r2r3r4r5r6r7r8r9r:r;r<r=r>r?r@rA __class__rrrD.sP U    zTTSDataset.__init__cCs@g}|jD]}t|^}}}tj|dd}||q|S)N)r.rospathgetsizeappend)rPlensr_r audio_lenrrrlengthss   zTTSDataset.lengthscC|jSN)rErPrrrr.szTTSDataset.samplescCsD||_t|dr ||j_t|dr||j_t|dr ||j_dSdS)NrLrNrJ)rEhasattrrLr.rNrJ)rP new_samplesrrrr.s    cC t|jSr^rr.r_rrr__len__ zTTSDataset.__len__cCs ||Sr^) load_data)rPidxrrr __getitem__rezTTSDataset.__getitem__levelreturncCXd|}tdt|dt|d|j|dt|dt|jdS)N  z> DataLoader initialization| > Tokenizer:r)| > Number of instances : printr/rOrr.rPriindentrrrrO zTTSDataset.print_logscCs|j|}|jdks J|SNr)r-load_wavsize)rPr&waveformrrrrvs zTTSDataset.load_wavcCsD|j|}||dksJ|d|dt|ddks J|S)Nrz != token_idsr)rJr)rPrgrout_dictrrr get_phonemess "zTTSDataset.get_phonemescC,|j|}|j|}|d|dksJ|SNaudio_unique_name)rLr.rPrgrzrrrrget_f0  zTTSDataset.get_f0cCr|r})rNr.rrrr get_energyrzTTSDataset.get_energycCs t|Sr^)rload)rrrr get_attn_masks zTTSDataset.get_attn_maskcCs6|jjr |||d}n|j|}tj|tjdS)Nrydtype)r/rHr{ text_to_idsrarrayint32)rPrgrryrrr get_token_idss zTTSDataset.get_token_idsc Cs|j|}|d}tj||dtjd}|jrt|}|||d}d}d|vr2||d}t ||j ks@t ||j krM|j d7_ | |j Sd}|jrY||d}d}|jre||d}|||||||d|d|d tj|d|d d } | S) Nr audio_fileralignment_filer)f0energyrlanguager~) raw_textryrpitchrattnitem_idxrr wav_file_namer~)r.rasarrayrvfloat32r?rrrrr7r8rFrfr1rr rrUrVbasename) rPrgrrrryrrrsamplerrrrfs>   zTTSDataset.load_datacCsPg}|D]!}tj|ddd}t|d}||d<||d<||g7}q|S)NrrSrTr audio_length text_length)rUrVrWr)r.rarr text_lenghtrrr_compute_lengths&s  zTTSDataset._compute_lengthsr\min_lenmax_lencCsRt|}g}g}|D]}||}||ks||kr||q ||q ||fSr^)rargsortrX)r\rridxs ignore_idxkeep_idxrglengthrrrfilter_by_length1s   zTTSDataset.filter_by_lengthcCsdd|D}t|}|S)NcSg|]}|dqSrr.0srrr @z-TTSDataset.sort_by_length..)rr)r. audio_lengthsrrrrsort_by_length>s zTTSDataset.sort_by_lengthcCsX|dksJtt||D]}||}||}|||}t|||||<q|Sru)rangerrshuffle)r.r5ioffset end_offset temp_itemsrrrcreate_bucketsDs   zTTSDataset.create_bucketscCs g}|D] }|||q|Sr^)rX)rr. samples_newrgrrr_select_samples_by_idxOsz!TTSDataset._select_samples_by_idxc Cs||j}dd|D}dd|D}|||j|j\}}|||j|j\}}tt|t|@}tt|t|B} | ||}| |} |j rZ| d} | d| d<| | d<| | |}t |dkrjt d|jdkrv|||j}dd|D}dd|D}||_|jrtd td t|td t|td t|td tdt|tdt|tdt|tdt | td|jdSdS)zSort `items` based on text length or audio length in ascending order. Filter out samples out or the length range. cSrrrrrrrrr]rz1TTSDataset.preprocess_samples..cSrrrrrrrr^rrz [!] No samples leftcSrrrrrrrryrcSrrrrrrrrzrz | > Preprocessing samplesz | > Max text length: {}z | > Min text length: {}z | > Avg text length: {}z | z | > Max audio length: {}z | > Min audio length: {}z | > Avg audio length: {}z' | > Num. instances discarded samples: z | > Batch group size: {}.N)rr.rr6r7r8r9listsetrrr@r RuntimeErrorr5rrArqformatrmaxminmean) rPr. text_lengthsrtext_ignore_idx text_keep_idxaudio_ignore_idxaudio_keep_idxrr sorted_idxs longest_idxsrrrpreprocess_samplesVsB       zTTSDataset.preprocess_samplescs6tjt|ddd\}}fdd|D||fS)zSort the batch by the input text length for RNN efficiency. Args: batch (Dict): Batch returned by `__getitem__`. text_lengths (List[int]): Lengths of the input character sequences. rT)dim descendingcsg|]}|qSrrrrgbatchrrrrz*TTSDataset._sort_batch..)torchsort LongTensor)rrids_sorted_decreasingrrr _sort_batchs zTTSDataset._sort_batchcs tdtjjrztddD}|\}}fdddDjdur9fdddD}nd}jdurPt d }fd d|D}nd}j rafd dd D}nd}fd ddD}dd|D} fdd|D} dd| D} t | j } t dtj} t|j }|ddd}t|}t| } t|}t| } t| } |durt|}|durt|}|durt|}d} jr fdddD} t| j } | ddd} |jd| jdksJt| } d}jrqdddD}t| jj}t|}ttdd|}tdD]3\}}| |}tj|djjj fdd}|d|jj}t |||ddd|jdf<q7|!ddj"rt d}|jd|jdksJd|jd|jt|dddddf}nd}j#rt d}|jd|jdksJd|jd|jt|dddddf}nd}d}dddur=fdd|D}t|D]=\}}|jd|jd}| jd|jd}|dkr|dksJd|d |t|d|gd|gg}|||<qt|j }t|$d}id!| d"|d#d d$| d%|d&| d'| d(d)d*|d+|d,|d-|d.d.d|d|d/|d0d St%d1&t'd)2z Perform preprocessing and create a final data batch: 1. Sort batch instances by text-length 2. Convert Audio signal to features. 3. PAD sequences wrt r. 4. Load to Torch. rcSg|]}t|dqSryr)rdrrrrz)TTSDataset.collate_fn..cs i|] fddDqS)csg|]}|qSrr)rdickrrrrz4TTSDataset.collate_fn...r)rrrr s z)TTSDataset.collate_fn..Ncg|]}j|qSr)r>)rlnr_rrrrr~csg|] }j|dqS) embedding)r=rwr_rrrcrr)r<)rsnr_rrrrrcg|] }j|dqSr)r-melspectrogramastyperr_rrrrcSg|]}|jdqSr)rrmrrrrrcsFg|]}|jdjr|jdj|jdjn|jdqSr)rr+rr_rrrs "cSs&g|]}tdg|ddgqS)r)g?)rr)rmel_lenrrrrs&ryr)crr)r- spectrogramrrr_rrrrcSrrrrrrrrredge)moderz[!] z vs rrcsg|] }d|jqS)r)Trrrrrrz[!] Negative padding - z and token_idtoken_id_lengths speaker_nameslinearmel mel_lengths stop_targets item_idxsr d_vectors speaker_idsattnsrxr language_idsaudio_unique_nameszUbatch must contain tensors, numbers, dicts or lists; found {})( isinstance collectionsabcMappingrrrr>r=rr<rr+rrrr transposerr FloatTensor contiguousr,rr4rr- hop_lengthzerosr enumeratepad from_numpy transpose_r1r unsqueeze TypeErrorrtype)rPrtoken_ids_lengthsrrembedding_keysrrrrmel_lengths_adjustedrryr wav_padded wav_lengths max_wav_lenrr mel_lengthrrrrgrpad2pad1r)rrPr collate_fns                &  ." ."(        zTTSDataset.collate_fnr)#__name__ __module__ __qualname__floatintboolr rrstrrDpropertyr\r.setterrdrhrOrvr{rr staticmethodrrrfrrrrrrrr __classcell__rrrQrr(-s          1    3  r(c@seZdZdZ ddeeeeefdddefddZd d Z d d Z d dZ ddZ dddZ ddZddeddfddZdS)rIaPhoneme Dataset for converting input text to phonemes and then token IDs At initialization, it pre-computes the phonemes under `cache_path` and loads them in training to reduce data loading latency. If `cache_path` is already present, it skips the pre-computation. Args: samples (Union[List[List], List[Dict]]): List of samples. Each sample is a list or a dict. tokenizer (TTSTokenizer): Tokenizer to convert input text to phonemes. cache_path (str): Path to cache phonemes. If `cache_path` is already present or None, it skips the pre-computation. precompute_num_workers (int): Number of workers used for pre-computing the phonemes. Defaults to 0. rr.r/r0rBcCsF||_||_||_|durtj|s!t|||dSdSdSr^)r.r/rBrUrVexistsmakedirs precompute)rPr.r/rBr;rrrrDJs zPhonemeDataset.__init__cCsJ|j|}|t|d|d|d}|j|}|d||t|dS)Nr~rr)rph_hatry token_ids_len)r.compute_or_loadr'r/ ids_to_textr)rPindexridsr+rrrrhXs  zPhonemeDataset.__getitem__cCrbr^rcr_rrrrd^rezPhonemeDataset.__len__cCsZd}tj|j||}zt|}W|Sty,|jj||d}t ||Y|Sw)zpCompute phonemes for the given text. If the phonemes are already cached, load them from cache. z _phoneme.npy)r) rUrVjoinrBrrFileNotFoundErrorr/rsave)rP file_namerrfile_extrBr0rrrr-as  zPhonemeDataset.compute_or_loadcCs|jjS)z%Get pad token ID for sequence padding)r/pad_idr_rrr get_pad_idoszPhonemeDataset.get_pad_idr)cCs|tdtjt|d(}|dkr|nd}tjjj||d||jd}|D]}||q$WddS1s7wYdS)zePrecompute phonemes for all samples. We use pytorch dataloader because we are lazy. z[*] Pre-computing phonemes...totalrr)F batch_sizedatasetr num_workersrN) rqtqdmrrutilsdata DataLoaderrupdate)rPr=pbarr; dataloderrZrrrr*ss  "zPhonemeDataset.precomputec Csdd|D}dd|D}dd|D}dd|D}t|}tt|||}t|D]\}} t||||d| f<q1|||dS)NcSrrrrrrrrrrz-PhonemeDataset.collate_fn..cSr)r,rrErrrrrcSr)rrrErrrrrcSr)r+rrErrrrr)rr+ryrrrrfill_r7r ) rPrr0ids_lenstexts texts_hat ids_lens_max ids_torchrids_lenrrrrs zPhonemeDataset.collate_fnrirjNcCrk)Nrlrmz> PhonemeDataset rnr)rorprrrrrrOrtzPhonemeDataset.print_logsrr)rrr__doc__rrrr#rDrhrdr-r7r*rr!rOrrrrrI6s"   rIc@seZdZdZ     d(deeeeefddd efd d Zd d Z ddZ d)ddZ ddZ e ddZe d*ddZe ddZddZddZddZd d!Zd"d#Zd)d$ed%dfd&d'ZdS)+rKaTF0 Dataset for computing F0 from wav files in CPU Pre-compute F0 values for all the samples at initialization if `cache_path` is not None or already present. It also computes the mean and std of F0 values if `normalize_f0` is True. Args: samples (Union[List[List], List[Dict]]): List of samples. Each sample is a list or a dict. ap (AudioProcessor): AudioProcessor to compute F0 from wav files. cache_path (str): Path to cache F0 values. If `cache_path` is already present or None, it skips the pre-computation. Defaults to None. precompute_num_workers (int): Number of workers used for pre-computing the F0 values. Defaults to 0. normalize_f0 (bool): Whether to normalize F0 values by mean and std. Defaults to True. NFrTr.r-r rBcCsn||_||_||_||_||_d|_d|_d|_|dur,tj |s,t || ||r5| |dSdSNr)r.r-rArB normalize_f0r6rstdrUrVr(r)r* load_stats)rPr.r- audio_configrArBr;rPrrrrDs   zF0Dataset.__init__cC\|j|}||dt|d}|jr'|jdur|jdus"Jd||}|d|dS)Nrr~" [!] Mean and STD is not available)r~r)r.r-r'rPrrQ normalize)rPrgrrrrrrh  zF0Dataset.__getitem__cCrbr^rcr_rrrrdrezF0Dataset.__len__c Ctdtjt|d@}|dkr|nd}|j}d|_tjjj||d||jd}g}|D]}|d}| dd |D| |q,||_Wdn1sOwY|jrzd d |D}| |\} } | | d } t j tj|jd | dddSdS)Nz[*] Pre-computing F0s...r8rr)Fr:rcs|]}|VqdSr^r)rfrrr z'F0Dataset.precompute..cSg|] }|D]}|qqSrrrrtensorrrrrrz(F0Dataset.precompute..rrQ pitch_statsT allow_pickle)rqr>rrPrr?r@rArrXrBcompute_pitch_statsrr3rUrVr1rB) rPr=rCr;rPrD computed_datarr pitch_mean pitch_stdrarrrr**    zF0Dataset.precomputecCr]r^r6r_rrrr7zF0Dataset.get_pad_idcCstj||d}|S)Nz _pitch.npy)rUrVr1)r4rB pitch_filerrrcreate_pitch_file_pathsz F0Dataset.create_pitch_file_pathcCs(||}||}|rt|||Sr^)rvr1rr3)r-rrkrrrrr_compute_and_save_pitchs   z!F0Dataset._compute_and_save_pitchcC2tdd|D}t|t|}}||fS)NcS"g|] }|t|dkdqSrrrwherervrrrr"z1F0Dataset.compute_pitch_stats..r concatenaterrQ) pitch_vecsnonzerosrrQrrrrdzF0Dataset.compute_pitch_statscCHtj|d}tj|dd}|dtj|_|dtj|_ dS)Nzpitch_stats.npyTrbrrQ rUrVr1rrrrrrrQrPrB stats_pathstatsrrrrRzF0Dataset.load_statscC2t|dkd}||j}||j}d||<|SNrrrrrrrQrPr zero_idxsrrrrV   zF0Dataset.normalizecC2t|dkd}||j9}||j7}d||<|SrrrrrQrrrrr denormalize rzF0Dataset.denormalizecCB|||j}tj|s||j||}nt|}| tj S)zH compute pitch and return a numpy array of pitch values ) rlrBrUrVr(rmr-rrrr)rPrr~rkrrrrr-    zF0Dataset.compute_or_loadc Cdd|D}dd|D}dd|D}t|}tt|||}t|D]\}}t||||d|f<q*|||dS)NcSrr~rrErrrrrz(F0Dataset.collate_fn..cSrrrrErrrrrcSrrrrErrrrr)r~rf0_lensrF) rPrr~f0sr f0_lens_max f0s_torchrf0_lenrrrr zF0Dataset.collate_fnrirjcC:d|}tdt|dt|dt|jdS)Nrlrmz > F0Dataset rorqrr.rrrrrrO&zF0Dataset.print_logs)NFNrTrr^)rrrrNrrrr#rDrhrdr*r7r&rlrmrdrRrVrr-rr!rOrrrrrKs<       rKc@seZdZdZ    d(deeeeefddd efd d Zd d Z ddZ d)ddZ ddZ e ddZe d*ddZe ddZddZddZddZd d!Zd"d#Zd)d$ed%dfd&d'ZdS)+rMa|Energy Dataset for computing Energy from wav files in CPU Pre-compute Energy values for all the samples at initialization if `cache_path` is not None or already present. It also computes the mean and std of Energy values if `normalize_Energy` is True. Args: samples (Union[List[List], List[Dict]]): List of samples. Each sample is a list or a dict. ap (AudioProcessor): AudioProcessor to compute Energy from wav files. cache_path (str): Path to cache Energy values. If `cache_path` is already present or None, it skips the pre-computation. Defaults to None. precompute_num_workers (int): Number of workers used for pre-computing the Energy values. Defaults to 0. normalize_Energy (bool): Whether to normalize Energy values by mean and std. Defaults to True. FNrTr.r-r rBcCsn||_||_||_||_||_d|_d|_d|_|dur,tj |s,t || ||r5| |dSdSrO)r.r-rArBnormalize_energyr6rrQrUrVr(r)r*rR)rPr.r-rArBr;rrrrrDEs   zEnergyDataset.__init__cCrT)Nrr~rU)r~r)r.r-r'rrrQrV)rPrgrrrrrrh\rWzEnergyDataset.__getitem__cCrbr^rcr_rrrrddrezEnergyDataset.__len__c CrX)Nz[*] Pre-computing energys...r8rr)Fr:rcsrYr^r)rerrrr[tr\z+EnergyDataset.precompute..cSr]rrr^rrrryrz,EnergyDataset.precompute..r` energy_statsTrb)rqr>rrrr?r@rArrXrBcompute_energy_statsrr3rUrVr1rB) rPr=rCr;rrDrerr energy_mean energy_stdrrrrr*grhzEnergyDataset.precomputecCr]r^rir_rrrr7~rjzEnergyDataset.get_pad_idcCs.tjtj|d}tj||d}|S)Nrz _energy.npy)rUrVsplitextrr1)rrBr4 energy_filerrrcreate_energy_file_pathsz%EnergyDataset.create_energy_file_pathcCs4||}t||j|j|jd}|rt|||S)N)fft_sizer  win_length)rvcalculate_energyrr rrr3)r-rrrrrrr_compute_and_save_energys  z&EnergyDataset._compute_and_save_energycCrn)NcSrorprqrsrrrrruz6EnergyDataset.compute_energy_stats..rv) energy_vecsryrrQrrrrrzz"EnergyDataset.compute_energy_statscCr{)Nzenergy_stats.npyTrbrrQr|r}rrrrRrzEnergyDataset.load_statscCrrrrPrrrrrrVrzEnergyDataset.normalizecCrrrrrrrrrzEnergyDataset.denormalizecCr)zJ compute energy and return a numpy array of energy values ) rrBrUrVr(rr-rrrr)rPrr~rrrrrr-rzEnergyDataset.compute_or_loadc Cr)NcSrrrrErrrrrz,EnergyDataset.collate_fn..cSrrrrErrrrrcSrrrrErrrrr)r~r energy_lensrF) rPrr~energysrenergy_lens_max energys_torchr energy_lenrrrrrzEnergyDataset.collate_fnrirjcCr)Nrlrmz> energyDataset rorrrrrrrOrzEnergyDataset.print_logs)FNrTrr^)rrrrNrrrr#rDrhrdr*r7r&rrrrRrVrr-rr!rOrrrrrM-s:       rM) r!rrUrtypingrrrnumpyrrr>torch.utils.datarTTS.tts.utils.datarrrTTS.utils.audior TTS.utils.audio.numpy_transformsr rmultiprocessingset_sharing_strategyrrr'r(rIrKrMrrrrs2     _