From 9af4d62cdfbb7842bb4b665fc90652039b71100c Mon Sep 17 00:00:00 2001 From: Cohee <18619528+Cohee1207@users.noreply.github.com> Date: Mon, 30 Sep 2024 21:51:24 +0300 Subject: [PATCH 1/7] Extend /sd command --- .../extensions/stable-diffusion/index.js | 317 ++++++++++++++---- .../extensions/stable-diffusion/settings.html | 15 +- src/endpoints/stable-diffusion.js | 69 ---- 3 files changed, 261 insertions(+), 140 deletions(-) diff --git a/public/scripts/extensions/stable-diffusion/index.js b/public/scripts/extensions/stable-diffusion/index.js index ab0b1bcb5..18da313b5 100644 --- a/public/scripts/extensions/stable-diffusion/index.js +++ b/public/scripts/extensions/stable-diffusion/index.js @@ -20,7 +20,7 @@ import { } from '../../../script.js'; import { getApiUrl, getContext, extension_settings, doExtrasFetch, modules, renderExtensionTemplateAsync, writeExtensionField } from '../../extensions.js'; import { selected_group } from '../../group-chats.js'; -import { stringFormat, initScrollHeight, resetScrollHeight, getCharaFilename, saveBase64AsFile, getBase64Async, delay, isTrueBoolean, debounce } from '../../utils.js'; +import { stringFormat, initScrollHeight, resetScrollHeight, getCharaFilename, saveBase64AsFile, getBase64Async, delay, isTrueBoolean, debounce, isFalseBoolean } from '../../utils.js'; import { getMessageTimeStamp, humanizedDateTime } from '../../RossAscends-mods.js'; import { SECRET_KEYS, secret_state, writeSecret } from '../../secrets.js'; import { getNovelUnlimitedImageGeneration, getNovelAnlas, loadNovelSubscriptionData } from '../../nai-settings.js'; @@ -31,6 +31,7 @@ import { ARGUMENT_TYPE, SlashCommandArgument, SlashCommandNamedArgument } from ' import { debounce_timeout } from '../../constants.js'; import { SlashCommandEnumValue } from '../../slash-commands/SlashCommandEnumValue.js'; import { POPUP_RESULT, POPUP_TYPE, Popup, callGenericPopup } from '../../popup.js'; +import { commonEnumProviders } from '../../slash-commands/SlashCommandCommonEnumsProvider.js'; export { MODULE_NAME }; const MODULE_NAME = 'sd'; @@ -221,7 +222,6 @@ const defaultSettings = { // Refine mode refine_mode: false, - expand: false, interactive_mode: false, multimodal_captioning: false, snap: false, @@ -240,7 +240,7 @@ const defaultSettings = { drawthings_auth: '', hr_upscaler: 'Latent', - hr_scale: 2.0, + hr_scale: 1.0, hr_scale_min: 1.0, hr_scale_max: 4.0, hr_scale_step: 0.1, @@ -260,10 +260,6 @@ const defaultSettings = { clip_skip: 1, // NovelAI settings - novel_upscale_ratio_min: 1.0, - novel_upscale_ratio_max: 4.0, - novel_upscale_ratio_step: 0.1, - novel_upscale_ratio: 1.0, novel_anlas_guard: false, novel_sm: false, novel_sm_dyn: false, @@ -416,7 +412,6 @@ async function loadSettings() { $('#sd_hr_scale').val(extension_settings.sd.hr_scale).trigger('input'); $('#sd_denoising_strength').val(extension_settings.sd.denoising_strength).trigger('input'); $('#sd_hr_second_pass_steps').val(extension_settings.sd.hr_second_pass_steps).trigger('input'); - $('#sd_novel_upscale_ratio').val(extension_settings.sd.novel_upscale_ratio).trigger('input'); $('#sd_novel_anlas_guard').prop('checked', extension_settings.sd.novel_anlas_guard); $('#sd_novel_sm').prop('checked', extension_settings.sd.novel_sm); $('#sd_novel_sm_dyn').prop('checked', extension_settings.sd.novel_sm_dyn); @@ -430,7 +425,6 @@ 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_multimodal_captioning').prop('checked', extension_settings.sd.multimodal_captioning); $('#sd_auto_url').val(extension_settings.sd.auto_url); $('#sd_auto_auth').val(extension_settings.sd.auto_auth); @@ -644,37 +638,13 @@ async function onSaveStyleClick() { saveSettingsDebounced(); } -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; - } -} - /** - * Modifies prompt based on auto-expansion and user inputs. + * Modifies prompt based on user inputs. * @param {string} prompt Prompt to refine - * @param {boolean} allowExpand Whether to allow auto-expansion * @param {boolean} isNegative Whether the prompt is a negative one * @returns {Promise} Refined prompt */ -async function refinePrompt(prompt, allowExpand, isNegative = false) { - if (allowExpand && extension_settings.sd.expand) { - prompt = await expandPrompt(prompt); - } - +async function refinePrompt(prompt, isNegative) { if (extension_settings.sd.refine_mode) { const text = isNegative ? '

Review and edit the negative prompt:

' : '

Review and edit the prompt:

'; const refinedPrompt = await callGenericPopup(text + 'Press "Cancel" to abort the image generation.', POPUP_TYPE.INPUT, prompt.trim(), { rows: 5, okButton: 'Continue' }); @@ -800,11 +770,6 @@ function combinePrefixes(str1, str2, macro = '') { return process(result); } -function onExpandInput() { - extension_settings.sd.expand = !!$(this).prop('checked'); - saveSettingsDebounced(); -} - function onRefineModeInput() { extension_settings.sd.refine_mode = !!$('#sd_refine_mode').prop('checked'); saveSettingsDebounced(); @@ -969,12 +934,6 @@ async function onViewAnlasClick() { toastr.info(`Free image generation: ${unlimitedGeneration ? 'Yes' : 'No'}`, `Anlas: ${anlas}`); } -function onNovelUpscaleRatioInput() { - extension_settings.sd.novel_upscale_ratio = Number($('#sd_novel_upscale_ratio').val()); - $('#sd_novel_upscale_ratio_value').val(extension_settings.sd.novel_upscale_ratio.toFixed(1)); - saveSettingsDebounced(); -} - function onNovelAnlasGuardInput() { extension_settings.sd.novel_anlas_guard = !!$('#sd_novel_anlas_guard').prop('checked'); saveSettingsDebounced(); @@ -2272,6 +2231,23 @@ function getRawLastMessage() { return `((${processReply(lastMessage.mes)})), (${processReply(character.scenario)}:0.7), (${processReply(character.description)}:0.5)`; } +/** + * Ensure that the selected option exists in the dropdown. + * @param {string} setting Setting key + * @param {string} selector Dropdown selector + * @returns {void} + */ +function ensureSelectionExists(setting, selector) { + /** @type {HTMLSelectElement} */ + const selectElement = document.querySelector(selector); + if (!selectElement) { + return; + } + if (selectElement.selectedOptions.length && !Array.from(selectElement.options).some(option => option.value === extension_settings.sd[setting])) { + extension_settings.sd[setting] = selectElement.selectedOptions[0].value; + } +} + /** * Generates an image based on the given trigger word. * @param {string} initiator The initiator of the image generation @@ -2292,8 +2268,8 @@ async function generatePicture(initiator, args, trigger, message, callback) { return; } - extension_settings.sd.sampler = $('#sd_sampler').find(':selected').val(); - extension_settings.sd.model = $('#sd_model').find(':selected').val(); + ensureSelectionExists('sampler', '#sd_sampler'); + ensureSelectionExists('model', '#sd_model'); trigger = trigger.trim(); const generationType = getGenerationType(trigger); @@ -2441,7 +2417,7 @@ async function getPrompt(generationType, message, trigger, quietPrompt, combineN } if (generationType !== generationMode.FREE) { - prompt = await refinePrompt(prompt, true); + prompt = await refinePrompt(prompt, false); } return prompt; @@ -3031,7 +3007,7 @@ async function generateNovelImage(prompt, negativePrompt, signal) { width: width, height: height, negative_prompt: negativePrompt, - upscale_ratio: extension_settings.sd.novel_upscale_ratio, + upscale_ratio: extension_settings.sd.hr_scale, decrisper: extension_settings.sd.novel_decrisper, sm: sm, sm_dyn: sm_dyn, @@ -3613,8 +3589,8 @@ async function sdMessageButton(e) { try { setBusyIcon(true); if (hasSavedImage) { - const prompt = await refinePrompt(message.extra.title, false, false); - const negative = hasSavedNegative ? await refinePrompt(message.extra.negative, false, true) : ''; + const prompt = await refinePrompt(message.extra.title, false); + const negative = hasSavedNegative ? await refinePrompt(message.extra.negative, true) : ''; message.extra.title = prompt; const generationType = message?.extra?.generationType ?? generationMode.FREE; @@ -3756,8 +3732,8 @@ async function onImageSwiped({ message, element, direction }) { eventSource.once(CUSTOM_STOP_EVENT, stopListener); const callback = () => { }; const hasNegative = message.extra.negative; - const prompt = await refinePrompt(message.extra.title, false, false); - const negativePromptPrefix = hasNegative ? await refinePrompt(message.extra.negative, false, true) : ''; + const prompt = await refinePrompt(message.extra.title, false); + const negativePromptPrefix = hasNegative ? await refinePrompt(message.extra.negative, true) : ''; const characterName = context.groupId ? context.groups[Object.keys(context.groups).filter(x => context.groups[x].id === context.groupId)[0]]?.id?.toString() : context.characters[context.characterId]?.name; @@ -3788,12 +3764,83 @@ async function onImageSwiped({ message, element, direction }) { await context.saveChat(); } +/** + * Applies the command arguments to the extension settings. + * @typedef {import('../../slash-commands/SlashCommand.js').NamedArguments} NamedArguments + * @typedef {import('../../slash-commands/SlashCommand.js').NamedArgumentsCapture} NamedArgumentsCapture + * @param {NamedArguments | NamedArgumentsCapture} args - Command arguments + * @returns {Record} - Current settings before applying the command arguments + */ +function applyCommandArguments(args) { + const overrideSettings = {}; + const currentSettings = {}; + const settingMap = { + 'edit': 'refine_mode', + 'extend': 'free_extend', + 'multimodal': 'multimodal_captioning', + 'seed': 'seed', + 'width': 'width', + 'height': 'height', + 'steps': 'steps', + 'cfg': 'scale', + 'skip': 'clip_skip', + 'model': 'model', + 'sampler': 'sampler', + 'scheduler': 'scheduler', + 'vae': 'vae', + 'upscaler': 'hr_upscaler', + 'scale': 'hr_scale', + 'hires': 'enable_hr', + 'denoise': 'denoising_strength', + '2ndpass': 'hr_second_pass_steps', + 'faces': 'restore_faces', + }; + + for (const [param, setting] of Object.entries(settingMap)) { + if (args[param] === undefined || defaultSettings[setting] === undefined) { + continue; + } + currentSettings[setting] = extension_settings.sd[setting]; + const value = String(args[param]); + const type = typeof defaultSettings[setting]; + switch (type) { + case 'boolean': + overrideSettings[setting] = isTrueBoolean(value) || !isFalseBoolean(value); + break; + case 'number': + overrideSettings[setting] = Number(value); + break; + default: + overrideSettings[setting] = value; + break; + } + } + + Object.assign(extension_settings.sd, overrideSettings); + return currentSettings; +} + jQuery(async () => { await addSDGenButtons(); SlashCommandParser.addCommandObject(SlashCommand.fromProps({ name: 'imagine', - callback: (args, trigger) => generatePicture(initiators.command, args, String(trigger)), + returns: 'URL of the generated image, or an empty string if the generation failed', + callback: async (args, trigger) => { + const currentSettings = applyCommandArguments(args); + + try { + return await generatePicture(initiators.command, args, String(trigger)); + } catch (error) { + console.error('Failed to generate image:', error); + return ''; + } finally { + if (Object.keys(currentSettings).length) { + Object.assign(extension_settings.sd, currentSettings); + saveSettingsDebounced(); + } + } + }, aliases: ['sd', 'img', 'image'], namedArgumentList: [ new SlashCommandNamedArgument( @@ -3803,6 +3850,164 @@ jQuery(async () => { name: 'negative', description: 'negative prompt prefix', typeList: [ARGUMENT_TYPE.STRING], + isRequired: false, + acceptsMultiple: false, + }), + SlashCommandNamedArgument.fromProps({ + name: 'extend', + description: 'auto-extend free mode prompts with the LLM', + typeList: [ARGUMENT_TYPE.BOOLEAN], + enumProvider: commonEnumProviders.boolean('trueFalse'), + isRequired: false, + acceptsMultiple: false, + }), + SlashCommandNamedArgument.fromProps({ + name: 'edit', + description: 'edit the prompt before generation', + typeList: [ARGUMENT_TYPE.BOOLEAN], + enumProvider: commonEnumProviders.boolean('trueFalse'), + isRequired: false, + acceptsMultiple: false, + }), + SlashCommandNamedArgument.fromProps({ + name: 'multimodal', + description: 'use multimodal captioning (for portraits only)', + typeList: [ARGUMENT_TYPE.BOOLEAN], + enumProvider: commonEnumProviders.boolean('trueFalse'), + isRequired: false, + acceptsMultiple: false, + }), + SlashCommandNamedArgument.fromProps({ + name: 'snap', + description: 'snap auto-adjusted dimensions to the nearest known resolution (portraits and backgrounds only)', + typeList: [ARGUMENT_TYPE.BOOLEAN], + enumProvider: commonEnumProviders.boolean('trueFalse'), + isRequired: false, + acceptsMultiple: false, + }), + SlashCommandNamedArgument.fromProps({ + name: 'seed', + description: 'random seed', + isRequired: false, + typeList: [ARGUMENT_TYPE.NUMBER], + acceptsMultiple: false, + }), + SlashCommandNamedArgument.fromProps({ + name: 'width', + description: 'image width', + isRequired: false, + typeList: [ARGUMENT_TYPE.NUMBER], + acceptsMultiple: false, + }), + SlashCommandNamedArgument.fromProps({ + name: 'height', + description: 'image height', + isRequired: false, + typeList: [ARGUMENT_TYPE.NUMBER], + acceptsMultiple: false, + }), + SlashCommandNamedArgument.fromProps({ + name: 'steps', + description: 'number of steps', + isRequired: false, + typeList: [ARGUMENT_TYPE.NUMBER], + acceptsMultiple: false, + }), + SlashCommandNamedArgument.fromProps({ + name: 'cfg', + description: 'CFG scale', + isRequired: false, + typeList: [ARGUMENT_TYPE.NUMBER], + acceptsMultiple: false, + }), + SlashCommandNamedArgument.fromProps({ + name: 'skip', + description: 'CLIP skip layers', + isRequired: false, + typeList: [ARGUMENT_TYPE.NUMBER], + acceptsMultiple: false, + }), + SlashCommandNamedArgument.fromProps({ + name: 'model', + description: 'model override', + isRequired: false, + typeList: [ARGUMENT_TYPE.STRING], + acceptsMultiple: false, + forceEnum: true, + enumProvider: () => Array.from(document.querySelectorAll('#sd_model > [value]')).map(o => new SlashCommandEnumValue(o.getAttribute('value'), o.textContent)), + }), + SlashCommandNamedArgument.fromProps({ + name: 'sampler', + description: 'sampler override', + isRequired: false, + typeList: [ARGUMENT_TYPE.STRING], + acceptsMultiple: false, + forceEnum: true, + enumProvider: () => Array.from(document.querySelectorAll('#sd_sampler > [value]')).map(o => new SlashCommandEnumValue(o.getAttribute('value'), o.textContent)), + }), + SlashCommandNamedArgument.fromProps({ + name: 'scheduler', + description: 'scheduler override', + isRequired: false, + typeList: [ARGUMENT_TYPE.STRING], + acceptsMultiple: false, + forceEnum: true, + enumProvider: () => Array.from(document.querySelectorAll('#sd_scheduler > [value]')).map(o => new SlashCommandEnumValue(o.getAttribute('value'), o.textContent)), + }), + SlashCommandNamedArgument.fromProps({ + name: 'vae', + description: 'VAE name override', + isRequired: false, + typeList: [ARGUMENT_TYPE.STRING], + acceptsMultiple: false, + forceEnum: true, + enumProvider: () => Array.from(document.querySelectorAll('#sd_vae > [value]')).map(o => new SlashCommandEnumValue(o.getAttribute('value'), o.textContent)), + }), + SlashCommandNamedArgument.fromProps({ + name: 'upscaler', + description: 'upscaler override', + isRequired: false, + typeList: [ARGUMENT_TYPE.STRING], + acceptsMultiple: false, + forceEnum: true, + enumProvider: () => Array.from(document.querySelectorAll('#sd_hr_upscaler > [value]')).map(o => new SlashCommandEnumValue(o.getAttribute('value'), o.textContent)), + }), + SlashCommandNamedArgument.fromProps({ + name: 'hires', + description: 'enable high-res fix', + isRequired: false, + typeList: [ARGUMENT_TYPE.BOOLEAN], + acceptsMultiple: false, + enumProvider: commonEnumProviders.boolean('trueFalse'), + }), + SlashCommandNamedArgument.fromProps({ + name: 'scale', + description: 'upscale amount', + isRequired: false, + typeList: [ARGUMENT_TYPE.NUMBER], + acceptsMultiple: false, + }), + SlashCommandNamedArgument.fromProps({ + name: 'denoise', + description: 'denoising strength', + isRequired: false, + typeList: [ARGUMENT_TYPE.NUMBER], + acceptsMultiple: false, + }), + SlashCommandNamedArgument.fromProps({ + name: '2ndpass', + description: 'second pass steps', + isRequired: false, + typeList: [ARGUMENT_TYPE.NUMBER], + acceptsMultiple: false, + }), + SlashCommandNamedArgument.fromProps({ + name: 'faces', + description: 'restore faces', + isRequired: false, + typeList: [ARGUMENT_TYPE.BOOLEAN], + acceptsMultiple: false, + enumProvider: commonEnumProviders.boolean('trueFalse'), }), ], unnamedArgumentList: [ @@ -3874,7 +4079,6 @@ jQuery(async () => { $('#sd_hr_scale').on('input', onHrScaleInput); $('#sd_denoising_strength').on('input', onDenoisingStrengthInput); $('#sd_hr_second_pass_steps').on('input', onHrSecondPassStepsInput); - $('#sd_novel_upscale_ratio').on('input', onNovelUpscaleRatioInput); $('#sd_novel_anlas_guard').on('input', onNovelAnlasGuardInput); $('#sd_novel_view_anlas').on('click', onViewAnlasClick); $('#sd_novel_sm').on('input', onNovelSmInput); @@ -3887,7 +4091,6 @@ jQuery(async () => { $('#sd_comfy_open_workflow_editor').on('click', onComfyOpenWorkflowEditorClick); $('#sd_comfy_new_workflow').on('click', onComfyNewWorkflowClick); $('#sd_comfy_delete_workflow').on('click', onComfyDeleteWorkflowClick); - $('#sd_expand').on('input', onExpandInput); $('#sd_style').on('change', onStyleSelect); $('#sd_save_style').on('click', onSaveStyleClick); $('#sd_delete_style').on('click', onDeleteStyleClick); diff --git a/public/scripts/extensions/stable-diffusion/settings.html b/public/scripts/extensions/stable-diffusion/settings.html index f43ed7246..2efa74d83 100644 --- a/public/scripts/extensions/stable-diffusion/settings.html +++ b/public/scripts/extensions/stable-diffusion/settings.html @@ -27,11 +27,6 @@ Extend free mode prompts (interactive/commands) -