rixty/SillyTavern
1
0
Fork 0
mirror of https://github.com/SillyTavern/SillyTavern.git synced 2025-06-05 21:59:27 +02:00
Code Issues Packages Projects Releases Wiki Activity

Add automatic image captioning mode

Browse Source
This commit is contained in:
Cohee
2024-06-30 00:06:17 +03:00
parent 2670709237
commit 7149f46c9a
4 changed files with 85 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

79
public/scripts/extensions/caption/index.js
View File

@ -1,6 +1,6 @@
import { ensureImageFormatSupported, getBase64Async, isTrueBoolean, saveBase64AsFile } from '../../utils.js';
import { getContext, getApiUrl, doExtrasFetch, extension_settings, modules, renderExtensionTemplateAsync } from '../../extensions.js';
import { callPopup, getRequestHeaders, saveSettingsDebounced, substituteParamsExtended } from '../../../script.js';
import { callPopup, eventSource, event_types, getRequestHeaders, saveSettingsDebounced, substituteParamsExtended } from '../../../script.js';
import { getMessageTimeStamp } from '../../RossAscends-mods.js';
import { SECRET_KEYS, secret_state } from '../../secrets.js';
import { getMultimodalCaption } from '../shared.js';
@ -84,12 +84,11 @@ async function setSpinnerIcon() {
}
/**
* Sends a captioned message to the chat.
* @param {string} caption Caption text
* @param {string} image Image URL
* Wraps a caption with a message template.
* @param {string} caption Raw caption
* @returns {Promise<string>} Wrapped caption
*/
async function sendCaptionedMessage(caption, image) {
const context = getContext();
async function wrapCaptionTemplate(caption) {
let template = extension_settings.caption.template || TEMPLATE_DEFAULT;
if (!/{{caption}}/i.test(template)) {
@ -101,7 +100,7 @@ async function sendCaptionedMessage(caption, image) {
if (extension_settings.caption.refine_mode) {
messageText = await callPopup(
'<h3>Review and edit the generated message:</h3>Press "Cancel" to abort the caption sending.',
'<h3>Review and edit the generated caption:</h3>Press "Cancel" to abort the caption sending.',
'input',
messageText,
{ rows: 5, okButton: 'Send' });
@ -111,6 +110,55 @@ async function sendCaptionedMessage(caption, image) {
}
}
return messageText;
}
/**
* Appends caption to an existing message.
* @param {Object} data Message data
* @returns {Promise<void>}
*/
async function captionExistingMessage(data) {
if (!(data?.extra?.image)) {
return;
}
const imageData = await fetch(data.extra.image);
const blob = await imageData.blob();
const type = imageData.headers.get('Content-Type');
const file = new File([blob], 'image.png', { type });
const caption = await getCaptionForFile(file, null, true);
if (!caption) {
console.warn('Failed to generate a caption for the image.');
return;
}
const wrappedCaption = await wrapCaptionTemplate(caption);
const messageText = String(data.mes).trim();
if (!messageText) {
data.extra.inline_image = false;
data.mes = wrappedCaption;
data.extra.title = wrappedCaption;
}
else {
data.extra.inline_image = true;
data.extra.append_title = true;
data.extra.title = wrappedCaption;
}
}
/**
* Sends a captioned message to the chat.
* @param {string} caption Caption text
* @param {string} image Image URL
*/
async function sendCaptionedMessage(caption, image) {
const messageText = await wrapCaptionTemplate(caption);
const context = getContext();
const message = {
name: context.name1,
is_user: true,
@ -423,6 +471,7 @@ jQuery(async function () {
$('#caption_refine_mode').prop('checked', !!(extension_settings.caption.refine_mode));
$('#caption_allow_reverse_proxy').prop('checked', !!(extension_settings.caption.allow_reverse_proxy));
$('#caption_prompt_ask').prop('checked', !!(extension_settings.caption.prompt_ask));
$('#caption_auto_mode').prop('checked', !!(extension_settings.caption.auto_mode));
$('#caption_source').val(extension_settings.caption.source);
$('#caption_prompt').val(extension_settings.caption.prompt);
$('#caption_template').val(extension_settings.caption.template);
@ -448,6 +497,22 @@ jQuery(async function () {
extension_settings.caption.prompt_ask = $('#caption_prompt_ask').prop('checked');
saveSettingsDebounced();
});
$('#caption_auto_mode').on('input', () => {
extension_settings.caption.auto_mode = !!$('#caption_auto_mode').prop('checked');
saveSettingsDebounced();
});
const onMessageEvent = async (index) => {
if (!extension_settings.caption.auto_mode) {
return;
}
const data = getContext().chat[index];
await captionExistingMessage(data);
};
eventSource.on(event_types.MESSAGE_SENT, onMessageEvent);
eventSource.on(event_types.MESSAGE_FILE_EMBEDDED, onMessageEvent);
SlashCommandParser.addCommandObject(SlashCommand.fromProps({ name: 'caption',
callback: captionCommandCallback,