Merge remote-tracking branch 'origin/main' into prefill_overlap

Author: wanghao7

lightllm/models/internvl/img_process.py

@@ -1,59 +1,56 @@
 import torch
-import torch.nn.functional as F
-from PIL import Image
+import math
 from torchvision import transforms as T
 from torchvision.transforms.functional import InterpolationMode
 
 
-def find_closest_aspect_ratio(aspect_ratio, target_ratios, width, height, image_size):
+def find_closest_aspect_ratio(width, height, min_num=1, max_num=6, image_size=448):
     """
     Find the closest aspect ratio from a list of target ratios to match the given aspect ratio.
     If the difference is the same, use the area to decide the better ratio.
     """
-    best_ratio_diff = float("inf")
-    best_ratio = (1, 1)
-    area = width * height
-    for ratio in target_ratios:
-        target_aspect_ratio = ratio[0] / ratio[1]
-        ratio_diff = abs(aspect_ratio - target_aspect_ratio)
-        if ratio_diff < best_ratio_diff:
-            best_ratio_diff = ratio_diff
-            best_ratio = ratio
-        elif ratio_diff == best_ratio_diff:
-            # Compare areas to decide the better ratio when the difference is the same
-            if area > 0.5 * image_size * image_size * ratio[0] * ratio[1]:
-                best_ratio = ratio
-    return best_ratio
-
-
-def dynamic_preprocess(image, min_num=1, max_num=6, image_size=448, use_thumbnail=False):
+    assert min_num == 1
+    log_ratio = math.log(width / height)
+    ratio = width * height / (image_size * image_size)
+    multiple = min(math.ceil(ratio), max_num)
+    if multiple <= 1:
+        return [1, 1]
+    candidate_split_grids_nums = []
+    for i in [multiple - 1, multiple, multiple + 1]:
+        if i > max_num:
+            continue
+        candidate_split_grids_nums.append(i)
+
+    candidate_grids = []
+    for split_grids_nums in candidate_split_grids_nums:
+        m = 1
+        while m <= split_grids_nums:
+            if split_grids_nums % m == 0:
+                candidate_grids.append([m, split_grids_nums // m])
+            m += 1
+    best_grid = [1, 1]
+    min_error = float("inf")
+    for grid in candidate_grids:
+        error = abs(log_ratio - math.log(grid[0] / grid[1]))
+        if error < min_error:
+            best_grid = grid
+            min_error = error
+
+    return best_grid
+
+
+def dynamic_preprocess(image, min_num=1, max_num=6, image_size=448, use_thumbnail=True):
     """
     Preprocess the image dynamically by finding the closest aspect ratio,
     resizing the image, and splitting it into smaller blocks.
     Optionally add a thumbnail version of the image.
     """
-    orig_width, orig_height = image.size
-    aspect_ratio = orig_width / orig_height
-
-    # Calculate the existing image aspect ratio
-    target_ratios = set(
-        (i, j)
-        for n in range(min_num, max_num + 1)
-        for i in range(1, n + 1)
-        for j in range(1, n + 1)
-        if i * j <= max_num and i * j >= min_num
-    )
-    target_ratios = sorted(target_ratios, key=lambda x: x[0] * x[1])
-
-    # Find the closest aspect ratio to the target
-    target_aspect_ratio = find_closest_aspect_ratio(aspect_ratio, target_ratios, orig_width, orig_height, image_size)
-
-    # Calculate the target width and height
+    original_width, original_height = image.size
+    target_aspect_ratio = find_closest_aspect_ratio(original_width, original_height, min_num, max_num, image_size)
     target_width = image_size * target_aspect_ratio[0]
     target_height = image_size * target_aspect_ratio[1]
     blocks = target_aspect_ratio[0] * target_aspect_ratio[1]
-
-    # Resize the image to the target dimensions
+    # resize the image
     resized_img = image.resize((target_width, target_height))
     processed_images = []
     for i in range(blocks):
@@ -63,40 +60,22 @@ def dynamic_preprocess(image, min_num=1, max_num=6, image_size=448, use_thumbnai
             ((i % (target_width // image_size)) + 1) * image_size,
             ((i // (target_width // image_size)) + 1) * image_size,
         )
-        # Split the image into blocks
+        # split the image
         split_img = resized_img.crop(box)
         processed_images.append(split_img)
-
     assert len(processed_images) == blocks
-
-    # Optionally add a thumbnail version of the image
     if use_thumbnail and len(processed_images) != 1:
         thumbnail_img = image.resize((image_size, image_size))
         processed_images.append(thumbnail_img)
-
     return processed_images
 
 
-def get_image_patch(orign_width, orign_height, min_num=1, max_num=6, image_size=448, use_thumbnail=False):
+def get_image_patch(orign_width, orign_height, min_num=1, max_num=6, image_size=448, use_thumbnail=True):
     """
     Calculate the number of image patches based on the closest aspect ratio
     and the given width and height of the original image.
     """
-    aspect_ratio = orign_width / orign_height
-
-    # calculate the existing image aspect ratio
-    target_ratios = set(
-        (i, j)
-        for n in range(min_num, max_num + 1)
-        for i in range(1, n + 1)
-        for j in range(1, n + 1)
-        if i * j <= max_num and i * j >= min_num
-    )
-    target_ratios = sorted(target_ratios, key=lambda x: x[0] * x[1])
-
-    # find the closest aspect ratio to the target
-    target_aspect_ratio = find_closest_aspect_ratio(aspect_ratio, target_ratios, orign_width, orign_height, image_size)
-
+    target_aspect_ratio = find_closest_aspect_ratio(orign_width, orign_height, min_num, max_num, image_size)
     blocks = target_aspect_ratio[0] * target_aspect_ratio[1]
     if use_thumbnail and blocks != 1:
         blocks += 1

lightllm/models/internvl/internvl_visual.py

@@ -8,6 +8,7 @@
 from torchvision import transforms as T
 from torchvision.transforms.functional import InterpolationMode
 from transformers import AutoModel, AutoTokenizer
+from lightllm.server.multimodal_params import MultimodalParams, ImageItem
 from lightllm.server.embed_cache.utils import read_shm, get_shm_name_data
 from io import BytesIO
 from lightllm.models.internvl.img_process import load_image
@@ -43,21 +44,21 @@ def load_model(self, weight_dir):
     def cuda(self):
         return self
 
-    def encode(self, image_uuids: List):
+    def encode(self, images: List[ImageItem]):
         img_tensors = []
         valid_ids = []
         valid_id = 0
         uuids = []
 
-        for i, url in enumerate(image_uuids):
-            if isinstance(url, int):
-                uuids.append(url)
-                image_data = read_shm(get_shm_name_data(url))
+        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))
-                t = self.load_image_func(image_data)
+                t = self.load_image_func(image_data, max_num=img.extra_params["image_patch_max_num"])
                 img_tensors.append(t)
             else:
-                raise Exception("Unsupport input types: {} for {}".format(type(url), url))
+                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])

lightllm/models/internvl/model.py

@@ -11,6 +11,7 @@
     InternVLLlamaPreAndPostLayerWeight,
     InternVLPhi3PreAndPostLayerWeight,
 )
+from lightllm.server.core.objs import SamplingParams
 from lightllm.models.internvl.layer_weights.pre_and_post_layer_weight import InternVLInternlm2PreAndPostLayerWeight
 from lightllm.models.llava.llava_visual import LlavaVisionModel
 
@@ -40,8 +41,32 @@ def __init__(self, tokenizer, model_cfg, **kwargs):
         self.image_end_id = tokenizer.convert_tokens_to_ids(self.image_end_tag)
         self.get_image_patch_func = get_image_patch_func(kwargs["weight_dir"])
 
+    def init_imageItem_extral_params(
+        self, img: ImageItem, multi_params: MultimodalParams, sampling_params: SamplingParams
+    ):
+        if sampling_params.image_max_patch_num >= 0:
+            img.extra_params["image_patch_max_num"] = sampling_params.image_max_patch_num
+            return
+        elif os.getenv("MAX_PATCH_NUM"):
+            img.extra_params["image_patch_max_num"] = int(os.getenv("MAX_PATCH_NUM"))
+            return
+        else:
+            num_images = len(multi_params.images)
+            if num_images == 1:
+                img.extra_params["image_patch_max_num"] = 12
+            elif num_images > 1 and num_images <= 6:
+                img.extra_params["image_patch_max_num"] = 6
+            elif num_images > 6:
+                img.extra_params["image_patch_max_num"] = 0
+        return
+
     def get_image_token_length(self, img: ImageItem):
-        return self.get_image_patch_func(img.image_w, img.image_h, use_thumbnail=True) * self.image_length
+        return (
+            self.get_image_patch_func(
+                img.image_w, img.image_h, max_num=img.extra_params["image_patch_max_num"], use_thumbnail=True
+            )
+            * self.image_length
+        )
 
     # only change the impl of the encode func:
     def encode(self, prompt, multimodal_params: MultimodalParams = None, **kwargs):

lightllm/models/llava/llava_visual.py

@@ -6,6 +6,7 @@
 from typing import List, Union
 from safetensors import safe_open
 from io import BytesIO
+from lightllm.server.multimodal_params import MultimodalParams, ImageItem
 from lightllm.server.embed_cache.utils import read_shm, get_shm_name_data
 from lightllm.utils.log_utils import init_logger
 
@@ -123,21 +124,21 @@ def forward(self, x):
         x = x.view(B, L, -1)
         return x
 
-    def encode(self, image_uuids: List):
+    def encode(self, images: List[ImageItem]):
         img_tensors = []
         uuids = []
         valid_id = 0
         valid_ids = []
 
-        for i, item in enumerate(image_uuids):
-            if isinstance(item, int):
-                uuids.append(item)
-                image_data = read_shm(get_shm_name_data(item))
+        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(item), item))
+                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])

lightllm/models/llava/model.py

@@ -6,6 +6,7 @@
 from lightllm.models.qwen_vl.layer_infer.pre_layer_infer import LlamaMultimodalPreLayerInfer
 from lightllm.models.llava.layer_weights.pre_and_post_layer_weight import LlavaPreAndPostLayerWeight
 from lightllm.server.multimodal_params import MultimodalParams, ImageItem
+from lightllm.server.core.objs import SamplingParams
 from lightllm.common.build_utils import repair_config
 from transformers import AutoConfig
 
@@ -33,6 +34,11 @@ def __init__(self, tokenizer, model_cfg):
         self.image_length = (image_size // patch_size) ** 2
         self.skip_start = model_cfg.get("skip_start", True)
 
+    def init_imageItem_extral_params(
+        self, img: ImageItem, multi_params: MultimodalParams, sampling_params: SamplingParams
+    ):
+        return
+
     def get_image_token_length(self, img: ImageItem):
         return self.image_length
 

lightllm/models/qwen2_vl/model.py

@@ -7,6 +7,7 @@
 from transformers.feature_extraction_utils import BatchFeature
 from transformers.image_utils import ImageInput
 from transformers.processing_utils import ProcessorMixin
+from lightllm.server.core.objs import SamplingParams
 from transformers.tokenization_utils_base import PaddingStrategy, PreTokenizedInput, TextInput, TruncationStrategy
 from typing import List, Optional, Union
 from transformers.utils import TensorType, logging
@@ -31,6 +32,11 @@ def __init__(self, tokenizer=None, image_processor=None, **kwargs):
         self.image_end_id = kwargs["model_cfg"]["vision_end_token_id"]
         self.image_token_id = kwargs["model_cfg"]["image_token_id"]
 
+    def init_imageItem_extral_params(
+        self, img: ImageItem, multi_params: MultimodalParams, sampling_params: SamplingParams
+    ):
+        return
+
     def get_image_token_length(self, img: ImageItem):
         width = img.image_w
         height = img.image_h

lightllm/models/qwen2_vl/qwen2_visual.py

@@ -41,6 +41,7 @@
 from transformers import AutoProcessor
 from safetensors import safe_open
 from transformers.utils import TensorType
+from lightllm.server.multimodal_params import MultimodalParams, ImageItem
 from lightllm.models.qwen2_vl.vision_process import Qwen2VLImageProcessor
 
 
@@ -425,17 +426,17 @@ def load_model(self, weight_dir):
 
         self.load_state_dict(weight_dict)
 
-    def encode(self, image_uuids: List):
+    def encode(self, images: List[ImageItem]):
         img_tensors = []
         valid_ids = []
         valid_id = 0
         img_grids = []
         uuids = []
 
-        for i, url in enumerate(image_uuids):
-            if isinstance(url, int):
-                uuids.append(url)
-                image_data = read_shm(get_shm_name_data(url))
+        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))
                 image_data = get_image(image_data)
                 image_inputs = self.processor.preprocess(images=image_data, return_tensors="pt")
@@ -444,7 +445,7 @@ def encode(self, image_uuids: List):
                 img_tensors.append(pixel_values)
                 img_grids.append(image_grid_thw)
             else:
-                raise Exception("Unsupport input types: {} for {}".format(type(url), url))
+                raise Exception("Unsupport input types: {} for {}".format(type(img), img))
 
             # must devide merge_length
             cur_num = img_tensors[-1].shape[0] // (self.spatial_merge_size ** 2)

lightllm/models/qwen_vl/model.py

@@ -1,6 +1,7 @@
 import json
 import numpy as np
 import unicodedata
+from lightllm.server.core.objs import SamplingParams
 from lightllm.models.qwen.model import QWenTpPartModel
 from .layer_infer.pre_layer_infer import LlamaMultimodalPreLayerInfer
 from lightllm.server.multimodal_params import MultimodalParams, ImageItem
@@ -19,6 +20,11 @@ def __init__(self, tokenizer, model_cfg):
         # <imgpad>: 151859
         self.image_length = model_cfg["visual"].get("n_queries", 256)
 
+    def init_imageItem_extral_params(
+        self, img: ImageItem, multi_params: MultimodalParams, sampling_params: SamplingParams
+    ):
+        return
+
     def _list_find(self, input_list, target, start_idx):
         cur_list = input_list[start_idx:]
         if target in cur_list:

lightllm/models/vit/layer_infer/post_layer_infer.py

@@ -3,6 +3,7 @@
 import torch.distributed as dist
 from lightllm.models.vit.layer_weights.pre_and_post_layer_weight import ViTPreAndPostLayerWeight
 from lightllm.utils.dist_utils import get_current_rank_in_dp, get_dp_world_size
+from lightllm.models.vit.triton_kernel.gelu_vit import gelu_fwd
 
 
 class ViTPostLayerInfer:
@@ -44,7 +45,8 @@ def forward(self, vit_embeds, layer_weight: ViTPreAndPostLayerWeight):
             layer_weight.mlp1_1_bias_, vit_embeds_norm.view(-1, vit_embeds_norm.shape[-1]), layer_weight.mlp1_1_weight_
         )
 
-        vit_embeds_gelu = torch.nn.functional.gelu(vit_embeds_1)
+        # vit_embeds_gelu = torch.nn.functional.gelu(vit_embeds_1)
+        vit_embeds_gelu = gelu_fwd(vit_embeds_1)
 
         vit_embeds_out = torch.addmm(
             layer_weight.mlp1_3_bias_,