Hugging Face inference API for image generation

This commit is contained in:
Alex Yancey 2024-08-19 23:02:15 -07:00
parent c70266984f
commit 135ba2336e
3 changed files with 88 additions and 0 deletions

View File

@ -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'));

View File

@ -49,6 +49,7 @@
<option value="auto">Stable Diffusion Web UI (AUTOMATIC1111)</option>
<option value="horde">Stable Horde</option>
<option value="togetherai">TogetherAI</option>
<option value="huggingface">HuggingFace (Image Inference Endpoint)</option>
</select>
<div data-sd-source="auto">
<label for="sd_auto_url">SD Web UI URL</label>
@ -82,6 +83,11 @@
<!-- (Original Text)<b>Important:</b> run DrawThings app with HTTP API switch enabled in the UI! The server must be accessible from the SillyTavern host machine. -->
<i><b data-i18n="Important:">Important:</b></i><i data-i18n="sd_drawthings_auth_txt"> run DrawThings app with HTTP API switch enabled in the UI! The server must be accessible from the SillyTavern host machine.</i>
</div>
<div data-sd-source="huggingface">
<i>Hint: Save an API key in the Hugging Face API settings to use it here.</i>
<label for="sd_huggingface_model_id" data-i18n="Model ID">Model ID</label>
<input id="sd_huggingface_model_id" type="text" class="text_pole" data-i18n="[placeholder]black-forest-labs/FLUX.1-dev" placeholder="black-forest-labs/FLUX.1-dev" value="" />
</div>
<div data-sd-source="vlad">
<label for="sd_vlad_url">SD.Next API URL</label>
<div class="flex-container flexnowrap">

View File

@ -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 };