From b440a33196322d5105c3cd8631731b8ffaf97222 Mon Sep 17 00:00:00 2001 From: shihao <3127647737@qq.com> Date: Wed, 12 Aug 2026 16:21:33 +0800 Subject: [PATCH] feat: add custom image API extension --- i18n/zh-cn.json | 29 ++ index.js | 755 ++++++++++++++++++++++++++++++++++++++++++++++++ manifest.json | 13 + settings.html | 81 ++++++ 4 files changed, 878 insertions(+) create mode 100644 i18n/zh-cn.json create mode 100644 index.js create mode 100644 manifest.json create mode 100644 settings.html diff --git a/i18n/zh-cn.json b/i18n/zh-cn.json new file mode 100644 index 0000000..88ddbd9 --- /dev/null +++ b/i18n/zh-cn.json @@ -0,0 +1,29 @@ +{ + "Image Auto Generation": "自动生成图片", + "Insert In Current Message": "插入当前消息", + "Create New Message": "创建新消息", + "Image Insert Type": "图片插入类型", + "Disable": "禁用", + "Enable Custom Image API": "启用自定义图片 API", + "Custom Image API Base URL": "自定义图片 API Base URL", + "Custom Image Model ID": "自定义图片模型 ID", + "Custom Image API Key": "图片 API Key", + "Clear Custom Image API Key": "清除图片 API Key", + "Custom Image API Key Hint": "留空时使用酒馆的通用 API Key。插件 Key 只保存在当前浏览器，不会打包进插件。", + "Image Size": "图片尺寸", + "Response Format": "返回格式", + "Auto": "自动", + "URL": "URL", + "Base64 JSON": "Base64 JSON", + "enabled": "已启用", + "disabled": "已禁用", + "Enable Prompt Injection": "启用提示注入", + "Prompt Template": "提示模板", + "Regex": "正则表达式", + "Inline Replace Mode": "行内替换模式", + "Position": "注入位置", + "Depth": "深度", + "at Depth System": "深度系统消息", + "at Depth User": "深度用户消息", + "at Depth AI": "深度助手消息" +} diff --git a/index.js b/index.js new file mode 100644 index 0000000..95b6e96 --- /dev/null +++ b/index.js @@ -0,0 +1,755 @@ +// The main script for the extension +// The following are examples of some basic extension functionality + +//You'll likely need to import extension_settings, getContext, and loadExtensionSettings from extensions.js +import { extension_settings, getContext } from '../../../extensions.js'; +//You'll likely need to import some other functions from the main script +import { + saveSettingsDebounced, + eventSource, + event_types, + updateMessageBlock, + ensureMessageMediaIsArray, +} from '../../../../script.js'; +import { appendMediaToMessage } from '../../../../script.js'; +import { regexFromString } from '../../../utils.js'; +import { SlashCommandParser } from '../../../slash-commands/SlashCommandParser.js'; +import { findSecret, SECRET_KEYS } from '../../../secrets.js'; + +// 扩展名称和路径 +const extensionName = 'st-image-auto-generation'; +// /scripts/extensions/third-party +const extensionFolderPath = `/scripts/extensions/third-party/${extensionName}`; +const apiKeyStorageKey = `${extensionName}.apiKey`; + +// 插入类型常量 +const INSERT_TYPE = { + DISABLED: 'disabled', + INLINE: 'inline', + NEW_MESSAGE: 'new', + REPLACE: 'replace', +}; + +/** + * Escapes characters for safe inclusion inside HTML attribute values. + * @param {string} value + * @returns {string} + */ +function escapeHtmlAttribute(value) { + if (typeof value !== 'string') { + return ''; + } + + return value + .replace(/&/g, '&') + .replace(/"/g, '"') + .replace(/'/g, ''') + .replace(//g, '>'); +} + +// 默认设置 +const defaultSettings = { + insertType: INSERT_TYPE.DISABLED, + customApi: { + enabled: true, + baseUrl: 'https://freeapi.dgbmc.top/v1', + model: 'grok-imagine-image', + size: '1024x1024', + responseFormat: 'auto', + }, + promptInjection: { + enabled: true, + prompt: ` +You must insert a at end of the reply. Prompts are used for stable diffusion image generation, based on the plot and character to output appropriate prompts to generate captivating images. +`, + regex: '/]*\\sprompt="([^"]*)"[^>]*?>/g', + position: 'deep_system', // deep_system, deep_user, deep_assistant + depth: 0, // 0表示添加到末尾，>0表示从末尾往前数第几个位置 + }, +}; + +/** + * Reads the API key stored for this extension in the current browser. + * It is deliberately kept out of extension_settings so shared settings + * exports and plugin packages do not contain the key. + * @returns {string} + */ +function getStoredApiKey() { + try { + return window.localStorage.getItem(apiKeyStorageKey) || ''; + } catch { + return ''; + } +} + +/** + * Stores or clears the API key for this extension in the current browser. + * @param {string} value + */ +function setStoredApiKey(value) { + try { + const normalizedValue = String(value || '').trim(); + if (normalizedValue) { + window.localStorage.setItem(apiKeyStorageKey, normalizedValue); + } else { + window.localStorage.removeItem(apiKeyStorageKey); + } + } catch (error) { + console.warn(`[${extensionName}] 无法保存图片 API Key:`, error); + } +} + +// 从设置更新UI +function updateUI() { + const settings = extension_settings[extensionName]; + + if (!settings) { + return; + } + + // 根据insertType设置开关状态 + $('#auto_generation').toggleClass( + 'selected', + settings.insertType !== INSERT_TYPE.DISABLED, + ); + + // 只在表单元素存在时更新它们 + if ($('#image_generation_insert_type').length) { + $('#image_generation_insert_type').val( + settings.insertType, + ); + $('#custom_image_api_enabled').prop('checked', settings.customApi.enabled); + $('#custom_image_api_base_url').val(settings.customApi.baseUrl); + $('#custom_image_api_model').val(settings.customApi.model); + $('#custom_image_api_size').val(settings.customApi.size); + $('#custom_image_api_response_format').val(settings.customApi.responseFormat); + $('#custom_image_api_key').val(getStoredApiKey()); + $('#custom_image_api_settings').toggle(settings.customApi.enabled); + $('#prompt_injection_enabled').prop( + 'checked', + settings.promptInjection.enabled, + ); + $('#prompt_injection_text').val( + settings.promptInjection.prompt, + ); + $('#prompt_injection_regex').val( + settings.promptInjection.regex, + ); + $('#prompt_injection_position').val( + settings.promptInjection.position, + ); + $('#prompt_injection_depth').val( + settings.promptInjection.depth, + ); + } +} + +// 加载设置 +async function loadSettings() { + extension_settings[extensionName] = extension_settings[extensionName] || {}; + const settings = extension_settings[extensionName]; + + if (settings.insertType === undefined) { + settings.insertType = defaultSettings.insertType; + } + + settings.customApi = { + ...defaultSettings.customApi, + ...(settings.customApi || {}), + }; + settings.customApi.enabled = settings.customApi.enabled !== false; + settings.customApi.baseUrl = String( + settings.customApi.baseUrl || defaultSettings.customApi.baseUrl, + ).trim(); + settings.customApi.model = String( + settings.customApi.model || defaultSettings.customApi.model, + ).trim(); + settings.customApi.size = String( + settings.customApi.size || defaultSettings.customApi.size, + ).trim(); + if (!['auto', 'url', 'b64_json'].includes(settings.customApi.responseFormat)) { + settings.customApi.responseFormat = defaultSettings.customApi.responseFormat; + } + + settings.promptInjection = { + ...defaultSettings.promptInjection, + ...(settings.promptInjection || {}), + }; + + updateUI(); +} + +// 创建设置页面 +async function createSettings(settingsHtml) { + // 创建一个容器来存放设置，确保其正确显示在扩展设置面板中 + if (!$('#image_auto_generation_container').length) { + $('#extensions_settings2').append( + '
', + ); + } + + // 使用传入的settingsHtml而不是重新获取 + $('#image_auto_generation_container').empty().append(settingsHtml); + + // 添加设置变更事件处理 + $('#image_generation_insert_type').on('change', function () { + const newValue = $(this).val(); + extension_settings[extensionName].insertType = newValue; + updateUI(); + saveSettingsDebounced(); + }); + + $('#custom_image_api_enabled').on('change', function () { + extension_settings[extensionName].customApi.enabled = $(this).prop('checked'); + updateUI(); + saveSettingsDebounced(); + }); + + $('#custom_image_api_base_url').on('input', function () { + extension_settings[extensionName].customApi.baseUrl = String($(this).val()).trim(); + saveSettingsDebounced(); + }); + + $('#custom_image_api_model').on('input', function () { + extension_settings[extensionName].customApi.model = String($(this).val()).trim(); + saveSettingsDebounced(); + }); + + $('#custom_image_api_size').on('change', function () { + extension_settings[extensionName].customApi.size = String($(this).val()).trim(); + saveSettingsDebounced(); + }); + + $('#custom_image_api_response_format').on('change', function () { + extension_settings[extensionName].customApi.responseFormat = $(this).val(); + saveSettingsDebounced(); + }); + + $('#custom_image_api_key').on('input', function () { + setStoredApiKey($(this).val()); + }); + + $('#custom_image_api_key_clear').on('click', function () { + setStoredApiKey(''); + $('#custom_image_api_key').val(''); + toastr.success('图片 API Key 已清除'); + }); + + // 添加提示词注入设置的事件处理 + $('#prompt_injection_enabled').on('change', function () { + extension_settings[extensionName].promptInjection.enabled = + $(this).prop('checked'); + saveSettingsDebounced(); + }); + + $('#prompt_injection_text').on('input', function () { + extension_settings[extensionName].promptInjection.prompt = + $(this).val(); + saveSettingsDebounced(); + }); + + $('#prompt_injection_regex').on('input', function () { + extension_settings[extensionName].promptInjection.regex = $(this).val(); + saveSettingsDebounced(); + }); + + $('#prompt_injection_position').on('change', function () { + extension_settings[extensionName].promptInjection.position = + $(this).val(); + saveSettingsDebounced(); + }); + + // 深度设置事件处理 + $('#prompt_injection_depth').on('input', function () { + const value = parseInt(String($(this).val())); + extension_settings[extensionName].promptInjection.depth = isNaN(value) + ? 0 + : value; + saveSettingsDebounced(); + }); + + // 初始化设置值 + updateUI(); +} + +// 设置变更处理函数 +function onExtensionButtonClick() { + // 直接访问扩展设置面板 + const extensionsDrawer = $('#extensions-settings-button .drawer-toggle'); + + // 如果抽屉是关闭的，点击打开它 + if ($('#rm_extensions_block').hasClass('closedDrawer')) { + extensionsDrawer.trigger('click'); + } + + // 等待抽屉打开后滚动到我们的设置容器 + setTimeout(() => { + // 找到我们的设置容器 + const container = $('#image_auto_generation_container'); + if (container.length) { + // 滚动到设置面板位置 + $('#rm_extensions_block').animate( + { + scrollTop: + container.offset().top - + $('#rm_extensions_block').offset().top + + $('#rm_extensions_block').scrollTop(), + }, + 500, + ); + + // 使用SillyTavern原生的抽屉展开方式 + // 检查抽屉内容是否可见 + const drawerContent = container.find('.inline-drawer-content'); + const drawerHeader = container.find('.inline-drawer-header'); + + // 只有当内容被隐藏时才触发展开 + if (drawerContent.is(':hidden') && drawerHeader.length) { + // 直接使用原生点击事件触发，而不做任何内部处理 + drawerHeader.trigger('click'); + } + } + }, 500); +} + +// 初始化扩展 +$(function () { + (async function () { + // 获取设置HTML (只获取一次) + const settingsHtml = await $.get( + `${extensionFolderPath}/settings.html`, + ); + + // 添加扩展到菜单 + $('#extensionsMenu') + .append(`
+
+ Image Auto Generation +
`); + + // 修改点击事件，打开设置面板而不是切换状态 + $('#auto_generation').off('click').on('click', onExtensionButtonClick); + + await loadSettings(); + + // 创建设置 - 将获取的HTML传递给createSettings + await createSettings(settingsHtml); + + // 确保设置面板可见时，设置值是正确的 + $('#extensions-settings-button').on('click', function () { + setTimeout(() => { + updateUI(); + }, 200); + }); + })(); +}); +// 获取消息角色 +function getMesRole() { + // 确保对象路径存在 + if ( + !extension_settings[extensionName] || + !extension_settings[extensionName].promptInjection || + !extension_settings[extensionName].promptInjection.position + ) { + return 'system'; // 默认返回system角色 + } + + switch (extension_settings[extensionName].promptInjection.position) { + case 'deep_system': + return 'system'; + case 'deep_user': + return 'user'; + case 'deep_assistant': + return 'assistant'; + default: + return 'system'; + } +} + +/** + * Builds an OpenAI-compatible image generation endpoint from a base URL. + * A full /images/generations URL is also accepted for convenience. + * @param {string} baseUrl + * @returns {string} + */ +function getImageGenerationUrl(baseUrl) { + const normalizedBaseUrl = String(baseUrl || '').trim().replace(/\/+$/, ''); + if (!normalizedBaseUrl) { + throw new Error('自定义图片 API 地址不能为空'); + } + + return /\/images\/generations$/i.test(normalizedBaseUrl) + ? normalizedBaseUrl + : `${normalizedBaseUrl}/images/generations`; +} + +/** + * Reads an error body without leaking the API key into the error message. + * @param {Response} response + * @returns {Promise} + */ +async function getImageApiError(response) { + const text = await response.text(); + if (!text) { + return `${response.status} ${response.statusText}`.trim(); + } + + try { + const data = JSON.parse(text); + return ( + data?.error?.message || + (typeof data?.error === 'string' ? data.error : '') || + data?.message || + text + ); + } catch { + return text; + } +} + +/** + * Extracts a browser-loadable image URL from an OpenAI-compatible response. + * @param {any} data + * @returns {string} + */ +function extractImageUrl(data) { + const item = Array.isArray(data?.data) + ? data.data[0] + : Array.isArray(data?.images) + ? data.images[0] + : data?.data; + + if (typeof item === 'string' && item.trim()) { + return item.trim(); + } + + const imageUrl = item?.url || item?.image_url || data?.url; + if (typeof imageUrl === 'string' && imageUrl.trim()) { + return imageUrl.trim(); + } + + const base64 = item?.b64_json || item?.base64 || data?.b64_json; + if (typeof base64 === 'string' && base64.trim()) { + if (base64.startsWith('data:')) { + return base64; + } + + const format = String(item?.format || data?.format || 'png') + .replace(/[^a-z0-9.+-]/gi, '') + .toLowerCase() || 'png'; + return `data:image/${format};base64,${base64}`; + } + + return ''; +} + +/** + * Generates an image through a user-configured OpenAI-compatible endpoint. + * @param {string} prompt + * @returns {Promise} + */ +async function generateCustomImage(prompt) { + const customApi = extension_settings[extensionName].customApi; + const apiKey = getStoredApiKey() || await findSecret(SECRET_KEYS.GENERIC); + const headers = { 'Content-Type': 'application/json' }; + const model = String(customApi.model || '').trim(); + const size = String(customApi.size || '').trim(); + + if (!model) { + throw new Error('自定义图片模型 ID 不能为空'); + } + + if (!size) { + throw new Error('图片尺寸不能为空'); + } + + if (apiKey) { + const normalizedKey = String(apiKey).trim(); + headers.Authorization = /^bearer\s/i.test(normalizedKey) + ? normalizedKey + : `Bearer ${normalizedKey}`; + } + + const body = { + prompt: prompt.trim(), + model, + size, + n: 1, + }; + + if (customApi.responseFormat !== 'auto') { + body.response_format = customApi.responseFormat; + } + + const response = await fetch(getImageGenerationUrl(customApi.baseUrl), { + method: 'POST', + headers, + body: JSON.stringify(body), + }); + + if (!response.ok) { + throw new Error(await getImageApiError(response)); + } + + const imageUrl = extractImageUrl(await response.json()); + if (!imageUrl) { + throw new Error('图片 API 返回中没有找到 url 或 b64_json'); + } + + return imageUrl; +} + +/** + * Generates an image using either the custom endpoint or SillyTavern's native + * stable-diffusion command. + * @param {string} prompt + * @param {string} insertType + * @returns {Promise} + */ +async function generateImage(prompt, insertType) { + const settings = extension_settings[extensionName]; + if (settings.customApi.enabled) { + return await generateCustomImage(prompt); + } + + const command = SlashCommandParser.commands['sd']; + if (!command?.callback) { + throw new Error('酒馆图片生成命令尚未加载，请刷新页面后重试'); + } + + return await command.callback( + { + quiet: insertType === INSERT_TYPE.NEW_MESSAGE ? 'false' : 'true', + }, + prompt, + ); +} + +/** + * Adds a custom-API result as a new assistant message. + * @param {any} context + * @param {string} prompt + * @param {string} imageUrl + * @returns {Promise} + */ +async function addImageMessage(context, prompt, imageUrl) { + const message = { + name: context.name2 || 'Image Generation', + is_user: false, + is_system: false, + send_date: new Date().toISOString(), + mes: '', + extra: { + media: [ + { + type: 'image', + url: imageUrl, + title: prompt, + source: 'generated', + }, + ], + media_display: 'gallery', + media_index: 0, + inline_image: false, + }, + }; + + context.chat.push(message); + const messageId = context.chat.length - 1; + await eventSource.emit(event_types.MESSAGE_RECEIVED, messageId, 'extension'); + context.addOneMessage(message); + await eventSource.emit(event_types.CHARACTER_MESSAGE_RENDERED, messageId, 'extension'); + await context.saveChat(); +} + +// 监听CHAT_COMPLETION_PROMPT_READY事件以注入提示词 +eventSource.on( + event_types.CHAT_COMPLETION_PROMPT_READY, + async function (eventData) { + try { + // 确保设置对象和promptInjection对象都存在 + if ( + !extension_settings[extensionName] || + !extension_settings[extensionName].promptInjection || + !extension_settings[extensionName].promptInjection.enabled || + extension_settings[extensionName].insertType === + INSERT_TYPE.DISABLED + ) { + return; + } + + const prompt = + extension_settings[extensionName].promptInjection.prompt; + const depth = + extension_settings[extensionName].promptInjection.depth || 0; + const role = getMesRole(); + + console.log( + `[${extensionName}] 准备注入提示词: 角色=${role}, 深度=${depth}`, + ); + console.log( + `[${extensionName}] 提示词内容: ${prompt.substring(0, 50)}...`, + ); + + // 根据depth参数决定插入位置 + if (depth === 0) { + // 添加到末尾 + eventData.chat.push({ role: role, content: prompt }); + console.log(`[${extensionName}] 提示词已添加到聊天末尾`); + } else { + // 从末尾向前插入 + eventData.chat.splice(-depth, 0, { + role: role, + content: prompt, + }); + console.log( + `[${extensionName}] 提示词已插入到聊天中，从末尾往前第 ${depth} 个位置`, + ); + } + } catch (error) { + console.error(`[${extensionName}] 提示词注入错误:`, error); + toastr.error(`提示词注入错误: ${error}`); + } + }, +); + +// 监听消息接收事件 +eventSource.on(event_types.MESSAGE_RECEIVED, handleIncomingMessage); +async function handleIncomingMessage() { + // 确保设置对象存在 + if ( + !extension_settings[extensionName] || + extension_settings[extensionName].insertType === INSERT_TYPE.DISABLED + ) { + return; + } + + const context = getContext(); + const message = context.chat[context.chat.length - 1]; + + // 检查是否是AI消息 + if (!message || message.is_user) { + return; + } + + // 确保promptInjection对象和regex属性存在 + if ( + !extension_settings[extensionName].promptInjection || + !extension_settings[extensionName].promptInjection.regex || + typeof message.mes !== 'string' + ) { + console.error('Prompt injection settings not properly initialized'); + return; + } + + // 使用正则表达式search + const imgTagRegex = regexFromString( + extension_settings[extensionName].promptInjection.regex, + ); + // const testRegex = regexFromString(extension_settings[extensionName].promptInjection.regex); + let matches; + if (imgTagRegex.global) { + matches = [...message.mes.matchAll(imgTagRegex)]; + } else { + const singleMatch = message.mes.match(imgTagRegex); + matches = singleMatch ? [singleMatch] : []; + } + console.log(imgTagRegex, matches); + if (matches.length > 0) { + // 延迟执行图片生成，确保消息首先显示出来 + setTimeout(async () => { + try { + toastr.info(`Generating ${matches.length} images...`); + const insertType = extension_settings[extensionName].insertType; + + const messageId = context.chat.indexOf(message); + if (messageId < 0 || context.chat[messageId] !== message) { + throw new Error('当前聊天已发生变化，已取消图片插入'); + } + + // 在当前消息中插入图片 + if (!message.extra) { + message.extra = {}; + } + ensureMessageMediaIsArray(message); + if (!Array.isArray(message.extra.media)) { + message.extra.media = []; + } + + // 获取消息元素用于稍后更新 + const messageElement = $(`.mes[mesid="${messageId}"]`); + + // 处理每个匹配的图片标签 + for (const match of matches) { + const prompt = + typeof match?.[1] === 'string' ? match[1] : ''; + if (!prompt.trim()) { + continue; + } + + // @ts-ignore + const imageUrl = await generateImage(prompt, insertType); + if (typeof imageUrl !== 'string' || !imageUrl.trim()) { + throw new Error('图片生成接口返回了空结果'); + } + + const customApiEnabled = extension_settings[extensionName].customApi.enabled; + + if (insertType === INSERT_TYPE.NEW_MESSAGE) { + // Native /sd already creates the new message when quiet=false. + if (customApiEnabled) { + await addImageMessage(context, prompt, imageUrl); + } + continue; + } + + // 统一插入到extra里 + if (insertType === INSERT_TYPE.INLINE) { + message.extra.media.push({ + type: 'image', + url: imageUrl, + title: prompt, + source: 'generated', + }); + message.extra.media_display = 'gallery'; + message.extra.media_index = message.extra.media.length - 1; + message.extra.inline_image = true; + + // 更新UI + appendMediaToMessage(message, messageElement); + + // 保存聊天记录 + await context.saveChat(); + } else if (insertType === INSERT_TYPE.REPLACE) { + // Find the original image tag in the message + const originalTag = + typeof match?.[0] === 'string' ? match[0] : ''; + if (!originalTag) { + continue; + } + + // Replace it with an actual image tag + const escapedUrl = escapeHtmlAttribute(imageUrl); + const escapedPrompt = escapeHtmlAttribute(prompt); + const newImageTag = `${escapedPrompt}`; + message.mes = message.mes.replace(originalTag, newImageTag); + + // Update the message display using updateMessageBlock + updateMessageBlock(messageId, message); + await eventSource.emit(event_types.MESSAGE_UPDATED, messageId); + + // Save the chat + await context.saveChat(); + } + } + toastr.success( + `${matches.length} images generated successfully`, + ); + } catch (error) { + toastr.error(`Image generation error: ${error}`); + console.error('Image generation error:', error); + } + }, 0); //防阻塞UI渲染 + } +} diff --git a/manifest.json b/manifest.json new file mode 100644 index 0000000..bad5aee --- /dev/null +++ b/manifest.json @@ -0,0 +1,13 @@ +{ + "display_name": "Image Auto Generation - Custom API", + "requires": [], + "optional": [], + "js": "index.js", + "author": "wickedcode01", + "version": "1.1.0", + "homePage": "https://github.com/wickedcode01/st-image-auto-generation", + "auto_update": false, + "i18n": { + "zh-cn": "i18n/zh-cn.json" + } +} diff --git a/settings.html b/settings.html new file mode 100644 index 0000000..66d7fa4 --- /dev/null +++ b/settings.html @@ -0,0 +1,81 @@ +
+
+ Image Auto Generation +
+
+
+
+ + +
+
+ +
+
+
+ + +
+
+ + +
+
+ +
+ + +
+ 留空时使用酒馆的通用 API Key。插件 Key 只保存在当前浏览器，不会打包进插件。 +
+
+ + +
+
+ + +
+
+
+ +
+
+
+ +
+
+
+ +
+
+ + +
+
+ + +
+
+