Spaces:
Runtime error
Runtime error
| import os | |
| import os.path as osp | |
| import re | |
| import sys | |
| import yaml | |
| import shutil | |
| import numpy as np | |
| import torch | |
| import click | |
| import warnings | |
| warnings.simplefilter("ignore") | |
| # load packages | |
| import random | |
| import yaml | |
| from munch import Munch | |
| import numpy as np | |
| import torch | |
| from torch import nn | |
| import torch.nn.functional as F | |
| import torchaudio | |
| import librosa | |
| from models import * | |
| from meldataset import build_dataloader | |
| from utils import * | |
| from losses import * | |
| from optimizers import build_optimizer | |
| import time | |
| from accelerate import Accelerator | |
| from accelerate.utils import LoggerType | |
| from accelerate import DistributedDataParallelKwargs | |
| from torch.utils.tensorboard import SummaryWriter | |
| import logging | |
| from accelerate.logging import get_logger | |
| logger = get_logger(__name__, log_level="DEBUG") | |
| def main(config_path): | |
| config = yaml.safe_load(open(config_path)) | |
| log_dir = config["log_dir"] | |
| if not osp.exists(log_dir): | |
| os.makedirs(log_dir, exist_ok=True) | |
| shutil.copy(config_path, osp.join(log_dir, osp.basename(config_path))) | |
| ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True) | |
| accelerator = Accelerator( | |
| project_dir=log_dir, split_batches=True, kwargs_handlers=[ddp_kwargs] | |
| ) | |
| if accelerator.is_main_process: | |
| writer = SummaryWriter(log_dir + "/tensorboard") | |
| # write logs | |
| file_handler = logging.FileHandler(osp.join(log_dir, "train.log")) | |
| file_handler.setLevel(logging.DEBUG) | |
| file_handler.setFormatter( | |
| logging.Formatter("%(levelname)s:%(asctime)s: %(message)s") | |
| ) | |
| logger.logger.addHandler(file_handler) | |
| batch_size = config.get("batch_size", 10) | |
| device = accelerator.device | |
| epochs = config.get("epochs_1st", 200) | |
| save_freq = config.get("save_freq", 2) | |
| log_interval = config.get("log_interval", 10) | |
| saving_epoch = config.get("save_freq", 2) | |
| data_params = config.get("data_params", None) | |
| sr = config["preprocess_params"].get("sr", 24000) | |
| train_path = data_params["train_data"] | |
| val_path = data_params["val_data"] | |
| root_path = data_params["root_path"] | |
| min_length = data_params["min_length"] | |
| OOD_data = data_params["OOD_data"] | |
| max_len = config.get("max_len", 200) | |
| # load data | |
| train_list, val_list = get_data_path_list(train_path, val_path) | |
| train_dataloader = build_dataloader( | |
| train_list, | |
| root_path, | |
| OOD_data=OOD_data, | |
| min_length=min_length, | |
| batch_size=batch_size, | |
| num_workers=2, | |
| dataset_config={}, | |
| device=device, | |
| ) | |
| val_dataloader = build_dataloader( | |
| val_list, | |
| root_path, | |
| OOD_data=OOD_data, | |
| min_length=min_length, | |
| batch_size=batch_size, | |
| validation=True, | |
| num_workers=0, | |
| device=device, | |
| dataset_config={}, | |
| ) | |
| with accelerator.main_process_first(): | |
| # load pretrained ASR model | |
| ASR_config = config.get("ASR_config", False) | |
| ASR_path = config.get("ASR_path", False) | |
| text_aligner = load_ASR_models(ASR_path, ASR_config) | |
| # load pretrained F0 model | |
| F0_path = config.get("F0_path", False) | |
| pitch_extractor = load_F0_models(F0_path) | |
| # load BERT model | |
| from Utils.PLBERT.util import load_plbert | |
| BERT_path = config.get("PLBERT_dir", False) | |
| plbert = load_plbert(BERT_path) | |
| scheduler_params = { | |
| "max_lr": float(config["optimizer_params"].get("lr", 1e-4)), | |
| "pct_start": float(config["optimizer_params"].get("pct_start", 0.0)), | |
| "epochs": epochs, | |
| "steps_per_epoch": len(train_dataloader), | |
| } | |
| model_params = recursive_munch(config["model_params"]) | |
| multispeaker = model_params.multispeaker | |
| model = build_model(model_params, text_aligner, pitch_extractor, plbert) | |
| best_loss = float("inf") # best test loss | |
| loss_train_record = list([]) | |
| loss_test_record = list([]) | |
| loss_params = Munch(config["loss_params"]) | |
| TMA_epoch = loss_params.TMA_epoch | |
| for k in model: | |
| model[k] = accelerator.prepare(model[k]) | |
| train_dataloader, val_dataloader = accelerator.prepare( | |
| train_dataloader, val_dataloader | |
| ) | |
| _ = [model[key].to(device) for key in model] | |
| # initialize optimizers after preparing models for compatibility with FSDP | |
| optimizer = build_optimizer( | |
| {key: model[key].parameters() for key in model}, | |
| scheduler_params_dict={key: scheduler_params.copy() for key in model}, | |
| lr=float(config["optimizer_params"].get("lr", 1e-4)), | |
| ) | |
| for k, v in optimizer.optimizers.items(): | |
| optimizer.optimizers[k] = accelerator.prepare(optimizer.optimizers[k]) | |
| optimizer.schedulers[k] = accelerator.prepare(optimizer.schedulers[k]) | |
| with accelerator.main_process_first(): | |
| if config.get("pretrained_model", "") != "": | |
| model, optimizer, start_epoch, iters = load_checkpoint( | |
| model, | |
| optimizer, | |
| config["pretrained_model"], | |
| load_only_params=config.get("load_only_params", True), | |
| ) | |
| else: | |
| start_epoch = 0 | |
| iters = 0 | |
| # in case not distributed | |
| try: | |
| n_down = model.text_aligner.module.n_down | |
| except: | |
| n_down = model.text_aligner.n_down | |
| # wrapped losses for compatibility with mixed precision | |
| stft_loss = MultiResolutionSTFTLoss().to(device) | |
| gl = GeneratorLoss(model.mpd, model.msd).to(device) | |
| dl = DiscriminatorLoss(model.mpd, model.msd).to(device) | |
| wl = WavLMLoss(model_params.slm.model, model.wd, sr, model_params.slm.sr).to(device) | |
| for epoch in range(start_epoch, epochs): | |
| running_loss = 0 | |
| start_time = time.time() | |
| _ = [model[key].train() for key in model] | |
| for i, batch in enumerate(train_dataloader): | |
| waves = batch[0] | |
| batch = [b.to(device) for b in batch[1:]] | |
| texts, input_lengths, _, _, mels, mel_input_length, _ = batch | |
| with torch.no_grad(): | |
| mask = length_to_mask(mel_input_length // (2**n_down)).to("cuda") | |
| text_mask = length_to_mask(input_lengths).to(texts.device) | |
| ppgs, s2s_pred, s2s_attn = model.text_aligner(mels, mask, texts) | |
| s2s_attn = s2s_attn.transpose(-1, -2) | |
| s2s_attn = s2s_attn[..., 1:] | |
| s2s_attn = s2s_attn.transpose(-1, -2) | |
| with torch.no_grad(): | |
| attn_mask = ( | |
| (~mask) | |
| .unsqueeze(-1) | |
| .expand(mask.shape[0], mask.shape[1], text_mask.shape[-1]) | |
| .float() | |
| .transpose(-1, -2) | |
| ) | |
| attn_mask = ( | |
| attn_mask.float() | |
| * (~text_mask) | |
| .unsqueeze(-1) | |
| .expand(text_mask.shape[0], text_mask.shape[1], mask.shape[-1]) | |
| .float() | |
| ) | |
| attn_mask = attn_mask < 1 | |
| s2s_attn.masked_fill_(attn_mask, 0.0) | |
| with torch.no_grad(): | |
| mask_ST = mask_from_lens( | |
| s2s_attn, input_lengths, mel_input_length // (2**n_down) | |
| ) | |
| s2s_attn_mono = maximum_path(s2s_attn, mask_ST) | |
| # encode | |
| t_en = model.text_encoder(texts, input_lengths, text_mask) | |
| # 50% of chance of using monotonic version | |
| if bool(random.getrandbits(1)): | |
| asr = t_en @ s2s_attn | |
| else: | |
| asr = t_en @ s2s_attn_mono | |
| # get clips | |
| mel_input_length_all = accelerator.gather( | |
| mel_input_length | |
| ) # for balanced load | |
| mel_len = min( | |
| [int(mel_input_length_all.min().item() / 2 - 1), max_len // 2] | |
| ) | |
| mel_len_st = int(mel_input_length.min().item() / 2 - 1) | |
| en = [] | |
| gt = [] | |
| wav = [] | |
| st = [] | |
| for bib in range(len(mel_input_length)): | |
| mel_length = int(mel_input_length[bib].item() / 2) | |
| random_start = np.random.randint(0, mel_length - mel_len) | |
| en.append(asr[bib, :, random_start : random_start + mel_len]) | |
| gt.append( | |
| mels[bib, :, (random_start * 2) : ((random_start + mel_len) * 2)] | |
| ) | |
| y = waves[bib][ | |
| (random_start * 2) * 300 : ((random_start + mel_len) * 2) * 300 | |
| ] | |
| wav.append(torch.from_numpy(y).to(device)) | |
| # style reference (better to be different from the GT) | |
| random_start = np.random.randint(0, mel_length - mel_len_st) | |
| st.append( | |
| mels[bib, :, (random_start * 2) : ((random_start + mel_len_st) * 2)] | |
| ) | |
| en = torch.stack(en) | |
| gt = torch.stack(gt).detach() | |
| st = torch.stack(st).detach() | |
| wav = torch.stack(wav).float().detach() | |
| # clip too short to be used by the style encoder | |
| if gt.shape[-1] < 80: | |
| continue | |
| with torch.no_grad(): | |
| real_norm = log_norm(gt.unsqueeze(1)).squeeze(1).detach() | |
| F0_real, _, _ = model.pitch_extractor(gt.unsqueeze(1)) | |
| s = model.style_encoder( | |
| st.unsqueeze(1) if multispeaker else gt.unsqueeze(1) | |
| ) | |
| y_rec = model.decoder(en, F0_real, real_norm, s) | |
| # discriminator loss | |
| if epoch >= TMA_epoch: | |
| optimizer.zero_grad() | |
| d_loss = dl(wav.detach().unsqueeze(1).float(), y_rec.detach()).mean() | |
| accelerator.backward(d_loss) | |
| optimizer.step("msd") | |
| optimizer.step("mpd") | |
| else: | |
| d_loss = 0 | |
| # generator loss | |
| optimizer.zero_grad() | |
| loss_mel = stft_loss(y_rec.squeeze(), wav.detach()) | |
| if epoch >= TMA_epoch: # start TMA training | |
| loss_s2s = 0 | |
| for _s2s_pred, _text_input, _text_length in zip( | |
| s2s_pred, texts, input_lengths | |
| ): | |
| loss_s2s += F.cross_entropy( | |
| _s2s_pred[:_text_length], _text_input[:_text_length] | |
| ) | |
| loss_s2s /= texts.size(0) | |
| loss_mono = F.l1_loss(s2s_attn, s2s_attn_mono) * 10 | |
| loss_gen_all = gl(wav.detach().unsqueeze(1).float(), y_rec).mean() | |
| loss_slm = wl(wav.detach(), y_rec).mean() | |
| g_loss = ( | |
| loss_params.lambda_mel * loss_mel | |
| + loss_params.lambda_mono * loss_mono | |
| + loss_params.lambda_s2s * loss_s2s | |
| + loss_params.lambda_gen * loss_gen_all | |
| + loss_params.lambda_slm * loss_slm | |
| ) | |
| else: | |
| loss_s2s = 0 | |
| loss_mono = 0 | |
| loss_gen_all = 0 | |
| loss_slm = 0 | |
| g_loss = loss_mel | |
| running_loss += accelerator.gather(loss_mel).mean().item() | |
| accelerator.backward(g_loss) | |
| optimizer.step("text_encoder") | |
| optimizer.step("style_encoder") | |
| optimizer.step("decoder") | |
| if epoch >= TMA_epoch: | |
| optimizer.step("text_aligner") | |
| optimizer.step("pitch_extractor") | |
| iters = iters + 1 | |
| if (i + 1) % log_interval == 0 and accelerator.is_main_process: | |
| log_print( | |
| "Epoch [%d/%d], Step [%d/%d], Mel Loss: %.5f, Gen Loss: %.5f, Disc Loss: %.5f, Mono Loss: %.5f, S2S Loss: %.5f, SLM Loss: %.5f" | |
| % ( | |
| epoch + 1, | |
| epochs, | |
| i + 1, | |
| len(train_list) // batch_size, | |
| running_loss / log_interval, | |
| loss_gen_all, | |
| d_loss, | |
| loss_mono, | |
| loss_s2s, | |
| loss_slm, | |
| ), | |
| logger, | |
| ) | |
| writer.add_scalar("train/mel_loss", running_loss / log_interval, iters) | |
| writer.add_scalar("train/gen_loss", loss_gen_all, iters) | |
| writer.add_scalar("train/d_loss", d_loss, iters) | |
| writer.add_scalar("train/mono_loss", loss_mono, iters) | |
| writer.add_scalar("train/s2s_loss", loss_s2s, iters) | |
| writer.add_scalar("train/slm_loss", loss_slm, iters) | |
| running_loss = 0 | |
| print("Time elasped:", time.time() - start_time) | |
| loss_test = 0 | |
| _ = [model[key].eval() for key in model] | |
| with torch.no_grad(): | |
| iters_test = 0 | |
| for batch_idx, batch in enumerate(val_dataloader): | |
| optimizer.zero_grad() | |
| waves = batch[0] | |
| batch = [b.to(device) for b in batch[1:]] | |
| texts, input_lengths, _, _, mels, mel_input_length, _ = batch | |
| with torch.no_grad(): | |
| mask = length_to_mask(mel_input_length // (2**n_down)).to("cuda") | |
| ppgs, s2s_pred, s2s_attn = model.text_aligner(mels, mask, texts) | |
| s2s_attn = s2s_attn.transpose(-1, -2) | |
| s2s_attn = s2s_attn[..., 1:] | |
| s2s_attn = s2s_attn.transpose(-1, -2) | |
| text_mask = length_to_mask(input_lengths).to(texts.device) | |
| attn_mask = ( | |
| (~mask) | |
| .unsqueeze(-1) | |
| .expand(mask.shape[0], mask.shape[1], text_mask.shape[-1]) | |
| .float() | |
| .transpose(-1, -2) | |
| ) | |
| attn_mask = ( | |
| attn_mask.float() | |
| * (~text_mask) | |
| .unsqueeze(-1) | |
| .expand(text_mask.shape[0], text_mask.shape[1], mask.shape[-1]) | |
| .float() | |
| ) | |
| attn_mask = attn_mask < 1 | |
| s2s_attn.masked_fill_(attn_mask, 0.0) | |
| # encode | |
| t_en = model.text_encoder(texts, input_lengths, text_mask) | |
| asr = t_en @ s2s_attn | |
| # get clips | |
| mel_input_length_all = accelerator.gather( | |
| mel_input_length | |
| ) # for balanced load | |
| mel_len = min( | |
| [int(mel_input_length.min().item() / 2 - 1), max_len // 2] | |
| ) | |
| en = [] | |
| gt = [] | |
| wav = [] | |
| for bib in range(len(mel_input_length)): | |
| mel_length = int(mel_input_length[bib].item() / 2) | |
| random_start = np.random.randint(0, mel_length - mel_len) | |
| en.append(asr[bib, :, random_start : random_start + mel_len]) | |
| gt.append( | |
| mels[ | |
| bib, :, (random_start * 2) : ((random_start + mel_len) * 2) | |
| ] | |
| ) | |
| y = waves[bib][ | |
| (random_start * 2) * 300 : ((random_start + mel_len) * 2) * 300 | |
| ] | |
| wav.append(torch.from_numpy(y).to("cuda")) | |
| wav = torch.stack(wav).float().detach() | |
| en = torch.stack(en) | |
| gt = torch.stack(gt).detach() | |
| F0_real, _, F0 = model.pitch_extractor(gt.unsqueeze(1)) | |
| s = model.style_encoder(gt.unsqueeze(1)) | |
| real_norm = log_norm(gt.unsqueeze(1)).squeeze(1) | |
| y_rec = model.decoder(en, F0_real, real_norm, s) | |
| loss_mel = stft_loss(y_rec.squeeze(), wav.detach()) | |
| loss_test += accelerator.gather(loss_mel).mean().item() | |
| iters_test += 1 | |
| if accelerator.is_main_process: | |
| print("Epochs:", epoch + 1) | |
| log_print( | |
| "Validation loss: %.3f" % (loss_test / iters_test) + "\n\n\n\n", logger | |
| ) | |
| print("\n\n\n") | |
| writer.add_scalar("eval/mel_loss", loss_test / iters_test, epoch + 1) | |
| attn_image = get_image(s2s_attn[0].cpu().numpy().squeeze()) | |
| writer.add_figure("eval/attn", attn_image, epoch) | |
| with torch.no_grad(): | |
| for bib in range(len(asr)): | |
| mel_length = int(mel_input_length[bib].item()) | |
| gt = mels[bib, :, :mel_length].unsqueeze(0) | |
| en = asr[bib, :, : mel_length // 2].unsqueeze(0) | |
| F0_real, _, _ = model.pitch_extractor(gt.unsqueeze(1)) | |
| F0_real = F0_real.unsqueeze(0) | |
| s = model.style_encoder(gt.unsqueeze(1)) | |
| real_norm = log_norm(gt.unsqueeze(1)).squeeze(1) | |
| y_rec = model.decoder(en, F0_real, real_norm, s) | |
| writer.add_audio( | |
| "eval/y" + str(bib), | |
| y_rec.cpu().numpy().squeeze(), | |
| epoch, | |
| sample_rate=sr, | |
| ) | |
| if epoch == 0: | |
| writer.add_audio( | |
| "gt/y" + str(bib), | |
| waves[bib].squeeze(), | |
| epoch, | |
| sample_rate=sr, | |
| ) | |
| if bib >= 6: | |
| break | |
| if epoch % saving_epoch == 0: | |
| if (loss_test / iters_test) < best_loss: | |
| best_loss = loss_test / iters_test | |
| print("Saving..") | |
| state = { | |
| "net": {key: model[key].state_dict() for key in model}, | |
| "optimizer": optimizer.state_dict(), | |
| "iters": iters, | |
| "val_loss": loss_test / iters_test, | |
| "epoch": epoch, | |
| } | |
| save_path = osp.join(log_dir, "epoch_1st_%05d.pth" % epoch) | |
| torch.save(state, save_path) | |
| if accelerator.is_main_process: | |
| print("Saving..") | |
| state = { | |
| "net": {key: model[key].state_dict() for key in model}, | |
| "optimizer": optimizer.state_dict(), | |
| "iters": iters, | |
| "val_loss": loss_test / iters_test, | |
| "epoch": epoch, | |
| } | |
| save_path = osp.join(log_dir, config.get("first_stage_path", "first_stage.pth")) | |
| torch.save(state, save_path) | |
| if __name__ == "__main__": | |
| main() | |