const express = require('express'); const fetch = require('node-fetch').default; const sanitize = require('sanitize-filename'); const { getBasicAuthHeader, delay, getHexString } = require('../util.js'); const fs = require('fs'); const path = require('path'); const writeFileAtomicSync = require('write-file-atomic').sync; const { jsonParser } = require('../express-common'); const { readSecret, SECRET_KEYS } = require('./secrets.js'); const FormData = require('form-data'); /** * Sanitizes a string. * @param {string} x String to sanitize * @returns {string} Sanitized string */ function safeStr(x) { x = String(x); x = x.replace(/ +/g, ' '); x = x.trim(); x = x.replace(/^[\s,.]+|[\s,.]+$/g, ''); return x; } const splitStrings = [ ', extremely', ', intricate,', ]; const dangerousPatterns = '[]【】()()|::'; /** * Removes patterns from a string. * @param {string} x String to sanitize * @param {string} pattern Pattern to remove * @returns {string} Sanitized string */ function removePattern(x, pattern) { for (let i = 0; i < pattern.length; i++) { let p = pattern[i]; let regex = new RegExp('\\' + p, 'g'); x = x.replace(regex, ''); } return x; } /** * Gets the comfy workflows. * @param {import('../users.js').UserDirectoryList} directories * @returns {string[]} List of comfy workflows */ function getComfyWorkflows(directories) { return fs .readdirSync(directories.comfyWorkflows) .filter(file => file[0] != '.' && file.toLowerCase().endsWith('.json')) .sort(Intl.Collator().compare); } const router = express.Router(); router.post('/ping', 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), }, }); if (!result.ok) { throw new Error('SD WebUI returned an error.'); } return response.sendStatus(200); } catch (error) { console.log(error); return response.sendStatus(500); } }); router.post('/upscalers', jsonParser, async (request, response) => { try { async function getUpscalerModels() { 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.'); } 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); } }); router.post('/vaes', jsonParser, async (request, response) => { try { const url = new URL(request.body.url); url.pathname = '/sdapi/v1/sd-vae'; 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.model_name); return response.send(names); } catch (error) { console.log(error); return response.sendStatus(500); } }); router.post('/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); } }); router.post('/schedulers', jsonParser, async (request, response) => { try { const url = new URL(request.body.url); url.pathname = '/sdapi/v1/schedulers'; 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); } }); 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 = '/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); } }); 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 controller = new AbortController(); request.socket.removeAllListeners('close'); request.socket.on('close', function () { if (!response.writableEnded) { const url = new URL(request.body.url); url.pathname = '/sdapi/v1/interrupt'; fetch(url, { method: 'POST', headers: { 'Authorization': getBasicAuthHeader(request.body.auth) } }); } controller.abort(); }); const result = await fetch(url, { method: 'POST', body: JSON.stringify(request.body), headers: { 'Content-Type': 'application/json', 'Authorization': getBasicAuthHeader(request.body.auth), }, timeout: 0, // @ts-ignore signal: controller.signal, }); 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(request.user.directories); return response.send(data); } catch (error) { console.log(error); return response.sendStatus(500); } }); comfy.post('/workflow', jsonParser, async (request, response) => { try { let filePath = path.join(request.user.directories.comfyWorkflows, sanitize(String(request.body.file_name))); if (!fs.existsSync(filePath)) { filePath = path.join(request.user.directories.comfyWorkflows, 'Default_Comfy_Workflow.json'); } const data = fs.readFileSync(filePath, { 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 { const filePath = path.join(request.user.directories.comfyWorkflows, sanitize(String(request.body.file_name))); writeFileAtomicSync(filePath, request.body.workflow, 'utf8'); const data = getComfyWorkflows(request.user.directories); return response.send(data); } catch (error) { console.log(error); return response.sendStatus(500); } }); comfy.post('/delete-workflow', jsonParser, async (request, response) => { try { const filePath = path.join(request.user.directories.comfyWorkflows, sanitize(String(request.body.file_name))); if (fs.existsSync(filePath)) { fs.unlinkSync(filePath); } 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 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, }); 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); } 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'; 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) { console.log(error); return response.sendStatus(500); } }); const together = express.Router(); together.post('/models', jsonParser, async (request, response) => { try { const key = readSecret(request.user.directories, SECRET_KEYS.TOGETHERAI); if (!key) { console.log('TogetherAI key not found.'); return response.sendStatus(400); } const modelsResponse = await fetch('https://api.together.xyz/api/models', { method: 'GET', headers: { 'Authorization': `Bearer ${key}`, }, }); if (!modelsResponse.ok) { console.log('TogetherAI returned an error.'); return response.sendStatus(500); } const data = await modelsResponse.json(); if (!Array.isArray(data)) { console.log('TogetherAI returned invalid data.'); return response.sendStatus(500); } const models = data .filter(x => x.display_type === 'image') .map(x => ({ value: x.name, text: x.display_name })); return response.send(models); } catch (error) { console.log(error); return response.sendStatus(500); } }); together.post('/generate', jsonParser, async (request, response) => { try { const key = readSecret(request.user.directories, SECRET_KEYS.TOGETHERAI); if (!key) { console.log('TogetherAI key not found.'); return response.sendStatus(400); } console.log('TogetherAI request:', request.body); const result = await fetch('https://api.together.xyz/api/inference', { method: 'POST', body: JSON.stringify({ request_type: 'image-model-inference', prompt: request.body.prompt, negative_prompt: request.body.negative_prompt, height: request.body.height, width: request.body.width, model: request.body.model, steps: request.body.steps, n: 1, // Limited to 10000 on playground, works fine with more. seed: request.body.seed >= 0 ? request.body.seed : Math.floor(Math.random() * 10_000_000), // Don't know if that's supposed to be random or not. It works either way. sessionKey: getHexString(40), }), headers: { 'Content-Type': 'application/json', 'Authorization': `Bearer ${key}`, }, }); if (!result.ok) { console.log('TogetherAI returned an error.'); return response.sendStatus(500); } const data = await result.json(); console.log('TogetherAI response:', data); if (data.status !== 'finished') { console.log('TogetherAI job failed.'); return response.sendStatus(500); } return response.send(data); } catch (error) { console.log(error); return response.sendStatus(500); } }); const drawthings = express.Router(); drawthings.post('/ping', jsonParser, async (request, response) => { try { const url = new URL(request.body.url); url.pathname = '/'; const result = await fetch(url, { method: 'HEAD', }); if (!result.ok) { throw new Error('SD DrawThings API returned an error.'); } return response.sendStatus(200); } catch (error) { console.log(error); return response.sendStatus(500); } }); drawthings.post('/get-model', jsonParser, async (request, response) => { try { const url = new URL(request.body.url); url.pathname = '/'; const result = await fetch(url, { method: 'GET', }); const data = await result.json(); return response.send(data['model']); } catch (error) { console.log(error); return response.sendStatus(500); } }); drawthings.post('/get-upscaler', jsonParser, async (request, response) => { try { const url = new URL(request.body.url); url.pathname = '/'; const result = await fetch(url, { method: 'GET', }); const data = await result.json(); return response.send(data['upscaler']); } catch (error) { console.log(error); return response.sendStatus(500); } }); drawthings.post('/generate', jsonParser, async (request, response) => { try { console.log('SD DrawThings API request:', request.body); const url = new URL(request.body.url); url.pathname = '/sdapi/v1/txt2img'; const body = { ...request.body }; const auth = getBasicAuthHeader(request.body.auth); delete body.url; delete body.auth; const result = await fetch(url, { method: 'POST', body: JSON.stringify(body), headers: { 'Content-Type': 'application/json', 'Authorization': auth, }, timeout: 0, }); if (!result.ok) { const text = await result.text(); throw new Error('SD DrawThings API returned an error.', { cause: text }); } const data = await result.json(); return response.send(data); } catch (error) { console.log(error); return response.sendStatus(500); } }); const pollinations = express.Router(); pollinations.post('/generate', jsonParser, async (request, response) => { try { const promptUrl = new URL(`https://image.pollinations.ai/prompt/${encodeURIComponent(request.body.prompt)}`); const params = new URLSearchParams({ model: String(request.body.model), negative_prompt: String(request.body.negative_prompt), seed: String(request.body.seed >= 0 ? request.body.seed : Math.floor(Math.random() * 10_000_000)), enhance: String(request.body.enhance ?? false), refine: String(request.body.refine ?? false), width: String(request.body.width ?? 1024), height: String(request.body.height ?? 1024), nologo: String(true), nofeed: String(true), referer: 'sillytavern', }); promptUrl.search = params.toString(); console.log('Pollinations request URL:', promptUrl.toString()); const result = await fetch(promptUrl); if (!result.ok) { console.log('Pollinations returned an error.', result.status, result.statusText); throw new Error('Pollinations request failed.'); } const buffer = await result.buffer(); const base64 = buffer.toString('base64'); return response.send({ image: base64 }); } catch (error) { console.log(error); return response.sendStatus(500); } }); const stability = express.Router(); stability.post('/generate', jsonParser, async (request, response) => { try { const key = readSecret(request.user.directories, SECRET_KEYS.STABILITY); if (!key) { console.log('Stability AI key not found.'); return response.sendStatus(400); } const { payload, model } = request.body; console.log('Stability AI request:', model, payload); const formData = new FormData(); for (const [key, value] of Object.entries(payload)) { if (value !== undefined) { formData.append(key, String(value)); } } let apiUrl; switch (model) { case 'stable-image-ultra': apiUrl = 'https://api.stability.ai/v2beta/stable-image/generate/ultra'; break; case 'stable-image-core': apiUrl = 'https://api.stability.ai/v2beta/stable-image/generate/core'; break; case 'stable-diffusion-3': apiUrl = 'https://api.stability.ai/v2beta/stable-image/generate/sd3'; break; default: throw new Error('Invalid Stability AI model selected'); } const result = await fetch(apiUrl, { method: 'POST', headers: { 'Authorization': `Bearer ${key}`, 'Accept': 'image/*', }, body: formData, timeout: 0, }); if (!result.ok) { const text = await result.text(); console.log('Stability AI returned an error.', result.status, result.statusText, text); return response.sendStatus(500); } const buffer = await result.buffer(); return response.send(buffer.toString('base64')); } catch (error) { console.log(error); return response.sendStatus(500); } }); const blockentropy = express.Router(); blockentropy.post('/models', jsonParser, async (request, response) => { try { const key = readSecret(request.user.directories, SECRET_KEYS.BLOCKENTROPY); if (!key) { console.log('Block Entropy key not found.'); return response.sendStatus(400); } const modelsResponse = await fetch('https://api.blockentropy.ai/sdapi/v1/sd-models', { method: 'GET', headers: { 'Authorization': `Bearer ${key}`, }, }); if (!modelsResponse.ok) { console.log('Block Entropy returned an error.'); return response.sendStatus(500); } const data = await modelsResponse.json(); if (!Array.isArray(data)) { console.log('Block Entropy returned invalid data.'); return response.sendStatus(500); } const models = data.map(x => ({ value: x.name, text: x.name })); return response.send(models); } catch (error) { console.log(error); return response.sendStatus(500); } }); blockentropy.post('/generate', jsonParser, async (request, response) => { try { const key = readSecret(request.user.directories, SECRET_KEYS.BLOCKENTROPY); if (!key) { console.log('Block Entropy key not found.'); return response.sendStatus(400); } console.log('Block Entropy request:', request.body); const result = await fetch('https://api.blockentropy.ai/sdapi/v1/txt2img', { method: 'POST', body: JSON.stringify({ prompt: request.body.prompt, negative_prompt: request.body.negative_prompt, model: request.body.model, steps: request.body.steps, width: request.body.width, height: request.body.height, // Random seed if negative. seed: request.body.seed >= 0 ? request.body.seed : Math.floor(Math.random() * 10_000_000), }), headers: { 'Content-Type': 'application/json', 'Authorization': `Bearer ${key}`, }, }); if (!result.ok) { console.log('Block Entropy returned an error.'); return response.sendStatus(500); } const data = await result.json(); console.log('Block Entropy response:', data); return response.send(data); } catch (error) { console.log(error); return response.sendStatus(500); } }); const huggingface = express.Router(); huggingface.post('/generate', jsonParser, async (request, response) => { try { const key = readSecret(request.user.directories, SECRET_KEYS.HUGGINGFACE); if (!key) { console.log('Hugging Face key not found.'); return response.sendStatus(400); } console.log('Hugging Face request:', request.body); const result = await fetch(`https://api-inference.huggingface.co/models/${request.body.model}`, { method: 'POST', body: JSON.stringify({ inputs: request.body.prompt, }), headers: { 'Content-Type': 'application/json', 'Authorization': `Bearer ${key}`, }, }); if (!result.ok) { console.log('Hugging Face returned an error.'); return response.sendStatus(500); } const buffer = await result.buffer(); return response.send({ image: buffer.toString('base64'), }); } catch (error) { console.log(error); return response.sendStatus(500); } }); router.use('/comfy', comfy); router.use('/together', together); router.use('/drawthings', drawthings); router.use('/pollinations', pollinations); router.use('/stability', stability); router.use('/blockentropy', blockentropy); router.use('/huggingface', huggingface); module.exports = { router };