Add a function tool for image generation

c853547b11a8d219fb9c3c04bd05737dc56678e7

Cohee <18619528+Cohee1207@users.noreply.github.com>

Parents: cc3cc58a06a7ecfd841e4ddad8f7190c61da5d98
2 files changed, +58 -0Ignore whitespace
public/scripts/extensions/stable-diffusion/index.js+54 -0
@@ -32,6 +32,7 @@ import { debounce_timeout } from '../../constants.js';
3232import { SlashCommandEnumValue } from '../../slash-commands/SlashCommandEnumValue.js';
3333import { POPUP_RESULT, POPUP_TYPE, Popup, callGenericPopup } from '../../popup.js';
3434import { commonEnumProviders } from '../../slash-commands/SlashCommandCommonEnumsProvider.js';
35+import { ToolManager } from '../../tool-calling.js';
3536export { MODULE_NAME };
3637
3738const MODULE_NAME = 'sd';
@@ -62,6 +63,7 @@ const initiators = {
6263 interactive: 'interactive',
6364 wand: 'wand',
6465 swipe: 'swipe',
66+ tool: 'tool',
6567};
6668
6769const generationMode = {
@@ -226,6 +228,7 @@ const defaultSettings = {
226228 multimodal_captioning: false,
227229 snap: false,
228230 free_extend: false,
231+ function_tool: false,
229232
230233 prompts: promptTemplates,
231234
@@ -291,6 +294,10 @@ const defaultSettings = {
291294const writePromptFieldsDebounced = debounce(writePromptFields, debounce_timeout.relaxed);
292295
293296function processTriggers(chat, _, abort) {
297+ if (extension_settings.sd.function_tool && ToolManager.isToolCallingSupported()) {
298+ return;
299+ }
300+
294301 if (!extension_settings.sd.interactive_mode) {
295302 return;
296303 }
@@ -447,6 +454,7 @@ async function loadSettings() {
447454 $('#sd_interactive_visible').prop('checked', extension_settings.sd.interactive_visible);
448455 $('#sd_stability_style_preset').val(extension_settings.sd.stability_style_preset);
449456 $('#sd_huggingface_model_id').val(extension_settings.sd.huggingface_model_id);
457+ $('#sd_function_tool').prop('checked', extension_settings.sd.function_tool);
450458
451459 for (const style of extension_settings.sd.styles) {
452460 const option = document.createElement('option');
@@ -461,6 +469,7 @@ async function loadSettings() {
461469
462470 toggleSourceControls();
463471 addPromptTemplates();
472+ registerFunctionTool();
464473
465474 await loadSettingOptions();
466475}
@@ -910,6 +919,12 @@ async function onSourceChange() {
910919 await loadSettingOptions();
911920}
912921
922+function onFunctionToolInput() {
923+ extension_settings.sd.function_tool = !!$(this).prop('checked');
924+ saveSettingsDebounced();
925+ registerFunctionTool();
926+}
927+
913928async function onOpenAiStyleSelect() {
914929 extension_settings.sd.openai_style = String($('#sd_openai_style').find(':selected').val());
915930 saveSettingsDebounced();
@@ -3822,6 +3837,44 @@ function applyCommandArguments(args) {
38223837 return currentSettings;
38233838}
38243839
3840+function registerFunctionTool() {
3841+ if (!extension_settings.sd.function_tool) {
3842+ return ToolManager.unregisterFunctionTool('GenerateImage');
3843+ }
3844+
3845+ ToolManager.registerFunctionTool({
3846+ name: 'GenerateImage',
3847+ displayName: 'Generate Image',
3848+ description: [
3849+ 'Generate an image from a given text prompt.',
3850+ 'Use when a user asks for an image, a selfie, to picture a scene, etc.',
3851+ ].join(' '),
3852+ parameters: Object.freeze({
3853+ $schema: 'http://json-schema.org/draft-04/schema#',
3854+ type: 'object',
3855+ properties: {
3856+ prompt: {
3857+ type: 'string',
3858+ description: [
3859+ 'The text prompt used to generate the image.',
3860+ 'Must represent an exhaustive description of the desired image that will allow an artist or a photographer to perfectly recreate it.',
3861+ ],
3862+ },
3863+ },
3864+ required: [
3865+ 'prompt',
3866+ ],
3867+ }),
3868+ action: async (args) => {
3869+ if (!isValidState()) throw new Error('Image generation is not configured.');
3870+ if (!args) throw new Error('Missing arguments');
3871+ if (!args.prompt) throw new Error('Missing prompt');
3872+ return generatePicture(initiators.tool, {}, args.prompt);
3873+ },
3874+ formatMessage: () => 'Generating an image...',
3875+ });
3876+}
3877+
38253878jQuery(async () => {
38263879 await addSDGenButtons();
38273880
@@ -4175,6 +4228,7 @@ jQuery(async () => {
41754228 $('#sd_stability_key').on('click', onStabilityKeyClick);
41764229 $('#sd_stability_style_preset').on('change', onStabilityStylePresetChange);
41774230 $('#sd_huggingface_model_id').on('input', onHFModelInput);
4231+ $('#sd_function_tool').on('input', onFunctionToolInput);
41784232
41794233 if (!CSS.supports('field-sizing', 'content')) {
41804234 $('.sd_settings .inline-drawer-toggle').on('click', function () {
public/scripts/extensions/stable-diffusion/settings.html+4 -0
@@ -18,6 +18,10 @@
1818 <input id="sd_interactive_mode" type="checkbox" />
1919 <span data-i18n="sd_interactive_mode_txt">Interactive mode</span>
2020 </label>
21+ <label for="sd_function_tool" class="checkbox_label" data-i18n="[title]sd_function_tool" title="Use the function tool to automatically detect intents to generate images.">
22+ <input id="sd_function_tool" type="checkbox" />
23+ <span data-i18n="sd_function_tool_txt">Enable function tool</span>
24+ </label>
2125 <label for="sd_multimodal_captioning" class="checkbox_label" data-i18n="[title]sd_multimodal_captioning" title="Use multimodal captioning to generate prompts for user and character portraits based on their avatars.">
2226 <input id="sd_multimodal_captioning" type="checkbox" />
2327 <span data-i18n="sd_multimodal_captioning_txt">Use multimodal captioning for portraits</span>