diff --git a/public/scripts/extensions/stable-diffusion/index.js b/public/scripts/extensions/stable-diffusion/index.js index 821023adb..0c16b4393 100644 --- a/public/scripts/extensions/stable-diffusion/index.js +++ b/public/scripts/extensions/stable-diffusion/index.js @@ -52,6 +52,7 @@ const sources = { pollinations: 'pollinations', stability: 'stability', blockentropy: 'blockentropy', + huggingface: 'huggingface', }; const initiators = { @@ -454,6 +455,7 @@ async function loadSettings() { $('#sd_command_visible').prop('checked', extension_settings.sd.command_visible); $('#sd_interactive_visible').prop('checked', extension_settings.sd.interactive_visible); $('#sd_stability_style_preset').val(extension_settings.sd.stability_style_preset); + $('#sd_huggingface_model_id').val(extension_settings.sd.huggingface_model_id); for (const style of extension_settings.sd.styles) { const option = document.createElement('option'); @@ -1091,6 +1093,11 @@ function onComfyUrlInput() { saveSettingsDebounced(); } +function onHFModelInput() { + extension_settings.sd.huggingface_model_id = $('#sd_huggingface_model_id').val(); + saveSettingsDebounced(); +} + function onComfyWorkflowChange() { extension_settings.sd.comfy_workflow = $('#sd_comfy_workflow').find(':selected').val(); saveSettingsDebounced(); @@ -1235,7 +1242,16 @@ async function onModelChange() { extension_settings.sd.model = $('#sd_model').find(':selected').val(); saveSettingsDebounced(); - const cloudSources = [sources.horde, sources.novel, sources.openai, sources.togetherai, sources.pollinations, sources.stability, sources.blockentropy]; + const cloudSources = [ + sources.horde, + sources.novel, + sources.openai, + sources.togetherai, + sources.pollinations, + sources.stability, + sources.blockentropy, + sources.huggingface, + ]; if (cloudSources.includes(extension_settings.sd.source)) { return; @@ -1450,6 +1466,9 @@ async function loadSamplers() { case sources.blockentropy: samplers = ['N/A']; break; + case sources.huggingface: + samplers = ['N/A']; + break; } for (const sampler of samplers) { @@ -1639,6 +1658,9 @@ async function loadModels() { case sources.blockentropy: models = await loadBlockEntropyModels(); break; + case sources.huggingface: + models = [{ value: '', text: '' }]; + break; } for (const model of models) { @@ -1986,6 +2008,9 @@ async function loadSchedulers() { case sources.blockentropy: schedulers = ['N/A']; break; + case sources.huggingface: + schedulers = ['N/A']; + break; } for (const scheduler of schedulers) { @@ -2065,6 +2090,9 @@ async function loadVaes() { case sources.blockentropy: vaes = ['N/A']; break; + case sources.huggingface: + vaes = ['N/A']; + break; } for (const vae of vaes) { @@ -2596,6 +2624,9 @@ async function sendGenerationRequest(generationType, prompt, additionalNegativeP case sources.blockentropy: result = await generateBlockEntropyImage(prefixedPrompt, negativePrompt, signal); break; + case sources.huggingface: + result = await generateHuggingFaceImage(prefixedPrompt, signal); + break; } if (!result.data) { @@ -3229,6 +3260,34 @@ async function generateComfyImage(prompt, negativePrompt, signal) { return { format: 'png', data: await promptResult.text() }; } + +/** + * Generates an image in Hugging Face Inference API using the provided prompt and configuration settings (model selected). + * @param {string} prompt - The main instruction used to guide the image generation. + * @param {AbortSignal} signal - An AbortSignal object that can be used to cancel the request. + * @returns {Promise<{format: string, data: string}>} - A promise that resolves when the image generation and processing are complete. + */ +async function generateHuggingFaceImage(prompt, signal) { + const result = await fetch('/api/sd/huggingface/generate', { + method: 'POST', + headers: getRequestHeaders(), + signal: signal, + body: JSON.stringify({ + model: extension_settings.sd.huggingface_model_id, + prompt: prompt, + }), + }); + + if (result.ok) { + const data = await result.json(); + return { format: 'jpg', data: data.image }; + } else { + const text = await result.text(); + throw new Error(text); + } +} + + async function onComfyOpenWorkflowEditorClick() { let workflow = await (await fetch('/api/sd/comfy/workflow', { method: 'POST', @@ -3508,6 +3567,8 @@ function isValidState() { return secret_state[SECRET_KEYS.STABILITY]; case sources.blockentropy: return secret_state[SECRET_KEYS.BLOCKENTROPY]; + case sources.huggingface: + return secret_state[SECRET_KEYS.HUGGINGFACE]; } } @@ -3848,6 +3909,7 @@ jQuery(async () => { $('#sd_swap_dimensions').on('click', onSwapDimensionsClick); $('#sd_stability_key').on('click', onStabilityKeyClick); $('#sd_stability_style_preset').on('change', onStabilityStylePresetChange); + $('#sd_huggingface_model_id').on('input', onHFModelInput); $('.sd_settings .inline-drawer-toggle').on('click', function () { initScrollHeight($('#sd_prompt_prefix')); diff --git a/public/scripts/extensions/stable-diffusion/settings.html b/public/scripts/extensions/stable-diffusion/settings.html index b9f8b0a4c..b89366b7c 100644 --- a/public/scripts/extensions/stable-diffusion/settings.html +++ b/public/scripts/extensions/stable-diffusion/settings.html @@ -41,6 +41,7 @@ + @@ -82,6 +83,11 @@ Important: run DrawThings app with HTTP API switch enabled in the UI! The server must be accessible from the SillyTavern host machine. +
+ Hint: Save an API key in the Hugging Face (Text Completion) API settings to use it here. + + +
diff --git a/src/endpoints/stable-diffusion.js b/src/endpoints/stable-diffusion.js index 333bcd24d..daf8bc792 100644 --- a/src/endpoints/stable-diffusion.js +++ b/src/endpoints/stable-diffusion.js @@ -991,11 +991,52 @@ blockentropy.post('/generate', jsonParser, async (request, response) => { }); +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 };