Spaces:
Running on Zero
Running on Zero
| import random | |
| import torch | |
| import torch.multiprocessing | |
| from torch.utils.data import ConcatDataset | |
| from tqdm import tqdm | |
| from TrainingInterfaces.Text_to_Spectrogram.FastSpeech2.FastSpeech2 import FastSpeech2 | |
| from TrainingInterfaces.Text_to_Spectrogram.FastSpeech2.meta_train_loop import train_loop | |
| from Utility.corpus_preparation import prepare_fastspeech_corpus | |
| from Utility.path_to_transcript_dicts import * | |
| def run(gpu_id, resume_checkpoint, finetune, model_dir, resume, remove_faulty_samples=False): | |
| # It is not recommended training this yourself or to finetune this, but you can. | |
| # The recommended use is to download the pretrained model from the github release | |
| # page and finetune to your desired data similar to how it is showcased in | |
| # FastSpeech2_finetune_to_German.py | |
| torch.manual_seed(131714) | |
| random.seed(131714) | |
| torch.random.manual_seed(131714) | |
| os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" | |
| os.environ["CUDA_VISIBLE_DEVICES"] = f"{gpu_id}" | |
| datasets = list() | |
| base_dir = os.path.join("Models", "FastSpeech2_Meta") | |
| if model_dir is not None: | |
| meta_save_dir = model_dir | |
| else: | |
| meta_save_dir = base_dir | |
| os.makedirs(meta_save_dir, exist_ok=True) | |
| print("Preparing") | |
| english_datasets = list() | |
| german_datasets = list() | |
| greek_datasets = list() | |
| spanish_datasets = list() | |
| finnish_datasets = list() | |
| russian_datasets = list() | |
| hungarian_datasets = list() | |
| dutch_datasets = list() | |
| french_datasets = list() | |
| portuguese_datasets = list() | |
| polish_datasets = list() | |
| italian_datasets = list() | |
| chinese_datasets = list() | |
| vietnamese_datasets = list() | |
| english_datasets.append(prepare_fastspeech_corpus(transcript_dict=build_path_to_transcript_dict_nancy(), | |
| corpus_dir=os.path.join("Corpora", "Nancy"), | |
| lang="en")) | |
| english_datasets.append(prepare_fastspeech_corpus(transcript_dict=build_path_to_transcript_dict_ljspeech(), | |
| corpus_dir=os.path.join("Corpora", "LJSpeech"), | |
| lang="en")) | |
| english_datasets.append(prepare_fastspeech_corpus(transcript_dict=build_path_to_transcript_dict_libritts_all_clean(), | |
| corpus_dir=os.path.join("Corpora", "libri_all_clean"), | |
| lang="en")) | |
| english_datasets.append(prepare_fastspeech_corpus(transcript_dict=build_path_to_transcript_dict_vctk(), | |
| corpus_dir=os.path.join("Corpora", "vctk"), | |
| lang="en")) | |
| english_datasets.append(prepare_fastspeech_corpus(transcript_dict=build_path_to_transcript_dict_nvidia_hifitts(), | |
| corpus_dir=os.path.join("Corpora", "hifi"), | |
| lang="en")) | |
| german_datasets.append(prepare_fastspeech_corpus(transcript_dict=build_path_to_transcript_dict_karlsson(), | |
| corpus_dir=os.path.join("Corpora", "Karlsson"), | |
| lang="de")) | |
| german_datasets.append(prepare_fastspeech_corpus(transcript_dict=build_path_to_transcript_dict_eva(), | |
| corpus_dir=os.path.join("Corpora", "Eva"), | |
| lang="de")) | |
| german_datasets.append(prepare_fastspeech_corpus(transcript_dict=build_path_to_transcript_dict_hokus(), | |
| corpus_dir=os.path.join("Corpora", "Hokus"), | |
| lang="de")) | |
| german_datasets.append(prepare_fastspeech_corpus(transcript_dict=build_path_to_transcript_dict_bernd(), | |
| corpus_dir=os.path.join("Corpora", "Bernd"), | |
| lang="de")) | |
| german_datasets.append(prepare_fastspeech_corpus(transcript_dict=build_path_to_transcript_dict_hui_others(), | |
| corpus_dir=os.path.join("Corpora", "hui_others"), | |
| lang="de")) | |
| german_datasets.append(prepare_fastspeech_corpus(transcript_dict=build_path_to_transcript_dict_thorsten(), | |
| corpus_dir=os.path.join("Corpora", "Thorsten"), | |
| lang="de")) | |
| greek_datasets.append(prepare_fastspeech_corpus(transcript_dict=build_path_to_transcript_dict_css10el(), | |
| corpus_dir=os.path.join("Corpora", "meta_Greek"), | |
| lang="el")) | |
| spanish_datasets.append(prepare_fastspeech_corpus(transcript_dict=build_path_to_transcript_dict_spanish_blizzard_train(), | |
| corpus_dir=os.path.join("Corpora", "spanish_blizzard"), | |
| lang="es")) | |
| spanish_datasets.append(prepare_fastspeech_corpus(transcript_dict=build_path_to_transcript_dict_css10es(), | |
| corpus_dir=os.path.join("Corpora", "meta_Spanish"), | |
| lang="es")) | |
| spanish_datasets.append(prepare_fastspeech_corpus(transcript_dict=build_path_to_transcript_dict_mls_spanish(), | |
| corpus_dir=os.path.join("Corpora", "mls_spanish"), | |
| lang="es")) | |
| finnish_datasets.append(prepare_fastspeech_corpus(transcript_dict=build_path_to_transcript_dict_css10fi(), | |
| corpus_dir=os.path.join("Corpora", "meta_Finnish"), | |
| lang="fi")) | |
| russian_datasets.append(prepare_fastspeech_corpus(transcript_dict=build_path_to_transcript_dict_css10ru(), | |
| corpus_dir=os.path.join("Corpora", "meta_Russian"), | |
| lang="ru")) | |
| hungarian_datasets.append(prepare_fastspeech_corpus(transcript_dict=build_path_to_transcript_dict_css10hu(), | |
| corpus_dir=os.path.join("Corpora", "meta_Hungarian"), | |
| lang="hu")) | |
| dutch_datasets.append(prepare_fastspeech_corpus(transcript_dict=build_path_to_transcript_dict_css10nl(), | |
| corpus_dir=os.path.join("Corpora", "meta_Dutch"), | |
| lang="nl")) | |
| dutch_datasets.append(prepare_fastspeech_corpus(transcript_dict=build_path_to_transcript_dict_mls_dutch(), | |
| corpus_dir=os.path.join("Corpora", "mls_dutch"), | |
| lang="nl")) | |
| french_datasets.append(prepare_fastspeech_corpus(transcript_dict=build_path_to_transcript_dict_css10fr(), | |
| corpus_dir=os.path.join("Corpora", "meta_French"), | |
| lang="fr")) | |
| french_datasets.append(prepare_fastspeech_corpus(transcript_dict=build_path_to_transcript_dict_mls_french(), | |
| corpus_dir=os.path.join("Corpora", "mls_french"), | |
| lang="fr")) | |
| portuguese_datasets.append(prepare_fastspeech_corpus(transcript_dict=build_path_to_transcript_dict_mls_portuguese(), | |
| corpus_dir=os.path.join("Corpora", "mls_porto"), | |
| lang="pt")) | |
| polish_datasets.append(prepare_fastspeech_corpus(transcript_dict=build_path_to_transcript_dict_mls_polish(), | |
| corpus_dir=os.path.join("Corpora", "mls_polish"), | |
| lang="pl")) | |
| italian_datasets.append(prepare_fastspeech_corpus(transcript_dict=build_path_to_transcript_dict_mls_italian(), | |
| corpus_dir=os.path.join("Corpora", "mls_italian"), | |
| lang="it")) | |
| chinese_datasets.append(prepare_fastspeech_corpus(transcript_dict=build_path_to_transcript_dict_css10cmn(), | |
| corpus_dir=os.path.join("Corpora", "css10_chinese"), | |
| lang="cmn")) | |
| chinese_datasets.append(prepare_fastspeech_corpus(transcript_dict=build_path_to_transcript_dict_aishell3(), | |
| corpus_dir=os.path.join("Corpora", "aishell3"), | |
| lang="cmn")) | |
| vietnamese_datasets.append(prepare_fastspeech_corpus(transcript_dict=build_path_to_transcript_dict_vietTTS(), | |
| corpus_dir=os.path.join("Corpora", "vietTTS"), | |
| lang="vi")) | |
| datasets.append(ConcatDataset(english_datasets)) | |
| datasets.append(ConcatDataset(german_datasets)) | |
| datasets.append(ConcatDataset(greek_datasets)) | |
| datasets.append(ConcatDataset(spanish_datasets)) | |
| datasets.append(ConcatDataset(finnish_datasets)) | |
| datasets.append(ConcatDataset(russian_datasets)) | |
| datasets.append(ConcatDataset(hungarian_datasets)) | |
| datasets.append(ConcatDataset(dutch_datasets)) | |
| datasets.append(ConcatDataset(french_datasets)) | |
| datasets.append(ConcatDataset(portuguese_datasets)) | |
| datasets.append(ConcatDataset(polish_datasets)) | |
| datasets.append(ConcatDataset(italian_datasets)) | |
| datasets.append(ConcatDataset(chinese_datasets)) | |
| datasets.append(ConcatDataset(vietnamese_datasets)) | |
| if remove_faulty_samples: | |
| find_and_remove_faulty_samples(net=FastSpeech2(lang_embs=100), | |
| datasets=english_datasets + | |
| german_datasets + | |
| greek_datasets + | |
| spanish_datasets + | |
| finnish_datasets + | |
| russian_datasets + | |
| hungarian_datasets + | |
| dutch_datasets + | |
| french_datasets + | |
| portuguese_datasets + | |
| polish_datasets + | |
| italian_datasets + | |
| chinese_datasets + | |
| vietnamese_datasets, | |
| device=torch.device("cuda"), | |
| path_to_checkpoint=resume_checkpoint) | |
| train_loop(net=FastSpeech2(lang_embs=100), | |
| device=torch.device("cuda"), | |
| datasets=datasets, | |
| batch_size=6, | |
| save_directory=meta_save_dir, | |
| steps=300000, | |
| steps_per_checkpoint=1000, | |
| lr=0.001, | |
| path_to_checkpoint=resume_checkpoint, | |
| resume=resume) | |
| def find_and_remove_faulty_samples(net, | |
| datasets, | |
| device, | |
| path_to_checkpoint): | |
| net = net.to(device) | |
| torch.multiprocessing.set_sharing_strategy('file_system') | |
| check_dict = torch.load(os.path.join(path_to_checkpoint), map_location=device) | |
| net.load_state_dict(check_dict["model"]) | |
| for dataset_index in range(len(datasets)): | |
| nan_ids = list() | |
| for datapoint_index in tqdm(range(len(datasets[dataset_index]))): | |
| loss = net(text_tensors=datasets[dataset_index][datapoint_index][0].unsqueeze(0).to(device), | |
| text_lengths=datasets[dataset_index][datapoint_index][1].to(device), | |
| gold_speech=datasets[dataset_index][datapoint_index][2].unsqueeze(0).to(device), | |
| speech_lengths=datasets[dataset_index][datapoint_index][3].to(device), | |
| gold_durations=datasets[dataset_index][datapoint_index][4].unsqueeze(0).to(device), | |
| gold_pitch=datasets[dataset_index][datapoint_index][6].unsqueeze(0).to(device), # mind the switched order | |
| gold_energy=datasets[dataset_index][datapoint_index][5].unsqueeze(0).to(device), # mind the switched order | |
| utterance_embedding=datasets[dataset_index][datapoint_index][7].unsqueeze(0).to(device), | |
| lang_ids=datasets[dataset_index][datapoint_index][8].unsqueeze(0).to(device), | |
| return_mels=False).squeeze() | |
| if torch.isnan(loss): | |
| print(f"NAN DETECTED: {dataset_index}, {datapoint_index}") | |
| nan_ids.append(datapoint_index) | |
| datasets[dataset_index].remove_samples(nan_ids) | |