KazEmoTTS/inference_EMA.py at master · aztechell/KazEmoTTS · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
import json
import numpy as np
import torch
from pydub import AudioSegment
import re
from num2words import num2words
from text import convert_text
from model import GradTTSWithEmo
import utils_data as utils
from attrdict import AttrDict
from models import Generator as HiFiGAN
from pathlib import Path
from datetime import datetime

HIFIGAN_CONFIG = './configs/hifigan-config.json'
HIFIGAN_CHECKPT = r'.\pre_trained\g_01720000'

if __name__ == '__main__':
    hps, args = utils.get_hparams_decode()
    device = torch.device('cpu' if not torch.cuda.is_available() else "cuda")
    ckpt = args.model
    model = GradTTSWithEmo(**hps.model).to(device)
    logger = utils.get_logger(hps.model_dir, "inference.log")
    utils.load_checkpoint(ckpt, model, None)
    _ = model.cuda().eval()

    print('Initializing HiFi-GAN...')
    with open(HIFIGAN_CONFIG) as f:
        h = AttrDict(json.load(f))
    vocoder = HiFiGAN(h)
    vocoder.load_state_dict(torch.load(HIFIGAN_CHECKPT, map_location=lambda loc, storage: loc)['generator'])
    _ = vocoder.cuda().eval()
    vocoder.remove_weight_norm()

    emos = sorted(["angry", "surprise", "fear", "happy", "neutral", "sad"])
    speakers = ['M1', 'F1', 'M2']

    with open(args.file, 'r', encoding='utf-8') as f:
        texts = [line.strip() for line in f.readlines()]

    replace_nums = []
    for i in texts:
        replace_nums.append(i.split('|', 1))

    nums2word = [re.sub('(\d+)', lambda m: num2words(m.group(), lang='kz'), sentence) for sentence in np.array(replace_nums)[:, 0]]
    text2speech = []
    for i, j in zip(nums2word, np.array(replace_nums)[:, 1]):
        text2speech.append(f'{i}|{j}')

    for i, line in enumerate(text2speech):
        emo_i = int(line.split('|')[1])
        control_spk_id = int(line.split('|')[2])
        control_emo_id = emos.index(emos[emo_i])
        text = line.split('|')[0]
        with torch.no_grad():
            ### define emotion
            emo = torch.LongTensor([control_emo_id]).to(device)
            sid = torch.LongTensor([control_spk_id]).to(device)
            text_padded, text_len = convert_text(text)
            y_enc, y_dec, attn = model.forward(text_padded, text_len,
                                        n_timesteps=args.timesteps,
                                        temperature=args.noise,
                                        stoc=args.stoc, spk=sid,emo=emo, length_scale=1.,
                                        classifier_free_guidance=args.guidance)
        res = y_dec.squeeze().cpu().numpy()
        x = torch.from_numpy(res).cuda().unsqueeze(0)
        y_g_hat = vocoder(x)
        audio = y_g_hat.squeeze()
        audio = audio * 32768.0
        audio = audio.detach().cpu().numpy().astype('int16')
        audio = AudioSegment(audio.data, frame_rate=22050, sample_width=2, channels=1)

        out_dir = Path(args.generated_path)
        out_dir.mkdir(parents=True, exist_ok=True)

        spk = speakers[int(line.split("|")[2])]
        ts = datetime.now().strftime("%Y%m%d-%H%M%S")
        fname = f"{emos[emo_i]}_{spk}_{ts}.mp3"
        audio.export(str(out_dir / fname), format="mp3", bitrate="192k")

        del y_enc, y_dec, attn, audio
        torch.cuda.empty_cache()