Add multimodal captioning for 01.ai

40ee236ca8302f5d9692b6a5b5dd33c96f27e5ed

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

4 files changed, +17 -3Ignore whitespace
public/scripts/extensions/caption/index.js+3 -3
@@ -8,13 +8,12 @@ import { textgen_types, textgenerationwebui_settings } from '../../textgen-setti
8import { SlashCommandParser } from '../../slash-commands/SlashCommandParser.js';8import { SlashCommandParser } from '../../slash-commands/SlashCommandParser.js';
9import { SlashCommand } from '../../slash-commands/SlashCommand.js';9import { SlashCommand } from '../../slash-commands/SlashCommand.js';
10import { ARGUMENT_TYPE, SlashCommandArgument, SlashCommandNamedArgument } from '../../slash-commands/SlashCommandArgument.js';10import { ARGUMENT_TYPE, SlashCommandArgument, SlashCommandNamedArgument } from '../../slash-commands/SlashCommandArgument.js';
11import { SlashCommandEnumValue } from '../../slash-commands/SlashCommandEnumValue.js';
12import { commonEnumProviders } from '../../slash-commands/SlashCommandCommonEnumsProvider.js';11import { commonEnumProviders } from '../../slash-commands/SlashCommandCommonEnumsProvider.js';
13export { MODULE_NAME };12export { MODULE_NAME };
1413
15const MODULE_NAME = 'caption';14const MODULE_NAME = 'caption';
1615
17const PROMPT_DEFAULT = 'What’s in this image?';16const PROMPT_DEFAULT = 'What\'s in this image?';
18const TEMPLATE_DEFAULT = '[{{user}} sends {{char}} a picture that contains: {{caption}}]';17const TEMPLATE_DEFAULT = '[{{user}} sends {{char}} a picture that contains: {{caption}}]';
1918
20/**19/**
@@ -334,7 +333,7 @@ async function getCaptionForFile(file, prompt, quiet) {
334 }333 }
335 catch (error) {334 catch (error) {
336 const errorMessage = error.message || 'Unknown error';335 const errorMessage = error.message || 'Unknown error';
337 toastr.error(errorMessage, "Failed to caption image.");336 toastr.error(errorMessage, 'Failed to caption image.');
338 console.error(error);337 console.error(error);
339 return '';338 return '';
340 }339 }
@@ -399,6 +398,7 @@ jQuery(async function () {
399 (modules.includes('caption') && extension_settings.caption.source === 'extras') ||398 (modules.includes('caption') && extension_settings.caption.source === 'extras') ||
400 (extension_settings.caption.source === 'multimodal' && extension_settings.caption.multimodal_api === 'openai' && (secret_state[SECRET_KEYS.OPENAI] || extension_settings.caption.allow_reverse_proxy)) ||399 (extension_settings.caption.source === 'multimodal' && extension_settings.caption.multimodal_api === 'openai' && (secret_state[SECRET_KEYS.OPENAI] || extension_settings.caption.allow_reverse_proxy)) ||
401 (extension_settings.caption.source === 'multimodal' && extension_settings.caption.multimodal_api === 'openrouter' && secret_state[SECRET_KEYS.OPENROUTER]) ||400 (extension_settings.caption.source === 'multimodal' && extension_settings.caption.multimodal_api === 'openrouter' && secret_state[SECRET_KEYS.OPENROUTER]) ||
401 (extension_settings.caption.source === 'multimodal' && extension_settings.caption.multimodal_api === 'zerooneai' && secret_state[SECRET_KEYS.ZEROONEAI]) ||
402 (extension_settings.caption.source === 'multimodal' && extension_settings.caption.multimodal_api === 'google' && (secret_state[SECRET_KEYS.MAKERSUITE] || extension_settings.caption.allow_reverse_proxy)) ||402 (extension_settings.caption.source === 'multimodal' && extension_settings.caption.multimodal_api === 'google' && (secret_state[SECRET_KEYS.MAKERSUITE] || extension_settings.caption.allow_reverse_proxy)) ||
403 (extension_settings.caption.source === 'multimodal' && extension_settings.caption.multimodal_api === 'anthropic' && (secret_state[SECRET_KEYS.CLAUDE] || extension_settings.caption.allow_reverse_proxy)) ||403 (extension_settings.caption.source === 'multimodal' && extension_settings.caption.multimodal_api === 'anthropic' && (secret_state[SECRET_KEYS.CLAUDE] || extension_settings.caption.allow_reverse_proxy)) ||
404 (extension_settings.caption.source === 'multimodal' && extension_settings.caption.multimodal_api === 'ollama' && textgenerationwebui_settings.server_urls[textgen_types.OLLAMA]) ||404 (extension_settings.caption.source === 'multimodal' && extension_settings.caption.multimodal_api === 'ollama' && textgenerationwebui_settings.server_urls[textgen_types.OLLAMA]) ||
public/scripts/extensions/caption/settings.html+2 -0
@@ -17,6 +17,7 @@
17 <div class="flex1 flex-container flexFlowColumn flexNoGap">17 <div class="flex1 flex-container flexFlowColumn flexNoGap">
18 <label for="caption_multimodal_api" data-i18n="API">API</label>18 <label for="caption_multimodal_api" data-i18n="API">API</label>
19 <select id="caption_multimodal_api" class="flex1 text_pole">19 <select id="caption_multimodal_api" class="flex1 text_pole">
20 <option value="zerooneai">01.AI (Yi)</option>
20 <option value="anthropic">Anthropic</option>21 <option value="anthropic">Anthropic</option>
21 <option value="custom" data-i18n="Custom (OpenAI-compatible)">Custom (OpenAI-compatible)</option>22 <option value="custom" data-i18n="Custom (OpenAI-compatible)">Custom (OpenAI-compatible)</option>
22 <option value="google">Google MakerSuite</option>23 <option value="google">Google MakerSuite</option>
@@ -32,6 +33,7 @@
32 <div class="flex1 flex-container flexFlowColumn flexNoGap">33 <div class="flex1 flex-container flexFlowColumn flexNoGap">
33 <label for="caption_multimodal_model" data-i18n="Model">Model</label>34 <label for="caption_multimodal_model" data-i18n="Model">Model</label>
34 <select id="caption_multimodal_model" class="flex1 text_pole">35 <select id="caption_multimodal_model" class="flex1 text_pole">
36 <option data-type="zerooneai" value="yi-vision">yi-vision</option>
35 <option data-type="openai" value="gpt-4-vision-preview">gpt-4-vision-preview</option>37 <option data-type="openai" value="gpt-4-vision-preview">gpt-4-vision-preview</option>
36 <option data-type="openai" value="gpt-4-turbo">gpt-4-turbo</option>38 <option data-type="openai" value="gpt-4-turbo">gpt-4-turbo</option>
37 <option data-type="openai" value="gpt-4o">gpt-4o</option>39 <option data-type="openai" value="gpt-4o">gpt-4o</option>
public/scripts/extensions/shared.js+4 -0
@@ -136,6 +136,10 @@ function throwIfInvalidModel(useReverseProxy) {
136 throw new Error('Anthropic (Claude) API key is not set.');136 throw new Error('Anthropic (Claude) API key is not set.');
137 }137 }
138138
139 if (extension_settings.caption.multimodal_api === 'zerooneai' && !secret_state[SECRET_KEYS.ZEROONEAI]) {
140 throw new Error('01.AI API key is not set.');
141 }
142
139 if (extension_settings.caption.multimodal_api === 'google' && !secret_state[SECRET_KEYS.MAKERSUITE] && !useReverseProxy) {143 if (extension_settings.caption.multimodal_api === 'google' && !secret_state[SECRET_KEYS.MAKERSUITE] && !useReverseProxy) {
140 throw new Error('MakerSuite API key is not set.');144 throw new Error('MakerSuite API key is not set.');
141 }145 }
src/endpoints/openai.js+8 -0
@@ -47,6 +47,10 @@ router.post('/caption-image', jsonParser, async (request, response) => {
47 key = readSecret(request.user.directories, SECRET_KEYS.VLLM);47 key = readSecret(request.user.directories, SECRET_KEYS.VLLM);
48 }48 }
4949
50 if (request.body.api === 'zerooneai') {
51 key = readSecret(request.user.directories, SECRET_KEYS.ZEROONEAI);
52 }
53
50 if (!key && !request.body.reverse_proxy && ['custom', 'ooba', 'koboldcpp', 'vllm'].includes(request.body.api) === false) {54 if (!key && !request.body.reverse_proxy && ['custom', 'ooba', 'koboldcpp', 'vllm'].includes(request.body.api) === false) {
51 console.log('No key found for API', request.body.api);55 console.log('No key found for API', request.body.api);
52 return response.sendStatus(400);56 return response.sendStatus(400);
@@ -100,6 +104,10 @@ router.post('/caption-image', jsonParser, async (request, response) => {
100 apiUrl = `${request.body.server_url}/chat/completions`;104 apiUrl = `${request.body.server_url}/chat/completions`;
101 }105 }
102106
107 if (request.body.api === 'zerooneai') {
108 apiUrl = 'https://api.01.ai/v1/chat/completions';
109 }
110
103 if (request.body.api === 'ooba') {111 if (request.body.api === 'ooba') {
104 apiUrl = `${trimV1(request.body.server_url)}/v1/chat/completions`;112 apiUrl = `${trimV1(request.body.server_url)}/v1/chat/completions`;
105 const imgMessage = body.messages.pop();113 const imgMessage = body.messages.pop();