449f51be7a
libs/tts was never tracked despite being a build dependency of the gamemaker example. tools/download_tts.sh fetches the public upstream weights (HuggingFace Kokoro-82M + whisper.cpp) and converts them locally with the in-repo stdlib-only converter; model artifacts are gitignored. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
101 lines
3.5 KiB
Python
101 lines
3.5 KiB
Python
#!/usr/bin/env python3
|
|
"""Run the reference Kokoro ONNX on tokens produced by *our* Rust phonemizer.
|
|
|
|
Ground truth for the port: same text, same tokens, same voice. Anything the Rust
|
|
graph produces later gets diffed against this.
|
|
|
|
/tmp/kref/bin/python libs/tts/tools/ref_infer.py "Escape the Gummer." out.wav
|
|
"""
|
|
|
|
import struct
|
|
import subprocess
|
|
import sys
|
|
|
|
import numpy as np
|
|
import onnxruntime as ort
|
|
|
|
MODEL = "kokoro_ref.onnx"
|
|
VOICE = "af_heart.mkvoice"
|
|
SAMPLE_RATE = 24_000
|
|
|
|
|
|
def rust_tokens(text):
|
|
"""Ask the Rust binary, so the reference sees exactly what Rust will feed it."""
|
|
out = subprocess.run(
|
|
["cargo", "run", "--release", "--quiet",
|
|
"--manifest-path", "libs/tts/Cargo.toml", "--bin", "g2p_test", "--", "--ids", text],
|
|
capture_output=True, text=True, check=True,
|
|
)
|
|
return [int(x) for x in out.stdout.strip().split(",")]
|
|
|
|
|
|
def load_voice(path):
|
|
"""Read the `.mktts` container: one [510, 1, 256] tensor."""
|
|
blob = open(path, "rb").read()
|
|
assert blob[:8] == b"MKTTS\0\0\0", "bad voice magic"
|
|
count = struct.unpack("<I", blob[12:16])[0]
|
|
assert count == 1, count
|
|
at = 16
|
|
name_len = struct.unpack("<I", blob[at:at + 4])[0]
|
|
at += 4 + name_len
|
|
_dtype, ndim = blob[at], blob[at + 1]
|
|
at += 2
|
|
shape = struct.unpack(f"<{ndim}I", blob[at:at + 4 * ndim])
|
|
at += 4 * ndim
|
|
offset, nbytes = struct.unpack("<QQ", blob[at:at + 16])
|
|
data = np.frombuffer(blob[offset:offset + nbytes], dtype=np.float32)
|
|
return data.reshape(shape)
|
|
|
|
|
|
def write_wav(path, samples, rate):
|
|
pcm = np.clip(samples, -1.0, 1.0)
|
|
pcm = (pcm * 32767.0).astype("<i2")
|
|
with open(path, "wb") as out:
|
|
out.write(b"RIFF")
|
|
out.write(struct.pack("<I", 36 + pcm.nbytes))
|
|
out.write(b"WAVEfmt ")
|
|
out.write(struct.pack("<IHHIIHH", 16, 1, 1, rate, rate * 2, 2, 16))
|
|
out.write(b"data")
|
|
out.write(struct.pack("<I", pcm.nbytes))
|
|
out.write(pcm.tobytes())
|
|
|
|
|
|
def main():
|
|
text = sys.argv[1] if len(sys.argv) > 1 else "Escape the Gummer, a squishy purple blob."
|
|
dst = sys.argv[2] if len(sys.argv) > 2 else "kokoro_ref.wav"
|
|
|
|
ids = rust_tokens(text)
|
|
phonemes = len(ids) - 2 # the ids are zero-padded at both ends
|
|
voice = load_voice(VOICE)
|
|
# One style vector per phoneme count. Row `phonemes`, both 128-halves.
|
|
# kokoro/pipeline.py uses `pack[len(ps) - 1]`. Row `phonemes` also sounds
|
|
# fine, which is exactly how this off-by-one survives a listening test.
|
|
style = voice[phonemes - 1]
|
|
print(f"text : {text}")
|
|
print(f"tokens : {len(ids)} ids ({phonemes} phonemes)")
|
|
print(f"style row: {phonemes - 1} of {voice.shape[0]} -> {style.shape}")
|
|
|
|
session = ort.InferenceSession(MODEL, providers=["CPUExecutionProvider"])
|
|
waveform = session.run(
|
|
None,
|
|
{
|
|
"input_ids": np.array([ids], dtype=np.int64),
|
|
"style": style.astype(np.float32).reshape(1, 256),
|
|
"speed": np.array([1.0], dtype=np.float32),
|
|
},
|
|
)[0][0]
|
|
|
|
print(f"waveform : {waveform.shape[0]} samples "
|
|
f"({waveform.shape[0]/SAMPLE_RATE:.2f}s), peak={np.abs(waveform).max():.4f}")
|
|
write_wav(dst, waveform, SAMPLE_RATE)
|
|
print(f"wrote : {dst}")
|
|
|
|
# A 16k copy so Whisper can score it with the same harness.
|
|
index = np.arange(0, len(waveform), SAMPLE_RATE / 16_000)
|
|
resampled = np.interp(index, np.arange(len(waveform)), waveform)
|
|
write_wav(dst.replace(".wav", "_16k.wav"), resampled, 16_000)
|
|
np.save(dst.replace(".wav", ".npy"), waveform)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|