OpenMOSS-Team/MOSS-VL-Realtime-NF4
15436
1# coding=utf-82# Copyright 2025 The FNLP Vision Team and The HuggingFace Inc. team. All rights reserved.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8# http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15"""16Processor class for Moss-VL.17"""18 19from typing import Any, Dict, List, Optional, Union20 21import numpy as np22import torch23from torchvision.transforms.v2 import functional as F24from PIL import Image25from transformers.feature_extraction_utils import BatchFeature26from transformers.image_utils import ImageInput, SizeDict, make_flat_list_of_images27from transformers.image_processing_utils_fast import group_images_by_shape, reorder_images28from transformers.utils import TensorType29from transformers.processing_utils import (30 ImagesKwargs,31 ProcessingKwargs,32 ProcessorMixin,33 Unpack,34 VideosKwargs,35)36from transformers.tokenization_utils_base import PreTokenizedInput, TextInput37from transformers.utils import logging38from transformers.models.qwen2_vl.image_processing_qwen2_vl_fast import Qwen2VLImageProcessorFast39from transformers.models.qwen2_vl.image_processing_qwen2_vl import smart_resize40 41 42logger = logging.get_logger(__name__)43 44 45class MossVLImageProcessorFast(Qwen2VLImageProcessorFast):46 """47 Custom image processor that overrides _preprocess to support multi_image_max_pixels.48 Inherits from Qwen2VLImageProcessorFast.49 """50 # Multi-image batch total pixels limit (read from config)51 multi_image_max_pixels = None52 53 54 def _preprocess(55 self,56 images: list["torch.Tensor"],57 do_resize: bool,58 size: SizeDict,59 interpolation: Optional["F.InterpolationMode"],60 do_rescale: bool,61 rescale_factor: float,62 do_normalize: bool,63 image_mean: Optional[Union[float, list[float]]],64 image_std: Optional[Union[float, list[float]]],65 patch_size: int,66 temporal_patch_size: int,67 merge_size: int,68 disable_grouping: Optional[bool],69 return_tensors: Optional[Union[str, TensorType]],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 allocated75 to each image based on its original pixel count. min_pixels remains a per-image76 constraint. multi_image_max_pixels can be configured separately from longest_edge.77 """78 min_pixels = size["shortest_edge"]79 max_pixels = size["longest_edge"] # Per-image upper limit80 # Use multi_image_max_pixels if configured, otherwise fall back to longest_edge81 multi_image_max_pixels = getattr(self, "multi_image_max_pixels", None) or max_pixels82 83 # Calculate total original pixels across all images in the batch84 # This is used to proportionally allocate max_pixels to each image85 total_original_pixels = sum(img.shape[-2] * img.shape[-1] for img in images)86 87 # Group images by size for batched resizing88 grouped_images, grouped_images_index = group_images_by_shape(images, disable_grouping=disable_grouping)89 resized_images_grouped = {}90 for shape, stacked_images in grouped_images.items():91 height, width = stacked_images.shape[-2:]92 if do_resize:93 # Calculate proportional max_pixels for images with this shape94 # Each image's max_pixels is allocated based on its proportion of total pixels95 original_pixels = height * width96 if total_original_pixels > 0:97 proportion = original_pixels / total_original_pixels98 proportional_max_pixels = int(multi_image_max_pixels * proportion)99 else:100 proportional_max_pixels = multi_image_max_pixels101 102 # Ensure proportional max_pixels is within [min_pixels, max_pixels] range103 # 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,111 factor=patch_size * merge_size,112 min_pixels=min_pixels,113 max_pixels=proportional_max_pixels,114 )115 stacked_images = self.resize(116 image=stacked_images,117 size=SizeDict(height=resized_height, width=resized_width),118 interpolation=interpolation,119 )120 resized_images_grouped[shape] = stacked_images121 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 constraint124 if len(images) > 1:125 total_resized_pixels = sum(img.shape[-2] * img.shape[-1] for img in resized_images)126 if total_resized_pixels > multi_image_max_pixels:127 logger.warning_once(128 f"Multi-image batch total pixels ({total_resized_pixels}) exceeds multi_image_max_pixels ({multi_image_max_pixels}). "129 f"This may happen when image_count * min_pixels > multi_image_max_pixels."130 )131 132 # Group images by size for further processing133 # Needed in case do_resize is False, or resize returns images with different sizes134 grouped_images, grouped_images_index = group_images_by_shape(resized_images, disable_grouping=disable_grouping)135 processed_images_grouped = {}136 processed_grids = {}137 for shape, stacked_images in grouped_images.items():138 resized_height, resized_width = stacked_images.shape[-2:]139 # Fused rescale and normalize140 patches = self.rescale_and_normalize(141 stacked_images, do_rescale, rescale_factor, do_normalize, image_mean, image_std142 )143 if patches.ndim == 4:144 # add a temporal dimension if we have images145 patches = patches.unsqueeze(1)146 if patches.shape[1] % temporal_patch_size != 0:147 repeats = patches[:, -1:].repeat(1, temporal_patch_size - 1, 1, 1, 1)148 patches = torch.cat([patches, repeats], dim=1)149 batch_size, grid_t, channel = patches.shape[:3]150 grid_t = grid_t // temporal_patch_size151 grid_h, grid_w = resized_height // patch_size, resized_width // patch_size152 153 patches = patches.view(154 batch_size,155 grid_t,156 temporal_patch_size,157 channel,158 grid_h // merge_size,159 merge_size,160 patch_size,161 grid_w // merge_size,162 merge_size,163 patch_size,164 )165 # Reorder dimensions to group grid and patch information for subsequent flattening.166 # (batch, grid_t, grid_h, grid_w, merge_h, merge_w, channel, temp_patch_size, patch_h, patch_w)167 # NPU ops support at most 8-D tensors; route the 10-D permute+reshape168 # through CPU there. CUDA handles 10-D natively โ keep it on-device.169 patches_device = patches.device170 if patches_device.type == "npu":171 patches = patches.cpu()172 patches = patches.permute(0, 1, 4, 7, 5, 8, 3, 2, 6, 9)173 flatten_patches = patches.reshape(174 batch_size,175 grid_t * grid_h * grid_w,176 channel * temporal_patch_size * patch_size * patch_size,177 ).to(patches_device)178 179 processed_images_grouped[shape] = flatten_patches180 processed_grids[shape] = [[grid_t, grid_h, grid_w]] * batch_size181 182 processed_images = reorder_images(processed_images_grouped, grouped_images_index)183 processed_grids = reorder_images(processed_grids, grouped_images_index)184 pixel_values = torch.cat(processed_images, dim=0)185 image_grid_thw = torch.tensor(processed_grids)186 187 return BatchFeature(188 data={"pixel_values": pixel_values, "image_grid_thw": image_grid_thw}, tensor_type=return_tensors189 )190 191def _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 primitive198 199 Returns:200 np.ndarray: NumPy array representation of the input201 """202 # Already numpy203 if isinstance(x, np.ndarray):204 return x205 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 first209 if hasattr(x, 'cpu'):210 return x.cpu().numpy()211 # For TensorFlow or already on CPU212 return x.numpy()213 214 # JAX arrays and other array-like objects that support __array__ protocol215 if hasattr(x, '__array__'):216 return np.asarray(x)217 218 # Python primitives (list, tuple, int, float)219 return np.array(x)220 221 222def _split_array_or_tensor(x, split_indices):223 """Split along the first dimension while preserving tensor/array type."""224 split_indices = [int(idx) for idx in split_indices]225 if isinstance(x, torch.Tensor):226 if not split_indices:227 return [x]228 chunks = []229 start = 0230 for end in split_indices:231 chunks.append(x[start:end])232 start = end233 chunks.append(x[start:])234 return chunks235 return np.split(x, split_indices)236 237 238def _concat_array_or_tensor(items, axis=0):239 """Concatenate while preserving tensor/array type and device."""240 if not items:241 return None242 243 if any(isinstance(item, torch.Tensor) for item in items):244 ref = next(item for item in items if isinstance(item, torch.Tensor))245 tensor_items = [246 item247 if isinstance(item, torch.Tensor)248 else torch.as_tensor(item, device=ref.device, dtype=ref.dtype)249 for item in items250 ]251 return torch.cat(tensor_items, dim=axis)252 253 return np.concatenate(items, axis=axis)254 255 256def _stack_array_or_tensor(items, axis=0):257 """Stack while preserving tensor/array type and device."""258 if not items:259 return None260 261 if any(isinstance(item, torch.Tensor) for item in items):262 ref = next(item for item in items if isinstance(item, torch.Tensor))263 tensor_items = [264 item265 if isinstance(item, torch.Tensor)266 else torch.as_tensor(item, device=ref.device, dtype=ref.dtype)267 for item in items268 ]269 return torch.stack(tensor_items, dim=axis)270 271 return np.stack(items, axis=axis)272 273 274class MossVLImagesKwargs(ImagesKwargs):275 min_pixels: Optional[int]276 max_pixels: Optional[int]277 patch_size: Optional[int]278 temporal_patch_size: Optional[int]279 merge_size: Optional[int]280 281 282 283class MossVLVideosKwargs(VideosKwargs, total=False):284 video_fps: Optional[Union[int, float]]285 min_frames: Optional[int]286 max_frames: Optional[int]287 num_extract_threads: Optional[int]288 289 290class MossVLProcessorKwargs(ProcessingKwargs, total=False):291 images_kwargs: MossVLImagesKwargs292 videos_kwargs: MossVLVideosKwargs293 # _defaults = {294 # "text_kwargs": {295 # "padding": True, # ๐ ๅฏ็จ padding296 # "padding_side": "left", # ๐ ๅทฆ padding297 # "pad_to_multiple_of": 8, # ๐ pad ๅฐ 8 ็ๅๆฐ298 # "return_token_type_ids": False,299 # "return_mm_token_type_ids": False,300 # },301 # "videos_kwargs": {"return_metadata": True},302 # }303 _defaults = {304 "text_kwargs": {305 "padding": False,306 "return_token_type_ids": False,307 "return_mm_token_type_ids": False,308 },309 "videos_kwargs": {"return_metadata": True},310 }311 312class MossVLProcessor(ProcessorMixin):313 r"""314 Constructs a Moss-VL processor which wraps a Qwen2VL image processor, Moss-VL video processor and a Qwen2 tokenizer315 into a single processor.316 317 [`MossVLProcessor`] offers all the functionalities of [`Qwen2VLImageProcessor`], [`MossVLVideoProcessor`] and [`Qwen2TokenizerFast`].318 See the [`~MossVLProcessor.__call__`] and [`~MossVLProcessor.decode`] for more information.319 320 Args:321 image_processor ([`Qwen2VLImageProcessor`], *optional*):322 The image processor is a required input.323 tokenizer ([`Qwen2TokenizerFast`], *optional*):324 The tokenizer is a required input.325 video_processor ([`MossVLVideoProcessor`], *optional*):326 The video processor is a required input.327 chat_template (`str`, *optional*): A Jinja template which will be used to convert lists of messages328 in a chat into a tokenizable string.329 """330 331 attributes = ["image_processor", "tokenizer", "video_processor"]332 image_processor_class = "AutoImageProcessor"333 video_processor_class = "AutoVideoProcessor"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 **kwargs343 ):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_token348 self.video_token = "<|video_pad|>" if not hasattr(tokenizer, "video_token") else tokenizer.video_token349 350 351 self.image_token_id = (352 tokenizer.image_token_id353 if getattr(tokenizer, "image_token_id", None)354 else tokenizer.convert_tokens_to_ids(self.image_token)355 )356 self.video_token_id = (357 tokenizer.video_token_id358 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_token364 )365 self.vision_end_token = (366 "<|vision_end|>" if not hasattr(tokenizer, "vision_end_token") else tokenizer.vision_end_token367 )368 369 # Placeholders used in input text370 self.image_placeholder = "<|image|>"371 self.video_placeholder = "<|video|>"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 387 388 def _process_media_per_sample(self, media, texts, kind, processor, kwargs):389 """Apply a complete media budget independently to each text sample.390 391 Lists within one sample still share the image/video budget (and video392 frame limit). A segmented video is one source item occupying several393 video placeholders; it must not cross text-sample boundaries.394 """395 if len(texts) == 1:396 return processor(**{kind: media}, **kwargs)397 if kind == "images":398 items = make_flat_list_of_images(media)399 widths = [1] * len(items)400 token = self.image_placeholder401 else:402 items = media if isinstance(media, list) else [media]403 widths = [404 len(item["segments"]) if isinstance(item, dict) and item.get("segments") else 1405 for item in items406 ]407 token = self.video_placeholder408 counts = [text.count(token) for text in texts]409 if sum(counts) != sum(widths):410 raise ValueError(411 f"{kind} placeholders do not match supplied media: "412 f"per-sample counts={counts}, media slots={sum(widths)}"413 )414 groups, offset = [], 0415 for count in counts:416 start, remaining = offset, count417 while remaining:418 width = widths[offset]419 if width > remaining:420 raise ValueError("A segmented video cannot cross text-sample boundaries")421 remaining -= width422 offset += 1423 groups.append(items[start:offset])424 # Keep media order, but never share resize/frame budgets across samples.425 results = [processor(**{kind: group}, **kwargs) for group in groups if group]426 if not results:427 return {}428 combined = {}429 for key in results[0]:430 values = [result[key] for result in results]431 if key == "video_metadata":432 combined[key] = [item for value in values for item in value]433 else:434 combined[key] = _concat_array_or_tensor(values, axis=0)435 return combined436 437 def __call__(438 self,439 text: Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]] = None,440 images: ImageInput = None,441 videos: Union[str, Dict[str, Any], List[Union[str, Dict[str, Any]]]] = None,442 labels_spans: Optional[Union[List[tuple], List[List[tuple]]]] = None,443 ignore_index: int = -100,444 **kwargs: Unpack[MossVLProcessorKwargs],445 ) -> BatchFeature:446 """447 Main method to prepare for the model one or several sequences(s) and image(s)/video(s).448 449 Args:450 text (`str`, `list[str]`, `list[list[str]]`):451 The sequence or batch of sequences to be encoded.452 images (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `list[PIL.Image.Image]`, `list[np.ndarray]`, `list[torch.Tensor]`):453 The image or batch of images to be prepared.454 videos (`str`, `Dict`, `list[str]`, `list[Dict]`):455 The video or batch of videos to be prepared. Each video can be:456 - A string path to a video file457 - A dict with keys:458 - "video_path": str, path to the video file459 - "segments": list of segments, where each segment is:460 - [start, end]: a time segment (left-closed, right-open interval in seconds)461 - [time]: a single frame at the specified time (in seconds)462 The number of segments should match the number of video placeholders in the text.463 labels_spans (`list[list[int]]`, `list[list[list[int]]]`, *optional*):464 Character-level spans indicating assistant regions in original text.465 Each span is a [start, end] list with inclusive start and exclusive end.466 Example: [[10, 50], [100, 150]] means characters [10:50) and [100:150) are assistant.467 Note: Use list (not tuple) for spans as they will be modified in place during processing.468 When provided, the processor will generate `labels` in the output, where:469 - Non-assistant tokens have value `ignore_index` (-100 by default)470 - Image tokens always have value `ignore_index` even in assistant part471 - Other assistant tokens have their token id as label472 ignore_index (`int`, *optional*, defaults to -100):473 Value for masked positions in labels.474 return_tensors (`str` or [`~utils.TensorType`], *optional*):475 If set, will return tensors of a particular framework. Acceptable values are:476 - `'tf'`: Return TensorFlow `tf.constant` objects.477 - `'pt'`: Return PyTorch `torch.Tensor` objects.478 - `'np'`: Return NumPy `np.ndarray` objects.479 - `'jax'`: Return JAX `jnp.ndarray` objects.480 481 482 Returns:483 [`BatchFeature`]: A [`BatchFeature`] with the following fields:484 - **input_ids** -- List of token ids to be fed to a model.485 - **attention_mask** -- List of indices specifying which tokens should be attended to by the model.486 - **pixel_values** -- Pixel values to be fed to a model (concatenation of images and videos).487 - **grid_thw** -- List of grid sizes (t, h, w) for each media item.488 - **media_nums_per_sample** -- List of number of media items per sample.489 - **labels** -- (Optional) Labels for training, only present when `labels_spans` is provided.490 """491 # Merge kwargs with defaults492 output_kwargs = self._merge_kwargs(493 MossVLProcessorKwargs,494 tokenizer_init_kwargs=self.tokenizer.init_kwargs,495 **kwargs,496 )497 498 # Establish sample boundaries before media preprocessing: each sample499 # receives its own complete image budget and video budget.500 if text is None or (isinstance(text, str) and not text.strip()):501 raise ValueError("Text input is required for MossVL processor and cannot be empty.")502 if not isinstance(text, list):503 text = [text]504 505 # Step 1: Process images if provided506 if images is not None:507 images_kwargs = output_kwargs["images_kwargs"].copy()508 images_kwargs["return_tensors"] = None509 image_inputs = self._process_media_per_sample(510 images, text, "images", self.image_processor, images_kwargs511 )512 image_grid_thw = image_inputs.get("image_grid_thw")513 else:514 image_inputs = {}515 image_grid_thw = None516 517 # Step 2: Process videos if provided518 if videos is not None:519 videos_kwargs = output_kwargs["videos_kwargs"].copy()520 videos_kwargs["return_tensors"] = None521 videos_inputs = self._process_media_per_sample(522 videos, text, "videos", self.video_processor, videos_kwargs523 )524 video_grid_thw = videos_inputs.get("video_grid_thw")525 # If user has not requested video metadata, pop it526 if "return_metadata" not in kwargs:527 video_metadata = videos_inputs.pop("video_metadata", [])528 else:529 video_metadata = videos_inputs.get("video_metadata", [])530 else:531 videos_inputs = {}532 video_grid_thw = None533 video_metadata = None534 535 # Step 3: Process text with placeholder replacement536 text = text.copy() # Copy to avoid in-place modifications537 538 # Prepare labels_spans if provided539 # labels_spans format: List[List[List[int]]] - batch of samples, each sample has multiple spans540 # Each span is [start, end] (list, not tuple) so it can be modified in place541 should_create_labels = labels_spans is not None542 if should_create_labels:543 # Ensure batch format: convert single sample spans to batch format544 # Single sample: [[start, end], [start, end], ...]545 # Batch: [[[start, end], ...], [[start, end], ...], ...]546 if labels_spans and isinstance(labels_spans[0], list) and len(labels_spans[0]) == 2 and isinstance(labels_spans[0][0], int):547 labels_spans = [labels_spans]548 549 # Step 3.0-pre: Check if we need to reorder (when both images and videos exist)550 # If only one media type exists, we can skip the expensive split+reorder+concat551 has_images = images is not None and "pixel_values" in image_inputs552 has_videos = videos is not None and "pixel_values_videos" in videos_inputs553 needs_reorder = has_images and has_videos554 555 image_pixel_values_list = []556 video_pixel_values_list = []557 558 # Step 3.0: Record the order of media in original text (before replacement)559 # This will be used later to correctly order pixel_values and grid_thw560 media_order_per_sample = []561 for i in range(len(text)):562 media_order = []563 temp_text = text[i]564 pos = 0565 while pos < len(temp_text):566 img_pos = temp_text.find(self.image_placeholder, pos)567 vid_pos = temp_text.find(self.video_placeholder, pos)568 569 if img_pos == -1 and vid_pos == -1:570 break571 572 if img_pos != -1 and (vid_pos == -1 or img_pos < vid_pos):573 media_order.append(("image", img_pos))574 pos = img_pos + len(self.image_placeholder)575 elif vid_pos != -1:576 media_order.append(("video", vid_pos))577 pos = vid_pos + len(self.video_placeholder)578 579 media_order_per_sample.append(media_order)580 581 # Step 3.0.1: Check if any sample has no media (empty samples need blank image)582 # If there are empty samples, we need to enter slow path to handle them properly583 has_empty_samples = any(len(order) == 0 for order in media_order_per_sample)584 if has_empty_samples:585 needs_reorder = True586 587 # Split pixel values for reordering if needed588 if needs_reorder:589 if has_images:590 flat_pixel_values = image_inputs["pixel_values"]591 flat_grid_thw = image_inputs["image_grid_thw"]592 # grid_thw is (t, h, w), num_patches = t * h * w593 patch_counts = [int(np.prod(_to_numpy(grid))) for grid in flat_grid_thw]594 if len(patch_counts) == 1:595 # Single image case: no need to split596 image_pixel_values_list = [flat_pixel_values]597 elif len(patch_counts) > 1:598 # Multiple images: split by cumulative counts599 split_indices = np.cumsum(patch_counts)[:-1]600 image_pixel_values_list = _split_array_or_tensor(601 flat_pixel_values, split_indices602 )603 604 if has_videos:605 flat_video_values = videos_inputs["pixel_values_videos"]606 flat_video_grid = videos_inputs["video_grid_thw"]607 video_patch_counts = [int(np.prod(_to_numpy(grid))) for grid in flat_video_grid]608 if len(video_patch_counts) == 1:609 # Single video case: no need to split610 video_pixel_values_list = [flat_video_values]611 elif len(video_patch_counts) > 1:612 # Multiple videos: split by cumulative counts613 split_indices = np.cumsum(video_patch_counts)[:-1]614 video_pixel_values_list = _split_array_or_tensor(615 flat_video_values, split_indices616 )617 618 # Step 3.1: Replace placeholders (simple replacement, no expansion yet)619 # In MossVL, one image placeholder = one image token620 # One video placeholder = one video token (will be expanded later)621 for i in range(len(text)):622 if should_create_labels:623 # Replace and update spans for image placeholders624 text[i], labels_spans[i] = self._replace_and_update_spans(625 text[i], self.image_placeholder, self.image_token, labels_spans[i]626 )627 # Replace and update spans for video placeholders628 text[i], labels_spans[i] = self._replace_and_update_spans(629 text[i], self.video_placeholder, self.video_token, labels_spans[i]630 )631 else:632 text[i] = text[i].replace(self.image_placeholder, self.image_token)633 text[i] = text[i].replace(self.video_placeholder, self.video_token)634 635 # Step 3.2: Validate token counts636 n_images_in_text = [t.count(self.image_token) for t in text]637 n_videos_in_text = [t.count(self.video_token) for t in text]638 639 # Count placeholders in text640 total_images_in_text = sum(n_images_in_text)641 total_videos_in_text = sum(n_videos_in_text)642 643 # Count actual images and videos provided644 total_images_provided = len(image_grid_thw) if image_grid_thw is not None else 0645 total_videos_provided = len(video_grid_thw) if video_grid_thw is not None else 0646 647 # Validate image counts648 if total_images_in_text != total_images_provided:649 raise ValueError(650 "Number of image tokens does not match number of images provided. "651 f"Found {total_images_in_text} image tokens in text and {total_images_provided} images."652 )653 654 # Validate video counts655 if total_videos_in_text != total_videos_provided:656 raise ValueError(657 "Number of video tokens does not match number of videos provided. "658 f"Found {total_videos_in_text} video tokens in text and {total_videos_provided} videos."659 )660 661 # Step 3.3: Expand video tokens with timestamps662 # Now expand each video token to multiple tokens (one per frame) with timestamps663 if video_grid_thw is not None:664 index = 0665 for i in range(len(text)):666 while self.video_token in text[i]:667 metadata = video_metadata[index]668 if metadata.fps is None:669 logger.warning_once(670 "MossVL requires frame timestamps to construct prompts, but the `fps` of the input video could not be inferred. "671 "Probably `video_metadata` was missing from inputs and you passed pre-sampled frames. "672 "Defaulting to `fps=24`. Please provide `video_metadata` for more accurate results."673 )674 metadata.fps = 24 if metadata.fps is None else metadata.fps675 676 # Calculate timestamps677 # Use actual_timestamps if available (for segments), otherwise use frames_indices678 actual_timestamps = getattr(metadata, 'actual_timestamps', None)679 curr_timestamp = self._calculate_timestamps(680 metadata.frames_indices,681 metadata.total_num_frames,682 metadata.fps,683 metadata.duration,684 self.video_processor.temporal_patch_size,685 actual_timestamps=actual_timestamps,686 )687 688 # Build video placeholder: one video token per frame with timestamp689 # video_grid_thw[index][0] is the temporal dimension (number of frames after merging)690 691 video_tokens = []692 for frame_idx in range(video_grid_thw[index][0]):693 curr_time = curr_timestamp[frame_idx]694 # Format: <|time_start|>X.X seconds<|time_end|><|image_pad|>695 video_tokens.append(696 f"{self.time_start_token}{curr_time:.1f} seconds{self.time_end_token}{self.image_token}"697 )698 699 # Wrap the entire video sequence with vision_start and vision_end tokens700 video_placeholder = f"{self.vision_start_token}{''.join(video_tokens)}{self.vision_end_token}"701 702 # Replace the video token with expanded sequence and update spans if needed703 if should_create_labels:704 text[i], labels_spans[i] = self._replace_and_update_spans(705 text[i], self.video_token, video_placeholder, labels_spans[i], replace_count=1706 )707 else:708 text[i] = text[i].replace(self.video_token, video_placeholder, 1)709 index += 1710 711 712 713 # Step 4: Tokenize text714 return_tensors = output_kwargs["text_kwargs"].pop("return_tensors", None)715 return_mm_token_type_ids = output_kwargs["text_kwargs"].pop("return_mm_token_type_ids", None)716 717 # Request offset_mapping if we need to create labels718 if should_create_labels:719 output_kwargs["text_kwargs"]["return_offsets_mapping"] = True720 721 text_inputs = self.tokenizer(text, **output_kwargs["text_kwargs"])722 723 # ignore check_special_mm_tokens nums in test and input ids.724 # self._check_special_mm_tokens(text, text_inputs, modalities=["image", "video"])725 726 # Create labels if labels_spans was provided727 if should_create_labels:728 offset_mapping = text_inputs.pop("offset_mapping")729 labels = self._create_labels_from_spans(730 text_inputs["input_ids"],731 offset_mapping,732 labels_spans,733 ignore_index734 )735 736 if return_mm_token_type_ids:737 array_ids = np.array(text_inputs["input_ids"])738 mm_token_type_ids = np.zeros_like(text_inputs["input_ids"])739 mm_token_type_ids[array_ids == self.image_token_id] = 1740 text_inputs["mm_token_type_ids"] = mm_token_type_ids.tolist()741 742 # Step 5: Concatenate pixel_values and grid_thw in sequence order743 # Prepare output744 output_data = {**text_inputs}745 746 if not needs_reorder:747 # Fast path: only one media type, no reordering needed748 final_pixel_values = []749 final_grid_thw = []750 751 if has_images:752 final_pixel_values.append(image_inputs["pixel_values"])753 final_grid_thw.extend(image_grid_thw)754 755 if has_videos:756 final_pixel_values.append(videos_inputs["pixel_values_videos"])757 final_grid_thw.extend(video_grid_thw)758 759 if final_pixel_values:760 output_data["pixel_values"] = np.concatenate(final_pixel_values, axis=0) if len(final_pixel_values) > 1 else final_pixel_values[0]761 762 if final_grid_thw:763 output_data["grid_thw"] = np.stack(final_grid_thw, axis=0)764 765 # Calculate media_nums_per_sample766 media_nums_per_sample = []767 for batch_idx in range(len(text)):768 media_order = media_order_per_sample[batch_idx]769 media_nums_per_sample.append(len(media_order) if len(media_order) > 0 else 1)770 771 # Don't add media_nums_per_sample to output_data yet772 # Will add it after BatchFeature to keep it as list773 774 else:775 # Slow path: both images and videos exist, need reordering776 final_pixel_values = []777 final_grid_thw = []778 media_nums_per_sample = []779 780 # Global indices to track position in flattened image/video arrays781 global_image_idx = 0782 global_video_idx = 0783 784 for batch_idx in range(len(text)):785 # Use the recorded media order from Step 3.0786 media_order = media_order_per_sample[batch_idx]787 788 if len(media_order) == 0:789 # If no media provided for this sample, add a blank image790 media_nums_per_sample.append(1)791 min_pixels = 128 * 128792 patch_size = getattr(self.image_processor, "patch_size", None) or 16793 temporal_patch_size = getattr(self.image_processor, "temporal_patch_size", None) or 1794 merge_size = getattr(self.image_processor, "merge_size", None) or 2795 796 factor = patch_size * merge_size797 side = int(np.ceil(np.sqrt(min_pixels) / factor) * factor)798 grid_h = side // patch_size799 grid_w = side // patch_size800 grid_t = 1801 802 # Channel = 3 (RGB)803 channel = 3804 dim = channel * temporal_patch_size * patch_size * patch_size805 num_patches = grid_t * grid_h * grid_w806 807 blank_pixel_values = np.zeros((num_patches, dim), dtype=np.float32)808 blank_grid_thw = np.array([grid_t, grid_h, grid_w], dtype=np.int64)809 810 final_pixel_values.append(blank_pixel_values)811 final_grid_thw.append(blank_grid_thw)812 else:813 media_nums_per_sample.append(len(media_order))814 815 # Collect media data according to the recorded order816 for media_type, _ in media_order:817 if media_type == "image" and image_grid_thw is not None:818 # Get image data819 if image_pixel_values_list:820 final_pixel_values.append(image_pixel_values_list[global_image_idx])821 final_grid_thw.append(image_grid_thw[global_image_idx])822 global_image_idx += 1823 elif media_type == "video" and video_grid_thw is not None:824 # Get video data825 if video_pixel_values_list:826 final_pixel_values.append(video_pixel_values_list[global_video_idx])827 final_grid_thw.append(video_grid_thw[global_video_idx])828 global_video_idx += 1829 830 # Concatenate/stack to unified format831 if final_pixel_values:832 output_data["pixel_values"] = _concat_array_or_tensor(833 final_pixel_values, axis=0834 )835 836 if final_grid_thw:837 output_data["grid_thw"] = _stack_array_or_tensor(838 final_grid_thw, axis=0839 )840 841 # Don't add media_nums_per_sample to output_data yet842 # Will add it after BatchFeature to keep it as list843 844 # Create cross_attention_mask using media_nums_per_sample845 if "input_ids" in output_data and "grid_thw" in output_data and media_nums_per_sample:846 cross_attention_mask = self._create_cross_attention_mask(847 output_data["input_ids"],848 output_data["grid_thw"],849 media_nums_per_sample,850 output_data.get("attention_mask", None)851 )852 output_data["cross_attention_mask"] = cross_attention_mask853 854 # Add labels to output if created855 if should_create_labels:856 output_data["labels"] = labels857 858 # BatchFeature will handle conversion to pt/tf/jax/np based on tensor_type859 batch_feature = BatchFeature(data=output_data, tensor_type=return_tensors)860 861 # Add media_nums_per_sample after BatchFeature to keep it as list (not tensor)862 if media_nums_per_sample:863 batch_feature["media_nums_per_sample"] = media_nums_per_sample864 865 return batch_feature866 867 def _create_cross_attention_mask(self, input_ids, grid_thw, media_nums_per_sample, attention_mask=None):868 """869 Create cross_attention_mask of shape (batch_size, 1, text_len, num_images).870 Video frames are treated as individual images.871 Mask values: True for masked, False for visible.872 Causal masking: text can see images that appear at or before the text position.873 874 Args:875 input_ids: List of token ids876 grid_thw: Grid sizes for each media item877 media_nums_per_sample: Number of media items per sample878 attention_mask: Optional attention mask to filter out padding positions879 """880 batch_size = len(input_ids)881 max_text_len = max(len(ids) for ids in input_ids)882 883 # Calculate total frames per sample to find max_num_frames884 total_frames_per_sample = []885 media_idx = 0886 for b in range(batch_size):887 num_media = media_nums_per_sample[b]888 if num_media == 0:889 total_frames_per_sample.append(0)890 continue891 892 sample_frames = 0893 for _ in range(num_media):894 # grid_thw is (N, 3) where first dim is t (num_frames)895 t = grid_thw[media_idx][0]896 if isinstance(t, torch.Tensor):897 t = int(t.item())898 else:899 t = int(t)900 sample_frames += t901 media_idx += 1902 total_frames_per_sample.append(sample_frames)903 904 max_num_frames = max(total_frames_per_sample) if total_frames_per_sample else 0905 906 if max_num_frames == 0:907 return None908 909 # Vectorized implementation for speed910 911 # 1. Pad input_ids to create a tensor912 # We use -1 as pad value since token ids are positive913 input_ids_tensor = torch.full((batch_size, max_text_len), -1, dtype=torch.long)914 for b, ids in enumerate(input_ids):915 l = len(ids)916 input_ids_tensor[b, :l] = torch.tensor(ids, dtype=torch.long)917 918 # 2. Identify image tokens919 is_image_token = (input_ids_tensor == self.image_token_id)920 921 # 3. Compute cumulative image tokens (how many image tokens appeared up to position t)922 # shape: (batch_size, text_len)923 cum_image_tokens = is_image_token.cumsum(dim=1)924 925 # 4. Create frame indices926 # shape: (1, 1, max_num_frames)927 frame_indices = torch.arange(max_num_frames).reshape(1, 1, -1)928 929 # 5. Determine visibility based on causal relationship930 # Text at `t` sees frame `i` if `cum_image_tokens[t] > i`931 # Because if frame `i` is the (i+1)-th image token, it becomes visible when count reaches i+1932 # shape: (batch_size, text_len, max_num_frames)933 visible_mask = cum_image_tokens.unsqueeze(-1) > frame_indices934 935 # 6. Apply attention_mask if provided936 if attention_mask is not None:937 # Convert to tensor if needed938 if isinstance(attention_mask, torch.Tensor):939 attn_mask_tensor = attention_mask940 else:941 # List of lists942 attn_mask_tensor = torch.zeros((batch_size, max_text_len), dtype=torch.long)943 for b, mask_row in enumerate(attention_mask):944 l = len(mask_row)945 attn_mask_tensor[b, :l] = torch.tensor(mask_row, dtype=torch.long)946 947 # shape: (batch_size, text_len, 1)948 valid_text = (attn_mask_tensor.unsqueeze(-1) == 1)949 visible_mask = visible_mask & valid_text950 951 # 7. Mask out frames that don't exist for a sample952 # shape: (batch_size, 1, 1)953 total_frames_tensor = torch.tensor(total_frames_per_sample).reshape(batch_size, 1, 1)954 # shape: (batch_size, 1, max_num_frames)955 valid_frames = frame_indices < total_frames_tensor956 957 visible_mask = visible_mask & valid_frames958 959 # 8. Create final mask (True for masked, False for visible)960 mask = ~visible_mask961 962 # 9. Add channel dimension: (batch_size, 1, text_len, max_num_frames)963 mask = mask.unsqueeze(1)964 965 return mask966 967 def _replace_and_update_spans(968 self,969 text: str,970 old_str: str,971 new_str: str,972 spans: List[List[int]],973 replace_count: int = -1974 ) -> tuple:975 """976 Replace occurrences of old_str with new_str and update spans accordingly.977 978 Args:979 text: The text to perform replacement on980 old_str: String to be replaced981 new_str: String to replace with982 spans: List of [start, end] spans to update (modified in place)983 replace_count: Maximum number of replacements (-1 for all)984 985 Returns:986 Tuple of (new_text, updated_spans)987 """988 delta = len(new_str) - len(old_str)989 result_text = text990 count = 0991 search_start = 0992 993 while True:994 pos = result_text.find(old_str, search_start)995 if pos == -1:996 break997 if replace_count != -1 and count >= replace_count:998 break999 1000 # Update all spans that come after this position1001 for span in spans:1002 if span[0] > pos:1003 # Span starts after replacement point1004 span[0] += delta1005 span[1] += delta1006 elif span[1] > pos:1007 # Span ends after replacement point (spans the replacement)1008 span[1] += delta1009 1010 # Perform the replacement1011 result_text = result_text[:pos] + new_str + result_text[pos + len(old_str):]1012 search_start = pos + len(new_str)1013 count += 11014 1015 return result_text, spans1016 1017 def _create_labels_from_spans(1018 self,1019 input_ids: List[List[int]],1020 offset_mapping: List[List[tuple]],1021 labels_spans: List[List[List[int]]],1022 ignore_index: int = -100,1023 mask_token_ids: Optional[set] = None1024 ) -> List[List[int]]:1025 """1026 Create labels from spans and offset_mapping.1027 1028 Args:1029 input_ids: Tokenized input ids1030 offset_mapping: Character offsets for each token from tokenizer (special tokens included)1031 labels_spans: Updated spans indicating assistant regions (after text transformations)1032 ignore_index: Value for masked positions1033 mask_token_ids: Set of token ids that should always be masked (set to ignore_index)1034 in labels, regardless of whether they fall inside a span.1035 Defaults to self.mask_token_ids if not provided.1036 1037 Returns:1038 labels: List of label ids, same shape as input_ids1039 1040 Note:1041 - Tokenizer's offset_mapping already includes correct offsets for special tokens in text1042 - Only need to mask tokens inside <|vision_start|>...<|vision_end|>1043 - Tokens whose id is in mask_token_ids are always masked1044 - All other tokens in spans (including special tokens like <|im_end|>) get labels1045 """1046 if mask_token_ids is None:1047 mask_token_ids = self.mask_token_ids1048 1049 batch_labels = []1050 1051 for batch_idx in range(len(input_ids)):1052 ids = input_ids[batch_idx]1053 offsets = offset_mapping[batch_idx]1054 spans = labels_spans[batch_idx]1055 1056 labels = [ignore_index] * len(ids)1057 1058 # Process each span: find token range and set labels1059 for span_start, span_end in spans:1060 in_vision = False1061 1062 # Find tokens that overlap with this span1063 for token_idx, (token_id, (char_start, char_end)) in enumerate(zip(ids, offsets)):1064 # Skip tokens completely before this span1065 if char_end <= span_start:1066 continue1067 # Stop when tokens are completely after this span1068 if char_start >= span_end:1069 break1070 1071 # Token overlaps with span, process it1072 # Track vision region: <|vision_start|> ... <|vision_end|>1073 if token_id == self.vision_start_token_id:1074 in_vision = True1075 continue1076 if token_id == self.vision_end_token_id:1077 in_vision = False1078 continue1079 1080 # Skip tokens inside vision region1081 if in_vision:1082 continue1083 1084 # Always mask special tokens that should never have labels1085 if token_id in mask_token_ids:1086 continue1087 1088 # Set label for this token1089 labels[token_idx] = token_id1090 1091 batch_labels.append(labels)1092 1093 return batch_labels1094 1095 def _calculate_timestamps(1096 self,1097 frames_indices: Optional[Union[List[int], np.ndarray]],1098 total_num_frames: int,1099 video_fps: float,1100 duration: float,1101 merge_size: int = 1,1102 actual_timestamps: Optional[List[float]] = None1103 ):1104 """1105 Calculate timestamps for video frames.1106 1107 Args:1108 frames_indices: Actual frame indices extracted (if available)1109 total_num_frames: Total number of sampled frames1110 video_fps: Video frames per second1111 duration: Video duration in seconds1112 merge_size: Temporal merge size1113 actual_timestamps: Pre-calculated actual timestamps (for segments)1114 1115 Returns:1116 List of timestamps (one per merged temporal patch)1117 """1118 # If actual timestamps are provided (from segment), use them directly1119 if actual_timestamps is not None:1120 timestamps = list(actual_timestamps)1121 1122 # Pad timestamps to be multiple of merge_size1123 if len(timestamps) % merge_size != 0:1124 timestamps.extend([timestamps[-1]] * (merge_size - len(timestamps) % merge_size))1125 1126 # Frames are merged by merge_size, so we average the timestamps within each temporal patch1127 timestamps = [1128 (timestamps[i] + timestamps[i + merge_size - 1]) / 21129 for i in range(0, len(timestamps), merge_size)1130 ]1131 return timestamps1132 1133 # Use frames_indices if available, otherwise generate uniformly sampled indices1134 if frames_indices is not None:1135 if isinstance(frames_indices, np.ndarray):1136 indices = frames_indices.tolist()1137 else:1138 indices = list(frames_indices)1139 else:1140 # Generate uniformly sampled frame indices1141 if total_num_frames <= 1:1142 indices = [0]1143 else:1144 # Uniformly sample frames across the video duration1145 indices = np.linspace(0, duration * video_fps - 1, total_num_frames).astype(np.int32).tolist()1146 1147 # Pad indices to be multiple of merge_size1148 if len(indices) % merge_size != 0:1149 indices.extend([indices[-1]] * (merge_size - len(indices) % merge_size))1150 1151 # Convert frame indices to timestamps1152 timestamps = [idx / video_fps for idx in indices]1153 1154 # Frames are merged by merge_size, so we average the timestamps within each temporal patch1155 timestamps = [1156 (timestamps[i] + timestamps[i + merge_size - 1]) / 21157 for i in range(0, len(timestamps), merge_size)1158 ]1159 return timestamps1160 1161 def batch_decode(self, *args, **kwargs):1162 """1163 This method forwards all its arguments to the tokenizer's batch_decode.1164 Please refer to the docstring of this method for more information.1165 """1166 return self.tokenizer.batch_decode(*args, **kwargs)1167 1168 def decode(self, *args, **kwargs):1169 """1170 This method forwards all its arguments to the tokenizer's decode.1171 Please refer to the docstring of this method for more information.1172 """1173 return self.tokenizer.decode(*args, **kwargs)1174 1175 def post_process_image_text_to_text(1176 self, generated_outputs, skip_special_tokens=True, clean_up_tokenization_spaces=False, **kwargs1177 ):1178 """1179 Post-process the output of the model to decode the text.1180 1181 Args:1182 generated_outputs (`torch.Tensor` or `np.ndarray`):1183 The output of the model `generate` function. The output is expected to be a tensor1184 of shape `(batch_size, sequence_length)` or `(sequence_length,)`.1185 skip_special_tokens (`bool`, *optional*, defaults to `True`):1186 Whether or not to remove special tokens in the output.1187 clean_up_tokenization_spaces (`bool`, *optional*, defaults to `False`):1188 Whether or not to clean up the tokenization spaces.1189 **kwargs:1190 Additional arguments to be passed to the tokenizer's `batch_decode` method.1191 1192 Returns:1193 `list[str]`: The decoded text.1194 """1195 return self.tokenizer.batch_decode(1196 generated_outputs,1197 skip_special_tokens=skip_special_tokens,1198 clean_up_tokenization_spaces=clean_up_tokenization_spaces,1199 **kwargs,1200 )