Update processing_moss_vl.py

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, PILImageResampling, SizeDict
-from transformers.image_transforms import group_images_by_shape, reorder_images
+from transformers.image_utils import ImageInput, SizeDict
+from transformers.image_processing_utils_fast import group_images_by_shape, reorder_images
 from transformers.utils import TensorType
 from transformers.processing_utils import (
     ImagesKwargs,
@@ -35,16 +35,17 @@ 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 import Qwen2VLImageProcessor, smart_resize
+from transformers.models.qwen2_vl.image_processing_qwen2_vl_fast import Qwen2VLImageProcessorFast
+from transformers.models.qwen2_vl.image_processing_qwen2_vl import smart_resize
 
 
 logger = logging.get_logger(__name__)
 
 
-class MossVLImageProcessor(Qwen2VLImageProcessor):
+class MossVLImageProcessorFast(Qwen2VLImageProcessorFast):
     """
     Custom image processor that overrides _preprocess to support multi_image_max_pixels.
-    Inherits from Qwen2VLImageProcessor.
+    Inherits from Qwen2VLImageProcessorFast.
     """
     # Multi-image batch total pixels limit (read from config)
     multi_image_max_pixels = None
@@ -55,7 +56,7 @@ class MossVLImageProcessor(Qwen2VLImageProcessor):
         images: list["torch.Tensor"],
         do_resize: bool,
         size: SizeDict,
-        resample: Optional[Union["PILImageResampling", "F.InterpolationMode", int]],
+        resample: Optional["F.InterpolationMode"],
         do_rescale: bool,
         rescale_factor: float,
         do_normalize: bool,
@@ -74,8 +75,6 @@ class MossVLImageProcessor(Qwen2VLImageProcessor):
         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
@@ -215,58 +214,6 @@ 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]
@@ -325,6 +272,8 @@ class MossVLProcessor(ProcessorMixin):
     """
 
     attributes = ["image_processor", "tokenizer", "video_processor"]
+    image_processor_class = "AutoImageProcessor"
+    video_processor_class = "AutoVideoProcessor"
     tokenizer_class = ("Qwen2Tokenizer", "Qwen2TokenizerFast")
 
     def __init__(
@@ -536,9 +485,7 @@ 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 = _split_array_or_tensor(
-                        flat_pixel_values, split_indices
-                    )
+                    image_pixel_values_list = np.split(flat_pixel_values, split_indices)
 
             if has_videos:
                 flat_video_values = videos_inputs["pixel_values_videos"]
@@ -550,9 +497,7 @@ 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 = _split_array_or_tensor(
-                        flat_video_values, split_indices
-                    )
+                    video_pixel_values_list = np.split(flat_video_values, split_indices)
         
         # Step 3.1: Replace placeholders (simple replacement, no expansion yet)
         # In MossVL, one image placeholder = one image token
@@ -768,14 +713,10 @@ class MossVLProcessor(ProcessorMixin):
             
             # Concatenate/stack to unified format
             if final_pixel_values:
-                output_data["pixel_values"] = _concat_array_or_tensor(
-                    final_pixel_values, axis=0
-                )
+                output_data["pixel_values"] = np.concatenate(final_pixel_values, axis=0)
             
             if final_grid_thw:
-                output_data["grid_thw"] = _stack_array_or_tensor(
-                    final_grid_thw, axis=0
-                )
+                output_data["grid_thw"] = np.stack(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
@@ -832,10 +773,6 @@ 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)
@@ -1138,4 +1075,5 @@ class MossVLProcessor(ProcessorMixin):
             **kwargs,
         )
 
-__all__ = ["MossVLProcessor", "MossVLImageProcessor"]
+
+__all__ = ["MossVLProcessor", "MossVLImageProcessorFast"]