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")

@click.command()

@click.option('-p', '--config_path', default='Configs/config.yml', type=str)

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

from Utils.PLBERT_all_languages.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()