style_bert_vits2/models/infer_onnx.py

from collections.abc import Sequence
from typing import Any

import numpy as np
import onnxruntime
from numpy.typing import NDArray
from pyopenjtalk import OpenJTalk

from style_bert_vits2.constants import Languages
from style_bert_vits2.models.hyper_parameters import HyperParameters
from style_bert_vits2.nlp import (
    clean_text_with_given_phone_tone,
    cleaned_text_to_sequence,
    extract_bert_feature_onnx,
)
from style_bert_vits2.utils import get_onnx_device_options


def __intersperse(lst: list[Any], item: Any) -> list[Any]:
    """
    リストの要素の間に特定のアイテムを挿入する
    style_bert_vits2.models.commons.intersperse と同一実装
    style_bert_vits2.models.commons モジュールは PyTorch に依存しているため、ONNX 推論時は import できない

    Args:
        lst (list[Any]): 元のリスト
        item (Any): 挿入するアイテム

    Returns:
        list[Any]: 新しいリスト
    """
    result = [item] * (len(lst) * 2 + 1)
    result[1::2] = lst
    return result


def get_text_onnx(
    text: str,
    language_str: Languages,
    hps: HyperParameters,
    onnx_providers: Sequence[str | tuple[str, dict[str, Any]]],
    assist_text: str | None = None,
    assist_text_weight: float = 0.7,
    given_phone: list[str] | None = None,
    given_tone: list[int] | None = None,
    jtalk: OpenJTalk | None = None,
) -> tuple[
    NDArray[Any], NDArray[Any], NDArray[Any], NDArray[Any], NDArray[Any], NDArray[Any]
]:
    use_jp_extra = hps.version.endswith("JP-Extra")
    norm_text, phone, tone, word2ph, sep_text, _, _ = clean_text_with_given_phone_tone(
        text,
        language_str,
        given_phone=given_phone,
        given_tone=given_tone,
        use_jp_extra=use_jp_extra,
        # 推論時のみ呼び出されるので、raise_yomi_error は False に設定
        raise_yomi_error=False,
        jtalk=jtalk,
    )
    phone, tone, language = cleaned_text_to_sequence(phone, tone, language_str)

    if hps.data.add_blank:
        phone = __intersperse(phone, 0)
        tone = __intersperse(tone, 0)
        language = __intersperse(language, 0)
        for i in range(len(word2ph)):
            word2ph[i] = word2ph[i] * 2
        word2ph[0] += 1
    bert_ori = extract_bert_feature_onnx(
        norm_text,
        word2ph,
        language_str,
        onnx_providers,
        assist_text,
        assist_text_weight,
        sep_text,  # clean_text_with_given_phone_tone() の中間生成物を再利用して効率向上を図る
    )
    del word2ph
    assert bert_ori.shape[-1] == len(phone), phone

    if language_str == Languages.ZH:
        bert = bert_ori
        ja_bert = np.zeros((1024, len(phone)), dtype=np.float32)
        en_bert = np.zeros((1024, len(phone)), dtype=np.float32)
    elif language_str == Languages.JP:
        bert = np.zeros((1024, len(phone)), dtype=np.float32)
        ja_bert = bert_ori
        en_bert = np.zeros((1024, len(phone)), dtype=np.float32)
    elif language_str == Languages.EN:
        bert = np.zeros((1024, len(phone)), dtype=np.float32)
        ja_bert = np.zeros((1024, len(phone)), dtype=np.float32)
        en_bert = bert_ori
    else:
        raise ValueError("language_str should be ZH, JP or EN")

    assert bert.shape[-1] == len(phone), (
        f"Bert seq len {bert.shape[-1]} != {len(phone)}"
    )

    phone = np.array(phone, dtype=np.int64)
    tone = np.array(tone, dtype=np.int64)
    language = np.array(language, dtype=np.int64)
    return bert, ja_bert, en_bert, phone, tone, language


def infer_onnx(
    text: str,
    style_vec: NDArray[Any],
    sdp_ratio: float,
    noise_scale: float,
    noise_scale_w: float,
    length_scale: float,
    sid: int,  # In the original Bert-VITS2, its speaker_name: str, but here it's id
    language: Languages,
    hps: HyperParameters,
    onnx_session: onnxruntime.InferenceSession,
    onnx_providers: Sequence[str | tuple[str, dict[str, Any]]],
    skip_start: bool = False,
    skip_end: bool = False,
    assist_text: str | None = None,
    assist_text_weight: float = 0.7,
    given_phone: list[str] | None = None,
    given_tone: list[int] | None = None,
    jtalk: OpenJTalk | None = None,
) -> NDArray[np.float32]:
    is_jp_extra = hps.version.endswith("JP-Extra")
    bert, ja_bert, en_bert, phones, tones, lang_ids = get_text_onnx(
        text,
        language,
        hps,
        onnx_providers=onnx_providers,
        assist_text=assist_text,
        assist_text_weight=assist_text_weight,
        given_phone=given_phone,
        given_tone=given_tone,
        jtalk=jtalk,
    )
    if skip_start:
        phones = phones[3:]
        tones = tones[3:]
        lang_ids = lang_ids[3:]
        bert = bert[:, 3:]
        ja_bert = ja_bert[:, 3:]
        en_bert = en_bert[:, 3:]
    if skip_end:
        phones = phones[:-2]
        tones = tones[:-2]
        lang_ids = lang_ids[:-2]
        bert = bert[:, :-2]
        ja_bert = ja_bert[:, :-2]
        en_bert = en_bert[:, :-2]

    x_tst = np.expand_dims(phones, axis=0)
    tones = np.expand_dims(tones, axis=0)
    lang_ids = np.expand_dims(lang_ids, axis=0)
    bert = np.expand_dims(bert, axis=0)
    ja_bert = np.expand_dims(ja_bert, axis=0)
    en_bert = np.expand_dims(en_bert, axis=0)
    x_tst_lengths = np.array([phones.shape[0]], dtype=np.int64)
    style_vec_tensor = np.expand_dims(style_vec, axis=0)
    del phones
    sid_tensor = np.array([sid], dtype=np.int64)

    input_names = [input.name for input in onnx_session.get_inputs()]
    output_name = onnx_session.get_outputs()[0].name
    if is_jp_extra:
        input_tensor = [
            x_tst,
            x_tst_lengths,
            sid_tensor,
            tones,
            lang_ids,
            ja_bert,
            style_vec_tensor,
            np.array(length_scale, dtype=np.float32),
            np.array(sdp_ratio, dtype=np.float32),
            np.array(noise_scale, dtype=np.float32),
            np.array(noise_scale_w, dtype=np.float32),
        ]
    else:
        input_tensor = [
            x_tst,
            x_tst_lengths,
            sid_tensor,
            tones,
            lang_ids,
            bert,
            ja_bert,
            en_bert,
            style_vec_tensor,
            np.array(length_scale, dtype=np.float32),
            np.array(sdp_ratio, dtype=np.float32),
            np.array(noise_scale, dtype=np.float32),
            np.array(noise_scale_w, dtype=np.float32),
        ]

    # 入力テンソルの転送に使用するデバイス種別, デバイス ID, 実行オプションを取得
    device_type, device_id, run_options = get_onnx_device_options(onnx_session, onnx_providers)  # fmt: skip

    # 推論デバイスに入力テンソルを割り当て
    ## GPU 推論の場合、device_type + device_id に対応する GPU デバイスに入力テンソルが割り当てられる
    io_binding = onnx_session.io_binding()
    for name, value in zip(input_names, input_tensor):
        gpu_tensor = onnxruntime.OrtValue.ortvalue_from_numpy(
            value, device_type, device_id
        )
        io_binding.bind_ortvalue_input(name, gpu_tensor)

    # 推論の実行
    io_binding.bind_output(output_name, device_type)
    onnx_session.run_with_iobinding(io_binding, run_options=run_options)
    output = io_binding.get_outputs()

    audio = output[0].numpy()[0, 0]

    del (
        x_tst,
        tones,
        lang_ids,
        bert,
        x_tst_lengths,
        sid_tensor,
        ja_bert,
        en_bert,
        style_vec,
    )  # , emo

    return audio