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

Extend /sd command

Browse Source
This commit is contained in:
Cohee
2024-09-30 21:51:24 +03:00
parent dedb96ec8d
commit 9af4d62cdf
3 changed files with 261 additions and 140 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@ -9,41 +9,6 @@ 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
@ -391,40 +356,6 @@ router.post('/sd-next/upscalers', jsonParser, async (request, response) => {
}
});
/**
* 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) => {