Interrupt Comfy gens on cancel

52497ea96ddec110091d7f3f50d72457e032f3db

Cohee <18619528+Cohee1207@users.noreply.github.com>

src/endpoints/stable-diffusion.js+15 -0 Viewed

         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) => {
594	}	605	}
595	await delay(100);	606	await delay(100);
596	}	607	}
		608	if (item.status.status_str === 'error') {
		609	throw new Error('ComfyUI generation did not succeed.');
		610	}
597	const imgInfo = Object.keys(item.outputs).map(it => item.outputs[it].images).flat()[0];	611	const imgInfo = Object.keys(item.outputs).map(it => item.outputs[it].images).flat()[0];
598	const imgUrl = new URL(request.body.url);	612	const imgUrl = new URL(request.body.url);
599	imgUrl.pathname = '/view';	613	imgUrl.pathname = '/view';

         const imgBuffer = await imgResponse.buffer();
         return response.send(imgBuffer.toString('base64'));
     } catch (error) {
+        console.log(error);
         return response.sendStatus(500);
+    }
 });