From 135ba2336e78ad3464decdbbbc87b22aad51ac5e Mon Sep 17 00:00:00 2001 From: Alex Yancey Date: Mon, 19 Aug 2024 23:02:15 -0700 Subject: [PATCH 1/4] 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 }; From 48f0c48a5bc7771d465ba2bc2af68b0781da04f1 Mon Sep 17 00:00:00 2001 From: Alex Yancey Date: Tue, 20 Aug 2024 01:31:57 -0700 Subject: [PATCH 2/4] Update public/scripts/extensions/stable-diffusion/index.js Co-authored-by: Cohee <18619528+Cohee1207@users.noreply.github.com> --- public/scripts/extensions/stable-diffusion/index.js | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/public/scripts/extensions/stable-diffusion/index.js b/public/scripts/extensions/stable-diffusion/index.js index 7f013eccc..11375b32e 100644 --- a/public/scripts/extensions/stable-diffusion/index.js +++ b/public/scripts/extensions/stable-diffusion/index.js @@ -3547,7 +3547,7 @@ function isValidState() { case sources.blockentropy: return secret_state[SECRET_KEYS.BLOCKENTROPY]; case sources.huggingface: - return true; + return secret_state[SECRET_KEYS.HUGGINGFACE]; } } From ea0508b91049a788d9c86e64ac4fefe7bdce8f4b Mon Sep 17 00:00:00 2001 From: Alex Yancey Date: Tue, 20 Aug 2024 01:33:14 -0700 Subject: [PATCH 3/4] Update settings.html --- public/scripts/extensions/stable-diffusion/settings.html | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/public/scripts/extensions/stable-diffusion/settings.html b/public/scripts/extensions/stable-diffusion/settings.html index e1a477a89..fc514001c 100644 --- a/public/scripts/extensions/stable-diffusion/settings.html +++ b/public/scripts/extensions/stable-diffusion/settings.html @@ -86,7 +86,7 @@
Hint: Save an API key in the Hugging Face API settings to use it here. - +
From a9700c0ff0764ed9fc8e1c93250e92f26df969cd Mon Sep 17 00:00:00 2001 From: Cohee <18619528+Cohee1207@users.noreply.github.com> Date: Tue, 20 Aug 2024 10:53:37 +0000 Subject: [PATCH 4/4] Small fixes --- .../extensions/stable-diffusion/index.js | 23 ++++++++++++++++++- .../extensions/stable-diffusion/settings.html | 6 ++--- 2 files changed, 25 insertions(+), 4 deletions(-) diff --git a/public/scripts/extensions/stable-diffusion/index.js b/public/scripts/extensions/stable-diffusion/index.js index 11375b32e..0c16b4393 100644 --- a/public/scripts/extensions/stable-diffusion/index.js +++ b/public/scripts/extensions/stable-diffusion/index.js @@ -1242,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; @@ -1457,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) { @@ -1646,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) { @@ -1993,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) { @@ -2072,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) { diff --git a/public/scripts/extensions/stable-diffusion/settings.html b/public/scripts/extensions/stable-diffusion/settings.html index fc514001c..b89366b7c 100644 --- a/public/scripts/extensions/stable-diffusion/settings.html +++ b/public/scripts/extensions/stable-diffusion/settings.html @@ -41,6 +41,7 @@ + @@ -49,7 +50,6 @@ -
@@ -84,9 +84,9 @@ 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. + Hint: Save an API key in the Hugging Face (Text Completion) API settings to use it here. - +