Upload processing_heron.py with huggingface_hub

processing_heron.py

@@ -1,7 +1,9 @@
 from collections.abc import Mapping
+from functools import lru_cache
 from typing import Unpack, cast
 
 import numpy as np
+import torch
 from PIL.Image import Image
 from transformers.feature_extraction_utils import BatchFeature
 from transformers.image_processing_utils import BaseImageProcessor
@@ -72,14 +74,16 @@ class HeronProcessor(ProcessorMixin):
             images = [to_pil_image(image) for image in images]
 
         image_inputs: Mapping = {}
+        num_image_tiles = None
         if images is not None:
-            image_inputs = self._process_images(images, **output_kwargs["images_kwargs"])
-            text = self._insert_media_placeholders(cast(list, text), image_inputs["pixel_values"])
+            if sum(s.count(self.image_token) for s in text) != len(images):
+                raise ValueError("the number of images does not match the number of image tokens in the text")
+            image_inputs = self.process_images(images, **output_kwargs["images_kwargs"])
+            num_image_tiles = image_inputs["pixel_values"].shape[1]
 
         return_tensors = output_kwargs["text_kwargs"].pop("return_tensors", None)
         return_mm_token_type_ids = output_kwargs["text_kwargs"].pop("return_mm_token_type_ids", None)
-        text_inputs = self._process_text(text, **output_kwargs["text_kwargs"])
-        self._check_special_mm_tokens(text, text_inputs, modalities=["image"])  # type: ignore[arg-type]
+        text_inputs = self.process_text(text, num_image_tiles, **output_kwargs["text_kwargs"])
 
         if return_mm_token_type_ids:
             array_ids = np.array(text_inputs["input_ids"])
@@ -89,45 +93,42 @@ class HeronProcessor(ProcessorMixin):
 
         return BatchFeature(data={**text_inputs, **image_inputs}, tensor_type=return_tensors)
 
-    def _insert_media_placeholders(self, text: list[str], image_pixel_values: list, separator="\n") -> list[str]:
-        image_feature_placeholder = self.image_token * self.num_image_features
+    def process_text(
+        self, text: list[str], num_image_tiles: int | None = None, **kwargs: Unpack[TextKwargs]
+    ) -> BatchEncoding:
+        if all(self.image_token not in prompt for prompt in text):
+            return self.tokenizer(text, **kwargs)
+
+        if num_image_tiles is None:
+            raise ValueError("num_image_tiles must be specified when processing image tokens")
 
-        # NOTE: Original implementation appends a newline after image features.
+        image_feature_placeholder = self.image_token * self.num_image_features
+        # NOTE: Original implementation appends an extra newline after image features.
         # https://github.com/NVlabs/VILA/blob/36f6adcd11a10be1580caeb7e647e1b6f8517f89/llava/model/encoders/image/basic.py#L38
         # Instead of appending a newline character, this implementation reserves 1 extra token to be replaced later.
         # This treatment is needed because "\n\n" is tokenized as [271] instead of [198, 198].
         assert self.tokenizer.eos_token is not None
         image_feature_placeholder += self.tokenizer.eos_token
 
+        # Expand image tokens according to the number of image features and tiles
         processed_text = []
-        image_iter = iter(image_pixel_values)
         for prompt in text:
             new_prompt = prompt
             assert "<placeholder>" not in new_prompt
             replace_strings = []
             while self.image_token in new_prompt:
-                try:
-                    image_tiles = next(image_iter)
-                except StopIteration:
-                    raise ValueError("insufficient number of images provided for the prompts") from None
-                replace_strings.append(separator.join([image_feature_placeholder] * len(image_tiles)))
+                replace_strings.append("\n".join([image_feature_placeholder] * num_image_tiles))
                 new_prompt = new_prompt.replace(self.image_token, "<placeholder>", 1)
             for s in replace_strings:
                 new_prompt = new_prompt.replace("<placeholder>", s, 1)
             processed_text.append(new_prompt)
 
-        if next(image_iter, None) is not None:
-            raise ValueError("excessive number of images provided for the prompts")
-        return processed_text
+        encoding = self.tokenizer(processed_text, **kwargs)
 
-    def _process_text(self, text: list[str], **kwargs: Unpack[TextKwargs]) -> BatchEncoding:
+        # Replace the last token of every image tile with the newline token
         token_ids = self.tokenizer.encode("\n")
         assert len(token_ids) == 1
         newline_token_id = token_ids[0]
-
-        encoding = self.tokenizer(text, **kwargs)
-
-        # Replace the last token of every image tile with the newline token
         for input_ids in encoding.input_ids:
             i, n = 0, len(input_ids)
             while i < n:
@@ -138,13 +139,14 @@ class HeronProcessor(ProcessorMixin):
                 input_ids[i] = newline_token_id
                 i += 1
 
+        self._check_special_mm_tokens(processed_text, encoding, modalities=["image"])  # type: ignore[arg-type]
         return encoding
 
-    def _process_images(
+    def process_images(
         self,
         images: list[Image],
-        min_tiles: int,
-        max_tiles: int,
+        min_tiles: int = 1,
+        max_tiles: int = 12,
         **kwargs: Unpack[ImagesKwargs],
     ) -> BatchFeature:
         assert isinstance(min_tiles, int) and isinstance(max_tiles, int)
@@ -156,19 +158,38 @@ class HeronProcessor(ProcessorMixin):
         pixel_values = []
         for image in images:
             image_tiles = _dynamic_preprocess(
-                image.convert("RGB"), min_num=min_tiles, max_num=max_tiles, image_size=crop_size["height"]
+                image, min_num=min_tiles, max_num=max_tiles, image_size=crop_size["height"]
             )
-            pixel_values.append(self.image_processor(image_tiles, **kwargs)["pixel_values"])
+            pixel_values.append(self.image_processor(image_tiles, **kwargs, return_tensors="pt")["pixel_values"])
 
-        return BatchFeature({"pixel_values": pixel_values}, tensor_type=return_tensors)
+        return BatchFeature({"pixel_values": torch.stack(pixel_values)}, tensor_type=return_tensors)
 
 
-# Copied from https://github.com/NVlabs/VILA/blob/36f6adcd11a10be1580caeb7e647e1b6f8517f89/llava/mm_utils.py#L296
+# Adapted from https://github.com/NVlabs/VILA/blob/36f6adcd11a10be1580caeb7e647e1b6f8517f89/llava/mm_utils.py#L296
 def _dynamic_preprocess(
     image: Image, min_num: int, max_num: int, image_size: int, use_thumbnail: bool = True
 ) -> list[Image]:
+    if image.mode != "RGB":
+        image = image.convert("RGB")
+
     orig_width, orig_height = image.size
-    aspect_ratio = orig_width / orig_height
+    (target_width, target_height), crop_boxes = _calculate_crops(orig_width, orig_height, image_size, min_num, max_num)
+
+    resized_img = image.resize((target_width, target_height))
+    processed_images = [resized_img.crop(box) for box in crop_boxes]
+
+    if use_thumbnail and len(processed_images) != 1:
+        thumbnail_img = image.resize((image_size, image_size))
+        processed_images.append(thumbnail_img)
+
+    return processed_images
+
+
+@lru_cache(maxsize=32)
+def _calculate_crops(
+    width: int, height: int, crop_size: int, min_num: int, max_num: int
+) -> tuple[tuple[int, int], list[tuple[int, int, int, int]]]:
+    aspect_ratio = width / height
 
     # calculate the existing image aspect ratio
     target_ratio_set = {
@@ -176,7 +197,7 @@ def _dynamic_preprocess(
         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
+        if min_num <= i * j <= max_num
     }
     target_ratios = sorted(target_ratio_set, key=lambda x: x[0] * x[1])
 
@@ -184,34 +205,27 @@ def _dynamic_preprocess(
     target_aspect_ratio = _find_closest_aspect_ratio(
         aspect_ratio,
         target_ratios,
-        orig_width,
-        orig_height,
-        image_size,
+        width,
+        height,
+        crop_size,
     )
 
     # calculate the target width and height
-    target_width = image_size * target_aspect_ratio[0]
-    target_height = image_size * target_aspect_ratio[1]
+    target_width = crop_size * target_aspect_ratio[0]
+    target_height = crop_size * target_aspect_ratio[1]
     blocks = target_aspect_ratio[0] * target_aspect_ratio[1]
 
-    # resize the image
-    resized_img = image.resize((target_width, target_height))
-    processed_images = []
+    crop_boxes = []
     for i in range(blocks):
         box = (
-            (i % (target_width // image_size)) * image_size,
-            (i // (target_width // image_size)) * image_size,
-            ((i % (target_width // image_size)) + 1) * image_size,
-            ((i // (target_width // image_size)) + 1) * image_size,
+            (i % (target_width // crop_size)) * crop_size,
+            (i // (target_width // crop_size)) * crop_size,
+            ((i % (target_width // crop_size)) + 1) * crop_size,
+            ((i // (target_width // crop_size)) + 1) * crop_size,
         )
-        # split the image
-        split_img = resized_img.crop(box)
-        processed_images.append(split_img)
-    assert len(processed_images) == blocks
-    if use_thumbnail and len(processed_images) != 1:
-        thumbnail_img = image.resize((image_size, image_size))
-        processed_images.append(thumbnail_img)
-    return processed_images
+        crop_boxes.append(box)
+
+    return (target_width, target_height), crop_boxes
 
 
 # Copied from https://github.com/NVlabs/VILA/blob/36f6adcd11a10be1580caeb7e647e1b6f8517f89/llava/mm_utils.py#L280