inference.py

#!/usr/bin/env python3
import os
import random
from pathlib import Path
from typing import Dict
from typing import List

import librosa
import numpy as np
import paddle
import soundfile as sf
import torch
from paddle import nn
from sedit_arg_parser import parse_args
from utils import build_vocoder_from_file
from utils import evaluate_durations
from utils import get_voc_out
from utils import is_chinese

from align import alignment
from align import alignment_zh
from align import words2phns
from align import words2phns_zh
from collect_fn import build_collate_fn
from mlm import build_model_from_file
from read_text import load_num_sequence_text
from read_text import read_2col_text
# from ParallelWaveGAN.parallel_wavegan.utils.utils import download_pretrained_model

random.seed(0)
np.random.seed(0)


def plot_mel_and_vocode_wav(wav_path: str,
                            source_lang: str='english',
                            target_lang: str='english',
                            model_name: str="paddle_checkpoint_en",
                            old_str: str="",
                            new_str: str="",
                            use_pt_vocoder: bool=False,
                            non_autoreg: bool=True):
    wav_org, output_feat, old_span_bdy, new_span_bdy, fs, hop_length = get_mlm_output(
        source_lang=source_lang,
        target_lang=target_lang,
        model_name=model_name,
        wav_path=wav_path,
        old_str=old_str,
        new_str=new_str,
        use_teacher_forcing=non_autoreg)

    masked_feat = output_feat[new_span_bdy[0]:new_span_bdy[1]]

    if target_lang == 'english':
        if use_pt_vocoder:
            output_feat = output_feat.cpu().numpy()
            output_feat = torch.tensor(output_feat, dtype=torch.float)
            vocoder = load_vocoder('vctk_parallel_wavegan.v1.long')
            replaced_wav = vocoder(output_feat).cpu().numpy()
        else:
            replaced_wav = get_voc_out(output_feat)

    elif target_lang == 'chinese':
        replaced_wav_only_mask_fst2_voc = get_voc_out(masked_feat)

    old_time_bdy = [hop_length * x for x in old_span_bdy]
    new_time_bdy = [hop_length * x for x in new_span_bdy]

    if target_lang == 'english':
        wav_org_replaced_paddle_voc = np.concatenate([
            wav_org[:old_time_bdy[0]],
            replaced_wav[new_time_bdy[0]:new_time_bdy[1]],
            wav_org[old_time_bdy[1]:]
        ])

        data_dict = {"origin": wav_org, "output": wav_org_replaced_paddle_voc}

    elif target_lang == 'chinese':
        wav_org_replaced_only_mask_fst2_voc = np.concatenate([
            wav_org[:old_time_bdy[0]], replaced_wav_only_mask_fst2_voc,
            wav_org[old_time_bdy[1]:]
        ])
        data_dict = {
            "origin": wav_org,
            "output": wav_org_replaced_only_mask_fst2_voc,
        }

    return data_dict, old_span_bdy


def load_vocoder(vocoder_tag: str="vctk_parallel_wavegan.v1.long"):
    vocoder_tag = vocoder_tag.replace("parallel_wavegan/", "")
    vocoder_file = download_pretrained_model(vocoder_tag)
    vocoder_config = Path(vocoder_file).parent / "config.yml"
    vocoder = build_vocoder_from_file(vocoder_config, vocoder_file, None, 'cpu')
    return vocoder


def load_model(model_name: str="paddle_checkpoint_en"):
    config_path = './pretrained_model/{}/config.yaml'.format(model_name)
    model_path = './pretrained_model/{}/model.pdparams'.format(model_name)
    mlm_model, conf = build_model_from_file(
        config_file=config_path, model_file=model_path)
    return mlm_model, conf


def read_data(uid: str, prefix: os.PathLike):
    # 获取 uid 对应的文本
    mfa_text = read_2col_text(prefix + '/text')[uid]
    # 获取 uid 对应的音频路径
    mfa_wav_path = read_2col_text(prefix + '/wav.scp')[uid]
    if not os.path.isabs(mfa_wav_path):
        mfa_wav_path = prefix + mfa_wav_path
    return mfa_text, mfa_wav_path


def get_align_data(uid: str, prefix: os.PathLike):
    mfa_path = prefix + "mfa_"
    mfa_text = read_2col_text(mfa_path + 'text')[uid]
    mfa_start = load_num_sequence_text(
        mfa_path + 'start', loader_type='text_float')[uid]
    mfa_end = load_num_sequence_text(
        mfa_path + 'end', loader_type='text_float')[uid]
    mfa_wav_path = read_2col_text(mfa_path + 'wav.scp')[uid]
    return mfa_text, mfa_start, mfa_end, mfa_wav_path


# 获取需要被 mask 的 mel 帧的范围
def get_masked_mel_bdy(mfa_start: List[float],
                       mfa_end: List[float],
                       fs: int,
                       hop_length: int,
                       span_to_repl: List[List[int]]):
    align_start = np.array(mfa_start)
    align_end = np.array(mfa_end)
    align_start = np.floor(fs * align_start / hop_length).astype('int')
    align_end = np.floor(fs * align_end / hop_length).astype('int')
    if span_to_repl[0] >= len(mfa_start):
        span_bdy = [align_end[-1], align_end[-1]]
    else:
        span_bdy = [
            align_start[span_to_repl[0]], align_end[span_to_repl[1] - 1]
        ]
    return span_bdy, align_start, align_end


def recover_dict(word2phns: Dict[str, str], tp_word2phns: Dict[str, str]):
    dic = {}
    keys_to_del = []
    exist_idx = []
    sp_count = 0
    add_sp_count = 0
    for key in word2phns.keys():
        idx, wrd = key.split('_')
        if wrd == 'sp':
            sp_count += 1
            exist_idx.append(int(idx))
        else:
            keys_to_del.append(key)

    for key in keys_to_del:
        del word2phns[key]

    cur_id = 0
    for key in tp_word2phns.keys():
        if cur_id in exist_idx:
            dic[str(cur_id) + "_sp"] = 'sp'
            cur_id += 1
            add_sp_count += 1
        idx, wrd = key.split('_')
        dic[str(cur_id) + "_" + wrd] = tp_word2phns[key]
        cur_id += 1

    if add_sp_count + 1 == sp_count:
        dic[str(cur_id) + "_sp"] = 'sp'
        add_sp_count += 1

    assert add_sp_count == sp_count, "sp are not added in dic"
    return dic


def get_max_idx(dic):
    return sorted([int(key.split('_')[0]) for key in dic.keys()])[-1]


def get_phns_and_spans(wav_path: str,
                       old_str: str="",
                       new_str: str="",
                       source_lang: str="english",
                       target_lang: str="english"):
    is_append = (old_str == new_str[:len(old_str)])
    old_phns, mfa_start, mfa_end = [], [], []
    # source
    if source_lang == "english":
        intervals, word2phns = alignment(wav_path, old_str)
    elif source_lang == "chinese":
        intervals, word2phns = alignment_zh(wav_path, old_str)
        _, tp_word2phns = words2phns_zh(old_str)

        for key, value in tp_word2phns.items():
            idx, wrd = key.split('_')
            cur_val = " ".join(value)
            tp_word2phns[key] = cur_val

        word2phns = recover_dict(word2phns, tp_word2phns)
    else:
        assert source_lang == "chinese" or source_lang == "english", \
            "source_lang is wrong..."

    for item in intervals:
        old_phns.append(item[0])
        mfa_start.append(float(item[1]))
        mfa_end.append(float(item[2]))
    # target
    if is_append and (source_lang != target_lang):
        cross_lingual_clone = True
    else:
        cross_lingual_clone = False

    if cross_lingual_clone:
        str_origin = new_str[:len(old_str)]
        str_append = new_str[len(old_str):]

        if target_lang == "chinese":
            phns_origin, origin_word2phns = words2phns(str_origin)
            phns_append, append_word2phns_tmp = words2phns_zh(str_append)

        elif target_lang == "english":
            # 原始句子
            phns_origin, origin_word2phns = words2phns_zh(str_origin)
            # clone 句子 
            phns_append, append_word2phns_tmp = words2phns(str_append)
        else:
            assert target_lang == "chinese" or target_lang == "english", \
                "cloning is not support for this language, please check it."

        new_phns = phns_origin + phns_append

        append_word2phns = {}
        length = len(origin_word2phns)
        for key, value in append_word2phns_tmp.items():
            idx, wrd = key.split('_')
            append_word2phns[str(int(idx) + length) + '_' + wrd] = value
        new_word2phns = origin_word2phns.copy()
        new_word2phns.update(append_word2phns)

    else:
        if source_lang == target_lang and target_lang == "english":
            new_phns, new_word2phns = words2phns(new_str)
        elif source_lang == target_lang and target_lang == "chinese":
            new_phns, new_word2phns = words2phns_zh(new_str)
        else:
            assert source_lang == target_lang, \
                "source language is not same with target language..."

    span_to_repl = [0, len(old_phns) - 1]
    span_to_add = [0, len(new_phns) - 1]
    left_idx = 0
    new_phns_left = []
    sp_count = 0
    # find the left different index
    for key in word2phns.keys():
        idx, wrd = key.split('_')
        if wrd == 'sp':
            sp_count += 1
            new_phns_left.append('sp')
        else:
            idx = str(int(idx) - sp_count)
            if idx + '_' + wrd in new_word2phns:
                left_idx += len(new_word2phns[idx + '_' + wrd])
                new_phns_left.extend(word2phns[key].split())
            else:
                span_to_repl[0] = len(new_phns_left)
                span_to_add[0] = len(new_phns_left)
                break

    # reverse word2phns and new_word2phns
    right_idx = 0
    new_phns_right = []
    sp_count = 0
    word2phns_max_idx = get_max_idx(word2phns)
    new_word2phns_max_idx = get_max_idx(new_word2phns)
    new_phns_mid = []
    if is_append:
        new_phns_right = []
        new_phns_mid = new_phns[left_idx:]
        span_to_repl[0] = len(new_phns_left)
        span_to_add[0] = len(new_phns_left)
        span_to_add[1] = len(new_phns_left) + len(new_phns_mid)
        span_to_repl[1] = len(old_phns) - len(new_phns_right)
    # speech edit
    else:
        for key in list(word2phns.keys())[::-1]:
            idx, wrd = key.split('_')
            if wrd == 'sp':
                sp_count += 1
                new_phns_right = ['sp'] + new_phns_right
            else:
                idx = str(new_word2phns_max_idx - (word2phns_max_idx - int(idx)
                                                   - sp_count))
                if idx + '_' + wrd in new_word2phns:
                    right_idx -= len(new_word2phns[idx + '_' + wrd])
                    new_phns_right = word2phns[key].split() + new_phns_right
                else:
                    span_to_repl[1] = len(old_phns) - len(new_phns_right)
                    new_phns_mid = new_phns[left_idx:right_idx]
                    span_to_add[1] = len(new_phns_left) + len(new_phns_mid)
                    if len(new_phns_mid) == 0:
                        span_to_add[1] = min(span_to_add[1] + 1, len(new_phns))
                        span_to_add[0] = max(0, span_to_add[0] - 1)
                        span_to_repl[0] = max(0, span_to_repl[0] - 1)
                        span_to_repl[1] = min(span_to_repl[1] + 1,
                                              len(old_phns))
                    break
    new_phns = new_phns_left + new_phns_mid + new_phns_right
    '''
    For that reason cover should not be given.
    For that reason cover is impossible to be given.
    span_to_repl: [17, 23] "should not"
    span_to_add: [17, 30]  "is impossible to"
    '''
    return mfa_start, mfa_end, old_phns, new_phns, span_to_repl, span_to_add


# mfa 获得的 duration 和 fs2 的 duration_predictor 获取的 duration 可能不同
# 此处获得一个缩放比例, 用于预测值和真实值之间的缩放
def duration_adjust_factor(orig_dur: List[int],
                           pred_dur: List[int],
                           phns: List[str]):
    length = 0
    factor_list = []
    for orig, pred, phn in zip(orig_dur, pred_dur, phns):
        if pred == 0 or phn == 'sp':
            continue
        else:
            factor_list.append(orig / pred)
    factor_list = np.array(factor_list)
    factor_list.sort()
    if len(factor_list) < 5:
        return 1
    length = 2
    avg = np.average(factor_list[length:-length])
    return avg


def prep_feats_with_dur(wav_path: str,
                        mlm_model: nn.Layer,
                        source_lang: str="English",
                        target_lang: str="English",
                        old_str: str="",
                        new_str: str="",
                        mask_reconstruct: bool=False,
                        duration_adjust: bool=True,
                        start_end_sp: bool=False,
                        fs: int=24000,
                        hop_length: int=300):
    '''
    Returns:
        np.ndarray: new wav, replace the part to be edited in original wav with 0
        List[str]: new phones
        List[float]: mfa start of new wav
        List[float]: mfa end of new wav
        List[int]: masked mel boundary of original wav
        List[int]: masked mel boundary of new wav
    '''
    wav_org, _ = librosa.load(wav_path, sr=fs)

    mfa_start, mfa_end, old_phns, new_phns, span_to_repl, span_to_add = get_phns_and_spans(
        wav_path=wav_path,
        old_str=old_str,
        new_str=new_str,
        source_lang=source_lang,
        target_lang=target_lang)

    if start_end_sp:
        if new_phns[-1] != 'sp':
            new_phns = new_phns + ['sp']
    # 中文的 phns 不一定都在 fastspeech2 的字典里, 用 sp 代替
    if target_lang == "english" or target_lang == "chinese":
        old_durs = evaluate_durations(old_phns, target_lang=source_lang)
    else:
        assert target_lang == "chinese" or target_lang == "english", \
            "calculate duration_predict is not support for this language..."

    orig_old_durs = [e - s for e, s in zip(mfa_end, mfa_start)]
    if '[MASK]' in new_str:
        new_phns = old_phns
        span_to_add = span_to_repl
        d_factor_left = duration_adjust_factor(
            orig_dur=orig_old_durs[:span_to_repl[0]],
            pred_dur=old_durs[:span_to_repl[0]],
            phns=old_phns[:span_to_repl[0]])
        d_factor_right = duration_adjust_factor(
            orig_dur=orig_old_durs[span_to_repl[1]:],
            pred_dur=old_durs[span_to_repl[1]:],
            phns=old_phns[span_to_repl[1]:])
        d_factor = (d_factor_left + d_factor_right) / 2
        new_durs_adjusted = [d_factor * i for i in old_durs]
    else:
        if duration_adjust:
            d_factor = duration_adjust_factor(
                orig_dur=orig_old_durs, pred_dur=old_durs, phns=old_phns)
            print("d_factor:", d_factor)
            d_factor = d_factor * 1.25
        else:
            d_factor = 1

        if target_lang == "english" or target_lang == "chinese":
            new_durs = evaluate_durations(new_phns, target_lang=target_lang)
        else:
            assert target_lang == "chinese" or target_lang == "english", \
                "calculate duration_predict is not support for this language..."

        new_durs_adjusted = [d_factor * i for i in new_durs]

    new_span_dur_sum = sum(new_durs_adjusted[span_to_add[0]:span_to_add[1]])
    old_span_dur_sum = sum(orig_old_durs[span_to_repl[0]:span_to_repl[1]])
    dur_offset = new_span_dur_sum - old_span_dur_sum
    new_mfa_start = mfa_start[:span_to_repl[0]]
    new_mfa_end = mfa_end[:span_to_repl[0]]
    for i in new_durs_adjusted[span_to_add[0]:span_to_add[1]]:
        if len(new_mfa_end) == 0:
            new_mfa_start.append(0)
            new_mfa_end.append(i)
        else:
            new_mfa_start.append(new_mfa_end[-1])
            new_mfa_end.append(new_mfa_end[-1] + i)
    new_mfa_start += [i + dur_offset for i in mfa_start[span_to_repl[1]:]]
    new_mfa_end += [i + dur_offset for i in mfa_end[span_to_repl[1]:]]

    # 3. get new wav
    # 在原始句子后拼接
    if span_to_repl[0] >= len(mfa_start):
        left_idx = len(wav_org)
        right_idx = left_idx
    # 在原始句子中间替换
    else:
        left_idx = int(np.floor(mfa_start[span_to_repl[0]] * fs))
        right_idx = int(np.ceil(mfa_end[span_to_repl[1] - 1] * fs))
    blank_wav = np.zeros(
        (int(np.ceil(new_span_dur_sum * fs)), ), dtype=wav_org.dtype)
    # 原始音频，需要编辑的部分替换成空音频，空音频的时间由 fs2 的 duration_predictor 决定
    new_wav = np.concatenate(
        [wav_org[:left_idx], blank_wav, wav_org[right_idx:]])

    # 4. get old and new mel span to be mask
    # [92, 92]

    old_span_bdy, mfa_start, mfa_end = get_masked_mel_bdy(
        mfa_start=mfa_start,
        mfa_end=mfa_end,
        fs=fs,
        hop_length=hop_length,
        span_to_repl=span_to_repl)
    # [92, 174]
    # new_mfa_start, new_mfa_end 时间级别的开始和结束时间 -> 帧级别
    new_span_bdy, new_mfa_start, new_mfa_end = get_masked_mel_bdy(
        mfa_start=new_mfa_start,
        mfa_end=new_mfa_end,
        fs=fs,
        hop_length=hop_length,
        span_to_repl=span_to_add)

    # old_span_bdy, new_span_bdy 是帧级别的范围
    return new_wav, new_phns, new_mfa_start, new_mfa_end, old_span_bdy, new_span_bdy


def prep_feats(mlm_model: nn.Layer,
               wav_path: str,
               source_lang: str="english",
               target_lang: str="english",
               old_str: str="",
               new_str: str="",
               duration_adjust: bool=True,
               start_end_sp: bool=False,
               mask_reconstruct: bool=False,
               fs: int=24000,
               hop_length: int=300,
               token_list: List[str]=[]):
    wav, phns, mfa_start, mfa_end, old_span_bdy, new_span_bdy = prep_feats_with_dur(
        source_lang=source_lang,
        target_lang=target_lang,
        mlm_model=mlm_model,
        old_str=old_str,
        new_str=new_str,
        wav_path=wav_path,
        duration_adjust=duration_adjust,
        start_end_sp=start_end_sp,
        mask_reconstruct=mask_reconstruct,
        fs=fs,
        hop_length=hop_length)

    token_to_id = {item: i for i, item in enumerate(token_list)}
    text = np.array(
        list(map(lambda x: token_to_id.get(x, token_to_id['<unk>']), phns)))
    span_bdy = np.array(new_span_bdy)

    batch = [('1', {
        "speech": wav,
        "align_start": mfa_start,
        "align_end": mfa_end,
        "text": text,
        "span_bdy": span_bdy
    })]

    return batch, old_span_bdy, new_span_bdy


def decode_with_model(mlm_model: nn.Layer,
                      collate_fn,
                      wav_path: str,
                      source_lang: str="english",
                      target_lang: str="english",
                      old_str: str="",
                      new_str: str="",
                      use_teacher_forcing: bool=False,
                      duration_adjust: bool=True,
                      start_end_sp: bool=False,
                      fs: int=24000,
                      hop_length: int=300,
                      token_list: List[str]=[]):
    batch, old_span_bdy, new_span_bdy = prep_feats(
        source_lang=source_lang,
        target_lang=target_lang,
        mlm_model=mlm_model,
        wav_path=wav_path,
        old_str=old_str,
        new_str=new_str,
        duration_adjust=duration_adjust,
        start_end_sp=start_end_sp,
        fs=fs,
        hop_length=hop_length,
        token_list=token_list)

    feats = collate_fn(batch)[1]

    if 'text_masked_pos' in feats.keys():
        feats.pop('text_masked_pos')

    output = mlm_model.inference(
        text=feats['text'],
        speech=feats['speech'],
        masked_pos=feats['masked_pos'],
        speech_mask=feats['speech_mask'],
        text_mask=feats['text_mask'],
        speech_seg_pos=feats['speech_seg_pos'],
        text_seg_pos=feats['text_seg_pos'],
        span_bdy=new_span_bdy,
        use_teacher_forcing=use_teacher_forcing)
    if 0 in output[0].shape and 0 not in output[-1].shape:
        output_feat = paddle.concat(
            output[1:-1] + [output[-1].squeeze()], axis=0)
    elif 0 not in output[0].shape and 0 in output[-1].shape:
        output_feat = paddle.concat(
            [output[0].squeeze()] + output[1:-1], axis=0)
    elif 0 in output[0].shape and 0 in output[-1].shape:
        output_feat = paddle.concat(output[1:-1], axis=0)
    else:
        output_feat = paddle.concat(
            [output[0].squeeze(0)] + output[1:-1] + [output[-1].squeeze(0)],
            axis=0)

    wav_org, _ = librosa.load(wav_path, sr=fs)
    return wav_org, output_feat, old_span_bdy, new_span_bdy, fs, hop_length


def get_mlm_output(wav_path: str,
                   model_name: str="paddle_checkpoint_en",
                   source_lang: str="english",
                   target_lang: str="english",
                   old_str: str="",
                   new_str: str="",
                   use_teacher_forcing: bool=False,
                   duration_adjust: bool=True,
                   start_end_sp: bool=False):
    mlm_model, train_conf = load_model(model_name)
    mlm_model.eval()

    collate_fn = build_collate_fn(
        sr=train_conf.feats_extract_conf['fs'],
        n_fft=train_conf.feats_extract_conf['n_fft'],
        hop_length=train_conf.feats_extract_conf['hop_length'],
        win_length=train_conf.feats_extract_conf['win_length'],
        n_mels=train_conf.feats_extract_conf['n_mels'],
        fmin=train_conf.feats_extract_conf['fmin'],
        fmax=train_conf.feats_extract_conf['fmax'],
        mlm_prob=train_conf['mlm_prob'],
        mean_phn_span=train_conf['mean_phn_span'],
        train=False,
        seg_emb=train_conf.encoder_conf['input_layer'] == 'sega_mlm')

    return decode_with_model(
        source_lang=source_lang,
        target_lang=target_lang,
        mlm_model=mlm_model,
        collate_fn=collate_fn,
        wav_path=wav_path,
        old_str=old_str,
        new_str=new_str,
        use_teacher_forcing=use_teacher_forcing,
        duration_adjust=duration_adjust,
        start_end_sp=start_end_sp,
        fs=train_conf.feats_extract_conf['fs'],
        hop_length=train_conf.feats_extract_conf['hop_length'],
        token_list=train_conf.token_list)


def evaluate(uid: str,
             source_lang: str="english",
             target_lang: str="english",
             use_pt_vocoder: bool=False,
             prefix: os.PathLike="./prompt/dev/",
             model_name: str="paddle_checkpoint_en",
             new_str: str="",
             prompt_decoding: bool=False,
             task_name: str=None):

    # get origin text and path of origin wav
    old_str, wav_path = read_data(uid=uid, prefix=prefix)

    if task_name == 'edit':
        new_str = new_str
    elif task_name == 'synthesize':
        new_str = old_str + new_str
    else:
        new_str = old_str + ' '.join([ch for ch in new_str if is_chinese(ch)])

    print('new_str is ', new_str)

    results_dict, old_span = plot_mel_and_vocode_wav(
        source_lang=source_lang,
        target_lang=target_lang,
        model_name=model_name,
        wav_path=wav_path,
        old_str=old_str,
        new_str=new_str,
        use_pt_vocoder=use_pt_vocoder)
    return results_dict


if __name__ == "__main__":
    # parse config and args
    args = parse_args()

    data_dict = evaluate(
        uid=args.uid,
        source_lang=args.source_lang,
        target_lang=args.target_lang,
        use_pt_vocoder=args.use_pt_vocoder,
        prefix=args.prefix,
        model_name=args.model_name,
        new_str=args.new_str,
        task_name=args.task_name)
    sf.write(args.output_name, data_dict['output'], samplerate=24000)
    print("finished...")