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
+