Kokorko ONNX Timestamped - Phoneme Level Timestamps

Phoneme.py

from pathlib import Path
from misaki import en
import numpy as np
import onnxruntime

def tokenize(phonemes):
    """
    Converts a list of phonemes (strings) into a list of their corresponding
    integer IDs based on the global VOCAB dictionary.
    Filters out any phonemes not found in the vocabulary.
    """
    return [i for i in map(VOCAB.get, phonemes) if i is not None]

def join_timestamps_phonemes_and_words(tokens, pred_dur):
    """
    Calculates and returns both phoneme-level and word-level timestamps
    based on the predicted durations from the ONNX model.

    Args:
        tokens (list): A list of Token objects from misaki.en.G2P,
                       each containing text, phonemes, and whitespace info.
        pred_dur (np.ndarray): A 1D numpy array of predicted durations (in model frames)
                               for each phoneme/unit, typically from the ONNX model output.

    Returns:
        tuple: A tuple containing two lists:
               - phoneme_timestamps (list of dict): Each dict contains 'phoneme',
                 'start_ts', 'end_ts', and 'word' for individual phonemes.
               - word_timestamps (list of dict): Each dict contains 'word',
                 'start_ts', and 'end_ts' for each word/token.
    """
    # Multiply by 600 to go from pred_dur frames to sample_rate 24000
    # Equivalent to dividing pred_dur frames by 40 to get timestamp in seconds
    # We will count nice round half-frames, so the divisor is 80 (24000 / 40 / 2)
    MAGIC_DIVISOR = 80

    if not tokens or len(pred_dur) < 3:
        # Return empty lists if input is invalid or too short
        return [], []

    phoneme_timestamps = [] # To store {'phoneme', 'start_ts', 'end_ts', 'word'}
    word_timestamps = []    # To store {'word', 'start_ts', 'end_ts'}

    # Initialize current timestamp pointer.
    # The 'max(0, pred_dur[0].item() - 3)' is an initial offset to account for
    # potential <bos> tokens or initial padding in pred_dur.
    left = right = max(0, pred_dur[0].item() - 3)
    i = 1 # Current index in pred_dur, starting after the initial offset/BOS

    for t_idx, t in enumerate(tokens):
        if i >= len(pred_dur) - 1:
            # Break if we've consumed all available duration predictions
            break

        word_start_ts = left / MAGIC_DIVISOR # Mark the start of the current word

        if not t.phonemes:
            # Handle non-phonemic tokens (like whitespace-only or punctuation without IPA)
            if t.whitespace and i < len(pred_dur):
                # If it's a whitespace token, consume its duration
                space_dur = pred_dur[i].item()
                left = right + space_dur
                right = left
                i += 1
                # For whitespace tokens, their duration is their own 'word' time
                word_timestamps.append({
                    "word": t.text,
                    "start_ts": word_start_ts,
                    "end_ts": left / MAGIC_DIVISOR
                })
            continue # Move to the next token if no phonemes to process

        # Process each phoneme within the current token
        for k, phoneme in enumerate(t.phonemes):
            if i >= len(pred_dur):
                # Ensure we don't go out of bounds of pred_dur
                break

            phoneme_start_ts = left / MAGIC_DIVISOR
            
            # Duration for the current phoneme (in model frames)
            phoneme_dur = pred_dur[i].item() 
            
            # Calculate phoneme end timestamp, converting frames to seconds
            # Multiply by 2 because 'left'/'right' are in 'half-frames'
            phoneme_end_ts = (left + (2 * phoneme_dur)) / MAGIC_DIVISOR 

            phoneme_timestamps.append({
                "phoneme": phoneme, 
                "start_ts": phoneme_start_ts, 
                "end_ts": phoneme_end_ts,
                "word": t.text # Link phoneme back to its parent word
            })
            
            # Update the current timestamp pointer for the next phoneme
            left = left + (2 * phoneme_dur) 
            right = left # For simplicity, right always matches left after a phoneme

            i += 1 # Move to the next duration prediction in pred_dur

        # After processing all phonemes for the current token, handle trailing whitespace
        # if the token itself indicates whitespace.
        if t.whitespace and i < len(pred_dur):
            space_dur = pred_dur[i].item()
            # The space duration is added after the phonemes of the current token
            left = right + space_dur 
            right = left 
            i += 1
        elif not t.whitespace and k == len(t.phonemes) - 1:
            # If no whitespace and this was the last phoneme of the token,
            # ensure 'right' is updated to where the next token would start
            right = left 
        
        # Mark the end of the current word after all its phonemes and trailing whitespace
        word_end_ts = left / MAGIC_DIVISOR
        word_timestamps.append({
            "word": t.text,
            "start_ts": word_start_ts,
            "end_ts": word_end_ts
        })
    
    return phoneme_timestamps, word_timestamps

# --- VOCABULARY DEFINITION (from your original code) ---
def get_vocab():
    _pad = "$"
    _punctuation = ';:,.!?¡¿—…"«»“” '
    _letters = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
    _letters_ipa = "ɑɐɒæɓʙβɔɕçɗɖðʤəɘɚɛɜɝɞɟʄɡɠɢʛɦɧħɥʜɨɪʝɭɬɫɮʟɱɯɰŋɳɲɴøɵɸθœɶʘɹɺɾɻʀʁɽʂʃʈʧʉʊʋⱱʌɣɤʍχʎʏʑʐʒʔʡʕʢǀǁǂǃˈˌːˑʼʴʰʱʲʷˠˤ˞↓↑→↗↘'̩'ᵻ"
    _symbols = [_pad] + list(_punctuation) + list(_letters) + list(_letters_ipa)
    dicts = {}
    for i in range(len((_symbols))):
        dicts[_symbols[i]] = i
    return dicts
VOCAB = get_vocab()

if __name__ == "__main__":
    g2p = en.G2P(trf=False, british=False, fallback=None)
    text = "Yesterday's the past, tomorrow's the future, but today is a gift. That's why it's called the present."
    
    # --- Load voice and prepare style input ---
    # NOTE: Paths are hardcoded here. Consider making them configurable.
    voice_path = "kokoro_models\\voices\\af_heart.bin"
    model_path = "kokoro_models\\model_q8f16.onnx"

    try:
        voice = Path(voice_path).read_bytes()
        voice = np.frombuffer(voice, dtype=np.float32)
    except FileNotFoundError:
        print(f"Error: Voice file not found at {voice_path}. Please check the path.")
        exit()

    sample_rate = 24000

    # Get phonemes (as string) and tokens (list of Token objects)
    phonemes_str, tokens = g2p(text) 
    
    # --- Initialize ONNX Runtime Session ---
    try:
        providers = ["CPUExecutionProvider"] # Consider "CUDAExecutionProvider" if you have a GPU
        sess = onnxruntime.InferenceSession(model_path, providers=providers)
    except Exception as e:
        print(f"Error initializing ONNX Runtime session: {e}")
        print(f"Please ensure the model file exists at {model_path} and ONNX Runtime is correctly installed.")
        exit()
    
    input_ids = tokenize(phonemes_str) # Convert phoneme string to integer IDs

    # --- Prepare style input for the ONNX model ---
    style_input_length = 256
    # Calculate offset for style input. This logic assumes the 'style' input
    # is a fixed-size segment from the beginning of the voice audio.
    # The 'len(tokens) * 40' is a rough heuristic if voice is very short.
    offset = min(len(voice) - style_input_length, len(tokens) * 40 if len(tokens) * 40 < len(voice) else len(voice) - style_input_length)
    offset = max(0, offset) # Ensure offset is non-negative
    voice_style_input = voice[offset : offset + style_input_length].reshape(1, style_input_length)

    # Add <bos> (0) and <eos> (0) tokens to the input_ids as expected by the model
    input_ids_for_model = [[0, *input_ids, 0]] 

    # --- Run Inference ---
    try:
        (audio, pred_dur) = sess.run(
            None, # Outputs to retrieve (None means all outputs)
            dict(input_ids=np.array(input_ids_for_model, dtype=np.int64), # Ensure correct dtype
                 style=voice_style_input, 
                 speed=np.ones(1, dtype=np.float32)), # Speed input, typically 1.0 for normal speed
        )
    except Exception as e:
        print(f"Error during ONNX model inference: {e}")
        print("Please check input shapes and dtypes expected by your ONNX model.")
        exit()

    # pred_dur is likely a batch output (e.g., (1, N)), so take the first element
    phoneme_times, word_times = join_timestamps_phonemes_and_words(tokens, pred_dur[0]) 

    # --- Print Results ---
    print("--- Phoneme-Level Timestamps ---")
    for p_info in phoneme_times:
        print(f"[{p_info['start_ts']:.4f} - {p_info['end_ts']:.4f}] {p_info['phoneme']} (Word: {p_info['word']})")

    print("\n--- Word-Level Timestamps ---")
    for w_info in word_times:
        print(f"[{w_info['start_ts']:.4f} - {w_info['end_ts']:.4f}] {w_info['word']}")