feat(server): cancel current generation on client disconnect

Co-authored-by: donington <jandastroy@gmail.com>

Author: wbruna

examples/server/main.cpp

@@ -2,6 +2,7 @@
 #include <chrono>
 #include <filesystem>
 #include <fstream>
+#include <future>
 #include <iomanip>
 #include <iostream>
 #include <mutex>
@@ -384,6 +385,18 @@ int main(int argc, const char** argv) {
         return httplib::Server::HandlerResponse::Unhandled;
     });
 
+    auto wait_for_generation = [](std::future<void>& ft, sd_ctx_t* sd_ctx, const httplib::Request& req) {
+        std::future_status ft_status;
+        do {
+            if (!ft.valid())
+                break;
+            ft_status = ft.wait_for(std::chrono::milliseconds(1000));
+            if (req.is_connection_closed()) {
+                sd_cancel_generation(sd_ctx, SD_CANCEL_ALL);
+            }
+        } while (ft_status != std::future_status::ready);
+    };
+
     // index html
     std::string index_html;
 #ifdef HAVE_INDEX_HTML
@@ -532,11 +545,13 @@ int main(int argc, const char** argv) {
             sd_image_t* results = nullptr;
             int num_results     = 0;
 
-            {
+            std::future<void> ft = std::async(std::launch::async, [&]() {
                 std::lock_guard<std::mutex> lock(sd_ctx_mutex);
                 results     = generate_image(sd_ctx, &img_gen_params);
                 num_results = gen_params.batch_count;
-            }
+            });
+
+            wait_for_generation(ft, sd_ctx, req);
 
             for (int i = 0; i < num_results; i++) {
                 if (results[i].data == nullptr) {
@@ -779,11 +794,13 @@ int main(int argc, const char** argv) {
             sd_image_t* results = nullptr;
             int num_results     = 0;
 
-            {
+            std::future<void> ft = std::async(std::launch::async, [&]() {
                 std::lock_guard<std::mutex> lock(sd_ctx_mutex);
                 results     = generate_image(sd_ctx, &img_gen_params);
                 num_results = gen_params.batch_count;
-            }
+            });
+
+            wait_for_generation(ft, sd_ctx, req);
 
             json out;
             out["created"]       = static_cast<long long>(std::time(nullptr));
@@ -1095,11 +1112,13 @@ int main(int argc, const char** argv) {
             sd_image_t* results = nullptr;
             int num_results     = 0;
 
-            {
+            std::future<void> ft = std::async(std::launch::async, [&]() {
                 std::lock_guard<std::mutex> lock(sd_ctx_mutex);
                 results     = generate_image(sd_ctx, &img_gen_params);
                 num_results = gen_params.batch_count;
-            }
+            });
+
+            wait_for_generation(ft, sd_ctx, req);
 
             json out;
             out["images"]     = json::array();