feat: enable SIGUSR1 to soft-cancel pending batches

This is mostly an example of using sd_cancel to asynchronously
cancel the current batch without discarding already generated
images.

Author: wbruna

examples/cli/main.cpp

@@ -474,6 +474,13 @@ bool save_results(const SDCliParams& cli_params,
     return sucessful_reults != 0;
 }
 
+#if defined(__unix__) || defined(__APPLE__) || defined(_POSIX_VERSION)
+#define SD_ENABLE_SIGNAL_HANDLER
+static void set_signal_cancel_handler(sd_ctx_t* sd_ctx);
+#else
+#define set_signal_cancel_handler(SD_CTX) ((void)SD_CTX)
+#endif
+
 int main(int argc, const char* argv[]) {
     if (argc > 1 && std::string(argv[1]) == "--version") {
         std::cout << version_string() << "\n";
@@ -711,6 +718,8 @@ int main(int argc, const char* argv[]) {
             return 1;
         }
 
+        set_signal_cancel_handler(sd_ctx);
+
         if (gen_params.sample_params.sample_method == SAMPLE_METHOD_COUNT) {
             gen_params.sample_params.sample_method = sd_get_default_sample_method(sd_ctx);
         }
@@ -783,6 +792,8 @@ int main(int argc, const char* argv[]) {
             results = generate_video(sd_ctx, &vid_gen_params, &num_results);
         }
 
+        set_signal_cancel_handler(nullptr);
+
         if (results == nullptr) {
             LOG_ERROR("generate failed");
             free_sd_ctx(sd_ctx);
@@ -836,3 +847,59 @@ int main(int argc, const char* argv[]) {
 
     return 0;
 }
+
+#ifdef SD_ENABLE_SIGNAL_HANDLER
+
+#include <atomic>
+#include <csignal>
+#include <thread>
+#include <unistd.h>
+
+// this lock is needed to avoid a race condition between
+// free_sd_ctx and a pending sd_cancel_generation call
+std::atomic_flag signal_lock = ATOMIC_FLAG_INIT;
+static int g_sigint_cnt;
+static sd_ctx_t* g_sd_ctx;
+
+static void sig_cancel_handler(int /* signum */)
+{
+    if (!signal_lock.test_and_set(std::memory_order_acquire)) {
+        if (g_sd_ctx != nullptr) {
+            if (g_sigint_cnt == 1) {
+                char msg[] = "\ngot cancel signal, cancelling new generations\n";
+                write(2, msg, sizeof(msg)-1);
+                /* first signal cancels only the remaining latents on a batch */
+                sd_cancel_generation(g_sd_ctx, SD_CANCEL_NEW_LATENTS);
+                ++g_sigint_cnt;
+            } else {
+                char msg[] = "\ngot cancel signal, cancelling everything\n";
+                write(2, msg, sizeof(msg)-1);
+                /* cancels everything */
+                sd_cancel_generation(g_sd_ctx, SD_CANCEL_ALL);
+            }
+        }
+        signal_lock.clear(std::memory_order_release);
+    }
+}
+
+static void set_signal_cancel_handler(sd_ctx_t* sd_ctx)
+{
+    if (g_sigint_cnt == 0) {
+        g_sigint_cnt++;
+        struct sigaction sa{};
+        sa.sa_handler = sig_cancel_handler;
+        sa.sa_flags   = SA_RESTART;
+        sigaction(SIGUSR1, &sa, nullptr);
+    }
+
+    while (signal_lock.test_and_set(std::memory_order_acquire)) {
+        std::this_thread::yield();
+    }
+
+    g_sd_ctx = sd_ctx;
+
+    signal_lock.clear(std::memory_order_release);
+}
+
+#endif
+