F5-TTS-ONNX-Inference-rknn2.py

# import ztu_ort_enable_debug
import re
import time
import jieba
import torch
import ztu_somemodelruntime_rknnlite2 as rknnort
import onnxruntime
import soundfile as sf
import numpy as np
from pydub import AudioSegment
from pypinyin import lazy_pinyin, Style
from tqdm import tqdm

vocab_path           = "vocab.txt"                             # The F5-TTS model vocab download path.     URL: https://huggingface.co/SWivid/F5-TTS/tree/main/F5TTS_v1_Base
onnx_model_A         = "F5_Preprocess.onnx"                     # The exported onnx model path.
onnx_model_B         = "F5_Transformer_opset19.onnx"                    # The exported onnx model path.
onnx_model_C         = "F5_Decode.onnx"                         # The exported onnx model path.
generated_audio      = "./generated_audio-rknn2.wav"
test_in_english = False

if test_in_english:
    reference_audio  = "basic_ref_en.wav"
    ref_text         = "Some call me nature, others call me mother nature."
    gen_text         = "Some call me Dake, others call me QQ."
else:
    reference_audio  = "basic_ref_zh.wav"        # The reference audio path.
    ref_text         = "对，这就是我，万人敬仰的太乙真人。"                                               # The ASR result of reference audio.
    gen_text         = "哇，这个辣鸡芯片居然也能运行模型。"                                               # The target TTS.


# RANDOM_SEED = 9527                                        # Set seed to reproduce the generated audio
NFE_STEP = 32                                             # F5-TTS model setting, 0~31
FUSE_NFE = 1                                              # Maintain the same values as the exported model.
SPEED = 1.0                                               # Set for talking speed. Only works with dynamic_axes=True
MODEL_SAMPLE_RATE = 24000                                 # Do not modify it.
HOP_LENGTH = 256                                          # It affects the generated audio length and speech speed.



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)

ort_session_A = onnxruntime.InferenceSession(onnx_model_A, 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


ort_session_B = rknnort.InferenceSession(onnx_model_B)

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

ort_session_C = onnxruntime.InferenceSession(onnx_model_C, 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) # Normalization should happen on the final output, not here
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


original_max_duration_val = ref_audio_len + int(ref_audio_len / ref_text_len * gen_text_len / SPEED)
fixed_max_duration_val = 1536 # The hardcoded duration
speed_adjustment_ratio = fixed_max_duration_val / original_max_duration_val

print(f"Original estimated duration: {original_max_duration_val} frames")
print(f"Fixed duration: {fixed_max_duration_val} frames")
print(f"Speed adjustment ratio: {speed_adjustment_ratio:.3f}")

# Force max_duration to the fixed value for the model input
max_duration = np.array([fixed_max_duration_val], dtype=np.int64)

print(f"zt: hack: force max_duration to {fixed_max_duration_val}")
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)


device_type = None

print("\n\nRun F5-TTS.")
start_count = time.time()
noise, rope_cos_q, rope_sin_q, rope_cos_k, rope_sin_k, cat_mel_text, cat_mel_text_drop, ref_signal_len = 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
        })
end_count = time.time()
print(f"\nPrepare input data time cost: {end_count - start_count:.3f} seconds")

print("NFE_STEP: 0")
for i in tqdm(range(0, NFE_STEP - 1, FUSE_NFE)):
    noise, time_step = 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: time_step
        })
    print(f"NFE_STEP: {i + FUSE_NFE}")

start_count = time.time()
generated_signal = ort_session_C.run(
        [out_name_C0],
        {
            in_name_C0: noise,
            in_name_C1: ref_signal_len
        })[0]
end_count = time.time()
print(f"Decode time cost: {end_count - start_count:.3f} seconds")

# Adjust speed of the generated audio
print("\nAdjusting audio speed...")
# Flatten signal and ensure it's float32 before normalization
generated_signal_flat = generated_signal.flatten().astype(np.float32)

# Normalize the generated signal to int16 for pydub
generated_signal_int16 = normalize_to_int16(generated_signal_flat)

# Create an AudioSegment from the int16 numpy array
try:
    segment = AudioSegment(
        generated_signal_int16.tobytes(),
        frame_rate=MODEL_SAMPLE_RATE,
        sample_width=generated_signal_int16.dtype.itemsize,
        channels=1
    )

    # Apply speed adjustment
    # Note: speedup > 1.0 means faster, < 1.0 means slower
    if speed_adjustment_ratio != 1.0:
         print(f"Applying speed change with factor: {speed_adjustment_ratio:.3f}")
         adjusted_segment = segment.speedup(playback_speed=speed_adjustment_ratio)
    else:
         print("No speed adjustment needed.")
         adjusted_segment = segment # No adjustment needed

    # Save the adjusted audio using pydub's export
    adjusted_segment.export(generated_audio, format="wav")
    print(f"\nSpeed-adjusted audio saved to: {generated_audio}")

except Exception as e:
    print(f"Error during audio speed adjustment or saving: {e}")
    print("Saving original audio instead.")
    # Fallback to saving the original audio if adjustment fails
    sf.write(generated_audio, generated_signal_flat, MODEL_SAMPLE_RATE, format='WAVEX')

print(f"\nAudio generation and processing is complete.\n\nONNXRuntime Time Cost in Seconds:\n{end_count - start_count:.3f}")