修改Open Webui用户上传图片大小,节省GPU显存实现性能优化及用户体验

在使用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
若本文对您有帮助,欢迎点赞打赏转发
您的支持将是作者更新最大的动力
点赞22打赏 分享
评论 抢沙发

请登录后发表评论

    暂无评论内容