Spaces:
Running
on
Zero
Running
on
Zero
| import torch.multiprocessing | |
| from Modules.Aligner.CodecAlignerDataset import CodecAlignerDataset | |
| from Modules.Aligner.autoaligner_train_loop import train_loop as train_aligner | |
| from Modules.ToucanTTS.TTSDataset import TTSDataset | |
| from Utility.path_to_transcript_dicts import * | |
| from Utility.storage_config import MODELS_DIR | |
| def prepare_aligner_corpus(transcript_dict, corpus_dir, lang, device, phone_input=False, | |
| gpu_count=1, | |
| rank=0): | |
| return CodecAlignerDataset(transcript_dict, | |
| cache_dir=corpus_dir, | |
| lang=lang, | |
| loading_processes=5, # this can be increased for massive clusters, but the overheads that are introduced are kind of not really worth it | |
| device=device, | |
| phone_input=phone_input, | |
| gpu_count=gpu_count, | |
| rank=rank) | |
| def prepare_tts_corpus(transcript_dict, | |
| corpus_dir, | |
| lang, | |
| # For small datasets it's best to turn this off and instead inspect the data with the scorer, if there are any issues. | |
| fine_tune_aligner=True, | |
| use_reconstruction=True, | |
| phone_input=False, | |
| save_imgs=False, | |
| gpu_count=1, | |
| rank=0): | |
| """ | |
| create an aligner dataset, | |
| fine-tune an aligner, | |
| create a TTS dataset, | |
| return it. | |
| Automatically skips parts that have been done before. | |
| """ | |
| if not os.path.exists(os.path.join(corpus_dir, "tts_train_cache.pt")): | |
| if fine_tune_aligner: | |
| aligner_dir = os.path.join(corpus_dir, "Aligner") | |
| aligner_loc = os.path.join(corpus_dir, "Aligner", "aligner.pt") | |
| if not os.path.exists(os.path.join(corpus_dir, "aligner_train_cache.pt")): | |
| prepare_aligner_corpus(transcript_dict, corpus_dir=corpus_dir, lang=lang, phone_input=phone_input, device=torch.device("cuda")) | |
| if not os.path.exists(os.path.join(aligner_dir, "aligner.pt")): | |
| aligner_datapoints = prepare_aligner_corpus(transcript_dict, corpus_dir=corpus_dir, lang=lang, phone_input=phone_input, device=torch.device("cuda")) | |
| if os.path.exists(os.path.join(MODELS_DIR, "Aligner", "aligner.pt")): | |
| train_aligner(train_dataset=aligner_datapoints, | |
| device=torch.device("cuda"), | |
| save_directory=aligner_dir, | |
| steps=min(len(aligner_datapoints) // 2, 10000), # relatively good finetuning heuristic | |
| batch_size=32 if len(aligner_datapoints) > 32 else len(aligner_datapoints) // 2, | |
| path_to_checkpoint=os.path.join(MODELS_DIR, "Aligner", "aligner.pt"), | |
| fine_tune=True, | |
| debug_img_path=aligner_dir, | |
| resume=False, | |
| use_reconstruction=use_reconstruction) | |
| else: | |
| train_aligner(train_dataset=aligner_datapoints, | |
| device=torch.device("cuda"), | |
| save_directory=aligner_dir, | |
| steps=len(aligner_datapoints) // 2, # relatively good heuristic | |
| batch_size=32 if len(aligner_datapoints) > 32 else len(aligner_datapoints) // 2, | |
| path_to_checkpoint=None, | |
| fine_tune=False, | |
| debug_img_path=aligner_dir, | |
| resume=False, | |
| use_reconstruction=use_reconstruction) | |
| else: | |
| aligner_loc = os.path.join(MODELS_DIR, "Aligner", "aligner.pt") | |
| else: | |
| aligner_loc = None | |
| return TTSDataset(transcript_dict, | |
| acoustic_checkpoint_path=aligner_loc, | |
| cache_dir=corpus_dir, | |
| device=torch.device("cuda"), | |
| lang=lang, | |
| save_imgs=save_imgs, | |
| gpu_count=gpu_count, | |
| rank=rank) | |