melotts_rknn.py

import numpy as np
import soundfile
import onnxruntime as ort
import argparse
import time
from utils import *
from rknnlite.api import RKNNLite



def get_args():
    parser = argparse.ArgumentParser(
        prog="melotts",
        description="Run TTS on input sentence"
    )
    parser.add_argument("--sentence", "-s", type=str, required=False, default="爱芯元智半导体股份有限公司，致力于打造世界领先的人工智能感知与边缘计算芯片。服务智慧城市、智能驾驶、机器人的海量普惠的应用")
    parser.add_argument("--wav", "-w", type=str, required=False, default="output.wav")
    parser.add_argument("--encoder", "-e", type=str, required=False, default="encoder.onnx")
    parser.add_argument("--decoder", "-d", type=str, required=False, default="decoder.rknn")
    parser.add_argument("--sample_rate", "-sr", type=int, required=False, default=44100)
    parser.add_argument("--speed", type=float, required=False, default=0.8)
    parser.add_argument("--lexicon", type=str, required=False, default="lexicon.txt")
    parser.add_argument("--token", type=str, required=False, default="tokens.txt")
    return parser.parse_args()


def audio_numpy_concat(segment_data_list, sr, speed=1.):
    audio_segments = []
    for segment_data in segment_data_list:
        audio_segments += segment_data.reshape(-1).tolist()
        audio_segments += [0] * int((sr * 0.05) / speed)
    audio_segments = np.array(audio_segments).astype(np.float32)
    return audio_segments


def merge_sub_audio(sub_audio_list, pad_size, audio_len):
    # Average pad part
    if pad_size > 0:
        for i in range(len(sub_audio_list) - 1):
            sub_audio_list[i][-pad_size:] += sub_audio_list[i+1][:pad_size]
            sub_audio_list[i][-pad_size:] /= 2
            if i > 0:
                sub_audio_list[i] = sub_audio_list[i][pad_size:]

    sub_audio = np.concatenate(sub_audio_list, axis=-1)
    return sub_audio[:audio_len]


def generate_pronounce_slice(yinjie_num):
    # 按照每个词的音节数产生发音切片
    pron_slice = []
    start = 0
    end = 0
    for i, n in enumerate(yinjie_num):
        end = start + n
        pron_slice.append(slice(start, end))
        start = end
    return pron_slice


def generate_word_pron_num(pron_lens, pron_slices):
    # 求每个词语的发音长度
    pron_num = []
    for s in pron_slices:
        pron_num.append(pron_lens[s].sum())
    return pron_num


def decode_long_word(sess_dec, z_p, g, dec_len):
    z_p_len = z_p.shape[-1]
    slice_num = int(np.ceil(z_p_len / dec_len))
    sub_audio_list = []

    for i in range(slice_num):
        z_p_slice = z_p[..., i * dec_len : (i + 1) * dec_len]
        sub_dec_len = z_p_slice.shape[-1]
        sub_audio_len = 512 * sub_dec_len

        if z_p_slice.shape[-1] < dec_len:
            z_p_slice = np.concatenate((z_p_slice, np.zeros((*z_p_slice.shape[:-1], dec_len - z_p_slice.shape[-1]), dtype=np.float32)), axis=-1)

        start = time.time()
        audio = sess_dec.inference(inputs=[z_p_slice, g])[0].flatten()
        audio = audio[:sub_audio_len]
        print(f"Long word slice[{i}]: decoder run take {1000 * (time.time() - start):.2f}ms")
        sub_audio_list.append(audio)
    return sub_audio_list


def generate_decode_slices(pron_num, dec_len):
    pron_slices = []
    zp_slices = []
    strip_flags = [] # head tail
    pron_lens = []
    is_long = []

    start = end = 0
    zp_start, zp_end = 0, 0
    while (end < len(pron_num)):
        if len(pron_lens) > 0 and pron_lens[-1] > 2 and pron_num[end] < dec_len:
            prev_end = pron_slices[-1][1]
            pad_size = pron_num[prev_end - 1] + pron_num[prev_end - 2]
            start = end - 2
            zp_start = zp_end - pad_size
            strip_head, strip_tail = True, False
            if is_long[-1]:
                strip_flags[-1][1] = False
            else:
                strip_flags[-1][1] = True
        else:
            pad_size = 0
            start = end
            zp_start = zp_end
            strip_head, strip_tail = False, False
            if len(strip_flags) > 0:
                strip_flags[-1][1] = False

        sub_dec_len = abs(zp_end - zp_start)
        while (end < len(pron_num) and sub_dec_len < dec_len):
            sub_dec_len += pron_num[end]
            end += 1

        if end - start == 1 and sub_dec_len > dec_len:
            # Long word
            is_long.append(True)
            if len(strip_flags) > 0:
                strip_flags[-1][1] = False
        else:
            is_long.append(False)
            if sub_dec_len > dec_len:
                sub_dec_len -= pron_num[end - 1]
                end -= 1

        zp_end = zp_start + sub_dec_len

        pron_slices.append([start, end])
        zp_slices.append([zp_start, zp_end])
        strip_flags.append([strip_head, strip_tail])
        pron_lens.append(end - start)

    return pron_slices, zp_slices, strip_flags, pron_lens, is_long


def main():
    args = get_args()
    sentence = args.sentence
    sample_rate = args.sample_rate
    lexicon_filename = args.lexicon
    token_filename = args.token
    enc_model = args.encoder # default="../models/encoder.onnx"
    dec_model = args.decoder # default="../models/decoder.axmodel"

    print(f"sentence: {sentence}")
    print(f"sample_rate: {sample_rate}")
    print(f"lexicon: {lexicon_filename}")
    print(f"token: {token_filename}")
    print(f"encoder: {enc_model}")
    print(f"decoder: {dec_model}")

    # Split sentence
    sens = split_sentences_zh(sentence)

    # Load lexicon
    lexicon = Lexicon(lexicon_filename, token_filename)

    # Load models
    start = time.time()
    sess_enc = ort.InferenceSession(enc_model, providers=["CPUExecutionProvider"], sess_options=ort.SessionOptions())
    
    # 初始化RKNNLite
    decoder = RKNNLite()
    ret = decoder.load_rknn(dec_model)
    if ret != 0:
        print('Load decoder RKNN model failed')
        exit(ret)
    ret = decoder.init_runtime()
    if ret != 0:
        print('Init runtime failed') 
        exit(ret)
    
    dec_len = 65536 // 512
    print(f"load models take {1000 * (time.time() - start)}ms")

    # Load static input
    g = np.fromfile("g.bin", dtype=np.float32).reshape(1, 256, 1)

    # Final wav
    audio_list = []

    # Iterate over splitted sentences
    for n, se in enumerate(sens):
        print(f"\nSentence[{n}]: {se}")
        # Convert sentence to phones and tones
        phone_str, yinjie_num, phones, tones = lexicon.convert(se)

        # Add blank between words
        phone_str = intersperse(phone_str, 0)
        phones = np.array(intersperse(phones, 0), dtype=np.int32)
        tones = np.array(intersperse(tones, 0), dtype=np.int32)
        yinjie_num = np.array(yinjie_num, dtype=np.int32) * 2
        yinjie_num[0] += 1
        assert (yinjie_num.sum() == phones.shape[0])
        pron_slices = generate_pronounce_slice(yinjie_num)

        # for s in pron_slices:
        #     print(phone_str[s])
        
        phone_len = phones.shape[-1]

        language = np.array([3] * phone_len, dtype=np.int32)

        start = time.time()
        z_p, pronoun_lens, audio_len = sess_enc.run(None, input_feed={
                                    'phone': phones, 'g': g,
                                    'tone': tones, 'language': language, 
                                    'noise_scale': np.array([0], dtype=np.float32),
                                    'length_scale': np.array([1.0 / args.speed], dtype=np.float32),
                                    'noise_scale_w': np.array([0], dtype=np.float32),
                                    'sdp_ratio': np.array([0], dtype=np.float32)})
        print(f"encoder run take {1000 * (time.time() - start):.2f}ms")
        
        audio_len = audio_len[0]
        actual_size = z_p.shape[-1]
        dec_slice_num = int(np.ceil(actual_size / dec_len))
        # print(f"origin z_p.shape: {z_p.shape}")
        z_p = np.pad(z_p, pad_width=((0,0),(0,0),(0, dec_slice_num * dec_len - actual_size)), mode="constant", constant_values=0)

        # print(f"phone_len: {phone_len}")
        # print(f"z_p.shape: {z_p.shape}")
        # print(f"dec_slice_num: {dec_slice_num}")
        # print(f"audio_len: {audio_len}")

        # 生成每个词的发音数量
        pron_num = generate_word_pron_num(pronoun_lens, pron_slices)
        # assert (sum(pron_num) == pronoun_lens.sum())

        sub_audio_list = []
        pron_num_slices, zp_slices, strip_flags, pron_lens, is_long = \
            generate_decode_slices(pron_num, dec_len)
        
        for i in range(len(pron_num_slices)):
            pron_start, pron_end = pron_num_slices[i]
            zp_start, zp_end = zp_slices[i]

            phone_strs = []
            for n in range(pron_start, pron_end):
                phone_strs.extend(phone_str[pron_slices[n]])
            # print(f"phone str: {phone_strs}")

            if is_long[i]:
                sub_audio_list.extend(decode_long_word(decoder, z_p[..., zp_start : zp_end], g, dec_len))
            else:
                sub_dec_len = zp_end - zp_start
                sub_audio_len = 512 * sub_dec_len

                zp_slice = z_p[..., zp_start : zp_end]
                if zp_slice.shape[-1] < dec_len:
                    zp_slice = np.concatenate((zp_slice, np.zeros((*zp_slice.shape[:-1], dec_len - zp_slice.shape[-1]), dtype=np.float32)), axis=-1)

                start = time.time()
                outputs = decoder.inference(inputs=[zp_slice, g])
                audio = outputs[0].flatten()
                audio = audio[:sub_audio_len]
                print(f"Sentence[{n}] Slice[{i}]: decoder run take {1000 * (time.time() - start):.2f}ms")

                if strip_flags[i][0]:
                    # strip head
                    head = 512 * pron_num[pron_start]
                    # print(f"Strip head: {phone_str[pron_slices[pron_start]]}")
                    audio = audio[head : ]
                if strip_flags[i][1]:
                    # strip tail
                    tail = 512 * pron_num[pron_end - 1]
                    # print(f"Strip tail: {phone_str[pron_slices[pron_end - 1]]}")
                    audio = audio[: -tail]

                sub_audio_list.append(audio)

        sub_audio = merge_sub_audio(sub_audio_list, 0, audio_len)
        audio_list.append(sub_audio)

    audio = audio_numpy_concat(audio_list, sr=sample_rate, speed=args.speed)
    soundfile.write(args.wav, audio, sample_rate)
    print(f"Save to {args.wav}")

    # 在最后添加资源释放
    decoder.release()

if __name__ == "__main__":
    main()