2

data_preparation.py

@@ -1,108 +0,0 @@
-import kaldiio
-import os
-import librosa
-from tqdm import tqdm
-import glob
-import json 
-from shutil import copyfile
-import pandas as pd
-import argparse
-from text import _clean_text, symbols
-from num2words import num2words
-import re
-from melspec import mel_spectrogram
-import torchaudio
-
-if __name__ == '__main__':
-    parser = argparse.ArgumentParser()
-    parser.add_argument('-d', '--data', type=str, required=True, help='path to the emotional dataset')
-    args = parser.parse_args()
-    dataset_path = args.data
-    filelists_path = 'filelists/all_spks/'
-    feats_scp_file = filelists_path + 'feats.scp'
-    feats_ark_file = filelists_path + 'feats.ark'
-
-
-    spks = ['1263201035', '805570882', '399172782']
-    train_files = []
-    eval_files = []
-    for spk in spks:
-        train_files += glob.glob(dataset_path + spk + "/train/*.wav")
-        eval_files += glob.glob(dataset_path + spk + "/eval/*.wav")
-
-    os.makedirs(filelists_path, exist_ok=True)
-
-    with open(filelists_path + 'train_utts.txt', 'w', encoding='utf-8') as f:
-        for wav_path in train_files:
-            wav_name = os.path.splitext(os.path.basename(wav_path))[0]
-            f.write(wav_name + '\n')
-    with open(filelists_path + 'eval_utts.txt', 'w', encoding='utf-8') as f:
-        for wav_path in eval_files:
-            wav_name = os.path.splitext(os.path.basename(wav_path))[0]
-            f.write(wav_name + '\n')
-
-    with open(feats_scp_file, 'w') as feats_scp, \
-        kaldiio.WriteHelper(f'ark,scp:{feats_ark_file},{feats_scp_file}') as writer:
-        for root, dirs, files in os.walk(dataset_path):
-            for file in tqdm(files):
-                if file.endswith('.wav'):
-                    # Get the file name and relative path to the root folder
-                    wav_path = os.path.join(root, file)
-                    rel_path = os.path.relpath(wav_path, dataset_path)
-                    wav_name = os.path.splitext(os.path.basename(wav_path))[0]
-                    signal, rate = torchaudio.load(wav_path)
-                    spec = mel_spectrogram(signal, 1024, 80, 22050, 256,
-                              1024, 0, 8000, center=False).squeeze()
-                    # Write the features to feats.ark and feats.scp
-                    writer[wav_name] = spec
-    
-
-    emotions = [os.path.basename(x).split("_")[1] for x in glob.glob(dataset_path + '/**/**/*')]
-    emotions = sorted(set(emotions))
-
-    utt2spk = {}
-    utt2emo = {}
-    wavs = glob.glob(dataset_path + '**/**/*.wav')
-    for wav_path in tqdm(wavs):
-        wav_name = os.path.splitext(os.path.basename(wav_path))[0]
-        emotion =  emotions.index(wav_name.split("_")[1])
-        if wav_path.split('/')[-3] == '1263201035':
-            spk = 0 ## labels should start with 0
-        elif wav_path.split('/')[-3] == '805570882':
-            spk = 1
-        else:
-            spk = 2
-        utt2spk[wav_name] = str(spk)
-        utt2emo[wav_name] = str(emotion)
-    utt2spk = dict(sorted(utt2spk.items()))
-    utt2emo = dict(sorted(utt2emo.items()))
-
-    with open(filelists_path + 'utt2emo.json', 'w') as fp:
-        json.dump(utt2emo, fp,  indent=4)
-    with open(filelists_path + 'utt2spk.json', 'w') as fp:
-        json.dump(utt2spk, fp,  indent=4) 
-    
-    txt_files = sorted(glob.glob(dataset_path + '/**/**/*.txt'))
-    count = 0
-    txt = []
-    basenames = []
-    utt2text = {}
-    flag = False
-    with open(filelists_path + 'text', 'w', encoding='utf-8') as write:
-        for txt_path in txt_files:
-            basename = os.path.basename(txt_path).replace('.txt', '')
-            with open(txt_path, 'r', encoding='utf-8') as f:
-                txt.append(_clean_text(f.read().strip("\n"), cleaner_names=["kazakh_cleaners"]).replace("'", ""))
-                basenames.append(basename) 
-    output_string = [re.sub('(\d+)', lambda m: num2words(m.group(), lang='kz'), sentence) for sentence in txt]
-    cleaned_txt = []
-    for t in output_string:
-        cleaned_txt.append(''.join([s for s in t if s in symbols]))               
-    utt2text = {basenames[i]: cleaned_txt[i] for i in range(len(cleaned_txt))}
-    utt2text = dict(sorted(utt2text.items()))
-
-    vocab = set()
-    with open(filelists_path + '/text', 'w', encoding='utf-8') as f:
-        for x, y in utt2text.items():
-            for c in y: vocab.add(c)
-            f.write(x + ' ' +  y + '\n')

melspec.py

@@ -1,40 +0,0 @@
-import torch
-import torchaudio
-import librosa
-
-mel_basis = {}
-hann_window = {}
-
-def dynamic_range_compression_torch(x, C=1, clip_val=1e-5):
-    return torch.log(torch.clamp(x, min=clip_val) * C)
-
-def spectral_normalize_torch(magnitudes):
-    output = dynamic_range_compression_torch(magnitudes)
-    return output
-
-
-
-def mel_spectrogram(y, n_fft, num_mels, sampling_rate, hop_size, win_size, fmin, fmax, center=False):
-    if torch.min(y) < -1.:
-        print('min value is ', torch.min(y))
-    if torch.max(y) > 1.:
-        print('max value is ', torch.max(y))
-
-    global mel_basis, hann_window
-    if fmax not in mel_basis:
-        mel = librosa.filters.mel(sr=sampling_rate, n_fft=n_fft, n_mels=num_mels, fmin=fmin, fmax=fmax)
-        mel_basis[str(fmax)+'_'+str(y.device)] = torch.from_numpy(mel).float().to(y.device)
-        hann_window[str(y.device)] = torch.hann_window(win_size).to(y.device)
-
-    y = torch.nn.functional.pad(y.unsqueeze(1), (int((n_fft-hop_size)/2), int((n_fft-hop_size)/2)), mode='reflect')
-    y = y.squeeze(1)
-
-    spec = torch.stft(y, n_fft, hop_length=hop_size, win_length=win_size, window=hann_window[str(y.device)],
-                      center=center, pad_mode='reflect', normalized=False, onesided=True)
-
-    spec = torch.sqrt(spec.pow(2).sum(-1)+(1e-9))
-
-    spec = torch.matmul(mel_basis[str(fmax)+'_'+str(y.device)], spec)
-    spec = spectral_normalize_torch(spec)
-
-    return spec.numpy()

requirements.txt

@@ -0,0 +1,8 @@
+scipy
+torch
+kaldiio
+tqdm
+einops
+matplotlib
+glob
+pydub

train_EMA.py

@@ -1,207 +0,0 @@
-import numpy as np
-from tqdm import tqdm
-
-from copy import deepcopy
-import torch
-from torch.utils.data import DataLoader
-from torch.utils.tensorboard import SummaryWriter
-import data_collate
-import data_loader
-from utils_data import plot_tensor, save_plot
-from model.utils import fix_len_compatibility
-from text.symbols import symbols
-import utils_data as utils
-
-
-class ModelEmaV2(torch.nn.Module):
-    def __init__(self, model, decay=0.9999, device=None):
-        super(ModelEmaV2, self).__init__()
-        self.model_state_dict = deepcopy(model.state_dict())
-        self.decay = decay
-        self.device = device  # perform ema on different device from model if set
-
-    def _update(self, model, update_fn):
-        with torch.no_grad():
-            for ema_v, model_v in zip(self.model_state_dict.values(), model.state_dict().values()):
-                if self.device is not None:
-                    model_v = model_v.to(device=self.device)
-                ema_v.copy_(update_fn(ema_v, model_v))
-
-    def update(self, model):
-        self._update(model, update_fn=lambda e, m: self.decay * e + (1. - self.decay) * m)
-
-    def set(self, model):
-        self._update(model, update_fn=lambda e, m: m)
-
-    def state_dict(self, destination=None, prefix='', keep_vars=False):
-        return self.model_state_dict
-
-
-if __name__ == "__main__":
-    hps = utils.get_hparams()
-    logger_text = utils.get_logger(hps.model_dir)
-    logger_text.info(hps)
-
-    out_size = fix_len_compatibility(2 * hps.data.sampling_rate // hps.data.hop_length)  # NOTE: 2-sec of mel-spec
-
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    torch.manual_seed(hps.train.seed)
-    np.random.seed(hps.train.seed)
-
-    print('Initializing logger...')
-    log_dir = hps.model_dir
-    logger = SummaryWriter(log_dir=log_dir)
-
-    train_dataset, collate, model = utils.get_correct_class(hps)
-    test_dataset, _, _ = utils.get_correct_class(hps, train=False)
-
-    print('Initializing data loaders...')
-
-    batch_collate = collate
-    loader = DataLoader(dataset=train_dataset, batch_size=hps.train.batch_size,
-                        collate_fn=batch_collate, drop_last=True,
-                        num_workers=4, shuffle=False)  # NOTE: if on server, worker can be 4
-
-    print('Initializing model...')
-    model = model(**hps.model).to(device)
-    print('Number of encoder + duration predictor parameters: %.2fm' % (model.encoder.nparams / 1e6))
-    print('Number of decoder parameters: %.2fm' % (model.decoder.nparams / 1e6))
-    print('Total parameters: %.2fm' % (model.nparams / 1e6))
-
-    use_gt_dur = getattr(hps.train, "use_gt_dur", False)
-    if use_gt_dur:
-        print("++++++++++++++> Using ground truth duration for training")
-
-    print('Initializing optimizer...')
-    optimizer = torch.optim.Adam(params=model.parameters(), lr=hps.train.learning_rate)
-
-    print('Logging test batch...')
-    test_batch = test_dataset.sample_test_batch(size=hps.train.test_size)
-    for i, item in enumerate(test_batch):
-        mel = item['mel']
-        logger.add_image(f'image_{i}/ground_truth', plot_tensor(mel.squeeze()),
-                         global_step=0, dataformats='HWC')
-        save_plot(mel.squeeze(), f'{log_dir}/original_{i}.png')
-
-    try:
-        model, optimizer, learning_rate, epoch_logged = utils.load_checkpoint(utils.latest_checkpoint_path(hps.model_dir, "grad_*.pt"), model, optimizer)
-        epoch_start = epoch_logged + 1
-        print(f"Loaded checkpoint from {epoch_logged} epoch, resuming training.")
-        global_step = epoch_logged * (len(train_dataset)/hps.train.batch_size)
-    except:
-        print(f"Cannot find trained checkpoint, begin to train from scratch")
-        epoch_start = 1
-        global_step = 0
-        learning_rate = hps.train.learning_rate
-
-    ema_model = ModelEmaV2(model, decay=0.9999)  # It's necessary that we put this after loading model.
-
-    print('Start training...')
-    used_items = set()
-    iteration = global_step
-    for epoch in range(epoch_start, hps.train.n_epochs + 1):
-        model.train()
-        dur_losses = []
-        prior_losses = []
-        diff_losses = []
-        with tqdm(loader, total=len(train_dataset) // hps.train.batch_size) as progress_bar:
-            for batch_idx, batch in enumerate(progress_bar):
-                model.zero_grad()
-                x, x_lengths = batch['text_padded'].to(device), \
-                               batch['input_lengths'].to(device)
-                y, y_lengths = batch['mel_padded'].to(device), \
-                               batch['output_lengths'].to(device)
-                if hps.xvector:
-                    spk = batch['xvector'].to(device)
-                else:
-                    spk = batch['spk_ids'].to(torch.long).to(device)
-                emo = batch['emo_ids'].to(torch.long).to(device)
-
-                dur_loss, prior_loss, diff_loss = model.compute_loss(x, x_lengths,
-                                                                     y, y_lengths,
-                                                                     spk=spk,
-                                                                     emo=emo,
-                                                                     out_size=out_size,
-                                                                     use_gt_dur=use_gt_dur,
-                                                                     durs=batch['dur_padded'].to(device) if use_gt_dur else None)
-                loss = sum([dur_loss, prior_loss, diff_loss])
-                loss.backward()
-
-                enc_grad_norm = torch.nn.utils.clip_grad_norm_(model.encoder.parameters(),
-                                                               max_norm=1)
-                dec_grad_norm = torch.nn.utils.clip_grad_norm_(model.decoder.parameters(),
-                                                               max_norm=1)
-                optimizer.step()
-                ema_model.update(model)
-
-                logger.add_scalar('training/duration_loss', dur_loss.item(),
-                                  global_step=iteration)
-                logger.add_scalar('training/prior_loss', prior_loss.item(),
-                                  global_step=iteration)
-                logger.add_scalar('training/diffusion_loss', diff_loss.item(),
-                                  global_step=iteration)
-                logger.add_scalar('training/encoder_grad_norm', enc_grad_norm,
-                                  global_step=iteration)
-                logger.add_scalar('training/decoder_grad_norm', dec_grad_norm,
-                                  global_step=iteration)
-
-                dur_losses.append(dur_loss.item())
-                prior_losses.append(prior_loss.item())
-                diff_losses.append(diff_loss.item())
-
-                if batch_idx % 5 == 0:
-                    msg = f'Epoch: {epoch}, iteration: {iteration} | dur_loss: {dur_loss.item()}, prior_loss: {prior_loss.item()}, diff_loss: {diff_loss.item()}'
-                    progress_bar.set_description(msg)
-
-                iteration += 1
-
-        log_msg = 'Epoch %d: duration loss = %.3f ' % (epoch, float(np.mean(dur_losses)))
-        log_msg += '| prior loss = %.3f ' % np.mean(prior_losses)
-        log_msg += '| diffusion loss = %.3f\n' % np.mean(diff_losses)
-        with open(f'{log_dir}/train.log', 'a') as f:
-            f.write(log_msg)
-
-        if epoch % hps.train.save_every > 0:
-            continue
-
-        model.eval()
-        print('Synthesis...')
-
-        with torch.no_grad():
-            for i, item in enumerate(test_batch):
-                if item['utt'] + "/truth" not in used_items:
-                    used_items.add(item['utt'] + "/truth")
-                x = item['text'].to(torch.long).unsqueeze(0).to(device)
-                if not hps.xvector:
-                    spk = item['spk_ids']
-                    spk = torch.LongTensor([spk]).to(device)
-                else:
-                    spk = item["xvector"]
-                    spk = spk.unsqueeze(0).to(device)
-                emo = item['emo_ids']
-                emo = torch.LongTensor([emo]).to(device)
-                
-                x_lengths = torch.LongTensor([x.shape[-1]]).to(device)
-                
-                y_enc, y_dec, attn = model(x, x_lengths, spk=spk, emo=emo, n_timesteps=10)           
-                logger.add_image(f'image_{i}/generated_enc',
-                                 plot_tensor(y_enc.squeeze().cpu()),
-                                 global_step=iteration, dataformats='HWC')
-                logger.add_image(f'image_{i}/generated_dec',
-                                 plot_tensor(y_dec.squeeze().cpu()),
-                                 global_step=iteration, dataformats='HWC')
-                logger.add_image(f'image_{i}/alignment',
-                                 plot_tensor(attn.squeeze().cpu()),
-                                 global_step=iteration, dataformats='HWC')
-                save_plot(y_enc.squeeze().cpu(),
-                          f'{log_dir}/generated_enc_{i}.png')
-                save_plot(y_dec.squeeze().cpu(),
-                          f'{log_dir}/generated_dec_{i}.png')
-                save_plot(attn.squeeze().cpu(),
-                          f'{log_dir}/alignment_{i}.png')
-
-        ckpt = model.state_dict()
-
-        utils.save_checkpoint(ema_model, optimizer, learning_rate, epoch, checkpoint_path=f"{log_dir}/EMA_grad_{epoch}.pt")
-        utils.save_checkpoint(model, optimizer, learning_rate, epoch, checkpoint_path=f"{log_dir}/grad_{epoch}.pt")
-