Merge branch 'vectors' into staging

This commit is contained in:
Cohee 2023-09-12 18:10:47 +03:00
commit abc1555c19
3 changed files with 72 additions and 21 deletions

View File

@ -19,6 +19,8 @@ const securityOverride = false;
const extras = { const extras = {
// Text classification model for sentiment analysis. HuggingFace ID of a model in ONNX format. // Text classification model for sentiment analysis. HuggingFace ID of a model in ONNX format.
classificationModel: 'Cohee/distilbert-base-uncased-go-emotions-onnx', classificationModel: 'Cohee/distilbert-base-uncased-go-emotions-onnx',
// Image captioning model. HuggingFace ID of a model in ONNX format.
captioningModel: 'Xenova/vit-gpt2-image-captioning',
}; };
// Request overrides for additional headers // Request overrides for additional headers

View File

@ -1,6 +1,6 @@
import { getBase64Async } from "../../utils.js"; import { getBase64Async, saveBase64AsFile } from "../../utils.js";
import { getContext, getApiUrl, doExtrasFetch, extension_settings } from "../../extensions.js"; import { getContext, getApiUrl, doExtrasFetch, extension_settings, modules } from "../../extensions.js";
import { callPopup, saveSettingsDebounced } from "../../../script.js"; import { callPopup, getRequestHeaders, saveSettingsDebounced } from "../../../script.js";
import { getMessageTimeStamp } from "../../RossAscends-mods.js"; import { getMessageTimeStamp } from "../../RossAscends-mods.js";
export { MODULE_NAME }; export { MODULE_NAME };
@ -8,7 +8,8 @@ const MODULE_NAME = 'caption';
const UPDATE_INTERVAL = 1000; const UPDATE_INTERVAL = 1000;
async function moduleWorker() { async function moduleWorker() {
$('#send_picture').toggle(getContext().onlineStatus !== 'no_connection'); const hasConnection = getContext().onlineStatus !== 'no_connection';
$('#send_picture').toggle(hasConnection);
} }
async function setImageIcon() { async function setImageIcon() {
@ -65,16 +66,21 @@ async function sendCaptionedMessage(caption, image) {
await context.generate('caption'); await context.generate('caption');
} }
async function onSelectImage(e) { async function doCaptionRequest(base64Img) {
setSpinnerIcon(); if (extension_settings.caption.local) {
const file = e.target.files[0]; const apiResult = await fetch('/api/extra/caption', {
method: 'POST',
headers: getRequestHeaders(),
body: JSON.stringify({ image: base64Img })
});
if (!file) { if (!apiResult.ok) {
return; throw new Error('Failed to caption image via local pipeline.');
} }
try { const data = await apiResult.json();
const base64Img = await getBase64Async(file); return data;
} else if (modules.includes('caption')) {
const url = new URL(getApiUrl()); const url = new URL(getApiUrl());
url.pathname = '/api/caption'; url.pathname = '/api/caption';
@ -84,17 +90,42 @@ async function onSelectImage(e) {
'Content-Type': 'application/json', 'Content-Type': 'application/json',
'Bypass-Tunnel-Reminder': 'bypass', 'Bypass-Tunnel-Reminder': 'bypass',
}, },
body: JSON.stringify({ image: base64Img.split(',')[1] }) body: JSON.stringify({ image: base64Img })
}); });
if (apiResult.ok) { if (!apiResult.ok) {
const data = await apiResult.json(); throw new Error('Failed to caption image via Extras.');
const caption = data.caption;
const imageToSave = data.thumbnail ? `data:image/jpeg;base64,${data.thumbnail}` : base64Img;
await sendCaptionedMessage(caption, imageToSave);
} }
const data = await apiResult.json();
return data;
} else {
throw new Error('No captioning module is available.');
}
}
async function onSelectImage(e) {
setSpinnerIcon();
const file = e.target.files[0];
if (!file || !(file instanceof File)) {
return;
}
try {
const context = getContext();
const fileData = await getBase64Async(file);
const base64Format = fileData.split(',')[0].split(';')[0].split('/')[1];
const base64Data = fileData.split(',')[1];
const data = await doCaptionRequest(base64Data);
const caption = data.caption;
const imageToSave = data.thumbnail ? data.thumbnail : base64Data;
const format = data.thumbnail ? 'jpeg' : base64Format;
const imagePath = await saveBase64AsFile(imageToSave, context.name2, '', format);
await sendCaptionedMessage(caption, imagePath);
} }
catch (error) { catch (error) {
toastr.error('Failed to caption image.');
console.log(error); console.log(error);
} }
finally { finally {
@ -118,7 +149,16 @@ jQuery(function () {
$('#extensionsMenu').prepend(sendButton); $('#extensionsMenu').prepend(sendButton);
$(sendButton).hide(); $(sendButton).hide();
$(sendButton).on('click', () => $('#img_file').trigger('click')); $(sendButton).on('click', () => {
const hasCaptionModule = modules.includes('caption') || extension_settings.caption.local;
if (!hasCaptionModule) {
toastr.error('No captioning module is available. Either enable the local captioning pipeline or connect to Extras.');
return;
}
$('#img_file').trigger('click');
});
} }
function addPictureSendForm() { function addPictureSendForm() {
const inputHtml = `<input id="img_file" type="file" accept="image/*">`; const inputHtml = `<input id="img_file" type="file" accept="image/*">`;
@ -138,6 +178,10 @@ jQuery(function () {
<div class="inline-drawer-icon fa-solid fa-circle-chevron-down down"></div> <div class="inline-drawer-icon fa-solid fa-circle-chevron-down down"></div>
</div> </div>
<div class="inline-drawer-content"> <div class="inline-drawer-content">
<label class="checkbox_label" for="caption_local">
<input id="caption_local" type="checkbox" class="checkbox">
Use local captioning pipeline
</label>
<label class="checkbox_label" for="caption_refine_mode"> <label class="checkbox_label" for="caption_refine_mode">
<input id="caption_refine_mode" type="checkbox" class="checkbox"> <input id="caption_refine_mode" type="checkbox" class="checkbox">
Edit captions before generation Edit captions before generation
@ -155,6 +199,11 @@ jQuery(function () {
setImageIcon(); setImageIcon();
moduleWorker(); moduleWorker();
$('#caption_refine_mode').prop('checked', !!(extension_settings.caption.refine_mode)); $('#caption_refine_mode').prop('checked', !!(extension_settings.caption.refine_mode));
$('#caption_local').prop('checked', !!(extension_settings.caption.local));
$('#caption_refine_mode').on('input', onRefineModeInput); $('#caption_refine_mode').on('input', onRefineModeInput);
$('#caption_local').on('input', () => {
extension_settings.caption.local = !!$('#caption_local').prop('checked');
saveSettingsDebounced();
});
setInterval(moduleWorker, UPDATE_INTERVAL); setInterval(moduleWorker, UPDATE_INTERVAL);
}); });

View File

@ -1,10 +1,10 @@
{ {
"display_name": "Image Captioning", "display_name": "Image Captioning",
"loading_order": 4, "loading_order": 4,
"requires": [ "requires": [],
"optional": [
"caption" "caption"
], ],
"optional": [],
"js": "index.js", "js": "index.js",
"css": "style.css", "css": "style.css",
"author": "Cohee#1207", "author": "Cohee#1207",