fix(diffusion): add real cancel/abort support to native generation

packages/lib-infer-diffusion/addon/src/handlers/SdCtxHandlers.hpp

@@ -104,7 +104,11 @@ struct SdCtxConfig {
   bool forceSDXLVaeConvScale = false; // force SDXL VAE conv scale (compat fix)
 
   // ── Internal ──────────────────────────────────────────────────────────────
-  bool freeParamsImmediately = true;
+  // Upstream defaults to true, which frees model weight buffers after each
+  // generate_image_internal() call. The addon reuses a single sd_ctx across
+  // multiple generations, so freeing params after the first run causes a
+  // use-after-free SIGSEGV on the second run (including cancel-then-rerun).
+  bool freeParamsImmediately = false;
 };
 
 // ─────────────────────────────────────────────────────────────────────────────

packages/lib-infer-diffusion/addon/src/model-interface/SdModel.cpp

@@ -29,6 +29,10 @@ struct ProgressCtx {
 };
 
 thread_local ProgressCtx tl_progressCtx;
+// Thread-local model pointer for abort callback routing — same pattern as
+// tl_progressCtx for progress.  Avoids relying on the process-global
+// sd_abort_cb_data when multiple SdModel instances could coexist.
+thread_local const SdModel* tl_abortModel = nullptr;
 
 void sdProgressCallback(int step, int steps, float /*time*/, void* /*data*/) {
   if (!tl_progressCtx.job || !tl_progressCtx.job->progressCallback)
@@ -46,6 +50,14 @@ void sdProgressCallback(int step, int steps, float /*time*/, void* /*data*/) {
   tl_progressCtx.job->progressCallback(oss.str());
 }
 
+// Abort callback — wired into sd_set_abort_callback() so that
+// generate_image() can be interrupted mid-denoising.
+// Reads from thread-local tl_abortModel (not the global sd_abort_cb_data)
+// to avoid concurrency issues when multiple SdModel instances coexist.
+bool sdAbortCallback(void* /*data*/) {
+  return tl_abortModel && tl_abortModel->isCancelRequested();
+}
+
 } // namespace
 
 // ---------------------------------------------------------------------------
@@ -206,6 +218,21 @@ std::any SdModel::process(const std::any& input) {
   tl_progressCtx.job = &job;
   tl_progressCtx.startTime = std::chrono::steady_clock::now();
   sd_set_progress_callback(sdProgressCallback, nullptr);
+  tl_abortModel = this;
+  sd_set_abort_callback(sdAbortCallback, nullptr);
+
+  // Scope guard: clear process-global callbacks on any exit path (including
+  // early exceptions from parsing/validation before generate_image runs).
+  auto clearCallbacks = [&]() {
+    tl_progressCtx.job = nullptr;
+    tl_abortModel = nullptr;
+    sd_set_progress_callback(nullptr, nullptr);
+    sd_set_abort_callback(nullptr, nullptr);
+  };
+  struct CallbackGuard {
+    std::function<void()> fn;
+    ~CallbackGuard() { fn(); }
+  } guard{clearCallbacks};
 
   // ── Parse JSON params ─────────────────────────────────────────────────────
   picojson::value v;
@@ -294,21 +321,37 @@ std::any SdModel::process(const std::any& input) {
   if (initImg.data)
     free(initImg.data);
 
+  const bool wasCancelled = cancelRequested_.load();
+
+  // Always free native image buffers, whether cancelled or not.
   int outputCount = 0;
   if (results) {
     for (int i = 0; i < gen.batchCount; ++i) {
-      if (results[i].data && !cancelRequested_.load()) {
-        auto png = encodeToPng(results[i]);
-        if (!png.empty() && job.outputCallback) {
-          job.outputCallback(png);
-          ++outputCount;
+      if (results[i].data) {
+        if (!wasCancelled) {
+          auto png = encodeToPng(results[i]);
+          if (!png.empty() && job.outputCallback) {
+            job.outputCallback(png);
+            ++outputCount;
+          }
         }
         free(results[i].data);
       }
     }
     free(results);
   }
 
+  // If cancelled, propagate as an exception so JobRunner emits
+  // queueException (error path), not queueResult + queueJobEnded.
+  //
+  // This intentionally differs from the LLM addon, which returns normally
+  // on cancel (partial text output is still useful).  Diffusion produces no
+  // partial images, so a "successful" completion with output_count=0 would
+  // be misleading — throwing gives the JS caller an explicit cancel signal.
+  if (wasCancelled) {
+    throw std::runtime_error("Job cancelled");
+  }
+
   const auto t1 = std::chrono::steady_clock::now();
   const double genMs =
       std::chrono::duration<double, std::milli>(t1 - t0).count();
@@ -334,7 +377,10 @@ std::any SdModel::process(const std::any& input) {
           ? (static_cast<double>(totalPixels_) / 1e6 / totalTimeSec)
           : 0.0;
 
-  // ── Build stats ────────────────────────────────────────────────────────────
+  // ── Build stats for runtimeStats() ─────────────────────────────────────────
+  // Stats are stored and emitted via queueJobEnded() → runtimeStats().
+  // process() returns std::any{} (empty) so images delivered via
+  // outputCallback are not duplicated as a queueResult event.
   lastStats_.clear();
 
   lastStats_.emplace_back("generation_time", genMs);
@@ -359,10 +405,9 @@ std::any SdModel::process(const std::any& input) {
   lastStats_.emplace_back("seed", gen.seed);
   lastStats_.emplace_back("output_count", static_cast<int64_t>(outputCount));
 
-  tl_progressCtx.job = nullptr;
-  sd_set_progress_callback(nullptr, nullptr);
-
-  return lastStats_;
+  // Return empty — images are already delivered via outputCallback,
+  // and stats are emitted by queueJobEnded() → runtimeStats().
+  return std::any{};
 }
 
 // ---------------------------------------------------------------------------

packages/lib-infer-diffusion/addon/src/model-interface/SdModel.hpp

@@ -99,6 +99,11 @@ class SdModel : public qvac_lib_inference_addon_cpp::model::IModel,
 
   void cancel() const final;
 
+  /** True if cancel() has been called since the last job started. */
+  [[nodiscard]] bool isCancelRequested() const noexcept {
+    return cancelRequested_.load();
+  }
+
   [[nodiscard]] qvac_lib_inference_addon_cpp::RuntimeStats
   runtimeStats() const final;
 

packages/lib-infer-diffusion/test/unit/CMakeLists.txt

@@ -15,6 +15,7 @@ add_executable(
   test_model_loading.cpp
   test_single_step_inference.cpp
   test_full_generation.cpp
+  test_cancel_context.cpp
   ${CMAKE_SOURCE_DIR}/addon/src/utils/BackendSelection.cpp
   ${CMAKE_SOURCE_DIR}/addon/src/utils/LoggingMacros.cpp
   ${CMAKE_SOURCE_DIR}/addon/src/model-interface/SdModel.cpp