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, 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,
-        resample: 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
@@ -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]
@@ -256,10 +309,8 @@ class MossVLProcessor(ProcessorMixin):
     r"""
     Constructs a Moss-VL processor which wraps a Qwen2VL image processor, Moss-VL video processor and a Qwen2 tokenizer
     into a single processor.
-
     [`MossVLProcessor`] offers all the functionalities of [`Qwen2VLImageProcessor`], [`MossVLVideoProcessor`] and [`Qwen2TokenizerFast`].
     See the [`~MossVLProcessor.__call__`] and [`~MossVLProcessor.decode`] for more information.
-
     Args:
         image_processor ([`Qwen2VLImageProcessor`], *optional*):
             The image processor is a required input.
@@ -272,8 +323,6 @@ class MossVLProcessor(ProcessorMixin):
     """
 
     attributes = ["image_processor", "tokenizer", "video_processor"]
-    image_processor_class = "AutoImageProcessor"
-    video_processor_class = "AutoVideoProcessor"
     tokenizer_class = ("Qwen2Tokenizer", "Qwen2TokenizerFast")
 
     def __init__(
@@ -338,7 +387,6 @@ class MossVLProcessor(ProcessorMixin):
     ) -> BatchFeature:
         """
         Main method to prepare for the model one or several sequences(s) and image(s)/video(s).
-
         Args:
             text (`str`, `list[str]`, `list[list[str]]`):
                 The sequence or batch of sequences to be encoded.
@@ -370,8 +418,6 @@ class MossVLProcessor(ProcessorMixin):
                 - `'pt'`: Return PyTorch `torch.Tensor` objects.
                 - `'np'`: Return NumPy `np.ndarray` objects.
                 - `'jax'`: Return JAX `jnp.ndarray` objects.
-
-
         Returns:
             [`BatchFeature`]: A [`BatchFeature`] with the following fields:
             - **input_ids** -- List of token ids to be fed to a model.
@@ -485,7 +531,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 +545,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 +763,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 +827,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)
@@ -1053,7 +1111,6 @@ class MossVLProcessor(ProcessorMixin):
     ):
         """
         Post-process the output of the model to decode the text.
-
         Args:
             generated_outputs (`torch.Tensor` or `np.ndarray`):
                 The output of the model `generate` function. The output is expected to be a tensor 
@@ -1064,7 +1121,6 @@ class MossVLProcessor(ProcessorMixin):
                 Whether or not to clean up the tokenization spaces.
             **kwargs:
                 Additional arguments to be passed to the tokenizer's `batch_decode` method.
-
         Returns:
             `list[str]`: The decoded text.
         """
@@ -1075,5 +1131,4 @@ class MossVLProcessor(ProcessorMixin):
             **kwargs,
         )
 
-
-__all__ = ["MossVLProcessor", "MossVLImageProcessorFast"]
+__all__ = ["MossVLProcessor", "MossVLImageProcessor"]