LabelOcr/model/qwen_ocr_vllm.py

import base64
from io import BytesIO
from PIL import Image
import json
import os
from pathlib import Path
from typing import List, Dict, Any, Optional

from qwen_vl_utils import process_vision_info
from transformers import AutoProcessor
import time
import torch

from config import MODEL_PATH, PROMPT_EXTRACT_NAME, PROMPT_EXTRACT_ICON, PROMPT_EXTRACT_COMPONENTS, PROMPT_EXTRACT_KEYWORD, PROMPT_EXTRACT_PREVENTION, PROMPT_EXTRACT_SUPPLIER

# vLLM imports
from vllm import LLM, SamplingParams
from vllm.multimodal.utils import fetch_image


def image_to_base64(pil_image, image_format="JPEG"):
    """将PIL Image图像转换为Base64编码"""
    buffered = BytesIO()
    pil_image.save(buffered, format=image_format)
    img_byte_array = buffered.getvalue()
    encode_image = base64.b64encode(img_byte_array).decode('utf-8')
    return encode_image


class QwenOcrVLLM:
    """基于vLLM加速框架的Qwen OCR推理类

    vLLM优势：
    1. PagedAttention技术 - 高效的KV cache管理
    2. 连续批处理 - 优化GPU利用率
    3. 快速模型执行 - CUDA/cuDNN kernel优化
    4. 支持量化 - AWQ, GPTQ等量化格式
    5. 张量并行 - 支持多GPU推理
    """

    def __init__(
        self,
        icon_dir: str = "./icon",
        tensor_parallel_size: int = 1,
        gpu_memory_utilization: float = 0.9,
        max_model_len: int = 8192,
        dtype: str = "bfloat16",
        trust_remote_code: bool = True,
    ):
        """初始化vLLM模型

        Args:
            icon_dir: icon参考图像目录
            tensor_parallel_size: 张量并行大小（多GPU推理）
            gpu_memory_utilization: GPU显存利用率（0.0-1.0）
            max_model_len: 最大模型序列长度
            dtype: 数据类型（"auto", "half", "float16", "bfloat16", "float", "float32"）
            trust_remote_code: 是否信任远程代码
        """
        print("=" * 60)
        print("初始化 vLLM 加速推理引擎...")
        print("=" * 60)

        # 初始化vLLM引擎
        self.llm = LLM(
            model=MODEL_PATH,
            tensor_parallel_size=tensor_parallel_size,
            gpu_memory_utilization=gpu_memory_utilization,
            max_model_len=max_model_len,
            dtype=dtype,
            trust_remote_code=trust_remote_code,
            # 视觉模型特定参数
            limit_mm_per_prompt={"image": 10},  # 每个prompt最多支持10张图像
        )

        # 加载processor用于消息模板处理
        self.processor = AutoProcessor.from_pretrained(
            MODEL_PATH,
            trust_remote_code=trust_remote_code
        )

        # 加载icon参考图像
        self.icon_dir = icon_dir
        # self.icon_images = self._load_icon_images()

        # 默认采样参数
        self.default_sampling_params = SamplingParams(
            temperature=0.0,      # 使用贪婪解码
            top_p=1.0,
            max_tokens=512,       # 最大生成token数
            stop_token_ids=None,
            skip_special_tokens=True,
        )

        print("=" * 60)
        print("vLLM 引擎初始化完成!")
        print(f"- 模型路径: {MODEL_PATH}")
        print(f"- 张量并行: {tensor_parallel_size} GPU(s)")
        print(f"- 显存利用率: {gpu_memory_utilization * 100:.1f}%")
        print(f"- 数据类型: {dtype}")
        print(f"- 最大序列长度: {max_model_len}")
        print("=" * 60)

        # 模型预热
        print("模型预热中...")
        self._warmup()
        print("模型预热完成!")
        print("=" * 60)

    def _load_icon_images(self) -> Dict[str, Image.Image]:
        """加载icon目录下的所有参考图像"""
        icon_images = {}
        icon_path = Path(self.icon_dir)

        if not icon_path.exists():
            print(f"警告: icon目录 {self.icon_dir} 不存在")
            return icon_images

        # 加载所有png图像文件
        for icon_file in icon_path.glob("*.png"):
            icon_name = icon_file.stem  # 获取文件名(不含扩展名), 如 GHS01
            try:
                icon_image = Image.open(icon_file).convert("RGB")
                icon_images[icon_name] = icon_image
                print(f"已加载icon参考图像: {icon_name}")
            except Exception as e:
                print(f"加载icon图像 {icon_file} 失败: {e}")

        return icon_images

    def _warmup(self):
        """预热模型以触发编译和优化"""
        dummy_image = Image.new('RGB', (224, 224), color='white')
        prompt = PROMPT_EXTRACT_NAME
        try:
            self.inference(dummy_image, prompt, warmup=True)
        except Exception as e:
            print(f"预热过程中出现警告(可忽略): {e}")

    def _build_messages(self, image: Image.Image, prompt: str) -> List[Dict]:
        """构建消息格式

        Args:
            image: PIL Image对象
            prompt: 提示词文本

        Returns:
            消息列表
        """
        messages = [
            {
                "role": "user",
                "content": [
                    {
                        "type": "image",
                        "image": image,
                    },
                    {
                        "type": "text",
                        "text": prompt
                    },
                ],
            }
        ]
        return messages

    def _prepare_inputs(
        self,
        messages: List[Dict]
    ) -> Dict[str, Any]:
        """准备vLLM输入格式

        Args:
            messages: 消息列表

        Returns:
            包含prompt和multi_modal_data的字典
        """
        # 应用chat模板
        text = self.processor.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True
        )

        # 处理视觉信息
        image_inputs, video_inputs = process_vision_info(messages)

        # vLLM 0.6.0+ 新版API格式
        # 直接返回包含文本和多模态数据的字典
        inputs = {
            "prompt": text,
            "multi_modal_data": {
                "image": image_inputs[0] if image_inputs else None
            }
        }

        return inputs

    def inference(
        self,
        image: Image.Image,
        prompt: str,
        warmup: bool = False,
        sampling_params: Optional[SamplingParams] = None
    ) -> List[str]:
        """OCR推理

        Args:
            image: PIL Image对象
            prompt: 提示词
            warmup: 是否为预热模式
            sampling_params: 自定义采样参数

        Returns:
            生成的文本列表
        """
        # 构建消息
        messages = self._build_messages(image, prompt)

        # 准备输入
        inputs = self._prepare_inputs(messages)

        # 使用默认或自定义采样参数
        params = sampling_params if sampling_params else self.default_sampling_params

        # vLLM 0.6.0+ 新版API：直接传递inputs字典
        outputs = self.llm.generate(
            inputs,
            params
        )

        # 提取生成的文本
        generated_texts = [output.outputs[0].text for output in outputs]

        if not warmup:
            return generated_texts

        return generated_texts

    def batch_inference(
        self,
        images: List[Image.Image],
        prompts: List[str],
        sampling_params: Optional[SamplingParams] = None
    ) -> List[str]:
        """批量OCR推理（vLLM的核心优势）

        Args:
            images: PIL Image对象列表
            prompts: 提示词列表
            sampling_params: 自定义采样参数

        Returns:
            生成的文本列表
        """
        if len(images) != len(prompts):
            raise ValueError(f"images数量({len(images)})和prompts数量({len(prompts)})不匹配")

        # 准备所有输入
        all_inputs = []
        for image, prompt in zip(images, prompts):
            messages = self._build_messages(image, prompt)
            inputs = self._prepare_inputs(messages)
            all_inputs.append(inputs)

        # 使用默认或自定义采样参数
        params = sampling_params if sampling_params else self.default_sampling_params

        # vLLM 0.6.0+ 新版API：直接传递inputs列表
        outputs = self.llm.generate(
            all_inputs,
            params
        )

        # 提取生成的文本
        generated_texts = [output.outputs[0].text for output in outputs]

        return generated_texts

    def extract_icons(self, image: Image.Image) -> List[str]:
        """识别图像中的象形图标识

        Args:
            image: PIL Image对象 - 待识别的化学品标签图像

        Returns:
            生成的文本列表
        """
        # 构建包含所有参考图像的messages
        messages = [
            {
                "role": "user",
                "content": []
            }
        ]

        # 添加所有icon参考图像
        content_list = messages[0]["content"]

        # 按GHS编号顺序添加参考图像
        sorted_icons = sorted(self.icon_images.items(), key=lambda x: x[0])
        for icon_name, icon_image in sorted_icons:
            content_list.append({
                "type": "image",
                "image": icon_image,
            })
            content_list.append({
                "type": "text",
                "text": f"参考图像：{icon_name}"
            })

        # 添加待识别的图像
        content_list.append({
            "type": "image",
            "image": image,
        })

        # 添加提示词
        content_list.append({
            "type": "text",
            "text": PROMPT_EXTRACT_ICON
        })

        # 准备输入
        inputs = self._prepare_inputs(messages)

        # vLLM 0.6.0+ 新版API：直接传递inputs字典
        outputs = self.llm.generate(
            inputs,
            self.default_sampling_params
        )

        # 提取生成的文本
        generated_texts = [output.outputs[0].text for output in outputs]

        return generated_texts

    def __del__(self):
        """析构函数，清理资源"""
        try:
            if hasattr(self, 'llm'):
                del self.llm
                print("vLLM引擎已释放")
        except:
            pass


if __name__ == '__main__':
    # 测试代码
    print("初始化 QwenOcrVLLM...")
    qwen_ocr = QwenOcrVLLM(
        tensor_parallel_size=1,      # 单GPU
        gpu_memory_utilization=0.9,  # 使用90%显存
        max_model_len=8192,          # 最大序列长度
        dtype="bfloat16"             # 使用bfloat16精度
    )

    # 测试单张图像推理
    print("\n" + "=" * 60)
    print("测试单张图像推理...")
    print("=" * 60)

    test_image_path = "./test3.jpg"
    if os.path.exists(test_image_path):
        image = Image.open(test_image_path).convert("RGB")

        # 测试提取名称
        start_time = time.time()
        result = qwen_ocr.inference(image, PROMPT_EXTRACT_PREVENTION)
        elapsed = time.time() - start_time

        print(f"\n推理耗时: {elapsed:.3f}秒")
        print(f"提取结果:\n{result[0]}")
    else:
        print(f"测试图像 {test_image_path} 不存在，跳过测试")

    # 测试批量推理
    print("\n" + "=" * 60)
    print("测试批量推理...")
    print("=" * 60)

    if os.path.exists(test_image_path):
        # 创建3张相同的测试图像
        images = [Image.open(test_image_path).convert("RGB") for _ in range(6)]
        prompts = [PROMPT_EXTRACT_ICON, PROMPT_EXTRACT_NAME, PROMPT_EXTRACT_COMPONENTS, PROMPT_EXTRACT_KEYWORD, PROMPT_EXTRACT_PREVENTION, PROMPT_EXTRACT_SUPPLIER]

        start_time = time.time()
        results = qwen_ocr.batch_inference(images, prompts)
        elapsed = time.time() - start_time
        print(results[0])
        print(results[1])
        print(results[2])
        print(results[3])
        print(results[4])
        print(results[5])

        print(f"\n批量推理耗时: {elapsed:.3f}秒")
        print(f"平均每张: {elapsed/3:.3f}秒")
        print(f"批量推理加速比: {(elapsed/3):.3f}秒/张 vs 单张推理")

    print("\n" + "=" * 60)
    print("测试完成!")
    print("=" * 60)