Update code for transformers 5.5.4

processing_moss_vl.py

@@ -23,8 +23,8 @@ import torch
 from torchvision.transforms.v2 import functional as F
 from PIL import Image
 from transformers.feature_extraction_utils import BatchFeature
-from transformers.image_utils import ImageInput, SizeDict
-from transformers.image_processing_utils_fast import group_images_by_shape, reorder_images
+from transformers.image_utils import ImageInput, PILImageResampling, SizeDict
+from transformers.image_transforms import group_images_by_shape, reorder_images
 from transformers.utils import TensorType
 from transformers.processing_utils import (
     ImagesKwargs,
@@ -35,17 +35,16 @@ from transformers.processing_utils import (
 )
 from transformers.tokenization_utils_base import PreTokenizedInput, TextInput
 from transformers.utils import logging
-from transformers.models.qwen2_vl.image_processing_qwen2_vl_fast import Qwen2VLImageProcessorFast
-from transformers.models.qwen2_vl.image_processing_qwen2_vl import smart_resize
+from transformers.models.qwen2_vl.image_processing_qwen2_vl import Qwen2VLImageProcessor, smart_resize
 
 
 logger = logging.get_logger(__name__)
 
 
-class MossVLImageProcessorFast(Qwen2VLImageProcessorFast):
+class MossVLImageProcessor(Qwen2VLImageProcessor):
     """
     Custom image processor that overrides _preprocess to support multi_image_max_pixels.
-    Inherits from Qwen2VLImageProcessorFast.
+    Inherits from Qwen2VLImageProcessor.
     """
     # Multi-image batch total pixels limit (read from config)
     multi_image_max_pixels = None
@@ -56,7 +55,7 @@ class MossVLImageProcessorFast(Qwen2VLImageProcessorFast):
         images: list["torch.Tensor"],
         do_resize: bool,
         size: SizeDict,
-        interpolation: Optional["F.InterpolationMode"],
+        resample: Optional[Union["PILImageResampling", "F.InterpolationMode", int]],
         do_rescale: bool,
         rescale_factor: float,
         do_normalize: bool,
@@ -75,6 +74,8 @@ class MossVLImageProcessorFast(Qwen2VLImageProcessorFast):
         to each image based on its original pixel count. min_pixels remains a per-image
         constraint. multi_image_max_pixels can be configured separately from longest_edge.
         """
+        if resample is None:
+            resample = kwargs.pop("interpolation", None)
         min_pixels = size["shortest_edge"]
         max_pixels = size["longest_edge"]  # Per-image upper limit
         # Use multi_image_max_pixels if configured, otherwise fall back to longest_edge
@@ -115,7 +116,7 @@ class MossVLImageProcessorFast(Qwen2VLImageProcessorFast):
                 stacked_images = self.resize(
                     image=stacked_images,
                     size=SizeDict(height=resized_height, width=resized_width),
-                    interpolation=interpolation,
+                    interpolation=resample,
                 )
             resized_images_grouped[shape] = stacked_images
         resized_images = reorder_images(resized_images_grouped, grouped_images_index)
@@ -214,6 +215,58 @@ def _to_numpy(x):
     return np.array(x)
 
 
+def _split_array_or_tensor(x, split_indices):
+    """Split along the first dimension while preserving tensor/array type."""
+    split_indices = [int(idx) for idx in split_indices]
+    if isinstance(x, torch.Tensor):
+        if not split_indices:
+            return [x]
+        chunks = []
+        start = 0
+        for end in split_indices:
+            chunks.append(x[start:end])
+            start = end
+        chunks.append(x[start:])
+        return chunks
+    return np.split(x, split_indices)
+
+
+def _concat_array_or_tensor(items, axis=0):
+    """Concatenate while preserving tensor/array type and device."""
+    if not items:
+        return None
+
+    if any(isinstance(item, torch.Tensor) for item in items):
+        ref = next(item for item in items if isinstance(item, torch.Tensor))
+        tensor_items = [
+            item
+            if isinstance(item, torch.Tensor)
+            else torch.as_tensor(item, device=ref.device, dtype=ref.dtype)
+            for item in items
+        ]
+        return torch.cat(tensor_items, dim=axis)
+
+    return np.concatenate(items, axis=axis)
+
+
+def _stack_array_or_tensor(items, axis=0):
+    """Stack while preserving tensor/array type and device."""
+    if not items:
+        return None
+
+    if any(isinstance(item, torch.Tensor) for item in items):
+        ref = next(item for item in items if isinstance(item, torch.Tensor))
+        tensor_items = [
+            item
+            if isinstance(item, torch.Tensor)
+            else torch.as_tensor(item, device=ref.device, dtype=ref.dtype)
+            for item in items
+        ]
+        return torch.stack(tensor_items, dim=axis)
+
+    return np.stack(items, axis=axis)
+
+
 class MossVLImagesKwargs(ImagesKwargs):
     min_pixels: Optional[int]
     max_pixels: Optional[int]
@@ -272,8 +325,6 @@ class MossVLProcessor(ProcessorMixin):
     """
 
     attributes = ["image_processor", "tokenizer", "video_processor"]
-    image_processor_class = "AutoImageProcessor"
-    video_processor_class = "AutoVideoProcessor"
     tokenizer_class = ("Qwen2Tokenizer", "Qwen2TokenizerFast")
 
     def __init__(
@@ -485,7 +536,9 @@ class MossVLProcessor(ProcessorMixin):
                 elif len(patch_counts) > 1:
                     # Multiple images: split by cumulative counts
                     split_indices = np.cumsum(patch_counts)[:-1]
-                    image_pixel_values_list = np.split(flat_pixel_values, split_indices)
+                    image_pixel_values_list = _split_array_or_tensor(
+                        flat_pixel_values, split_indices
+                    )
 
             if has_videos:
                 flat_video_values = videos_inputs["pixel_values_videos"]
@@ -497,7 +550,9 @@ class MossVLProcessor(ProcessorMixin):
                 elif len(video_patch_counts) > 1:
                     # Multiple videos: split by cumulative counts
                     split_indices = np.cumsum(video_patch_counts)[:-1]
-                    video_pixel_values_list = np.split(flat_video_values, split_indices)
+                    video_pixel_values_list = _split_array_or_tensor(
+                        flat_video_values, split_indices
+                    )
         
         # Step 3.1: Replace placeholders (simple replacement, no expansion yet)
         # In MossVL, one image placeholder = one image token
@@ -713,10 +768,14 @@ class MossVLProcessor(ProcessorMixin):
             
             # Concatenate/stack to unified format
             if final_pixel_values:
-                output_data["pixel_values"] = np.concatenate(final_pixel_values, axis=0)
+                output_data["pixel_values"] = _concat_array_or_tensor(
+                    final_pixel_values, axis=0
+                )
             
             if final_grid_thw:
-                output_data["grid_thw"] = np.stack(final_grid_thw, axis=0)
+                output_data["grid_thw"] = _stack_array_or_tensor(
+                    final_grid_thw, axis=0
+                )
             
             # Don't add media_nums_per_sample to output_data yet
             # Will add it after BatchFeature to keep it as list
@@ -773,6 +832,10 @@ class MossVLProcessor(ProcessorMixin):
             for _ in range(num_media):
                 # grid_thw is (N, 3) where first dim is t (num_frames)
                 t = grid_thw[media_idx][0]
+                if isinstance(t, torch.Tensor):
+                    t = int(t.item())
+                else:
+                    t = int(t)
                 sample_frames += t
                 media_idx += 1
             total_frames_per_sample.append(sample_frames)
@@ -1075,5 +1138,4 @@ class MossVLProcessor(ProcessorMixin):
             **kwargs,
         )
 
-
-__all__ = ["MossVLProcessor", "MossVLImageProcessorFast"]
+__all__ = ["MossVLProcessor", "MossVLImageProcessor"]