Fix binding of video frames to video placeholder in InternVL model (#41237)

* Fix binding video frames to video placeholder in prompt

Signed-off-by: Daniel Bershatsky <daniel.bershatsky@gmail.com>

* Add test on binding video frames to prompt

Signed-off-by: Daniel Bershatsky <daniel.bershatsky@gmail.com>

* Fix code style issues

Signed-off-by: Daniel Bershatsky <daniel.bershatsky@gmail.com>

* Fix broken tests on `InternVLProcessor`

Signed-off-by: Daniel Bershatsky <daniel.bershatsky@gmail.com>

* Add `return_tensors` to video processor defaults

Signed-off-by: Daniel Bershatsky <daniel.bershatsky@gmail.com>

---------

Signed-off-by: Daniel Bershatsky <daniel.bershatsky@gmail.com>

Author: daskol

src/transformers/models/internvl/processing_internvl.py

@@ -40,7 +40,9 @@ class InternVLProcessorKwargs(ProcessingKwargs, total=False):
         "images_kwargs": {
             "crop_to_patches": True,
         },
-        "videos_kwargs": {},
+        "videos_kwargs": {
+            "return_tensors": "pt",
+        },
     }
 
 
@@ -132,10 +134,10 @@ def _insert_media_placeholders(
                     # Get the slice of patches corresponding to the current video
                     # Here we need to account for both the multiple video frames and the potential multiple patches per frame
                     # As of now, InternVL only supports one patch per frame, but we keep the code flexible for future updates
-                    current_patch_index = video_patch_indices[video_index - 1] if video_index > 0 else 0
-                    end_patch_index = video_patch_indices[video_index]
-                    start_index = video_num_patches_indices[current_patch_index] if video_index > 0 else 0
-                    end_index = video_num_patches_indices[end_patch_index - 1]
+                    current_patch_index = video_patch_indices[video_index]
+                    end_patch_index = video_patch_indices[video_index + 1]
+                    start_index = video_num_patches_indices[current_patch_index]
+                    end_index = video_num_patches_indices[end_patch_index]
                     image_video_patches.append(video_pixel_values[start_index:end_index])
                     # Get the number of patches per frame and replace the video placeholder with the correct number of image tokens
                     num_patches = list(video_num_patches[current_patch_index:end_patch_index])
@@ -206,31 +208,38 @@ def __call__(
 
         # Process images and videos separately, as videos don't support crop_to_patches
         image_num_patches = []
-        video_num_patches = []
-        image_videos_inputs = {}
         image_pixel_values = None
-        video_pixel_values = None
         image_num_patches_indices = np.array([0])
-        video_patch_indices = np.array([0])
-        video_num_patches_indices = np.array([0])
         if images is not None:
             images = self.image_processor.fetch_images(images)
             images = make_flat_list_of_images(images)
             image_inputs = self.image_processor(images=images, **output_kwargs["images_kwargs"])
             image_num_patches = image_inputs.pop("num_patches")
             image_pixel_values = image_inputs.pop("pixel_values")
             image_num_patches_indices = np.cumsum(image_num_patches)
+
+        video_num_patches = []  # per frame
+        video_pixel_values = None
+        video_patch_indices = np.array([0])
+        video_num_patches_indices = np.array([0])
         if videos is not None:
-            video_inputs = self.video_processor(videos=videos, **output_kwargs["videos_kwargs"])
+            video_kwargs = output_kwargs["videos_kwargs"]
+            video_inputs = self.video_processor(videos=videos, **video_kwargs)
             video_pixel_values = video_inputs.pop("pixel_values_videos")
 
-            # Obtain per frame information first and then flatten to (BS * T, ...)
-            num_frames_per_video = [len(video) for video in video_pixel_values]
-            video_num_patches = [1 for frames in num_frames_per_video for _ in range(frames)]
-            video_patch_indices = np.cumsum(num_frames_per_video)
-            video_num_patches_indices = np.cumsum(video_num_patches)
+            batch_size, num_frames, *_ = video_pixel_values.shape
+            num_frames_per_video = np.full(batch_size, num_frames)
+            num_frames = sum(num_frames_per_video)  # total
+            video_patch_indices = np.empty(batch_size + 1, int)
+            video_patch_indices[0] = 0
+            video_patch_indices[1:] = np.cumsum(num_frames_per_video)
+            video_num_patches = [1] * num_frames
+            video_num_patches_indices = np.empty(num_frames + 1, int)
+            video_num_patches_indices[0] = 0
+            video_num_patches_indices[1:] = np.cumsum(video_num_patches)
             video_pixel_values = video_pixel_values.flatten(0, 1)
 
+        image_videos_inputs = {}
         if images is not None or videos is not None:
             text, image_video_patches, image_index, video_index = self._insert_media_placeholders(
                 text,

tests/models/internvl/test_processing_internvl.py

@@ -17,6 +17,8 @@
 import tempfile
 import unittest
 
+from parameterized import parameterized
+
 from transformers import AutoProcessor, AutoTokenizer, InternVLProcessor
 from transformers.testing_utils import require_av, require_torch, require_vision
 from transformers.utils import is_torch_available, is_vision_available
@@ -345,25 +347,30 @@ def _test_apply_chat_template(
         for idx, url in enumerate(input_data[:batch_size]):
             batch_messages[idx][0]["content"] = [batch_messages[idx][0]["content"][0], {"type": modality, "url": url}]
 
+        num_frames = 2  # by default no more than 2 frames, otherwise too slow
         out_dict = processor.apply_chat_template(
             batch_messages,
             add_generation_prompt=True,
             tokenize=True,
             return_dict=True,
             return_tensors="pt",
-            num_frames=2,  # by default no more than 2 frames, otherwise too slow
+            num_frames=num_frames,
         )
         self.assertTrue(self.videos_input_name in out_dict)
         self.assertEqual(len(out_dict["input_ids"]), batch_size)
         self.assertEqual(len(out_dict["attention_mask"]), batch_size)
 
         # InternVL internally collects frames from all the videos in a batch and flattens the batch dimension (B T C H W) -> (B*T C H W) then patches and removes the frames
         # hence output length does not equal batch size
-        # removed hardcoded video length check video_len = 2 if batch_size == 1 else 3
-        # from experiment video_len looks like batch_size + 1
-        # TODO: update expected video_len calculation based on the internal processing logic of InternVLProcessor
-        output_len = batch_size + 1 if modality == "video" else batch_size
-        self.assertEqual(len(out_dict[self.videos_input_name]), output_len)
+        num_pixel_planes = 0  # i.e. images + video frames
+        for message_thread in batch_messages:
+            for message in message_thread:
+                for content in message.get("content", []):
+                    if (content_type := content.get("type")) == "image":
+                        num_pixel_planes += 1
+                    elif content_type == "video":
+                        num_pixel_planes += num_frames
+        self.assertEqual(len(out_dict[self.videos_input_name]), num_pixel_planes)
         for k in out_dict:
             self.assertIsInstance(out_dict[k], torch.Tensor)
 
@@ -377,3 +384,25 @@ def _test_apply_chat_template(
         continue_prompt = processor.apply_chat_template(batch_messages, continue_final_message=True, tokenize=False)
         for prompt in continue_prompt:
             self.assertTrue(prompt.endswith("It is the sound of"))  # no `eos` token at the end
+
+    @parameterized.expand([(1,), (2,)])
+    @require_torch
+    def test_frames_binding(self, batch_size: int):
+        texts = [
+            "<video>\nAre there any cyan objects that enter the scene?\nno",
+            "<video>\nAre there any red spheres that enter the scene?\nno",
+        ]
+        frames = torch.ones((4, 448, 448, 3), dtype=torch.float32)
+        videos = [frames, frames]
+
+        processor = self.get_processor()
+        inputs = processor(
+            text=texts[:batch_size],
+            return_tensors="pt",
+            videos=videos[:batch_size],
+            videos_kwargs={"size": (448, 448)},
+        )
+
+        actual_num_frames = inputs.pixel_values.shape[0]
+        expected_num_frames = sum(x.shape[0] for x in videos[:batch_size])
+        assert actual_num_frames == expected_num_frames