Add multimodal captioning for 01.ai

40ee236ca8302f5d9692b6a5b5dd33c96f27e5ed

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

Parents: 7b9eb97c7f3032bee19f8cf98899250a9859cd6d
4 files changed, +17 -3Showing whitespace changes
public/scripts/extensions/caption/index.js+3 -3
@@ -8,13 +8,12 @@ import { textgen_types, textgenerationwebui_settings } from '../../textgen-setti
88import { SlashCommandParser } from '../../slash-commands/SlashCommandParser.js';
99import { SlashCommand } from '../../slash-commands/SlashCommand.js';
1010import { ARGUMENT_TYPE, SlashCommandArgument, SlashCommandNamedArgument } from '../../slash-commands/SlashCommandArgument.js';
11-import { SlashCommandEnumValue } from '../../slash-commands/SlashCommandEnumValue.js';
1211import { commonEnumProviders } from '../../slash-commands/SlashCommandCommonEnumsProvider.js';
1312export { MODULE_NAME };
1413
1514const MODULE_NAME = 'caption';
1615
1716const PROMPT_DEFAULT = 'What’\'s in this image?';
1817const TEMPLATE_DEFAULT = '[{{user}} sends {{char}} a picture that contains: {{caption}}]';
1918
2019/**
@@ -334,7 +333,7 @@ async function getCaptionForFile(file, prompt, quiet) {
334333 }
335334 catch (error) {
336335 const errorMessage = error.message || 'Unknown error';
337336 toastr.error(errorMessage, "'Failed to caption image."');
338337 console.error(error);
339338 return '';
340339 }
@@ -399,6 +398,7 @@ jQuery(async function () {
399398 (modules.includes('caption') && extension_settings.caption.source === 'extras') ||
400399 (extension_settings.caption.source === 'multimodal' && extension_settings.caption.multimodal_api === 'openai' && (secret_state[SECRET_KEYS.OPENAI] || extension_settings.caption.allow_reverse_proxy)) ||
401400 (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]) ||
402402 (extension_settings.caption.source === 'multimodal' && extension_settings.caption.multimodal_api === 'google' && (secret_state[SECRET_KEYS.MAKERSUITE] || extension_settings.caption.allow_reverse_proxy)) ||
403403 (extension_settings.caption.source === 'multimodal' && extension_settings.caption.multimodal_api === 'anthropic' && (secret_state[SECRET_KEYS.CLAUDE] || extension_settings.caption.allow_reverse_proxy)) ||
404404 (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 @@
1717 <div class="flex1 flex-container flexFlowColumn flexNoGap">
1818 <label for="caption_multimodal_api" data-i18n="API">API</label>
1919 <select id="caption_multimodal_api" class="flex1 text_pole">
20+ <option value="zerooneai">01.AI (Yi)</option>
2021 <option value="anthropic">Anthropic</option>
2122 <option value="custom" data-i18n="Custom (OpenAI-compatible)">Custom (OpenAI-compatible)</option>
2223 <option value="google">Google MakerSuite</option>
@@ -32,6 +33,7 @@
3233 <div class="flex1 flex-container flexFlowColumn flexNoGap">
3334 <label for="caption_multimodal_model" data-i18n="Model">Model</label>
3435 <select id="caption_multimodal_model" class="flex1 text_pole">
36+ <option data-type="zerooneai" value="yi-vision">yi-vision</option>
3537 <option data-type="openai" value="gpt-4-vision-preview">gpt-4-vision-preview</option>
3638 <option data-type="openai" value="gpt-4-turbo">gpt-4-turbo</option>
3739 <option data-type="openai" value="gpt-4o">gpt-4o</option>
public/scripts/extensions/shared.js+4 -0
@@ -136,6 +136,10 @@ function throwIfInvalidModel(useReverseProxy) {
136136 throw new Error('Anthropic (Claude) API key is not set.');
137137 }
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+
139143 if (extension_settings.caption.multimodal_api === 'google' && !secret_state[SECRET_KEYS.MAKERSUITE] && !useReverseProxy) {
140144 throw new Error('MakerSuite API key is not set.');
141145 }
src/endpoints/openai.js+8 -0
@@ -47,6 +47,10 @@ router.post('/caption-image', jsonParser, async (request, response) => {
4747 key = readSecret(request.user.directories, SECRET_KEYS.VLLM);
4848 }
4949
50+ if (request.body.api === 'zerooneai') {
51+ key = readSecret(request.user.directories, SECRET_KEYS.ZEROONEAI);
52+ }
53+
5054 if (!key && !request.body.reverse_proxy && ['custom', 'ooba', 'koboldcpp', 'vllm'].includes(request.body.api) === false) {
5155 console.log('No key found for API', request.body.api);
5256 return response.sendStatus(400);
@@ -100,6 +104,10 @@ router.post('/caption-image', jsonParser, async (request, response) => {
100104 apiUrl = `${request.body.server_url}/chat/completions`;
101105 }
102106
107+ if (request.body.api === 'zerooneai') {
108+ apiUrl = 'https://api.01.ai/v1/chat/completions';
109+ }
110+
103111 if (request.body.api === 'ooba') {
104112 apiUrl = `${trimV1(request.body.server_url)}/v1/chat/completions`;
105113 const imgMessage = body.messages.pop();