Add vLLM as multimodal captioning source

003066a036f8640d875b44230b32b585e82445fc

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

Parents: a287ccfca268aae06d7c1e4a0cf38ddf8e320c9c
4 files changed, +26 -2Showing whitespace changes
public/scripts/extensions/caption/index.js+1 -0
@@ -356,6 +356,7 @@ jQuery(async function () {
356356 (extension_settings.caption.source === 'multimodal' && extension_settings.caption.multimodal_api === 'llamacpp' && textgenerationwebui_settings.server_urls[textgen_types.LLAMACPP]) ||
357357 (extension_settings.caption.source === 'multimodal' && extension_settings.caption.multimodal_api === 'ooba' && textgenerationwebui_settings.server_urls[textgen_types.OOBA]) ||
358358 (extension_settings.caption.source === 'multimodal' && extension_settings.caption.multimodal_api === 'koboldcpp' && textgenerationwebui_settings.server_urls[textgen_types.KOBOLDCPP]) ||
359+ (extension_settings.caption.source === 'multimodal' && extension_settings.caption.multimodal_api === 'vllm' && textgenerationwebui_settings.server_urls[textgen_types.VLLM]) ||
359360 (extension_settings.caption.source === 'multimodal' && extension_settings.caption.multimodal_api === 'custom') ||
360361 extension_settings.caption.source === 'local' ||
361362 extension_settings.caption.source === 'horde';
public/scripts/extensions/caption/settings.html+2 -0
@@ -26,6 +26,7 @@
2626 <option value="openai">OpenAI</option>
2727 <option value="openrouter">OpenRouter</option>
2828 <option value="ooba" data-i18n="Text Generation WebUI (oobabooga)">Text Generation WebUI (oobabooga)</option>
29+ <option value="vllm">vLLM</option>
2930 </select>
3031 </div>
3132 <div class="flex1 flex-container flexFlowColumn flexNoGap">
@@ -66,6 +67,7 @@
6667 <option data-type="llamacpp" value="llamacpp_current" data-i18n="currently_loaded">[Currently loaded]</option>
6768 <option data-type="ooba" value="ooba_current" data-i18n="currently_loaded">[Currently loaded]</option>
6869 <option data-type="koboldcpp" value="koboldcpp_current" data-i18n="currently_loaded">[Currently loaded]</option>
70+ <option data-type="vllm" value="vllm_current" data-i18n="currently_selected">[Currently selected]</option>
6971 <option data-type="custom" value="custom_current" data-i18n="currently_selected">[Currently selected]</option>
7072 </select>
7173 </div>
public/scripts/extensions/shared.js+17 -0
@@ -34,6 +34,7 @@ export async function getMultimodalCaption(base64Img, prompt) {
3434 const isCustom = extension_settings.caption.multimodal_api === 'custom';
3535 const isOoba = extension_settings.caption.multimodal_api === 'ooba';
3636 const isKoboldCpp = extension_settings.caption.multimodal_api === 'koboldcpp';
37+ const isVllm = extension_settings.caption.multimodal_api === 'vllm';
3738 const base64Bytes = base64Img.length * 0.75;
3839 const compressionLimit = 2 * 1024 * 1024;
3940 if ((['google', 'openrouter'].includes(extension_settings.caption.multimodal_api) && base64Bytes > compressionLimit) || isOoba || isKoboldCpp) {
@@ -65,6 +66,14 @@ export async function getMultimodalCaption(base64Img, prompt) {
6566 requestBody.server_url = textgenerationwebui_settings.server_urls[textgen_types.OLLAMA];
6667 }
6768
69+ if (isVllm) {
70+ if (extension_settings.caption.multimodal_model === 'vllm_current') {
71+ requestBody.model = textgenerationwebui_settings.vllm_model;
72+ }
73+
74+ requestBody.server_url = textgenerationwebui_settings.server_urls[textgen_types.VLLM];
75+ }
76+
6877 if (isLlamaCpp) {
6978 requestBody.server_url = textgenerationwebui_settings.server_urls[textgen_types.LLAMACPP];
7079 }
@@ -151,6 +160,14 @@ function throwIfInvalidModel(useReverseProxy) {
151160 throw new Error('KoboldCpp server URL is not set.');
152161 }
153162
163+ if (extension_settings.caption.multimodal_api === 'vllm' && !textgenerationwebui_settings.server_urls[textgen_types.VLLM]) {
164+ throw new Error('vLLM server URL is not set.');
165+ }
166+
167+ if (extension_settings.caption.multimodal_api === 'vllm' && extension_settings.caption.multimodal_model === 'vllm_current' && !textgenerationwebui_settings.vllm_model) {
168+ throw new Error('vLLM model is not set.');
169+ }
170+
154171 if (extension_settings.caption.multimodal_api === 'custom' && !oai_settings.custom_url) {
155172 throw new Error('Custom API URL is not set.');
156173 }
src/endpoints/openai.js+6 -2
@@ -43,7 +43,11 @@ router.post('/caption-image', jsonParser, async (request, response) => {
4343 key = readSecret(request.user.directories, SECRET_KEYS.KOBOLDCPP);
4444 }
4545
46- if (!key && !request.body.reverse_proxy && ['custom', 'ooba', 'koboldcpp'].includes(request.body.api) === false) {
46+ if (request.body.api === 'vllm') {
47+ key = readSecret(request.user.directories, SECRET_KEYS.VLLM);
48+ }
49+
50+ if (!key && !request.body.reverse_proxy && ['custom', 'ooba', 'koboldcpp', 'vllm'].includes(request.body.api) === false) {
4751 console.log('No key found for API', request.body.api);
4852 return response.sendStatus(400);
4953 }
@@ -110,7 +114,7 @@ router.post('/caption-image', jsonParser, async (request, response) => {
110114 });
111115 }
112116
113117 if (request.body.api === 'koboldcpp' || request.body.api === 'vllm') {
114118 apiUrl = `${trimV1(request.body.server_url)}/v1/chat/completions`;
115119 }
116120