Extend /sd command

This commit is contained in:
Cohee
2024-09-30 21:51:24 +03:00
parent dedb96ec8d
commit 9af4d62cdf
3 changed files with 261 additions and 140 deletions

View File

@ -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<string>} 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 ? '<h3>Review and edit the <i>negative</i> prompt:</h3>' : '<h3>Review and edit the prompt:</h3>';
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<string, any>} - 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);