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

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