fix(kokoro): poll synthesis cancel between stages

tools/kokoro/include/kokoro-decoder.h

@@ -8,6 +8,8 @@
 
 #pragma once
 
+#include "kokoro.h"
+
 #include <string>
 #include <vector>
 
@@ -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<float> & audio,
-        std::string & err);
+        std::string & err,
+        kokoro_cancel_callback cancel = nullptr,
+        void * cancel_user_data = nullptr);
 
 } // namespace eliza_kokoro

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;

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<float> & 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);

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<int32_t> 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<std::mutex> lk(const_cast<kokoro_model *>(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 {

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)

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 <cstdio>
+#include <cstdlib>
+#include <cstring>
+#include <string>
+
+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<CancelAfter *>(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<float> 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;
+}

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<std::mutex> 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");