diff --git a/default/config.conf b/default/config.conf index 4c629e402..1e3c73b43 100644 --- a/default/config.conf +++ b/default/config.conf @@ -26,6 +26,8 @@ const extras = { captioningModel: 'Xenova/vit-gpt2-image-captioning', // Feature extraction model. HuggingFace ID of a model in ONNX format. embeddingModel: 'Xenova/all-mpnet-base-v2', + // GPT-2 text generation model. HuggingFace ID of a model in ONNX format. + promptExpansionModel: 'Cohee/fooocus_expansion-onnx', }; // Request overrides for additional headers diff --git a/public/scripts/extensions/stable-diffusion/index.js b/public/scripts/extensions/stable-diffusion/index.js index 523889ef8..2d962aacc 100644 --- a/public/scripts/extensions/stable-diffusion/index.js +++ b/public/scripts/extensions/stable-diffusion/index.js @@ -160,6 +160,7 @@ const defaultSettings = { // Refine mode refine_mode: false, + expand: false, prompts: promptTemplates, @@ -257,6 +258,7 @@ async function loadSettings() { $('#sd_restore_faces').prop('checked', extension_settings.sd.restore_faces); $('#sd_enable_hr').prop('checked', extension_settings.sd.enable_hr); $('#sd_refine_mode').prop('checked', extension_settings.sd.refine_mode); + $('#sd_expand').prop('checked', extension_settings.sd.expand); $('#sd_auto_url').val(extension_settings.sd.auto_url); $('#sd_auto_auth').val(extension_settings.sd.auto_auth); $('#sd_vlad_url').val(extension_settings.sd.vlad_url); @@ -300,7 +302,30 @@ function addPromptTemplates() { } } +async function expandPrompt(prompt) { + try { + const response = await fetch('/api/sd/expand', { + method: 'POST', + headers: getRequestHeaders(), + body: JSON.stringify({ prompt: prompt }), + }); + + if (!response.ok) { + throw new Error('API returned an error.'); + } + + const data = await response.json(); + return data.prompt; + } catch { + return prompt; + } +} + async function refinePrompt(prompt) { + if (extension_settings.sd.expand) { + prompt = await expandPrompt(prompt); + } + if (extension_settings.sd.refine_mode) { const refinedPrompt = await callPopup('

Review and edit the prompt:

Press "Cancel" to abort the image generation.', 'input', prompt.trim(), { rows: 5, okButton: 'Generate' }); @@ -361,6 +386,11 @@ function combinePrefixes(str1, str2) { return result; } +function onExpandInput() { + extension_settings.sd.expand = !!$(this).prop('checked'); + saveSettingsDebounced(); +} + function onRefineModeInput() { extension_settings.sd.refine_mode = !!$('#sd_refine_mode').prop('checked'); saveSettingsDebounced(); @@ -1610,6 +1640,7 @@ jQuery(async () => { $('#sd_novel_upscale_ratio').on('input', onNovelUpscaleRatioInput); $('#sd_novel_anlas_guard').on('input', onNovelAnlasGuardInput); $('#sd_novel_view_anlas').on('click', onViewAnlasClick); + $('#sd_expand').on('input', onExpandInput); $('#sd_character_prompt_block').hide(); $('.sd_settings .inline-drawer-toggle').on('click', function () { diff --git a/public/scripts/extensions/stable-diffusion/settings.html b/public/scripts/extensions/stable-diffusion/settings.html index 5e37fdc31..437f5909f 100644 --- a/public/scripts/extensions/stable-diffusion/settings.html +++ b/public/scripts/extensions/stable-diffusion/settings.html @@ -12,6 +12,10 @@ Edit prompts before generation +