在使用Open Webui构建AI后端,生产环境使用vllm时,用户图片过大将直接让Token暴增,从而挤压输出Tokens,甚至超出模型Tokens上限出错,使用以下函数可解决:
"""
title: Image Size Filter
description: Resize large base64 images to prevent token limit exceeded errors. Automatically compresses images that are too large.
author: ulin
version: 1.0
"""
import base64
from io import BytesIO
from PIL import Image
from pydantic import BaseModel, Field
from typing import Optional, Callable, Any, Awaitable
class Filter:
class Valves(BaseModel):
priority: int = Field(default=0, description="Priority level")
max_image_size_kb: int = Field(
default=500,
description="Maximum allowed image size in KB before compression",
)
target_image_width: int = Field(
default=768,
description="Target width for resizing large images (maintains aspect ratio)",
)
quality: int = Field(default=85, description="JPEG compression quality (1-100)")
class UserValves(BaseModel):
pass
def __init__(self):
self.valves = self.Valves()
async def inlet(
self,
body: dict,
__event_emitter__: Callable[[Any], Awaitable[None]],
__model__: Optional[dict] = None,
) -> dict:
messages = body.get("messages", [])
if not messages:
return body
processed_messages = []
image_processed = False
for message in messages:
content = message.get("content", "")
role = message.get("role", "")
if role == "user" and content:
processed_content = await self.process_content(
content, __event_emitter__
)
if processed_content != content:
image_processed = True
message["content"] = processed_content
processed_messages.append(message)
body["messages"] = processed_messages
return body
async def process_content(self, content, event_emitter) -> Any:
"""处理消息内容中的图片"""
if isinstance(content, list):
1. 多模态内容(文本+图片)
processed_items = []
for item in content:
if isinstance(item, dict) and item.get("type") == "image_url":
image_url = item.get("image_url", {}).get("url", "")
if image_url.startswith("data:image"):
processed_url = await self.compress_base64_image(
image_url, event_emitter
)
item["image_url"]["url"] = processed_url
processed_items.append(item)
return processed_items
elif isinstance(content, str) and content.startswith("data:image"):
1. 纯图片消息
return await self.compress_base64_image(content, event_emitter)
else:
1. 纯文本或其他类型内容
return content
async def compress_base64_image(self, base64_str: str, event_emitter) -> str:
"""压缩base64编码的图片"""
try:
1. 分离MIME类型和base64数据
header, data = base64_str.split(",", 1)
1. 解码base64
image_data = base64.b64decode(data)
1. 检查图片大小
size_kb = len(image_data) / 1024
if size_kb <= self.valves.max_image_size_kb:
return base64_str # 图片大小合适,无需压缩
1. 打开图片并处理
image = Image.open(BytesIO(image_data))
original_format = image.format
width, height = image.size
1. 计算新尺寸(保持宽高比)
if width > self.valves.target_image_width:
ratio = self.valves.target_image_width / width
new_width = self.valves.target_image_width
new_height = int(height * ratio)
image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)
1. 转换为RGB模式(如果必要)
if image.mode in ("RGBA", "P"):
image = image.convert("RGB")
1. 压缩并重新编码为base64
output_buffer = BytesIO()
1. 优先使用原始格式,否则使用JPEG
save_format = (
original_format if original_format in ["JPEG", "PNG"] else "JPEG"
)
if save_format == "JPEG":
image.save(
output_buffer,
format="JPEG",
quality=self.valves.quality,
optimize=True,
)
else:
image.save(output_buffer, format=save_format, optimize=True)
compressed_data = output_buffer.getvalue()
new_base64 = base64.b64encode(compressed_data).decode("utf-8")
1. 返回新的base64 URL
mime_type = f"image/{save_format.lower()}"
return f"data:{mime_type};base64,{new_base64}"
except Exception as e:
1. 如果处理失败,返回原始图片并记录错误
print(f"图片压缩失败: {e}")
return base64_str © 版权声明
原创文章版权归红影堂所有,收集文章经过翻译整理,转载需征得本站同意。
THE END






暂无评论内容