CCCCyx commited on
Commit
5b97f7a
·
verified ·
1 Parent(s): 1e6a45b

Unify MOSS-VL processor and inference preprocessing defaults

Browse files

Use the shared processor and tokenizer, set image/video pixel budgets to 201326592, and disable vision tail padding with vision_seq_pad_multiple=1. Preserve model implementation and weights.

preprocessor_config.json CHANGED
@@ -7,7 +7,7 @@
7
  "longest_edge": 16777216,
8
  "shortest_edge": 4096
9
  },
10
- "multi_image_max_pixels": 65536000,
11
  "patch_size": 16,
12
  "temporal_patch_size": 1,
13
  "merge_size": 2,
 
7
  "longest_edge": 16777216,
8
  "shortest_edge": 4096
9
  },
10
+ "multi_image_max_pixels": 201326592,
11
  "patch_size": 16,
12
  "temporal_patch_size": 1,
13
  "merge_size": 2,
processing_moss_vl.py CHANGED
@@ -49,7 +49,7 @@ class MossVLImageProcessorFast(Qwen2VLImageProcessorFast):
49
  """
50
  # Multi-image batch total pixels limit (read from config)
51
  multi_image_max_pixels = None
52
-
53
 
54
  def _preprocess(
55
  self,
@@ -70,7 +70,7 @@ class MossVLImageProcessorFast(Qwen2VLImageProcessorFast):
70
  **kwargs,
71
  ):
72
  """Override _preprocess to use custom smart_resize with batch-level max_pixels.
73
-
74
  multi_image_max_pixels is treated as a batch-level total budget, proportionally allocated
75
  to each image based on its original pixel count. min_pixels remains a per-image
76
  constraint. multi_image_max_pixels can be configured separately from longest_edge.
@@ -79,7 +79,7 @@ class MossVLImageProcessorFast(Qwen2VLImageProcessorFast):
79
  max_pixels = size["longest_edge"] # Per-image upper limit
80
  # Use multi_image_max_pixels if configured, otherwise fall back to longest_edge
81
  multi_image_max_pixels = getattr(self, "multi_image_max_pixels", None) or max_pixels
82
-
83
  # Calculate total original pixels across all images in the batch
84
  # This is used to proportionally allocate max_pixels to each image
85
  total_original_pixels = sum(img.shape[-2] * img.shape[-1] for img in images)
@@ -98,13 +98,13 @@ class MossVLImageProcessorFast(Qwen2VLImageProcessorFast):
98
  proportional_max_pixels = int(multi_image_max_pixels * proportion)
99
  else:
100
  proportional_max_pixels = multi_image_max_pixels
101
-
102
  # Ensure proportional max_pixels is within [min_pixels, max_pixels] range
103
  # min_pixels: per-image lower limit (shortest_edge)
104
  # max_pixels: per-image upper limit (longest_edge)
105
  proportional_max_pixels = max(proportional_max_pixels, min_pixels)
106
  proportional_max_pixels = min(proportional_max_pixels, max_pixels)
107
-
108
  resized_height, resized_width = smart_resize(
109
  height,
110
  width,
@@ -119,7 +119,7 @@ class MossVLImageProcessorFast(Qwen2VLImageProcessorFast):
119
  )
120
  resized_images_grouped[shape] = stacked_images
121
  resized_images = reorder_images(resized_images_grouped, grouped_images_index)
122
-
123
  # Warn if multi-image batch exceeds multi_image_max_pixels due to min_pixels constraint
124
  if len(images) > 1:
125
  total_resized_pixels = sum(img.shape[-2] * img.shape[-1] for img in resized_images)
@@ -192,17 +192,17 @@ def _to_numpy(x):
192
  """
193
  Convert various tensor types to numpy array.
194
  Supports torch.Tensor, tf.Tensor, jax.Array, np.ndarray, lists, and primitives.
195
-
196
  Args:
197
  x: Input value that can be a tensor from various frameworks or a Python primitive
198
-
199
  Returns:
200
  np.ndarray: NumPy array representation of the input
201
  """
202
  # Already numpy
203
  if isinstance(x, np.ndarray):
204
  return x
205
-
206
  # Torch tensor or TensorFlow tensor (both have .numpy() method)
207
  if hasattr(x, 'numpy'):
208
  # For torch tensors on CUDA, need to move to CPU first
@@ -210,11 +210,11 @@ def _to_numpy(x):
210
  return x.cpu().numpy()
211
  # For TensorFlow or already on CPU
212
  return x.numpy()
213
-
214
  # JAX arrays and other array-like objects that support __array__ protocol
215
  if hasattr(x, '__array__'):
216
  return np.asarray(x)
217
-
218
  # Python primitives (list, tuple, int, float)
219
  return np.array(x)
220
 
@@ -334,20 +334,20 @@ class MossVLProcessor(ProcessorMixin):
334
  tokenizer_class = ("Qwen2Tokenizer", "Qwen2TokenizerFast")
335
 
336
  def __init__(
337
- self,
338
- image_processor=None,
339
- tokenizer=None,
340
- video_processor=None,
341
  chat_template=None,
342
  **kwargs
343
  ):
344
  super().__init__(image_processor, tokenizer, video_processor, chat_template=chat_template)
345
-
346
 
347
  self.image_token = "<|image_pad|>" if not hasattr(tokenizer, "image_token") else tokenizer.image_token
348
  self.video_token = "<|video_pad|>" if not hasattr(tokenizer, "video_token") else tokenizer.video_token
349
 
350
-
351
  self.image_token_id = (
352
  tokenizer.image_token_id
353
  if getattr(tokenizer, "image_token_id", None)
@@ -358,7 +358,7 @@ class MossVLProcessor(ProcessorMixin):
358
  if getattr(tokenizer, "video_token_id", None)
359
  else tokenizer.convert_tokens_to_ids(self.video_token)
360
  )
361
-
362
  self.vision_start_token = (
363
  "<|vision_start|>" if not hasattr(tokenizer, "vision_start_token") else tokenizer.vision_start_token
364
  )
@@ -372,15 +372,15 @@ class MossVLProcessor(ProcessorMixin):
372
 
373
  self.time_start_token = "<|time_start|>"
374
  self.time_end_token = "<|time_end|>"
375
-
376
  # EOS token for labels generation (assistant's response should end with this)
377
  self.im_end_token = "<|im_end|>"
378
  self.im_end_token_id = tokenizer.convert_tokens_to_ids(self.im_end_token)
379
-
380
  # Vision-related token ids (all should be masked in labels)
381
  self.vision_start_token_id = tokenizer.convert_tokens_to_ids(self.vision_start_token)
382
  self.vision_end_token_id = tokenizer.convert_tokens_to_ids(self.vision_end_token)
383
-
384
  # Token ids that should always be masked in labels (e.g. <|image_pad|>)
385
  self.mask_token_ids = {self.image_token_id}
386
 
@@ -444,7 +444,7 @@ class MossVLProcessor(ProcessorMixin):
444
  tokenizer_init_kwargs=self.tokenizer.init_kwargs,
445
  **kwargs,
446
  )
447
-
448
  # Step 1: Process images if provided
449
  if images is not None:
450
  images_kwargs = output_kwargs["images_kwargs"].copy()
@@ -474,12 +474,12 @@ class MossVLProcessor(ProcessorMixin):
474
  # Step 3: Process text with placeholder replacement
475
  if text is None or (isinstance(text, str) and len(text.strip()) == 0):
476
  raise ValueError("Text input is required for MossVL processor and cannot be empty.")
477
-
478
  if not isinstance(text, list):
479
  text = [text]
480
-
481
  text = text.copy() # Copy to avoid in-place modifications
482
-
483
  # Prepare labels_spans if provided
484
  # labels_spans format: List[List[List[int]]] - batch of samples, each sample has multiple spans
485
  # Each span is [start, end] (list, not tuple) so it can be modified in place
@@ -496,10 +496,10 @@ class MossVLProcessor(ProcessorMixin):
496
  has_images = images is not None and "pixel_values" in image_inputs
497
  has_videos = videos is not None and "pixel_values_videos" in videos_inputs
498
  needs_reorder = has_images and has_videos
499
-
500
  image_pixel_values_list = []
501
  video_pixel_values_list = []
502
-
503
  # Step 3.0: Record the order of media in original text (before replacement)
504
  # This will be used later to correctly order pixel_values and grid_thw
505
  media_order_per_sample = []
@@ -510,25 +510,25 @@ class MossVLProcessor(ProcessorMixin):
510
  while pos < len(temp_text):
511
  img_pos = temp_text.find(self.image_placeholder, pos)
512
  vid_pos = temp_text.find(self.video_placeholder, pos)
513
-
514
  if img_pos == -1 and vid_pos == -1:
515
  break
516
-
517
  if img_pos != -1 and (vid_pos == -1 or img_pos < vid_pos):
518
  media_order.append(("image", img_pos))
519
  pos = img_pos + len(self.image_placeholder)
520
  elif vid_pos != -1:
521
  media_order.append(("video", vid_pos))
522
  pos = vid_pos + len(self.video_placeholder)
523
-
524
  media_order_per_sample.append(media_order)
525
-
526
  # Step 3.0.1: Check if any sample has no media (empty samples need blank image)
527
  # If there are empty samples, we need to enter slow path to handle them properly
528
  has_empty_samples = any(len(order) == 0 for order in media_order_per_sample)
529
  if has_empty_samples:
530
  needs_reorder = True
531
-
532
  # Split pixel values for reordering if needed
533
  if needs_reorder:
534
  if has_images:
@@ -559,7 +559,7 @@ class MossVLProcessor(ProcessorMixin):
559
  video_pixel_values_list = _split_array_or_tensor(
560
  flat_video_values, split_indices
561
  )
562
-
563
  # Step 3.1: Replace placeholders (simple replacement, no expansion yet)
564
  # In MossVL, one image placeholder = one image token
565
  # One video placeholder = one video token (will be expanded later)
@@ -576,33 +576,33 @@ class MossVLProcessor(ProcessorMixin):
576
  else:
577
  text[i] = text[i].replace(self.image_placeholder, self.image_token)
578
  text[i] = text[i].replace(self.video_placeholder, self.video_token)
579
-
580
- # Step 3.2: Validate token counts
581
  n_images_in_text = [t.count(self.image_token) for t in text]
582
  n_videos_in_text = [t.count(self.video_token) for t in text]
583
-
584
  # Count placeholders in text
585
  total_images_in_text = sum(n_images_in_text)
586
  total_videos_in_text = sum(n_videos_in_text)
587
-
588
  # Count actual images and videos provided
589
  total_images_provided = len(image_grid_thw) if image_grid_thw is not None else 0
590
  total_videos_provided = len(video_grid_thw) if video_grid_thw is not None else 0
591
-
592
  # Validate image counts
593
  if total_images_in_text != total_images_provided:
594
  raise ValueError(
595
  "Number of image tokens does not match number of images provided. "
596
  f"Found {total_images_in_text} image tokens in text and {total_images_provided} images."
597
  )
598
-
599
  # Validate video counts
600
  if total_videos_in_text != total_videos_provided:
601
  raise ValueError(
602
  "Number of video tokens does not match number of videos provided. "
603
  f"Found {total_videos_in_text} video tokens in text and {total_videos_provided} videos."
604
  )
605
-
606
  # Step 3.3: Expand video tokens with timestamps
607
  # Now expand each video token to multiple tokens (one per frame) with timestamps
608
  if video_grid_thw is not None:
@@ -640,10 +640,10 @@ class MossVLProcessor(ProcessorMixin):
640
  video_tokens.append(
641
  f"{self.time_start_token}{curr_time:.1f} seconds{self.time_end_token}{self.image_token}"
642
  )
643
-
644
  # Wrap the entire video sequence with vision_start and vision_end tokens
645
  video_placeholder = f"{self.vision_start_token}{''.join(video_tokens)}{self.vision_end_token}"
646
-
647
  # Replace the video token with expanded sequence and update spans if needed
648
  if should_create_labels:
649
  text[i], labels_spans[i] = self._replace_and_update_spans(
@@ -652,22 +652,22 @@ class MossVLProcessor(ProcessorMixin):
652
  else:
653
  text[i] = text[i].replace(self.video_token, video_placeholder, 1)
654
  index += 1
655
-
656
 
657
 
658
  # Step 4: Tokenize text
659
  return_tensors = output_kwargs["text_kwargs"].pop("return_tensors", None)
660
  return_mm_token_type_ids = output_kwargs["text_kwargs"].pop("return_mm_token_type_ids", None)
661
-
662
  # Request offset_mapping if we need to create labels
663
  if should_create_labels:
664
  output_kwargs["text_kwargs"]["return_offsets_mapping"] = True
665
-
666
  text_inputs = self.tokenizer(text, **output_kwargs["text_kwargs"])
667
-
668
  # ignore check_special_mm_tokens nums in test and input ids.
669
  # self._check_special_mm_tokens(text, text_inputs, modalities=["image", "video"])
670
-
671
  # Create labels if labels_spans was provided
672
  if should_create_labels:
673
  offset_mapping = text_inputs.pop("offset_mapping")
@@ -687,49 +687,49 @@ class MossVLProcessor(ProcessorMixin):
687
  # Step 5: Concatenate pixel_values and grid_thw in sequence order
688
  # Prepare output
689
  output_data = {**text_inputs}
690
-
691
  if not needs_reorder:
692
  # Fast path: only one media type, no reordering needed
693
  final_pixel_values = []
694
  final_grid_thw = []
695
-
696
  if has_images:
697
  final_pixel_values.append(image_inputs["pixel_values"])
698
  final_grid_thw.extend(image_grid_thw)
699
-
700
  if has_videos:
701
  final_pixel_values.append(videos_inputs["pixel_values_videos"])
702
  final_grid_thw.extend(video_grid_thw)
703
-
704
  if final_pixel_values:
705
  output_data["pixel_values"] = np.concatenate(final_pixel_values, axis=0) if len(final_pixel_values) > 1 else final_pixel_values[0]
706
-
707
  if final_grid_thw:
708
  output_data["grid_thw"] = np.stack(final_grid_thw, axis=0)
709
-
710
  # Calculate media_nums_per_sample
711
  media_nums_per_sample = []
712
  for batch_idx in range(len(text)):
713
  media_order = media_order_per_sample[batch_idx]
714
  media_nums_per_sample.append(len(media_order) if len(media_order) > 0 else 1)
715
-
716
  # Don't add media_nums_per_sample to output_data yet
717
  # Will add it after BatchFeature to keep it as list
718
-
719
  else:
720
  # Slow path: both images and videos exist, need reordering
721
  final_pixel_values = []
722
  final_grid_thw = []
723
  media_nums_per_sample = []
724
-
725
  # Global indices to track position in flattened image/video arrays
726
  global_image_idx = 0
727
  global_video_idx = 0
728
-
729
  for batch_idx in range(len(text)):
730
  # Use the recorded media order from Step 3.0
731
  media_order = media_order_per_sample[batch_idx]
732
-
733
  if len(media_order) == 0:
734
  # If no media provided for this sample, add a blank image
735
  media_nums_per_sample.append(1)
@@ -737,26 +737,26 @@ class MossVLProcessor(ProcessorMixin):
737
  patch_size = getattr(self.image_processor, "patch_size", None) or 16
738
  temporal_patch_size = getattr(self.image_processor, "temporal_patch_size", None) or 1
739
  merge_size = getattr(self.image_processor, "merge_size", None) or 2
740
-
741
  factor = patch_size * merge_size
742
  side = int(np.ceil(np.sqrt(min_pixels) / factor) * factor)
743
  grid_h = side // patch_size
744
  grid_w = side // patch_size
745
  grid_t = 1
746
-
747
  # Channel = 3 (RGB)
748
  channel = 3
749
  dim = channel * temporal_patch_size * patch_size * patch_size
750
  num_patches = grid_t * grid_h * grid_w
751
-
752
  blank_pixel_values = np.zeros((num_patches, dim), dtype=np.float32)
753
  blank_grid_thw = np.array([grid_t, grid_h, grid_w], dtype=np.int64)
754
-
755
  final_pixel_values.append(blank_pixel_values)
756
  final_grid_thw.append(blank_grid_thw)
757
  else:
758
  media_nums_per_sample.append(len(media_order))
759
-
760
  # Collect media data according to the recorded order
761
  for media_type, _ in media_order:
762
  if media_type == "image" and image_grid_thw is not None:
@@ -771,18 +771,18 @@ class MossVLProcessor(ProcessorMixin):
771
  final_pixel_values.append(video_pixel_values_list[global_video_idx])
772
  final_grid_thw.append(video_grid_thw[global_video_idx])
773
  global_video_idx += 1
774
-
775
  # Concatenate/stack to unified format
776
  if final_pixel_values:
777
  output_data["pixel_values"] = _concat_array_or_tensor(
778
  final_pixel_values, axis=0
779
  )
780
-
781
  if final_grid_thw:
782
  output_data["grid_thw"] = _stack_array_or_tensor(
783
  final_grid_thw, axis=0
784
  )
785
-
786
  # Don't add media_nums_per_sample to output_data yet
787
  # Will add it after BatchFeature to keep it as list
788
 
@@ -795,18 +795,18 @@ class MossVLProcessor(ProcessorMixin):
795
  output_data.get("attention_mask", None)
796
  )
797
  output_data["cross_attention_mask"] = cross_attention_mask
798
-
799
  # Add labels to output if created
800
  if should_create_labels:
801
  output_data["labels"] = labels
802
 
803
  # BatchFeature will handle conversion to pt/tf/jax/np based on tensor_type
804
  batch_feature = BatchFeature(data=output_data, tensor_type=return_tensors)
805
-
806
  # Add media_nums_per_sample after BatchFeature to keep it as list (not tensor)
807
  if media_nums_per_sample:
808
  batch_feature["media_nums_per_sample"] = media_nums_per_sample
809
-
810
  return batch_feature
811
 
812
  def _create_cross_attention_mask(self, input_ids, grid_thw, media_nums_per_sample, attention_mask=None):
@@ -815,7 +815,7 @@ class MossVLProcessor(ProcessorMixin):
815
  Video frames are treated as individual images.
816
  Mask values: True for masked, False for visible.
817
  Causal masking: text can see images that appear at or before the text position.
818
-
819
  Args:
820
  input_ids: List of token ids
821
  grid_thw: Grid sizes for each media item
@@ -824,7 +824,7 @@ class MossVLProcessor(ProcessorMixin):
824
  """
825
  batch_size = len(input_ids)
826
  max_text_len = max(len(ids) for ids in input_ids)
827
-
828
  # Calculate total frames per sample to find max_num_frames
829
  total_frames_per_sample = []
830
  media_idx = 0
@@ -833,7 +833,7 @@ class MossVLProcessor(ProcessorMixin):
833
  if num_media == 0:
834
  total_frames_per_sample.append(0)
835
  continue
836
-
837
  sample_frames = 0
838
  for _ in range(num_media):
839
  # grid_thw is (N, 3) where first dim is t (num_frames)
@@ -845,38 +845,38 @@ class MossVLProcessor(ProcessorMixin):
845
  sample_frames += t
846
  media_idx += 1
847
  total_frames_per_sample.append(sample_frames)
848
-
849
  max_num_frames = max(total_frames_per_sample) if total_frames_per_sample else 0
850
-
851
  if max_num_frames == 0:
852
  return None
853
-
854
  # Vectorized implementation for speed
855
-
856
  # 1. Pad input_ids to create a tensor
857
  # We use -1 as pad value since token ids are positive
858
  input_ids_tensor = torch.full((batch_size, max_text_len), -1, dtype=torch.long)
859
  for b, ids in enumerate(input_ids):
860
  l = len(ids)
861
  input_ids_tensor[b, :l] = torch.tensor(ids, dtype=torch.long)
862
-
863
  # 2. Identify image tokens
864
  is_image_token = (input_ids_tensor == self.image_token_id)
865
-
866
  # 3. Compute cumulative image tokens (how many image tokens appeared up to position t)
867
  # shape: (batch_size, text_len)
868
  cum_image_tokens = is_image_token.cumsum(dim=1)
869
-
870
  # 4. Create frame indices
871
  # shape: (1, 1, max_num_frames)
872
  frame_indices = torch.arange(max_num_frames).reshape(1, 1, -1)
873
-
874
  # 5. Determine visibility based on causal relationship
875
  # Text at `t` sees frame `i` if `cum_image_tokens[t] > i`
876
  # Because if frame `i` is the (i+1)-th image token, it becomes visible when count reaches i+1
877
  # shape: (batch_size, text_len, max_num_frames)
878
  visible_mask = cum_image_tokens.unsqueeze(-1) > frame_indices
879
-
880
  # 6. Apply attention_mask if provided
881
  if attention_mask is not None:
882
  # Convert to tensor if needed
@@ -888,25 +888,25 @@ class MossVLProcessor(ProcessorMixin):
888
  for b, mask_row in enumerate(attention_mask):
889
  l = len(mask_row)
890
  attn_mask_tensor[b, :l] = torch.tensor(mask_row, dtype=torch.long)
891
-
892
  # shape: (batch_size, text_len, 1)
893
  valid_text = (attn_mask_tensor.unsqueeze(-1) == 1)
894
  visible_mask = visible_mask & valid_text
895
-
896
  # 7. Mask out frames that don't exist for a sample
897
  # shape: (batch_size, 1, 1)
898
  total_frames_tensor = torch.tensor(total_frames_per_sample).reshape(batch_size, 1, 1)
899
  # shape: (batch_size, 1, max_num_frames)
900
  valid_frames = frame_indices < total_frames_tensor
901
-
902
  visible_mask = visible_mask & valid_frames
903
-
904
  # 8. Create final mask (True for masked, False for visible)
905
  mask = ~visible_mask
906
-
907
  # 9. Add channel dimension: (batch_size, 1, text_len, max_num_frames)
908
  mask = mask.unsqueeze(1)
909
-
910
  return mask
911
 
912
  def _replace_and_update_spans(
@@ -919,14 +919,14 @@ class MossVLProcessor(ProcessorMixin):
919
  ) -> tuple:
920
  """
921
  Replace occurrences of old_str with new_str and update spans accordingly.
922
-
923
  Args:
924
  text: The text to perform replacement on
925
  old_str: String to be replaced
926
  new_str: String to replace with
927
  spans: List of [start, end] spans to update (modified in place)
928
  replace_count: Maximum number of replacements (-1 for all)
929
-
930
  Returns:
931
  Tuple of (new_text, updated_spans)
932
  """
@@ -934,14 +934,14 @@ class MossVLProcessor(ProcessorMixin):
934
  result_text = text
935
  count = 0
936
  search_start = 0
937
-
938
  while True:
939
  pos = result_text.find(old_str, search_start)
940
  if pos == -1:
941
  break
942
  if replace_count != -1 and count >= replace_count:
943
  break
944
-
945
  # Update all spans that come after this position
946
  for span in spans:
947
  if span[0] > pos:
@@ -951,12 +951,12 @@ class MossVLProcessor(ProcessorMixin):
951
  elif span[1] > pos:
952
  # Span ends after replacement point (spans the replacement)
953
  span[1] += delta
954
-
955
  # Perform the replacement
956
  result_text = result_text[:pos] + new_str + result_text[pos + len(old_str):]
957
  search_start = pos + len(new_str)
958
  count += 1
959
-
960
  return result_text, spans
961
 
962
  def _create_labels_from_spans(
@@ -969,7 +969,7 @@ class MossVLProcessor(ProcessorMixin):
969
  ) -> List[List[int]]:
970
  """
971
  Create labels from spans and offset_mapping.
972
-
973
  Args:
974
  input_ids: Tokenized input ids
975
  offset_mapping: Character offsets for each token from tokenizer (special tokens included)
@@ -978,10 +978,10 @@ class MossVLProcessor(ProcessorMixin):
978
  mask_token_ids: Set of token ids that should always be masked (set to ignore_index)
979
  in labels, regardless of whether they fall inside a span.
980
  Defaults to self.mask_token_ids if not provided.
981
-
982
  Returns:
983
  labels: List of label ids, same shape as input_ids
984
-
985
  Note:
986
  - Tokenizer's offset_mapping already includes correct offsets for special tokens in text
987
  - Only need to mask tokens inside <|vision_start|>...<|vision_end|>
@@ -990,20 +990,20 @@ class MossVLProcessor(ProcessorMixin):
990
  """
991
  if mask_token_ids is None:
992
  mask_token_ids = self.mask_token_ids
993
-
994
  batch_labels = []
995
-
996
  for batch_idx in range(len(input_ids)):
997
  ids = input_ids[batch_idx]
998
  offsets = offset_mapping[batch_idx]
999
  spans = labels_spans[batch_idx]
1000
-
1001
  labels = [ignore_index] * len(ids)
1002
-
1003
  # Process each span: find token range and set labels
1004
  for span_start, span_end in spans:
1005
  in_vision = False
1006
-
1007
  # Find tokens that overlap with this span
1008
  for token_idx, (token_id, (char_start, char_end)) in enumerate(zip(ids, offsets)):
1009
  # Skip tokens completely before this span
@@ -1012,7 +1012,7 @@ class MossVLProcessor(ProcessorMixin):
1012
  # Stop when tokens are completely after this span
1013
  if char_start >= span_end:
1014
  break
1015
-
1016
  # Token overlaps with span, process it
1017
  # Track vision region: <|vision_start|> ... <|vision_end|>
1018
  if token_id == self.vision_start_token_id:
@@ -1021,34 +1021,34 @@ class MossVLProcessor(ProcessorMixin):
1021
  if token_id == self.vision_end_token_id:
1022
  in_vision = False
1023
  continue
1024
-
1025
  # Skip tokens inside vision region
1026
  if in_vision:
1027
  continue
1028
-
1029
  # Always mask special tokens that should never have labels
1030
  if token_id in mask_token_ids:
1031
  continue
1032
-
1033
  # Set label for this token
1034
  labels[token_idx] = token_id
1035
-
1036
  batch_labels.append(labels)
1037
-
1038
  return batch_labels
1039
 
1040
  def _calculate_timestamps(
1041
- self,
1042
- frames_indices: Optional[Union[List[int], np.ndarray]],
1043
- total_num_frames: int,
1044
- video_fps: float,
1045
- duration: float,
1046
  merge_size: int = 1,
1047
  actual_timestamps: Optional[List[float]] = None
1048
  ):
1049
  """
1050
  Calculate timestamps for video frames.
1051
-
1052
  Args:
1053
  frames_indices: Actual frame indices extracted (if available)
1054
  total_num_frames: Total number of sampled frames
@@ -1056,25 +1056,25 @@ class MossVLProcessor(ProcessorMixin):
1056
  duration: Video duration in seconds
1057
  merge_size: Temporal merge size
1058
  actual_timestamps: Pre-calculated actual timestamps (for segments)
1059
-
1060
  Returns:
1061
  List of timestamps (one per merged temporal patch)
1062
  """
1063
  # If actual timestamps are provided (from segment), use them directly
1064
  if actual_timestamps is not None:
1065
  timestamps = list(actual_timestamps)
1066
-
1067
  # Pad timestamps to be multiple of merge_size
1068
  if len(timestamps) % merge_size != 0:
1069
  timestamps.extend([timestamps[-1]] * (merge_size - len(timestamps) % merge_size))
1070
-
1071
  # Frames are merged by merge_size, so we average the timestamps within each temporal patch
1072
  timestamps = [
1073
- (timestamps[i] + timestamps[i + merge_size - 1]) / 2
1074
  for i in range(0, len(timestamps), merge_size)
1075
  ]
1076
  return timestamps
1077
-
1078
  # Use frames_indices if available, otherwise generate uniformly sampled indices
1079
  if frames_indices is not None:
1080
  if isinstance(frames_indices, np.ndarray):
@@ -1088,24 +1088,24 @@ class MossVLProcessor(ProcessorMixin):
1088
  else:
1089
  # Uniformly sample frames across the video duration
1090
  indices = np.linspace(0, duration * video_fps - 1, total_num_frames).astype(np.int32).tolist()
1091
-
1092
  # Pad indices to be multiple of merge_size
1093
  if len(indices) % merge_size != 0:
1094
  indices.extend([indices[-1]] * (merge_size - len(indices) % merge_size))
1095
-
1096
  # Convert frame indices to timestamps
1097
  timestamps = [idx / video_fps for idx in indices]
1098
-
1099
  # Frames are merged by merge_size, so we average the timestamps within each temporal patch
1100
  timestamps = [
1101
- (timestamps[i] + timestamps[i + merge_size - 1]) / 2
1102
  for i in range(0, len(timestamps), merge_size)
1103
  ]
1104
  return timestamps
1105
 
1106
  def batch_decode(self, *args, **kwargs):
1107
  """
1108
- This method forwards all its arguments to the tokenizer's batch_decode.
1109
  Please refer to the docstring of this method for more information.
1110
  """
1111
  return self.tokenizer.batch_decode(*args, **kwargs)
@@ -1125,7 +1125,7 @@ class MossVLProcessor(ProcessorMixin):
1125
 
1126
  Args:
1127
  generated_outputs (`torch.Tensor` or `np.ndarray`):
1128
- The output of the model `generate` function. The output is expected to be a tensor
1129
  of shape `(batch_size, sequence_length)` or `(sequence_length,)`.
1130
  skip_special_tokens (`bool`, *optional*, defaults to `True`):
1131
  Whether or not to remove special tokens in the output.
 
49
  """
50
  # Multi-image batch total pixels limit (read from config)
51
  multi_image_max_pixels = None
52
+
53
 
54
  def _preprocess(
55
  self,
 
70
  **kwargs,
71
  ):
72
  """Override _preprocess to use custom smart_resize with batch-level max_pixels.
73
+
74
  multi_image_max_pixels is treated as a batch-level total budget, proportionally allocated
75
  to each image based on its original pixel count. min_pixels remains a per-image
76
  constraint. multi_image_max_pixels can be configured separately from longest_edge.
 
79
  max_pixels = size["longest_edge"] # Per-image upper limit
80
  # Use multi_image_max_pixels if configured, otherwise fall back to longest_edge
81
  multi_image_max_pixels = getattr(self, "multi_image_max_pixels", None) or max_pixels
82
+
83
  # Calculate total original pixels across all images in the batch
84
  # This is used to proportionally allocate max_pixels to each image
85
  total_original_pixels = sum(img.shape[-2] * img.shape[-1] for img in images)
 
98
  proportional_max_pixels = int(multi_image_max_pixels * proportion)
99
  else:
100
  proportional_max_pixels = multi_image_max_pixels
101
+
102
  # Ensure proportional max_pixels is within [min_pixels, max_pixels] range
103
  # min_pixels: per-image lower limit (shortest_edge)
104
  # max_pixels: per-image upper limit (longest_edge)
105
  proportional_max_pixels = max(proportional_max_pixels, min_pixels)
106
  proportional_max_pixels = min(proportional_max_pixels, max_pixels)
107
+
108
  resized_height, resized_width = smart_resize(
109
  height,
110
  width,
 
119
  )
120
  resized_images_grouped[shape] = stacked_images
121
  resized_images = reorder_images(resized_images_grouped, grouped_images_index)
122
+
123
  # Warn if multi-image batch exceeds multi_image_max_pixels due to min_pixels constraint
124
  if len(images) > 1:
125
  total_resized_pixels = sum(img.shape[-2] * img.shape[-1] for img in resized_images)
 
192
  """
193
  Convert various tensor types to numpy array.
194
  Supports torch.Tensor, tf.Tensor, jax.Array, np.ndarray, lists, and primitives.
195
+
196
  Args:
197
  x: Input value that can be a tensor from various frameworks or a Python primitive
198
+
199
  Returns:
200
  np.ndarray: NumPy array representation of the input
201
  """
202
  # Already numpy
203
  if isinstance(x, np.ndarray):
204
  return x
205
+
206
  # Torch tensor or TensorFlow tensor (both have .numpy() method)
207
  if hasattr(x, 'numpy'):
208
  # For torch tensors on CUDA, need to move to CPU first
 
210
  return x.cpu().numpy()
211
  # For TensorFlow or already on CPU
212
  return x.numpy()
213
+
214
  # JAX arrays and other array-like objects that support __array__ protocol
215
  if hasattr(x, '__array__'):
216
  return np.asarray(x)
217
+
218
  # Python primitives (list, tuple, int, float)
219
  return np.array(x)
220
 
 
334
  tokenizer_class = ("Qwen2Tokenizer", "Qwen2TokenizerFast")
335
 
336
  def __init__(
337
+ self,
338
+ image_processor=None,
339
+ tokenizer=None,
340
+ video_processor=None,
341
  chat_template=None,
342
  **kwargs
343
  ):
344
  super().__init__(image_processor, tokenizer, video_processor, chat_template=chat_template)
345
+
346
 
347
  self.image_token = "<|image_pad|>" if not hasattr(tokenizer, "image_token") else tokenizer.image_token
348
  self.video_token = "<|video_pad|>" if not hasattr(tokenizer, "video_token") else tokenizer.video_token
349
 
350
+
351
  self.image_token_id = (
352
  tokenizer.image_token_id
353
  if getattr(tokenizer, "image_token_id", None)
 
358
  if getattr(tokenizer, "video_token_id", None)
359
  else tokenizer.convert_tokens_to_ids(self.video_token)
360
  )
361
+
362
  self.vision_start_token = (
363
  "<|vision_start|>" if not hasattr(tokenizer, "vision_start_token") else tokenizer.vision_start_token
364
  )
 
372
 
373
  self.time_start_token = "<|time_start|>"
374
  self.time_end_token = "<|time_end|>"
375
+
376
  # EOS token for labels generation (assistant's response should end with this)
377
  self.im_end_token = "<|im_end|>"
378
  self.im_end_token_id = tokenizer.convert_tokens_to_ids(self.im_end_token)
379
+
380
  # Vision-related token ids (all should be masked in labels)
381
  self.vision_start_token_id = tokenizer.convert_tokens_to_ids(self.vision_start_token)
382
  self.vision_end_token_id = tokenizer.convert_tokens_to_ids(self.vision_end_token)
383
+
384
  # Token ids that should always be masked in labels (e.g. <|image_pad|>)
385
  self.mask_token_ids = {self.image_token_id}
386
 
 
444
  tokenizer_init_kwargs=self.tokenizer.init_kwargs,
445
  **kwargs,
446
  )
447
+
448
  # Step 1: Process images if provided
449
  if images is not None:
450
  images_kwargs = output_kwargs["images_kwargs"].copy()
 
474
  # Step 3: Process text with placeholder replacement
475
  if text is None or (isinstance(text, str) and len(text.strip()) == 0):
476
  raise ValueError("Text input is required for MossVL processor and cannot be empty.")
477
+
478
  if not isinstance(text, list):
479
  text = [text]
480
+
481
  text = text.copy() # Copy to avoid in-place modifications
482
+
483
  # Prepare labels_spans if provided
484
  # labels_spans format: List[List[List[int]]] - batch of samples, each sample has multiple spans
485
  # Each span is [start, end] (list, not tuple) so it can be modified in place
 
496
  has_images = images is not None and "pixel_values" in image_inputs
497
  has_videos = videos is not None and "pixel_values_videos" in videos_inputs
498
  needs_reorder = has_images and has_videos
499
+
500
  image_pixel_values_list = []
501
  video_pixel_values_list = []
502
+
503
  # Step 3.0: Record the order of media in original text (before replacement)
504
  # This will be used later to correctly order pixel_values and grid_thw
505
  media_order_per_sample = []
 
510
  while pos < len(temp_text):
511
  img_pos = temp_text.find(self.image_placeholder, pos)
512
  vid_pos = temp_text.find(self.video_placeholder, pos)
513
+
514
  if img_pos == -1 and vid_pos == -1:
515
  break
516
+
517
  if img_pos != -1 and (vid_pos == -1 or img_pos < vid_pos):
518
  media_order.append(("image", img_pos))
519
  pos = img_pos + len(self.image_placeholder)
520
  elif vid_pos != -1:
521
  media_order.append(("video", vid_pos))
522
  pos = vid_pos + len(self.video_placeholder)
523
+
524
  media_order_per_sample.append(media_order)
525
+
526
  # Step 3.0.1: Check if any sample has no media (empty samples need blank image)
527
  # If there are empty samples, we need to enter slow path to handle them properly
528
  has_empty_samples = any(len(order) == 0 for order in media_order_per_sample)
529
  if has_empty_samples:
530
  needs_reorder = True
531
+
532
  # Split pixel values for reordering if needed
533
  if needs_reorder:
534
  if has_images:
 
559
  video_pixel_values_list = _split_array_or_tensor(
560
  flat_video_values, split_indices
561
  )
562
+
563
  # Step 3.1: Replace placeholders (simple replacement, no expansion yet)
564
  # In MossVL, one image placeholder = one image token
565
  # One video placeholder = one video token (will be expanded later)
 
576
  else:
577
  text[i] = text[i].replace(self.image_placeholder, self.image_token)
578
  text[i] = text[i].replace(self.video_placeholder, self.video_token)
579
+
580
+ # Step 3.2: Validate token counts
581
  n_images_in_text = [t.count(self.image_token) for t in text]
582
  n_videos_in_text = [t.count(self.video_token) for t in text]
583
+
584
  # Count placeholders in text
585
  total_images_in_text = sum(n_images_in_text)
586
  total_videos_in_text = sum(n_videos_in_text)
587
+
588
  # Count actual images and videos provided
589
  total_images_provided = len(image_grid_thw) if image_grid_thw is not None else 0
590
  total_videos_provided = len(video_grid_thw) if video_grid_thw is not None else 0
591
+
592
  # Validate image counts
593
  if total_images_in_text != total_images_provided:
594
  raise ValueError(
595
  "Number of image tokens does not match number of images provided. "
596
  f"Found {total_images_in_text} image tokens in text and {total_images_provided} images."
597
  )
598
+
599
  # Validate video counts
600
  if total_videos_in_text != total_videos_provided:
601
  raise ValueError(
602
  "Number of video tokens does not match number of videos provided. "
603
  f"Found {total_videos_in_text} video tokens in text and {total_videos_provided} videos."
604
  )
605
+
606
  # Step 3.3: Expand video tokens with timestamps
607
  # Now expand each video token to multiple tokens (one per frame) with timestamps
608
  if video_grid_thw is not None:
 
640
  video_tokens.append(
641
  f"{self.time_start_token}{curr_time:.1f} seconds{self.time_end_token}{self.image_token}"
642
  )
643
+
644
  # Wrap the entire video sequence with vision_start and vision_end tokens
645
  video_placeholder = f"{self.vision_start_token}{''.join(video_tokens)}{self.vision_end_token}"
646
+
647
  # Replace the video token with expanded sequence and update spans if needed
648
  if should_create_labels:
649
  text[i], labels_spans[i] = self._replace_and_update_spans(
 
652
  else:
653
  text[i] = text[i].replace(self.video_token, video_placeholder, 1)
654
  index += 1
655
+
656
 
657
 
658
  # Step 4: Tokenize text
659
  return_tensors = output_kwargs["text_kwargs"].pop("return_tensors", None)
660
  return_mm_token_type_ids = output_kwargs["text_kwargs"].pop("return_mm_token_type_ids", None)
661
+
662
  # Request offset_mapping if we need to create labels
663
  if should_create_labels:
664
  output_kwargs["text_kwargs"]["return_offsets_mapping"] = True
665
+
666
  text_inputs = self.tokenizer(text, **output_kwargs["text_kwargs"])
667
+
668
  # ignore check_special_mm_tokens nums in test and input ids.
669
  # self._check_special_mm_tokens(text, text_inputs, modalities=["image", "video"])
670
+
671
  # Create labels if labels_spans was provided
672
  if should_create_labels:
673
  offset_mapping = text_inputs.pop("offset_mapping")
 
687
  # Step 5: Concatenate pixel_values and grid_thw in sequence order
688
  # Prepare output
689
  output_data = {**text_inputs}
690
+
691
  if not needs_reorder:
692
  # Fast path: only one media type, no reordering needed
693
  final_pixel_values = []
694
  final_grid_thw = []
695
+
696
  if has_images:
697
  final_pixel_values.append(image_inputs["pixel_values"])
698
  final_grid_thw.extend(image_grid_thw)
699
+
700
  if has_videos:
701
  final_pixel_values.append(videos_inputs["pixel_values_videos"])
702
  final_grid_thw.extend(video_grid_thw)
703
+
704
  if final_pixel_values:
705
  output_data["pixel_values"] = np.concatenate(final_pixel_values, axis=0) if len(final_pixel_values) > 1 else final_pixel_values[0]
706
+
707
  if final_grid_thw:
708
  output_data["grid_thw"] = np.stack(final_grid_thw, axis=0)
709
+
710
  # Calculate media_nums_per_sample
711
  media_nums_per_sample = []
712
  for batch_idx in range(len(text)):
713
  media_order = media_order_per_sample[batch_idx]
714
  media_nums_per_sample.append(len(media_order) if len(media_order) > 0 else 1)
715
+
716
  # Don't add media_nums_per_sample to output_data yet
717
  # Will add it after BatchFeature to keep it as list
718
+
719
  else:
720
  # Slow path: both images and videos exist, need reordering
721
  final_pixel_values = []
722
  final_grid_thw = []
723
  media_nums_per_sample = []
724
+
725
  # Global indices to track position in flattened image/video arrays
726
  global_image_idx = 0
727
  global_video_idx = 0
728
+
729
  for batch_idx in range(len(text)):
730
  # Use the recorded media order from Step 3.0
731
  media_order = media_order_per_sample[batch_idx]
732
+
733
  if len(media_order) == 0:
734
  # If no media provided for this sample, add a blank image
735
  media_nums_per_sample.append(1)
 
737
  patch_size = getattr(self.image_processor, "patch_size", None) or 16
738
  temporal_patch_size = getattr(self.image_processor, "temporal_patch_size", None) or 1
739
  merge_size = getattr(self.image_processor, "merge_size", None) or 2
740
+
741
  factor = patch_size * merge_size
742
  side = int(np.ceil(np.sqrt(min_pixels) / factor) * factor)
743
  grid_h = side // patch_size
744
  grid_w = side // patch_size
745
  grid_t = 1
746
+
747
  # Channel = 3 (RGB)
748
  channel = 3
749
  dim = channel * temporal_patch_size * patch_size * patch_size
750
  num_patches = grid_t * grid_h * grid_w
751
+
752
  blank_pixel_values = np.zeros((num_patches, dim), dtype=np.float32)
753
  blank_grid_thw = np.array([grid_t, grid_h, grid_w], dtype=np.int64)
754
+
755
  final_pixel_values.append(blank_pixel_values)
756
  final_grid_thw.append(blank_grid_thw)
757
  else:
758
  media_nums_per_sample.append(len(media_order))
759
+
760
  # Collect media data according to the recorded order
761
  for media_type, _ in media_order:
762
  if media_type == "image" and image_grid_thw is not None:
 
771
  final_pixel_values.append(video_pixel_values_list[global_video_idx])
772
  final_grid_thw.append(video_grid_thw[global_video_idx])
773
  global_video_idx += 1
774
+
775
  # Concatenate/stack to unified format
776
  if final_pixel_values:
777
  output_data["pixel_values"] = _concat_array_or_tensor(
778
  final_pixel_values, axis=0
779
  )
780
+
781
  if final_grid_thw:
782
  output_data["grid_thw"] = _stack_array_or_tensor(
783
  final_grid_thw, axis=0
784
  )
785
+
786
  # Don't add media_nums_per_sample to output_data yet
787
  # Will add it after BatchFeature to keep it as list
788
 
 
795
  output_data.get("attention_mask", None)
796
  )
797
  output_data["cross_attention_mask"] = cross_attention_mask
798
+
799
  # Add labels to output if created
800
  if should_create_labels:
801
  output_data["labels"] = labels
802
 
803
  # BatchFeature will handle conversion to pt/tf/jax/np based on tensor_type
804
  batch_feature = BatchFeature(data=output_data, tensor_type=return_tensors)
805
+
806
  # Add media_nums_per_sample after BatchFeature to keep it as list (not tensor)
807
  if media_nums_per_sample:
808
  batch_feature["media_nums_per_sample"] = media_nums_per_sample
809
+
810
  return batch_feature
811
 
812
  def _create_cross_attention_mask(self, input_ids, grid_thw, media_nums_per_sample, attention_mask=None):
 
815
  Video frames are treated as individual images.
816
  Mask values: True for masked, False for visible.
817
  Causal masking: text can see images that appear at or before the text position.
818
+
819
  Args:
820
  input_ids: List of token ids
821
  grid_thw: Grid sizes for each media item
 
824
  """
825
  batch_size = len(input_ids)
826
  max_text_len = max(len(ids) for ids in input_ids)
827
+
828
  # Calculate total frames per sample to find max_num_frames
829
  total_frames_per_sample = []
830
  media_idx = 0
 
833
  if num_media == 0:
834
  total_frames_per_sample.append(0)
835
  continue
836
+
837
  sample_frames = 0
838
  for _ in range(num_media):
839
  # grid_thw is (N, 3) where first dim is t (num_frames)
 
845
  sample_frames += t
846
  media_idx += 1
847
  total_frames_per_sample.append(sample_frames)
848
+
849
  max_num_frames = max(total_frames_per_sample) if total_frames_per_sample else 0
850
+
851
  if max_num_frames == 0:
852
  return None
853
+
854
  # Vectorized implementation for speed
855
+
856
  # 1. Pad input_ids to create a tensor
857
  # We use -1 as pad value since token ids are positive
858
  input_ids_tensor = torch.full((batch_size, max_text_len), -1, dtype=torch.long)
859
  for b, ids in enumerate(input_ids):
860
  l = len(ids)
861
  input_ids_tensor[b, :l] = torch.tensor(ids, dtype=torch.long)
862
+
863
  # 2. Identify image tokens
864
  is_image_token = (input_ids_tensor == self.image_token_id)
865
+
866
  # 3. Compute cumulative image tokens (how many image tokens appeared up to position t)
867
  # shape: (batch_size, text_len)
868
  cum_image_tokens = is_image_token.cumsum(dim=1)
869
+
870
  # 4. Create frame indices
871
  # shape: (1, 1, max_num_frames)
872
  frame_indices = torch.arange(max_num_frames).reshape(1, 1, -1)
873
+
874
  # 5. Determine visibility based on causal relationship
875
  # Text at `t` sees frame `i` if `cum_image_tokens[t] > i`
876
  # Because if frame `i` is the (i+1)-th image token, it becomes visible when count reaches i+1
877
  # shape: (batch_size, text_len, max_num_frames)
878
  visible_mask = cum_image_tokens.unsqueeze(-1) > frame_indices
879
+
880
  # 6. Apply attention_mask if provided
881
  if attention_mask is not None:
882
  # Convert to tensor if needed
 
888
  for b, mask_row in enumerate(attention_mask):
889
  l = len(mask_row)
890
  attn_mask_tensor[b, :l] = torch.tensor(mask_row, dtype=torch.long)
891
+
892
  # shape: (batch_size, text_len, 1)
893
  valid_text = (attn_mask_tensor.unsqueeze(-1) == 1)
894
  visible_mask = visible_mask & valid_text
895
+
896
  # 7. Mask out frames that don't exist for a sample
897
  # shape: (batch_size, 1, 1)
898
  total_frames_tensor = torch.tensor(total_frames_per_sample).reshape(batch_size, 1, 1)
899
  # shape: (batch_size, 1, max_num_frames)
900
  valid_frames = frame_indices < total_frames_tensor
901
+
902
  visible_mask = visible_mask & valid_frames
903
+
904
  # 8. Create final mask (True for masked, False for visible)
905
  mask = ~visible_mask
906
+
907
  # 9. Add channel dimension: (batch_size, 1, text_len, max_num_frames)
908
  mask = mask.unsqueeze(1)
909
+
910
  return mask
911
 
912
  def _replace_and_update_spans(
 
919
  ) -> tuple:
920
  """
921
  Replace occurrences of old_str with new_str and update spans accordingly.
922
+
923
  Args:
924
  text: The text to perform replacement on
925
  old_str: String to be replaced
926
  new_str: String to replace with
927
  spans: List of [start, end] spans to update (modified in place)
928
  replace_count: Maximum number of replacements (-1 for all)
929
+
930
  Returns:
931
  Tuple of (new_text, updated_spans)
932
  """
 
934
  result_text = text
935
  count = 0
936
  search_start = 0
937
+
938
  while True:
939
  pos = result_text.find(old_str, search_start)
940
  if pos == -1:
941
  break
942
  if replace_count != -1 and count >= replace_count:
943
  break
944
+
945
  # Update all spans that come after this position
946
  for span in spans:
947
  if span[0] > pos:
 
951
  elif span[1] > pos:
952
  # Span ends after replacement point (spans the replacement)
953
  span[1] += delta
954
+
955
  # Perform the replacement
956
  result_text = result_text[:pos] + new_str + result_text[pos + len(old_str):]
957
  search_start = pos + len(new_str)
958
  count += 1
959
+
960
  return result_text, spans
961
 
962
  def _create_labels_from_spans(
 
969
  ) -> List[List[int]]:
970
  """
971
  Create labels from spans and offset_mapping.
972
+
973
  Args:
974
  input_ids: Tokenized input ids
975
  offset_mapping: Character offsets for each token from tokenizer (special tokens included)
 
978
  mask_token_ids: Set of token ids that should always be masked (set to ignore_index)
979
  in labels, regardless of whether they fall inside a span.
980
  Defaults to self.mask_token_ids if not provided.
981
+
982
  Returns:
983
  labels: List of label ids, same shape as input_ids
984
+
985
  Note:
986
  - Tokenizer's offset_mapping already includes correct offsets for special tokens in text
987
  - Only need to mask tokens inside <|vision_start|>...<|vision_end|>
 
990
  """
991
  if mask_token_ids is None:
992
  mask_token_ids = self.mask_token_ids
993
+
994
  batch_labels = []
995
+
996
  for batch_idx in range(len(input_ids)):
997
  ids = input_ids[batch_idx]
998
  offsets = offset_mapping[batch_idx]
999
  spans = labels_spans[batch_idx]
1000
+
1001
  labels = [ignore_index] * len(ids)
1002
+
1003
  # Process each span: find token range and set labels
1004
  for span_start, span_end in spans:
1005
  in_vision = False
1006
+
1007
  # Find tokens that overlap with this span
1008
  for token_idx, (token_id, (char_start, char_end)) in enumerate(zip(ids, offsets)):
1009
  # Skip tokens completely before this span
 
1012
  # Stop when tokens are completely after this span
1013
  if char_start >= span_end:
1014
  break
1015
+
1016
  # Token overlaps with span, process it
1017
  # Track vision region: <|vision_start|> ... <|vision_end|>
1018
  if token_id == self.vision_start_token_id:
 
1021
  if token_id == self.vision_end_token_id:
1022
  in_vision = False
1023
  continue
1024
+
1025
  # Skip tokens inside vision region
1026
  if in_vision:
1027
  continue
1028
+
1029
  # Always mask special tokens that should never have labels
1030
  if token_id in mask_token_ids:
1031
  continue
1032
+
1033
  # Set label for this token
1034
  labels[token_idx] = token_id
1035
+
1036
  batch_labels.append(labels)
1037
+
1038
  return batch_labels
1039
 
1040
  def _calculate_timestamps(
1041
+ self,
1042
+ frames_indices: Optional[Union[List[int], np.ndarray]],
1043
+ total_num_frames: int,
1044
+ video_fps: float,
1045
+ duration: float,
1046
  merge_size: int = 1,
1047
  actual_timestamps: Optional[List[float]] = None
1048
  ):
1049
  """
1050
  Calculate timestamps for video frames.
1051
+
1052
  Args:
1053
  frames_indices: Actual frame indices extracted (if available)
1054
  total_num_frames: Total number of sampled frames
 
1056
  duration: Video duration in seconds
1057
  merge_size: Temporal merge size
1058
  actual_timestamps: Pre-calculated actual timestamps (for segments)
1059
+
1060
  Returns:
1061
  List of timestamps (one per merged temporal patch)
1062
  """
1063
  # If actual timestamps are provided (from segment), use them directly
1064
  if actual_timestamps is not None:
1065
  timestamps = list(actual_timestamps)
1066
+
1067
  # Pad timestamps to be multiple of merge_size
1068
  if len(timestamps) % merge_size != 0:
1069
  timestamps.extend([timestamps[-1]] * (merge_size - len(timestamps) % merge_size))
1070
+
1071
  # Frames are merged by merge_size, so we average the timestamps within each temporal patch
1072
  timestamps = [
1073
+ (timestamps[i] + timestamps[i + merge_size - 1]) / 2
1074
  for i in range(0, len(timestamps), merge_size)
1075
  ]
1076
  return timestamps
1077
+
1078
  # Use frames_indices if available, otherwise generate uniformly sampled indices
1079
  if frames_indices is not None:
1080
  if isinstance(frames_indices, np.ndarray):
 
1088
  else:
1089
  # Uniformly sample frames across the video duration
1090
  indices = np.linspace(0, duration * video_fps - 1, total_num_frames).astype(np.int32).tolist()
1091
+
1092
  # Pad indices to be multiple of merge_size
1093
  if len(indices) % merge_size != 0:
1094
  indices.extend([indices[-1]] * (merge_size - len(indices) % merge_size))
1095
+
1096
  # Convert frame indices to timestamps
1097
  timestamps = [idx / video_fps for idx in indices]
1098
+
1099
  # Frames are merged by merge_size, so we average the timestamps within each temporal patch
1100
  timestamps = [
1101
+ (timestamps[i] + timestamps[i + merge_size - 1]) / 2
1102
  for i in range(0, len(timestamps), merge_size)
1103
  ]
1104
  return timestamps
1105
 
1106
  def batch_decode(self, *args, **kwargs):
1107
  """
1108
+ This method forwards all its arguments to the tokenizer's batch_decode.
1109
  Please refer to the docstring of this method for more information.
1110
  """
1111
  return self.tokenizer.batch_decode(*args, **kwargs)
 
1125
 
1126
  Args:
1127
  generated_outputs (`torch.Tensor` or `np.ndarray`):
1128
+ The output of the model `generate` function. The output is expected to be a tensor
1129
  of shape `(batch_size, sequence_length)` or `(sequence_length,)`.
1130
  skip_special_tokens (`bool`, *optional*, defaults to `True`):
1131
  Whether or not to remove special tokens in the output.
video_preprocessor_config.json CHANGED
@@ -7,7 +7,7 @@
7
  "longest_edge": 16777216,
8
  "shortest_edge": 4096
9
  },
10
- "video_max_pixels": 65536000,
11
  "patch_size": 16,
12
  "temporal_patch_size": 1,
13
  "merge_size": 2,
 
7
  "longest_edge": 16777216,
8
  "shortest_edge": 4096
9
  },
10
+ "video_max_pixels": 201326592,
11
  "patch_size": 16,
12
  "temporal_patch_size": 1,
13
  "merge_size": 2,
video_processing_moss_vl.py CHANGED
@@ -14,12 +14,15 @@
14
  # limitations under the License.
15
  """video processor class for Moss-VL."""
16
 
 
 
17
  import json
18
  import logging as system_logging
19
  import math
20
  import os
21
  import re
22
  import subprocess
 
23
  import traceback
24
  from functools import lru_cache
25
  from typing import Any, Dict, List, Optional, Union
@@ -49,10 +52,10 @@ TORCHCODEC_TIMESTAMP_EPSILON = 1e-6
49
  def check_video_for_extra_streams_and_errors(video_path: str) -> dict:
50
  """
51
  Check if video file has abnormal streams or errors reported by ffprobe.
52
-
53
  Args:
54
  video_path: Path to the video file.
55
-
56
  Returns:
57
  A dictionary containing:
58
  - 'has_extra_streams': bool, whether there are streams other than video and audio.
@@ -70,7 +73,7 @@ def check_video_for_extra_streams_and_errors(video_path: str) -> dict:
70
  'stream_details': [],
71
  'num_streams': 0
72
  }
73
-
74
  command = [
75
  "ffprobe",
76
  "-v", "error",
@@ -79,7 +82,7 @@ def check_video_for_extra_streams_and_errors(video_path: str) -> dict:
79
  "-of", "json",
80
  video_path
81
  ]
82
-
83
  try:
84
  process = subprocess.run(
85
  command,
@@ -88,12 +91,12 @@ def check_video_for_extra_streams_and_errors(video_path: str) -> dict:
88
  check=False
89
  )
90
  result['ffprobe_successful'] = (process.returncode == 0)
91
-
92
  if process.stderr:
93
  result['ffprobe_output_error'] = process.stderr
94
  unsupported_codec_pattern = re.compile(r"Unsupported codec with id \d+ for input stream \d+")
95
  result['unsupported_codec_errors'] = unsupported_codec_pattern.findall(process.stderr)
96
-
97
  if process.stdout:
98
  ffprobe_data = json.loads(process.stdout)
99
  if 'streams' in ffprobe_data:
@@ -104,7 +107,7 @@ def check_video_for_extra_streams_and_errors(video_path: str) -> dict:
104
  result['stream_details'].append({'index': stream_index, 'codec_type': stream_type})
105
  if stream_type not in ['video', 'audio']:
106
  result['has_extra_streams'] = True
107
-
108
  if 'format' in ffprobe_data and 'nb_streams' in ffprobe_data['format']:
109
  if result['num_streams'] == 0:
110
  result['num_streams'] = ffprobe_data['format']['nb_streams']
@@ -122,18 +125,18 @@ def check_video_for_extra_streams_and_errors(video_path: str) -> dict:
122
  except Exception as e:
123
  result['ffprobe_output_error'] = f"An unexpected error occurred: {e}"
124
  result['ffprobe_successful'] = False
125
-
126
  return result
127
 
128
 
129
  def remove_video_extra_stream_ffmpeg(input_video: str, output_video: str) -> bool:
130
  """
131
  Remove extra streams from video using ffmpeg.
132
-
133
  Args:
134
  input_video: Path to input video.
135
  output_video: Path to output video.
136
-
137
  Returns:
138
  bool: True if successful, False otherwise.
139
  """
@@ -149,7 +152,7 @@ def remove_video_extra_stream_ffmpeg(input_video: str, output_video: str) -> boo
149
  "-movflags", "faststart",
150
  output_video,
151
  ]
152
-
153
  try:
154
  subprocess.run(command_list, shell=False, check=True, capture_output=True)
155
  return True
@@ -165,13 +168,27 @@ def remove_video_extra_stream_ffmpeg(input_video: str, output_video: str) -> boo
165
  return False
166
 
167
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
168
  def clean_video_streams(video_path: str) -> str:
169
  """
170
  Clean video streams if extra streams are detected.
171
-
172
  Args:
173
  video_path: Path to the video file.
174
-
175
  Returns:
176
  str: Path to cleaned video (or original if no cleaning needed).
177
  """
@@ -182,12 +199,31 @@ def clean_video_streams(video_path: str) -> str:
182
  file_name_without_ext, file_ext = os.path.splitext(base_name)
183
  new_base_name = f"{file_name_without_ext}_fix{file_ext}"
184
  video_path_output = os.path.join(output_folder, new_base_name)
185
-
186
- process_flag = remove_video_extra_stream_ffmpeg(video_path, video_path_output)
187
- if not process_flag:
188
- logger.warning("Failed to remove extra streams with ffmpeg")
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
189
  return video_path
190
- return video_path_output
191
  return video_path
192
 
193
 
@@ -203,6 +239,9 @@ def clamp_timestamps_for_torchcodec(timestamps: List[float], torchcodec_metadata
203
  min_pts = torchcodec_metadata.begin_stream_seconds_from_content
204
  if min_pts is None:
205
  min_pts = 0.0
 
 
 
206
 
207
  max_pts_candidates = []
208
  if torchcodec_metadata.num_frames_from_content and torchcodec_metadata.average_fps:
@@ -216,24 +255,30 @@ def clamp_timestamps_for_torchcodec(timestamps: List[float], torchcodec_metadata
216
  max_pts_candidates.append(torchcodec_metadata.duration_seconds - TORCHCODEC_TIMESTAMP_EPSILON)
217
 
218
  if max_pts_candidates:
219
- max_pts = max(min_pts, min(max_pts_candidates))
220
- return [max(min_pts, min(float(t), max_pts)) for t in timestamps]
221
- if min_pts > 0:
222
- return [max(min_pts, float(t)) for t in timestamps]
223
  return [float(t) for t in timestamps]
224
 
225
 
226
  def split_indices(indices: List[Union[int, float]], num_chunks: int) -> List[List[Union[int, float]]]:
227
  """
228
  Split an index list into roughly equal chunks.
229
-
230
  Args:
231
  indices: List of indices to split.
232
  num_chunks: Number of chunks to create.
233
-
234
  Returns:
235
  List of index chunks.
236
  """
 
 
 
 
 
 
237
  chunk_size = len(indices) // num_chunks
238
  chunks = []
239
  for i in range(num_chunks - 1):
@@ -245,12 +290,12 @@ def split_indices(indices: List[Union[int, float]], num_chunks: int) -> List[Lis
245
  def decode_sequentially(indices: List[int], video_path: str, ffmpeg_threads: int = 0):
246
  """
247
  Decode frames sequentially from a video.
248
-
249
  Args:
250
  indices: List of frame indices to decode.
251
  video_path: Path to the video file.
252
  ffmpeg_threads: Number of ffmpeg threads to use.
253
-
254
  Returns:
255
  FrameBatch from torchcodec.
256
  """
@@ -264,20 +309,20 @@ def decode_sequentially(indices: List[int], video_path: str, ffmpeg_threads: int
264
  def decode_with_multithreading(indices: List[int], num_threads: int, video_path: str) -> dict:
265
  """
266
  Decode frames using multithreading with joblib.
267
-
268
  Args:
269
  indices: List of frame indices to decode.
270
  num_threads: Number of threads to use.
271
  video_path: Path to the video file.
272
-
273
  Returns:
274
  dict: Contains 'data', 'duration_seconds', 'pts_seconds' tensors.
275
  """
276
  chunks = split_indices(indices, num_chunks=num_threads)
277
- results = Parallel(n_jobs=num_threads, prefer="threads", verbose=0)(
278
  delayed(decode_sequentially)(chunk, video_path) for chunk in chunks
279
  )
280
-
281
  return {
282
  "data": torch.cat([frame_batch.data for frame_batch in results], dim=0),
283
  "duration_seconds": torch.cat([frame_batch.duration_seconds for frame_batch in results], dim=0),
@@ -288,12 +333,12 @@ def decode_with_multithreading(indices: List[int], num_threads: int, video_path:
288
  def decode_sequentially_timestamp(timestamp_list: List[float], video_path: str, ffmpeg_threads: int = 0):
289
  """
290
  Decode frames sequentially from a video based on timestamps.
291
-
292
  Args:
293
  timestamp_list: List of timestamps (in seconds) to decode.
294
  video_path: Path to the video file.
295
  ffmpeg_threads: Number of ffmpeg threads to use.
296
-
297
  Returns:
298
  FrameBatch from torchcodec.
299
  """
@@ -310,29 +355,29 @@ def decode_sequentially_timestamp(timestamp_list: List[float], video_path: str,
310
  def timestamp_decode_with_multithreading(timestamp_list: List[float], num_threads: int, video_path: str) -> dict:
311
  """
312
  Decode frames using multithreading based on timestamps.
313
-
314
  Args:
315
  timestamp_list: List of timestamps (in seconds) to decode.
316
  num_threads: Number of threads to use.
317
  video_path: Path to the video file.
318
-
319
  Returns:
320
  dict: Contains 'data', 'duration_seconds', 'pts_seconds' tensors.
321
  """
322
  chunks = split_indices(timestamp_list, num_chunks=num_threads)
323
- results = Parallel(n_jobs=num_threads, prefer="threads", verbose=0)(
324
  delayed(decode_sequentially_timestamp)(chunk, video_path) for chunk in chunks
325
  )
326
-
327
  # Concatenate results from all threads
328
  data_list = [frame_batch.data for frame_batch in results]
329
  duration_list = [frame_batch.duration_seconds for frame_batch in results]
330
  pts_list = [frame_batch.pts_seconds for frame_batch in results]
331
-
332
  if not data_list:
333
  logger.warning("No frames were successfully decoded.")
334
  return {"data": torch.empty(0), "duration_seconds": torch.empty(0), "pts_seconds": torch.empty(0)}
335
-
336
  return {
337
  "data": torch.cat(data_list, dim=0),
338
  "duration_seconds": torch.cat(duration_list, dim=0),
@@ -348,13 +393,13 @@ def extract_frames_with_torchcodec(
348
  ) -> Optional[dict]:
349
  """
350
  Extract frames from video using torchcodec with multithreading.
351
-
352
  Args:
353
  video_path: Path to the video file.
354
  sample_frames_count: Number of frames to sample.
355
  num_threads: Number of threads to use for extraction.
356
  sampling_method: Sampling method, either "index" (uniform frame indices) or "timestamp" (uniform timestamps).
357
-
358
  Returns:
359
  dict: Contains 'data' (N, C, H, W), 'duration_seconds' (N,), 'pts_seconds' (N,) tensors.
360
  Returns None if extraction fails.
@@ -366,24 +411,24 @@ def extract_frames_with_torchcodec(
366
 
367
 
368
  total_frames_in_video = metadata.num_frames_from_content
369
-
370
  effective_sample_count = min(sample_frames_count, total_frames_in_video)
371
  if effective_sample_count == 0:
372
  logger.error("Cannot extract frames: video has 0 frames or specified frame count is 0")
373
  return None
374
-
375
  # Generate uniform frame indices
376
  frame_indices = np.linspace(0, total_frames_in_video - 1, effective_sample_count).astype(np.int32)
377
  # Ensure indices are valid and remove duplicates
378
  frame_indices = np.unique(np.clip(frame_indices, 0, total_frames_in_video - 1))
379
-
380
  result = decode_with_multithreading(frame_indices.tolist(), num_threads=num_threads, video_path=video_path)
381
  # Add frame_indices to the result for later use
382
  result["frame_indices"] = frame_indices
383
  return result
384
 
385
 
386
-
387
  except Exception:
388
  traceback.print_exc()
389
  return None
@@ -519,6 +564,25 @@ class MossVLVideoProcessor(BaseVideoProcessor):
519
  return video_input["video_path"]
520
  return video_input
521
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
522
  def _get_video_duration_seconds(self, video_input: Union[str, Dict[str, Any]]) -> float:
523
  """Get video duration in seconds for weighted frame-budget allocation."""
524
  video_path = cached_clean_video_streams(self._get_video_path_from_input(video_input))
@@ -606,13 +670,13 @@ class MossVLVideoProcessor(BaseVideoProcessor):
606
  ) -> int:
607
  """
608
  Calculate the number of frames to sample using fps-based logic with min/max constraints.
609
-
610
  Logic:
611
  1. Calculate target_frames based on fps and video duration
612
  2. Apply min_frames and max_frames constraints
613
  3. Apply max_allowed_frames protection (rough cap from total video_max_pixels budget)
614
  4. Return the number of frames to sample
615
-
616
  Args:
617
  metadata (`VideoMetadata`):
618
  Metadata of the video containing information about total duration, fps and total number of frames.
@@ -632,11 +696,11 @@ class MossVLVideoProcessor(BaseVideoProcessor):
632
  raise ValueError("`num_frames` and `fps` are mutually exclusive arguments, please use only one!")
633
 
634
  total_num_frames = metadata.total_num_frames
635
-
636
  # Use provided min/max or fall back to defaults
637
  effective_min_frames = min_frames if min_frames is not None else self.min_frames
638
  effective_max_frames = max_frames if max_frames is not None else self.max_frames
639
-
640
  # Rough per-video frame cap derived from the multi-video total budget
641
  # (exact allocation happens later in _preprocess via weighted distribution)
642
  per_frame_min_pixels = self.size.get("shortest_edge", None) if self.size else None
@@ -644,7 +708,7 @@ class MossVLVideoProcessor(BaseVideoProcessor):
644
  if per_frame_min_pixels is not None and video_max_pixels is not None and per_frame_min_pixels > 0:
645
  max_allowed_frames = video_max_pixels // per_frame_min_pixels
646
  effective_max_frames = min(effective_max_frames, max_allowed_frames)
647
-
648
  # Get video duration
649
  if hasattr(metadata, 'duration') and metadata.duration is not None:
650
  duration = metadata.duration
@@ -662,12 +726,12 @@ class MossVLVideoProcessor(BaseVideoProcessor):
662
 
663
  # Use provided fps or default
664
  target_fps = fps if fps is not None else self.video_fps
665
-
666
  # Calculate target frames based on fps and duration
667
  if num_frames is None:
668
  # Calculate how many frames we should sample based on target fps
669
  target_total_frames = int(math.ceil(duration * target_fps - 1e-6))
670
-
671
  # Apply min/max constraints
672
  sample_frames = max(target_total_frames, effective_min_frames)
673
  sample_frames = min(sample_frames, effective_max_frames, total_num_frames)
@@ -816,14 +880,14 @@ class MossVLVideoProcessor(BaseVideoProcessor):
816
  ):
817
  """
818
  Fetch video frames for a specific segment.
819
-
820
  Args:
821
  video_path: Path to the video file
822
  segment: [start, end] for a segment (left-closed, right-open) or [time] for a single frame
823
  min_frames: Minimum frames for this segment (weighted). Defaults to self.min_frames. Must be >= 1.
824
  max_frames: Maximum frames for this segment (weighted). Defaults to self.max_frames. Must be >= 1.
825
  video_fps: Target frames per second for video sampling. If None, uses self.video_fps.
826
-
827
  Returns:
828
  Tuple of (video_tensor, video_metadata)
829
  """
@@ -832,21 +896,21 @@ class MossVLVideoProcessor(BaseVideoProcessor):
832
  max_frames = max(1, max_frames if max_frames is not None else self.max_frames)
833
  # Use provided video_fps or fall back to self.video_fps
834
  target_video_fps = video_fps if video_fps is not None else self.video_fps
835
-
836
- video_path = clean_video_streams(video_path)
837
  decoder = VideoDecoder(video_path, num_ffmpeg_threads=0)
838
  try:
839
  torchcodec_metadata = decoder.metadata
840
-
841
  video_fps = torchcodec_metadata.average_fps
842
-
843
  # Calculate duration
844
  duration = None
845
  if torchcodec_metadata.end_stream_seconds_from_content is not None and torchcodec_metadata.begin_stream_seconds_from_content is not None:
846
  duration = torchcodec_metadata.end_stream_seconds_from_content - torchcodec_metadata.begin_stream_seconds_from_content
847
  if duration is None or duration <= 0:
848
  duration = torchcodec_metadata.duration_seconds
849
-
850
  if len(segment) == 1:
851
  # Single frame at specified time
852
  actual_timestamps = self._clamp_timestamps_for_decoder([segment[0]], torchcodec_metadata)
@@ -857,26 +921,26 @@ class MossVLVideoProcessor(BaseVideoProcessor):
857
  # Segment [start, end) - left-closed, right-open interval
858
  start_time, end_time = segment
859
  segment_duration = end_time - start_time
860
-
861
  # Calculate number of frames to sample for this segment
862
  target_frames = int(math.ceil(segment_duration * target_video_fps))
863
  target_frames = max(target_frames, min_frames)
864
  target_frames = min(target_frames, max_frames)
865
-
866
  # Generate timestamps for uniform sampling within segment
867
  if target_frames == 1:
868
  actual_timestamps = [start_time] # Use start_time for single frame
869
  else:
870
  # Sample uniformly within [start, end), endpoint=False for left-closed right-open
871
  actual_timestamps = np.linspace(start_time, end_time, target_frames, endpoint=False).tolist()
872
-
873
  actual_timestamps = self._clamp_timestamps_for_decoder(actual_timestamps, torchcodec_metadata)
874
 
875
  # Use multithreading for extraction
876
  result = timestamp_decode_with_multithreading(actual_timestamps, self.num_extract_threads, video_path)
877
  video_tensor = result["data"]
878
  sample_count = len(actual_timestamps)
879
-
880
  # Create VideoMetadata
881
  video_metadata = VideoMetadata(
882
  total_num_frames=sample_count,
@@ -887,39 +951,40 @@ class MossVLVideoProcessor(BaseVideoProcessor):
887
  width=torchcodec_metadata.width,
888
  frames_indices=None
889
  )
890
-
891
  # Store actual timestamps as a custom attribute for _calculate_timestamps to use
892
  video_metadata.actual_timestamps = actual_timestamps
893
-
894
  return video_tensor, video_metadata
895
  finally:
896
  del decoder
897
 
898
  def fetch_videos(
899
- self,
900
- video_url_or_urls: Union[str, Dict[str, Any], List[Union[str, Dict[str, Any]]]],
901
- sample_indices_fn=None,
902
  video_fps: Optional[float] = None,
903
  min_frames: Optional[int] = None,
904
  max_frames: Optional[int] = None,
905
  ):
906
  """
907
  Override fetch_videos to use torchcodec for frame extraction.
908
-
909
  This method uses torchcodec with multithreading for efficient frame extraction.
910
  Frame count is calculated by the calculate_num_frames method
911
  (fps-based with min/max constraints).
912
-
913
  Args:
914
  video_url_or_urls: Can be one of:
915
  - str: Single video path
916
- - Dict: Video with segments {"video_path": str, "segments": List[List[float]]}
917
- - List[Union[str, Dict]]: List of video paths or segment dicts
 
918
  sample_indices_fn: (Not used) Kept for compatibility with base class signature.
919
  video_fps: Target frames per second for video sampling. If None, uses self.video_fps.
920
  min_frames: Minimum number of frames to sample. If None, uses self.min_frames.
921
  max_frames: Maximum number of frames to sample. If None, uses self.max_frames.
922
-
923
  Returns:
924
  Tuple of (videos, metadata) where videos are torch.Tensors and metadata are VideoMetadata objects.
925
  """
@@ -939,9 +1004,11 @@ class MossVLVideoProcessor(BaseVideoProcessor):
939
  effective_max_frames,
940
  )
941
  for x, allocated_max_frames in zip(video_url_or_urls, per_video_max_frames):
 
 
942
  result = self.fetch_videos(
943
- x,
944
- video_fps=effective_video_fps,
945
  min_frames=effective_min_frames,
946
  max_frames=allocated_max_frames,
947
  )
@@ -953,23 +1020,29 @@ class MossVLVideoProcessor(BaseVideoProcessor):
953
  all_videos.append(result[0])
954
  all_metadata.append(result[1])
955
  return all_videos, all_metadata
956
-
957
- # Handle dict with segments - returns lists (one per segment)
 
958
  if isinstance(video_url_or_urls, dict):
959
  video_path = video_url_or_urls["video_path"]
960
- segments = video_url_or_urls["segments"]
961
-
962
- return self._fetch_video_segments_batched(
963
- video_path,
964
- segments,
965
- min_frames=effective_min_frames,
966
- max_frames=effective_max_frames,
967
- video_fps=effective_video_fps,
968
- )
969
-
 
 
 
 
 
970
  # Single video path
971
  video_path = video_url_or_urls
972
-
973
  # Clean video streams first (remove extra streams if needed)
974
  video_path = cached_clean_video_streams(video_path)
975
 
@@ -982,13 +1055,13 @@ class MossVLVideoProcessor(BaseVideoProcessor):
982
  duration = None
983
  if torchcodec_metadata.end_stream_seconds_from_content is not None and torchcodec_metadata.begin_stream_seconds_from_content is not None:
984
  duration = torchcodec_metadata.end_stream_seconds_from_content - torchcodec_metadata.begin_stream_seconds_from_content
985
-
986
  if duration is None or duration <= 0:
987
  duration = torchcodec_metadata.duration_seconds
988
-
989
  # Use num_frames_from_content for accurate frame count (consistent with extraction)
990
  total_frames_in_video = torchcodec_metadata.num_frames_from_content
991
-
992
  # Create VideoMetadata object for sample_frames method
993
  temp_metadata = VideoMetadata(
994
  total_num_frames=total_frames_in_video,
@@ -999,31 +1072,31 @@ class MossVLVideoProcessor(BaseVideoProcessor):
999
  width=torchcodec_metadata.width,
1000
  frames_indices=None
1001
  )
1002
-
1003
  # Use calculate_num_frames method to get the number of frames to sample
1004
  sample_frames_count = self.calculate_num_frames(
1005
- temp_metadata,
1006
  fps=effective_video_fps,
1007
  min_frames=effective_min_frames,
1008
  max_frames=effective_max_frames,
1009
  )
1010
-
1011
  # Ensure sample count is valid
1012
  effective_sample_count = min(sample_frames_count, total_frames_in_video)
1013
  if effective_sample_count == 0:
1014
  raise ValueError(f"Cannot extract frames: video has 0 frames or specified frame count is 0")
1015
-
1016
  # Generate uniform frame indices
1017
  frame_indices = np.linspace(0, total_frames_in_video - 1, effective_sample_count).astype(np.int32)
1018
  # Ensure indices are valid and remove duplicates
1019
  frame_indices = np.unique(np.clip(frame_indices, 0, total_frames_in_video - 1))
1020
-
1021
  # Extract frames using multithreading (decoder is created inside each thread for thread safety)
1022
  result = decode_with_multithreading(frame_indices.tolist(), num_threads=self.num_extract_threads, video_path=video_path)
1023
-
1024
  # Extract frame tensor (N, C, H, W)
1025
  frames_tensor = result["data"]
1026
-
1027
  # Create final VideoMetadata object
1028
  video_metadata = VideoMetadata(
1029
  total_num_frames=len(frame_indices),
@@ -1034,15 +1107,15 @@ class MossVLVideoProcessor(BaseVideoProcessor):
1034
  width=torchcodec_metadata.width,
1035
  frames_indices=frame_indices
1036
  )
1037
-
1038
  # Ensure frames are in (T, C, H, W) format
1039
  if frames_tensor.dim() == 4: # (N, C, H, W)
1040
  video_tensor = frames_tensor
1041
  else:
1042
  raise ValueError(f"Unexpected frame tensor shape: {frames_tensor.shape}")
1043
-
1044
  return video_tensor, video_metadata
1045
-
1046
  except Exception as e:
1047
  logger.error(f"Error loading video {video_path}: {e}")
1048
  traceback.print_exc()
@@ -1180,15 +1253,15 @@ class MossVLVideoProcessor(BaseVideoProcessor):
1180
  ) -> BatchFeature:
1181
  """
1182
  Preprocess videos for the model.
1183
-
1184
  This method overrides the base class to handle two video input formats:
1185
  1. String path: "path/to/video.mp4"
1186
- 2. Dict with segments: {"video_path": "...", "segment": [[start, end], [time], ...]}
1187
-
1188
  Args:
1189
  videos: Video input(s) in one of the supported formats.
1190
  **kwargs: Additional arguments passed to _preprocess.
1191
-
1192
  Returns:
1193
  BatchFeature with pixel_values_videos, video_grid_thw, and optionally video_metadata.
1194
  """
@@ -1197,11 +1270,11 @@ class MossVLVideoProcessor(BaseVideoProcessor):
1197
  captured_kwargs=kwargs.keys(),
1198
  valid_processor_keys=list(self.valid_kwargs.__annotations__.keys()) + ["return_tensors"],
1199
  )
1200
-
1201
  # Set default kwargs from self
1202
  for kwarg_name in self.valid_kwargs.__annotations__:
1203
  kwargs.setdefault(kwarg_name, getattr(self, kwarg_name, None))
1204
-
1205
  # Pop kwargs that are handled separately
1206
  return_tensors = kwargs.pop("return_tensors", None)
1207
  return_metadata = kwargs.pop("return_metadata", False)
@@ -1210,44 +1283,43 @@ class MossVLVideoProcessor(BaseVideoProcessor):
1210
  kwargs.pop("video_metadata", None) # We generate our own metadata
1211
  kwargs.pop("do_sample_frames", None) # We handle sampling ourselves
1212
  kwargs.pop("data_format", None) # Not used
1213
-
1214
  # Normalize input to list format
1215
  if not isinstance(videos, list):
1216
  videos = [videos]
1217
-
1218
  # Get video processing params from kwargs (may be passed explicitly for per-batch configuration)
1219
  video_fps = kwargs.pop("video_fps", None)
1220
  min_frames = kwargs.pop("min_frames", None)
1221
  max_frames = kwargs.pop("max_frames", None)
1222
-
1223
  # Use fetch_videos to handle both string and dict formats
1224
  video_tensors, video_metadata = self.fetch_videos(
1225
- videos,
1226
  video_fps=video_fps,
1227
  min_frames=min_frames,
1228
  max_frames=max_frames,
1229
  )
1230
-
1231
  # Prepare video tensors using _prepare_input_videos
1232
  prepared_videos = self._prepare_input_videos(
1233
  videos=video_tensors,
1234
  input_data_format=input_data_format,
1235
  device=device,
1236
  )
1237
-
1238
  # Process kwargs for _preprocess
1239
  kwargs = self._further_process_kwargs(**kwargs)
1240
  self._validate_preprocess_kwargs(**kwargs)
1241
-
1242
  # Call _preprocess with prepared videos
1243
  result = self._preprocess(videos=prepared_videos, return_tensors=return_tensors, **kwargs)
1244
-
1245
  # Add metadata if requested
1246
  if return_metadata:
1247
  result["video_metadata"] = video_metadata
1248
-
1249
  return result
1250
 
1251
 
1252
  __all__ = ["MossVLVideoProcessor"]
1253
-
 
14
  # limitations under the License.
15
  """video processor class for Moss-VL."""
16
 
17
+ import fcntl
18
+ import hashlib
19
  import json
20
  import logging as system_logging
21
  import math
22
  import os
23
  import re
24
  import subprocess
25
+ import tempfile
26
  import traceback
27
  from functools import lru_cache
28
  from typing import Any, Dict, List, Optional, Union
 
52
  def check_video_for_extra_streams_and_errors(video_path: str) -> dict:
53
  """
54
  Check if video file has abnormal streams or errors reported by ffprobe.
55
+
56
  Args:
57
  video_path: Path to the video file.
58
+
59
  Returns:
60
  A dictionary containing:
61
  - 'has_extra_streams': bool, whether there are streams other than video and audio.
 
73
  'stream_details': [],
74
  'num_streams': 0
75
  }
76
+
77
  command = [
78
  "ffprobe",
79
  "-v", "error",
 
82
  "-of", "json",
83
  video_path
84
  ]
85
+
86
  try:
87
  process = subprocess.run(
88
  command,
 
91
  check=False
92
  )
93
  result['ffprobe_successful'] = (process.returncode == 0)
94
+
95
  if process.stderr:
96
  result['ffprobe_output_error'] = process.stderr
97
  unsupported_codec_pattern = re.compile(r"Unsupported codec with id \d+ for input stream \d+")
98
  result['unsupported_codec_errors'] = unsupported_codec_pattern.findall(process.stderr)
99
+
100
  if process.stdout:
101
  ffprobe_data = json.loads(process.stdout)
102
  if 'streams' in ffprobe_data:
 
107
  result['stream_details'].append({'index': stream_index, 'codec_type': stream_type})
108
  if stream_type not in ['video', 'audio']:
109
  result['has_extra_streams'] = True
110
+
111
  if 'format' in ffprobe_data and 'nb_streams' in ffprobe_data['format']:
112
  if result['num_streams'] == 0:
113
  result['num_streams'] = ffprobe_data['format']['nb_streams']
 
125
  except Exception as e:
126
  result['ffprobe_output_error'] = f"An unexpected error occurred: {e}"
127
  result['ffprobe_successful'] = False
128
+
129
  return result
130
 
131
 
132
  def remove_video_extra_stream_ffmpeg(input_video: str, output_video: str) -> bool:
133
  """
134
  Remove extra streams from video using ffmpeg.
135
+
136
  Args:
137
  input_video: Path to input video.
138
  output_video: Path to output video.
139
+
140
  Returns:
141
  bool: True if successful, False otherwise.
142
  """
 
152
  "-movflags", "faststart",
153
  output_video,
154
  ]
155
+
156
  try:
157
  subprocess.run(command_list, shell=False, check=True, capture_output=True)
158
  return True
 
168
  return False
169
 
170
 
171
+ def _video_clean_lock_path(video_path: str) -> str:
172
+ lock_dir = os.path.join(tempfile.gettempdir(), "mossvl_video_clean_locks")
173
+ os.makedirs(lock_dir, exist_ok=True)
174
+ path_key = os.path.realpath(video_path).encode("utf-8")
175
+ return os.path.join(lock_dir, f"{hashlib.sha256(path_key).hexdigest()}.lock")
176
+
177
+
178
+ def _is_reusable_clean_video(video_path: str) -> bool:
179
+ if not os.path.isfile(video_path) or os.path.getsize(video_path) <= 0:
180
+ return False
181
+ probe = check_video_for_extra_streams_and_errors(video_path)
182
+ return probe["ffprobe_successful"] and not probe["has_extra_streams"]
183
+
184
+
185
  def clean_video_streams(video_path: str) -> str:
186
  """
187
  Clean video streams if extra streams are detected.
188
+
189
  Args:
190
  video_path: Path to the video file.
191
+
192
  Returns:
193
  str: Path to cleaned video (or original if no cleaning needed).
194
  """
 
199
  file_name_without_ext, file_ext = os.path.splitext(base_name)
200
  new_base_name = f"{file_name_without_ext}_fix{file_ext}"
201
  video_path_output = os.path.join(output_folder, new_base_name)
202
+
203
+ try:
204
+ with open(_video_clean_lock_path(video_path), "a") as lock_file:
205
+ fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX)
206
+ if _is_reusable_clean_video(video_path_output):
207
+ return video_path_output
208
+
209
+ temp_fd, temp_output = tempfile.mkstemp(
210
+ prefix=f".{file_name_without_ext}_fix.",
211
+ suffix=file_ext,
212
+ dir=output_folder,
213
+ )
214
+ os.close(temp_fd)
215
+ try:
216
+ if not remove_video_extra_stream_ffmpeg(video_path, temp_output):
217
+ logger.warning("Failed to remove extra streams with ffmpeg")
218
+ return video_path
219
+ os.replace(temp_output, video_path_output)
220
+ return video_path_output
221
+ finally:
222
+ if os.path.exists(temp_output):
223
+ os.unlink(temp_output)
224
+ except OSError as exc:
225
+ logger.warning(f"Failed to publish cleaned video for {video_path}: {exc}")
226
  return video_path
 
227
  return video_path
228
 
229
 
 
239
  min_pts = torchcodec_metadata.begin_stream_seconds_from_content
240
  if min_pts is None:
241
  min_pts = 0.0
242
+ # TorchCodec can reject timestamps exactly equal to the reported stream
243
+ # begin due to tiny metadata/decoder precision differences.
244
+ safe_min_pts = min_pts + TORCHCODEC_TIMESTAMP_EPSILON
245
 
246
  max_pts_candidates = []
247
  if torchcodec_metadata.num_frames_from_content and torchcodec_metadata.average_fps:
 
255
  max_pts_candidates.append(torchcodec_metadata.duration_seconds - TORCHCODEC_TIMESTAMP_EPSILON)
256
 
257
  if max_pts_candidates:
258
+ max_pts = max(safe_min_pts, min(max_pts_candidates))
259
+ return [max(safe_min_pts, min(float(t), max_pts)) for t in timestamps]
260
+ if safe_min_pts > 0:
261
+ return [max(safe_min_pts, float(t)) for t in timestamps]
262
  return [float(t) for t in timestamps]
263
 
264
 
265
  def split_indices(indices: List[Union[int, float]], num_chunks: int) -> List[List[Union[int, float]]]:
266
  """
267
  Split an index list into roughly equal chunks.
268
+
269
  Args:
270
  indices: List of indices to split.
271
  num_chunks: Number of chunks to create.
272
+
273
  Returns:
274
  List of index chunks.
275
  """
276
+ if len(indices) == 0:
277
+ raise ValueError("indices must not be empty")
278
+ if num_chunks <= 0:
279
+ raise ValueError("num_chunks must be positive")
280
+
281
+ num_chunks = min(num_chunks, len(indices))
282
  chunk_size = len(indices) // num_chunks
283
  chunks = []
284
  for i in range(num_chunks - 1):
 
290
  def decode_sequentially(indices: List[int], video_path: str, ffmpeg_threads: int = 0):
291
  """
292
  Decode frames sequentially from a video.
293
+
294
  Args:
295
  indices: List of frame indices to decode.
296
  video_path: Path to the video file.
297
  ffmpeg_threads: Number of ffmpeg threads to use.
298
+
299
  Returns:
300
  FrameBatch from torchcodec.
301
  """
 
309
  def decode_with_multithreading(indices: List[int], num_threads: int, video_path: str) -> dict:
310
  """
311
  Decode frames using multithreading with joblib.
312
+
313
  Args:
314
  indices: List of frame indices to decode.
315
  num_threads: Number of threads to use.
316
  video_path: Path to the video file.
317
+
318
  Returns:
319
  dict: Contains 'data', 'duration_seconds', 'pts_seconds' tensors.
320
  """
321
  chunks = split_indices(indices, num_chunks=num_threads)
322
+ results = Parallel(n_jobs=len(chunks), prefer="threads", verbose=0)(
323
  delayed(decode_sequentially)(chunk, video_path) for chunk in chunks
324
  )
325
+
326
  return {
327
  "data": torch.cat([frame_batch.data for frame_batch in results], dim=0),
328
  "duration_seconds": torch.cat([frame_batch.duration_seconds for frame_batch in results], dim=0),
 
333
  def decode_sequentially_timestamp(timestamp_list: List[float], video_path: str, ffmpeg_threads: int = 0):
334
  """
335
  Decode frames sequentially from a video based on timestamps.
336
+
337
  Args:
338
  timestamp_list: List of timestamps (in seconds) to decode.
339
  video_path: Path to the video file.
340
  ffmpeg_threads: Number of ffmpeg threads to use.
341
+
342
  Returns:
343
  FrameBatch from torchcodec.
344
  """
 
355
  def timestamp_decode_with_multithreading(timestamp_list: List[float], num_threads: int, video_path: str) -> dict:
356
  """
357
  Decode frames using multithreading based on timestamps.
358
+
359
  Args:
360
  timestamp_list: List of timestamps (in seconds) to decode.
361
  num_threads: Number of threads to use.
362
  video_path: Path to the video file.
363
+
364
  Returns:
365
  dict: Contains 'data', 'duration_seconds', 'pts_seconds' tensors.
366
  """
367
  chunks = split_indices(timestamp_list, num_chunks=num_threads)
368
+ results = Parallel(n_jobs=len(chunks), prefer="threads", verbose=0)(
369
  delayed(decode_sequentially_timestamp)(chunk, video_path) for chunk in chunks
370
  )
371
+
372
  # Concatenate results from all threads
373
  data_list = [frame_batch.data for frame_batch in results]
374
  duration_list = [frame_batch.duration_seconds for frame_batch in results]
375
  pts_list = [frame_batch.pts_seconds for frame_batch in results]
376
+
377
  if not data_list:
378
  logger.warning("No frames were successfully decoded.")
379
  return {"data": torch.empty(0), "duration_seconds": torch.empty(0), "pts_seconds": torch.empty(0)}
380
+
381
  return {
382
  "data": torch.cat(data_list, dim=0),
383
  "duration_seconds": torch.cat(duration_list, dim=0),
 
393
  ) -> Optional[dict]:
394
  """
395
  Extract frames from video using torchcodec with multithreading.
396
+
397
  Args:
398
  video_path: Path to the video file.
399
  sample_frames_count: Number of frames to sample.
400
  num_threads: Number of threads to use for extraction.
401
  sampling_method: Sampling method, either "index" (uniform frame indices) or "timestamp" (uniform timestamps).
402
+
403
  Returns:
404
  dict: Contains 'data' (N, C, H, W), 'duration_seconds' (N,), 'pts_seconds' (N,) tensors.
405
  Returns None if extraction fails.
 
411
 
412
 
413
  total_frames_in_video = metadata.num_frames_from_content
414
+
415
  effective_sample_count = min(sample_frames_count, total_frames_in_video)
416
  if effective_sample_count == 0:
417
  logger.error("Cannot extract frames: video has 0 frames or specified frame count is 0")
418
  return None
419
+
420
  # Generate uniform frame indices
421
  frame_indices = np.linspace(0, total_frames_in_video - 1, effective_sample_count).astype(np.int32)
422
  # Ensure indices are valid and remove duplicates
423
  frame_indices = np.unique(np.clip(frame_indices, 0, total_frames_in_video - 1))
424
+
425
  result = decode_with_multithreading(frame_indices.tolist(), num_threads=num_threads, video_path=video_path)
426
  # Add frame_indices to the result for later use
427
  result["frame_indices"] = frame_indices
428
  return result
429
 
430
 
431
+
432
  except Exception:
433
  traceback.print_exc()
434
  return None
 
564
  return video_input["video_path"]
565
  return video_input
566
 
567
+ def _get_video_fps_from_input(
568
+ self,
569
+ video_input: Union[str, Dict[str, Any]],
570
+ default_fps: Optional[float],
571
+ ) -> Optional[float]:
572
+ if not isinstance(video_input, dict) or "fps" not in video_input:
573
+ return default_fps
574
+
575
+ raw_fps = video_input["fps"]
576
+ if raw_fps is None:
577
+ return default_fps
578
+ try:
579
+ fps = float(raw_fps)
580
+ except (TypeError, ValueError) as exc:
581
+ raise ValueError(f"Invalid video fps value: {raw_fps!r}") from exc
582
+ if fps <= 0:
583
+ raise ValueError(f"Invalid video fps value: {raw_fps!r}; fps must be positive")
584
+ return fps
585
+
586
  def _get_video_duration_seconds(self, video_input: Union[str, Dict[str, Any]]) -> float:
587
  """Get video duration in seconds for weighted frame-budget allocation."""
588
  video_path = cached_clean_video_streams(self._get_video_path_from_input(video_input))
 
670
  ) -> int:
671
  """
672
  Calculate the number of frames to sample using fps-based logic with min/max constraints.
673
+
674
  Logic:
675
  1. Calculate target_frames based on fps and video duration
676
  2. Apply min_frames and max_frames constraints
677
  3. Apply max_allowed_frames protection (rough cap from total video_max_pixels budget)
678
  4. Return the number of frames to sample
679
+
680
  Args:
681
  metadata (`VideoMetadata`):
682
  Metadata of the video containing information about total duration, fps and total number of frames.
 
696
  raise ValueError("`num_frames` and `fps` are mutually exclusive arguments, please use only one!")
697
 
698
  total_num_frames = metadata.total_num_frames
699
+
700
  # Use provided min/max or fall back to defaults
701
  effective_min_frames = min_frames if min_frames is not None else self.min_frames
702
  effective_max_frames = max_frames if max_frames is not None else self.max_frames
703
+
704
  # Rough per-video frame cap derived from the multi-video total budget
705
  # (exact allocation happens later in _preprocess via weighted distribution)
706
  per_frame_min_pixels = self.size.get("shortest_edge", None) if self.size else None
 
708
  if per_frame_min_pixels is not None and video_max_pixels is not None and per_frame_min_pixels > 0:
709
  max_allowed_frames = video_max_pixels // per_frame_min_pixels
710
  effective_max_frames = min(effective_max_frames, max_allowed_frames)
711
+
712
  # Get video duration
713
  if hasattr(metadata, 'duration') and metadata.duration is not None:
714
  duration = metadata.duration
 
726
 
727
  # Use provided fps or default
728
  target_fps = fps if fps is not None else self.video_fps
729
+
730
  # Calculate target frames based on fps and duration
731
  if num_frames is None:
732
  # Calculate how many frames we should sample based on target fps
733
  target_total_frames = int(math.ceil(duration * target_fps - 1e-6))
734
+
735
  # Apply min/max constraints
736
  sample_frames = max(target_total_frames, effective_min_frames)
737
  sample_frames = min(sample_frames, effective_max_frames, total_num_frames)
 
880
  ):
881
  """
882
  Fetch video frames for a specific segment.
883
+
884
  Args:
885
  video_path: Path to the video file
886
  segment: [start, end] for a segment (left-closed, right-open) or [time] for a single frame
887
  min_frames: Minimum frames for this segment (weighted). Defaults to self.min_frames. Must be >= 1.
888
  max_frames: Maximum frames for this segment (weighted). Defaults to self.max_frames. Must be >= 1.
889
  video_fps: Target frames per second for video sampling. If None, uses self.video_fps.
890
+
891
  Returns:
892
  Tuple of (video_tensor, video_metadata)
893
  """
 
896
  max_frames = max(1, max_frames if max_frames is not None else self.max_frames)
897
  # Use provided video_fps or fall back to self.video_fps
898
  target_video_fps = video_fps if video_fps is not None else self.video_fps
899
+
900
+ video_path = cached_clean_video_streams(video_path)
901
  decoder = VideoDecoder(video_path, num_ffmpeg_threads=0)
902
  try:
903
  torchcodec_metadata = decoder.metadata
904
+
905
  video_fps = torchcodec_metadata.average_fps
906
+
907
  # Calculate duration
908
  duration = None
909
  if torchcodec_metadata.end_stream_seconds_from_content is not None and torchcodec_metadata.begin_stream_seconds_from_content is not None:
910
  duration = torchcodec_metadata.end_stream_seconds_from_content - torchcodec_metadata.begin_stream_seconds_from_content
911
  if duration is None or duration <= 0:
912
  duration = torchcodec_metadata.duration_seconds
913
+
914
  if len(segment) == 1:
915
  # Single frame at specified time
916
  actual_timestamps = self._clamp_timestamps_for_decoder([segment[0]], torchcodec_metadata)
 
921
  # Segment [start, end) - left-closed, right-open interval
922
  start_time, end_time = segment
923
  segment_duration = end_time - start_time
924
+
925
  # Calculate number of frames to sample for this segment
926
  target_frames = int(math.ceil(segment_duration * target_video_fps))
927
  target_frames = max(target_frames, min_frames)
928
  target_frames = min(target_frames, max_frames)
929
+
930
  # Generate timestamps for uniform sampling within segment
931
  if target_frames == 1:
932
  actual_timestamps = [start_time] # Use start_time for single frame
933
  else:
934
  # Sample uniformly within [start, end), endpoint=False for left-closed right-open
935
  actual_timestamps = np.linspace(start_time, end_time, target_frames, endpoint=False).tolist()
936
+
937
  actual_timestamps = self._clamp_timestamps_for_decoder(actual_timestamps, torchcodec_metadata)
938
 
939
  # Use multithreading for extraction
940
  result = timestamp_decode_with_multithreading(actual_timestamps, self.num_extract_threads, video_path)
941
  video_tensor = result["data"]
942
  sample_count = len(actual_timestamps)
943
+
944
  # Create VideoMetadata
945
  video_metadata = VideoMetadata(
946
  total_num_frames=sample_count,
 
951
  width=torchcodec_metadata.width,
952
  frames_indices=None
953
  )
954
+
955
  # Store actual timestamps as a custom attribute for _calculate_timestamps to use
956
  video_metadata.actual_timestamps = actual_timestamps
957
+
958
  return video_tensor, video_metadata
959
  finally:
960
  del decoder
961
 
962
  def fetch_videos(
963
+ self,
964
+ video_url_or_urls: Union[str, Dict[str, Any], List[Union[str, Dict[str, Any]]]],
965
+ sample_indices_fn=None,
966
  video_fps: Optional[float] = None,
967
  min_frames: Optional[int] = None,
968
  max_frames: Optional[int] = None,
969
  ):
970
  """
971
  Override fetch_videos to use torchcodec for frame extraction.
972
+
973
  This method uses torchcodec with multithreading for efficient frame extraction.
974
  Frame count is calculated by the calculate_num_frames method
975
  (fps-based with min/max constraints).
976
+
977
  Args:
978
  video_url_or_urls: Can be one of:
979
  - str: Single video path
980
+ - Dict: Video descriptor with ``video_path`` and optional ``fps`` /
981
+ ``segments`` (``segment`` is accepted as an alias)
982
+ - List[Union[str, Dict]]: List of video paths or descriptors
983
  sample_indices_fn: (Not used) Kept for compatibility with base class signature.
984
  video_fps: Target frames per second for video sampling. If None, uses self.video_fps.
985
  min_frames: Minimum number of frames to sample. If None, uses self.min_frames.
986
  max_frames: Maximum number of frames to sample. If None, uses self.max_frames.
987
+
988
  Returns:
989
  Tuple of (videos, metadata) where videos are torch.Tensors and metadata are VideoMetadata objects.
990
  """
 
1004
  effective_max_frames,
1005
  )
1006
  for x, allocated_max_frames in zip(video_url_or_urls, per_video_max_frames):
1007
+ explicit_video_fps = self._get_video_fps_from_input(x, None)
1008
+ item_video_fps = explicit_video_fps if explicit_video_fps is not None else effective_video_fps
1009
  result = self.fetch_videos(
1010
+ x,
1011
+ video_fps=item_video_fps,
1012
  min_frames=effective_min_frames,
1013
  max_frames=allocated_max_frames,
1014
  )
 
1020
  all_videos.append(result[0])
1021
  all_metadata.append(result[1])
1022
  return all_videos, all_metadata
1023
+
1024
+ # Dict inputs may override FPS and optionally select segments. Without
1025
+ # segments, the dict describes a full video.
1026
  if isinstance(video_url_or_urls, dict):
1027
  video_path = video_url_or_urls["video_path"]
1028
+ input_video_fps = self._get_video_fps_from_input(video_url_or_urls, effective_video_fps)
1029
+ segments = video_url_or_urls.get("segments", video_url_or_urls.get("segment"))
1030
+
1031
+ if segments is not None:
1032
+ return self._fetch_video_segments_batched(
1033
+ video_path,
1034
+ segments,
1035
+ min_frames=effective_min_frames,
1036
+ max_frames=effective_max_frames,
1037
+ video_fps=input_video_fps,
1038
+ )
1039
+
1040
+ video_url_or_urls = video_path
1041
+ effective_video_fps = input_video_fps
1042
+
1043
  # Single video path
1044
  video_path = video_url_or_urls
1045
+
1046
  # Clean video streams first (remove extra streams if needed)
1047
  video_path = cached_clean_video_streams(video_path)
1048
 
 
1055
  duration = None
1056
  if torchcodec_metadata.end_stream_seconds_from_content is not None and torchcodec_metadata.begin_stream_seconds_from_content is not None:
1057
  duration = torchcodec_metadata.end_stream_seconds_from_content - torchcodec_metadata.begin_stream_seconds_from_content
1058
+
1059
  if duration is None or duration <= 0:
1060
  duration = torchcodec_metadata.duration_seconds
1061
+
1062
  # Use num_frames_from_content for accurate frame count (consistent with extraction)
1063
  total_frames_in_video = torchcodec_metadata.num_frames_from_content
1064
+
1065
  # Create VideoMetadata object for sample_frames method
1066
  temp_metadata = VideoMetadata(
1067
  total_num_frames=total_frames_in_video,
 
1072
  width=torchcodec_metadata.width,
1073
  frames_indices=None
1074
  )
1075
+
1076
  # Use calculate_num_frames method to get the number of frames to sample
1077
  sample_frames_count = self.calculate_num_frames(
1078
+ temp_metadata,
1079
  fps=effective_video_fps,
1080
  min_frames=effective_min_frames,
1081
  max_frames=effective_max_frames,
1082
  )
1083
+
1084
  # Ensure sample count is valid
1085
  effective_sample_count = min(sample_frames_count, total_frames_in_video)
1086
  if effective_sample_count == 0:
1087
  raise ValueError(f"Cannot extract frames: video has 0 frames or specified frame count is 0")
1088
+
1089
  # Generate uniform frame indices
1090
  frame_indices = np.linspace(0, total_frames_in_video - 1, effective_sample_count).astype(np.int32)
1091
  # Ensure indices are valid and remove duplicates
1092
  frame_indices = np.unique(np.clip(frame_indices, 0, total_frames_in_video - 1))
1093
+
1094
  # Extract frames using multithreading (decoder is created inside each thread for thread safety)
1095
  result = decode_with_multithreading(frame_indices.tolist(), num_threads=self.num_extract_threads, video_path=video_path)
1096
+
1097
  # Extract frame tensor (N, C, H, W)
1098
  frames_tensor = result["data"]
1099
+
1100
  # Create final VideoMetadata object
1101
  video_metadata = VideoMetadata(
1102
  total_num_frames=len(frame_indices),
 
1107
  width=torchcodec_metadata.width,
1108
  frames_indices=frame_indices
1109
  )
1110
+
1111
  # Ensure frames are in (T, C, H, W) format
1112
  if frames_tensor.dim() == 4: # (N, C, H, W)
1113
  video_tensor = frames_tensor
1114
  else:
1115
  raise ValueError(f"Unexpected frame tensor shape: {frames_tensor.shape}")
1116
+
1117
  return video_tensor, video_metadata
1118
+
1119
  except Exception as e:
1120
  logger.error(f"Error loading video {video_path}: {e}")
1121
  traceback.print_exc()
 
1253
  ) -> BatchFeature:
1254
  """
1255
  Preprocess videos for the model.
1256
+
1257
  This method overrides the base class to handle two video input formats:
1258
  1. String path: "path/to/video.mp4"
1259
+ 2. Dict descriptor with ``video_path`` and optional ``fps`` / ``segments``
1260
+
1261
  Args:
1262
  videos: Video input(s) in one of the supported formats.
1263
  **kwargs: Additional arguments passed to _preprocess.
1264
+
1265
  Returns:
1266
  BatchFeature with pixel_values_videos, video_grid_thw, and optionally video_metadata.
1267
  """
 
1270
  captured_kwargs=kwargs.keys(),
1271
  valid_processor_keys=list(self.valid_kwargs.__annotations__.keys()) + ["return_tensors"],
1272
  )
1273
+
1274
  # Set default kwargs from self
1275
  for kwarg_name in self.valid_kwargs.__annotations__:
1276
  kwargs.setdefault(kwarg_name, getattr(self, kwarg_name, None))
1277
+
1278
  # Pop kwargs that are handled separately
1279
  return_tensors = kwargs.pop("return_tensors", None)
1280
  return_metadata = kwargs.pop("return_metadata", False)
 
1283
  kwargs.pop("video_metadata", None) # We generate our own metadata
1284
  kwargs.pop("do_sample_frames", None) # We handle sampling ourselves
1285
  kwargs.pop("data_format", None) # Not used
1286
+
1287
  # Normalize input to list format
1288
  if not isinstance(videos, list):
1289
  videos = [videos]
1290
+
1291
  # Get video processing params from kwargs (may be passed explicitly for per-batch configuration)
1292
  video_fps = kwargs.pop("video_fps", None)
1293
  min_frames = kwargs.pop("min_frames", None)
1294
  max_frames = kwargs.pop("max_frames", None)
1295
+
1296
  # Use fetch_videos to handle both string and dict formats
1297
  video_tensors, video_metadata = self.fetch_videos(
1298
+ videos,
1299
  video_fps=video_fps,
1300
  min_frames=min_frames,
1301
  max_frames=max_frames,
1302
  )
1303
+
1304
  # Prepare video tensors using _prepare_input_videos
1305
  prepared_videos = self._prepare_input_videos(
1306
  videos=video_tensors,
1307
  input_data_format=input_data_format,
1308
  device=device,
1309
  )
1310
+
1311
  # Process kwargs for _preprocess
1312
  kwargs = self._further_process_kwargs(**kwargs)
1313
  self._validate_preprocess_kwargs(**kwargs)
1314
+
1315
  # Call _preprocess with prepared videos
1316
  result = self._preprocess(videos=prepared_videos, return_tensors=return_tensors, **kwargs)
1317
+
1318
  # Add metadata if requested
1319
  if return_metadata:
1320
  result["video_metadata"] = video_metadata
1321
+
1322
  return result
1323
 
1324
 
1325
  __all__ = ["MossVLVideoProcessor"]