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'})")