Replace processor with 0708 version (Fast image/video processor, transformers 4.57+); keep 0408 pixel caps

preprocessor_config.json

@@ -1,26 +1,26 @@
 {
-    "auto_map": {
-        "AutoProcessor": "processing_moss_vl.MossVLProcessor",
-        "AutoImageProcessor": "processing_moss_vl.MossVLImageProcessor"
-    },
-    "size": {
-        "longest_edge": 16777216,
-        "shortest_edge": 4096
-    },
-    "multi_image_max_pixels": 65536000,
-    "patch_size": 16,
-    "temporal_patch_size": 1,
-    "merge_size": 2,
-    "image_mean": [
-        0.5,
-        0.5,
-        0.5
-    ],
-    "image_std": [
-        0.5,
-        0.5,
-        0.5
-    ],
-    "processor_class": "MossVLProcessor",
-    "image_processor_type": "MossVLImageProcessor"
+  "auto_map": {
+    "AutoProcessor": "processing_moss_vl.MossVLProcessor",
+    "AutoImageProcessor": "processing_moss_vl.MossVLImageProcessorFast"
+  },
+  "size": {
+    "longest_edge": 16777216,
+    "shortest_edge": 4096
+  },
+  "multi_image_max_pixels": 65536000,
+  "patch_size": 16,
+  "temporal_patch_size": 1,
+  "merge_size": 2,
+  "image_mean": [
+    0.5,
+    0.5,
+    0.5
+  ],
+  "image_std": [
+    0.5,
+    0.5,
+    0.5
+  ],
+  "processor_class": "MossVLProcessor",
+  "image_processor_type": "MossVLImageProcessorFast"
 }

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,27 +35,28 @@ 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
-    
+
 
     def _preprocess(
         self,
         images: list["torch.Tensor"],
         do_resize: bool,
         size: SizeDict,
-        resample: Optional[Union["PILImageResampling", "F.InterpolationMode", int]],
+        interpolation: Optional["F.InterpolationMode"],
         do_rescale: bool,
         rescale_factor: float,
         do_normalize: bool,
@@ -69,18 +70,16 @@ class MossVLImageProcessor(Qwen2VLImageProcessor):
         **kwargs,
     ):
         """Override _preprocess to use custom smart_resize with batch-level max_pixels.
-        
+
         multi_image_max_pixels is treated as a batch-level total budget, proportionally allocated
         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
         multi_image_max_pixels = getattr(self, "multi_image_max_pixels", None) or max_pixels
-        
+
         # Calculate total original pixels across all images in the batch
         # This is used to proportionally allocate max_pixels to each image
         total_original_pixels = sum(img.shape[-2] * img.shape[-1] for img in images)
@@ -99,13 +98,13 @@ class MossVLImageProcessor(Qwen2VLImageProcessor):
                     proportional_max_pixels = int(multi_image_max_pixels * proportion)
                 else:
                     proportional_max_pixels = multi_image_max_pixels
-                
+
                 # Ensure proportional max_pixels is within [min_pixels, max_pixels] range
                 # min_pixels: per-image lower limit (shortest_edge)
                 # max_pixels: per-image upper limit (longest_edge)
                 proportional_max_pixels = max(proportional_max_pixels, min_pixels)
                 proportional_max_pixels = min(proportional_max_pixels, max_pixels)
-                
+
                 resized_height, resized_width = smart_resize(
                     height,
                     width,
@@ -116,11 +115,11 @@ class MossVLImageProcessor(Qwen2VLImageProcessor):
                 stacked_images = self.resize(
                     image=stacked_images,
                     size=SizeDict(height=resized_height, width=resized_width),
-                    resample=resample,
+                    interpolation=interpolation,
                 )
             resized_images_grouped[shape] = stacked_images
         resized_images = reorder_images(resized_images_grouped, grouped_images_index)
-        
+
         # Warn if multi-image batch exceeds multi_image_max_pixels due to min_pixels constraint
         if len(images) > 1:
             total_resized_pixels = sum(img.shape[-2] * img.shape[-1] for img in resized_images)
@@ -165,12 +164,17 @@ class MossVLImageProcessor(Qwen2VLImageProcessor):
             )
             # Reorder dimensions to group grid and patch information for subsequent flattening.
             # (batch, grid_t, grid_h, grid_w, merge_h, merge_w, channel, temp_patch_size, patch_h, patch_w)
+            # NPU ops support at most 8-D tensors; route the 10-D permute+reshape
+            # through CPU there. CUDA handles 10-D natively — keep it on-device.
+            patches_device = patches.device
+            if patches_device.type == "npu":
+                patches = patches.cpu()
             patches = patches.permute(0, 1, 4, 7, 5, 8, 3, 2, 6, 9)
             flatten_patches = patches.reshape(
                 batch_size,
                 grid_t * grid_h * grid_w,
                 channel * temporal_patch_size * patch_size * patch_size,
-            )
+            ).to(patches_device)
 
             processed_images_grouped[shape] = flatten_patches
             processed_grids[shape] = [[grid_t, grid_h, grid_w]] * batch_size
@@ -188,17 +192,17 @@ def _to_numpy(x):
     """
     Convert various tensor types to numpy array.
     Supports torch.Tensor, tf.Tensor, jax.Array, np.ndarray, lists, and primitives.
-    
+
     Args:
         x: Input value that can be a tensor from various frameworks or a Python primitive
-        
+
     Returns:
         np.ndarray: NumPy array representation of the input
     """
     # Already numpy
     if isinstance(x, np.ndarray):
         return x
-    
+
     # Torch tensor or TensorFlow tensor (both have .numpy() method)
     if hasattr(x, 'numpy'):
         # For torch tensors on CUDA, need to move to CPU first
@@ -206,11 +210,11 @@ def _to_numpy(x):
             return x.cpu().numpy()
         # For TensorFlow or already on CPU
         return x.numpy()
-    
+
     # JAX arrays and other array-like objects that support __array__ protocol
     if hasattr(x, '__array__'):
         return np.asarray(x)
-    
+
     # Python primitives (list, tuple, int, float)
     return np.array(x)
 
@@ -325,23 +329,25 @@ class MossVLProcessor(ProcessorMixin):
     """
 
     attributes = ["image_processor", "tokenizer", "video_processor"]
+    image_processor_class = "AutoImageProcessor"
+    video_processor_class = "AutoVideoProcessor"
     tokenizer_class = ("Qwen2Tokenizer", "Qwen2TokenizerFast")
 
     def __init__(
-        self, 
-        image_processor=None, 
-        tokenizer=None, 
-        video_processor=None, 
+        self,
+        image_processor=None,
+        tokenizer=None,
+        video_processor=None,
         chat_template=None,
         **kwargs
     ):
         super().__init__(image_processor, tokenizer, video_processor, chat_template=chat_template)
-        
+
 
         self.image_token = "<|image_pad|>" if not hasattr(tokenizer, "image_token") else tokenizer.image_token
         self.video_token = "<|video_pad|>" if not hasattr(tokenizer, "video_token") else tokenizer.video_token
 
-        
+
         self.image_token_id = (
             tokenizer.image_token_id
             if getattr(tokenizer, "image_token_id", None)
@@ -352,7 +358,7 @@ class MossVLProcessor(ProcessorMixin):
             if getattr(tokenizer, "video_token_id", None)
             else tokenizer.convert_tokens_to_ids(self.video_token)
         )
-        
+
         self.vision_start_token = (
             "<|vision_start|>" if not hasattr(tokenizer, "vision_start_token") else tokenizer.vision_start_token
         )
@@ -366,15 +372,15 @@ class MossVLProcessor(ProcessorMixin):
 
         self.time_start_token = "<|time_start|>"
         self.time_end_token = "<|time_end|>"
-        
+
         # EOS token for labels generation (assistant's response should end with this)
         self.im_end_token = "<|im_end|>"
         self.im_end_token_id = tokenizer.convert_tokens_to_ids(self.im_end_token)
-        
+
         # Vision-related token ids (all should be masked in labels)
         self.vision_start_token_id = tokenizer.convert_tokens_to_ids(self.vision_start_token)
         self.vision_end_token_id = tokenizer.convert_tokens_to_ids(self.vision_end_token)
-        
+
         # Token ids that should always be masked in labels (e.g. <|image_pad|>)
         self.mask_token_ids = {self.image_token_id}
 
@@ -438,7 +444,7 @@ class MossVLProcessor(ProcessorMixin):
             tokenizer_init_kwargs=self.tokenizer.init_kwargs,
             **kwargs,
         )
-        
+
         # Step 1: Process images if provided
         if images is not None:
             images_kwargs = output_kwargs["images_kwargs"].copy()
@@ -468,12 +474,12 @@ class MossVLProcessor(ProcessorMixin):
         # 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
         # labels_spans format: List[List[List[int]]] - batch of samples, each sample has multiple spans
         # Each span is [start, end] (list, not tuple) so it can be modified in place
@@ -490,10 +496,10 @@ class MossVLProcessor(ProcessorMixin):
         has_images = images is not None and "pixel_values" in image_inputs
         has_videos = videos is not None and "pixel_values_videos" in videos_inputs
         needs_reorder = has_images and has_videos
-        
+
         image_pixel_values_list = []
         video_pixel_values_list = []
-        
+
         # Step 3.0: Record the order of media in original text (before replacement)
         # This will be used later to correctly order pixel_values and grid_thw
         media_order_per_sample = []
@@ -504,25 +510,25 @@ class MossVLProcessor(ProcessorMixin):
             while pos < len(temp_text):
                 img_pos = temp_text.find(self.image_placeholder, pos)
                 vid_pos = temp_text.find(self.video_placeholder, pos)
-                
+
                 if img_pos == -1 and vid_pos == -1:
                     break
-                
+
                 if img_pos != -1 and (vid_pos == -1 or img_pos < vid_pos):
                     media_order.append(("image", img_pos))
                     pos = img_pos + len(self.image_placeholder)
                 elif vid_pos != -1:
                     media_order.append(("video", vid_pos))
                     pos = vid_pos + len(self.video_placeholder)
-            
+
             media_order_per_sample.append(media_order)
-        
+
         # Step 3.0.1: Check if any sample has no media (empty samples need blank image)
         # If there are empty samples, we need to enter slow path to handle them properly
         has_empty_samples = any(len(order) == 0 for order in media_order_per_sample)
         if has_empty_samples:
             needs_reorder = True
-        
+
         # Split pixel values for reordering if needed
         if needs_reorder:
             if has_images:
@@ -553,7 +559,7 @@ class MossVLProcessor(ProcessorMixin):
                     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
         # One video placeholder = one video token (will be expanded later)
@@ -570,33 +576,33 @@ class MossVLProcessor(ProcessorMixin):
             else:
                 text[i] = text[i].replace(self.image_placeholder, self.image_token)
                 text[i] = text[i].replace(self.video_placeholder, self.video_token)
-        
-        # Step 3.2: Validate token counts 
+
+        # Step 3.2: Validate token counts
         n_images_in_text = [t.count(self.image_token) for t in text]
         n_videos_in_text = [t.count(self.video_token) for t in text]
-        
+
         # Count placeholders in text
         total_images_in_text = sum(n_images_in_text)
         total_videos_in_text = sum(n_videos_in_text)
-        
+
         # Count actual images and videos provided
         total_images_provided = len(image_grid_thw) if image_grid_thw is not None else 0
         total_videos_provided = len(video_grid_thw) if video_grid_thw is not None else 0
-        
+
         # Validate image counts
         if total_images_in_text != total_images_provided:
             raise ValueError(
                 "Number of image tokens does not match number of images provided. "
                 f"Found {total_images_in_text} image tokens in text and {total_images_provided} images."
             )
-        
+
         # Validate video counts
         if total_videos_in_text != total_videos_provided:
             raise ValueError(
                 "Number of video tokens does not match number of videos provided. "
                 f"Found {total_videos_in_text} video tokens in text and {total_videos_provided} videos."
             )
-        
+
         # Step 3.3: Expand video tokens with timestamps
         # Now expand each video token to multiple tokens (one per frame) with timestamps
         if video_grid_thw is not None:
@@ -634,10 +640,10 @@ class MossVLProcessor(ProcessorMixin):
                         video_tokens.append(
                             f"{self.time_start_token}{curr_time:.1f} seconds{self.time_end_token}{self.image_token}"
                         )
-                    
+
                     # Wrap the entire video sequence with vision_start and vision_end tokens
                     video_placeholder = f"{self.vision_start_token}{''.join(video_tokens)}{self.vision_end_token}"
-                    
+
                     # Replace the video token with expanded sequence and update spans if needed
                     if should_create_labels:
                         text[i], labels_spans[i] = self._replace_and_update_spans(
@@ -646,22 +652,22 @@ class MossVLProcessor(ProcessorMixin):
                     else:
                         text[i] = text[i].replace(self.video_token, video_placeholder, 1)
                     index += 1
-                
+
 
 
         # Step 4: Tokenize text
         return_tensors = output_kwargs["text_kwargs"].pop("return_tensors", None)
         return_mm_token_type_ids = output_kwargs["text_kwargs"].pop("return_mm_token_type_ids", None)
-        
+
         # Request offset_mapping if we need to create labels
         if should_create_labels:
             output_kwargs["text_kwargs"]["return_offsets_mapping"] = True
-        
+
         text_inputs = self.tokenizer(text, **output_kwargs["text_kwargs"])
-        
+
         # ignore check_special_mm_tokens nums in test and input ids.
         # self._check_special_mm_tokens(text, text_inputs, modalities=["image", "video"])
-        
+
         # Create labels if labels_spans was provided
         if should_create_labels:
             offset_mapping = text_inputs.pop("offset_mapping")
@@ -681,49 +687,49 @@ class MossVLProcessor(ProcessorMixin):
         # Step 5: Concatenate pixel_values and grid_thw in sequence order
         # Prepare output
         output_data = {**text_inputs}
-        
+
         if not needs_reorder:
             # Fast path: only one media type, no reordering needed
             final_pixel_values = []
             final_grid_thw = []
-            
+
             if has_images:
                 final_pixel_values.append(image_inputs["pixel_values"])
                 final_grid_thw.extend(image_grid_thw)
-            
+
             if has_videos:
                 final_pixel_values.append(videos_inputs["pixel_values_videos"])
                 final_grid_thw.extend(video_grid_thw)
-            
+
             if final_pixel_values:
                 output_data["pixel_values"] = np.concatenate(final_pixel_values, axis=0) if len(final_pixel_values) > 1 else final_pixel_values[0]
-            
+
             if final_grid_thw:
                 output_data["grid_thw"] = np.stack(final_grid_thw, axis=0)
-            
+
             # Calculate media_nums_per_sample
             media_nums_per_sample = []
             for batch_idx in range(len(text)):
                 media_order = media_order_per_sample[batch_idx]
                 media_nums_per_sample.append(len(media_order) if len(media_order) > 0 else 1)
-            
+
             # Don't add media_nums_per_sample to output_data yet
             # Will add it after BatchFeature to keep it as list
-            
+
         else:
             # Slow path: both images and videos exist, need reordering
             final_pixel_values = []
             final_grid_thw = []
             media_nums_per_sample = []
-            
+
             # Global indices to track position in flattened image/video arrays
             global_image_idx = 0
             global_video_idx = 0
-            
+
             for batch_idx in range(len(text)):
                 # Use the recorded media order from Step 3.0
                 media_order = media_order_per_sample[batch_idx]
-                
+
                 if len(media_order) == 0:
                     # If no media provided for this sample, add a blank image
                     media_nums_per_sample.append(1)
@@ -731,26 +737,26 @@ class MossVLProcessor(ProcessorMixin):
                     patch_size = getattr(self.image_processor, "patch_size", None) or 16
                     temporal_patch_size = getattr(self.image_processor, "temporal_patch_size", None) or 1
                     merge_size = getattr(self.image_processor, "merge_size", None) or 2
-                    
+
                     factor = patch_size * merge_size
                     side = int(np.ceil(np.sqrt(min_pixels) / factor) * factor)
                     grid_h = side // patch_size
                     grid_w = side // patch_size
                     grid_t = 1
-                    
+
                     # Channel = 3 (RGB)
                     channel = 3
                     dim = channel * temporal_patch_size * patch_size * patch_size
                     num_patches = grid_t * grid_h * grid_w
-                    
+
                     blank_pixel_values = np.zeros((num_patches, dim), dtype=np.float32)
                     blank_grid_thw = np.array([grid_t, grid_h, grid_w], dtype=np.int64)
-                    
+
                     final_pixel_values.append(blank_pixel_values)
                     final_grid_thw.append(blank_grid_thw)
                 else:
                     media_nums_per_sample.append(len(media_order))
-                    
+
                     # Collect media data according to the recorded order
                     for media_type, _ in media_order:
                         if media_type == "image" and image_grid_thw is not None:
@@ -765,18 +771,18 @@ class MossVLProcessor(ProcessorMixin):
                                 final_pixel_values.append(video_pixel_values_list[global_video_idx])
                             final_grid_thw.append(video_grid_thw[global_video_idx])
                             global_video_idx += 1
-            
+
             # Concatenate/stack to unified format
             if final_pixel_values:
                 output_data["pixel_values"] = _concat_array_or_tensor(
                     final_pixel_values, axis=0
                 )
-            
+
             if final_grid_thw:
                 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
 
@@ -789,18 +795,18 @@ class MossVLProcessor(ProcessorMixin):
                 output_data.get("attention_mask", None)
             )
             output_data["cross_attention_mask"] = cross_attention_mask
-        
+
         # Add labels to output if created
         if should_create_labels:
             output_data["labels"] = labels
 
         # BatchFeature will handle conversion to pt/tf/jax/np based on tensor_type
         batch_feature = BatchFeature(data=output_data, tensor_type=return_tensors)
-        
+
         # Add media_nums_per_sample after BatchFeature to keep it as list (not tensor)
         if media_nums_per_sample:
             batch_feature["media_nums_per_sample"] = media_nums_per_sample
-        
+
         return batch_feature
 
     def _create_cross_attention_mask(self, input_ids, grid_thw, media_nums_per_sample, attention_mask=None):
@@ -809,7 +815,7 @@ class MossVLProcessor(ProcessorMixin):
         Video frames are treated as individual images.
         Mask values: True for masked, False for visible.
         Causal masking: text can see images that appear at or before the text position.
-        
+
         Args:
             input_ids: List of token ids
             grid_thw: Grid sizes for each media item
@@ -818,7 +824,7 @@ class MossVLProcessor(ProcessorMixin):
         """
         batch_size = len(input_ids)
         max_text_len = max(len(ids) for ids in input_ids)
-        
+
         # Calculate total frames per sample to find max_num_frames
         total_frames_per_sample = []
         media_idx = 0
@@ -827,7 +833,7 @@ class MossVLProcessor(ProcessorMixin):
             if num_media == 0:
                 total_frames_per_sample.append(0)
                 continue
-                
+
             sample_frames = 0
             for _ in range(num_media):
                 # grid_thw is (N, 3) where first dim is t (num_frames)
@@ -839,38 +845,38 @@ class MossVLProcessor(ProcessorMixin):
                 sample_frames += t
                 media_idx += 1
             total_frames_per_sample.append(sample_frames)
-            
+
         max_num_frames = max(total_frames_per_sample) if total_frames_per_sample else 0
-        
+
         if max_num_frames == 0:
             return None
-            
+
         # Vectorized implementation for speed
-        
+
         # 1. Pad input_ids to create a tensor
         # We use -1 as pad value since token ids are positive
         input_ids_tensor = torch.full((batch_size, max_text_len), -1, dtype=torch.long)
         for b, ids in enumerate(input_ids):
             l = len(ids)
             input_ids_tensor[b, :l] = torch.tensor(ids, dtype=torch.long)
-            
+
         # 2. Identify image tokens
         is_image_token = (input_ids_tensor == self.image_token_id)
-        
+
         # 3. Compute cumulative image tokens (how many image tokens appeared up to position t)
         # shape: (batch_size, text_len)
         cum_image_tokens = is_image_token.cumsum(dim=1)
-        
+
         # 4. Create frame indices
         # shape: (1, 1, max_num_frames)
         frame_indices = torch.arange(max_num_frames).reshape(1, 1, -1)
-        
+
         # 5. Determine visibility based on causal relationship
         # Text at `t` sees frame `i` if `cum_image_tokens[t] > i`
         # Because if frame `i` is the (i+1)-th image token, it becomes visible when count reaches i+1
         # shape: (batch_size, text_len, max_num_frames)
         visible_mask = cum_image_tokens.unsqueeze(-1) > frame_indices
-        
+
         # 6. Apply attention_mask if provided
         if attention_mask is not None:
             # Convert to tensor if needed
@@ -882,25 +888,25 @@ class MossVLProcessor(ProcessorMixin):
                  for b, mask_row in enumerate(attention_mask):
                      l = len(mask_row)
                      attn_mask_tensor[b, :l] = torch.tensor(mask_row, dtype=torch.long)
-            
+
             # shape: (batch_size, text_len, 1)
             valid_text = (attn_mask_tensor.unsqueeze(-1) == 1)
             visible_mask = visible_mask & valid_text
-            
+
         # 7. Mask out frames that don't exist for a sample
         # shape: (batch_size, 1, 1)
         total_frames_tensor = torch.tensor(total_frames_per_sample).reshape(batch_size, 1, 1)
         # shape: (batch_size, 1, max_num_frames)
         valid_frames = frame_indices < total_frames_tensor
-        
+
         visible_mask = visible_mask & valid_frames
-        
+
         # 8. Create final mask (True for masked, False for visible)
         mask = ~visible_mask
-        
+
         # 9. Add channel dimension: (batch_size, 1, text_len, max_num_frames)
         mask = mask.unsqueeze(1)
-        
+
         return mask
 
     def _replace_and_update_spans(
@@ -913,14 +919,14 @@ class MossVLProcessor(ProcessorMixin):
     ) -> tuple:
         """
         Replace occurrences of old_str with new_str and update spans accordingly.
-        
+
         Args:
             text: The text to perform replacement on
             old_str: String to be replaced
             new_str: String to replace with
             spans: List of [start, end] spans to update (modified in place)
             replace_count: Maximum number of replacements (-1 for all)
-            
+
         Returns:
             Tuple of (new_text, updated_spans)
         """
@@ -928,14 +934,14 @@ class MossVLProcessor(ProcessorMixin):
         result_text = text
         count = 0
         search_start = 0
-        
+
         while True:
             pos = result_text.find(old_str, search_start)
             if pos == -1:
                 break
             if replace_count != -1 and count >= replace_count:
                 break
-            
+
             # Update all spans that come after this position
             for span in spans:
                 if span[0] > pos:
@@ -945,12 +951,12 @@ class MossVLProcessor(ProcessorMixin):
                 elif span[1] > pos:
                     # Span ends after replacement point (spans the replacement)
                     span[1] += delta
-            
+
             # Perform the replacement
             result_text = result_text[:pos] + new_str + result_text[pos + len(old_str):]
             search_start = pos + len(new_str)
             count += 1
-        
+
         return result_text, spans
 
     def _create_labels_from_spans(
@@ -963,7 +969,7 @@ class MossVLProcessor(ProcessorMixin):
     ) -> List[List[int]]:
         """
         Create labels from spans and offset_mapping.
-        
+
         Args:
             input_ids: Tokenized input ids
             offset_mapping: Character offsets for each token from tokenizer (special tokens included)
@@ -972,10 +978,10 @@ class MossVLProcessor(ProcessorMixin):
             mask_token_ids: Set of token ids that should always be masked (set to ignore_index)
                 in labels, regardless of whether they fall inside a span.
                 Defaults to self.mask_token_ids if not provided.
-            
+
         Returns:
             labels: List of label ids, same shape as input_ids
-        
+
         Note:
             - Tokenizer's offset_mapping already includes correct offsets for special tokens in text
             - Only need to mask tokens inside <|vision_start|>...<|vision_end|>
@@ -984,20 +990,20 @@ class MossVLProcessor(ProcessorMixin):
         """
         if mask_token_ids is None:
             mask_token_ids = self.mask_token_ids
-        
+
         batch_labels = []
-        
+
         for batch_idx in range(len(input_ids)):
             ids = input_ids[batch_idx]
             offsets = offset_mapping[batch_idx]
             spans = labels_spans[batch_idx]
-            
+
             labels = [ignore_index] * len(ids)
-            
+
             # Process each span: find token range and set labels
             for span_start, span_end in spans:
                 in_vision = False
-                
+
                 # Find tokens that overlap with this span
                 for token_idx, (token_id, (char_start, char_end)) in enumerate(zip(ids, offsets)):
                     # Skip tokens completely before this span
@@ -1006,7 +1012,7 @@ class MossVLProcessor(ProcessorMixin):
                     # Stop when tokens are completely after this span
                     if char_start >= span_end:
                         break
-                    
+
                     # Token overlaps with span, process it
                     # Track vision region: <|vision_start|> ... <|vision_end|>
                     if token_id == self.vision_start_token_id:
@@ -1015,34 +1021,34 @@ class MossVLProcessor(ProcessorMixin):
                     if token_id == self.vision_end_token_id:
                         in_vision = False
                         continue
-                    
+
                     # Skip tokens inside vision region
                     if in_vision:
                         continue
-                    
+
                     # Always mask special tokens that should never have labels
                     if token_id in mask_token_ids:
                         continue
-                    
+
                     # Set label for this token
                     labels[token_idx] = token_id
-            
+
             batch_labels.append(labels)
-        
+
         return batch_labels
 
     def _calculate_timestamps(
-        self, 
-        frames_indices: Optional[Union[List[int], np.ndarray]], 
-        total_num_frames: int, 
-        video_fps: float, 
-        duration: float, 
+        self,
+        frames_indices: Optional[Union[List[int], np.ndarray]],
+        total_num_frames: int,
+        video_fps: float,
+        duration: float,
         merge_size: int = 1,
         actual_timestamps: Optional[List[float]] = None
     ):
         """
         Calculate timestamps for video frames.
-        
+
         Args:
             frames_indices: Actual frame indices extracted (if available)
             total_num_frames: Total number of sampled frames
@@ -1050,25 +1056,25 @@ class MossVLProcessor(ProcessorMixin):
             duration: Video duration in seconds
             merge_size: Temporal merge size
             actual_timestamps: Pre-calculated actual timestamps (for segments)
-            
+
         Returns:
             List of timestamps (one per merged temporal patch)
         """
         # If actual timestamps are provided (from segment), use them directly
         if actual_timestamps is not None:
             timestamps = list(actual_timestamps)
-            
+
             # Pad timestamps to be multiple of merge_size
             if len(timestamps) % merge_size != 0:
                 timestamps.extend([timestamps[-1]] * (merge_size - len(timestamps) % merge_size))
-            
+
             # Frames are merged by merge_size, so we average the timestamps within each temporal patch
             timestamps = [
-                (timestamps[i] + timestamps[i + merge_size - 1]) / 2 
+                (timestamps[i] + timestamps[i + merge_size - 1]) / 2
                 for i in range(0, len(timestamps), merge_size)
             ]
             return timestamps
-        
+
         # Use frames_indices if available, otherwise generate uniformly sampled indices
         if frames_indices is not None:
             if isinstance(frames_indices, np.ndarray):
@@ -1082,24 +1088,24 @@ class MossVLProcessor(ProcessorMixin):
             else:
                 # Uniformly sample frames across the video duration
                 indices = np.linspace(0, duration * video_fps - 1, total_num_frames).astype(np.int32).tolist()
-        
+
         # Pad indices to be multiple of merge_size
         if len(indices) % merge_size != 0:
             indices.extend([indices[-1]] * (merge_size - len(indices) % merge_size))
-        
+
         # Convert frame indices to timestamps
         timestamps = [idx / video_fps for idx in indices]
-        
+
         # Frames are merged by merge_size, so we average the timestamps within each temporal patch
         timestamps = [
-            (timestamps[i] + timestamps[i + merge_size - 1]) / 2 
+            (timestamps[i] + timestamps[i + merge_size - 1]) / 2
             for i in range(0, len(timestamps), merge_size)
         ]
         return timestamps
 
     def batch_decode(self, *args, **kwargs):
         """
-        This method forwards all its arguments to the tokenizer's batch_decode. 
+        This method forwards all its arguments to the tokenizer's batch_decode.
         Please refer to the docstring of this method for more information.
         """
         return self.tokenizer.batch_decode(*args, **kwargs)
@@ -1119,7 +1125,7 @@ class MossVLProcessor(ProcessorMixin):
 
         Args:
             generated_outputs (`torch.Tensor` or `np.ndarray`):
-                The output of the model `generate` function. The output is expected to be a tensor 
+                The output of the model `generate` function. The output is expected to be a tensor
                 of shape `(batch_size, sequence_length)` or `(sequence_length,)`.
             skip_special_tokens (`bool`, *optional*, defaults to `True`):
                 Whether or not to remove special tokens in the output.
@@ -1138,4 +1144,5 @@ class MossVLProcessor(ProcessorMixin):
             **kwargs,
         )
 
-__all__ = ["MossVLProcessor", "MossVLImageProcessor"]
+
+__all__ = ["MossVLProcessor", "MossVLImageProcessorFast"]

video_preprocessor_config.json

@@ -1,30 +1,30 @@
 {
-    "auto_map": {
-        "AutoProcessor": "processing_moss_vl.MossVLProcessor",
-        "AutoVideoProcessor": "video_processing_moss_vl.MossVLVideoProcessor"
-    },
-    "size": {
-        "longest_edge": 16777216,
-        "shortest_edge": 4096
-    },
-    "video_max_pixels": 65536000,
-    "patch_size": 16,
-    "temporal_patch_size": 1,
-    "merge_size": 2,
-    "video_fps": 1.0,
-    "min_frames": 1,
-    "max_frames": 256,
-    "num_extract_threads": 4,
-    "image_mean": [
-        0.5,
-        0.5,
-        0.5
-    ],
-    "image_std": [
-        0.5,
-        0.5,
-        0.5
-    ],
-    "processor_class": "MossVLProcessor",
-    "video_processor_type": "MossVLVideoProcessor"
+  "auto_map": {
+    "AutoProcessor": "processing_moss_vl.MossVLProcessor",
+    "AutoVideoProcessor": "video_processing_moss_vl.MossVLVideoProcessor"
+  },
+  "size": {
+    "longest_edge": 16777216,
+    "shortest_edge": 4096
+  },
+  "video_max_pixels": 65536000,
+  "patch_size": 16,
+  "temporal_patch_size": 1,
+  "merge_size": 2,
+  "video_fps": 1.0,
+  "min_frames": 1,
+  "max_frames": 256,
+  "num_extract_threads": 4,
+  "image_mean": [
+    0.5,
+    0.5,
+    0.5
+  ],
+  "image_std": [
+    0.5,
+    0.5,
+    0.5
+  ],
+  "processor_class": "MossVLProcessor",
+  "video_processor_type": "MossVLVideoProcessor"
 }

video_processing_moss_vl.py

@@ -14,13 +14,17 @@
 # limitations under the License.
 """video processor class for Moss-VL."""
 
+import fcntl
+import hashlib
 import json
 import logging as system_logging
 import math
 import os
 import re
 import subprocess
+import tempfile
 import traceback
+from functools import lru_cache
 from typing import Any, Dict, List, Optional, Union
 
 import numpy as np
@@ -38,6 +42,8 @@ from transformers.video_utils import VideoMetadata, group_videos_by_shape, reord
 
 logger = logging.get_logger(__name__)
 
+TORCHCODEC_TIMESTAMP_EPSILON = 1e-6
+
 
 # -----------------------------------------------------------------------------
 # Torchcodec video frame extraction utilities
@@ -46,10 +52,10 @@ logger = logging.get_logger(__name__)
 def check_video_for_extra_streams_and_errors(video_path: str) -> dict:
     """
     Check if video file has abnormal streams or errors reported by ffprobe.
-    
+
     Args:
         video_path: Path to the video file.
-        
+
     Returns:
         A dictionary containing:
         - 'has_extra_streams': bool, whether there are streams other than video and audio.
@@ -67,7 +73,7 @@ def check_video_for_extra_streams_and_errors(video_path: str) -> dict:
         'stream_details': [],
         'num_streams': 0
     }
-    
+
     command = [
         "ffprobe",
         "-v", "error",
@@ -76,7 +82,7 @@ def check_video_for_extra_streams_and_errors(video_path: str) -> dict:
         "-of", "json",
         video_path
     ]
-    
+
     try:
         process = subprocess.run(
             command,
@@ -85,12 +91,12 @@ def check_video_for_extra_streams_and_errors(video_path: str) -> dict:
             check=False
         )
         result['ffprobe_successful'] = (process.returncode == 0)
-        
+
         if process.stderr:
             result['ffprobe_output_error'] = process.stderr
             unsupported_codec_pattern = re.compile(r"Unsupported codec with id \d+ for input stream \d+")
             result['unsupported_codec_errors'] = unsupported_codec_pattern.findall(process.stderr)
-        
+
         if process.stdout:
             ffprobe_data = json.loads(process.stdout)
             if 'streams' in ffprobe_data:
@@ -101,7 +107,7 @@ def check_video_for_extra_streams_and_errors(video_path: str) -> dict:
                     result['stream_details'].append({'index': stream_index, 'codec_type': stream_type})
                     if stream_type not in ['video', 'audio']:
                         result['has_extra_streams'] = True
-                        
+
             if 'format' in ffprobe_data and 'nb_streams' in ffprobe_data['format']:
                 if result['num_streams'] == 0:
                     result['num_streams'] = ffprobe_data['format']['nb_streams']
@@ -119,18 +125,18 @@ def check_video_for_extra_streams_and_errors(video_path: str) -> dict:
     except Exception as e:
         result['ffprobe_output_error'] = f"An unexpected error occurred: {e}"
         result['ffprobe_successful'] = False
-        
+
     return result
 
 
 def remove_video_extra_stream_ffmpeg(input_video: str, output_video: str) -> bool:
     """
     Remove extra streams from video using ffmpeg.
-    
+
     Args:
         input_video: Path to input video.
         output_video: Path to output video.
-        
+
     Returns:
         bool: True if successful, False otherwise.
     """
@@ -146,7 +152,7 @@ def remove_video_extra_stream_ffmpeg(input_video: str, output_video: str) -> boo
         "-movflags", "faststart",
         output_video,
     ]
-    
+
     try:
         subprocess.run(command_list, shell=False, check=True, capture_output=True)
         return True
@@ -162,13 +168,27 @@ def remove_video_extra_stream_ffmpeg(input_video: str, output_video: str) -> boo
         return False
 
 
+def _video_clean_lock_path(video_path: str) -> str:
+    lock_dir = os.path.join(tempfile.gettempdir(), "mossvl_video_clean_locks")
+    os.makedirs(lock_dir, exist_ok=True)
+    path_key = os.path.realpath(video_path).encode("utf-8")
+    return os.path.join(lock_dir, f"{hashlib.sha256(path_key).hexdigest()}.lock")
+
+
+def _is_reusable_clean_video(video_path: str) -> bool:
+    if not os.path.isfile(video_path) or os.path.getsize(video_path) <= 0:
+        return False
+    probe = check_video_for_extra_streams_and_errors(video_path)
+    return probe["ffprobe_successful"] and not probe["has_extra_streams"]
+
+
 def clean_video_streams(video_path: str) -> str:
     """
     Clean video streams if extra streams are detected.
-    
+
     Args:
         video_path: Path to the video file.
-        
+
     Returns:
         str: Path to cleaned video (or original if no cleaning needed).
     """
@@ -179,26 +199,86 @@ def clean_video_streams(video_path: str) -> str:
         file_name_without_ext, file_ext = os.path.splitext(base_name)
         new_base_name = f"{file_name_without_ext}_fix{file_ext}"
         video_path_output = os.path.join(output_folder, new_base_name)
-        
-        process_flag = remove_video_extra_stream_ffmpeg(video_path, video_path_output)
-        if not process_flag:
-            logger.warning("Failed to remove extra streams with ffmpeg")
+
+        try:
+            with open(_video_clean_lock_path(video_path), "a") as lock_file:
+                fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX)
+                if _is_reusable_clean_video(video_path_output):
+                    return video_path_output
+
+                temp_fd, temp_output = tempfile.mkstemp(
+                    prefix=f".{file_name_without_ext}_fix.",
+                    suffix=file_ext,
+                    dir=output_folder,
+                )
+                os.close(temp_fd)
+                try:
+                    if not remove_video_extra_stream_ffmpeg(video_path, temp_output):
+                        logger.warning("Failed to remove extra streams with ffmpeg")
+                        return video_path
+                    os.replace(temp_output, video_path_output)
+                    return video_path_output
+                finally:
+                    if os.path.exists(temp_output):
+                        os.unlink(temp_output)
+        except OSError as exc:
+            logger.warning(f"Failed to publish cleaned video for {video_path}: {exc}")
             return video_path
-        return video_path_output
     return video_path
 
 
+@lru_cache(maxsize=8192)
+def cached_clean_video_streams(video_path: str) -> str:
+    return clean_video_streams(video_path)
+
+
+def clamp_timestamps_for_torchcodec(timestamps: List[float], torchcodec_metadata) -> List[float]:
+    if not timestamps:
+        return timestamps
+
+    min_pts = torchcodec_metadata.begin_stream_seconds_from_content
+    if min_pts is None:
+        min_pts = 0.0
+    # TorchCodec can reject timestamps exactly equal to the reported stream
+    # begin due to tiny metadata/decoder precision differences.
+    safe_min_pts = min_pts + TORCHCODEC_TIMESTAMP_EPSILON
+
+    max_pts_candidates = []
+    if torchcodec_metadata.num_frames_from_content and torchcodec_metadata.average_fps:
+        max_pts_candidates.append(
+            (torchcodec_metadata.num_frames_from_content - 1) / torchcodec_metadata.average_fps + min_pts
+        )
+    if torchcodec_metadata.end_stream_seconds_from_content is not None:
+        # TorchCodec requires requested PTS to be strictly smaller than the content end.
+        max_pts_candidates.append(torchcodec_metadata.end_stream_seconds_from_content - TORCHCODEC_TIMESTAMP_EPSILON)
+    if not max_pts_candidates and torchcodec_metadata.duration_seconds is not None:
+        max_pts_candidates.append(torchcodec_metadata.duration_seconds - TORCHCODEC_TIMESTAMP_EPSILON)
+
+    if max_pts_candidates:
+        max_pts = max(safe_min_pts, min(max_pts_candidates))
+        return [max(safe_min_pts, min(float(t), max_pts)) for t in timestamps]
+    if safe_min_pts > 0:
+        return [max(safe_min_pts, float(t)) for t in timestamps]
+    return [float(t) for t in timestamps]
+
+
 def split_indices(indices: List[Union[int, float]], num_chunks: int) -> List[List[Union[int, float]]]:
     """
     Split an index list into roughly equal chunks.
-    
+
     Args:
         indices: List of indices to split.
         num_chunks: Number of chunks to create.
-        
+
     Returns:
         List of index chunks.
     """
+    if len(indices) == 0:
+        raise ValueError("indices must not be empty")
+    if num_chunks <= 0:
+        raise ValueError("num_chunks must be positive")
+
+    num_chunks = min(num_chunks, len(indices))
     chunk_size = len(indices) // num_chunks
     chunks = []
     for i in range(num_chunks - 1):
@@ -210,12 +290,12 @@ def split_indices(indices: List[Union[int, float]], num_chunks: int) -> List[Lis
 def decode_sequentially(indices: List[int], video_path: str, ffmpeg_threads: int = 0):
     """
     Decode frames sequentially from a video.
-    
+
     Args:
         indices: List of frame indices to decode.
         video_path: Path to the video file.
         ffmpeg_threads: Number of ffmpeg threads to use.
-        
+
     Returns:
         FrameBatch from torchcodec.
     """
@@ -229,20 +309,20 @@ def decode_sequentially(indices: List[int], video_path: str, ffmpeg_threads: int
 def decode_with_multithreading(indices: List[int], num_threads: int, video_path: str) -> dict:
     """
     Decode frames using multithreading with joblib.
-    
+
     Args:
         indices: List of frame indices to decode.
         num_threads: Number of threads to use.
         video_path: Path to the video file.
-        
+
     Returns:
         dict: Contains 'data', 'duration_seconds', 'pts_seconds' tensors.
     """
     chunks = split_indices(indices, num_chunks=num_threads)
-    results = Parallel(n_jobs=num_threads, prefer="threads", verbose=0)(
+    results = Parallel(n_jobs=len(chunks), prefer="threads", verbose=0)(
         delayed(decode_sequentially)(chunk, video_path) for chunk in chunks
     )
-    
+
     return {
         "data": torch.cat([frame_batch.data for frame_batch in results], dim=0),
         "duration_seconds": torch.cat([frame_batch.duration_seconds for frame_batch in results], dim=0),
@@ -253,35 +333,19 @@ def decode_with_multithreading(indices: List[int], num_threads: int, video_path:
 def decode_sequentially_timestamp(timestamp_list: List[float], video_path: str, ffmpeg_threads: int = 0):
     """
     Decode frames sequentially from a video based on timestamps.
-    
+
     Args:
         timestamp_list: List of timestamps (in seconds) to decode.
         video_path: Path to the video file.
         ffmpeg_threads: Number of ffmpeg threads to use.
-        
+
     Returns:
         FrameBatch from torchcodec.
     """
     decoder = VideoDecoder(video_path, num_ffmpeg_threads=ffmpeg_threads)
     try:
         metadata = decoder.metadata
-
-        min_pts = metadata.begin_stream_seconds_from_content
-        if min_pts is None:
-            min_pts = 0.0
-
-        max_pts = None
-        if metadata.num_frames_from_content and metadata.average_fps:
-            max_pts = (metadata.num_frames_from_content - 1) / metadata.average_fps + min_pts
-        elif metadata.end_stream_seconds_from_content is not None:
-            max_pts = metadata.end_stream_seconds_from_content
-        else:
-            max_pts = metadata.duration_seconds
-
-        if max_pts is not None and max_pts > 0:
-            timestamp_list = [max(min_pts, min(t, max_pts)) for t in timestamp_list]
-        elif min_pts > 0:
-            timestamp_list = [max(min_pts, t) for t in timestamp_list]
+        timestamp_list = clamp_timestamps_for_torchcodec(timestamp_list, metadata)
 
         return decoder.get_frames_played_at(timestamp_list)
     finally:
@@ -291,29 +355,29 @@ def decode_sequentially_timestamp(timestamp_list: List[float], video_path: str,
 def timestamp_decode_with_multithreading(timestamp_list: List[float], num_threads: int, video_path: str) -> dict:
     """
     Decode frames using multithreading based on timestamps.
-    
+
     Args:
         timestamp_list: List of timestamps (in seconds) to decode.
         num_threads: Number of threads to use.
         video_path: Path to the video file.
-        
+
     Returns:
         dict: Contains 'data', 'duration_seconds', 'pts_seconds' tensors.
     """
     chunks = split_indices(timestamp_list, num_chunks=num_threads)
-    results = Parallel(n_jobs=num_threads, prefer="threads", verbose=0)(
+    results = Parallel(n_jobs=len(chunks), prefer="threads", verbose=0)(
         delayed(decode_sequentially_timestamp)(chunk, video_path) for chunk in chunks
     )
-    
+
     # Concatenate results from all threads
     data_list = [frame_batch.data for frame_batch in results]
     duration_list = [frame_batch.duration_seconds for frame_batch in results]
     pts_list = [frame_batch.pts_seconds for frame_batch in results]
-    
+
     if not data_list:
         logger.warning("No frames were successfully decoded.")
         return {"data": torch.empty(0), "duration_seconds": torch.empty(0), "pts_seconds": torch.empty(0)}
-    
+
     return {
         "data": torch.cat(data_list, dim=0),
         "duration_seconds": torch.cat(duration_list, dim=0),
@@ -329,42 +393,42 @@ def extract_frames_with_torchcodec(
 ) -> Optional[dict]:
     """
     Extract frames from video using torchcodec with multithreading.
-    
+
     Args:
         video_path: Path to the video file.
         sample_frames_count: Number of frames to sample.
         num_threads: Number of threads to use for extraction.
         sampling_method: Sampling method, either "index" (uniform frame indices) or "timestamp" (uniform timestamps).
-        
+
     Returns:
         dict: Contains 'data' (N, C, H, W), 'duration_seconds' (N,), 'pts_seconds' (N,) tensors.
               Returns None if extraction fails.
     """
     try:
-        video_path = clean_video_streams(video_path)
+        video_path = cached_clean_video_streams(video_path)
         decoder = VideoDecoder(video_path, num_ffmpeg_threads=0)
         metadata = decoder.metadata
 
 
         total_frames_in_video = metadata.num_frames_from_content
-        
+
         effective_sample_count = min(sample_frames_count, total_frames_in_video)
         if effective_sample_count == 0:
             logger.error("Cannot extract frames: video has 0 frames or specified frame count is 0")
             return None
-        
+
         # Generate uniform frame indices
         frame_indices = np.linspace(0, total_frames_in_video - 1, effective_sample_count).astype(np.int32)
         # Ensure indices are valid and remove duplicates
         frame_indices = np.unique(np.clip(frame_indices, 0, total_frames_in_video - 1))
-        
+
         result = decode_with_multithreading(frame_indices.tolist(), num_threads=num_threads, video_path=video_path)
         # Add frame_indices to the result for later use
         result["frame_indices"] = frame_indices
         return result
 
 
-            
+
     except Exception:
         traceback.print_exc()
         return None
@@ -500,9 +564,28 @@ class MossVLVideoProcessor(BaseVideoProcessor):
             return video_input["video_path"]
         return video_input
 
+    def _get_video_fps_from_input(
+        self,
+        video_input: Union[str, Dict[str, Any]],
+        default_fps: Optional[float],
+    ) -> Optional[float]:
+        if not isinstance(video_input, dict) or "fps" not in video_input:
+            return default_fps
+
+        raw_fps = video_input["fps"]
+        if raw_fps is None:
+            return default_fps
+        try:
+            fps = float(raw_fps)
+        except (TypeError, ValueError) as exc:
+            raise ValueError(f"Invalid video fps value: {raw_fps!r}") from exc
+        if fps <= 0:
+            raise ValueError(f"Invalid video fps value: {raw_fps!r}; fps must be positive")
+        return fps
+
     def _get_video_duration_seconds(self, video_input: Union[str, Dict[str, Any]]) -> float:
         """Get video duration in seconds for weighted frame-budget allocation."""
-        video_path = clean_video_streams(self._get_video_path_from_input(video_input))
+        video_path = cached_clean_video_streams(self._get_video_path_from_input(video_input))
         decoder = VideoDecoder(video_path, num_ffmpeg_threads=0)
         try:
             metadata = decoder.metadata
@@ -587,13 +670,13 @@ class MossVLVideoProcessor(BaseVideoProcessor):
     ) -> int:
         """
         Calculate the number of frames to sample using fps-based logic with min/max constraints.
-        
+
         Logic:
         1. Calculate target_frames based on fps and video duration
         2. Apply min_frames and max_frames constraints
         3. Apply max_allowed_frames protection (rough cap from total video_max_pixels budget)
         4. Return the number of frames to sample
-        
+
         Args:
             metadata (`VideoMetadata`):
                 Metadata of the video containing information about total duration, fps and total number of frames.
@@ -613,11 +696,11 @@ class MossVLVideoProcessor(BaseVideoProcessor):
             raise ValueError("`num_frames` and `fps` are mutually exclusive arguments, please use only one!")
 
         total_num_frames = metadata.total_num_frames
-        
+
         # Use provided min/max or fall back to defaults
         effective_min_frames = min_frames if min_frames is not None else self.min_frames
         effective_max_frames = max_frames if max_frames is not None else self.max_frames
-        
+
         # Rough per-video frame cap derived from the multi-video total budget
         # (exact allocation happens later in _preprocess via weighted distribution)
         per_frame_min_pixels = self.size.get("shortest_edge", None) if self.size else None
@@ -625,7 +708,7 @@ class MossVLVideoProcessor(BaseVideoProcessor):
         if per_frame_min_pixels is not None and video_max_pixels is not None and per_frame_min_pixels > 0:
             max_allowed_frames = video_max_pixels // per_frame_min_pixels
             effective_max_frames = min(effective_max_frames, max_allowed_frames)
-        
+
         # Get video duration
         if hasattr(metadata, 'duration') and metadata.duration is not None:
             duration = metadata.duration
@@ -643,12 +726,12 @@ class MossVLVideoProcessor(BaseVideoProcessor):
 
         # Use provided fps or default
         target_fps = fps if fps is not None else self.video_fps
-        
+
         # Calculate target frames based on fps and duration
         if num_frames is None:
             # Calculate how many frames we should sample based on target fps
             target_total_frames = int(math.ceil(duration * target_fps - 1e-6))
-            
+
             # Apply min/max constraints
             sample_frames = max(target_total_frames, effective_min_frames)
             sample_frames = min(sample_frames, effective_max_frames, total_num_frames)
@@ -659,6 +742,134 @@ class MossVLVideoProcessor(BaseVideoProcessor):
         return sample_frames
 
 
+    def _decode_timestamps_with_decoder(
+        self,
+        decoder: VideoDecoder,
+        timestamps: List[float],
+        chunk_size: int = 128,
+    ) -> torch.Tensor:
+        if not timestamps:
+            return torch.empty(0)
+
+        frame_chunks = []
+        for start in range(0, len(timestamps), chunk_size):
+            frame_batch = decoder.get_frames_played_at(timestamps[start:start + chunk_size])
+            frame_chunks.append(frame_batch.data)
+
+        if len(frame_chunks) == 1:
+            return frame_chunks[0]
+        return torch.cat(frame_chunks, dim=0)
+
+    def _clamp_timestamps_for_decoder(
+        self,
+        timestamps: List[float],
+        torchcodec_metadata,
+    ) -> List[float]:
+        return clamp_timestamps_for_torchcodec(timestamps, torchcodec_metadata)
+
+    def _fetch_video_segments_batched(
+        self,
+        video_path: str,
+        segments: List[List[float]],
+        min_frames: Optional[int] = None,
+        max_frames: Optional[int] = None,
+        video_fps: Optional[float] = None,
+    ):
+        min_frames = max(1, min_frames if min_frames is not None else self.min_frames)
+        max_frames = max(1, max_frames if max_frames is not None else self.max_frames)
+        target_video_fps = video_fps if video_fps is not None else self.video_fps
+
+        video_path = cached_clean_video_streams(video_path)
+        decoder = VideoDecoder(video_path, num_ffmpeg_threads=0)
+        try:
+            torchcodec_metadata = decoder.metadata
+            source_video_fps = torchcodec_metadata.average_fps
+
+            duration = None
+            if (
+                torchcodec_metadata.end_stream_seconds_from_content is not None
+                and torchcodec_metadata.begin_stream_seconds_from_content is not None
+            ):
+                duration = (
+                    torchcodec_metadata.end_stream_seconds_from_content
+                    - torchcodec_metadata.begin_stream_seconds_from_content
+                )
+            if duration is None or duration <= 0:
+                duration = torchcodec_metadata.duration_seconds
+
+            segment_durations = [
+                segment[1] - segment[0] if len(segment) == 2 else None
+                for segment in segments
+            ]
+            total_segment_duration = sum(d for d in segment_durations if d is not None)
+            num_range_segments = sum(1 for d in segment_durations if d is not None)
+
+            segment_timestamps = []
+            decode_timestamps = []
+            for i, segment in enumerate(segments):
+                if len(segment) == 1:
+                    actual_timestamps = self._clamp_timestamps_for_decoder([segment[0]], torchcodec_metadata)
+                    segment_timestamps.append(actual_timestamps)
+                    decode_timestamps.extend(actual_timestamps)
+                    continue
+
+                start_time, end_time = segment
+                segment_duration = end_time - start_time
+                target_frames = int(math.ceil(segment_duration * target_video_fps))
+
+                if total_segment_duration > 0:
+                    weight = segment_durations[i] / total_segment_duration
+                else:
+                    weight = 1.0 / num_range_segments if num_range_segments > 0 else 1.0
+
+                weighted_min_frames = max(1, int(round(min_frames * weight)))
+                weighted_max_frames = max(1, int(round(max_frames * weight)))
+                target_frames = max(target_frames, weighted_min_frames)
+                target_frames = min(target_frames, weighted_max_frames)
+
+                if target_frames == 1:
+                    actual_timestamps = [start_time]
+                else:
+                    actual_timestamps = np.linspace(
+                        start_time,
+                        end_time,
+                        target_frames,
+                        endpoint=False,
+                    ).tolist()
+
+                actual_timestamps = self._clamp_timestamps_for_decoder(actual_timestamps, torchcodec_metadata)
+                segment_timestamps.append(actual_timestamps)
+                decode_timestamps.extend(actual_timestamps)
+
+            flat_frames = self._decode_timestamps_with_decoder(decoder, decode_timestamps)
+
+            videos = []
+            metadata = []
+            frame_offset = 0
+            for actual_timestamps in segment_timestamps:
+                sample_count = len(actual_timestamps)
+                video_tensor = flat_frames[frame_offset:frame_offset + sample_count]
+                frame_offset += sample_count
+
+                video_metadata = VideoMetadata(
+                    total_num_frames=sample_count,
+                    fps=source_video_fps,
+                    duration=duration,
+                    video_backend="torchcodec",
+                    height=torchcodec_metadata.height,
+                    width=torchcodec_metadata.width,
+                    frames_indices=None,
+                )
+                video_metadata.actual_timestamps = actual_timestamps
+
+                videos.append(video_tensor)
+                metadata.append(video_metadata)
+
+            return videos, metadata
+        finally:
+            del decoder
+
+
     def _fetch_video_segment(
         self,
         video_path: str,
@@ -669,14 +880,14 @@ class MossVLVideoProcessor(BaseVideoProcessor):
     ):
         """
         Fetch video frames for a specific segment.
-        
+
         Args:
             video_path: Path to the video file
             segment: [start, end] for a segment (left-closed, right-open) or [time] for a single frame
             min_frames: Minimum frames for this segment (weighted). Defaults to self.min_frames. Must be >= 1.
             max_frames: Maximum frames for this segment (weighted). Defaults to self.max_frames. Must be >= 1.
             video_fps: Target frames per second for video sampling. If None, uses self.video_fps.
-        
+
         Returns:
             Tuple of (video_tensor, video_metadata)
         """
@@ -685,50 +896,51 @@ class MossVLVideoProcessor(BaseVideoProcessor):
         max_frames = max(1, max_frames if max_frames is not None else self.max_frames)
         # Use provided video_fps or fall back to self.video_fps
         target_video_fps = video_fps if video_fps is not None else self.video_fps
-        
-        video_path = clean_video_streams(video_path)
+
+        video_path = cached_clean_video_streams(video_path)
         decoder = VideoDecoder(video_path, num_ffmpeg_threads=0)
         try:
             torchcodec_metadata = decoder.metadata
-            
+
             video_fps = torchcodec_metadata.average_fps
-            
+
             # Calculate duration
             duration = None
             if torchcodec_metadata.end_stream_seconds_from_content is not None and torchcodec_metadata.begin_stream_seconds_from_content is not None:
                 duration = torchcodec_metadata.end_stream_seconds_from_content - torchcodec_metadata.begin_stream_seconds_from_content
             if duration is None or duration <= 0:
                 duration = torchcodec_metadata.duration_seconds
-            
+
             if len(segment) == 1:
                 # Single frame at specified time
-                timestamp = segment[0]
-                frame_batch = decoder.get_frames_played_at([timestamp])
+                actual_timestamps = self._clamp_timestamps_for_decoder([segment[0]], torchcodec_metadata)
+                frame_batch = decoder.get_frames_played_at(actual_timestamps)
                 video_tensor = frame_batch.data
-                actual_timestamps = [timestamp]
                 sample_count = 1
             else:
                 # Segment [start, end) - left-closed, right-open interval
                 start_time, end_time = segment
                 segment_duration = end_time - start_time
-                
+
                 # Calculate number of frames to sample for this segment
                 target_frames = int(math.ceil(segment_duration * target_video_fps))
                 target_frames = max(target_frames, min_frames)
                 target_frames = min(target_frames, max_frames)
-                
+
                 # Generate timestamps for uniform sampling within segment
                 if target_frames == 1:
                     actual_timestamps = [start_time]  # Use start_time for single frame
                 else:
                     # Sample uniformly within [start, end), endpoint=False for left-closed right-open
                     actual_timestamps = np.linspace(start_time, end_time, target_frames, endpoint=False).tolist()
-                
+
+                actual_timestamps = self._clamp_timestamps_for_decoder(actual_timestamps, torchcodec_metadata)
+
                 # Use multithreading for extraction
                 result = timestamp_decode_with_multithreading(actual_timestamps, self.num_extract_threads, video_path)
                 video_tensor = result["data"]
                 sample_count = len(actual_timestamps)
-            
+
             # Create VideoMetadata
             video_metadata = VideoMetadata(
                 total_num_frames=sample_count,
@@ -739,39 +951,40 @@ class MossVLVideoProcessor(BaseVideoProcessor):
                 width=torchcodec_metadata.width,
                 frames_indices=None
             )
-            
+
             # Store actual timestamps as a custom attribute for _calculate_timestamps to use
             video_metadata.actual_timestamps = actual_timestamps
-            
+
             return video_tensor, video_metadata
         finally:
             del decoder
 
     def fetch_videos(
-        self, 
-        video_url_or_urls: Union[str, Dict[str, Any], List[Union[str, Dict[str, Any]]]], 
-        sample_indices_fn=None, 
+        self,
+        video_url_or_urls: Union[str, Dict[str, Any], List[Union[str, Dict[str, Any]]]],
+        sample_indices_fn=None,
         video_fps: Optional[float] = None,
         min_frames: Optional[int] = None,
         max_frames: Optional[int] = None,
     ):
         """
         Override fetch_videos to use torchcodec for frame extraction.
-        
+
         This method uses torchcodec with multithreading for efficient frame extraction.
         Frame count is calculated by the calculate_num_frames method
         (fps-based with min/max constraints).
-        
+
         Args:
             video_url_or_urls: Can be one of:
                 - str: Single video path
-                - Dict: Video with segments {"video_path": str, "segments": List[List[float]]}
-                - List[Union[str, Dict]]: List of video paths or segment dicts
+                - Dict: Video descriptor with ``video_path`` and optional ``fps`` /
+                  ``segments`` (``segment`` is accepted as an alias)
+                - List[Union[str, Dict]]: List of video paths or descriptors
             sample_indices_fn: (Not used) Kept for compatibility with base class signature.
             video_fps: Target frames per second for video sampling. If None, uses self.video_fps.
             min_frames: Minimum number of frames to sample. If None, uses self.min_frames.
             max_frames: Maximum number of frames to sample. If None, uses self.max_frames.
-        
+
         Returns:
             Tuple of (videos, metadata) where videos are torch.Tensors and metadata are VideoMetadata objects.
         """
@@ -791,9 +1004,11 @@ class MossVLVideoProcessor(BaseVideoProcessor):
                     effective_max_frames,
                 )
             for x, allocated_max_frames in zip(video_url_or_urls, per_video_max_frames):
+                explicit_video_fps = self._get_video_fps_from_input(x, None)
+                item_video_fps = explicit_video_fps if explicit_video_fps is not None else effective_video_fps
                 result = self.fetch_videos(
-                    x, 
-                    video_fps=effective_video_fps,
+                    x,
+                    video_fps=item_video_fps,
                     min_frames=effective_min_frames,
                     max_frames=allocated_max_frames,
                 )
@@ -805,57 +1020,31 @@ class MossVLVideoProcessor(BaseVideoProcessor):
                     all_videos.append(result[0])
                     all_metadata.append(result[1])
             return all_videos, all_metadata
-        
-        # Handle dict with segments - returns lists (one per segment)
+
+        # Dict inputs may override FPS and optionally select segments. Without
+        # segments, the dict describes a full video.
         if isinstance(video_url_or_urls, dict):
             video_path = video_url_or_urls["video_path"]
-            segments = video_url_or_urls["segments"]
-            
-            # Calculate total duration of all time-range segments (len == 2) for weighted min/max frames
-            # Single-frame segments (len == 1) are excluded from weighting
-            segment_durations = []
-            for seg in segments:
-                if len(seg) == 2:
-                    segment_durations.append(seg[1] - seg[0])
-                else:
-                    segment_durations.append(None)  # Single frame, no weighting
-            
-            total_segment_duration = sum(d for d in segment_durations if d is not None)
-            
-            videos = []
-            metadata = []
-            for i, segment in enumerate(segments):
-                if len(segment) == 1:
-                    # Single frame - no weighted min/max, just extract directly
-                    video, meta = self._fetch_video_segment(video_path, segment, video_fps=effective_video_fps)
-                else:
-                    # Time-range segment - apply weighted min/max frames
-                    if total_segment_duration > 0:
-                        weight = segment_durations[i] / total_segment_duration
-                    else:
-                        # Fallback: equal weight among time-range segments
-                        num_range_segments = sum(1 for d in segment_durations if d is not None)
-                        weight = 1.0 / num_range_segments if num_range_segments > 0 else 1.0
-                    
-                    # Calculate weighted min/max frames (ensure >= 1)
-                    weighted_min_frames = max(1, int(round(effective_min_frames * weight)))
-                    weighted_max_frames = max(1, int(round(effective_max_frames * weight)))
-                    
-                    video, meta = self._fetch_video_segment(
-                        video_path, segment,
-                        min_frames=weighted_min_frames,
-                        max_frames=weighted_max_frames,
-                        video_fps=effective_video_fps,
-                    )
-                videos.append(video)
-                metadata.append(meta)
-            return videos, metadata
-        
+            input_video_fps = self._get_video_fps_from_input(video_url_or_urls, effective_video_fps)
+            segments = video_url_or_urls.get("segments", video_url_or_urls.get("segment"))
+
+            if segments is not None:
+                return self._fetch_video_segments_batched(
+                    video_path,
+                    segments,
+                    min_frames=effective_min_frames,
+                    max_frames=effective_max_frames,
+                    video_fps=input_video_fps,
+                )
+
+            video_url_or_urls = video_path
+            effective_video_fps = input_video_fps
+
         # Single video path
         video_path = video_url_or_urls
-        
+
         # Clean video streams first (remove extra streams if needed)
-        video_path = clean_video_streams(video_path)
+        video_path = cached_clean_video_streams(video_path)
 
         decoder = None
         try:
@@ -866,13 +1055,13 @@ class MossVLVideoProcessor(BaseVideoProcessor):
             duration = None
             if torchcodec_metadata.end_stream_seconds_from_content is not None and torchcodec_metadata.begin_stream_seconds_from_content is not None:
                 duration = torchcodec_metadata.end_stream_seconds_from_content - torchcodec_metadata.begin_stream_seconds_from_content
-            
+
             if duration is None or duration <= 0:
                 duration = torchcodec_metadata.duration_seconds
-            
+
             # Use num_frames_from_content for accurate frame count (consistent with extraction)
             total_frames_in_video = torchcodec_metadata.num_frames_from_content
-            
+
             # Create VideoMetadata object for sample_frames method
             temp_metadata = VideoMetadata(
                 total_num_frames=total_frames_in_video,
@@ -883,31 +1072,31 @@ class MossVLVideoProcessor(BaseVideoProcessor):
                 width=torchcodec_metadata.width,
                 frames_indices=None
             )
-            
+
             # Use calculate_num_frames method to get the number of frames to sample
             sample_frames_count = self.calculate_num_frames(
-                temp_metadata, 
+                temp_metadata,
                 fps=effective_video_fps,
                 min_frames=effective_min_frames,
                 max_frames=effective_max_frames,
             )
-            
+
             # Ensure sample count is valid
             effective_sample_count = min(sample_frames_count, total_frames_in_video)
             if effective_sample_count == 0:
                 raise ValueError(f"Cannot extract frames: video has 0 frames or specified frame count is 0")
-            
+
             # Generate uniform frame indices
             frame_indices = np.linspace(0, total_frames_in_video - 1, effective_sample_count).astype(np.int32)
             # Ensure indices are valid and remove duplicates
             frame_indices = np.unique(np.clip(frame_indices, 0, total_frames_in_video - 1))
-            
+
             # Extract frames using multithreading (decoder is created inside each thread for thread safety)
             result = decode_with_multithreading(frame_indices.tolist(), num_threads=self.num_extract_threads, video_path=video_path)
-            
+
             # Extract frame tensor (N, C, H, W)
             frames_tensor = result["data"]
-            
+
             # Create final VideoMetadata object
             video_metadata = VideoMetadata(
                 total_num_frames=len(frame_indices),
@@ -918,15 +1107,15 @@ class MossVLVideoProcessor(BaseVideoProcessor):
                 width=torchcodec_metadata.width,
                 frames_indices=frame_indices
             )
-            
+
             # Ensure frames are in (T, C, H, W) format
             if frames_tensor.dim() == 4:  # (N, C, H, W)
                 video_tensor = frames_tensor
             else:
                 raise ValueError(f"Unexpected frame tensor shape: {frames_tensor.shape}")
-            
+
             return video_tensor, video_metadata
-            
+
         except Exception as e:
             logger.error(f"Error loading video {video_path}: {e}")
             traceback.print_exc()
@@ -941,7 +1130,7 @@ class MossVLVideoProcessor(BaseVideoProcessor):
         do_convert_rgb: bool = True,
         do_resize: bool = True,
         size: Optional[SizeDict] = None,
-        resample: PILImageResampling = PILImageResampling.BICUBIC,
+        interpolation: PILImageResampling = PILImageResampling.BICUBIC,
         do_rescale: bool = True,
         rescale_factor: float = 1 / 255.0,
         do_normalize: bool = True,
@@ -991,7 +1180,7 @@ class MossVLVideoProcessor(BaseVideoProcessor):
                 stacked_videos = self.resize(
                     stacked_videos,
                     size=SizeDict(height=resized_height, width=resized_width),
-                    resample=resample,
+                    interpolation=interpolation,
                 )
                 stacked_videos = stacked_videos.view(B, T, C, resized_height, resized_width)
             resized_videos_grouped[shape] = stacked_videos
@@ -1031,12 +1220,17 @@ class MossVLVideoProcessor(BaseVideoProcessor):
                 merge_size,
                 patch_size,
             )
+            patches_device = patches.device
+            # NPU: max 8D tensors — route the 10D permute+reshape through CPU.
+            # CUDA handles 10D natively — keep it on-device.
+            if patches_device.type == "npu":
+                patches = patches.cpu()
             patches = patches.permute(0, 1, 4, 7, 5, 8, 3, 2, 6, 9)
             flatten_patches = patches.reshape(
                 batch_size,
                 grid_t * grid_h * grid_w,
                 channel * temporal_patch_size * patch_size * patch_size,
-            )
+            ).to(patches_device)
 
             processed_videos_grouped[shape] = flatten_patches
             processed_grids[shape] = [[grid_t, grid_h, grid_w]] * batch_size
@@ -1059,15 +1253,15 @@ class MossVLVideoProcessor(BaseVideoProcessor):
     ) -> BatchFeature:
         """
         Preprocess videos for the model.
-        
+
         This method overrides the base class to handle two video input formats:
         1. String path: "path/to/video.mp4"
-        2. Dict with segments: {"video_path": "...", "segment": [[start, end], [time], ...]}
-        
+        2. Dict descriptor with ``video_path`` and optional ``fps`` / ``segments``
+
         Args:
             videos: Video input(s) in one of the supported formats.
             **kwargs: Additional arguments passed to _preprocess.
-            
+
         Returns:
             BatchFeature with pixel_values_videos, video_grid_thw, and optionally video_metadata.
         """
@@ -1076,11 +1270,11 @@ class MossVLVideoProcessor(BaseVideoProcessor):
             captured_kwargs=kwargs.keys(),
             valid_processor_keys=list(self.valid_kwargs.__annotations__.keys()) + ["return_tensors"],
         )
-        
+
         # Set default kwargs from self
         for kwarg_name in self.valid_kwargs.__annotations__:
             kwargs.setdefault(kwarg_name, getattr(self, kwarg_name, None))
-        
+
         # Pop kwargs that are handled separately
         return_tensors = kwargs.pop("return_tensors", None)
         return_metadata = kwargs.pop("return_metadata", False)
@@ -1089,44 +1283,43 @@ class MossVLVideoProcessor(BaseVideoProcessor):
         kwargs.pop("video_metadata", None)  # We generate our own metadata
         kwargs.pop("do_sample_frames", None)  # We handle sampling ourselves
         kwargs.pop("data_format", None)  # Not used
-        
+
         # Normalize input to list format
         if not isinstance(videos, list):
             videos = [videos]
-        
+
         # Get video processing params from kwargs (may be passed explicitly for per-batch configuration)
         video_fps = kwargs.pop("video_fps", None)
         min_frames = kwargs.pop("min_frames", None)
         max_frames = kwargs.pop("max_frames", None)
-        
+
         # Use fetch_videos to handle both string and dict formats
         video_tensors, video_metadata = self.fetch_videos(
-            videos, 
+            videos,
             video_fps=video_fps,
             min_frames=min_frames,
             max_frames=max_frames,
         )
-        
+
         # Prepare video tensors using _prepare_input_videos
         prepared_videos = self._prepare_input_videos(
             videos=video_tensors,
             input_data_format=input_data_format,
             device=device,
         )
-        
+
         # Process kwargs for _preprocess
         kwargs = self._further_process_kwargs(**kwargs)
         self._validate_preprocess_kwargs(**kwargs)
-        
+
         # Call _preprocess with prepared videos
         result = self._preprocess(videos=prepared_videos, return_tensors=return_tensors, **kwargs)
-        
+
         # Add metadata if requested
         if return_metadata:
             result["video_metadata"] = video_metadata
-        
+
         return result
 
 
 __all__ = ["MossVLVideoProcessor"]
-