support Intern-S1 video (#5514)

* wip video

* support video

* test

* si video

* test

* fix test input

* annotation

Author: hjh0119

swift/llm/template/template/internvl.py

@@ -1,4 +1,5 @@
 # Copyright (c) Alibaba, Inc. and its affiliates.
+from ast import Tuple
 from functools import partial
 from typing import Any, Dict, List, Literal, Optional
 
@@ -183,49 +184,126 @@ class InternS1Template(Internvl2Template, ThinkingTemplate):
                                       'making your solution path and reasoning clear to others. '
                                       'Please put your thinking process within <think>...</think> tags.')
 
+    def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int,
+                    inputs: StdTemplateInputs) -> List[Context]:
+        assert media_type in ['image', 'video']
+        if media_type == 'video':
+            if self.mode == 'vllm':
+                return ['<video>']
+            else:
+                return [[-200]]
+        return super().replace_tag(media_type, index, inputs)
+
     def _swift_encode(self, inputs: StdTemplateInputs):
         if inputs.system is None and self.template_meta.response_prefix == '<think>':
             inputs.system = self.InternS1DefaultThinkinngSystem
 
         return super()._swift_encode(inputs)
 
     def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]:
-        from transformers.image_utils import make_flat_list_of_images
+        from transformers.image_utils import make_flat_list_of_images, concatenate_list
+        from transformers.video_utils import make_batched_videos
+        from swift.llm.template.vision_utils import load_video_hf
         import numpy as np
         encoded = super(InternvlTemplate, self)._encode(inputs)
         input_ids = encoded['input_ids']
-        idx_list = findall(input_ids, -100)
         labels = encoded['labels']
         loss_scale = encoded.get('loss_scale', None)
         images = inputs.images
-        if inputs.videos:
-            # TODO
-            raise NotImplementedError('Video is not supported yet.')
+        videos = inputs.videos
+        image_num_patches_indices = np.array([0])
+        video_num_patches_indices = np.array([0])
+        video_patch_indices = np.array([0])
+        image_num_patches = []
+        video_num_patches = []
+        image_video_patches = []
+        image_idx_list = []
+        video_idx_list = []
+        image_pixel_values = None
+        video_pixel_values = None
+
         if images:
             # InternS1Processor
+            image_idx_list = findall(input_ids, -100)
             images = make_flat_list_of_images(images)
             image_inputs = self.processor.image_processor(images=images, crop_to_patches=True, return_tensors='pt')
             image_num_patches = image_inputs.pop('num_patches')
-            pixel_values = image_inputs.pop('pixel_values')
+            image_pixel_values = image_inputs.pop('pixel_values')
             image_num_patches_indices = np.cumsum(image_num_patches)
-            # has_video = bool(inputs.videos) # TODO:video
-        else:
-            pixel_values = None
-            image_num_patches_indices = []
-        assert len(image_num_patches_indices) == len(
-            idx_list), f'len(num_patches): {len(num_patches)}, len(idx_list): {len(idx_list)}'
+        if videos:
+            video_idx_list = findall(input_ids, -200)
+            videos, _ = load_video_hf(videos)
+            videos = make_batched_videos(videos)
+            video_inputs = self.processor.video_processor(videos=videos, return_tensors='pt')
+            video_pixel_values = video_inputs.pop('pixel_values_videos')
+            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)
+            video_pixel_values = video_pixel_values.flatten(0, 1)
+
+        def merge_and_sort(image_idx_list: List[int], video_idx_list: List[int]) -> tuple:
+            """Merge and sort image and video index lists while preserving their relative order."""
+            merged = []
+            is_image_list = []
+            i, j = 0, 0
+
+            while i < len(image_idx_list) and j < len(video_idx_list):
+                if image_idx_list[i] < video_idx_list[j]:
+                    merged.append(image_idx_list[i])
+                    i += 1
+                    is_image_list.append(True)
+                else:
+                    merged.append(video_idx_list[j])
+                    j += 1
+                    is_image_list.append(False)
+            # Add remaining elements
+            merged.extend(image_idx_list[i:])
+            is_image_list.extend([True] * (len(image_idx_list) - i))
+            merged.extend(video_idx_list[j:])
+            is_image_list.extend([False] * (len(video_idx_list) - j))
+            return merged, is_image_list
+
+        # Merge and sort the index lists
+        idx_list, is_image_list = merge_and_sort(image_idx_list, video_idx_list)
+
+        # Validate the lengths
+        if images and len(image_idx_list) > 0:
+            assert len(image_num_patches_indices) == len(image_idx_list)
+        if videos and len(video_idx_list) > 0:
+            assert len(video_patch_indices) == len(video_idx_list)
 
         def _get_new_tokens(i):
-            start = image_num_patches_indices[i - 1] if i > 0 else 0
-            end = image_num_patches_indices[i]
-            image_seq_length = self.processor.image_seq_length
-            img_tokens: List[int] = self.processor.encode(
-                '<IMG_CONTEXT>', add_special_tokens=False) * image_seq_length * image_num_patches[start:end]
+            if is_image_list[i]:
+                # Find the corresponding image index
+                image_idx = sum(is_image_list[:i])
+                start = image_num_patches_indices[image_idx - 1] if image_idx > 0 else 0
+                end = image_num_patches_indices[image_idx]
+                image_seq_length = self.processor.image_seq_length
+                image_video_patches.append(image_pixel_values[start:end])
+                img_tokens: List[int] = self.processor.encode(
+                    '<IMG_CONTEXT>', add_special_tokens=False) * image_seq_length * image_num_patches[image_idx]
+            else:
+                # Find the corresponding video index
+                video_idx = i - sum(is_image_list[:i])
+                current_patch = video_patch_indices[video_idx - 1] if video_idx > 0 else 0
+                end_patch = video_patch_indices[video_idx]
+
+                start = video_num_patches_indices[current_patch] if video_idx > 0 else 0
+                end = video_num_patches_indices[end_patch - 1]
+                image_video_patches.append(video_pixel_values[start:end])
+                image_seq_length = self.processor.image_seq_length
+                num_patches = list(video_num_patches[current_patch:end_patch])
+                video_prompt = '\n'.join(
+                    f"Frame{i + 1}: <img>{'<IMG_CONTEXT>' * image_seq_length * num_patches[i]}</img>"
+                    for i in range(len(num_patches)))
+                img_tokens = self.processor.encode(video_prompt, add_special_tokens=False)
             return img_tokens
 
         encoded['input_ids'], encoded['labels'], encoded['loss_scale'] = self._extend_tokens(
             input_ids, labels, loss_scale, idx_list, _get_new_tokens)
-        encoded['pixel_values'] = pixel_values
+        if images or videos:
+            encoded['pixel_values'] = concatenate_list(image_video_patches)
         return encoded
 
     def _post_encode(self, model: nn.Module, inputs: Dict[str, Any]) -> Dict[str, Any]:
@@ -247,8 +325,6 @@ def _post_encode(self, model: nn.Module, inputs: Dict[str, Any]) -> Dict[str, An
                 pixel_values, vision_feature_layer=-1, vision_feature_select_strategy='default')
             image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype)
             inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_features)
-
-            inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_features)
         elif is_deepspeed_enabled():
             dummy_pixel_values = torch.zeros((1, 3, 32, 32), device=device, dtype=inputs_embeds.dtype)
             vit_embeds = model.model.vision_tower.embeddings(dummy_pixel_values)[0].to(device=device)

tests/test_align/test_template/test_video.py

@@ -190,6 +190,16 @@ def test_ovis2_5():
     print(f'response: {response}')
 
 
+def test_interns1():
+    pt_engine = PtEngine('Shanghai_AI_Laboratory/Intern-S1-mini')
+    messages = [{'role': 'user', 'content': '<video>Describe this video in detail.'}]
+    videos = ['https://modelscope-open.oss-cn-hangzhou.aliyuncs.com/images/baby.mp4']
+    response = _infer_model(pt_engine, messages=messages, videos=videos)
+    pt_engine.default_template.template_backend = 'jinja'
+    response2 = _infer_model(pt_engine, messages=messages, videos=videos)
+    assert response == response2
+
+
 if __name__ == '__main__':
     from swift.llm import PtEngine, RequestConfig
     from swift.utils import get_logger, seed_everything
@@ -207,4 +217,5 @@ def test_ovis2_5():
     # test_glm4_1v()  # bug now, wait model fix
     # test_keye_vl()
     # test_glm4_5v()
-    test_ovis2_5()
+    # test_ovis2_5()
+    test_interns1()