Update new model

processing_llava.py

@@ -31,27 +31,72 @@ from transformers.utils import TensorType
 import torch
 from open_clip.transform import PreprocessCfg, image_transform_v2
 from modeling_llava import LlavaForConditionalGeneration
+from PIL import Image
+import math
 
 
 class OpenCLIPImageProcessor:
-    def __init__(self, config):
+    def __init__(self, config, crop_size=384, max_tokens=100):
         cfg = PreprocessCfg(**config)
         transform = image_transform_v2(cfg=cfg, is_train=False)
         self.transform = transform
+        self.crop_size = crop_size
+        self.max_tokens = max_tokens
 
-    def __call__(self, image, return_tensors):
-        if isinstance(image, list):
-            outputs = []
-            for item in image:
-                outputs.append(self.transform(item))
-            return {
-                "pixel_values": torch.tensor(outputs),
-            }
-        output = self.transform(image)
+    def __call__(self, image: Image.Image):
+        output = self.transform_func(image)
         return {
-            "pixel_values": output.unsqueeze(0),
+            "pixel_values": output,
         }
 
+    def transform_func(self, image: Image.Image):
+        outputs = []
+        outputs.append(self.transform(image))
+        width, height = image.size
+        crop_size = self.crop_size
+        if width <= crop_size and height <= crop_size:
+            outputs = torch.stack(outputs, dim=0)
+            return outputs
+        total_tokens = math.inf
+        while total_tokens > self.max_tokens:
+            total_tokens = math.floor(
+                (2 * width - crop_size)
+                / crop_size
+                * (2 * height - crop_size)
+                / crop_size
+            )
+            if total_tokens > self.max_tokens:
+                crop_size += 10
+        stride = crop_size // 2
+        x_steps = int(round((2 * width - crop_size) / crop_size))
+        if x_steps < 1:
+            x_steps = 1
+        y_steps = int(round((2 * height - crop_size) / crop_size))
+        if y_steps < 1:
+            y_steps = 1
+        x_coords = []
+        y_coords = []
+        for i in range(x_steps):
+            x_coords.append([i * stride, i * stride + crop_size])
+        if x_coords[-1][1] != width:
+            x_coords[-1][1] = width
+        for i in range(y_steps):
+            y_coords.append([i * stride, i * stride + crop_size])
+        if y_coords[-1][1] != height:
+            y_coords[-1][1] = height
+        image_parts = []
+        for i in range(len(x_coords)):
+            for j in range(len(y_coords)):
+                image_parts.append(
+                    image.crop(
+                        (x_coords[i][0], y_coords[j][0], x_coords[i][1], y_coords[j][1])
+                    )
+                )
+        for image_part in image_parts:
+            outputs.append(self.transform(image_part))
+        outputs = torch.stack(outputs, dim=0)
+        return outputs
+
     @property
     def model_input_names(self):
         return ["pixel_values"]
@@ -75,12 +120,13 @@ class LlavaProcessor:
         return_tensors: Optional[Union[str, TensorType]] = TensorType.PYTORCH,
     ) -> BatchFeature:
         if images is not None:
-            pixel_values = self.image_processor(images, return_tensors=return_tensors)[
+            pixel_values = self.image_processor(images)[
                 "pixel_values"
             ]
             pixel_values = pixel_values.to(model.device).to(model.dtype)
             image_outputs = model.vision_model(pixel_values)
             image_features = model.multi_modal_projector(image_outputs)
+            image_features = image_features.unsqueeze(0)
         else:
             image_features = None
         text_inputs = self.tokenizer(