rixty
/
SillyTavern
mirror of https://github.com/SillyTavern/SillyTavern.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Interrupt Comfy gens on cancel

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
Show Stats 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);
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);
}
});