varsan-g/plapre-nano-ONNX

onnx/
  model_prompt.onnx          # Prompt encoder (FP32)
  model_prompt.onnx.data
  model_prompt_fp16.onnx     # Prompt encoder (FP16)
  model_prompt_fp16.onnx.data
  model_prompt_int8.onnx     # Prompt encoder (INT8)
  model_prompt_int4.onnx     # Prompt encoder (INT4)
  model_prompt_int4.onnx.data
  model_decode.onnx          # Autoregressive decoder (FP32)
  model_decode.onnx.data
  model_decode_fp16.onnx     # Autoregressive decoder (FP16)
  model_decode_fp16.onnx.data
  model_decode_int8.onnx     # Autoregressive decoder (INT8)
  model_decode_int4.onnx     # Autoregressive decoder (INT4)
  model_decode_int4.onnx.data
  speaker_proj.onnx          # Speaker embedding projection
  speaker_proj.onnx.data
config.json
tokenizer.json
tokenizer_config.json
special_tokens_map.json
speakers.json
embed_tokens.npy             # Embedding lookup table (numpy)

pip install onnxruntime numpy transformers soundfile kanade-tokenizer

python inference_onnx.py "Hej, hvordan har du det?" --model nano
python inference_onnx.py "Godmorgen!" --model nano --speaker ida --output greeting.wav
python inference_onnx.py "Hej verden." --model nano --quant int8  # INT8 quantized

import json
import numpy as np
import onnxruntime as ort
from transformers import PreTrainedTokenizerFast

# --- Load models (use _int8 / _fp16 / _int4 suffix for quantized variants) ---
prompt_sess = ort.InferenceSession("onnx/model_prompt.onnx")
decode_sess = ort.InferenceSession("onnx/model_decode.onnx")
speaker_sess = ort.InferenceSession("onnx/speaker_proj.onnx")
embed_table = np.load("embed_tokens.npy")

# --- Load tokenizer (AutoTokenizer won't work due to custom tokenizer_class) ---
with open("tokenizer_config.json") as f:
    tok_cfg = json.load(f)
tokenizer = PreTrainedTokenizerFast(
    tokenizer_file="tokenizer.json",
    eos_token=tok_cfg.get("eos_token", "<eos>"),
    pad_token=tok_cfg.get("pad_token"),
)
TEXT_TAG = tokenizer.convert_tokens_to_ids("<text>")
AUDIO_TAG = tokenizer.convert_tokens_to_ids("<audio>")
AUDIO_START = tokenizer.convert_tokens_to_ids("<audio_0>")
AUDIO_END = tokenizer.convert_tokens_to_ids("<audio_12799>")
EOS_ID = 0  # <eos> token, MUST be 0, not 2

# --- Load speaker embedding ---
with open("speakers.json") as f:
    speakers = json.load(f)
speaker_emb = np.array(speakers["dv"], dtype=np.float32)  # or "ida"

# --- Sampling parameters (these are critical!) ---
TEMPERATURE = 0.6
TOP_K = 50
TOP_P = 0.95
REPETITION_PENALTY = 1.3
MAX_TOKENS = 500

# --- 1. Project speaker embedding to hidden size ---
speaker_hidden = speaker_sess.run(None, {"input": speaker_emb[np.newaxis, :]})[0]

# --- 2. Tokenize and build inputs_embeds ---
text = "Hej, hvordan har du det?"
text_ids = tokenizer.encode(text, add_special_tokens=False)
prompt_ids = [TEXT_TAG] + text_ids + [AUDIO_TAG]
token_embeds = embed_table[prompt_ids]  # [seq_len, H]

inputs_embeds = np.concatenate([
    speaker_hidden[:, np.newaxis, :],   # [1, 1, H]
    token_embeds[np.newaxis, :, :],     # [1, seq_len, H]
], axis=1).astype(np.float32)           # [1, 1+seq_len, H]

total_len = inputs_embeds.shape[1]
attention_mask = np.ones((1, total_len), dtype=np.int64)

# --- 3. Prompt pass ---
outputs = prompt_sess.run(None, {
    "inputs_embeds": inputs_embeds,
    "attention_mask": attention_mask,
})
logits = outputs[0]
past_kvs = outputs[1:]

# Read KV cache input names from decode session
kv_names = [inp.name for inp in decode_sess.get_inputs()][3:]

# --- 4. Autoregressive decode loop ---
generated = []
expected_audio = len(text_ids) * 10
eos_boost_threshold = max(int(expected_audio * 1.5), 50)
audio_count = 0

for step in range(MAX_TOKENS):
    current_logits = logits[0, -1, :].copy()

    # EOS boost: after expected duration, progressively boost EOS logit
    if audio_count > eos_boost_threshold:
        current_logits[EOS_ID] += (audio_count - eos_boost_threshold) * 2.0

    # Repetition penalty: penalize already-generated tokens
    if REPETITION_PENALTY > 1.0 and generated:
        for tid in set(generated):
            if current_logits[tid] > 0:
                current_logits[tid] /= REPETITION_PENALTY
            else:
                current_logits[tid] *= REPETITION_PENALTY

    # Temperature + top-k + top-p sampling
    current_logits /= TEMPERATURE
    if TOP_K > 0:
        top_k_idx = np.argsort(current_logits)[-TOP_K:]
        mask = np.full_like(current_logits, -np.inf)
        mask[top_k_idx] = current_logits[top_k_idx]
        current_logits = mask
    current_logits -= np.max(current_logits)
    probs = np.exp(current_logits) / np.sum(np.exp(current_logits))
    if TOP_P < 1.0:
        sorted_idx = np.argsort(probs)[::-1]
        cumsum = np.cumsum(probs[sorted_idx])
        keep = sorted_idx[:np.searchsorted(cumsum, TOP_P) + 1]
        filtered = np.zeros_like(probs)
        filtered[keep] = probs[keep]
        probs = filtered / filtered.sum()

    next_token = int(np.random.choice(len(probs), p=probs))

    # Stop on EOS
    if next_token == EOS_ID:
        break
    generated.append(next_token)

    # Track audio tokens
    if AUDIO_START <= next_token <= AUDIO_END:
        audio_count += 1

    # Loop detection: 8 identical tokens = degenerate loop, stop
    if len(generated) >= 8 and len(set(generated[-8:])) == 1:
        break

    # Prepare next decode step
    new_embed = embed_table[next_token][np.newaxis, np.newaxis, :].astype(np.float32)
    total_len += 1
    attention_mask = np.ones((1, total_len), dtype=np.int64)
    cache_pos = np.array([total_len - 1], dtype=np.int64)

    feed = {"inputs_embeds": new_embed, "attention_mask": attention_mask, "cache_position": cache_pos}
    for name, kv in zip(kv_names, past_kvs):
        feed[name] = kv

    outputs = decode_sess.run(None, feed)
    logits = outputs[0]
    past_kvs = outputs[1:]

# --- 5. Extract audio tokens ---
audio_tokens = [t - AUDIO_START for t in generated if AUDIO_START <= t <= AUDIO_END]
print(f"Generated {len(generated)} tokens ({len(audio_tokens)} audio)")

# --- 6. Decode to audio (requires PyTorch + kanade-tokenizer) ---
import torch
from kanade_tokenizer import KanadeModel, load_vocoder, vocode

kanade = KanadeModel.from_pretrained("frothywater/kanade-25hz-clean").eval()
vocoder = load_vocoder(kanade.config.vocoder_name)

with torch.no_grad():
    mel = kanade.decode(
        content_token_indices=torch.tensor(audio_tokens),
        global_embedding=torch.tensor(speaker_emb),
    )
    waveform = vocode(vocoder, mel.unsqueeze(0)).squeeze().numpy()

import soundfile as sf
sf.write("output.wav", waveform, 24000)