Use Express router for stable diffusion endpoint
This commit is contained in:
valadaptive
parent
414c9bd5fb
commit
7f129df7b8
|
|
@ -3606,7 +3606,7 @@ require('./src/endpoints/sprites').registerEndpoints(app, jsonParser, urlencoded
|
|||
require('./src/endpoints/content-manager').registerEndpoints(app, jsonParser);
|
||||
|
||||
// Stable Diffusion generation
|
||||
require('./src/endpoints/stable-diffusion').registerEndpoints(app, jsonParser);
|
||||
app.use('/api/sd', require('./src/endpoints/stable-diffusion').router);
|
||||
|
||||
// LLM and SD Horde generation
|
||||
require('./src/endpoints/horde').registerEndpoints(app, jsonParser);
|
||||
|
|
|
|||
|
|
@ -1,9 +1,11 @@
|
|||
const express = require('express');
|
||||
const fetch = require('node-fetch').default;
|
||||
const sanitize = require('sanitize-filename');
|
||||
const { getBasicAuthHeader, delay } = require('../util.js');
|
||||
const fs = require('fs');
|
||||
const { DIRECTORIES } = require('../constants.js');
|
||||
const writeFileAtomicSync = require('write-file-atomic').sync;
|
||||
const { jsonParser } = require('../express-common');
|
||||
|
||||
/**
|
||||
* Sanitizes a string.
|
||||
|
|
@ -47,252 +49,34 @@ function getComfyWorkflows() {
|
|||
.sort(Intl.Collator().compare);
|
||||
}
|
||||
|
||||
/**
|
||||
* Registers the endpoints for the Stable Diffusion API extension.
|
||||
* @param {import("express").Express} app Express app
|
||||
* @param {any} jsonParser JSON parser middleware
|
||||
*/
|
||||
function registerEndpoints(app, jsonParser) {
|
||||
app.post('/api/sd/ping', jsonParser, async (request, response) => {
|
||||
try {
|
||||
const url = new URL(request.body.url);
|
||||
url.pathname = '/sdapi/v1/options';
|
||||
const router = express.Router();
|
||||
|
||||
const result = await fetch(url, {
|
||||
method: 'GET',
|
||||
headers: {
|
||||
'Authorization': getBasicAuthHeader(request.body.auth),
|
||||
},
|
||||
});
|
||||
router.post('/ping', jsonParser, async (request, response) => {
|
||||
try {
|
||||
const url = new URL(request.body.url);
|
||||
url.pathname = '/sdapi/v1/options';
|
||||
|
||||
if (!result.ok) {
|
||||
throw new Error('SD WebUI returned an error.');
|
||||
}
|
||||
const result = await fetch(url, {
|
||||
method: 'GET',
|
||||
headers: {
|
||||
'Authorization': getBasicAuthHeader(request.body.auth),
|
||||
},
|
||||
});
|
||||
|
||||
return response.sendStatus(200);
|
||||
} catch (error) {
|
||||
console.log(error);
|
||||
return response.sendStatus(500);
|
||||
if (!result.ok) {
|
||||
throw new Error('SD WebUI returned an error.');
|
||||
}
|
||||
});
|
||||
|
||||
app.post('/api/sd/upscalers', jsonParser, async (request, response) => {
|
||||
try {
|
||||
async function getUpscalerModels() {
|
||||
const url = new URL(request.body.url);
|
||||
url.pathname = '/sdapi/v1/upscalers';
|
||||
return response.sendStatus(200);
|
||||
} catch (error) {
|
||||
console.log(error);
|
||||
return response.sendStatus(500);
|
||||
}
|
||||
});
|
||||
|
||||
const result = await fetch(url, {
|
||||
method: 'GET',
|
||||
headers: {
|
||||
'Authorization': getBasicAuthHeader(request.body.auth),
|
||||
},
|
||||
});
|
||||
|
||||
if (!result.ok) {
|
||||
throw new Error('SD WebUI returned an error.');
|
||||
}
|
||||
|
||||
const data = await result.json();
|
||||
const names = data.map(x => x.name);
|
||||
return names;
|
||||
}
|
||||
|
||||
async function getLatentUpscalers() {
|
||||
const url = new URL(request.body.url);
|
||||
url.pathname = '/sdapi/v1/latent-upscale-modes';
|
||||
|
||||
const result = await fetch(url, {
|
||||
method: 'GET',
|
||||
headers: {
|
||||
'Authorization': getBasicAuthHeader(request.body.auth),
|
||||
},
|
||||
});
|
||||
|
||||
if (!result.ok) {
|
||||
throw new Error('SD WebUI returned an error.');
|
||||
}
|
||||
|
||||
const data = await result.json();
|
||||
const names = data.map(x => x.name);
|
||||
return names;
|
||||
}
|
||||
|
||||
const [upscalers, latentUpscalers] = await Promise.all([getUpscalerModels(), getLatentUpscalers()]);
|
||||
|
||||
// 0 = None, then Latent Upscalers, then Upscalers
|
||||
upscalers.splice(1, 0, ...latentUpscalers);
|
||||
|
||||
return response.send(upscalers);
|
||||
} catch (error) {
|
||||
console.log(error);
|
||||
return response.sendStatus(500);
|
||||
}
|
||||
});
|
||||
|
||||
app.post('/api/sd/samplers', jsonParser, async (request, response) => {
|
||||
try {
|
||||
const url = new URL(request.body.url);
|
||||
url.pathname = '/sdapi/v1/samplers';
|
||||
|
||||
const result = await fetch(url, {
|
||||
method: 'GET',
|
||||
headers: {
|
||||
'Authorization': getBasicAuthHeader(request.body.auth),
|
||||
},
|
||||
});
|
||||
|
||||
if (!result.ok) {
|
||||
throw new Error('SD WebUI returned an error.');
|
||||
}
|
||||
|
||||
const data = await result.json();
|
||||
const names = data.map(x => x.name);
|
||||
return response.send(names);
|
||||
|
||||
} catch (error) {
|
||||
console.log(error);
|
||||
return response.sendStatus(500);
|
||||
}
|
||||
});
|
||||
|
||||
app.post('/api/sd/models', jsonParser, async (request, response) => {
|
||||
try {
|
||||
const url = new URL(request.body.url);
|
||||
url.pathname = '/sdapi/v1/sd-models';
|
||||
|
||||
const result = await fetch(url, {
|
||||
method: 'GET',
|
||||
headers: {
|
||||
'Authorization': getBasicAuthHeader(request.body.auth),
|
||||
},
|
||||
});
|
||||
|
||||
if (!result.ok) {
|
||||
throw new Error('SD WebUI returned an error.');
|
||||
}
|
||||
|
||||
const data = await result.json();
|
||||
const models = data.map(x => ({ value: x.title, text: x.title }));
|
||||
return response.send(models);
|
||||
} catch (error) {
|
||||
console.log(error);
|
||||
return response.sendStatus(500);
|
||||
}
|
||||
});
|
||||
|
||||
app.post('/api/sd/get-model', jsonParser, async (request, response) => {
|
||||
try {
|
||||
const url = new URL(request.body.url);
|
||||
url.pathname = '/sdapi/v1/options';
|
||||
|
||||
const result = await fetch(url, {
|
||||
method: 'GET',
|
||||
headers: {
|
||||
'Authorization': getBasicAuthHeader(request.body.auth),
|
||||
},
|
||||
});
|
||||
const data = await result.json();
|
||||
return response.send(data['sd_model_checkpoint']);
|
||||
} catch (error) {
|
||||
console.log(error);
|
||||
return response.sendStatus(500);
|
||||
}
|
||||
});
|
||||
|
||||
app.post('/api/sd/set-model', jsonParser, async (request, response) => {
|
||||
try {
|
||||
async function getProgress() {
|
||||
const url = new URL(request.body.url);
|
||||
url.pathname = '/sdapi/v1/progress';
|
||||
|
||||
const result = await fetch(url, {
|
||||
method: 'GET',
|
||||
headers: {
|
||||
'Authorization': getBasicAuthHeader(request.body.auth),
|
||||
},
|
||||
timeout: 0,
|
||||
});
|
||||
const data = await result.json();
|
||||
return data;
|
||||
}
|
||||
|
||||
const url = new URL(request.body.url);
|
||||
url.pathname = '/sdapi/v1/options';
|
||||
|
||||
const options = {
|
||||
sd_model_checkpoint: request.body.model,
|
||||
};
|
||||
|
||||
const result = await fetch(url, {
|
||||
method: 'POST',
|
||||
body: JSON.stringify(options),
|
||||
headers: {
|
||||
'Content-Type': 'application/json',
|
||||
'Authorization': getBasicAuthHeader(request.body.auth),
|
||||
},
|
||||
timeout: 0,
|
||||
});
|
||||
|
||||
if (!result.ok) {
|
||||
throw new Error('SD WebUI returned an error.');
|
||||
}
|
||||
|
||||
const MAX_ATTEMPTS = 10;
|
||||
const CHECK_INTERVAL = 2000;
|
||||
|
||||
for (let attempt = 0; attempt < MAX_ATTEMPTS; attempt++) {
|
||||
const progressState = await getProgress();
|
||||
|
||||
const progress = progressState['progress'];
|
||||
const jobCount = progressState['state']['job_count'];
|
||||
if (progress == 0.0 && jobCount === 0) {
|
||||
break;
|
||||
}
|
||||
|
||||
console.log(`Waiting for SD WebUI to finish model loading... Progress: ${progress}; Job count: ${jobCount}`);
|
||||
await delay(CHECK_INTERVAL);
|
||||
}
|
||||
|
||||
return response.sendStatus(200);
|
||||
} catch (error) {
|
||||
console.log(error);
|
||||
return response.sendStatus(500);
|
||||
}
|
||||
});
|
||||
|
||||
app.post('/api/sd/generate', jsonParser, async (request, response) => {
|
||||
try {
|
||||
console.log('SD WebUI request:', request.body);
|
||||
|
||||
const url = new URL(request.body.url);
|
||||
url.pathname = '/sdapi/v1/txt2img';
|
||||
|
||||
const result = await fetch(url, {
|
||||
method: 'POST',
|
||||
body: JSON.stringify(request.body),
|
||||
headers: {
|
||||
'Content-Type': 'application/json',
|
||||
'Authorization': getBasicAuthHeader(request.body.auth),
|
||||
},
|
||||
timeout: 0,
|
||||
});
|
||||
|
||||
if (!result.ok) {
|
||||
const text = await result.text();
|
||||
throw new Error('SD WebUI returned an error.', { cause: text });
|
||||
}
|
||||
|
||||
const data = await result.json();
|
||||
return response.send(data);
|
||||
} catch (error) {
|
||||
console.log(error);
|
||||
return response.sendStatus(500);
|
||||
}
|
||||
});
|
||||
|
||||
app.post('/api/sd/sd-next/upscalers', jsonParser, async (request, response) => {
|
||||
try {
|
||||
router.post('/upscalers', jsonParser, async (request, response) => {
|
||||
try {
|
||||
async function getUpscalerModels() {
|
||||
const url = new URL(request.body.url);
|
||||
url.pathname = '/sdapi/v1/upscalers';
|
||||
|
||||
|
|
@ -307,245 +91,460 @@ function registerEndpoints(app, jsonParser) {
|
|||
throw new Error('SD WebUI returned an error.');
|
||||
}
|
||||
|
||||
// Vlad doesn't provide Latent Upscalers in the API, so we have to hardcode them here
|
||||
const latentUpscalers = ['Latent', 'Latent (antialiased)', 'Latent (bicubic)', 'Latent (bicubic antialiased)', 'Latent (nearest)', 'Latent (nearest-exact)'];
|
||||
const data = await result.json();
|
||||
const names = data.map(x => x.name);
|
||||
return names;
|
||||
}
|
||||
|
||||
async function getLatentUpscalers() {
|
||||
const url = new URL(request.body.url);
|
||||
url.pathname = '/sdapi/v1/latent-upscale-modes';
|
||||
|
||||
const result = await fetch(url, {
|
||||
method: 'GET',
|
||||
headers: {
|
||||
'Authorization': getBasicAuthHeader(request.body.auth),
|
||||
},
|
||||
});
|
||||
|
||||
if (!result.ok) {
|
||||
throw new Error('SD WebUI returned an error.');
|
||||
}
|
||||
|
||||
const data = await result.json();
|
||||
const names = data.map(x => x.name);
|
||||
|
||||
// 0 = None, then Latent Upscalers, then Upscalers
|
||||
names.splice(1, 0, ...latentUpscalers);
|
||||
|
||||
return response.send(names);
|
||||
} catch (error) {
|
||||
console.log(error);
|
||||
return response.sendStatus(500);
|
||||
}
|
||||
});
|
||||
|
||||
/**
|
||||
* SD prompt expansion using GPT-2 text generation model.
|
||||
* Adapted from: https://github.com/lllyasviel/Fooocus/blob/main/modules/expansion.py
|
||||
*/
|
||||
app.post('/api/sd/expand', jsonParser, async (request, response) => {
|
||||
const originalPrompt = request.body.prompt;
|
||||
|
||||
if (!originalPrompt) {
|
||||
console.warn('No prompt provided for SD expansion.');
|
||||
return response.send({ prompt: '' });
|
||||
return names;
|
||||
}
|
||||
|
||||
console.log('Refine prompt input:', originalPrompt);
|
||||
const splitString = splitStrings[Math.floor(Math.random() * splitStrings.length)];
|
||||
let prompt = safeStr(originalPrompt) + splitString;
|
||||
const [upscalers, latentUpscalers] = await Promise.all([getUpscalerModels(), getLatentUpscalers()]);
|
||||
|
||||
try {
|
||||
const task = 'text-generation';
|
||||
const module = await import('../transformers.mjs');
|
||||
const pipe = await module.default.getPipeline(task);
|
||||
// 0 = None, then Latent Upscalers, then Upscalers
|
||||
upscalers.splice(1, 0, ...latentUpscalers);
|
||||
|
||||
const result = await pipe(prompt, { num_beams: 1, max_new_tokens: 256, do_sample: true });
|
||||
return response.send(upscalers);
|
||||
} catch (error) {
|
||||
console.log(error);
|
||||
return response.sendStatus(500);
|
||||
}
|
||||
});
|
||||
|
||||
const newText = result[0].generated_text;
|
||||
const newPrompt = safeStr(removePattern(newText, dangerousPatterns));
|
||||
console.log('Refine prompt output:', newPrompt);
|
||||
router.post('/samplers', jsonParser, async (request, response) => {
|
||||
try {
|
||||
const url = new URL(request.body.url);
|
||||
url.pathname = '/sdapi/v1/samplers';
|
||||
|
||||
return response.send({ prompt: newPrompt });
|
||||
} catch {
|
||||
console.warn('Failed to load transformers.js pipeline.');
|
||||
return response.send({ prompt: originalPrompt });
|
||||
const result = await fetch(url, {
|
||||
method: 'GET',
|
||||
headers: {
|
||||
'Authorization': getBasicAuthHeader(request.body.auth),
|
||||
},
|
||||
});
|
||||
|
||||
if (!result.ok) {
|
||||
throw new Error('SD WebUI returned an error.');
|
||||
}
|
||||
});
|
||||
|
||||
app.post('/api/sd/comfy/ping', jsonParser, async (request, response) => {
|
||||
try {
|
||||
const data = await result.json();
|
||||
const names = data.map(x => x.name);
|
||||
return response.send(names);
|
||||
|
||||
} catch (error) {
|
||||
console.log(error);
|
||||
return response.sendStatus(500);
|
||||
}
|
||||
});
|
||||
|
||||
router.post('/models', jsonParser, async (request, response) => {
|
||||
try {
|
||||
const url = new URL(request.body.url);
|
||||
url.pathname = '/sdapi/v1/sd-models';
|
||||
|
||||
const result = await fetch(url, {
|
||||
method: 'GET',
|
||||
headers: {
|
||||
'Authorization': getBasicAuthHeader(request.body.auth),
|
||||
},
|
||||
});
|
||||
|
||||
if (!result.ok) {
|
||||
throw new Error('SD WebUI returned an error.');
|
||||
}
|
||||
|
||||
const data = await result.json();
|
||||
const models = data.map(x => ({ value: x.title, text: x.title }));
|
||||
return response.send(models);
|
||||
} catch (error) {
|
||||
console.log(error);
|
||||
return response.sendStatus(500);
|
||||
}
|
||||
});
|
||||
|
||||
router.post('/get-model', jsonParser, async (request, response) => {
|
||||
try {
|
||||
const url = new URL(request.body.url);
|
||||
url.pathname = '/sdapi/v1/options';
|
||||
|
||||
const result = await fetch(url, {
|
||||
method: 'GET',
|
||||
headers: {
|
||||
'Authorization': getBasicAuthHeader(request.body.auth),
|
||||
},
|
||||
});
|
||||
const data = await result.json();
|
||||
return response.send(data['sd_model_checkpoint']);
|
||||
} catch (error) {
|
||||
console.log(error);
|
||||
return response.sendStatus(500);
|
||||
}
|
||||
});
|
||||
|
||||
router.post('/set-model', jsonParser, async (request, response) => {
|
||||
try {
|
||||
async function getProgress() {
|
||||
const url = new URL(request.body.url);
|
||||
url.pathname = '/system_stats';
|
||||
url.pathname = '/sdapi/v1/progress';
|
||||
|
||||
const result = await fetch(url);
|
||||
if (!result.ok) {
|
||||
throw new Error('ComfyUI returned an error.');
|
||||
}
|
||||
|
||||
return response.sendStatus(200);
|
||||
} catch (error) {
|
||||
console.log(error);
|
||||
return response.sendStatus(500);
|
||||
}
|
||||
});
|
||||
|
||||
app.post('/api/sd/comfy/samplers', jsonParser, async (request, response) => {
|
||||
try {
|
||||
const url = new URL(request.body.url);
|
||||
url.pathname = '/object_info';
|
||||
|
||||
const result = await fetch(url);
|
||||
if (!result.ok) {
|
||||
throw new Error('ComfyUI returned an error.');
|
||||
}
|
||||
|
||||
const data = await result.json();
|
||||
return response.send(data.KSampler.input.required.sampler_name[0]);
|
||||
} catch (error) {
|
||||
console.log(error);
|
||||
return response.sendStatus(500);
|
||||
}
|
||||
});
|
||||
|
||||
app.post('/api/sd/comfy/models', jsonParser, async (request, response) => {
|
||||
try {
|
||||
const url = new URL(request.body.url);
|
||||
url.pathname = '/object_info';
|
||||
|
||||
const result = await fetch(url);
|
||||
if (!result.ok) {
|
||||
throw new Error('ComfyUI returned an error.');
|
||||
}
|
||||
const data = await result.json();
|
||||
return response.send(data.CheckpointLoaderSimple.input.required.ckpt_name[0].map(it => ({ value: it, text: it })));
|
||||
} catch (error) {
|
||||
console.log(error);
|
||||
return response.sendStatus(500);
|
||||
}
|
||||
});
|
||||
|
||||
app.post('/api/sd/comfy/schedulers', jsonParser, async (request, response) => {
|
||||
try {
|
||||
const url = new URL(request.body.url);
|
||||
url.pathname = '/object_info';
|
||||
|
||||
const result = await fetch(url);
|
||||
if (!result.ok) {
|
||||
throw new Error('ComfyUI returned an error.');
|
||||
}
|
||||
|
||||
const data = await result.json();
|
||||
return response.send(data.KSampler.input.required.scheduler[0]);
|
||||
} catch (error) {
|
||||
console.log(error);
|
||||
return response.sendStatus(500);
|
||||
}
|
||||
});
|
||||
|
||||
app.post('/api/sd/comfy/vaes', jsonParser, async (request, response) => {
|
||||
try {
|
||||
const url = new URL(request.body.url);
|
||||
url.pathname = '/object_info';
|
||||
|
||||
const result = await fetch(url);
|
||||
if (!result.ok) {
|
||||
throw new Error('ComfyUI returned an error.');
|
||||
}
|
||||
|
||||
const data = await result.json();
|
||||
return response.send(data.VAELoader.input.required.vae_name[0]);
|
||||
} catch (error) {
|
||||
console.log(error);
|
||||
return response.sendStatus(500);
|
||||
}
|
||||
});
|
||||
|
||||
app.post('/api/sd/comfy/workflows', jsonParser, async (request, response) => {
|
||||
try {
|
||||
const data = getComfyWorkflows();
|
||||
return response.send(data);
|
||||
} catch (error) {
|
||||
console.log(error);
|
||||
return response.sendStatus(500);
|
||||
}
|
||||
});
|
||||
|
||||
app.post('/api/sd/comfy/workflow', jsonParser, async (request, response) => {
|
||||
try {
|
||||
let path = `${DIRECTORIES.comfyWorkflows}/${sanitize(String(request.body.file_name))}`;
|
||||
if (!fs.existsSync(path)) {
|
||||
path = `${DIRECTORIES.comfyWorkflows}/Default_Comfy_Workflow.json`;
|
||||
}
|
||||
const data = fs.readFileSync(
|
||||
path,
|
||||
{ encoding: 'utf-8' },
|
||||
);
|
||||
return response.send(JSON.stringify(data));
|
||||
} catch (error) {
|
||||
console.log(error);
|
||||
return response.sendStatus(500);
|
||||
}
|
||||
});
|
||||
|
||||
app.post('/api/sd/comfy/save-workflow', jsonParser, async (request, response) => {
|
||||
try {
|
||||
writeFileAtomicSync(
|
||||
`${DIRECTORIES.comfyWorkflows}/${sanitize(String(request.body.file_name))}`,
|
||||
request.body.workflow,
|
||||
'utf8',
|
||||
);
|
||||
const data = getComfyWorkflows();
|
||||
return response.send(data);
|
||||
} catch (error) {
|
||||
console.log(error);
|
||||
return response.sendStatus(500);
|
||||
}
|
||||
});
|
||||
|
||||
app.post('/api/sd/comfy/delete-workflow', jsonParser, async (request, response) => {
|
||||
try {
|
||||
let path = `${DIRECTORIES.comfyWorkflows}/${sanitize(String(request.body.file_name))}`;
|
||||
if (fs.existsSync(path)) {
|
||||
fs.unlinkSync(path);
|
||||
}
|
||||
return response.sendStatus(200);
|
||||
} catch (error) {
|
||||
console.log(error);
|
||||
return response.sendStatus(500);
|
||||
}
|
||||
});
|
||||
|
||||
app.post('/api/sd/comfy/generate', jsonParser, async (request, response) => {
|
||||
try {
|
||||
const url = new URL(request.body.url);
|
||||
url.pathname = '/prompt';
|
||||
|
||||
const promptResult = await fetch(url, {
|
||||
method: 'POST',
|
||||
body: request.body.prompt,
|
||||
const result = await fetch(url, {
|
||||
method: 'GET',
|
||||
headers: {
|
||||
'Authorization': getBasicAuthHeader(request.body.auth),
|
||||
},
|
||||
timeout: 0,
|
||||
});
|
||||
if (!promptResult.ok) {
|
||||
throw new Error('ComfyUI returned an error.');
|
||||
}
|
||||
|
||||
const data = await promptResult.json();
|
||||
const id = data.prompt_id;
|
||||
let item;
|
||||
const historyUrl = new URL(request.body.url);
|
||||
historyUrl.pathname = '/history';
|
||||
while (true) {
|
||||
const result = await fetch(historyUrl);
|
||||
if (!result.ok) {
|
||||
throw new Error('ComfyUI returned an error.');
|
||||
}
|
||||
const history = await result.json();
|
||||
item = history[id];
|
||||
if (item) {
|
||||
break;
|
||||
}
|
||||
await delay(100);
|
||||
}
|
||||
const imgInfo = Object.keys(item.outputs).map(it => item.outputs[it].images).flat()[0];
|
||||
const imgUrl = new URL(request.body.url);
|
||||
imgUrl.pathname = '/view';
|
||||
imgUrl.search = `?filename=${imgInfo.filename}&subfolder=${imgInfo.subfolder}&type=${imgInfo.type}`;
|
||||
const imgResponse = await fetch(imgUrl);
|
||||
if (!imgResponse.ok) {
|
||||
throw new Error('ComfyUI returned an error.');
|
||||
}
|
||||
const imgBuffer = await imgResponse.buffer();
|
||||
return response.send(imgBuffer.toString('base64'));
|
||||
} catch (error) {
|
||||
return response.sendStatus(500);
|
||||
const data = await result.json();
|
||||
return data;
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
module.exports = {
|
||||
registerEndpoints,
|
||||
};
|
||||
const url = new URL(request.body.url);
|
||||
url.pathname = '/sdapi/v1/options';
|
||||
|
||||
const options = {
|
||||
sd_model_checkpoint: request.body.model,
|
||||
};
|
||||
|
||||
const result = await fetch(url, {
|
||||
method: 'POST',
|
||||
body: JSON.stringify(options),
|
||||
headers: {
|
||||
'Content-Type': 'application/json',
|
||||
'Authorization': getBasicAuthHeader(request.body.auth),
|
||||
},
|
||||
timeout: 0,
|
||||
});
|
||||
|
||||
if (!result.ok) {
|
||||
throw new Error('SD WebUI returned an error.');
|
||||
}
|
||||
|
||||
const MAX_ATTEMPTS = 10;
|
||||
const CHECK_INTERVAL = 2000;
|
||||
|
||||
for (let attempt = 0; attempt < MAX_ATTEMPTS; attempt++) {
|
||||
const progressState = await getProgress();
|
||||
|
||||
const progress = progressState['progress'];
|
||||
const jobCount = progressState['state']['job_count'];
|
||||
if (progress == 0.0 && jobCount === 0) {
|
||||
break;
|
||||
}
|
||||
|
||||
console.log(`Waiting for SD WebUI to finish model loading... Progress: ${progress}; Job count: ${jobCount}`);
|
||||
await delay(CHECK_INTERVAL);
|
||||
}
|
||||
|
||||
return response.sendStatus(200);
|
||||
} catch (error) {
|
||||
console.log(error);
|
||||
return response.sendStatus(500);
|
||||
}
|
||||
});
|
||||
|
||||
router.post('/generate', jsonParser, async (request, response) => {
|
||||
try {
|
||||
console.log('SD WebUI request:', request.body);
|
||||
|
||||
const url = new URL(request.body.url);
|
||||
url.pathname = '/sdapi/v1/txt2img';
|
||||
|
||||
const result = await fetch(url, {
|
||||
method: 'POST',
|
||||
body: JSON.stringify(request.body),
|
||||
headers: {
|
||||
'Content-Type': 'application/json',
|
||||
'Authorization': getBasicAuthHeader(request.body.auth),
|
||||
},
|
||||
timeout: 0,
|
||||
});
|
||||
|
||||
if (!result.ok) {
|
||||
const text = await result.text();
|
||||
throw new Error('SD WebUI returned an error.', { cause: text });
|
||||
}
|
||||
|
||||
const data = await result.json();
|
||||
return response.send(data);
|
||||
} catch (error) {
|
||||
console.log(error);
|
||||
return response.sendStatus(500);
|
||||
}
|
||||
});
|
||||
|
||||
router.post('/sd-next/upscalers', jsonParser, async (request, response) => {
|
||||
try {
|
||||
const url = new URL(request.body.url);
|
||||
url.pathname = '/sdapi/v1/upscalers';
|
||||
|
||||
const result = await fetch(url, {
|
||||
method: 'GET',
|
||||
headers: {
|
||||
'Authorization': getBasicAuthHeader(request.body.auth),
|
||||
},
|
||||
});
|
||||
|
||||
if (!result.ok) {
|
||||
throw new Error('SD WebUI returned an error.');
|
||||
}
|
||||
|
||||
// Vlad doesn't provide Latent Upscalers in the API, so we have to hardcode them here
|
||||
const latentUpscalers = ['Latent', 'Latent (antialiased)', 'Latent (bicubic)', 'Latent (bicubic antialiased)', 'Latent (nearest)', 'Latent (nearest-exact)'];
|
||||
|
||||
const data = await result.json();
|
||||
const names = data.map(x => x.name);
|
||||
|
||||
// 0 = None, then Latent Upscalers, then Upscalers
|
||||
names.splice(1, 0, ...latentUpscalers);
|
||||
|
||||
return response.send(names);
|
||||
} catch (error) {
|
||||
console.log(error);
|
||||
return response.sendStatus(500);
|
||||
}
|
||||
});
|
||||
|
||||
/**
|
||||
* SD prompt expansion using GPT-2 text generation model.
|
||||
* Adapted from: https://github.com/lllyasviel/Fooocus/blob/main/modules/expansion.py
|
||||
*/
|
||||
router.post('/expand', jsonParser, async (request, response) => {
|
||||
const originalPrompt = request.body.prompt;
|
||||
|
||||
if (!originalPrompt) {
|
||||
console.warn('No prompt provided for SD expansion.');
|
||||
return response.send({ prompt: '' });
|
||||
}
|
||||
|
||||
console.log('Refine prompt input:', originalPrompt);
|
||||
const splitString = splitStrings[Math.floor(Math.random() * splitStrings.length)];
|
||||
let prompt = safeStr(originalPrompt) + splitString;
|
||||
|
||||
try {
|
||||
const task = 'text-generation';
|
||||
const module = await import('../transformers.mjs');
|
||||
const pipe = await module.default.getPipeline(task);
|
||||
|
||||
const result = await pipe(prompt, { num_beams: 1, max_new_tokens: 256, do_sample: true });
|
||||
|
||||
const newText = result[0].generated_text;
|
||||
const newPrompt = safeStr(removePattern(newText, dangerousPatterns));
|
||||
console.log('Refine prompt output:', newPrompt);
|
||||
|
||||
return response.send({ prompt: newPrompt });
|
||||
} catch {
|
||||
console.warn('Failed to load transformers.js pipeline.');
|
||||
return response.send({ prompt: originalPrompt });
|
||||
}
|
||||
});
|
||||
|
||||
const comfy = express.Router();
|
||||
|
||||
comfy.post('/ping', jsonParser, async (request, response) => {
|
||||
try {
|
||||
const url = new URL(request.body.url);
|
||||
url.pathname = '/system_stats';
|
||||
|
||||
const result = await fetch(url);
|
||||
if (!result.ok) {
|
||||
throw new Error('ComfyUI returned an error.');
|
||||
}
|
||||
|
||||
return response.sendStatus(200);
|
||||
} catch (error) {
|
||||
console.log(error);
|
||||
return response.sendStatus(500);
|
||||
}
|
||||
});
|
||||
|
||||
comfy.post('/samplers', jsonParser, async (request, response) => {
|
||||
try {
|
||||
const url = new URL(request.body.url);
|
||||
url.pathname = '/object_info';
|
||||
|
||||
const result = await fetch(url);
|
||||
if (!result.ok) {
|
||||
throw new Error('ComfyUI returned an error.');
|
||||
}
|
||||
|
||||
const data = await result.json();
|
||||
return response.send(data.KSampler.input.required.sampler_name[0]);
|
||||
} catch (error) {
|
||||
console.log(error);
|
||||
return response.sendStatus(500);
|
||||
}
|
||||
});
|
||||
|
||||
comfy.post('/models', jsonParser, async (request, response) => {
|
||||
try {
|
||||
const url = new URL(request.body.url);
|
||||
url.pathname = '/object_info';
|
||||
|
||||
const result = await fetch(url);
|
||||
if (!result.ok) {
|
||||
throw new Error('ComfyUI returned an error.');
|
||||
}
|
||||
const data = await result.json();
|
||||
return response.send(data.CheckpointLoaderSimple.input.required.ckpt_name[0].map(it => ({ value: it, text: it })));
|
||||
} catch (error) {
|
||||
console.log(error);
|
||||
return response.sendStatus(500);
|
||||
}
|
||||
});
|
||||
|
||||
comfy.post('/schedulers', jsonParser, async (request, response) => {
|
||||
try {
|
||||
const url = new URL(request.body.url);
|
||||
url.pathname = '/object_info';
|
||||
|
||||
const result = await fetch(url);
|
||||
if (!result.ok) {
|
||||
throw new Error('ComfyUI returned an error.');
|
||||
}
|
||||
|
||||
const data = await result.json();
|
||||
return response.send(data.KSampler.input.required.scheduler[0]);
|
||||
} catch (error) {
|
||||
console.log(error);
|
||||
return response.sendStatus(500);
|
||||
}
|
||||
});
|
||||
|
||||
comfy.post('/vaes', jsonParser, async (request, response) => {
|
||||
try {
|
||||
const url = new URL(request.body.url);
|
||||
url.pathname = '/object_info';
|
||||
|
||||
const result = await fetch(url);
|
||||
if (!result.ok) {
|
||||
throw new Error('ComfyUI returned an error.');
|
||||
}
|
||||
|
||||
const data = await result.json();
|
||||
return response.send(data.VAELoader.input.required.vae_name[0]);
|
||||
} catch (error) {
|
||||
console.log(error);
|
||||
return response.sendStatus(500);
|
||||
}
|
||||
});
|
||||
|
||||
comfy.post('/workflows', jsonParser, async (request, response) => {
|
||||
try {
|
||||
const data = getComfyWorkflows();
|
||||
return response.send(data);
|
||||
} catch (error) {
|
||||
console.log(error);
|
||||
return response.sendStatus(500);
|
||||
}
|
||||
});
|
||||
|
||||
comfy.post('/workflow', jsonParser, async (request, response) => {
|
||||
try {
|
||||
let path = `${DIRECTORIES.comfyWorkflows}/${sanitize(String(request.body.file_name))}`;
|
||||
if (!fs.existsSync(path)) {
|
||||
path = `${DIRECTORIES.comfyWorkflows}/Default_Comfy_Workflow.json`;
|
||||
}
|
||||
const data = fs.readFileSync(
|
||||
path,
|
||||
{ encoding: 'utf-8' },
|
||||
);
|
||||
return response.send(JSON.stringify(data));
|
||||
} catch (error) {
|
||||
console.log(error);
|
||||
return response.sendStatus(500);
|
||||
}
|
||||
});
|
||||
|
||||
comfy.post('/save-workflow', jsonParser, async (request, response) => {
|
||||
try {
|
||||
writeFileAtomicSync(
|
||||
`${DIRECTORIES.comfyWorkflows}/${sanitize(String(request.body.file_name))}`,
|
||||
request.body.workflow,
|
||||
'utf8',
|
||||
);
|
||||
const data = getComfyWorkflows();
|
||||
return response.send(data);
|
||||
} catch (error) {
|
||||
console.log(error);
|
||||
return response.sendStatus(500);
|
||||
}
|
||||
});
|
||||
|
||||
comfy.post('/delete-workflow', jsonParser, async (request, response) => {
|
||||
try {
|
||||
let path = `${DIRECTORIES.comfyWorkflows}/${sanitize(String(request.body.file_name))}`;
|
||||
if (fs.existsSync(path)) {
|
||||
fs.unlinkSync(path);
|
||||
}
|
||||
return response.sendStatus(200);
|
||||
} catch (error) {
|
||||
console.log(error);
|
||||
return response.sendStatus(500);
|
||||
}
|
||||
});
|
||||
|
||||
comfy.post('/generate', jsonParser, async (request, response) => {
|
||||
try {
|
||||
const url = new URL(request.body.url);
|
||||
url.pathname = '/prompt';
|
||||
|
||||
const promptResult = await fetch(url, {
|
||||
method: 'POST',
|
||||
body: request.body.prompt,
|
||||
});
|
||||
if (!promptResult.ok) {
|
||||
throw new Error('ComfyUI returned an error.');
|
||||
}
|
||||
|
||||
const data = await promptResult.json();
|
||||
const id = data.prompt_id;
|
||||
let item;
|
||||
const historyUrl = new URL(request.body.url);
|
||||
historyUrl.pathname = '/history';
|
||||
while (true) {
|
||||
const result = await fetch(historyUrl);
|
||||
if (!result.ok) {
|
||||
throw new Error('ComfyUI returned an error.');
|
||||
}
|
||||
const history = await result.json();
|
||||
item = history[id];
|
||||
if (item) {
|
||||
break;
|
||||
}
|
||||
await delay(100);
|
||||
}
|
||||
const imgInfo = Object.keys(item.outputs).map(it => item.outputs[it].images).flat()[0];
|
||||
const imgUrl = new URL(request.body.url);
|
||||
imgUrl.pathname = '/view';
|
||||
imgUrl.search = `?filename=${imgInfo.filename}&subfolder=${imgInfo.subfolder}&type=${imgInfo.type}`;
|
||||
const imgResponse = await fetch(imgUrl);
|
||||
if (!imgResponse.ok) {
|
||||
throw new Error('ComfyUI returned an error.');
|
||||
}
|
||||
const imgBuffer = await imgResponse.buffer();
|
||||
return response.send(imgBuffer.toString('base64'));
|
||||
} catch (error) {
|
||||
return response.sendStatus(500);
|
||||
}
|
||||
});
|
||||
|
||||
router.use('/comfy', comfy);
|
||||
|
||||
module.exports = { router };
|
||||
|
|
|
|||
Loading…
Reference in New Issue