From 86e629445d4ffe3c2de85ee6e50309628c6f4e6b Mon Sep 17 00:00:00 2001 From: Shaw Date: Wed, 14 Oct 2026 00:38:08 -0400 Subject: [PATCH] Integrate reviewed native speech and ABI 18 on current main Assisted-by: Codex --- CMakeLists.txt | 21 + tools/CMakeLists.txt | 12 + tools/kokoro/CMakeLists.txt | 21 +- tools/kokoro/include/kokoro-decoder-front.h | 50 +- tools/kokoro/include/kokoro-decoder.h | 42 +- tools/kokoro/include/kokoro-generator.h | 81 +- tools/kokoro/include/kokoro-layers.h | 449 +- tools/kokoro/include/kokoro-model-internal.h | 23 + tools/kokoro/include/kokoro-predictor.h | 82 +- tools/kokoro/include/kokoro-profile.h | 18 +- tools/kokoro/include/kokoro.h | 62 +- tools/kokoro/src/kokoro-decoder.cpp | 205 +- tools/kokoro/src/kokoro-generator.cpp | 288 +- tools/kokoro/src/kokoro-predictor.cpp | 368 +- tools/kokoro/src/kokoro.cpp | 464 +- tools/kokoro/tests/CMakeLists.txt | 16 + tools/kokoro/tests/test_kokoro_backend.cpp | 34 + tools/kokoro/tests/test_kokoro_cancel.cpp | 114 + tools/kokoro/tests/test_kokoro_g2p_espeak.cpp | 23 +- .../tests/test_kokoro_layer_options.cpp | 123 + .../tests/test_kokoro_layers_fast_impl.cpp | 2 + .../tests/test_kokoro_layers_portable.cpp | 13 +- .../tests/test_kokoro_layers_ref_impl.cpp | 2 + .../tests/test_kokoro_layers_test_api.h | 15 + tools/kokoro/tests/test_kokoro_phonemes.cpp | 55 +- .../kokoro/tests/test_kokoro_tensor_names.cpp | 20 +- tools/misaki-g2p/CMakeLists.txt | 87 + .../misaki-g2p/LICENSES/eSpeak-NG-NOTICE.txt | 23 + .../misaki-g2p/LICENSES/misaki-Apache-2.0.txt | 201 + tools/misaki-g2p/LICENSES/misaki-NOTICE.txt | 16 + tools/misaki-g2p/LICENSES/spaCy-MIT.txt | 21 + tools/misaki-g2p/README.md | 121 + tools/misaki-g2p/convert_misaki_lexicon.py | 273 + tools/misaki-g2p/include/misaki-g2p.h | 115 + tools/misaki-g2p/pins.json | 117 + tools/misaki-g2p/scripts/gen_tables.py | 212 + tools/misaki-g2p/src/espeak.cpp | 311 + tools/misaki-g2p/src/g2p.cpp | 1115 +++ tools/misaki-g2p/src/lexicon.cpp | 173 + tools/misaki-g2p/src/misaki-g2p.cpp | 238 + tools/misaki-g2p/src/misaki-internal.h | 192 + tools/misaki-g2p/src/numbers.cpp | 116 + tools/misaki-g2p/src/spacy-en-rules.inc | 1353 ++++ tools/misaki-g2p/src/tagger.cpp | 632 ++ tools/misaki-g2p/src/tokenizer.cpp | 596 ++ tools/misaki-g2p/src/unicode-tables.inc | 6949 +++++++++++++++++ tools/misaki-g2p/src/unicode.cpp | 258 + tools/misaki-g2p/tests/CMakeLists.txt | 75 + tools/misaki-g2p/tests/compare_misaki.py | 343 + tools/misaki-g2p/tests/corpus/agent_style.txt | 70 + .../tests/corpus/eliza_android_sentences.txt | 25 + .../tests/corpus/paradee_samples.txt | 5 + .../tests/corpus/wikitext_heldout_200.txt | 200 + tools/misaki-g2p/tests/data/e2m.tsv | 24 + tools/misaki-g2p/tests/data/mini_gold.json | 86 + tools/misaki-g2p/tests/data/mini_silver.json | 4 + tools/misaki-g2p/tests/data/stress.tsv | 18 + tools/misaki-g2p/tests/data/subtokenize.tsv | 27 + tools/misaki-g2p/tests/data/tokenizer.tsv | 20 + .../misaki-g2p/tests/fixtures/agent_style.txt | 953 +++ .../tests/fixtures/agent_style.txt.sha256 | 1 + tools/misaki-g2p/tests/fixtures/reference.txt | 4951 ++++++++++++ .../tests/fixtures/reference.txt.sha256 | 1 + .../misaki-g2p/tests/gen_unit_expectations.py | 122 + tools/misaki-g2p/tests/test_agreement.cpp | 271 + tools/misaki-g2p/tests/test_misaki_g2p.cpp | 274 + tools/misaki-g2p/tools/misaki-g2p-cli.cpp | 122 + tools/omnivoice/CMakeLists.txt | 107 +- tools/omnivoice/compat/msvc/sys/mman.h | 2 +- tools/omnivoice/compat/msvc/unistd.h | 2 +- tools/omnivoice/include/eliza-inference-ffi.h | 327 +- tools/omnivoice/src/eliza-inference-ffi.cpp | 1294 ++- tools/omnivoice/tests/dac-parity.cpp | 13 +- tools/omnivoice/tests/diarizer-metadata.c | 41 +- tools/omnivoice/tests/ffi-abi.c | 259 + tools/omnivoice/tests/ffi-paradee-asr.cpp | 627 ++ .../paradee-contract/fake-paradee-control.h | 28 + .../tests/paradee-contract/fake-paradee.cpp | 183 + .../tests/paradee-contract/misaki-g2p.h | 41 + .../tests/paradee-contract/paradee.h | 79 + tools/paradee/CMakeLists.txt | 47 + tools/paradee/README.md | 125 + tools/paradee/convert_paradee_pt_to_gguf.py | 444 ++ tools/paradee/include/paradee.h | 215 + tools/paradee/pins.json | 93 + tools/paradee/src/paradee-phase-lock.cpp | 226 + tools/paradee/src/paradee.cpp | 670 ++ tools/paradee/tests/CMakeLists.txt | 46 + .../tests/fixtures/phase_lock_synthetic.f32 | Bin 0 -> 76800 bytes .../tests/fixtures/sentences.misaki.json | 758 ++ tools/paradee/tests/test_paradee_api.cpp | 186 + tools/paradee/tests/test_paradee_parity.cpp | 204 + .../paradee/tests/test_paradee_phase_lock.cpp | 134 + tools/paradee/tools/dump_paradee_reference.py | 561 ++ tools/paradee/tools/paradee-stage-dump.cpp | 91 + tools/paradee/tools/paradee-tts.cpp | 113 + ty.toml | 4 + 97 files changed, 29535 insertions(+), 699 deletions(-) create mode 100644 tools/kokoro/include/kokoro-model-internal.h create mode 100644 tools/kokoro/tests/test_kokoro_backend.cpp create mode 100644 tools/kokoro/tests/test_kokoro_cancel.cpp create mode 100644 tools/kokoro/tests/test_kokoro_layer_options.cpp create mode 100644 tools/kokoro/tests/test_kokoro_layers_test_api.h create mode 100644 tools/misaki-g2p/CMakeLists.txt create mode 100644 tools/misaki-g2p/LICENSES/eSpeak-NG-NOTICE.txt create mode 100644 tools/misaki-g2p/LICENSES/misaki-Apache-2.0.txt create mode 100644 tools/misaki-g2p/LICENSES/misaki-NOTICE.txt create mode 100644 tools/misaki-g2p/LICENSES/spaCy-MIT.txt create mode 100644 tools/misaki-g2p/README.md create mode 100644 tools/misaki-g2p/convert_misaki_lexicon.py create mode 100644 tools/misaki-g2p/include/misaki-g2p.h create mode 100644 tools/misaki-g2p/pins.json create mode 100644 tools/misaki-g2p/scripts/gen_tables.py create mode 100644 tools/misaki-g2p/src/espeak.cpp create mode 100644 tools/misaki-g2p/src/g2p.cpp create mode 100644 tools/misaki-g2p/src/lexicon.cpp create mode 100644 tools/misaki-g2p/src/misaki-g2p.cpp create mode 100644 tools/misaki-g2p/src/misaki-internal.h create mode 100644 tools/misaki-g2p/src/numbers.cpp create mode 100644 tools/misaki-g2p/src/spacy-en-rules.inc create mode 100644 tools/misaki-g2p/src/tagger.cpp create mode 100644 tools/misaki-g2p/src/tokenizer.cpp create mode 100644 tools/misaki-g2p/src/unicode-tables.inc create mode 100644 tools/misaki-g2p/src/unicode.cpp create mode 100644 tools/misaki-g2p/tests/CMakeLists.txt create mode 100644 tools/misaki-g2p/tests/compare_misaki.py create mode 100644 tools/misaki-g2p/tests/corpus/agent_style.txt create mode 100644 tools/misaki-g2p/tests/corpus/eliza_android_sentences.txt create mode 100644 tools/misaki-g2p/tests/corpus/paradee_samples.txt create mode 100644 tools/misaki-g2p/tests/corpus/wikitext_heldout_200.txt create mode 100644 tools/misaki-g2p/tests/data/e2m.tsv create mode 100644 tools/misaki-g2p/tests/data/mini_gold.json create mode 100644 tools/misaki-g2p/tests/data/mini_silver.json create mode 100644 tools/misaki-g2p/tests/data/stress.tsv create mode 100644 tools/misaki-g2p/tests/data/subtokenize.tsv create mode 100644 tools/misaki-g2p/tests/data/tokenizer.tsv create mode 100644 tools/misaki-g2p/tests/fixtures/agent_style.txt create mode 100644 tools/misaki-g2p/tests/fixtures/agent_style.txt.sha256 create mode 100644 tools/misaki-g2p/tests/fixtures/reference.txt create mode 100644 tools/misaki-g2p/tests/fixtures/reference.txt.sha256 create mode 100644 tools/misaki-g2p/tests/gen_unit_expectations.py create mode 100644 tools/misaki-g2p/tests/test_agreement.cpp create mode 100644 tools/misaki-g2p/tests/test_misaki_g2p.cpp create mode 100644 tools/misaki-g2p/tools/misaki-g2p-cli.cpp create mode 100644 tools/omnivoice/tests/ffi-abi.c create mode 100644 tools/omnivoice/tests/ffi-paradee-asr.cpp create mode 100644 tools/omnivoice/tests/paradee-contract/fake-paradee-control.h create mode 100644 tools/omnivoice/tests/paradee-contract/fake-paradee.cpp create mode 100644 tools/omnivoice/tests/paradee-contract/misaki-g2p.h create mode 100644 tools/omnivoice/tests/paradee-contract/paradee.h create mode 100644 tools/paradee/CMakeLists.txt create mode 100644 tools/paradee/README.md create mode 100644 tools/paradee/convert_paradee_pt_to_gguf.py create mode 100644 tools/paradee/include/paradee.h create mode 100644 tools/paradee/pins.json create mode 100644 tools/paradee/src/paradee-phase-lock.cpp create mode 100644 tools/paradee/src/paradee.cpp create mode 100644 tools/paradee/tests/CMakeLists.txt create mode 100644 tools/paradee/tests/fixtures/phase_lock_synthetic.f32 create mode 100644 tools/paradee/tests/fixtures/sentences.misaki.json create mode 100644 tools/paradee/tests/test_paradee_api.cpp create mode 100644 tools/paradee/tests/test_paradee_parity.cpp create mode 100644 tools/paradee/tests/test_paradee_phase_lock.cpp create mode 100644 tools/paradee/tools/dump_paradee_reference.py create mode 100644 tools/paradee/tools/paradee-stage-dump.cpp create mode 100644 tools/paradee/tools/paradee-tts.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index 81f23d7e70..5590e31ebd 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -115,6 +115,12 @@ option(LLAMA_USE_PREBUILT_UI "llama: use prebuilt UI from HF Bucket when availab option(LLAMA_TOOLS_INSTALL "llama: install tools" ${LLAMA_TOOLS_INSTALL_DEFAULT}) option(LLAMA_TESTS_INSTALL "llama: install tests" ON) +option(LLAMA_BUILD_KOKORO "Build native Kokoro speech" ON) +option(LLAMA_BUILD_PARADEE "Build native Paradee speech" ON) +option(LLAMA_BUILD_MISAKI_G2P "Build native misaki G2P" ON) +option(LLAMA_BUILD_OMNIVOICE "Build native OmniVoice and fused inference" ON) +option(OMNIVOICE_SHARED "Build libomnivoice for ABI consumers" OFF) + # 3rd party libs option(LLAMA_OPENSSL "llama: use openssl to support HTTPS" ON) option(LLAMA_LLGUIDANCE "llama-common: include LLGuidance library for structured output in common utils" OFF) @@ -232,6 +238,21 @@ if (LLAMA_BUILD_MTMD AND NOT (LLAMA_BUILD_COMMON AND LLAMA_BUILD_TOOLS)) add_subdirectory(tools/mtmd) endif() +if (NOT (LLAMA_BUILD_COMMON AND LLAMA_BUILD_TOOLS)) + if (LLAMA_BUILD_KOKORO) + add_subdirectory(tools/kokoro) + endif() + if (LLAMA_BUILD_MISAKI_G2P) + add_subdirectory(tools/misaki-g2p) + endif() + if (LLAMA_BUILD_PARADEE AND LLAMA_BUILD_KOKORO) + add_subdirectory(tools/paradee) + endif() + if (LLAMA_BUILD_OMNIVOICE) + add_subdirectory(tools/omnivoice) + endif() +endif() + # # install # diff --git a/tools/CMakeLists.txt b/tools/CMakeLists.txt index 780df32661..9d547d5a03 100644 --- a/tools/CMakeLists.txt +++ b/tools/CMakeLists.txt @@ -30,6 +30,18 @@ else() add_subdirectory(parser) add_subdirectory(tts) add_subdirectory(mtmd) + if (LLAMA_BUILD_KOKORO) + add_subdirectory(kokoro) + endif() + if (LLAMA_BUILD_MISAKI_G2P) + add_subdirectory(misaki-g2p) + endif() + if (LLAMA_BUILD_PARADEE AND LLAMA_BUILD_KOKORO) + add_subdirectory(paradee) + endif() + if (LLAMA_BUILD_OMNIVOICE) + add_subdirectory(omnivoice) + endif() if (GGML_RPC) add_subdirectory(rpc) endif() diff --git a/tools/kokoro/CMakeLists.txt b/tools/kokoro/CMakeLists.txt index d14468ba01..db3b7446b3 100644 --- a/tools/kokoro/CMakeLists.txt +++ b/tools/kokoro/CMakeLists.txt @@ -19,6 +19,7 @@ set(KOKORO_PUBLIC_HEADERS include/kokoro-server-mount.h include/kokoro-tensor-names.h include/kokoro-layers.h + include/kokoro-model-internal.h include/kokoro-predictor.h) set(KOKORO_CORE_SOURCES @@ -32,6 +33,8 @@ set(KOKORO_CORE_SOURCES add_library(kokoro_lib STATIC ${KOKORO_CORE_SOURCES} ${KOKORO_PUBLIC_HEADERS}) +# The fused inference shared library links this static archive. +set_target_properties(kokoro_lib PROPERTIES POSITION_INDEPENDENT_CODE ON) target_include_directories(kokoro_lib PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/include) @@ -52,11 +55,9 @@ target_link_libraries(kokoro_lib PUBLIC ggml) if(TARGET ggml-base) target_link_libraries(kokoro_lib PUBLIC ggml-base) endif() -if(TARGET ggml-cpu) - target_link_libraries(kokoro_lib PUBLIC ggml-cpu) -endif() target_compile_features(kokoro_lib PUBLIC cxx_std_17) +target_compile_definitions(kokoro_lib PRIVATE KOKORO_USE_GGML_BACKEND) # kokoro-layers.h routes the Conv1d/ConvTranspose1d/Linear/LSTM hot loops # through the Accelerate BLAS (im2col + sgemm/sgemv) on Apple platforms — @@ -70,12 +71,18 @@ else() find_package(Threads REQUIRED) target_link_libraries(kokoro_lib PUBLIC Threads::Threads) endif() +if(ANDROID) + target_link_libraries(kokoro_lib PUBLIC log) +endif() # Real G2P via libespeak-ng. When present, kokoro_phonemize() drives # espeak_TextToPhonemes() (en-us IPA) and maps codepoints to Kokoro vocab ids, -# reproducing the reference token sequence. When absent, the build falls back -# to the degraded ASCII grapheme mapping and the TS layer must supply IPA. -# Override the search with -DKOKORO_ESPEAK_ROOT= (e.g. Homebrew). +# reproducing the reference token sequence. When absent, the text path falls +# back to the degraded ASCII grapheme mapping — so an espeak-less build reports +# `eliza_inference_kokoro_g2p_kind() == ELIZA_KOKORO_G2P_ASCII` and the TS voice +# layer supplies espeak-ng (WASM) IPA through +# `eliza_inference_kokoro_synthesize_ipa()` (ABI v14, #11776) to stay +# intelligible. Override the search with -DKOKORO_ESPEAK_ROOT= (Homebrew). option(KOKORO_ENABLE_ESPEAK "Link libespeak-ng for real Kokoro G2P" ON) if(KOKORO_ENABLE_ESPEAK) find_path(ESPEAK_NG_INCLUDE_DIR espeak-ng/speak_lib.h @@ -90,7 +97,7 @@ if(KOKORO_ENABLE_ESPEAK) target_compile_definitions(kokoro_lib PRIVATE KOKORO_USE_ESPEAK) message(STATUS "Kokoro G2P: libespeak-ng found (${ESPEAK_NG_LIBRARY}) — real IPA path enabled") else() - message(STATUS "Kokoro G2P: libespeak-ng not found — falling back to ASCII grapheme mapping (TS layer must supply IPA)") + message(STATUS "Kokoro G2P: libespeak-ng not found — ASCII grapheme fallback; g2p_kind() reports ASCII so the TS layer feeds IPA via kokoro_synthesize_ipa (#11776)") endif() endif() diff --git a/tools/kokoro/include/kokoro-decoder-front.h b/tools/kokoro/include/kokoro-decoder-front.h index 4a89770c5d..11244bba8a 100644 --- a/tools/kokoro/include/kokoro-decoder-front.h +++ b/tools/kokoro/include/kokoro-decoder-front.h @@ -1,5 +1,9 @@ // SPDX-License-Identifier: MIT // kokoro-decoder-front.h — Decoder.forward up to the generator (validated port, #9588). +// +// Channel widths come from the weights struct (filled from the GGUF tensor +// shapes by kokoro-decoder.cpp), so width-scaled students of Kokoro-82M +// (Paradee-8M: AdaIN width 256, style 16) share this code. #pragma once #include #include @@ -24,6 +28,12 @@ struct DecAdainResBlk { const float * conv1x1_b = nullptr; // [Cout] (null — conv1x1 bias=False) const float * pool_w = nullptr; // [Cin, 1, 3] (upsample only) const float * pool_b = nullptr; // [Cin] (upsample only) + // Trained InstanceNorm affine inside AdaIN1d (students only; null keeps + // Kokoro-82M's identity affine). + const float * norm1_w = nullptr; // [Cin] + const float * norm1_b = nullptr; // [Cin] + const float * norm2_w = nullptr; // [Cout] + const float * norm2_b = nullptr; // [Cout] }; // AdainResBlk1d (decode-block flavor: leaky_relu 0.2; pool/shortcut; @@ -35,7 +45,7 @@ inline void dec_adainresblk1d_forward( // residual branch: norm1 -> leaky_relu(0.2) -> [pool] -> conv1 -> norm2 -> leaky_relu -> conv2 std::vector r(x, x + (size_t)Cin * T_in); - adain1d_forward(r.data(), Cin, T_in, s, Sdim, w.norm1_fc_w, w.norm1_fc_b); + adain1d_forward(r.data(), Cin, T_in, s, Sdim, w.norm1_fc_w, w.norm1_fc_b, w.norm1_w, w.norm1_b); for (size_t i = 0; i < r.size(); ++i) if (r[i] < 0) r[i] *= 0.2f; int T_pool = T_in; @@ -47,7 +57,7 @@ inline void dec_adainresblk1d_forward( } std::vector r3((size_t)Cout * T_pool); conv1d_forward(r.data(), Cin, T_pool, w.conv1_w, w.conv1_b, Cout, 3, 1, 1, 1, r3.data(), T_pool); - adain1d_forward(r3.data(), Cout, T_pool, s, Sdim, w.norm2_fc_w, w.norm2_fc_b); + adain1d_forward(r3.data(), Cout, T_pool, s, Sdim, w.norm2_fc_w, w.norm2_fc_b, w.norm2_w, w.norm2_b); for (size_t i = 0; i < r3.size(); ++i) if (r3[i] < 0) r3[i] *= 0.2f; std::vector r4((size_t)Cout * T_pool); conv1d_forward(r3.data(), Cout, T_pool, w.conv2_w, w.conv2_b, Cout, 3, 1, 1, 1, r4.data(), T_pool); @@ -82,15 +92,16 @@ struct DecoderFrontWeights { const float * F0_conv_b = nullptr; // [1] const float * N_conv_w = nullptr; // [1,1,3] const float * N_conv_b = nullptr; // [1] - const float * asr_res_w = nullptr; // [64,512,1] - const float * asr_res_b = nullptr; // [64] - DecAdainResBlk encode; // 514 -> 1024, learned_sc - DecAdainResBlk decode[4]; // 1090->1024 (x3), 1090->512 upsample + const float * asr_res_w = nullptr; // [asr_res_C, Cin_asr, 1] + const float * asr_res_b = nullptr; // [asr_res_C] + int asr_res_C = 64; // asr_res output channels (64 for Kokoro and Paradee) + DecAdainResBlk encode; // Kokoro: 514 -> 1024, learned_sc + DecAdainResBlk decode[4]; // Kokoro: 1090->1024 (x3), 1090->512 upsample }; // Decoder.forward up to (not including) the generator. -// asr[512,T_asr] (T_asr=132), F0_curve[2*T_asr], N[2*T_asr], s[128] -// Output: x_out [512, 2*T_asr] (== generator_in_0); also returns the +// asr[Cin_asr,T_asr], F0_curve[2*T_asr], N[2*T_asr], s[Sdim] +// Output: x_out [decode[3].Cout, 2*T_asr] (== generator_in_0); also returns the // stride-2 conv outputs F0_down[T_asr], N_down[T_asr] (caller passes them, // together with the ORIGINAL F0_curve, into the generator). inline void decoder_front( @@ -100,36 +111,37 @@ inline void decoder_front( std::vector & x_out, std::vector & F0_down, std::vector & N_down) { - const int Tc = 2 * T_asr; // 264 + const int Tc = 2 * T_asr; F0_down.assign(T_asr, 0.0f); conv1d_forward(F0_curve, 1, Tc, W.F0_conv_w, W.F0_conv_b, 1, 3, 2, 1, 1, F0_down.data(), T_asr); N_down.assign(T_asr, 0.0f); conv1d_forward(N_in, 1, Tc, W.N_conv_w, W.N_conv_b, 1, 3, 2, 1, 1, N_down.data(), T_asr); - // x = cat([asr, F0, N], dim=channels) -> [514, T_asr] + // x = cat([asr, F0, N], dim=channels) -> [Cin_asr + 2, T_asr] std::vector xcat((size_t)(Cin_asr + 2) * T_asr); std::memcpy(xcat.data(), asr, sizeof(float) * (size_t)Cin_asr * T_asr); std::memcpy(xcat.data() + (size_t)Cin_asr * T_asr, F0_down.data(), sizeof(float) * T_asr); std::memcpy(xcat.data() + (size_t)(Cin_asr + 1) * T_asr, N_down.data(), sizeof(float) * T_asr); std::vector x; int T_x; - dec_adainresblk1d_forward(W.encode, xcat.data(), T_asr, s, x, T_x); // encode 514->1024 + dec_adainresblk1d_forward(W.encode, xcat.data(), T_asr, s, x, T_x); // encode Cin_asr+2 -> W - std::vector asr_res((size_t)64 * T_asr); // asr_res Conv1d k1 512->64 - conv1d_forward(asr, Cin_asr, T_asr, W.asr_res_w, W.asr_res_b, 64, 1, 1, 0, 1, asr_res.data(), T_asr); + const int Cr = W.asr_res_C; + std::vector asr_res((size_t)Cr * T_asr); // asr_res Conv1d k1 Cin_asr->Cr + conv1d_forward(asr, Cin_asr, T_asr, W.asr_res_w, W.asr_res_b, Cr, 1, 1, 0, 1, asr_res.data(), T_asr); bool res = true; for (int b = 0; b < 4; ++b) { std::vector blk_in; - if (res) { // cat([x, asr_res, F0, N]) -> 1024+64+1+1 = 1090 + if (res) { // cat([x, asr_res, F0, N]) -> Kokoro 1024+64+1+1 = 1090 const int Cx = (int)(x.size() / T_x); - const int Cin_blk = Cx + 64 + 1 + 1; + const int Cin_blk = Cx + Cr + 1 + 1; blk_in.assign((size_t)Cin_blk * T_x, 0.0f); std::memcpy(blk_in.data(), x.data(), sizeof(float) * (size_t)Cx * T_x); - std::memcpy(blk_in.data() + (size_t)Cx * T_x, asr_res.data(), sizeof(float) * 64 * T_x); - std::memcpy(blk_in.data() + (size_t)(Cx + 64) * T_x, F0_down.data(), sizeof(float) * T_x); - std::memcpy(blk_in.data() + (size_t)(Cx + 65) * T_x, N_down.data(), sizeof(float) * T_x); + std::memcpy(blk_in.data() + (size_t)Cx * T_x, asr_res.data(), sizeof(float) * (size_t)Cr * T_x); + std::memcpy(blk_in.data() + (size_t)(Cx + Cr) * T_x, F0_down.data(), sizeof(float) * T_x); + std::memcpy(blk_in.data() + (size_t)(Cx + Cr + 1) * T_x, N_down.data(), sizeof(float) * T_x); } else { blk_in.assign(x.begin(), x.end()); } @@ -138,7 +150,7 @@ inline void decoder_front( x.swap(y); T_x = T_y; if (W.decode[b].upsample) res = false; // decode3 upsamples -> res stops } - x_out.swap(x); // [512, 264] + x_out.swap(x); // [decode[3].Cout, 2*T_asr] } } // namespace eliza_kokoro \ No newline at end of file diff --git a/tools/kokoro/include/kokoro-decoder.h b/tools/kokoro/include/kokoro-decoder.h index 3de5d2562d..312e13ed57 100644 --- a/tools/kokoro/include/kokoro-decoder.h +++ b/tools/kokoro/include/kokoro-decoder.h @@ -5,27 +5,67 @@ // Wires the validated decoder_front (kokoro-decoder-front.h) + Generator // (kokoro-generator.h) against the model's all-F32 ggml context. Replaces the // J2-ship placeholder spectrogram in kokoro_synthesize (#9588). +// +// Every channel width is read from the GGUF tensor shapes, so the same code +// runs Kokoro-82M and width-scaled students of it (Paradee-8M). #pragma once +#include "kokoro.h" + #include #include +struct ggml_context; + namespace eliza_kokoro { struct kokoro_model; +struct DecoderConfig { + std::string ns = "kokoro"; // tensor namespace (".decoder.*", ".gen.*") + bool instnorm_affine = false; // AdaIN1d uses the trained InstanceNorm affine (normN.norm.*, adainN.j.norm.*) + bool canonical_source_stft = false; // see GeneratorConfig::canonical_source_stft + bool precise_source = false; // see GeneratorConfig::precise_source + int up_rates[2] = { 10, 6 }; // generator upsample rates (kernels and widths come from the tensors) + int istft_hop = 5; // generator iSTFT hop (n_fft comes from conv_post) + std::vector * har_source_out = nullptr; // generator m_source output + const std::vector * har_source_in = nullptr; // parity hook: replace m_source's output + bool (*cancel)(void * user) = nullptr; + void * cancel_user = nullptr; + bool cancelled = false; // out: set when cancel() stopped the forward +}; + +// Generic decoder over a loaded F32 context. +// asr_ct : [asr_dim, T_frame] channel-major +// F0, N : [2*T_frame] +// style : [Sdim] decoder style (Sdim read from the GGUF) +// Returns false on error (err set) or cancellation (cfg.cancelled, err "cancelled"). +bool decoder_forward( + ggml_context * ctx, + DecoderConfig & cfg, + const float * asr_ct, int asr_dim, int T_frame, + const float * F0, const float * N, + const float * style, + std::vector & audio, + std::string & err); + // Run the full decoder. Inputs come from kokoro_predictor_forward: // asr_ct : [512, T_frame] channel-major (transpose of PredictorOut.asr [T,512]) // F0, N : [2*T_frame] (PredictorOut.F0_pred / N_pred — the up-2x curves) // ref_s_dec: [128] decoder-half style (ref_s[:128]) // Output: audio (24 kHz mono), resized to (2*T_frame)*300. +// `cancel` (optional) is polled between the decoder front and the generator; +// when it returns true the call stops early, returns false with err +// "cancelled", and leaves `audio` empty. bool kokoro_decoder_forward( const kokoro_model * model, const float * asr_ct, int T_frame, const float * F0, const float * N, const float * ref_s_dec, std::vector & audio, - std::string & err); + std::string & err, + kokoro_cancel_callback cancel = nullptr, + void * cancel_user_data = nullptr); } // namespace eliza_kokoro diff --git a/tools/kokoro/include/kokoro-generator.h b/tools/kokoro/include/kokoro-generator.h index d1d88f6968..86a733b71d 100644 --- a/tools/kokoro/include/kokoro-generator.h +++ b/tools/kokoro/include/kokoro-generator.h @@ -2,15 +2,18 @@ // // kokoro-generator.h — iSTFTNet Generator.forward (StyleTTS-2 decoder back-end). // -// The generator turns the decoder body output `x` [512, 264], the style -// vector `s` [128], and the (un-downsampled) F0 curve `f0_curve` [264] into -// `audio` [79200] (24 kHz). +// The generator turns the decoder body output `x` [Cin, T0], the style +// vector `s` [Sdim], and the (un-downsampled) F0 curve `f0_curve` [T0] into +// `audio` [T0 * prod(upsample_rates) * hop] (24 kHz). +// +// Kokoro-82M: Cin 512, stage channels 256/128, Sdim 128. +// Paradee-8M: Cin 128, stage channels 64/32, Sdim 16. // // Weights are raw float pointers (PyTorch row-major, weight_norm-fused): // Conv1d weight [Cout, Cin, K] // ConvTranspose1d wt [Cin, Cout, K] // Linear weight [out, in] -// AdaIN1d fc.weight [2C, style_dim] (style_dim = 128) +// AdaIN1d fc.weight [2C, style_dim] // Snake alpha [C] // The caller supplies them via GeneratorWeights so the function composes with // any weight-loading boundary (GGUF tensor lookup, raw .f32 fixtures, …). @@ -28,12 +31,18 @@ struct GenSubBlockWeights { const float * conv1_b = nullptr; // [C] const float * conv2_w = nullptr; // [C, C, K] const float * conv2_b = nullptr; // [C] - const float * adain1_fc_w = nullptr; // [2C, 128] + const float * adain1_fc_w = nullptr; // [2C, Sdim] const float * adain1_fc_b = nullptr; // [2C] - const float * adain2_fc_w = nullptr; // [2C, 128] + const float * adain2_fc_w = nullptr; // [2C, Sdim] const float * adain2_fc_b = nullptr; // [2C] const float * alpha1 = nullptr; // [C] const float * alpha2 = nullptr; // [C] + // Trained InstanceNorm affine inside AdaIN1d (students only; null keeps + // Kokoro-82M's identity affine). + const float * adain1_norm_w = nullptr; // [C] + const float * adain1_norm_b = nullptr; // [C] + const float * adain2_norm_w = nullptr; // [C] + const float * adain2_norm_b = nullptr; // [C] }; struct GenAdaResBlockWeights { @@ -41,30 +50,78 @@ struct GenAdaResBlockWeights { }; struct GeneratorWeights { - // m_source.l_linear: Linear(9 -> 1). - const float * l_linear_w = nullptr; // [1, 9] + // m_source.l_linear: Linear(harmonics -> 1). + const float * l_linear_w = nullptr; // [1, harmonics] const float * l_linear_b = nullptr; // [1] // ups[0], ups[1]: ConvTranspose1d. weight [Cin, Cout, K], bias [Cout]. const float * ups_w[2] = { nullptr, nullptr }; const float * ups_b[2] = { nullptr, nullptr }; - // noise_convs[0], noise_convs[1]: Conv1d. weight [Cout, 22, K], bias [Cout]. + // noise_convs[0], noise_convs[1]: Conv1d. weight [Cout, n_fft+2, K], bias [Cout]. const float * noise_convs_w[2] = { nullptr, nullptr }; const float * noise_convs_b[2] = { nullptr, nullptr }; // noise_res[0] (k=7), noise_res[1] (k=11): AdaINResBlock1. GenAdaResBlockWeights noise_res[2]; - // resblocks[0..5]: AdaINResBlock1 (stage0: k=3,7,11 ch=256; stage1: ch=128). + // resblocks[0..5]: AdaINResBlock1 (3 kernels per stage, 2 stages). GenAdaResBlockWeights resblocks[6]; - // conv_post: Conv1d(128 -> 22, k=7, pad=3). weight [22, 128, 7], bias [22]. + // conv_post: Conv1d(ch[1] -> n_fft+2, k=7, pad=3). const float * conv_post_w = nullptr; const float * conv_post_b = nullptr; }; -// Generator.forward. audio is resized to T0 * 300 (== 79200 for T0=264). +// Generator shape and behaviour. The defaults are Kokoro-82M v1.0. +struct GeneratorConfig { + int Cin = 512; // input channels (upsample_initial_channel) + int ch[2] = { 256, 128 }; // channels after ups[0], ups[1] + int Sdim = 128; // style width + int up_rates[2] = { 10, 6 }; + int up_kernels[2] = { 20, 12 }; + int rb_kernels[3] = { 3, 7, 11 }; + int dil[3] = { 1, 3, 5 }; + int n_fft = 20; + int hop = 5; + int harmonics = 9; // harmonic_num + 1 + float sine_amp = 0.1f; + float voiced_threshold = 10.0f; + float sample_rate = 24000.0f; + + // Source STFT convention. false (Kokoro-82M as shipped): a direct double + // DFT, whose rounding noise decides the sign of near-zero imaginary parts + // and therefore the phase of near-empty bins. true (Paradee): imaginary + // parts and magnitudes below the DFT's rounding floor are set to exactly + // zero, so the DC/Nyquist bins of a real signal and the empty bins of a + // constant (unvoiced, noise-free) source get torch.stft's phase (0, or pi + // for a negative real part) instead of a sign picked by rounding noise. + bool canonical_source_stft = false; + // Harmonic source precision. false (Kokoro-82M as shipped): rad/phase + // pass through float between the SineGen steps. true (Paradee): every + // SineGen step runs in double and only sine_waves is rounded to float. + bool precise_source = false; + bool instnorm_affine = false; // documents that the AdaIN1d norm_* pointers are set + + // Optional outputs / hooks. + std::vector * har_source_out = nullptr; // m_source output [T0 * upsample_scale] + const std::vector * har_source_in = nullptr; // parity hook: use this source instead of m_source + bool (*cancel)(void * user) = nullptr; // checked between generator stages + void * cancel_user = nullptr; +}; + +// Generator.forward with an explicit configuration. Returns false only when +// `cfg.cancel` asked it to stop (audio is then cleared). +bool generator_forward( + const float * x_in, // [cfg.Cin, T0] channel-major + int T0, // input time (== 2 * predictor T_frame) + const float * s, // [cfg.Sdim] + const float * f0_curve, // [T0] + const GeneratorWeights & w, + const GeneratorConfig & cfg, + std::vector & audio); + +// Kokoro-82M Generator.forward. audio is resized to T0 * 300 (== 79200 for T0=264). void kokoro_generator_forward( const float * x_in, // [512, T0] channel-major int T0, // input time (== 2 * predictor T_frame) diff --git a/tools/kokoro/include/kokoro-layers.h b/tools/kokoro/include/kokoro-layers.h index f19f9c032d..398b560bfb 100644 --- a/tools/kokoro/include/kokoro-layers.h +++ b/tools/kokoro/include/kokoro-layers.h @@ -1,10 +1,10 @@ // SPDX-License-Identifier: MIT // -// kokoro-layers.h — small neural-net layer primitives used by the Kokoro -// predictor + decoder forwards. No GGML graph / backend dispatch: the forward -// pass is line-for-line matched to the upstream `kokoro` Python package -// (modules.py + istftnet.py) so the trained weights produce numerically -// equivalent output. +// Neural-network primitives for the Kokoro predictor and decoder. Production +// convolution layers dispatch through GGML while the portable and scalar +// implementations remain numerical references for parity testing. The forward +// equations match the upstream `kokoro` Python package (modules.py + +// istftnet.py), so backend selection cannot change the trained model contract. // // The O(C^2*K*T) hot loops (Conv1d, ConvTranspose1d, Linear, LSTM gate // matvecs) have three implementations, selected at compile time: @@ -36,9 +36,17 @@ #include #include #include +#include #include +#include #include +#if defined(KOKORO_USE_GGML_BACKEND) +#include "ggml-alloc.h" +#include "ggml-backend.h" +#include "ggml.h" +#endif + #if defined(__APPLE__) && !defined(KOKORO_NO_ACCELERATE) && !defined(KOKORO_FORCE_SCALAR) #define KOKORO_USE_ACCELERATE 1 #if !defined(ACCELERATE_NEW_LAPACK) @@ -49,7 +57,6 @@ #define KOKORO_USE_PORTABLE_FAST 1 #include #include -#include #include #include #include @@ -61,6 +68,170 @@ namespace eliza_kokoro { +#if defined(KOKORO_USE_GGML_BACKEND) +namespace detail { + +struct CachedF16Weight { + size_t elements = 0; + std::vector storage; + + ggml_fp16_t * data() { + const uintptr_t address = reinterpret_cast(storage.data()); + const uintptr_t aligned = + (address + GGML_MEM_ALIGN - 1) & ~(uintptr_t) (GGML_MEM_ALIGN - 1); + return reinterpret_cast(aligned); + } +}; + +struct KokoroComputeState { + ggml_backend_t backend = nullptr; + std::unordered_map f16_weights; +}; + +inline thread_local KokoroComputeState * kokoro_compute_state = nullptr; + +class KokoroComputeBackendScope { +public: + explicit KokoroComputeBackendScope(KokoroComputeState * state) + : previous_(kokoro_compute_state) { + kokoro_compute_state = state; + } + + ~KokoroComputeBackendScope() { + kokoro_compute_state = previous_; + } + + KokoroComputeBackendScope(const KokoroComputeBackendScope &) = delete; + KokoroComputeBackendScope & operator=(const KokoroComputeBackendScope &) = delete; + +private: + KokoroComputeState * previous_; +}; + +inline bool ggml_compute_graph_and_copy( + ggml_backend_t backend, ggml_context * ctx, ggml_tensor * output, + float * destination, size_t elements) { + ggml_cgraph * graph = ggml_new_graph_custom(ctx, 256, false); + ggml_build_forward_expand(graph, output); + ggml_gallocr_t allocator = ggml_gallocr_new( + ggml_backend_get_default_buffer_type(backend)); + if (!allocator) return false; + const bool allocated = ggml_gallocr_alloc_graph(allocator, graph); + const enum ggml_status status = allocated + ? ggml_backend_graph_compute(backend, graph) + : GGML_STATUS_ALLOC_FAILED; + if (status == GGML_STATUS_SUCCESS) { + ggml_backend_tensor_get( + output, destination, 0, sizeof(float) * elements); + } + ggml_gallocr_free(allocator); + return status == GGML_STATUS_SUCCESS; +} + +inline const ggml_fp16_t * cached_f16_weight( + KokoroComputeState & state, const float * weight, size_t elements) { + auto & cached = state.f16_weights[weight]; + if (cached.elements != elements) { + cached.elements = elements; + cached.storage.resize( + elements * sizeof(ggml_fp16_t) + GGML_MEM_ALIGN - 1); + ggml_fp32_to_fp16_row(weight, cached.data(), elements); + } + return cached.data(); +} + +inline bool ggml_conv1d_f32( + const float * x, int Cin, int T, + const float * W, const float * b, int Cout, int K, + int stride, int pad, int dilation, + float * y, int T_out) { + KokoroComputeState * state = kokoro_compute_state; + if (!state || !state->backend) return false; + ggml_backend_t backend = state->backend; + + ggml_init_params params = { + /* mem_size = */ 2 * 1024 * 1024, + /* mem_buffer = */ nullptr, + /* no_alloc = */ true, + }; + ggml_context * ctx = ggml_init(params); + if (!ctx) return false; + + ggml_tensor * weight = + ggml_new_tensor_3d(ctx, GGML_TYPE_F16, K, Cin, Cout); + ggml_tensor * input = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, T, Cin, 1); + weight->data = const_cast(cached_f16_weight( + *state, W, (size_t) K * Cin * Cout)); + input->data = const_cast(x); + + ggml_tensor * output = ggml_conv_1d( + ctx, weight, input, stride, pad, dilation); + if (b) { + ggml_tensor * bias = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, Cout); + bias->data = const_cast(b); + output = ggml_add(ctx, output, bias); + } + + const bool ok = + output->ne[0] == T_out && + output->ne[1] == Cout && + ggml_compute_graph_and_copy( + backend, ctx, output, y, (size_t) Cout * T_out); + ggml_free(ctx); + return ok; +} + +inline bool ggml_convtranspose1d_f32( + const float * x, int Cin, int T, + const float * W, const float * b, int Cout, int K, + int stride, int pad, int output_pad, + float * y, int T_out) { + KokoroComputeState * state = kokoro_compute_state; + if (!state || !state->backend) return false; + ggml_backend_t backend = state->backend; + + ggml_init_params params = { + /* mem_size = */ 2 * 1024 * 1024, + /* mem_buffer = */ nullptr, + /* no_alloc = */ true, + }; + ggml_context * ctx = ggml_init(params); + if (!ctx) return false; + + ggml_tensor * weight = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, K, Cout, Cin); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, T, Cin); + weight->data = const_cast(W); + input->data = const_cast(x); + + ggml_tensor * output = ggml_conv_transpose_1d(ctx, weight, input, stride, 0, 1); + if (pad > 0) { + const int64_t cropped = output->ne[0] - 2 * (int64_t) pad; + output = ggml_view_2d( + ctx, output, cropped, output->ne[1], output->nb[1], + (size_t) pad * output->nb[0]); + output = ggml_cont(ctx, output); + } + if (output_pad > 0) { + output = ggml_pad(ctx, output, output_pad, 0, 0, 0); + } + if (b) { + ggml_tensor * bias = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, Cout); + bias->data = const_cast(b); + output = ggml_add(ctx, output, bias); + } + + const bool ok = + output->ne[0] == T_out && + output->ne[1] == Cout && + ggml_compute_graph_and_copy( + backend, ctx, output, y, (size_t) Cout * T_out); + ggml_free(ctx); + return ok; +} + +} // namespace detail +#endif + #if defined(KOKORO_USE_PORTABLE_FAST) namespace detail { @@ -213,6 +384,154 @@ inline void kokoro_axpy_f32(float a, const float * x, float * y, int n) { } // namespace detail #endif // KOKORO_USE_PORTABLE_FAST +// ================================================================= +// Opt-in layer strategies. Kokoro-82M never sets them, so its arithmetic is +// unchanged; Paradee enables them for its synthesis calls. +// tap_gemm_conv: stride-1 Conv1d on Accelerate as one SGEMM per kernel +// tap over a zero-padded input, instead of im2col (whose +// column matrix is Cin*K*T floats, ~70 MB for a 32-channel, +// K=11, 10 s generator stage). On the portable path, +// register-tiled Conv1d / ConvTranspose1d kernels instead of +// per-tap AXPY loops. +// fast_snake: Snake1D with a polynomial sin^2 (|error| < 5e-7) that +// vectorises, instead of one libm sin call per element. +// ================================================================= +namespace detail { +struct KokoroLayerOptions { + bool tap_gemm_conv = false; + bool fast_snake = false; +}; +inline KokoroLayerOptions & kokoro_layer_options() { + static thread_local KokoroLayerOptions options; + return options; +} + +class KokoroLayerOptionsScope { +public: + explicit KokoroLayerOptionsScope(const KokoroLayerOptions & o) : previous_(kokoro_layer_options()) { + kokoro_layer_options() = o; + } + ~KokoroLayerOptionsScope() { kokoro_layer_options() = previous_; } + KokoroLayerOptionsScope(const KokoroLayerOptionsScope &) = delete; + KokoroLayerOptionsScope & operator=(const KokoroLayerOptionsScope &) = delete; +private: + KokoroLayerOptions previous_; +}; + +// sin(y)^2 for float y: reduction by pi in double (sin^2 has period pi; a +// float reduction loses ~1e-7 per 100 rad), then an odd degree-11 +// polynomial for sin on [-pi/2, pi/2] in float. +inline float fast_sin2(float y) { + const double n = std::nearbyint((double) y * 0.318309886183790671538); + const float r = (float) ((double) y - n * 3.14159265358979323846); + const float r2 = r * r; + float p = -2.50521083854417187751e-08f; + p = p * r2 + 2.75573192239858906526e-06f; + p = p * r2 - 1.98412698412698412526e-04f; + p = p * r2 + 8.33333333333333321769e-03f; + p = p * r2 - 1.66666666666666657415e-01f; + const float sn = r + r * r2 * p; + return sn * sn; +} + +#if defined(KOKORO_USE_PORTABLE_FAST) && defined(KOKORO_USE_NEON) +#define KOKORO_HAVE_TILED_KERNELS 1 +// Register-tiled NEON kernels for the portable path (tap_gemm_conv). A tile +// is 4 output rows x 16 time steps held in 16 q-register accumulators across +// the whole reduction, with the 4 weights of each (ci, k) packed into one +// vector, so each 16-float input load feeds 64 FMAs (the AXPY path does one +// load, one FMA and one store per 4 MACs). Tiles run on the layer pool. +constexpr int KT_MR = 4; +constexpr int KT_NR = 16; + +// a: KT_MR-wide packed A columns (a[p*4 + r]), b: B rows (b + p*ldb), Kd steps. +inline void kt_micro(const float * a, const float * b, size_t ldb, int Kd, float32x4_t acc[KT_MR][4]) { + for (int p = 0; p < Kd; ++p) { + const float * bp = b + (size_t) p * ldb; + const float32x4_t x0 = vld1q_f32(bp), x1 = vld1q_f32(bp + 4), x2 = vld1q_f32(bp + 8), x3 = vld1q_f32(bp + 12); + const float32x4_t w = vld1q_f32(a + (size_t) p * KT_MR); +#define KT_ROW(r) \ + acc[r][0] = vfmaq_laneq_f32(acc[r][0], x0, w, r); acc[r][1] = vfmaq_laneq_f32(acc[r][1], x1, w, r); \ + acc[r][2] = vfmaq_laneq_f32(acc[r][2], x2, w, r); acc[r][3] = vfmaq_laneq_f32(acc[r][3], x3, w, r); + KT_ROW(0) KT_ROW(1) KT_ROW(2) KT_ROW(3) +#undef KT_ROW + } +} + +inline void kt_store(float32x4_t acc[KT_MR][4], float * c, size_t ldc, int mr, int nr) { + for (int r = 0; r < mr; ++r) { + float tmp[KT_NR]; + vst1q_f32(tmp, acc[r][0]); vst1q_f32(tmp + 4, acc[r][1]); + vst1q_f32(tmp + 8, acc[r][2]); vst1q_f32(tmp + 12, acc[r][3]); + std::memcpy(c + (size_t) r * ldc, tmp, sizeof(float) * (size_t) nr); + } +} + +// y[Cout, T_out] = b + sum_{ci,k} W[co, ci, k] * xp[ci, to + k*dil] +// xp: [Cin, ldx] zero-padded input, readable up to to + (K-1)*dil + KT_NR. +inline void conv1d_s1_tiled(const float * xp, int ldx, int Cin, + const float * W, const float * b, int Cout, int K, int dil, + float * y, int T_out) { + const int n_mb = (Cout + KT_MR - 1) / KT_MR; + const int n_nb = (T_out + KT_NR - 1) / KT_NR; + // Pack weights per row block: wp[mb][ci][k][r]. + std::vector wp((size_t) n_mb * Cin * K * KT_MR, 0.0f); + for (int co = 0; co < Cout; ++co) + for (int ci = 0; ci < Cin; ++ci) + for (int k = 0; k < K; ++k) + wp[(((size_t) (co / KT_MR) * Cin + ci) * K + k) * KT_MR + co % KT_MR] = + W[((size_t) co * Cin + ci) * K + k]; + const size_t total_macs = (size_t) Cout * T_out * Cin * K; + kokoro_parallel_for(total_macs, n_mb * n_nb, [&](int tile) { + const int mb = tile / n_nb; + const int co0 = mb * KT_MR; + const int t0 = (tile % n_nb) * KT_NR; + const int mr = std::min(KT_MR, Cout - co0); + const int nr = std::min(KT_NR, T_out - t0); + float32x4_t acc[KT_MR][4]; + for (int r = 0; r < KT_MR; ++r) { + const float32x4_t bv = vdupq_n_f32((b && r < mr) ? b[co0 + r] : 0.0f); + acc[r][0] = acc[r][1] = acc[r][2] = acc[r][3] = bv; + } + const float * wb = wp.data() + (size_t) mb * Cin * K * KT_MR; + for (int ci = 0; ci < Cin; ++ci) { + const float * xr = xp + (size_t) ci * ldx + t0; + if (dil == 1) { + // taps k are consecutive rows of a Hankel view: ldb = 1 + kt_micro(wb + (size_t) ci * K * KT_MR, xr, 1, K, acc); + } else { + kt_micro(wb + (size_t) ci * K * KT_MR, xr, (size_t) dil, K, acc); + } + } + kt_store(acc, y + (size_t) co0 * T_out + t0, (size_t) T_out, mr, nr); + }); +} + +// C[M, N] = A[M, Kd] (row-major, lda) * B[Kd, N] (row-major, ldb, readable +// to N rounded up to KT_NR). +inline void sgemm_tiled(int M, int N, int Kd, const float * A, int lda, const float * B, int ldb, + float * C, int ldc) { + const int n_mb = (M + KT_MR - 1) / KT_MR; + const int n_nb = (N + KT_NR - 1) / KT_NR; + std::vector ap((size_t) n_mb * Kd * KT_MR, 0.0f); + for (int i = 0; i < M; ++i) + for (int p = 0; p < Kd; ++p) + ap[((size_t) (i / KT_MR) * Kd + p) * KT_MR + i % KT_MR] = A[(size_t) i * lda + p]; + kokoro_parallel_for((size_t) M * N * Kd, n_mb * n_nb, [&](int tile) { + const int mb = tile / n_nb; + const int i0 = mb * KT_MR; + const int j0 = (tile % n_nb) * KT_NR; + const int mr = std::min(KT_MR, M - i0); + const int nr = std::min(KT_NR, N - j0); + float32x4_t acc[KT_MR][4]; + for (int r = 0; r < KT_MR; ++r) acc[r][0] = acc[r][1] = acc[r][2] = acc[r][3] = vdupq_n_f32(0.0f); + kt_micro(ap.data() + (size_t) mb * Kd * KT_MR, B + j0, (size_t) ldb, Kd, acc); + kt_store(acc, C + (size_t) i0 * ldc + j0, (size_t) ldc, mr, nr); + }); +} +#endif +} // namespace detail + // ================================================================= // Tensor (lightweight host tensor — view or owned) // ================================================================= @@ -274,7 +593,40 @@ inline void conv1d_forward( const float * W, const float * b, int Cout, int K, int stride, int pad, int dilation, float * y, int T_out) { +#if defined(KOKORO_USE_GGML_BACKEND) + if (detail::ggml_conv1d_f32( + x, Cin, T, W, b, Cout, K, stride, pad, dilation, y, T_out)) { + return; + } +#endif #if defined(KOKORO_USE_ACCELERATE) + if (detail::kokoro_layer_options().tap_gemm_conv && stride == 1 && K > 1) { + // y = sum_k W[:, :, k] · x_pad[:, k*dil : k*dil + T_out] + const int Tp = T + 2 * pad; + std::vector xp((size_t) Cin * Tp, 0.0f); + for (int ci = 0; ci < Cin; ++ci) { + std::memcpy(xp.data() + (size_t) ci * Tp + pad, x + (size_t) ci * T, sizeof(float) * (size_t) T); + } + std::vector wk((size_t) K * Cout * Cin); + for (int co = 0; co < Cout; ++co) + for (int ci = 0; ci < Cin; ++ci) + for (int k = 0; k < K; ++k) + wk[((size_t) k * Cout + co) * Cin + ci] = W[((size_t) co * Cin + ci) * K + k]; + if (b) { + for (int co = 0; co < Cout; ++co) { + float * yc = y + (size_t) co * T_out; + for (int to = 0; to < T_out; ++to) yc[to] = b[co]; + } + } + for (int k = 0; k < K; ++k) { + cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans, + Cout, T_out, Cin, + 1.0f, wk.data() + (size_t) k * Cout * Cin, Cin, + xp.data() + (size_t) k * dilation, Tp, + (b || k > 0) ? 1.0f : 0.0f, y, T_out); + } + return; + } // im2col + one SGEMM. The PyTorch weight layout [Cout, Cin, K] is already // a row-major [Cout, Cin*K] GEMM operand; the column matrix is // [Cin*K, T_out] with zero padding baked in, so @@ -314,6 +666,40 @@ inline void conv1d_forward( 1.0f, W, Cin * K, cols.data(), T_out, beta, y, T_out); #elif defined(KOKORO_USE_PORTABLE_FAST) +#if defined(KOKORO_HAVE_TILED_KERNELS) + if (detail::kokoro_layer_options().tap_gemm_conv && stride == 1) { + const int ldx = T + 2 * pad + detail::KT_NR; + std::vector xp((size_t) Cin * ldx, 0.0f); + for (int ci = 0; ci < Cin; ++ci) { + std::memcpy(xp.data() + (size_t) ci * ldx + pad, x + (size_t) ci * T, sizeof(float) * (size_t) T); + } + detail::conv1d_s1_tiled(xp.data(), ldx, Cin, W, b, Cout, K, dilation, y, T_out); + return; + } + if (detail::kokoro_layer_options().tap_gemm_conv && stride > 1) { + // Strided: im2col [Cin*K, T_out] (small for these shapes) + tiled GEMM. + const int ldc = T_out + detail::KT_NR; + std::vector cols((size_t) Cin * K * ldc, 0.0f); + for (int ci = 0; ci < Cin; ++ci) { + const float * xi = x + (size_t) ci * T; + for (int k = 0; k < K; ++k) { + float * col = cols.data() + ((size_t) ci * K + k) * ldc; + for (int to = 0; to < T_out; ++to) { + const int ti = to * stride + k * dilation - pad; + if (ti >= 0 && ti < T) col[to] = xi[ti]; + } + } + } + detail::sgemm_tiled(Cout, T_out, Cin * K, W, Cin * K, cols.data(), ldc, y, T_out); + if (b) { + for (int co = 0; co < Cout; ++co) { + float * yc = y + (size_t) co * T_out; + for (int to = 0; to < T_out; ++to) yc[to] += b[co]; + } + } + return; + } +#endif // Parallel over output channels (each thread owns disjoint y rows), and // per (ci, k) tap a branch-free AXPY over the valid `to` range — the // per-element bounds check of the reference loop is hoisted into the @@ -424,6 +810,12 @@ inline void convtranspose1d_forward( const float * W, const float * b, int Cout, int K, int stride, int pad, int output_pad, float * y, int T_out) { +#if defined(KOKORO_USE_GGML_BACKEND) + if (detail::ggml_convtranspose1d_f32( + x, Cin, T, W, b, Cout, K, stride, pad, output_pad, y, T_out)) { + return; + } +#endif // output_pad affects T_out (computed by caller) but does not change // the per-element formula here. (void)output_pad; @@ -457,6 +849,32 @@ inline void convtranspose1d_forward( } } #elif defined(KOKORO_USE_PORTABLE_FAST) +#if defined(KOKORO_HAVE_TILED_KERNELS) + if (detail::kokoro_layer_options().tap_gemm_conv) { + // tmp[Cout*K, T] = W^T x, then col2im (as the Accelerate path). + std::vector wt((size_t) Cout * K * Cin); + for (int ci = 0; ci < Cin; ++ci) + for (int j = 0; j < Cout * K; ++j) + wt[(size_t) j * Cin + ci] = W[(size_t) ci * Cout * K + j]; + const int ldb = T + detail::KT_NR; + std::vector xb((size_t) Cin * ldb, 0.0f); + for (int ci = 0; ci < Cin; ++ci) std::memcpy(xb.data() + (size_t) ci * ldb, x + (size_t) ci * T, sizeof(float) * (size_t) T); + std::vector tmp((size_t) Cout * K * T); + detail::sgemm_tiled(Cout * K, T, Cin, wt.data(), Cin, xb.data(), ldb, tmp.data(), T); + for (int co = 0; co < Cout; ++co) { + float * yc = y + (size_t) co * T_out; + for (int k = 0; k < K; ++k) { + const float * row = tmp.data() + ((size_t) co * K + k) * T; + const int t_off = k - pad; + for (int t_in = 0; t_in < T; ++t_in) { + const int to = t_in * stride + t_off; + if (to >= 0 && to < T_out) yc[to] += row[t_in]; + } + } + } + return; + } +#endif // Parallel over output channels (each thread scatters into disjoint y // rows — no races). Two branch-free AXPY shapes, valid ranges hoisted // out of the inner loop: @@ -705,7 +1123,9 @@ inline void bilstm_forward( inline void adain1d_forward( float * x, int C, int T, const float * s, int Sdim, - const float * fc_W /* [2C, Sdim] */, const float * fc_b /* [2C] */) { + const float * fc_W /* [2C, Sdim] */, const float * fc_b /* [2C] */, + const float * norm_W = nullptr /* [C], InstanceNorm affine weight */, + const float * norm_b = nullptr /* [C], InstanceNorm affine bias */) { std::vector h(2 * C); linear_forward(s, Sdim, fc_W, fc_b, 2 * C, h.data()); // gamma = h[0..C], beta = h[C..2C] @@ -722,7 +1142,11 @@ inline void adain1d_forward( // saved state_dict shows no norm.weight key under norm1/norm2 of // AdainResBlk1d, only fc.weight/fc.bias. The default-affine // weight=1/bias=0 is correct then. - instance_norm1d_forward(x, C, T, 1e-5f, nullptr, nullptr); + // + // Models trained with the affine parameters live (Paradee's students) + // pass them as norm_W / norm_b; Kokoro-82M passes nullptr and keeps the + // identity affine, so its arithmetic is unchanged. + instance_norm1d_forward(x, C, T, 1e-5f, norm_W, norm_b); for (int c = 0; c < C; ++c) { const float gamma = h[c]; const float beta = h[C + c]; @@ -809,6 +1233,15 @@ struct AdainResBlk1dWeights { // Snake1D activation: x = x + (1/a) * sin(a*x)^2 // ================================================================= inline void snake1d_forward(float * x, int C, int T, const float * a /* [C] */) { + if (detail::kokoro_layer_options().fast_snake) { + for (int c = 0; c < C; ++c) { + const float av = std::abs(a[c]) > 1e-8f ? a[c] : 1e-8f; + const float inv_a = 1.0f / av; + float * xc = x + (size_t)c * T; + for (int t = 0; t < T; ++t) xc[t] = xc[t] + inv_a * detail::fast_sin2(av * xc[t]); + } + return; + } for (int c = 0; c < C; ++c) { const float av = std::abs(a[c]) > 1e-8f ? a[c] : 1e-8f; const float inv_a = 1.0f / av; diff --git a/tools/kokoro/include/kokoro-model-internal.h b/tools/kokoro/include/kokoro-model-internal.h new file mode 100644 index 0000000000..9e2dbe8da1 --- /dev/null +++ b/tools/kokoro/include/kokoro-model-internal.h @@ -0,0 +1,23 @@ +// SPDX-License-Identifier: MIT + +/** + * Shares the loaded Kokoro tensor context across the owned predictor and + * decoder translation units. Tensor-name compatibility stays at this boundary + * so every stage consumes the same canonical or published schema. + */ + +#pragma once + +#include + +struct ggml_context; +struct ggml_tensor; + +namespace eliza_kokoro { + +struct kokoro_model; + +ggml_context * kokoro_model_ggml_ctx(const kokoro_model * model); +ggml_tensor * kokoro_find_tensor_compat(ggml_context * ctx, const std::string & name); + +} diff --git a/tools/kokoro/include/kokoro-predictor.h b/tools/kokoro/include/kokoro-predictor.h index efea20d07c..f28fce6046 100644 --- a/tools/kokoro/include/kokoro-predictor.h +++ b/tools/kokoro/include/kokoro-predictor.h @@ -9,14 +9,21 @@ // numerically equivalent activations to the upstream Python reference. // // Stages (all forwarded by `kokoro_predictor_forward`): -// 1. Albert BERT (12 layers, num_hidden_groups=1) over input phoneme ids. -// 2. bert_encoder Linear 768→512 → d_en [T, 512]. +// 1. Albert BERT (n_layer, num_hidden_groups=1) over input phoneme ids. +// 2. bert_encoder Linear d_model→hid → d_en [T, hid]. // 3. predictor.text_encoder (DurationEncoder): 3 × (BiLSTM + AdaLayerNorm), -// conditioned on ref_s[:, 128:] (predictor-half style). -// Output `d` has shape [T, 512+128=640] (concat of LSTM output + style). -// 4. predictor.lstm BiLSTM over `d` → x [T, 512]. -// 5. duration_proj Linear (512→50) → duration logits → sigmoid → sum/speed +// conditioned on the predictor style. +// Output `d` has shape [T, hid+sty] (concat of LSTM output + style). +// 4. predictor.lstm BiLSTM over `d` → x [T, hid]. +// 5. duration_proj Linear (hid→max_dur) → duration logits → sigmoid → sum/speed // → per-phoneme float duration. round.clamp(min=1).long() → pred_dur. +// +// Widths come from the GGUF tensor shapes, so the same code runs Kokoro-82M +// (hid 512, style 128) and width-scaled students of it such as Paradee-8M +// (hid 192, style 32). `predictor_forward` takes a PredictorConfig for the +// student-only parts (tensor namespace, asr projection MLP, InstanceNorm +// affine, torch rounding, cancellation); `kokoro_predictor_forward` is the +// Kokoro-82M entry and keeps its arithmetic unchanged. #pragma once @@ -26,6 +33,8 @@ #include #include +struct ggml_context; + namespace eliza_kokoro { // Forward declaration — kokoro_model holds the actual tensor pointers. @@ -36,22 +45,69 @@ struct kokoro_model; struct PredictorOut { int T_phon = 0; // padded input length (BOS+phonemes+EOS) int T_frame = 0; // total frames after length-regulate (sum of pred_dur) + int hid = 0; // predictor hidden width (Kokoro 512) + int sty = 0; // predictor style width (Kokoro 128) + int asr_dim = 0; // width of `asr` rows (Kokoro 512 = hid) - std::vector d; // [T_phon, 640] = DurationEncoder output (LSTM cat style) - std::vector x_lstm; // [T_phon, 512] = predictor.lstm output (forward+reverse concat) + std::vector bert; // [T_phon, d_model] (only when PredictorConfig::keep_bert) + std::vector d; // [T_phon, hid+sty] = DurationEncoder output (LSTM cat style) + std::vector x_lstm; // [T_phon, hid] = predictor.lstm output (forward+reverse concat) std::vector duration; // [T_phon] = sigmoid+sum (pre-round, post-speed) std::vector pred_dur; // [T_phon] = round.clamp(1).long() // After length-regulate: - std::vector en; // [T_frame, 640] = d transposed and matmul'd with pred_aln_trg - std::vector t_en; // [T_phon, 512] = text_encoder forward (separate from BERT branch) - std::vector asr; // [T_frame, 512] = t_en @ pred_aln_trg + std::vector en; // [T_frame, hid+sty] = d transposed and matmul'd with pred_aln_trg + std::vector t_en; // [T_phon, hid] = text_encoder forward (separate from BERT branch) + std::vector asr; // [T_frame, asr_dim] = (asr_proj(t_en) or t_en) @ pred_aln_trg // F0 / N predictions (run via predictor.F0Ntrain on `en`). - std::vector F0_pred; // [T_frame] - std::vector N_pred; // [T_frame] + std::vector F0_pred; // [2*T_frame] + std::vector N_pred; // [2*T_frame] + + bool cancelled = false; // set when PredictorConfig::cancel stopped the forward +}; + +// Cancellation hook: return true to stop. Checked between stages. +using predictor_cancel_fn = bool (*)(void * user); + +struct PredictorConfig { + // Tensor namespace: ".bert.*", ".predictor.*", ... ("kokoro" for + // Kokoro-82M, "paradee" for Paradee-8M). + std::string ns = "kokoro"; + + // Albert text branch (from the GGUF hparams). + int text_n_layer = 12; + int text_n_head = 12; + int text_d_model = 768; + int text_d_ff = 2048; + + int n_layer = 3; // DurationEncoder / TextEncoder depth + int max_dur = 50; // duration_proj width + int text_kernel = 5; // TextEncoder conv kernel + + // Predictor style vector, [sty] floats (sty read from the GGUF). + const float * style = nullptr; + + // Student-only behaviour; all false reproduces Kokoro-82M exactly. + bool torch_round = false; // round half to even like torch.round (Kokoro: std::round) + bool instnorm_affine = false; // AdaIN1d applies the trained InstanceNorm affine (.normN.norm.*) + bool asr_proj_mlp = false; // asr = Linear(.asr_proj.1)(GELU(Linear(.asr_proj.0)(t_en))) + bool keep_bert = false; // keep the BERT output in PredictorOut::bert (stage dumps) + + predictor_cancel_fn cancel = nullptr; + void * cancel_user = nullptr; }; +// Generic predictor forward over a loaded F32 ggml context. Returns false on +// error (`err_out` set) or cancellation (`out.cancelled`, err_out "cancelled"). +bool predictor_forward( + ggml_context * ctx, + const PredictorConfig & cfg, + const std::vector & phoneme_ids, + float speed, + PredictorOut & out, + std::string & err_out); + // Run the predictor + text_encoder + length-regulation. `phoneme_ids` is the // raw phoneme sequence (BOS+phon+EOS); `ref_s` is the 256-dim style vector // (split internally into predictor + decoder halves). diff --git a/tools/kokoro/include/kokoro-profile.h b/tools/kokoro/include/kokoro-profile.h index e7174f5946..8f34630b2e 100644 --- a/tools/kokoro/include/kokoro-profile.h +++ b/tools/kokoro/include/kokoro-profile.h @@ -12,6 +12,10 @@ #include #include +#if defined(__ANDROID__) +#include +#endif + namespace eliza_kokoro { inline bool kokoro_profile_enabled() { @@ -22,8 +26,16 @@ inline bool kokoro_profile_enabled() { return on; } -// RAII phase timer: logs "