From 52497ea96ddec110091d7f3f50d72457e032f3db Mon Sep 17 00:00:00 2001
From: Cohee <18619528+Cohee1207@users.noreply.github.com>
Date: Fri, 19 Jul 2024 00:39:21 +0300
Subject: [PATCH] Interrupt Comfy gens on cancel

---
 src/endpoints/stable-diffusion.js | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/src/endpoints/stable-diffusion.js b/src/endpoints/stable-diffusion.js
index 9f4a4e741..fe4a9604c 100644
--- a/src/endpoints/stable-diffusion.js
+++ b/src/endpoints/stable-diffusion.js
@@ -569,6 +569,17 @@ comfy.post('/generate', jsonParser, async (request, response) => {
         const url = new URL(request.body.url);
         url.pathname = '/prompt';
 
+        const controller = new AbortController();
+        request.socket.removeAllListeners('close');
+        request.socket.on('close', function () {
+            if (!response.writableEnded && !item) {
+                const interruptUrl = new URL(request.body.url);
+                interruptUrl.pathname = '/interrupt';
+                fetch(interruptUrl, { method: 'POST', headers: { 'Authorization': getBasicAuthHeader(request.body.auth) } });
+            }
+            controller.abort();
+        });
+
         const promptResult = await fetch(url, {
             method: 'POST',
             body: request.body.prompt,
@@ -594,6 +605,9 @@ comfy.post('/generate', jsonParser, async (request, response) => {
             }
             await delay(100);
         }
+        if (item.status.status_str === 'error') {
+            throw new Error('ComfyUI generation did not succeed.');
+        }
         const imgInfo = Object.keys(item.outputs).map(it => item.outputs[it].images).flat()[0];
         const imgUrl = new URL(request.body.url);
         imgUrl.pathname = '/view';
@@ -605,6 +619,7 @@ comfy.post('/generate', jsonParser, async (request, response) => {
         const imgBuffer = await imgResponse.buffer();
         return response.send(imgBuffer.toString('base64'));
     } catch (error) {
+        console.log(error);
         return response.sendStatus(500);
     }
 });