adapt mineru2 visual

Author: zhhangBian

lightllm/models/mineru2_qwen/image_processing_mineru2.py

@@ -1,5 +1,4 @@
 import ast
-import math
 import re
 from functools import partial, reduce
 from typing import Dict, Optional, Union
@@ -94,33 +93,6 @@ def get_anyres_image_grid_shape(image_size, grid_pinpoints, patch_size):
     return width // patch_size, height // patch_size
 
 
-# This functions is not used.
-def resize_and_pad_image(image, target_resolution):
-    original_width, original_height = image.size
-    target_width, target_height = target_resolution
-
-    scale_w = target_width / original_width
-    scale_h = target_height / original_height
-
-    if scale_w < scale_h:
-        new_width = target_width
-        new_height = min(math.ceil(original_height * scale_w), target_height)
-    else:
-        new_height = target_height
-        new_width = min(math.ceil(original_width * scale_h), target_width)
-
-    # Resize the image
-    resized_image = image.resize((new_width, new_height))
-
-    new_image = Image.new("RGB", (target_width, target_height), (0, 0, 0))
-    paste_x = (target_width - new_width) // 2
-    paste_y = (target_height - new_height) // 2
-    new_image.paste(resized_image, (paste_x, paste_y))
-
-    return new_image
-
-
-# DIFFERENT from sglang.srt.mm_utils.process_anyres_image
 def process_anyres_image(image, processor, grid_pinpoints):
     if isinstance(grid_pinpoints, str) and "x" in grid_pinpoints:
         patch_size = processor.crop_size["height"]

lightllm/models/mineru2_qwen/mineru2_visual.py

@@ -0,0 +1,113 @@
+import re
+import os
+import json
+
+from typing import List
+from io import BytesIO
+from PIL import Image
+from safetensors import safe_open
+
+import torch
+import torch.nn as nn
+from transformers import (
+    CLIPVisionModel,
+    CLIPVisionConfig,
+    SiglipVisionConfig,
+    SiglipVisionModel,
+)
+
+from .configuration_mineru2 import Mineru2QwenConfig
+from .image_processing_mineru2 import Mineru2ImageProcessor
+
+from lightllm.server.multimodal_params import ImageItem
+from lightllm.server.embed_cache.utils import read_shm, get_shm_name_data
+from lightllm.utils.log_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+def build_vision_tower(config: Mineru2QwenConfig):
+    vision_tower = getattr(config, "mm_vision_tower", getattr(config, "vision_tower", ""))
+    model_path = getattr(config, "_name_or_path", "")
+
+    if "clip" in vision_tower.lower():
+        if model_path:
+            vision_config = CLIPVisionConfig.from_pretrained(f"{model_path}/{vision_tower}")
+            return CLIPVisionModel(vision_config)
+        else:
+            vision_config = CLIPVisionConfig.from_pretrained(vision_tower)
+            return CLIPVisionModel(vision_config)
+    elif "siglip" in vision_tower.lower():
+        if model_path:
+            vision_config = SiglipVisionConfig.from_pretrained(f"{model_path}/{vision_tower}")
+            return SiglipVisionModel(vision_config)
+        else:
+            vision_config = SiglipVisionConfig.from_pretrained(vision_tower)
+            return SiglipVisionModel(vision_config)
+    else:
+        raise ValueError(f"Unknown vision tower: {model_path}")
+
+
+def build_vision_projector(config: Mineru2QwenConfig):
+    projector_type = getattr(config, "mm_projector_type", "linear")
+
+    if projector_type == "linear":
+        return nn.Linear(config.mm_hidden_size, config.hidden_size)
+
+    mlp_gelu_match = re.match(r"^mlp(\d+)x_gelu$", projector_type)
+    if mlp_gelu_match:
+        mlp_depth = int(mlp_gelu_match.group(1))
+        modules = [nn.Linear(config.mm_hidden_size, config.hidden_size)]
+        for _ in range(1, mlp_depth):
+            modules.append(nn.GELU())
+            modules.append(nn.Linear(config.hidden_size, config.hidden_size))
+        return nn.Sequential(*modules)
+
+    if projector_type == "identity":
+        return nn.Identity()
+
+    raise ValueError(f"Unknown projector type: {projector_type}")
+
+
+class Mineru2VisionModel:
+    def __init__(self):
+        pass
+
+    def load_model(self, weight_dir):
+        # config_file = os.path.join(weight_dir, "config.json")
+        vision_config = Mineru2QwenConfig.from_pretrained(weight_dir)
+
+        self.vision_tower = build_vision_tower(vision_config)
+        self.projector = build_vision_projector(vision_config)
+        self.image_processor = Mineru2ImageProcessor()
+
+    def forward(self, x):
+        return self.projector(self.vision_tower(x))
+
+    def encode(self, images: List[ImageItem]):
+        img_tensors = []
+        uuids = []
+        valid_id = 0
+        valid_ids = []
+
+        for i, img in enumerate(images):
+            if isinstance(img, ImageItem):
+                uuids.append(img.uuid)
+                image_data = read_shm(get_shm_name_data(img.uuid))
+                image_data = Image.open(BytesIO(image_data)).convert("RGB")
+                t = self.image_processor.preprocess(image_data, return_tensors="pt")["pixel_values"]
+                img_tensors.append(t)
+            else:
+                raise Exception("Unsupport input types: {} for {}".format(type(img), img))
+
+            cur_num = img_tensors[-1].shape[0]
+            valid_ids.append([valid_id, valid_id + cur_num])
+            valid_id += cur_num
+
+        if len(img_tensors) <= 0:
+            return None
+
+        img = torch.cat(img_tensors, dim=0)
+        all_img_embeds = self.forward(img)
+
+        return all_img_embeds, uuids, valid_ids

lightllm/models/mineru2_qwen/model.py

@@ -1,12 +1,8 @@
 from lightllm.models.qwen2.model import Qwen2TpPartModel
 from lightllm.models.registry import ModelRegistry
-from .configuration_mineru2 import Mineru2QwenConfig
 
 
 @ModelRegistry("mineru2_qwen", is_multimodal=True)
 class Mineru2QwenForCausalLM(Qwen2TpPartModel):
-    # a new config class is not necessary
-    config_class = Mineru2QwenConfig
-
     def __init__(self, kvargs):
         super().__init__(kvargs)