From 135ba2336e78ad3464decdbbbc87b22aad51ac5e Mon Sep 17 00:00:00 2001 From: Alex Yancey Date: Mon, 19 Aug 2024 23:02:15 -0700 Subject: [PATCH] Hugging Face inference API for image generation --- .../extensions/stable-diffusion/index.js | 41 +++++++++++++++++++ .../extensions/stable-diffusion/settings.html | 6 +++ src/endpoints/stable-diffusion.js | 41 +++++++++++++++++++ 3 files changed, 88 insertions(+) diff --git a/public/scripts/extensions/stable-diffusion/index.js b/public/scripts/extensions/stable-diffusion/index.js index 821023adb..7f013eccc 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(); @@ -2596,6 +2603,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 +3239,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 +3546,8 @@ function isValidState() { return secret_state[SECRET_KEYS.STABILITY]; case sources.blockentropy: return secret_state[SECRET_KEYS.BLOCKENTROPY]; + case sources.huggingface: + return true; } } @@ -3848,6 +3888,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..e1a477a89 100644 --- a/public/scripts/extensions/stable-diffusion/settings.html +++ b/public/scripts/extensions/stable-diffusion/settings.html @@ -49,6 +49,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 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 };