diff --git a/public/scripts/extensions/stable-diffusion/index.js b/public/scripts/extensions/stable-diffusion/index.js index ad12381a7..c4af71d2b 100644 --- a/public/scripts/extensions/stable-diffusion/index.js +++ b/public/scripts/extensions/stable-diffusion/index.js @@ -3,9 +3,11 @@ import { saveSettingsDebounced, systemUserName, hideSwipeButtons, - showSwipeButtons + showSwipeButtons, + callPopup, + getRequestHeaders } from "../../../script.js"; -import { getApiUrl, getContext, extension_settings, defaultRequestArgs } from "../../extensions.js"; +import { getApiUrl, getContext, extension_settings, defaultRequestArgs, modules } from "../../extensions.js"; import { stringFormat, initScrollHeight, resetScrollHeight } from "../../utils.js"; export { MODULE_NAME }; @@ -94,6 +96,9 @@ const defaultSettings = { negative_prompt: 'lowres, bad anatomy, bad hands, text, error, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry', sampler: 'DDIM', model: '', + + horde: false, + horde_nsfw: false, } async function loadSettings() { @@ -107,11 +112,10 @@ async function loadSettings() { $('#sd_negative_prompt').val(extension_settings.sd.negative_prompt).trigger('input'); $('#sd_width').val(extension_settings.sd.width).trigger('input'); $('#sd_height').val(extension_settings.sd.height).trigger('input'); - + $('#sd_horde').prop('checked', extension_settings.sd.horde); + $('#sd_horde_nsfw').prop('checked', extension_settings.sd.horde_nsfw); await Promise.all([loadSamplers(), loadModels()]); - - } function onScaleInput() { @@ -155,10 +159,29 @@ function onHeightInput() { saveSettingsDebounced(); } +async function onHordeInput() { + extension_settings.sd.model = null; + extension_settings.sd.sampler = null; + extension_settings.sd.horde = !!$(this).prop('checked'); + saveSettingsDebounced(); + await Promise.all([loadModels(), loadSamplers()]); +} + +async function onHordeNsfwInput() { + extension_settings.sd.horde_nsfw = !!$(this).prop('checked'); + saveSettingsDebounced(); +} + async function onModelChange() { extension_settings.sd.model = $('#sd_model').find(':selected').val(); saveSettingsDebounced(); + if (extension_settings.sd.horde == false) { + await updateExtrasRemoteModel(); + } +} + +async function updateExtrasRemoteModel() { const url = new URL(getApiUrl()); url.pathname = '/api/image/model'; const getCurrentModelResult = await fetch(url, { @@ -173,25 +196,86 @@ async function onModelChange() { } async function loadSamplers() { + $('#sd_sampler').empty(); + let samplers = []; + + if (extension_settings.sd.horde) { + samplers = await loadHordeSamplers(); + } else { + samplers = await loadExtrasSamplers(); + } + + for (const sampler of samplers) { + const option = document.createElement('option'); + option.innerText = sampler; + option.value = sampler; + option.selected = sampler === extension_settings.sd.sampler; + $('#sd_sampler').append(option); + } +} + +async function loadHordeSamplers() { + const result = await fetch('/horde_samplers', { + method: 'POST', + headers: getRequestHeaders(), + }); + + if (result.ok) { + const data = await result.json(); + return data; + } + + return []; +} + +async function loadExtrasSamplers() { const url = new URL(getApiUrl()); url.pathname = '/api/image/samplers'; const result = await fetch(url, defaultRequestArgs); if (result.ok) { const data = await result.json(); - const samplers = data.samplers; - - for (const sampler of samplers) { - const option = document.createElement('option'); - option.innerText = sampler; - option.value = sampler; - option.selected = sampler === extension_settings.sd.sampler; - $('#sd_sampler').append(option); - } + return data.samplers; } + + return []; } async function loadModels() { + $('#sd_model').empty(); + let models = []; + + if (extension_settings.sd.horde) { + models = await loadHordeModels(); + } else { + models = await loadExtrasModels(); + } + + for (const model of models) { + const option = document.createElement('option'); + option.innerText = model.text; + option.value = model.value; + option.selected = model.value === extension_settings.sd.model; + $('#sd_model').append(option); + } +} + +async function loadHordeModels() { + const result = await fetch('/horde_models', { + method: 'POST', + headers: getRequestHeaders(), + }); + + if (result.ok) { + const data = await result.json(); + const models = data.map(x => ({ value: x.name, text: `${x.name} (ETA: ${x.eta}s, Queue: ${x.queued}, Workers: ${x.count})` })); + return models; + } + + return []; +} + +async function loadExtrasModels() { const url = new URL(getApiUrl()); url.pathname = '/api/image/model'; const getCurrentModelResult = await fetch(url, defaultRequestArgs); @@ -206,16 +290,11 @@ async function loadModels() { if (getModelsResult.ok) { const data = await getModelsResult.json(); - const models = data.models; - - for (const model of models) { - const option = document.createElement('option'); - option.innerText = model; - option.value = model; - option.selected = model === extension_settings.sd.model; - $('#sd_model').append(option); - } + const view_models = data.models.map(x => ({ value: x, text: x })); + return view_models; } + + return []; } function getGenerationType(prompt) { @@ -257,6 +336,14 @@ async function generatePicture(_, trigger) { return; } + if (!modules.includes('sd') && !extension_settings.sd.horde) { + callPopup("Extensions API is not connected or doesn't provide SD module. Enable Stable Horde to generate images.", 'text'); + return; + } + + extension_settings.sd.sampler = $('#sd_sampler').find(':selected').val(); + extension_settings.sd.model = $('#sd_model').find(':selected').val(); + trigger = trigger.trim(); const generationMode = getGenerationType(trigger); console.log('Generation mode', generationMode, 'triggered with', trigger); @@ -279,30 +366,10 @@ async function generatePicture(_, trigger) { console.log('Processed Stable Diffusion prompt:', prompt); - const url = new URL(getApiUrl()); - url.pathname = '/api/image'; - const result = await fetch(url, { - method: 'POST', - headers: postHeaders, - body: JSON.stringify({ - prompt: prompt, - sampler: extension_settings.sd.sampler, - steps: extension_settings.sd.steps, - scale: extension_settings.sd.scale, - width: extension_settings.sd.width, - height: extension_settings.sd.height, - prompt_prefix: extension_settings.sd.prompt_prefix, - negative_prompt: extension_settings.sd.negative_prompt, - restore_faces: true, - face_restoration_model: 'GFPGAN', - - }), - }); - - if (result.ok) { - const data = await result.json(); - const base64Image = `data:image/jpeg;base64,${data.image}`; - sendMessage(prompt, base64Image); + if (extension_settings.sd.horde) { + await generateHordeImage(prompt); + } else { + await generateExtrasImage(prompt); } } catch (err) { console.trace(err); @@ -314,6 +381,59 @@ async function generatePicture(_, trigger) { } } +async function generateExtrasImage(prompt) { + const url = new URL(getApiUrl()); + url.pathname = '/api/image'; + const result = await fetch(url, { + method: 'POST', + headers: postHeaders, + body: JSON.stringify({ + prompt: prompt, + sampler: extension_settings.sd.sampler, + steps: extension_settings.sd.steps, + scale: extension_settings.sd.scale, + width: extension_settings.sd.width, + height: extension_settings.sd.height, + prompt_prefix: extension_settings.sd.prompt_prefix, + negative_prompt: extension_settings.sd.negative_prompt, + restore_faces: true, + face_restoration_model: 'GFPGAN', + }), + }); + + + if (result.ok) { + const data = await result.json(); + const base64Image = `data:image/jpeg;base64,${data.image}`; + sendMessage(prompt, base64Image); + } +} + +async function generateHordeImage(prompt) { + const result = await fetch('/horde_generateimage', { + method: 'POST', + headers: getRequestHeaders(), + body: JSON.stringify({ + prompt: prompt, + sampler: extension_settings.sd.sampler, + steps: extension_settings.sd.steps, + scale: extension_settings.sd.scale, + width: extension_settings.sd.width, + height: extension_settings.sd.height, + prompt_prefix: extension_settings.sd.prompt_prefix, + negative_prompt: extension_settings.sd.negative_prompt, + model: extension_settings.sd.model, + nsfw: extension_settings.sd.horde_nsfw, + }), + }); + + if (result.ok) { + const data = await result.text(); + const base64Image = `data:image/webp;base64,${data}`; + sendMessage(prompt, base64Image); + } +} + async function sendMessage(prompt, image) { const context = getContext(); const messageText = `[${context.name2} sends a picture that contains: ${prompt}]`; @@ -386,16 +506,6 @@ function addSDGenButtons() { async function moduleWorker() { const context = getContext(); - /* if (context.onlineStatus === 'no_connection') { - $('#sd_gen').hide(200); - } else if ($("#send_but").css('display') === 'flex') { - $('#sd_gen').show(200); - $("#sd_gen_wait").hide(200); - } else { - $('#sd_gen').hide(200); - $("#sd_gen_wait").show(200); - } */ - context.onlineStatus === 'no_connection' ? $('#sd_gen').hide(200) : $('#sd_gen').show(200) @@ -444,6 +554,16 @@ jQuery(async () => {