7f59912916
The model code was spread across eight crates that had grown into each other:
ggml and cuda and mlx each owned part of a tensor runtime, llama and tts and
voice2 each owned part of a model, and libs/diffusion owned everything else.
They are now one tree with an explicit shape:
libs/ai/cuda — kernels and launch surface
libs/ai/metal — Metal shaders and the shim
libs/ai/llm — the language-model runtime (sessions, lanes, contexts,
the CUDA and Metal executors, the compiled Metal path)
libs/ai/models/ — common, flux, h3, music, paint, speech, stems, vision
libs/diffusion is not deleted but demoted: what remains is the VALIDATOR
crate — several dozen `*_validate.rs` oracles that check a native
implementation against a reference, which is where they belong now that the
implementations live next door.
The functional work inside the move is mostly in the LLM runtime: N lanes that
draft while one verify batch serves all of them, per-slot prefill over a shared
folded attention arena, speculation that survives batching, and a scheduler
that reports rather than publishes. And in the CUDA build: a machine without
usable CUDA must still LINK (and say so), the default kernel arch is the
building machine's GPU, `NO_CUDA` forces the stub even where the toolkit
exists, and kernels compile in parallel with progress.
libs/video_flow is new here: classical optical flow estimation and the `mkfl`
motion-field payload — a flow field measured from a clip without a model,
which is what drives free-rate bounce-looping playback and the uprez/tween
enhance pipe.
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()
|