diff --git a/tools/kokoro/include/kokoro-decoder.h b/tools/kokoro/include/kokoro-decoder.h index 3de5d2562..7920b5199 100644 --- a/tools/kokoro/include/kokoro-decoder.h +++ b/tools/kokoro/include/kokoro-decoder.h @@ -8,6 +8,8 @@ #pragma once +#include "kokoro.h" + #include #include @@ -20,12 +22,17 @@ struct kokoro_model; // 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.h b/tools/kokoro/include/kokoro.h index aacb54e9d..9c22d3a87 100644 --- a/tools/kokoro/include/kokoro.h +++ b/tools/kokoro/include/kokoro.h @@ -91,6 +91,7 @@ enum kokoro_status { KOKORO_E_OOM = 5, KOKORO_E_NOT_IMPLEMENTED = 6, KOKORO_E_RUNTIME = 7, + KOKORO_E_CANCELLED = 8, }; const char * kokoro_status_str(kokoro_status st) noexcept; @@ -131,6 +132,14 @@ enum kokoro_g2p_kind { }; kokoro_g2p_kind kokoro_g2p_kind_of_build() noexcept; +// Optional cancellation hook for the synthesis entries below. Return true to +// stop the call. It is polled on the synthesizing thread at stage boundaries: +// before the predictor, between the predictor and the decoder, between the +// decoder front and the iSTFTNet generator, and after the generator. A stage +// already running is not interrupted. A cancelled call returns +// KOKORO_E_CANCELLED with `out.samples` empty. +using kokoro_cancel_callback = bool (*)(void * user_data); + // Synthesize a single utterance. `text` is the natural-language input, // `voice` is the loaded ref_s preset. Output PCM lands in `out`. The // `speed_mult` parameter scales the predicted durations (1.0 = native rate). @@ -142,7 +151,9 @@ kokoro_status kokoro_synthesize( const std::string & text, float speed_mult, kokoro_audio & out, - std::string & err_out) noexcept; + std::string & err_out, + kokoro_cancel_callback cancel = nullptr, + void * cancel_user_data = nullptr) noexcept; // Synthesize from a precomputed espeak-ng IPA string (UTF-8) instead of raw // text. The IPA is mapped straight to Kokoro vocab ids via `ipa_to_token_ids` @@ -156,7 +167,9 @@ kokoro_status kokoro_synthesize_ipa( const std::string & ipa, float speed_mult, kokoro_audio & out, - std::string & err_out) noexcept; + std::string & err_out, + kokoro_cancel_callback cancel = nullptr, + void * cancel_user_data = nullptr) noexcept; // Returns the model's audio sample rate (24000 for v1.0). int kokoro_sample_rate(const kokoro_model * model) noexcept; diff --git a/tools/kokoro/src/kokoro-decoder.cpp b/tools/kokoro/src/kokoro-decoder.cpp index a9a74402a..07deda551 100644 --- a/tools/kokoro/src/kokoro-decoder.cpp +++ b/tools/kokoro/src/kokoro-decoder.cpp @@ -73,7 +73,9 @@ bool kokoro_decoder_forward( const float * F0, const float * N, const float * ref_s_dec, std::vector & audio, - std::string & err) { + std::string & err, + kokoro_cancel_callback cancel, + void * cancel_user_data) { audio.clear(); if (!model) { err = "null model"; return false; } if (T_frame <= 0) { err = "non-positive T_frame"; return false; } @@ -136,6 +138,11 @@ bool kokoro_decoder_forward( err = "decoder_front output width mismatch"; return false; } + // The generator is the most expensive stage; skip it when cancelled. + if (cancel && cancel(cancel_user_data)) { + err = "cancelled"; + return false; + } { kokoro_phase_timer t("generator (iSTFTNet)"); kokoro_generator_forward(x.data(), T0, ref_s_dec, F0, G, audio); diff --git a/tools/kokoro/src/kokoro.cpp b/tools/kokoro/src/kokoro.cpp index c8213a9fa..4618698bf 100644 --- a/tools/kokoro/src/kokoro.cpp +++ b/tools/kokoro/src/kokoro.cpp @@ -70,6 +70,7 @@ const char * kokoro_status_str(kokoro_status st) noexcept { case KOKORO_E_OOM: return "out of memory"; case KOKORO_E_NOT_IMPLEMENTED: return "feature not implemented"; case KOKORO_E_RUNTIME: return "runtime error"; + case KOKORO_E_CANCELLED: return "cancelled"; } return "unknown"; } @@ -615,8 +616,23 @@ static kokoro_status kokoro_synthesize_from_input_ids( std::vector phonemes, float speed_mult, kokoro_audio & out, - std::string & err_out) noexcept { + std::string & err_out, + kokoro_cancel_callback cancel, + void * cancel_user_data) noexcept { std::lock_guard lk(const_cast(model)->mu); + const auto cancelled = [&]() { + if (!cancel || !cancel(cancel_user_data)) return false; + err_out = "cancelled"; + out.samples.clear(); + return true; + }; + // Polled at stage boundaries: before the predictor, between the predictor + // and the decoder, between the decoder front and the iSTFTNet generator + // (inside kokoro_decoder_forward), and after the generator. There is no + // hook inside a GGML graph: a failed conv graph falls back to the CPU + // loops in kokoro-layers.h, so a backend abort callback would not stop + // the work. + if (cancelled()) return KOKORO_E_CANCELLED; out.sample_rate = model->hparams.sample_rate; @@ -659,6 +675,7 @@ static kokoro_status kokoro_synthesize_from_input_ids( return KOKORO_E_RUNTIME; } } + if (cancelled()) return KOKORO_E_CANCELLED; const int T = pred.T_frame; if (T <= 0 || (int) pred.asr.size() != T * 512) { err_out = "predictor produced empty/invalid asr (T=" + std::to_string(T) + ")"; @@ -678,11 +695,14 @@ static kokoro_status kokoro_synthesize_from_input_ids( kokoro_phase_timer t("decoder forward"); if (!kokoro_decoder_forward(model, asr_ct.data(), T, pred.F0_pred.data(), pred.N_pred.data(), - ref_s, out.samples, err_out)) { + ref_s, out.samples, err_out, + cancel, cancel_user_data)) { + if (cancelled()) return KOKORO_E_CANCELLED; if (err_out.empty()) err_out = "decoder forward failed"; return KOKORO_E_RUNTIME; } } + if (cancelled()) return KOKORO_E_CANCELLED; return KOKORO_OK; } @@ -710,7 +730,9 @@ kokoro_status kokoro_synthesize( const std::string & text, float speed_mult, kokoro_audio & out, - std::string & err_out) noexcept { + std::string & err_out, + kokoro_cancel_callback cancel, + void * cancel_user_data) noexcept { err_out.clear(); out.samples.clear(); out.sample_rate = 24000; @@ -724,7 +746,8 @@ kokoro_status kokoro_synthesize( // 1. Phonemize (espeak-ng IPA when linked, else lossy ASCII grapheme map). return kokoro_synthesize_from_input_ids( - model, voice, kokoro_phonemize(text), speed_mult, out, err_out); + model, voice, kokoro_phonemize(text), speed_mult, out, err_out, + cancel, cancel_user_data); } kokoro_status kokoro_synthesize_ipa( @@ -733,7 +756,9 @@ kokoro_status kokoro_synthesize_ipa( const std::string & ipa, float speed_mult, kokoro_audio & out, - std::string & err_out) noexcept { + std::string & err_out, + kokoro_cancel_callback cancel, + void * cancel_user_data) noexcept { err_out.clear(); out.samples.clear(); out.sample_rate = 24000; @@ -749,7 +774,7 @@ kokoro_status kokoro_synthesize_ipa( // wrapped as the model input_ids [PAD, *ids, PAD] — no internal G2P. return kokoro_synthesize_from_input_ids( model, voice, wrap_input_ids(ipa_to_token_ids(ipa)), speed_mult, out, - err_out); + err_out, cancel, cancel_user_data); } int kokoro_sample_rate(const kokoro_model * model) noexcept { diff --git a/tools/kokoro/tests/CMakeLists.txt b/tools/kokoro/tests/CMakeLists.txt index 43b22504d..a7725c941 100644 --- a/tools/kokoro/tests/CMakeLists.txt +++ b/tools/kokoro/tests/CMakeLists.txt @@ -38,3 +38,9 @@ add_test(NAME test-kokoro-layers-portable COMMAND test-kokoro-layers-portable) add_executable(test-kokoro-backend test_kokoro_backend.cpp) target_link_libraries(test-kokoro-backend PRIVATE kokoro_lib) add_test(NAME test-kokoro-backend COMMAND test-kokoro-backend) + +# Synthesis cancellation. Weight-free checks always run; set KOKORO_TEST_MODEL +# and KOKORO_TEST_VOICE to also cancel a real synthesis at every stage boundary. +add_executable(test-kokoro-cancel test_kokoro_cancel.cpp) +target_link_libraries(test-kokoro-cancel PRIVATE kokoro_lib) +add_test(NAME test-kokoro-cancel COMMAND test-kokoro-cancel) diff --git a/tools/kokoro/tests/test_kokoro_cancel.cpp b/tools/kokoro/tests/test_kokoro_cancel.cpp new file mode 100644 index 000000000..c57a419b5 --- /dev/null +++ b/tools/kokoro/tests/test_kokoro_cancel.cpp @@ -0,0 +1,114 @@ +// SPDX-License-Identifier: MIT +// +// Kokoro synthesis cancellation. The weight-free checks always run. With +// KOKORO_TEST_MODEL (GGUF) and KOKORO_TEST_VOICE (voice .bin) set, a real +// synthesis is cancelled at each stage boundary in turn. + +#include "kokoro.h" +#include "kokoro-decoder.h" + +#include +#include +#include +#include + +namespace { + +int failures = 0; + +void check(bool ok, const char * what) { + if (!ok) { + std::fprintf(stderr, "FAIL: %s\n", what); + ++failures; + } +} + +// Returns true from the `cancel_at`-th poll onwards; counts every poll. +struct CancelAfter { + int cancel_at; + int polls; +}; + +bool cancel_after(void * user_data) { + CancelAfter * state = static_cast(user_data); + ++state->polls; + return state->cancel_at > 0 && state->polls >= state->cancel_at; +} + +void check_without_weights() { + using namespace eliza_kokoro; + check(std::strcmp(kokoro_status_str(KOKORO_E_CANCELLED), "cancelled") == 0, + "KOKORO_E_CANCELLED has a status string"); + + // Argument validation runs before the first poll: a null model is an + // invalid argument, not a cancellation, and the callback is never called. + CancelAfter state{1, 0}; + kokoro_voice_preset voice; + kokoro_audio audio; + std::string err; + check(kokoro_synthesize(nullptr, voice, "hi", 1.0f, audio, err, + cancel_after, &state) == KOKORO_E_INVALID_ARG, + "null model is reported as an invalid argument"); + check(kokoro_synthesize_ipa(nullptr, voice, "hi", 1.0f, audio, err, + cancel_after, &state) == KOKORO_E_INVALID_ARG, + "null model is reported as an invalid argument (ipa)"); + check(state.polls == 0, "cancel callback is not polled before validation"); + + std::vector pcm; + check(!kokoro_decoder_forward(nullptr, nullptr, 1, nullptr, nullptr, nullptr, + pcm, err, cancel_after, &state), + "decoder rejects a null model"); +} + +int check_with_weights(const char * model_path, const char * voice_path) { + using namespace eliza_kokoro; + std::string err; + kokoro_model_ptr model = kokoro_load_model(model_path, err); + if (!model) { + std::fprintf(stderr, "FAIL: load %s: %s\n", model_path, err.c_str()); + return 1; + } + const kokoro_hparams * hp = kokoro_get_hparams(model.get()); + kokoro_voice_preset voice; + if (kokoro_load_voice_preset(voice_path, hp ? hp->style_dim : 256, voice, err) != KOKORO_OK) { + std::fprintf(stderr, "FAIL: voice %s: %s\n", voice_path, err.c_str()); + return 1; + } + + // Uncancelled: count the polls a full synthesis makes. + CancelAfter full{0, 0}; + kokoro_audio audio; + check(kokoro_synthesize(model.get(), voice, "Hello there.", 1.0f, audio, err, + cancel_after, &full) == KOKORO_OK, + "uncancelled synthesis succeeds"); + check(!audio.samples.empty(), "uncancelled synthesis produces audio"); + check(full.polls >= 4, "synthesis polls at every stage boundary"); + + // Cancel at each poll in turn: every one stops the call with no audio. + for (int at = 1; at <= full.polls; ++at) { + CancelAfter state{at, 0}; + kokoro_audio cancelled_audio; + const kokoro_status st = kokoro_synthesize( + model.get(), voice, "Hello there.", 1.0f, cancelled_audio, err, + cancel_after, &state); + check(st == KOKORO_E_CANCELLED, "cancel at a stage boundary returns KOKORO_E_CANCELLED"); + check(cancelled_audio.samples.empty(), "cancelled synthesis returns no audio"); + check(state.polls <= at + 1, "synthesis stops at the poll that cancelled it"); + } + return 0; +} + +} // namespace + +int main() { + check_without_weights(); + const char * model_path = std::getenv("KOKORO_TEST_MODEL"); + const char * voice_path = std::getenv("KOKORO_TEST_VOICE"); + if (model_path && *model_path && voice_path && *voice_path) { + failures += check_with_weights(model_path, voice_path); + } else { + std::printf("test-kokoro-cancel: KOKORO_TEST_MODEL/KOKORO_TEST_VOICE unset; " + "skipping the weight-backed stage checks\n"); + } + return failures == 0 ? 0 : 1; +} diff --git a/tools/omnivoice/src/eliza-inference-ffi.cpp b/tools/omnivoice/src/eliza-inference-ffi.cpp index 8ec811696..b1ebeb0d0 100644 --- a/tools/omnivoice/src/eliza-inference-ffi.cpp +++ b/tools/omnivoice/src/eliza-inference-ffi.cpp @@ -1890,6 +1890,13 @@ int eliza_inference_kokoro_load( } #ifdef ELIZA_ENABLE_KOKORO +/* Passed to kokoro_lib as its cancel callback. Reads only the Kokoro flag, so + * an OmniVoice or Paradee call starting on another thread cannot clear it. */ +static bool eliza_kokoro_cancel_requested(void * user_data) { + EliInferenceContext * ctx = (EliInferenceContext *) user_data; + return ctx && ctx->kokoro_cancel.load(std::memory_order_acquire); +} + static int eliza_kokoro_synthesize_audio( EliInferenceContext * ctx, const char * input, @@ -1908,10 +1915,11 @@ static int eliza_kokoro_synthesize_audio( } std::lock_guard lock(ctx->kokoro_mutex); /* A cancel_tts issued before this call held the lock targets an earlier - * call; start this one clean. kokoro_lib has no in-graph cancel hook, so - * the flag is checked after the (possibly lazy) reload and after the - * forward pass: a cancelled call discards its audio and reports - * ELIZA_ERR_CANCELLED. */ + * call; start this one clean. The flag is checked after the (possibly + * lazy) reload, by kokoro_lib between its stages (predictor, decoder + * front, generator), and after the forward pass. The exchange below + * clears it on the way out. A cancelled call discards its audio and + * reports ELIZA_ERR_CANCELLED. */ ctx->kokoro_cancel.store(false, std::memory_order_release); const int load_rc = eliza_kokoro_ensure_loaded_locked(ctx, operation, out_error); if (load_rc != ELIZA_OK) return load_rc; @@ -1925,10 +1933,13 @@ static int eliza_kokoro_synthesize_audio( std::string err; const eliza_kokoro::kokoro_status st = ipa_input ? eliza_kokoro::kokoro_synthesize_ipa( - ctx->kokoro_model.get(), ctx->kokoro_voice, in, speed, audio, err) + ctx->kokoro_model.get(), ctx->kokoro_voice, in, speed, audio, err, + eliza_kokoro_cancel_requested, ctx) : eliza_kokoro::kokoro_synthesize( - ctx->kokoro_model.get(), ctx->kokoro_voice, in, speed, audio, err); - if (ctx->kokoro_cancel.exchange(false, std::memory_order_acq_rel)) { + ctx->kokoro_model.get(), ctx->kokoro_voice, in, speed, audio, err, + eliza_kokoro_cancel_requested, ctx); + if (ctx->kokoro_cancel.exchange(false, std::memory_order_acq_rel) || + st == eliza_kokoro::KOKORO_E_CANCELLED) { audio.samples.clear(); eliza_set_error(out_error, std::string("[libelizainference] ") + operation + ": cancelled");