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.
109 lines
3.3 KiB
Python
109 lines
3.3 KiB
Python
#!/usr/bin/env python3
|
|
"""Expose Kokoro's internal ONNX tensors so the Rust port can be diffed stage by stage.
|
|
|
|
Without this, a wrong sign somewhere in the prosody predictor shows up only as
|
|
"the audio sounds a bit off" — 2,463 nodes downstream. With it, each stage is
|
|
checked where it is written.
|
|
|
|
# list candidate tensors
|
|
/tmp/kref/bin/python libs/tts/tools/ref_dump.py --list lstm
|
|
|
|
# dump matching tensors as .npy under refdump/
|
|
/tmp/kref/bin/python libs/tts/tools/ref_dump.py --dump text_encoder "Hello there."
|
|
"""
|
|
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
|
|
import numpy as np
|
|
import onnx
|
|
import onnxruntime as ort
|
|
|
|
MODEL = "kokoro_ref.onnx"
|
|
VOICE = "af_heart.mkvoice"
|
|
OUT_DIR = "refdump"
|
|
|
|
|
|
def rust_tokens(text):
|
|
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():
|
|
sys.path.insert(0, os.path.dirname(__file__))
|
|
from ref_infer import load_voice as load # reuse the container reader
|
|
|
|
return load(VOICE)
|
|
|
|
|
|
def candidates(model):
|
|
"""Every value produced by a node, in graph order."""
|
|
return [out for node in model.graph.node for out in node.output if out]
|
|
|
|
|
|
def main():
|
|
if len(sys.argv) < 3:
|
|
raise SystemExit(__doc__)
|
|
mode, pattern = sys.argv[1], sys.argv[2]
|
|
text = sys.argv[3] if len(sys.argv) > 3 else "Escape the Gummer."
|
|
|
|
model = onnx.load(MODEL)
|
|
names = [n for n in candidates(model) if pattern.lower() in n.lower()]
|
|
|
|
if mode == "--list":
|
|
print(f"{len(names)} tensors matching {pattern!r}:")
|
|
for name in names[:60]:
|
|
print(f" {name}")
|
|
if len(names) > 60:
|
|
print(f" ... {len(names)-60} more")
|
|
return
|
|
|
|
if mode != "--dump":
|
|
raise SystemExit(__doc__)
|
|
if not names:
|
|
raise SystemExit(f"nothing matches {pattern!r}")
|
|
|
|
# Promote the chosen tensors to graph outputs.
|
|
existing = {o.name for o in model.graph.output}
|
|
for name in names:
|
|
if name not in existing:
|
|
model.graph.output.extend([onnx.ValueInfoProto(name=name)])
|
|
|
|
ids = rust_tokens(text)
|
|
# `pack[len(ps) - 1]`, where len(ps) == len(ids) - 2 (the zero pads).
|
|
style = load_voice()[len(ids) - 3].astype(np.float32).reshape(1, 256)
|
|
|
|
session = ort.InferenceSession(
|
|
model.SerializeToString(), providers=["CPUExecutionProvider"]
|
|
)
|
|
wanted = [o.name for o in session.get_outputs()]
|
|
values = session.run(
|
|
wanted,
|
|
{
|
|
"input_ids": np.array([ids], dtype=np.int64),
|
|
"style": style,
|
|
"speed": np.array([1.0], dtype=np.float32),
|
|
},
|
|
)
|
|
|
|
os.makedirs(OUT_DIR, exist_ok=True)
|
|
print(f"text: {text}\ntokens: {len(ids)}\n")
|
|
for name, value in zip(wanted, values):
|
|
if name not in names:
|
|
continue
|
|
safe = name.strip("/").replace("/", "_")
|
|
np.save(f"{OUT_DIR}/{safe}.npy", value)
|
|
array = np.asarray(value)
|
|
print(f" {name}")
|
|
print(f" shape={array.shape} dtype={array.dtype} "
|
|
f"mean={array.mean():+.5f} std={array.std():.5f} "
|
|
f"min={array.min():+.4f} max={array.max():+.4f}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|