mirror of
https://github.com/theroyallab/tabbyAPI.git
synced 2026-03-14 15:57:27 +00:00
242ff4f892
The vision module from the ExllamaV2 backend is used in files outside the backends contained folder. Therefore, import ExllamaV2 as an optional dependency here. Signed-off-by: kingbri <bdashore3@proton.me>
78 lines
2.3 KiB
Python
78 lines
2.3 KiB
Python
"""Vision utilities for ExLlamaV2."""
|
|
|
|
import aiohttp
|
|
import base64
|
|
import io
|
|
import re
|
|
from async_lru import alru_cache
|
|
from fastapi import HTTPException
|
|
from PIL import Image
|
|
|
|
from common import model
|
|
from common.networking import (
|
|
handle_request_error,
|
|
)
|
|
from common.optional_dependencies import dependencies
|
|
from common.tabby_config import config
|
|
|
|
# Since this is used outside the Exl2 backend, the dependency
|
|
# may be optional
|
|
if dependencies.exllamav2:
|
|
from exllamav2.generator import ExLlamaV2MMEmbedding
|
|
|
|
|
|
async def get_image(url: str) -> Image:
|
|
if url.startswith("data:image"):
|
|
# Handle base64 image
|
|
match = re.match(r"^data:image\/[a-zA-Z0-9]+;base64,(.*)$", url)
|
|
if match:
|
|
base64_image = match.group(1)
|
|
bytes_image = base64.b64decode(base64_image)
|
|
else:
|
|
error_message = handle_request_error(
|
|
"Failed to read base64 image input.",
|
|
exc_info=False,
|
|
).error.message
|
|
|
|
raise HTTPException(400, error_message)
|
|
|
|
else:
|
|
# Handle image URL
|
|
if config.network.disable_fetch_requests:
|
|
error_message = handle_request_error(
|
|
f"Failed to fetch image from {url} as fetch requests are disabled.",
|
|
exc_info=False,
|
|
).error.message
|
|
|
|
raise HTTPException(400, error_message)
|
|
|
|
async with aiohttp.ClientSession() as session:
|
|
async with session.get(url) as response:
|
|
if response.status == 200:
|
|
bytes_image = await response.read()
|
|
else:
|
|
error_message = handle_request_error(
|
|
f"Failed to fetch image from {url}.",
|
|
exc_info=False,
|
|
).error.message
|
|
|
|
raise HTTPException(400, error_message)
|
|
|
|
return Image.open(io.BytesIO(bytes_image))
|
|
|
|
|
|
# Fetch the return type on runtime
|
|
@alru_cache(20)
|
|
async def get_image_embedding(url: str) -> "ExLlamaV2MMEmbedding":
|
|
image = await get_image(url)
|
|
return model.container.vision_model.get_image_embeddings(
|
|
model=model.container.model,
|
|
tokenizer=model.container.tokenizer,
|
|
image=image,
|
|
text_alias=None,
|
|
)
|
|
|
|
|
|
def clear_image_embedding_cache():
|
|
get_image_embedding.cache_clear()
|