File size: 7,607 Bytes
0b85fb9
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
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