rixty/SillyTavern
1
0
Fork 0
mirror of https://github.com/SillyTavern/SillyTavern.git synced 2025-06-05 21:59:27 +02:00
Code Issues Packages Projects Releases Wiki Activity

Interrupt Comfy gens on cancel

Browse Source
This commit is contained in:
Cohee
2024-07-19 00:39:21 +03:00
parent 7e4abaabff
commit 52497ea96d
1 changed files with 15 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
src/endpoints/stable-diffusion.js
View File

@ -569,6 +569,17 @@ comfy.post('/generate', jsonParser, async (request, response) => {
const url = new URL(request.body.url); const url = new URL(request.body.url);
url.pathname = '/prompt'; 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, { const promptResult = await fetch(url, {
method: 'POST', method: 'POST',
body: request.body.prompt, body: request.body.prompt,
@ -594,6 +605,9 @@ comfy.post('/generate', jsonParser, async (request, response) => {
} }
await delay(100); 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 imgInfo = Object.keys(item.outputs).map(it => item.outputs[it].images).flat()[0];
const imgUrl = new URL(request.body.url); const imgUrl = new URL(request.body.url);
imgUrl.pathname = '/view'; imgUrl.pathname = '/view';
@ -605,6 +619,7 @@ comfy.post('/generate', jsonParser, async (request, response) => {
const imgBuffer = await imgResponse.buffer(); const imgBuffer = await imgResponse.buffer();
return response.send(imgBuffer.toString('base64')); return response.send(imgBuffer.toString('base64'));
} catch (error) { } catch (error) {
console.log(error);
return response.sendStatus(500); return response.sendStatus(500);
} }
}); });