Hitonabi
09a1c98514
- pocket_server.py (Produktions-TTS mit Stimmen-Waechter), text_norm, Bench-/Diag-Skripte - lucy-f5: f5_server/f5_test/bench_dml (DirectML-Experiment, Phase C/D offen) - .gitignore: venvs/Modelle/Audio/Logs der beiden Ordner + box_recon/gemma_swap-Scratch Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
269 lines
13 KiB
Python
269 lines
13 KiB
Python
import re
|
||
import site
|
||
import time
|
||
import jieba
|
||
import torch
|
||
import onnxruntime
|
||
import soundfile as sf
|
||
import numpy as np
|
||
from pydub import AudioSegment
|
||
from pypinyin import lazy_pinyin, Style
|
||
python_package_path = site.getsitepackages()[-1]
|
||
|
||
BASE = r"F:\Coding Stuff\mission-control-2\client\lucy-f5"
|
||
vocab_path = BASE + r"\vocab_v1.txt"
|
||
onnx_model_A = BASE + r"\onnx_f32\F5_Preprocess.onnx"
|
||
onnx_model_B = BASE + r"\onnx_f32\F5_Transformer.onnx"
|
||
onnx_model_C = BASE + r"\onnx_f32\F5_Decode.onnx"
|
||
generated_audio = BASE + r"\bench_out.wav"
|
||
test_in_english = True # eingebaute engl. Referenz -> kein Chinesisch/Pinyin; Tempo ist sprachunabhängig
|
||
|
||
if test_in_english:
|
||
reference_audio = python_package_path + "/f5_tts/infer/examples/basic/basic_ref_en.wav"
|
||
ref_text = "Some call me nature, others call me mother nature."
|
||
# realistischer Satz-Längen-Benchmark (wie eine echte Lucy-Antwort):
|
||
gen_text = "Of course, Commander. I will restart the service, check the logs, and let you know once everything is running again."
|
||
else:
|
||
reference_audio = python_package_path + "/f5_tts/infer/examples/basic/basic_ref_zh.wav" # The reference audio path.
|
||
ref_text = "对，这就是我，万人敬仰的太乙真人。" # The ASR result of reference audio.
|
||
gen_text = "对，这就是我，万人敬仰的大可奇奇。" # The target TTS.
|
||
|
||
|
||
import os as _os
|
||
ORT_Accelerate_Providers = [_os.environ.get("BENCH_PROVIDER", "DmlExecutionProvider")] # DML (9070 XT) oder CPUExecutionProvider
|
||
# else keep empty.
|
||
RANDOM_SEED = 9527 # Set seed to reproduce the generated audio
|
||
NFE_STEP = int(_os.environ.get("BENCH_NFE", "32")) if (_os := __import__("os")) else 32 # via env testbar
|
||
FUSE_NFE = 1 # Maintain the same values as the exported model.
|
||
SPEED = 1.0 # Set for talking speed. Only works with dynamic_axes=True
|
||
MAX_THREADS = 8 # Max CPU parallel threads.
|
||
DEVICE_ID = 0 # The GPU id, default to 0.
|
||
MODEL_SAMPLE_RATE = 24000 # Do not modify it.
|
||
HOP_LENGTH = 256 # It affects the generated audio length and speech speed.
|
||
|
||
if "OpenVINOExecutionProvider" in ORT_Accelerate_Providers:
|
||
provider_options = [
|
||
{
|
||
'device_type': 'CPU', # [CPU, NPU, GPU, GPU.0, GPU.1]]
|
||
'precision': 'ACCURACY', # [FP32, FP16, ACCURACY]
|
||
'num_of_threads': MAX_THREADS,
|
||
'num_streams': 1,
|
||
'enable_opencl_throttling': True,
|
||
'enable_qdq_optimizer': False # Enable it carefully
|
||
}
|
||
]
|
||
elif "CUDAExecutionProvider" in ORT_Accelerate_Providers:
|
||
provider_options = [
|
||
{
|
||
'device_id': DEVICE_ID,
|
||
'gpu_mem_limit': 8 * 1024 * 1024 * 1024, # 8 GB
|
||
'arena_extend_strategy': 'kNextPowerOfTwo',
|
||
'cudnn_conv_algo_search': 'EXHAUSTIVE',
|
||
'cudnn_conv_use_max_workspace': '1',
|
||
'do_copy_in_default_stream': '1',
|
||
'cudnn_conv1d_pad_to_nc1d': '1',
|
||
'enable_cuda_graph': '0', # Set to '0' to avoid potential errors when enabled.
|
||
'use_tf32': '0'
|
||
}
|
||
]
|
||
else:
|
||
# Please config by yourself for others providers.
|
||
provider_options = None
|
||
|
||
|
||
with open(vocab_path, "r", encoding="utf-8") as f:
|
||
vocab_char_map = {}
|
||
for i, char in enumerate(f):
|
||
vocab_char_map[char[:-1]] = i
|
||
vocab_size = len(vocab_char_map)
|
||
|
||
|
||
# From the official code
|
||
def convert_char_to_pinyin(text_list, polyphone=True):
|
||
if jieba.dt.initialized is False:
|
||
jieba.default_logger.setLevel(50) # CRITICAL
|
||
jieba.initialize()
|
||
|
||
final_text_list = []
|
||
custom_trans = str.maketrans(
|
||
{";": ",", "“": '"', "”": '"', "‘": "'", "’": "'"}
|
||
) # add custom trans here, to address oov
|
||
|
||
def is_chinese(c):
|
||
return (
|
||
"\u3100" <= c <= "\u9fff" # common chinese characters
|
||
)
|
||
|
||
for text in text_list:
|
||
char_list = []
|
||
text = text.translate(custom_trans)
|
||
for seg in jieba.cut(text):
|
||
seg_byte_len = len(bytes(seg, "UTF-8"))
|
||
if seg_byte_len == len(seg): # if pure alphabets and symbols
|
||
if char_list and seg_byte_len > 1 and char_list[-1] not in " :'\"":
|
||
char_list.append(" ")
|
||
char_list.extend(seg)
|
||
elif polyphone and seg_byte_len == 3 * len(seg): # if pure east asian characters
|
||
seg_ = lazy_pinyin(seg, style=Style.TONE3, tone_sandhi=True)
|
||
for i, c in enumerate(seg):
|
||
if is_chinese(c):
|
||
char_list.append(" ")
|
||
char_list.append(seg_[i])
|
||
else: # if mixed characters, alphabets and symbols
|
||
for c in seg:
|
||
if ord(c) < 256:
|
||
char_list.extend(c)
|
||
elif is_chinese(c):
|
||
char_list.append(" ")
|
||
char_list.extend(lazy_pinyin(c, style=Style.TONE3, tone_sandhi=True))
|
||
else:
|
||
char_list.append(c)
|
||
final_text_list.append(char_list)
|
||
return final_text_list
|
||
|
||
|
||
# From the official code
|
||
def list_str_to_idx(
|
||
text: list[str] | list[list[str]],
|
||
vocab_char_map: dict[str, int], # {char: idx}
|
||
padding_value=-1
|
||
):
|
||
get_idx = vocab_char_map.get
|
||
list_idx_tensors = [torch.tensor([get_idx(c, 0) for c in t], dtype=torch.int32) for t in text]
|
||
text = torch.nn.utils.rnn.pad_sequence(list_idx_tensors, padding_value=padding_value, batch_first=True)
|
||
return text
|
||
|
||
|
||
def normalize_to_int16(audio):
|
||
max_val = np.max(np.abs(audio))
|
||
scaling_factor = 32767.0 / max_val if max_val > 0 else 1.0
|
||
return (audio * float(scaling_factor)).astype(np.int16)
|
||
|
||
|
||
# ONNX Runtime settings
|
||
onnxruntime.set_seed(RANDOM_SEED)
|
||
session_opts = onnxruntime.SessionOptions()
|
||
session_opts.log_severity_level = 4 # fatal level = 4, it an adjustable value.
|
||
session_opts.log_verbosity_level = 4 # fatal level = 4, it an adjustable value.
|
||
session_opts.inter_op_num_threads = MAX_THREADS # Run different nodes with num_threads. Set 0 for auto.
|
||
session_opts.intra_op_num_threads = MAX_THREADS # Under the node, execute the operators with num_threads. Set 0 for auto.
|
||
session_opts.enable_cpu_mem_arena = True # True for execute speed; False for less memory usage.
|
||
session_opts.execution_mode = onnxruntime.ExecutionMode.ORT_SEQUENTIAL
|
||
session_opts.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL
|
||
session_opts.add_session_config_entry("session.intra_op.allow_spinning", "1")
|
||
session_opts.add_session_config_entry("session.inter_op.allow_spinning", "1")
|
||
session_opts.add_session_config_entry("session.set_denormal_as_zero", "1")
|
||
|
||
session_opts.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL
|
||
ort_session_A = onnxruntime.InferenceSession(onnx_model_A, sess_options=session_opts, providers=['CPUExecutionProvider'], provider_options=None)
|
||
model_type = ort_session_A._inputs_meta[0].type
|
||
in_name_A = ort_session_A.get_inputs()
|
||
out_name_A = ort_session_A.get_outputs()
|
||
in_name_A0 = in_name_A[0].name
|
||
in_name_A1 = in_name_A[1].name
|
||
in_name_A2 = in_name_A[2].name
|
||
out_name_A0 = out_name_A[0].name
|
||
out_name_A1 = out_name_A[1].name
|
||
out_name_A2 = out_name_A[2].name
|
||
out_name_A3 = out_name_A[3].name
|
||
out_name_A4 = out_name_A[4].name
|
||
out_name_A5 = out_name_A[5].name
|
||
out_name_A6 = out_name_A[6].name
|
||
out_name_A7 = out_name_A[7].name
|
||
|
||
if "CPUExecutionProvider" in ORT_Accelerate_Providers or not ORT_Accelerate_Providers:
|
||
session_opts.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL
|
||
else:
|
||
session_opts.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_BASIC
|
||
ort_session_B = onnxruntime.InferenceSession(onnx_model_B, sess_options=session_opts, providers=ORT_Accelerate_Providers, provider_options=provider_options)
|
||
ORT_Accelerate_Providers = ort_session_B.get_providers()[0]
|
||
# For Windows DirectML + Intel/AMD/Nvidia GPU,
|
||
# pip install onnxruntime-directml --upgrade
|
||
# ort_session_B = onnxruntime.InferenceSession(onnx_model_B, sess_options=session_opts, providers=['DmlExecutionProvider'])
|
||
print(f"\nUsable Providers: {ORT_Accelerate_Providers}")
|
||
model_dtype = ort_session_B._inputs_meta[0].type
|
||
in_name_B = ort_session_B.get_inputs()
|
||
out_name_B = ort_session_B.get_outputs()
|
||
in_name_B0 = in_name_B[0].name
|
||
in_name_B1 = in_name_B[1].name
|
||
in_name_B2 = in_name_B[2].name
|
||
in_name_B3 = in_name_B[3].name
|
||
in_name_B4 = in_name_B[4].name
|
||
in_name_B5 = in_name_B[5].name
|
||
in_name_B6 = in_name_B[6].name
|
||
in_name_B7 = in_name_B[7].name
|
||
out_name_B0 = out_name_B[0].name
|
||
out_name_B1 = out_name_B[1].name
|
||
|
||
session_opts.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL
|
||
ort_session_C = onnxruntime.InferenceSession(onnx_model_C, sess_options=session_opts, providers=['CPUExecutionProvider'], provider_options=None)
|
||
in_name_C = ort_session_C.get_inputs()
|
||
out_name_C = ort_session_C.get_outputs()
|
||
in_name_C0 = in_name_C[0].name
|
||
in_name_C1 = in_name_C[1].name
|
||
out_name_C0 = out_name_C[0].name
|
||
|
||
# Load the input audio
|
||
print(f"\nReference Audio: {reference_audio}")
|
||
audio = np.array(AudioSegment.from_file(reference_audio).set_channels(1).set_frame_rate(MODEL_SAMPLE_RATE).get_array_of_samples(), dtype=np.float32)
|
||
audio = normalize_to_int16(audio)
|
||
audio_len = len(audio)
|
||
audio = audio.reshape(1, 1, -1)
|
||
|
||
zh_pause_punc = r"。，、；：？！"
|
||
ref_text_len = len(ref_text.encode('utf-8')) + 3 * len(re.findall(zh_pause_punc, ref_text))
|
||
gen_text_len = len(gen_text.encode('utf-8')) + 3 * len(re.findall(zh_pause_punc, gen_text))
|
||
ref_audio_len = audio_len // HOP_LENGTH + 1
|
||
max_duration = np.array([ref_audio_len + int(ref_audio_len / ref_text_len * gen_text_len / SPEED)], dtype=np.int64)
|
||
gen_text = convert_char_to_pinyin([ref_text + gen_text])
|
||
text_ids = list_str_to_idx(gen_text, vocab_char_map).numpy()
|
||
time_step = np.array([0], dtype=np.int32)
|
||
|
||
if "CPUExecutionProvider" in ORT_Accelerate_Providers or not ORT_Accelerate_Providers:
|
||
device_type = 'cpu'
|
||
elif "CUDAExecutionProvider" in ORT_Accelerate_Providers or "TensorrtExecutionProvider" in ORT_Accelerate_Providers:
|
||
device_type = 'cuda'
|
||
elif "DmlExecutionProvider" in ORT_Accelerate_Providers:
|
||
device_type = 'dml'
|
||
else:
|
||
device_type = None
|
||
|
||
def run_pipeline():
|
||
a_out = ort_session_A.run(
|
||
[out_name_A0, out_name_A1, out_name_A2, out_name_A3, out_name_A4, out_name_A5, out_name_A6, out_name_A7],
|
||
{in_name_A0: audio, in_name_A1: text_ids, in_name_A2: max_duration})
|
||
noise, rope_cos_q, rope_sin_q, rope_cos_k, rope_sin_k, cat_mel_text, cat_mel_text_drop, ref_signal_len = a_out
|
||
ts = np.array([0], dtype=np.int32)
|
||
if device_type:
|
||
inputs = [onnxruntime.OrtValue.ortvalue_from_numpy(x, device_type, DEVICE_ID) for x in
|
||
(noise, rope_cos_q, rope_sin_q, rope_cos_k, rope_sin_k, cat_mel_text, cat_mel_text_drop, ts)]
|
||
outputs = [inputs[0], inputs[-1]]
|
||
iob = ort_session_B.io_binding()
|
||
for i in range(len(inputs)):
|
||
iob.bind_ortvalue_input(name=in_name_B[i].name, ortvalue=inputs[i])
|
||
for i in range(len(outputs)):
|
||
iob.bind_ortvalue_output(name=out_name_B[i].name, ortvalue=outputs[i])
|
||
for _ in range(0, NFE_STEP, FUSE_NFE):
|
||
ort_session_B.run_with_iobinding(iob)
|
||
noise = onnxruntime.OrtValue.numpy(iob.get_outputs()[0])
|
||
else:
|
||
for _ in range(0, NFE_STEP - 1, FUSE_NFE):
|
||
noise, ts = ort_session_B.run(
|
||
[out_name_B0, out_name_B1],
|
||
{in_name_B0: noise, in_name_B1: rope_cos_q, in_name_B2: rope_sin_q, in_name_B3: rope_cos_k,
|
||
in_name_B4: rope_sin_k, in_name_B5: cat_mel_text, in_name_B6: cat_mel_text_drop, in_name_B7: ts})
|
||
return ort_session_C.run([out_name_C0], {in_name_C0: noise, in_name_C1: ref_signal_len})[0]
|
||
|
||
print(f"\nProvider={ORT_Accelerate_Providers} device_type={device_type} NFE={NFE_STEP}")
|
||
print("Warmup (DML kompiliert beim 1. Lauf die Shader) ...")
|
||
t0 = time.time(); _ = run_pipeline(); print(f" warmup gen = {time.time()-t0:.2f}s")
|
||
best = 1e9
|
||
for k in range(2):
|
||
t0 = time.time(); gen = run_pipeline(); dt = time.time() - t0; best = min(best, dt)
|
||
print(f" run {k+1}: gen = {dt:.2f}s")
|
||
audio_s = gen.reshape(-1).shape[0] / MODEL_SAMPLE_RATE
|
||
sf.write(generated_audio, gen.reshape(-1), MODEL_SAMPLE_RATE, format='WAVEX')
|
||
rtf = best / max(audio_s, 0.01)
|
||
print(f"\n=== ERGEBNIS NFE={NFE_STEP} === audio={audio_s:.2f}s gen(best)={best:.2f}s RTF={rtf:.2f} "
|
||
f"({'REAL-TIME' if rtf < 1 else 'zu langsam'})")
|