per-sample media budgets in processor; per-sample MRoPE position ids for batched vision inputs

processing_moss_vl.py

@@ -23,7 +23,7 @@ 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_utils import ImageInput, SizeDict, make_flat_list_of_images
 from transformers.image_processing_utils_fast import group_images_by_shape, reorder_images
 from transformers.utils import TensorType
 from transformers.processing_utils import (
@@ -384,6 +384,56 @@ class MossVLProcessor(ProcessorMixin):
         # Token ids that should always be masked in labels (e.g. <|image_pad|>)
         self.mask_token_ids = {self.image_token_id}
 
+
+    def _process_media_per_sample(self, media, texts, kind, processor, kwargs):
+        """Apply a complete media budget independently to each text sample.
+
+        Lists within one sample still share the image/video budget (and video
+        frame limit). A segmented video is one source item occupying several
+        video placeholders; it must not cross text-sample boundaries.
+        """
+        if len(texts) == 1:
+            return processor(**{kind: media}, **kwargs)
+        if kind == "images":
+            items = make_flat_list_of_images(media)
+            widths = [1] * len(items)
+            token = self.image_placeholder
+        else:
+            items = media if isinstance(media, list) else [media]
+            widths = [
+                len(item["segments"]) if isinstance(item, dict) and item.get("segments") else 1
+                for item in items
+            ]
+            token = self.video_placeholder
+        counts = [text.count(token) for text in texts]
+        if sum(counts) != sum(widths):
+            raise ValueError(
+                f"{kind} placeholders do not match supplied media: "
+                f"per-sample counts={counts}, media slots={sum(widths)}"
+            )
+        groups, offset = [], 0
+        for count in counts:
+            start, remaining = offset, count
+            while remaining:
+                width = widths[offset]
+                if width > remaining:
+                    raise ValueError("A segmented video cannot cross text-sample boundaries")
+                remaining -= width
+                offset += 1
+            groups.append(items[start:offset])
+        # Keep media order, but never share resize/frame budgets across samples.
+        results = [processor(**{kind: group}, **kwargs) for group in groups if group]
+        if not results:
+            return {}
+        combined = {}
+        for key in results[0]:
+            values = [result[key] for result in results]
+            if key == "video_metadata":
+                combined[key] = [item for value in values for item in value]
+            else:
+                combined[key] = _concat_array_or_tensor(values, axis=0)
+        return combined
+
     def __call__(
         self,
         text: Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]] = None,
@@ -445,12 +495,21 @@ class MossVLProcessor(ProcessorMixin):
             **kwargs,
         )
 
+        # Establish sample boundaries before media preprocessing: each sample
+        # receives its own complete image budget and video budget.
+        if text is None or (isinstance(text, str) and not text.strip()):
+            raise ValueError("Text input is required for MossVL processor and cannot be empty.")
+        if not isinstance(text, list):
+            text = [text]
+
         # Step 1: Process images if provided
         if images is not None:
             images_kwargs = output_kwargs["images_kwargs"].copy()
             images_kwargs["return_tensors"] = None
-            image_inputs = self.image_processor(images=images, **images_kwargs)
-            image_grid_thw = image_inputs["image_grid_thw"]
+            image_inputs = self._process_media_per_sample(
+                images, text, "images", self.image_processor, images_kwargs
+            )
+            image_grid_thw = image_inputs.get("image_grid_thw")
         else:
             image_inputs = {}
             image_grid_thw = None
@@ -459,25 +518,21 @@ class MossVLProcessor(ProcessorMixin):
         if videos is not None:
             videos_kwargs = output_kwargs["videos_kwargs"].copy()
             videos_kwargs["return_tensors"] = None
-            videos_inputs = self.video_processor(videos=videos, **videos_kwargs)
-            video_grid_thw = videos_inputs["video_grid_thw"]
+            videos_inputs = self._process_media_per_sample(
+                videos, text, "videos", self.video_processor, videos_kwargs
+            )
+            video_grid_thw = videos_inputs.get("video_grid_thw")
             # If user has not requested video metadata, pop it
             if "return_metadata" not in kwargs:
-                video_metadata = videos_inputs.pop("video_metadata")
+                video_metadata = videos_inputs.pop("video_metadata", [])
             else:
-                video_metadata = videos_inputs["video_metadata"]
+                video_metadata = videos_inputs.get("video_metadata", [])
         else:
             videos_inputs = {}
             video_grid_thw = None
             video_metadata = None
 
         # Step 3: Process text with placeholder replacement
-        if text is None or (isinstance(text, str) and len(text.strip()) == 0):
-            raise ValueError("Text input is required for MossVL processor and cannot be empty.")
-
-        if not isinstance(text, list):
-            text = [text]
-
         text = text.copy()  # Copy to avoid in-place modifications
 
         # Prepare labels_spans if provided