| 1 | import { ensureImageFormatSupported, getBase64Async, getFileExtension, isTrueBoolean, saveBase64AsFile } from '../../utils.js'; |
| 2 | import { getContext, getApiUrl, doExtrasFetch, extension_settings, modules, renderExtensionTemplateAsync } from '../../extensions.js'; |
| 3 | import { appendMediaToMessage, chat_metadata, eventSource, event_types, getRequestHeaders, saveChatConditional, saveSettingsDebounced, substituteParams } from '../../../script.js'; |
| 4 | import { getMessageTimeStamp } from '../../RossAscends-mods.js'; |
| 5 | import { SECRET_KEYS, secret_state } from '../../secrets.js'; |
| 6 | import { oai_settings } from '../../openai.js'; |
| 7 | import { getMultimodalCaption } from '../shared.js'; |
| 8 | import { textgen_types, textgenerationwebui_settings } from '../../textgen-settings.js'; |
| 9 | import { SlashCommandParser } from '../../slash-commands/SlashCommandParser.js'; |
| 10 | import { SlashCommand } from '../../slash-commands/SlashCommand.js'; |
| 11 | import { ARGUMENT_TYPE, SlashCommandArgument, SlashCommandNamedArgument } from '../../slash-commands/SlashCommandArgument.js'; |
| 12 | import { commonEnumProviders } from '../../slash-commands/SlashCommandCommonEnumsProvider.js'; |
| 13 | import { callGenericPopup, Popup, POPUP_TYPE } from '../../popup.js'; |
| 14 | import { debounce_timeout, MEDIA_DISPLAY, MEDIA_SOURCE, MEDIA_TYPE, SCROLL_BEHAVIOR } from '../../constants.js'; |
| 15 | export { MODULE_NAME }; |
| 16 | |
| 17 | const MODULE_NAME = 'caption'; |
| 18 | |
| 19 | const PROMPT_DEFAULT = 'What\'s in this image?'; |
| 20 | const TEMPLATE_DEFAULT = '[{{user}} sends {{char}} a picture that contains: {{caption}}]'; |
| 21 | |
| 22 | /** |
| 23 | * Migrates old extension settings to the new format. |
| 24 | * Must keep this function for compatibility with old settings. |
| 25 | */ |
| 26 | function migrateSettings() { |
| 27 | if (extension_settings.caption.local !== undefined) { |
| 28 | extension_settings.caption.source = extension_settings.caption.local ? 'local' : 'extras'; |
| 29 | } |
| 30 | |
| 31 | delete extension_settings.caption.local; |
| 32 | |
| 33 | if (!extension_settings.caption.source) { |
| 34 | extension_settings.caption.source = 'extras'; |
| 35 | } |
| 36 | |
| 37 | if (extension_settings.caption.source === 'openai') { |
| 38 | extension_settings.caption.source = 'multimodal'; |
| 39 | extension_settings.caption.multimodal_api = 'openai'; |
| 40 | extension_settings.caption.multimodal_model = 'gpt-4-turbo'; |
| 41 | } |
| 42 | |
| 43 | if (!extension_settings.caption.multimodal_api) { |
| 44 | extension_settings.caption.multimodal_api = 'openai'; |
| 45 | } |
| 46 | |
| 47 | if (!extension_settings.caption.multimodal_model) { |
| 48 | extension_settings.caption.multimodal_model = 'gpt-4-turbo'; |
| 49 | } |
| 50 | |
| 51 | if (!extension_settings.caption.prompt) { |
| 52 | extension_settings.caption.prompt = PROMPT_DEFAULT; |
| 53 | } |
| 54 | |
| 55 | if (!extension_settings.caption.template) { |
| 56 | extension_settings.caption.template = TEMPLATE_DEFAULT; |
| 57 | } |
| 58 | |
| 59 | if (!extension_settings.caption.show_in_chat) { |
| 60 | extension_settings.caption.show_in_chat = false; |
| 61 | } |
| 62 | } |
| 63 | |
| 64 | /** |
| 65 | * Sets an image icon for the send button. |
| 66 | */ |
| 67 | async function setImageIcon() { |
| 68 | try { |
| 69 | const sendButton = $('#send_picture .extensionsMenuExtensionButton'); |
| 70 | sendButton.addClass('fa-image'); |
| 71 | sendButton.removeClass('fa-hourglass-half'); |
| 72 | } catch (error) { |
| 73 | console.log(error); |
| 74 | } |
| 75 | } |
| 76 | |
| 77 | /** |
| 78 | * Sets a spinner icon for the send button. |
| 79 | */ |
| 80 | async function setSpinnerIcon() { |
| 81 | try { |
| 82 | const sendButton = $('#send_picture .extensionsMenuExtensionButton'); |
| 83 | sendButton.removeClass('fa-image'); |
| 84 | sendButton.addClass('fa-hourglass-half'); |
| 85 | } catch (error) { |
| 86 | console.log(error); |
| 87 | } |
| 88 | } |
| 89 | |
| 90 | /** |
| 91 | * Wraps a caption with a message template. |
| 92 | * @param {string} caption Raw caption |
| 93 | * @returns {Promise<string>} Wrapped caption |
| 94 | */ |
| 95 | async function wrapCaptionTemplate(caption) { |
| 96 | let template = extension_settings.caption.template || TEMPLATE_DEFAULT; |
| 97 | |
| 98 | if (!/{{caption}}/i.test(template)) { |
| 99 | console.warn('Poka-yoke: Caption template does not contain {{caption}}. Appending it.'); |
| 100 | template += ' {{caption}}'; |
| 101 | } |
| 102 | |
| 103 | let messageText = substituteParams(template, { dynamicMacros: { caption: caption } }); |
| 104 | |
| 105 | if (extension_settings.caption.refine_mode) { |
| 106 | messageText = await Popup.show.input( |
| 107 | 'Review and edit the generated caption:', |
| 108 | 'Press "Cancel" to abort the caption sending.', |
| 109 | messageText, |
| 110 | { rows: 8, okButton: 'Send' }); |
| 111 | |
| 112 | if (!messageText) { |
| 113 | throw new Error('User aborted the caption sending.'); |
| 114 | } |
| 115 | } |
| 116 | |
| 117 | return messageText; |
| 118 | } |
| 119 | |
| 120 | /** |
| 121 | * Appends caption to an existing message. |
| 122 | * @param {ChatMessage} message Message data |
| 123 | * @param {number} mediaIndex Index of the image to caption |
| 124 | * @returns {Promise<void>} |
| 125 | */ |
| 126 | async function captionExistingMessage(message, mediaIndex) { |
| 127 | if (!Array.isArray(message?.extra?.media) || message.extra.media.length === 0) { |
| 128 | return; |
| 129 | } |
| 130 | |
| 131 | if (mediaIndex === undefined || isNaN(mediaIndex) || mediaIndex < 0 || mediaIndex >= message.extra.media.length) { |
| 132 | mediaIndex = 0; |
| 133 | } |
| 134 | |
| 135 | const mediaAttachment = message.extra.media[mediaIndex]; |
| 136 | |
| 137 | if (!mediaAttachment || !mediaAttachment.url || mediaAttachment.type === MEDIA_TYPE.AUDIO) { |
| 138 | return; |
| 139 | } |
| 140 | |
| 141 | if (mediaAttachment.type === MEDIA_TYPE.VIDEO && !isVideoCaptioningAvailable()) { |
| 142 | throw new Error('Captioning videos is not supported for the current source.'); |
| 143 | } |
| 144 | |
| 145 | const imageData = await fetch(mediaAttachment.url); |
| 146 | const blob = await imageData.blob(); |
| 147 | const fileName = mediaAttachment.url.split('/').pop().split('?')[0] || 'image.jpg'; |
| 148 | const file = new File([blob], fileName, { type: blob.type }); |
| 149 | const caption = await getCaptionForFile(file, null, true); |
| 150 | |
| 151 | if (!caption) { |
| 152 | console.warn('Failed to generate a caption for the image.'); |
| 153 | return; |
| 154 | } |
| 155 | |
| 156 | const wrappedCaption = await wrapCaptionTemplate(caption); |
| 157 | |
| 158 | const messageText = String(message.mes).trim(); |
| 159 | |
| 160 | if (!messageText) { |
| 161 | message.extra.inline_image = false; |
| 162 | message.mes = wrappedCaption; |
| 163 | mediaAttachment.title = wrappedCaption; |
| 164 | mediaAttachment.captioned = true; |
| 165 | } else { |
| 166 | message.extra.inline_image = true; |
| 167 | mediaAttachment.append_title = true; |
| 168 | mediaAttachment.title = wrappedCaption; |
| 169 | mediaAttachment.captioned = true; |
| 170 | } |
| 171 | } |
| 172 | |
| 173 | /** |
| 174 | * Sends a captioned message to the chat. |
| 175 | * @param {string} caption Caption text |
| 176 | * @param {string} image Image URL |
| 177 | * @param {string} mimeType Image MIME type |
| 178 | * @returns {Promise<void>} |
| 179 | */ |
| 180 | async function sendCaptionedMessage(caption, image, mimeType) { |
| 181 | const messageText = await wrapCaptionTemplate(caption); |
| 182 | |
| 183 | const context = getContext(); |
| 184 | |
| 185 | /** @type {MediaAttachment} */ |
| 186 | const mediaAttachment = { |
| 187 | url: image, |
| 188 | type: MEDIA_TYPE.getFromMime(mimeType) || MEDIA_TYPE.IMAGE, |
| 189 | title: messageText, |
| 190 | captioned: true, |
| 191 | source: MEDIA_SOURCE.CAPTIONED, |
| 192 | }; |
| 193 | /** @type {ChatMessage} */ |
| 194 | const message = { |
| 195 | name: context.name1, |
| 196 | is_user: true, |
| 197 | send_date: getMessageTimeStamp(), |
| 198 | mes: messageText, |
| 199 | extra: { |
| 200 | media: [mediaAttachment], |
| 201 | media_display: MEDIA_DISPLAY.GALLERY, |
| 202 | media_index: 0, |
| 203 | inline_image: !!extension_settings.caption.show_in_chat, |
| 204 | }, |
| 205 | }; |
| 206 | chat_metadata.tainted = true; |
| 207 | context.chat.push(message); |
| 208 | const messageId = context.chat.length - 1; |
| 209 | await eventSource.emit(event_types.MESSAGE_SENT, messageId); |
| 210 | context.addOneMessage(message); |
| 211 | await eventSource.emit(event_types.USER_MESSAGE_RENDERED, messageId); |
| 212 | await context.saveChat(); |
| 213 | setTimeout(() => context.scrollOnMediaLoad(), debounce_timeout.short); |
| 214 | } |
| 215 | |
| 216 | /** |
| 217 | * Generates a caption for an image using a selected source. |
| 218 | * @param {string} base64Img Base64 encoded image without the data:image/...;base64, prefix |
| 219 | * @param {string} fileData Base64 encoded image with the data:image/...;base64, prefix |
| 220 | * @param {string} externalPrompt Caption prompt |
| 221 | * @returns {Promise<{caption: string}>} Generated caption |
| 222 | */ |
| 223 | async function doCaptionRequest(base64Img, fileData, externalPrompt) { |
| 224 | switch (extension_settings.caption.source) { |
| 225 | case 'local': |
| 226 | return await captionLocal(base64Img); |
| 227 | case 'extras': |
| 228 | return await captionExtras(base64Img); |
| 229 | case 'horde': |
| 230 | return await captionHorde(base64Img); |
| 231 | case 'multimodal': |
| 232 | return await captionMultimodal(fileData, externalPrompt); |
| 233 | default: |
| 234 | throw new Error('Unknown caption source.'); |
| 235 | } |
| 236 | } |
| 237 | |
| 238 | /** |
| 239 | * Generates a caption for an image using Extras API. |
| 240 | * @param {string} base64Img Base64 encoded image without the data:image/...;base64, prefix |
| 241 | * @returns {Promise<{caption: string}>} Generated caption |
| 242 | */ |
| 243 | async function captionExtras(base64Img) { |
| 244 | if (!modules.includes('caption')) { |
| 245 | throw new Error('No captioning module is available.'); |
| 246 | } |
| 247 | |
| 248 | const url = new URL(getApiUrl()); |
| 249 | url.pathname = '/api/caption'; |
| 250 | |
| 251 | const apiResult = await doExtrasFetch(url, { |
| 252 | method: 'POST', |
| 253 | headers: { |
| 254 | 'Content-Type': 'application/json', |
| 255 | 'Bypass-Tunnel-Reminder': 'bypass', |
| 256 | }, |
| 257 | body: JSON.stringify({ image: base64Img }), |
| 258 | }); |
| 259 | |
| 260 | if (!apiResult.ok) { |
| 261 | throw new Error('Failed to caption image via Extras.'); |
| 262 | } |
| 263 | |
| 264 | const data = await apiResult.json(); |
| 265 | return data; |
| 266 | } |
| 267 | |
| 268 | /** |
| 269 | * Generates a caption for an image using a local model. |
| 270 | * @param {string} base64Img Base64 encoded image without the data:image/...;base64, prefix |
| 271 | * @returns {Promise<{caption: string}>} Generated caption |
| 272 | */ |
| 273 | async function captionLocal(base64Img) { |
| 274 | const apiResult = await fetch('/api/extra/caption', { |
| 275 | method: 'POST', |
| 276 | headers: getRequestHeaders(), |
| 277 | body: JSON.stringify({ image: base64Img }), |
| 278 | }); |
| 279 | |
| 280 | if (!apiResult.ok) { |
| 281 | throw new Error('Failed to caption image via local pipeline.'); |
| 282 | } |
| 283 | |
| 284 | const data = await apiResult.json(); |
| 285 | return data; |
| 286 | } |
| 287 | |
| 288 | /** |
| 289 | * Generates a caption for an image using a Horde model. |
| 290 | * @param {string} base64Img Base64 encoded image without the data:image/...;base64, prefix |
| 291 | * @returns {Promise<{caption: string}>} Generated caption |
| 292 | */ |
| 293 | async function captionHorde(base64Img) { |
| 294 | const apiResult = await fetch('/api/horde/caption-image', { |
| 295 | method: 'POST', |
| 296 | headers: getRequestHeaders(), |
| 297 | body: JSON.stringify({ image: base64Img }), |
| 298 | }); |
| 299 | |
| 300 | if (!apiResult.ok) { |
| 301 | throw new Error('Failed to caption image via Horde.'); |
| 302 | } |
| 303 | |
| 304 | const data = await apiResult.json(); |
| 305 | return data; |
| 306 | } |
| 307 | |
| 308 | /** |
| 309 | * Generates a caption for an image using a multimodal model. |
| 310 | * @param {string} base64Img Base64 encoded image with the data:image/...;base64, prefix |
| 311 | * @param {string} externalPrompt Caption prompt |
| 312 | * @returns {Promise<{caption: string}>} Generated caption |
| 313 | */ |
| 314 | async function captionMultimodal(base64Img, externalPrompt) { |
| 315 | let prompt = externalPrompt || extension_settings.caption.prompt || PROMPT_DEFAULT; |
| 316 | |
| 317 | if (!externalPrompt && extension_settings.caption.prompt_ask) { |
| 318 | const customPrompt = await callGenericPopup('Enter a comment or question:', POPUP_TYPE.INPUT, prompt, { rows: 4 }); |
| 319 | if (!customPrompt) { |
| 320 | throw new Error('User aborted the caption sending.'); |
| 321 | } |
| 322 | prompt = String(customPrompt).trim(); |
| 323 | } |
| 324 | |
| 325 | prompt = substituteParams(prompt); |
| 326 | |
| 327 | const caption = await getMultimodalCaption(base64Img, prompt); |
| 328 | return { caption }; |
| 329 | } |
| 330 | |
| 331 | /** |
| 332 | * Handles the image selection event. |
| 333 | * @param {Event} e Input event |
| 334 | * @param {string} prompt Caption prompt |
| 335 | * @param {boolean} quiet Suppresses sending a message |
| 336 | * @returns {Promise<string>} Generated caption |
| 337 | */ |
| 338 | async function onSelectImage(e, prompt, quiet) { |
| 339 | if (!(e.target instanceof HTMLInputElement)) { |
| 340 | return ''; |
| 341 | } |
| 342 | |
| 343 | const file = e.target.files[0]; |
| 344 | const form = e.target.form; |
| 345 | |
| 346 | if (!file || !(file instanceof File)) { |
| 347 | form && form.reset(); |
| 348 | return ''; |
| 349 | } |
| 350 | |
| 351 | const caption = await getCaptionForFile(file, prompt, quiet); |
| 352 | form && form.reset(); |
| 353 | return caption; |
| 354 | } |
| 355 | |
| 356 | /** |
| 357 | * Gets a caption for an image file. |
| 358 | * @param {File} file Input file |
| 359 | * @param {string} prompt Caption prompt |
| 360 | * @param {boolean} quiet Suppresses sending a message |
| 361 | * @returns {Promise<string>} Generated caption |
| 362 | */ |
| 363 | async function getCaptionForFile(file, prompt, quiet) { |
| 364 | try { |
| 365 | if (file.type.startsWith('video/') && !isVideoCaptioningAvailable()) { |
| 366 | throw new Error('Video captioning is not available for the current source.'); |
| 367 | } |
| 368 | |
| 369 | setSpinnerIcon(); |
| 370 | const context = getContext(); |
| 371 | const fileData = await getBase64Async(await ensureImageFormatSupported(file)); |
| 372 | const extension = getFileExtension(file); |
| 373 | const base64Data = fileData.split(',')[1]; |
| 374 | const { caption } = await doCaptionRequest(base64Data, fileData, prompt); |
| 375 | if (!quiet) { |
| 376 | const imagePath = await saveBase64AsFile(base64Data, context.name2, '', extension); |
| 377 | await sendCaptionedMessage(caption, imagePath, file.type); |
| 378 | } |
| 379 | return caption; |
| 380 | } catch (error) { |
| 381 | const errorMessage = error.message || 'Unknown error'; |
| 382 | toastr.error(errorMessage, 'Failed to caption'); |
| 383 | console.error(error); |
| 384 | return ''; |
| 385 | } finally { |
| 386 | setImageIcon(); |
| 387 | } |
| 388 | } |
| 389 | |
| 390 | function onRefineModeInput() { |
| 391 | extension_settings.caption.refine_mode = $('#caption_refine_mode').prop('checked'); |
| 392 | saveSettingsDebounced(); |
| 393 | } |
| 394 | |
| 395 | /** |
| 396 | * Callback for the /caption command. |
| 397 | * @param {object} args Named parameters |
| 398 | * @param {string} prompt Caption prompt |
| 399 | */ |
| 400 | async function captionCommandCallback(args, prompt) { |
| 401 | const quiet = isTrueBoolean(args?.quiet); |
| 402 | const messageId = args?.mesId ?? args?.id; |
| 403 | const index = Number(args?.index ?? 0); |
| 404 | |
| 405 | if (!isNaN(Number(messageId))) { |
| 406 | /** @type {ChatMessage} */ |
| 407 | const message = getContext().chat[messageId]; |
| 408 | if (Array.isArray(message?.extra?.media) && message.extra.media.length > 0) { |
| 409 | try { |
| 410 | const mediaAttachment = message.extra.media[index] || message.extra.media[0]; |
| 411 | if (!mediaAttachment || !mediaAttachment.url) { |
| 412 | toastr.error('The specified message does not contain an image.'); |
| 413 | return ''; |
| 414 | } |
| 415 | if (mediaAttachment.type === MEDIA_TYPE.AUDIO) { |
| 416 | toastr.error('The specified media is an audio file. Captioning audio files is not supported.'); |
| 417 | return ''; |
| 418 | } |
| 419 | if (mediaAttachment.type === MEDIA_TYPE.VIDEO && !isVideoCaptioningAvailable()) { |
| 420 | toastr.error('The specified media is a video. Captioning videos is not supported for the current source.'); |
| 421 | return ''; |
| 422 | } |
| 423 | const fetchResult = await fetch(mediaAttachment.url); |
| 424 | const blob = await fetchResult.blob(); |
| 425 | const fileName = mediaAttachment.url.split('/').pop().split('?')[0] || 'image.jpg'; |
| 426 | const file = new File([blob], fileName, { type: blob.type }); |
| 427 | return await getCaptionForFile(file, prompt, quiet); |
| 428 | } catch (error) { |
| 429 | toastr.error('Failed to get image from the message. Make sure the image is accessible.'); |
| 430 | return ''; |
| 431 | } |
| 432 | } |
| 433 | } |
| 434 | |
| 435 | return new Promise(resolve => { |
| 436 | const input = document.createElement('input'); |
| 437 | input.type = 'file'; |
| 438 | input.accept = 'image/*,video/*'; |
| 439 | input.onchange = async (e) => { |
| 440 | const caption = await onSelectImage(e, prompt, quiet); |
| 441 | resolve(caption); |
| 442 | }; |
| 443 | input.oncancel = () => resolve(''); |
| 444 | input.click(); |
| 445 | }); |
| 446 | } |
| 447 | |
| 448 | /** |
| 449 | * Checks if video captioning is available for the current source. |
| 450 | * @returns {boolean} True if video captioning is supported for the current source. |
| 451 | */ |
| 452 | function isVideoCaptioningAvailable() { |
| 453 | if (extension_settings.caption.source !== 'multimodal') { |
| 454 | return false; |
| 455 | } |
| 456 | |
| 457 | return ['google', 'vertexai', 'zai'].includes(extension_settings.caption.multimodal_api); |
| 458 | } |
| 459 | |
| 460 | export async function init() { |
| 461 | function addSendPictureButton() { |
| 462 | const sendButton = $(` |
| 463 | <div id="send_picture" class="list-group-item flex-container flexGap5"> |
| 464 | <div class="fa-solid fa-image extensionsMenuExtensionButton"></div> |
| 465 | <span data-i18n="Generate Caption">Generate Caption</span> |
| 466 | </div>`); |
| 467 | |
| 468 | $('#caption_wand_container').append(sendButton); |
| 469 | $(sendButton).on('click', () => { |
| 470 | const hasCaptionModule = (() => { |
| 471 | const settings = extension_settings.caption; |
| 472 | |
| 473 | // Handle non-multimodal sources |
| 474 | if (settings.source === 'extras' && modules.includes('caption')) return true; |
| 475 | if (settings.source === 'local' || settings.source === 'horde') return true; |
| 476 | |
| 477 | // Handle multimodal sources |
| 478 | if (settings.source === 'multimodal') { |
| 479 | const api = settings.multimodal_api; |
| 480 | const altEndpointEnabled = settings.alt_endpoint_enabled; |
| 481 | const altEndpointUrl = settings.alt_endpoint_url; |
| 482 | |
| 483 | // APIs that support reverse proxy |
| 484 | const reverseProxyApis = { |
| 485 | 'openai': SECRET_KEYS.OPENAI, |
| 486 | 'mistral': SECRET_KEYS.MISTRALAI, |
| 487 | 'google': SECRET_KEYS.MAKERSUITE, |
| 488 | 'vertexai': SECRET_KEYS.VERTEXAI, |
| 489 | 'anthropic': SECRET_KEYS.CLAUDE, |
| 490 | 'xai': SECRET_KEYS.XAI, |
| 491 | 'zai': SECRET_KEYS.ZAI, |
| 492 | 'moonshot': SECRET_KEYS.MOONSHOT, |
| 493 | }; |
| 494 | |
| 495 | if (reverseProxyApis[api]) { |
| 496 | if (secret_state[reverseProxyApis[api]] || settings.allow_reverse_proxy) { |
| 497 | return true; |
| 498 | } |
| 499 | } |
| 500 | |
| 501 | const chatCompletionApis = { |
| 502 | 'openrouter': SECRET_KEYS.OPENROUTER, |
| 503 | 'groq': SECRET_KEYS.GROQ, |
| 504 | 'cohere': SECRET_KEYS.COHERE, |
| 505 | 'aimlapi': SECRET_KEYS.AIMLAPI, |
| 506 | 'nanogpt': SECRET_KEYS.NANOGPT, |
| 507 | 'chutes': SECRET_KEYS.CHUTES, |
| 508 | 'electronhub': SECRET_KEYS.ELECTRONHUB, |
| 509 | 'pollinations': SECRET_KEYS.POLLINATIONS, |
| 510 | 'workers_ai': SECRET_KEYS.WORKERS_AI, |
| 511 | }; |
| 512 | |
| 513 | if (chatCompletionApis[api] && secret_state[chatCompletionApis[api]]) { |
| 514 | return true; |
| 515 | } |
| 516 | |
| 517 | const textCompletionApis = { |
| 518 | 'ollama': textgen_types.OLLAMA, |
| 519 | 'llamacpp': textgen_types.LLAMACPP, |
| 520 | 'ooba': textgen_types.OOBA, |
| 521 | 'koboldcpp': textgen_types.KOBOLDCPP, |
| 522 | 'vllm': textgen_types.VLLM, |
| 523 | }; |
| 524 | |
| 525 | if (textCompletionApis[api] && altEndpointEnabled && altEndpointUrl) { |
| 526 | return true; |
| 527 | } |
| 528 | |
| 529 | if (textCompletionApis[api] && !altEndpointEnabled && textgenerationwebui_settings.server_urls[textCompletionApis[api]]) { |
| 530 | return true; |
| 531 | } |
| 532 | |
| 533 | // Custom API doesn't need additional checks |
| 534 | if (api === 'custom') { |
| 535 | return true; |
| 536 | } |
| 537 | } |
| 538 | |
| 539 | return false; |
| 540 | })(); |
| 541 | |
| 542 | if (!hasCaptionModule) { |
| 543 | toastr.error('Choose other captioning source in the extension settings.', 'Captioning is not available'); |
| 544 | return; |
| 545 | } |
| 546 | |
| 547 | $('#img_file').trigger('click'); |
| 548 | }); |
| 549 | } |
| 550 | function addPictureSendForm() { |
| 551 | const imgInput = document.createElement('input'); |
| 552 | imgInput.type = 'file'; |
| 553 | imgInput.id = 'img_file'; |
| 554 | imgInput.accept = 'image/*,video/*'; |
| 555 | imgInput.hidden = true; |
| 556 | imgInput.addEventListener('change', (e) => onSelectImage(e, '', false)); |
| 557 | const imgForm = document.createElement('form'); |
| 558 | imgForm.id = 'img_form'; |
| 559 | imgForm.appendChild(imgInput); |
| 560 | imgForm.hidden = true; |
| 561 | $('#form_sheld').append(imgForm); |
| 562 | } |
| 563 | async function switchMultimodalBlocks() { |
| 564 | await addRemoteEndpointModels(); |
| 565 | const isMultimodal = extension_settings.caption.source === 'multimodal'; |
| 566 | if (!extension_settings.caption.multimodal_model) { |
| 567 | const dropdown = $('#caption_multimodal_model'); |
| 568 | const options = dropdown.find(`option[data-type="${extension_settings.caption.multimodal_api}"]`); |
| 569 | extension_settings.caption.multimodal_model = String(options.first().val()); |
| 570 | } |
| 571 | $('#caption_multimodal_block').toggle(isMultimodal); |
| 572 | $('#caption_prompt_block').toggle(isMultimodal); |
| 573 | $('#caption_multimodal_api').val(extension_settings.caption.multimodal_api); |
| 574 | $('#caption_multimodal_model').val(extension_settings.caption.multimodal_model); |
| 575 | $('#caption_multimodal_block [data-type]').each(function () { |
| 576 | const type = $(this).data('type'); |
| 577 | const types = type.split(','); |
| 578 | $(this).toggle(types.includes(extension_settings.caption.multimodal_api)); |
| 579 | }); |
| 580 | } |
| 581 | async function addSettings() { |
| 582 | const html = await renderExtensionTemplateAsync('caption', 'settings', { TEMPLATE_DEFAULT, PROMPT_DEFAULT }); |
| 583 | $('#caption_container').append(html); |
| 584 | } |
| 585 | |
| 586 | async function addRemoteEndpointModels() { |
| 587 | async function processEndpoint(api, url, additionalParams = {}) { |
| 588 | const dropdown = document.getElementById('caption_multimodal_model'); |
| 589 | if (!(dropdown instanceof HTMLSelectElement)) { |
| 590 | return; |
| 591 | } |
| 592 | if (extension_settings.caption.source !== 'multimodal' || extension_settings.caption.multimodal_api !== api) { |
| 593 | return; |
| 594 | } |
| 595 | const options = Array.from(dropdown.options); |
| 596 | const response = await fetch(url, { |
| 597 | method: 'POST', |
| 598 | headers: getRequestHeaders(), |
| 599 | body: JSON.stringify(additionalParams), |
| 600 | }); |
| 601 | if (!response.ok) { |
| 602 | return; |
| 603 | } |
| 604 | const modelIds = await response.json(); |
| 605 | if (Array.isArray(modelIds) && modelIds.length > 0) { |
| 606 | modelIds.sort().forEach((modelId) => { |
| 607 | if (!modelId || typeof modelId !== 'string' || options.some(o => o.value === modelId && o.dataset.type === api)) { |
| 608 | return; |
| 609 | } |
| 610 | const option = document.createElement('option'); |
| 611 | option.value = modelId; |
| 612 | option.textContent = modelId; |
| 613 | option.dataset.type = api; |
| 614 | dropdown.add(option); |
| 615 | }); |
| 616 | } |
| 617 | } |
| 618 | |
| 619 | await processEndpoint('openrouter', '/api/openrouter/models/multimodal'); |
| 620 | await processEndpoint('aimlapi', '/api/backends/chat-completions/multimodal-models/aimlapi'); |
| 621 | await processEndpoint('pollinations', '/api/backends/chat-completions/multimodal-models/pollinations'); |
| 622 | await processEndpoint('nanogpt', '/api/backends/chat-completions/multimodal-models/nanogpt'); |
| 623 | await processEndpoint('chutes', '/api/backends/chat-completions/multimodal-models/chutes'); |
| 624 | await processEndpoint('electronhub', '/api/backends/chat-completions/multimodal-models/electronhub'); |
| 625 | await processEndpoint('mistral', '/api/backends/chat-completions/multimodal-models/mistral'); |
| 626 | await processEndpoint('xai', '/api/backends/chat-completions/multimodal-models/xai'); |
| 627 | await processEndpoint('moonshot', '/api/backends/chat-completions/multimodal-models/moonshot'); |
| 628 | await processEndpoint('workers_ai', '/api/backends/chat-completions/multimodal-models/workers_ai', { workers_ai_account_id: oai_settings.workers_ai_account_id }); |
| 629 | } |
| 630 | |
| 631 | await addSettings(); |
| 632 | addPictureSendForm(); |
| 633 | addSendPictureButton(); |
| 634 | setImageIcon(); |
| 635 | migrateSettings(); |
| 636 | await switchMultimodalBlocks(); |
| 637 | |
| 638 | $('#caption_refine_mode').prop('checked', !!(extension_settings.caption.refine_mode)); |
| 639 | $('#caption_allow_reverse_proxy').prop('checked', !!(extension_settings.caption.allow_reverse_proxy)); |
| 640 | $('#caption_prompt_ask').prop('checked', !!(extension_settings.caption.prompt_ask)); |
| 641 | $('#caption_auto_mode').prop('checked', !!(extension_settings.caption.auto_mode)); |
| 642 | $('#caption_source').val(extension_settings.caption.source); |
| 643 | $('#caption_prompt').val(extension_settings.caption.prompt); |
| 644 | $('#caption_template').val(extension_settings.caption.template); |
| 645 | $('#caption_refine_mode').on('input', onRefineModeInput); |
| 646 | $('#caption_source').on('change', async () => { |
| 647 | extension_settings.caption.source = String($('#caption_source').val()); |
| 648 | await switchMultimodalBlocks(); |
| 649 | saveSettingsDebounced(); |
| 650 | }); |
| 651 | $('#caption_prompt').on('input', () => { |
| 652 | extension_settings.caption.prompt = String($('#caption_prompt').val()); |
| 653 | saveSettingsDebounced(); |
| 654 | }); |
| 655 | $('#caption_template').on('input', () => { |
| 656 | extension_settings.caption.template = String($('#caption_template').val()); |
| 657 | saveSettingsDebounced(); |
| 658 | }); |
| 659 | $('#caption_allow_reverse_proxy').on('input', () => { |
| 660 | extension_settings.caption.allow_reverse_proxy = $('#caption_allow_reverse_proxy').prop('checked'); |
| 661 | saveSettingsDebounced(); |
| 662 | }); |
| 663 | $('#caption_prompt_ask').on('input', () => { |
| 664 | extension_settings.caption.prompt_ask = $('#caption_prompt_ask').prop('checked'); |
| 665 | saveSettingsDebounced(); |
| 666 | }); |
| 667 | $('#caption_auto_mode').on('input', () => { |
| 668 | extension_settings.caption.auto_mode = !!$('#caption_auto_mode').prop('checked'); |
| 669 | saveSettingsDebounced(); |
| 670 | }); |
| 671 | $('#caption_ollama_pull').on('click', (e) => { |
| 672 | const selectedModel = extension_settings.caption.multimodal_model; |
| 673 | const staticModels = { 'ollama_current': textgenerationwebui_settings.ollama_model, 'ollama_custom': extension_settings.caption.ollama_custom_model }; |
| 674 | const presetModel = staticModels[selectedModel] || selectedModel; |
| 675 | e.preventDefault(); |
| 676 | $('#ollama_download_model').trigger('click'); |
| 677 | $('.popup .popup-input').val(presetModel); |
| 678 | }); |
| 679 | $('#caption_multimodal_api').on('change', async () => { |
| 680 | const api = String($('#caption_multimodal_api').val()); |
| 681 | extension_settings.caption.multimodal_api = api; |
| 682 | extension_settings.caption.multimodal_model = ''; |
| 683 | await switchMultimodalBlocks(); |
| 684 | saveSettingsDebounced(); |
| 685 | }); |
| 686 | $('#caption_multimodal_model').on('change', () => { |
| 687 | extension_settings.caption.multimodal_model = String($('#caption_multimodal_model').val()); |
| 688 | saveSettingsDebounced(); |
| 689 | }); |
| 690 | $('#caption_altEndpoint_url').val(extension_settings.caption.alt_endpoint_url).on('input', () => { |
| 691 | extension_settings.caption.alt_endpoint_url = String($('#caption_altEndpoint_url').val()); |
| 692 | saveSettingsDebounced(); |
| 693 | }); |
| 694 | $('#caption_altEndpoint_enabled').prop('checked', !!(extension_settings.caption.alt_endpoint_enabled)).on('input', () => { |
| 695 | extension_settings.caption.alt_endpoint_enabled = !!$('#caption_altEndpoint_enabled').prop('checked'); |
| 696 | saveSettingsDebounced(); |
| 697 | }); |
| 698 | $('#caption_show_in_chat').prop('checked', !!(extension_settings.caption.show_in_chat)).on('input', () => { |
| 699 | extension_settings.caption.show_in_chat = !!$('#caption_show_in_chat').prop('checked'); |
| 700 | saveSettingsDebounced(); |
| 701 | }); |
| 702 | $('#caption_ollama_custom_model').val(extension_settings.caption.ollama_custom_model || '').on('input', () => { |
| 703 | extension_settings.caption.ollama_custom_model = String($('#caption_ollama_custom_model').val()).trim(); |
| 704 | saveSettingsDebounced(); |
| 705 | }); |
| 706 | $('#caption_custom_model').val(extension_settings.caption.custom_model || '').on('input', () => { |
| 707 | extension_settings.caption.custom_model = String($('#caption_custom_model').val()).trim(); |
| 708 | saveSettingsDebounced(); |
| 709 | }); |
| 710 | $('#caption_refresh_models').on('click', async () => { |
| 711 | extension_settings.caption.multimodal_model = ''; |
| 712 | await switchMultimodalBlocks(); |
| 713 | saveSettingsDebounced(); |
| 714 | }); |
| 715 | |
| 716 | const onMessageEvent = async (/** @type {number} */ messageId) => { |
| 717 | if (!extension_settings.caption.auto_mode) { |
| 718 | return; |
| 719 | } |
| 720 | |
| 721 | const message = getContext().chat[messageId]; |
| 722 | if (Array.isArray(message?.extra?.media) && message.extra.media.length > 0) { |
| 723 | for (let mediaIndex = 0; mediaIndex < message.extra.media.length; mediaIndex++) { |
| 724 | const mediaAttachment = message.extra.media[mediaIndex]; |
| 725 | if (mediaAttachment.type === MEDIA_TYPE.VIDEO && !isVideoCaptioningAvailable()) { |
| 726 | continue; |
| 727 | } |
| 728 | if (mediaAttachment.type === MEDIA_TYPE.AUDIO) { |
| 729 | continue; |
| 730 | } |
| 731 | // Skip already captioned images and non-uploaded (generated, etc.) images |
| 732 | if (mediaAttachment.source !== MEDIA_SOURCE.UPLOAD || mediaAttachment.captioned) { |
| 733 | continue; |
| 734 | } |
| 735 | try { |
| 736 | await captionExistingMessage(message, mediaIndex); |
| 737 | } catch (e) { |
| 738 | console.error(`Auto-captioning failed for message ID ${messageId}, media index ${mediaIndex}`, e); |
| 739 | continue; |
| 740 | } |
| 741 | } |
| 742 | } |
| 743 | }; |
| 744 | |
| 745 | eventSource.on(event_types.MESSAGE_SENT, onMessageEvent); |
| 746 | eventSource.on(event_types.MESSAGE_FILE_EMBEDDED, onMessageEvent); |
| 747 | |
| 748 | $(document).on('click', '.mes_img_caption', async function () { |
| 749 | const animationClass = 'fa-fade'; |
| 750 | const messageBlock = $(this).closest('.mes'); |
| 751 | const mediaContainer = $(this).closest('.mes_media_container'); |
| 752 | const messageMedia = mediaContainer.find('.mes_img, .mes_video'); |
| 753 | if (messageMedia.hasClass(animationClass)) return; |
| 754 | messageMedia.addClass(animationClass); |
| 755 | try { |
| 756 | const messageId = Number(messageBlock.attr('mesid')); |
| 757 | const mediaIndex = Number(mediaContainer.attr('data-index')); |
| 758 | const data = getContext().chat[messageId]; |
| 759 | await captionExistingMessage(data, mediaIndex); |
| 760 | appendMediaToMessage(data, messageBlock, SCROLL_BEHAVIOR.KEEP); |
| 761 | await saveChatConditional(); |
| 762 | } catch (e) { |
| 763 | console.error('Message image recaption failed', e); |
| 764 | toastr.error(e.message || 'Unknown error', 'Failed to caption'); |
| 765 | } finally { |
| 766 | messageMedia.removeClass(animationClass); |
| 767 | } |
| 768 | }); |
| 769 | |
| 770 | SlashCommandParser.addCommandObject(SlashCommand.fromProps({ |
| 771 | name: 'caption', |
| 772 | callback: captionCommandCallback, |
| 773 | returns: 'caption', |
| 774 | namedArgumentList: [ |
| 775 | new SlashCommandNamedArgument( |
| 776 | 'quiet', 'suppress sending a captioned message', [ARGUMENT_TYPE.BOOLEAN], false, false, 'false', |
| 777 | ), |
| 778 | SlashCommandNamedArgument.fromProps({ |
| 779 | name: 'mesId', |
| 780 | description: 'get image from a message with this ID', |
| 781 | typeList: [ARGUMENT_TYPE.NUMBER], |
| 782 | enumProvider: commonEnumProviders.messages(), |
| 783 | }), |
| 784 | SlashCommandNamedArgument.fromProps({ |
| 785 | name: 'index', |
| 786 | description: 'index of the image in the message to caption (starting from 0)', |
| 787 | typeList: [ARGUMENT_TYPE.NUMBER], |
| 788 | enumProvider: commonEnumProviders.messageMedia(), |
| 789 | }), |
| 790 | ], |
| 791 | unnamedArgumentList: [ |
| 792 | new SlashCommandArgument( |
| 793 | 'prompt', [ARGUMENT_TYPE.STRING], false, |
| 794 | ), |
| 795 | ], |
| 796 | helpString: ` |
| 797 | <div> |
| 798 | Caption an image with an optional prompt and passes the caption down the pipe. |
| 799 | </div> |
| 800 | <div> |
| 801 | Only multimodal sources support custom prompts. |
| 802 | </div> |
| 803 | <div> |
| 804 | Provide a message ID to get an image from a message instead of uploading one. |
| 805 | </div> |
| 806 | <div> |
| 807 | Set the "quiet" argument to true to suppress sending a captioned message, default: false. |
| 808 | </div> |
| 809 | `, |
| 810 | })); |
| 811 | |
| 812 | document.body.classList.add('caption'); |
| 813 | } |