Skip to content

Commit 34ba8cd

Browse files
committed
fix(kokoro): accept published GGUF tensor names + emit F16 weights (#9588)
The fused-lib loader required unprefixed dev tensor names (bert.embd.tok.weight, pred.F0_proj.weight, dec.gen.conv_post.weight) while the published elizaOS Kokoro bundles use the kokoro.* namespace and mainline llama.cpp names (kokoro.bert.layer.attn_q.weight). The loader, the shipped GGUF, and the in-tree converter all disagreed, so kokoro_init_from_file failed its weight sanity check for every available GGUF and the engine silently fell through to OmniVoice/stub. Two root causes, two fixes: 1. Tensor-name mismatch. Centralize the accepted name variants in kokoro-tensor-names.h (published + mainline + legacy dev) and look them up via require_tensor_any(). Required tensors are now a HARD load error instead of the old "non-fatal during J2 — treat absent tensors as zero" path, which produced shape-correct but acoustically degraded (noise) output and masked this bug. 2. All-F32 GGUFs load but synthesize noise. The converter now emits weight matrices / conv kernels (ndim >= 2) as F16 and keeps biases/norms F32, matching the dtype layout the fused forward pass expects. The stub emitter also writes kokoro.gen.conv_post.{weight,bias} so it passes the new required-tensor check. Adds test_kokoro_tensor_names.cpp (LLAMA_BUILD_TESTS-gated) asserting the alias picker resolves published, mainline, and legacy schemas and returns null when a tensor is genuinely absent. Closes the loader half of #9588; regenerating + republishing the bundle GGUF with this converter is the remaining ops step.
1 parent 02020a6 commit 34ba8cd

6 files changed

Lines changed: 232 additions & 8 deletions

File tree

tools/kokoro/CMakeLists.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ set(KOKORO_PUBLIC_HEADERS
1717
include/kokoro-istft.h
1818
include/kokoro-phonemes.h
1919
include/kokoro-server-mount.h
20+
include/kokoro-tensor-names.h
2021
include/kokoro-layers.h
2122
include/kokoro-predictor.h)
2223

tools/kokoro/convert_kokoro_pth_to_gguf.py

Lines changed: 13 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -95,11 +95,20 @@
9595

9696

9797
def _add_tensor(writer: gguf.GGUFWriter, name: str, data: np.ndarray) -> None:
98-
"""Add a tensor as fp32 (cast from fp64/fp16 if needed; ensure c-contig)."""
99-
if data.dtype != np.float32:
98+
"""Add tensors with the dtype layout the Kokoro forward pass expects.
99+
100+
Weight matrices and convolution kernels (ndim >= 2) are emitted as F16;
101+
biases, norms, and other vectors stay F32. All-F32 GGUFs can load but
102+
synthesize noise in the fused runtime path.
103+
"""
104+
if data.dtype not in (np.float32, np.float16):
100105
data = data.astype(np.float32)
101106
if not data.flags["C_CONTIGUOUS"]:
102107
data = np.ascontiguousarray(data)
108+
if data.ndim >= 2:
109+
data = data.astype(np.float16)
110+
elif data.dtype != np.float32:
111+
data = data.astype(np.float32)
103112
writer.add_tensor(name, data)
104113

105114

@@ -281,6 +290,8 @@ def emit_stub(out_path: str, hp: dict) -> None:
281290
_add_tensor(writer, "kokoro.predictor.F0_proj.bias", np.zeros((1,), dtype=np.float32))
282291
_add_tensor(writer, "kokoro.predictor.N_proj.weight", rng.standard_normal((1, hid//2, 1), dtype=np.float32) * scale)
283292
_add_tensor(writer, "kokoro.predictor.N_proj.bias", np.zeros((1,), dtype=np.float32))
293+
_add_tensor(writer, "kokoro.gen.conv_post.weight", rng.standard_normal((22, 128, 7), dtype=np.float32) * scale)
294+
_add_tensor(writer, "kokoro.gen.conv_post.bias", np.zeros((22,), dtype=np.float32))
284295

285296
writer.write_header_to_file()
286297
writer.write_kv_data_to_file()
Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,72 @@
1+
// SPDX-License-Identifier: MIT
2+
//
3+
// kokoro-tensor-names.h — tensor-name compatibility for Kokoro GGUFs.
4+
//
5+
// The published elizaOS Kokoro bundles use the `kokoro.*` namespace emitted by
6+
// tools/kokoro/convert_kokoro_pth_to_gguf.py. Older local/dev GGUFs used short
7+
// unprefixed names. Keep the aliases centralized so desktop and iOS builds load
8+
// the same schema through the same kokoro_lib code path.
9+
10+
#pragma once
11+
12+
namespace eliza_kokoro {
13+
14+
inline constexpr const char * KOKORO_TENSOR_BERT_TOKEN_EMBD[] = {
15+
"kokoro.bert.token_embd.weight",
16+
"kokoro.bert.embd.tok.weight",
17+
"bert.embd.tok.weight",
18+
"kokoro.token_embd.weight",
19+
nullptr,
20+
};
21+
22+
inline constexpr const char * KOKORO_TENSOR_BERT_ATTN_Q[] = {
23+
"kokoro.bert.layer.attn_q.weight",
24+
"kokoro.bert.attn_q.weight",
25+
"bert.layer.attn_q.weight",
26+
"bert.attn_q.weight",
27+
nullptr,
28+
};
29+
30+
inline constexpr const char * KOKORO_TENSOR_DURATION_PROJ[] = {
31+
"kokoro.predictor.duration_proj.weight",
32+
"kokoro.predictor.duration.weight",
33+
"predictor.duration_proj.weight",
34+
"pred.duration_proj.weight",
35+
"pred.duration.weight",
36+
nullptr,
37+
};
38+
39+
inline constexpr const char * KOKORO_TENSOR_F0_PROJ[] = {
40+
"kokoro.predictor.F0_proj.weight",
41+
"predictor.F0_proj.weight",
42+
"pred.F0_proj.weight",
43+
nullptr,
44+
};
45+
46+
inline constexpr const char * KOKORO_TENSOR_N_PROJ[] = {
47+
"kokoro.predictor.N_proj.weight",
48+
"predictor.N_proj.weight",
49+
"pred.N_proj.weight",
50+
nullptr,
51+
};
52+
53+
inline constexpr const char * KOKORO_TENSOR_GEN_CONV_POST[] = {
54+
"kokoro.gen.conv_post.weight",
55+
"kokoro.decoder.gen.conv_post.weight",
56+
"decoder.gen.conv_post.weight",
57+
"dec.gen.conv_post.weight",
58+
nullptr,
59+
};
60+
61+
inline const char * kokoro_pick_tensor_name(
62+
const char * const * aliases,
63+
bool (*has_tensor)(const char * name, void * user_data),
64+
void * user_data) {
65+
if (!aliases || !has_tensor) return nullptr;
66+
for (const char * const * p = aliases; *p; ++p) {
67+
if (has_tensor(*p, user_data)) return *p;
68+
}
69+
return nullptr;
70+
}
71+
72+
} // namespace eliza_kokoro

tools/kokoro/src/kokoro.cpp

Lines changed: 64 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@
3030
#include "kokoro.h"
3131
#include "kokoro-istft.h"
3232
#include "kokoro-phonemes.h"
33+
#include "kokoro-tensor-names.h"
3334

3435
#include "ggml.h"
3536
#include "ggml-alloc.h"
@@ -164,6 +165,39 @@ static ggml_tensor * find_tensor(ggml_context * ctx, const std::string & name) {
164165
return ggml_get_tensor(ctx, name.c_str());
165166
}
166167

168+
static bool has_tensor_alias(const char * name, void * user_data) {
169+
return name && ggml_get_tensor((ggml_context *) user_data, name) != nullptr;
170+
}
171+
172+
static ggml_tensor * find_tensor_any(ggml_context * ctx, const char * const * aliases) {
173+
const char * name = kokoro_pick_tensor_name(aliases, has_tensor_alias, ctx);
174+
return name ? ggml_get_tensor(ctx, name) : nullptr;
175+
}
176+
177+
static std::string format_aliases(const char * const * aliases) {
178+
std::string out;
179+
for (const char * const * p = aliases; p && *p; ++p) {
180+
if (!out.empty()) out += ", ";
181+
out += "'";
182+
out += *p;
183+
out += "'";
184+
}
185+
return out;
186+
}
187+
188+
static ggml_tensor * require_tensor_any(
189+
ggml_context * ctx,
190+
const char * const * aliases,
191+
const char * label,
192+
std::string & err_out) {
193+
ggml_tensor * t = find_tensor_any(ctx, aliases);
194+
if (!t) {
195+
err_out = std::string("required tensor missing for ") + label
196+
+ " (accepted names: " + format_aliases(aliases) + ")";
197+
}
198+
return t;
199+
}
200+
167201
} // namespace
168202

169203
// ---------------------------------------------------------------------------
@@ -259,14 +293,38 @@ kokoro_model_ptr kokoro_load_model(
259293
}
260294
}
261295

262-
// Bind canonical tensors. Missing tensors are non-fatal during the J2
263-
// ship phase — the synthesis path treats absent tensors as zero, which
264-
// produces shape-correct but acoustically degraded output. See the
265-
// J2-kokoro-port-notes.md gap log.
266-
model->tok_embd = find_tensor(model->ctx, "kokoro.token_embd.weight");
296+
// Bind the published Kokoro GGUF schema, while accepting the older
297+
// unprefixed dev names from pre-publication GGUFs. Missing required
298+
// tensors are a hard load error: otherwise the synth path can appear to
299+
// work while silently skipping the real model weights.
300+
model->tok_embd = require_tensor_any(
301+
model->ctx,
302+
KOKORO_TENSOR_BERT_TOKEN_EMBD,
303+
"BERT token embedding",
304+
err_out);
305+
if (!model->tok_embd) return {nullptr, kokoro_model_deleter{}};
306+
307+
if (!require_tensor_any(model->ctx, KOKORO_TENSOR_BERT_ATTN_Q, "BERT attention Q", err_out)) {
308+
return {nullptr, kokoro_model_deleter{}};
309+
}
310+
if (!require_tensor_any(model->ctx, KOKORO_TENSOR_F0_PROJ, "F0 projection", err_out)) {
311+
return {nullptr, kokoro_model_deleter{}};
312+
}
313+
if (!require_tensor_any(model->ctx, KOKORO_TENSOR_N_PROJ, "noise projection", err_out)) {
314+
return {nullptr, kokoro_model_deleter{}};
315+
}
316+
if (!require_tensor_any(model->ctx, KOKORO_TENSOR_GEN_CONV_POST, "generator post convolution", err_out)) {
317+
return {nullptr, kokoro_model_deleter{}};
318+
}
319+
267320
model->mel_proj = find_tensor(model->ctx, "kokoro.decoder.mel_proj.weight");
268321
model->phase_proj = find_tensor(model->ctx, "kokoro.decoder.phase_proj.weight");
269-
model->dur_proj = find_tensor(model->ctx, "kokoro.predictor.duration.weight");
322+
model->dur_proj = require_tensor_any(
323+
model->ctx,
324+
KOKORO_TENSOR_DURATION_PROJ,
325+
"duration projection",
326+
err_out);
327+
if (!model->dur_proj) return {nullptr, kokoro_model_deleter{}};
270328
model->style_proj = find_tensor(model->ctx, "kokoro.style.proj.weight");
271329
model->out_norm = find_tensor(model->ctx, "kokoro.text.out_norm.weight");
272330

tools/kokoro/tests/CMakeLists.txt

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,3 +10,7 @@ add_test(NAME test-kokoro-phonemes COMMAND test-kokoro-phonemes)
1010
add_executable(test-kokoro-istft test_kokoro_istft.cpp)
1111
target_link_libraries(test-kokoro-istft PRIVATE kokoro_lib)
1212
add_test(NAME test-kokoro-istft COMMAND test-kokoro-istft)
13+
14+
add_executable(test-kokoro-tensor-names test_kokoro_tensor_names.cpp)
15+
target_link_libraries(test-kokoro-tensor-names PRIVATE kokoro_lib)
16+
add_test(NAME test-kokoro-tensor-names COMMAND test-kokoro-tensor-names)
Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,78 @@
1+
// SPDX-License-Identifier: MIT
2+
//
3+
// test_kokoro_tensor_names.cpp — regression coverage for issue #9588.
4+
//
5+
// macOS and iOS both link kokoro_lib into the fused libelizainference target.
6+
// If these aliases drift from the published GGUF schema, both platforms can
7+
// export Kokoro symbols yet fail or silently skip the real weights at load time.
8+
9+
#include "kokoro-tensor-names.h"
10+
11+
#include <cassert>
12+
#include <cstdio>
13+
#include <cstring>
14+
#include <set>
15+
#include <string>
16+
17+
namespace {
18+
19+
bool has_name(const char * name, void * user_data) {
20+
const auto * names = static_cast<const std::set<std::string> *>(user_data);
21+
return names->find(name) != names->end();
22+
}
23+
24+
void expect_pick(
25+
const char * const * aliases,
26+
const std::set<std::string> & available,
27+
const char * expected) {
28+
const char * actual = eliza_kokoro::kokoro_pick_tensor_name(
29+
aliases,
30+
has_name,
31+
(void *) &available);
32+
assert(actual != nullptr);
33+
assert(std::strcmp(actual, expected) == 0);
34+
}
35+
36+
} // namespace
37+
38+
int main() {
39+
using namespace eliza_kokoro;
40+
41+
const std::set<std::string> published_schema = {
42+
"kokoro.bert.token_embd.weight",
43+
"kokoro.bert.layer.attn_q.weight",
44+
"kokoro.predictor.duration_proj.weight",
45+
"kokoro.predictor.F0_proj.weight",
46+
"kokoro.predictor.N_proj.weight",
47+
"kokoro.gen.conv_post.weight",
48+
};
49+
50+
expect_pick(KOKORO_TENSOR_BERT_TOKEN_EMBD, published_schema, "kokoro.bert.token_embd.weight");
51+
expect_pick(KOKORO_TENSOR_BERT_ATTN_Q, published_schema, "kokoro.bert.layer.attn_q.weight");
52+
expect_pick(KOKORO_TENSOR_DURATION_PROJ, published_schema, "kokoro.predictor.duration_proj.weight");
53+
expect_pick(KOKORO_TENSOR_F0_PROJ, published_schema, "kokoro.predictor.F0_proj.weight");
54+
expect_pick(KOKORO_TENSOR_N_PROJ, published_schema, "kokoro.predictor.N_proj.weight");
55+
expect_pick(KOKORO_TENSOR_GEN_CONV_POST, published_schema, "kokoro.gen.conv_post.weight");
56+
57+
const std::set<std::string> legacy_schema = {
58+
"bert.embd.tok.weight",
59+
"bert.layer.attn_q.weight",
60+
"pred.duration_proj.weight",
61+
"pred.F0_proj.weight",
62+
"pred.N_proj.weight",
63+
"dec.gen.conv_post.weight",
64+
};
65+
66+
expect_pick(KOKORO_TENSOR_BERT_TOKEN_EMBD, legacy_schema, "bert.embd.tok.weight");
67+
expect_pick(KOKORO_TENSOR_BERT_ATTN_Q, legacy_schema, "bert.layer.attn_q.weight");
68+
expect_pick(KOKORO_TENSOR_DURATION_PROJ, legacy_schema, "pred.duration_proj.weight");
69+
expect_pick(KOKORO_TENSOR_F0_PROJ, legacy_schema, "pred.F0_proj.weight");
70+
expect_pick(KOKORO_TENSOR_N_PROJ, legacy_schema, "pred.N_proj.weight");
71+
expect_pick(KOKORO_TENSOR_GEN_CONV_POST, legacy_schema, "dec.gen.conv_post.weight");
72+
73+
const std::set<std::string> empty_schema;
74+
assert(kokoro_pick_tensor_name(KOKORO_TENSOR_BERT_TOKEN_EMBD, has_name, (void *) &empty_schema) == nullptr);
75+
76+
std::printf("test_kokoro_tensor_names: OK\n");
77+
return 0;
78+
}

0 commit comments

Comments
 (0)