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>
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()
|