| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Multimodal preprocessing for Dots3Note checkpoints.""" |
| 4 | |
| 5 | import math |
| 6 | from collections.abc import Mapping, Sequence |
| 7 | from functools import cached_property |
| 8 | from typing import Any, cast |
| 9 | |
| 10 | import numpy as np |
| 11 | import torch |
| 12 | from PIL import Image |
| 13 | from transformers import BatchFeature |
| 14 | |
| 15 | from vllm.config.multimodal import ( |
| 16 | MultiModalDummyOptions, |
| 17 | ) |
| 18 | from vllm.inputs import MultiModalDataDict |
| 19 | from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargsItems |
| 20 | from vllm.multimodal.media import MediaWithBytes |
| 21 | from vllm.multimodal.parse import ( |
| 22 | MultiModalDataItems, |
| 23 | MultiModalDataParser, |
| 24 | VideoProcessorItems, |
| 25 | ) |
| 26 | from vllm.multimodal.processing import ( |
| 27 | BaseDummyInputsBuilder, |
| 28 | BaseMultiModalProcessor, |
| 29 | BaseProcessingInfo, |
| 30 | PromptReplacement, |
| 31 | PromptUpdate, |
| 32 | PromptUpdateDetails, |
| 33 | ) |
| 34 | from vllm.multimodal.processing.processor import HFMultiModalInputs |
| 35 | from vllm.transformers_utils.repo_utils import get_hf_file_to_dict |
| 36 | |
| 37 | from .video import preprocess_dots3_note_video |
| 38 | |
| 39 | IMAGE_START = "<|img|>" |
| 40 | IMAGE_PAD = "<|imgpad|>" |
| 41 | IMAGE_END = "<|endofimg|>" |
| 42 | AUDIO_START = "<|audio_comp_start|>" |
| 43 | AUDIO_PAD = "<|audio_comp_pad|>" |
| 44 | AUDIO_END = "<|audio_comp_end|>" |
| 45 | VIDEO_PLACEHOLDER = "<|video_pad|>" |
| 46 | |
| 47 | _HOP_LENGTH = 160 |
| 48 | _DEFAULT_SAMPLE_RATE = 16000 |
| 49 | |
| 50 | |
| 51 | def load_note_config_section( |
| 52 | model: str, |
| 53 | revision: str | None, |
| 54 | section: str, |
| 55 | ) -> dict[str, Any] | None: |
| 56 | config = get_hf_file_to_dict("config.json", model, revision) |
| 57 | value = (config or {}).get(section) |
| 58 | return value if isinstance(value, dict) else None |
| 59 | |
| 60 | |
| 61 | class Dots3NoteImageProcessor: |
| 62 | """CPU image preprocessing matching the native Cybertron ViT.""" |
| 63 | |
| 64 | def __init__( |
| 65 | self, |
| 66 | config: Mapping[str, Any], |
| 67 | image_details: Mapping[str, Any] | None = None, |
| 68 | ) -> None: |
| 69 | self.min_pixels = int(config["min_pixels"]) |
| 70 | self.max_pixels = int(config["max_pixels"]) |
| 71 | self.patch_size = int(config["patch_size"]) |
| 72 | self.temporal_patch_size = int(config["temporal_patch_size"]) |
| 73 | self.merge_size = int(config["merge_size"]) |
| 74 | self.pre_pixel_shuffle = bool(config.get("pre_pixel_shuffle", True)) |
| 75 | self.image_mean = np.asarray(config["image_mean"], dtype=np.float32) |
| 76 | self.image_std = np.asarray(config["image_std"], dtype=np.float32) |
| 77 | self.image_details = dict(image_details or {}) |
| 78 | |
| 79 | @property |
| 80 | def factor(self) -> int: |
| 81 | return self.patch_size * self.merge_size |
| 82 | |
| 83 | @staticmethod |
| 84 | def _round_by_factor(value: int, factor: int) -> int: |
| 85 | return round(value / factor) * factor |
| 86 | |
| 87 | @staticmethod |
| 88 | def _ceil_by_factor(value: float, factor: int) -> int: |
| 89 | return math.ceil(value / factor) * factor |
| 90 | |
| 91 | @staticmethod |
| 92 | def _floor_by_factor(value: float, factor: int) -> int: |
| 93 | return math.floor(value / factor) * factor |
| 94 | |
| 95 | def resized_size( |
| 96 | self, |
| 97 | width: int, |
| 98 | height: int, |
| 99 | *, |
| 100 | detail: str = "auto", |
| 101 | min_pixels: int | None = None, |
| 102 | max_pixels: int | None = None, |
| 103 | target_height: int | None = None, |
| 104 | target_width: int | None = None, |
| 105 | ) -> tuple[int, int]: |
| 106 | detail_config = self.image_details.get(detail, {}) |
| 107 | min_pixels = int( |
| 108 | min_pixels |
| 109 | if min_pixels is not None |
| 110 | else detail_config.get("min_pixels", self.min_pixels) |
| 111 | ) |
| 112 | max_pixels = int( |
| 113 | max_pixels |
| 114 | if max_pixels is not None |
| 115 | else detail_config.get("max_pixels", self.max_pixels) |
| 116 | ) |
| 117 | height = int(target_height or detail_config.get("target_height") or height) |
| 118 | width = int(target_width or detail_config.get("target_width") or width) |
| 119 | |
| 120 | factor = self.factor |
| 121 | if min(height, width) < factor // 4: |
| 122 | raise ValueError( |
| 123 | f"Image height and width must be at least {factor // 4}, " |
| 124 | f"got {height}x{width}" |
| 125 | ) |
| 126 | if max(height, width) / min(height, width) > 200: |
| 127 | raise ValueError("Image aspect ratio must be smaller than 200") |
| 128 | |
| 129 | resized_h = max(factor, self._round_by_factor(height, factor)) |
| 130 | resized_w = max(factor, self._round_by_factor(width, factor)) |
| 131 | if resized_h * resized_w > max_pixels: |
| 132 | beta = math.sqrt(height * width / max_pixels) |
| 133 | resized_h = max(factor, self._floor_by_factor(height / beta, factor)) |
| 134 | resized_w = max(factor, self._floor_by_factor(width / beta, factor)) |
| 135 | elif resized_h * resized_w < min_pixels: |
| 136 | beta = math.sqrt(min_pixels / (height * width)) |
| 137 | resized_h = self._ceil_by_factor(height * beta, factor) |
| 138 | resized_w = self._ceil_by_factor(width * beta, factor) |
| 139 | if resized_h * resized_w > max_pixels: |
| 140 | beta = math.sqrt(resized_h * resized_w / max_pixels) |
| 141 | resized_h = max(factor, self._floor_by_factor(resized_h / beta, factor)) |
| 142 | resized_w = max(factor, self._floor_by_factor(resized_w / beta, factor)) |
| 143 | return resized_h, resized_w |
| 144 | |
| 145 | def preprocess( |
| 146 | self, |
| 147 | image: Image.Image, |
| 148 | *, |
| 149 | detail: str = "auto", |
| 150 | min_pixels: int | None = None, |
| 151 | max_pixels: int | None = None, |
| 152 | target_height: int | None = None, |
| 153 | target_width: int | None = None, |
| 154 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 155 | if not isinstance(image, Image.Image): |
| 156 | raise TypeError(f"Expected a PIL image, got {type(image)}") |
| 157 | if image.mode == "RGBA": |
| 158 | background = Image.new("RGB", image.size, (255, 255, 255)) |
| 159 | background.paste(image, mask=image.getchannel("A")) |
| 160 | image = background |
| 161 | elif image.mode != "RGB": |
| 162 | image = image.convert("RGB") |
| 163 | |
| 164 | resized_h, resized_w = self.resized_size( |
| 165 | *image.size, |
| 166 | detail=detail, |
| 167 | min_pixels=min_pixels, |
| 168 | max_pixels=max_pixels, |
| 169 | target_height=target_height, |
| 170 | target_width=target_width, |
| 171 | ) |
| 172 | image = image.resize((resized_w, resized_h), Image.Resampling.BICUBIC) |
| 173 | array = np.asarray(image, dtype=np.float32) / 255.0 |
| 174 | array = (array - self.image_mean) / self.image_std |
| 175 | patches = array.transpose(2, 0, 1)[None] |
| 176 | if patches.shape[0] == 1: |
| 177 | patches = np.tile(patches, (self.temporal_patch_size, 1, 1, 1)) |
| 178 | |
| 179 | channel = patches.shape[1] |
| 180 | grid_t = patches.shape[0] // self.temporal_patch_size |
| 181 | grid_h = resized_h // self.patch_size |
| 182 | grid_w = resized_w // self.patch_size |
| 183 | if self.pre_pixel_shuffle: |
| 184 | patches = patches.reshape( |
| 185 | grid_t, |
| 186 | self.temporal_patch_size, |
| 187 | channel, |
| 188 | grid_h // self.merge_size, |
| 189 | self.merge_size, |
| 190 | self.patch_size, |
| 191 | grid_w // self.merge_size, |
| 192 | self.merge_size, |
| 193 | self.patch_size, |
| 194 | ) |
| 195 | patches = patches.transpose(0, 3, 6, 4, 7, 2, 1, 5, 8) |
| 196 | else: |
| 197 | patches = patches.reshape( |
| 198 | grid_t, |
| 199 | self.temporal_patch_size, |
| 200 | channel, |
| 201 | grid_h, |
| 202 | self.patch_size, |
| 203 | grid_w, |
| 204 | self.patch_size, |
| 205 | ) |
| 206 | patches = patches.transpose(0, 3, 5, 2, 1, 4, 6) |
| 207 | |
| 208 | pixel_values = torch.from_numpy( |
| 209 | patches.reshape( |
| 210 | grid_t * grid_h * grid_w, |
| 211 | channel * self.temporal_patch_size * self.patch_size * self.patch_size, |
| 212 | ) |
| 213 | ) |
| 214 | grid_thw = torch.tensor([grid_t, grid_h, grid_w], dtype=torch.long) |
| 215 | return pixel_values, grid_thw |
| 216 | |
| 217 | |
| 218 | class Dots3NoteProcessor: |
| 219 | """Small HF-like processor used by vLLM's multimodal frontend.""" |
| 220 | |
| 221 | def __init__( |
| 222 | self, |
| 223 | tokenizer, |
| 224 | image_processor: Dots3NoteImageProcessor | None, |
| 225 | *, |
| 226 | max_model_len: int, |
| 227 | video_audio_enabled: bool, |
| 228 | audio_token_stride: int, |
| 229 | ) -> None: |
| 230 | self.tokenizer = tokenizer |
| 231 | self.image_processor = image_processor |
| 232 | self.max_model_len = max_model_len |
| 233 | self.video_audio_enabled = video_audio_enabled |
| 234 | self.audio_token_stride = audio_token_stride |
| 235 | |
| 236 | def _token_ids(self, text: str) -> list[int]: |
| 237 | if hasattr(self.tokenizer, "encode"): |
| 238 | return self.tokenizer.encode(text, add_special_tokens=False) |
| 239 | token_ids = self.tokenizer(text, add_special_tokens=False)["input_ids"] |
| 240 | if token_ids and isinstance(token_ids[0], list): |
| 241 | return token_ids[0] |
| 242 | return token_ids |
| 243 | |
| 244 | def _process_video( |
| 245 | self, |
| 246 | video, |
| 247 | *, |
| 248 | detail: str, |
| 249 | size_overrides: Mapping[str, int | None], |
| 250 | seq: int, |
| 251 | output_reserve: int | None, |
| 252 | audio_cap: float, |
| 253 | audio_sample_rate: int, |
| 254 | k_mode: str, |
| 255 | max_new_tokens: int, |
| 256 | question: str, |
| 257 | ) -> dict[str, torch.Tensor]: |
| 258 | if self.image_processor is None: |
| 259 | raise ValueError("This NOTE checkpoint has no vision encoder") |
| 260 | if not self.video_audio_enabled: |
| 261 | audio_cap = 0.0 |
| 262 | parts = preprocess_dots3_note_video( |
| 263 | video, |
| 264 | tokenizer=self.tokenizer, |
| 265 | question=question, |
| 266 | seq=seq, |
| 267 | output_reserve=output_reserve, |
| 268 | audio_cap=audio_cap, |
| 269 | audio_sample_rate=audio_sample_rate, |
| 270 | k_mode=k_mode, |
| 271 | max_new_tokens=max_new_tokens, |
| 272 | ) |
| 273 | |
| 274 | fragments: list[str] = [] |
| 275 | pixel_values: list[torch.Tensor] = [] |
| 276 | image_grids: list[torch.Tensor] = [] |
| 277 | audio_values: list[torch.Tensor] = [] |
| 278 | modalities: list[int] = [] |
| 279 | image_pad_count = 0 |
| 280 | audio_pad_count = 0 |
| 281 | |
| 282 | for part in parts: |
| 283 | if part.kind == "text": |
| 284 | fragments.append(cast(str, part.value)) |
| 285 | continue |
| 286 | if part.kind == "image": |
| 287 | pixels, grid = self.image_processor.preprocess( |
| 288 | cast(Image.Image, part.value), |
| 289 | detail=detail, |
| 290 | min_pixels=size_overrides.get("min_pixels"), |
| 291 | max_pixels=size_overrides.get("max_pixels"), |
| 292 | target_height=size_overrides.get("target_height"), |
| 293 | target_width=size_overrides.get("target_width"), |
| 294 | ) |
| 295 | num_tokens = int(grid.prod()) // self.image_processor.merge_size**2 |
| 296 | fragments.append(f"{IMAGE_START}{IMAGE_PAD * num_tokens}{IMAGE_END}") |
| 297 | pixel_values.append(pixels) |
| 298 | image_grids.append(grid) |
| 299 | modalities.append(0) |
| 300 | image_pad_count += num_tokens |
| 301 | continue |
| 302 | |
| 303 | waveform = torch.from_numpy( |
| 304 | np.ascontiguousarray(part.value, dtype=np.float32) |
| 305 | ).view(torch.int32) |
| 306 | num_tokens = math.ceil(waveform.numel() / self.audio_token_stride) |
| 307 | fragments.append(f"{AUDIO_START}{AUDIO_PAD * num_tokens}{AUDIO_END}") |
| 308 | audio_values.append(waveform) |
| 309 | modalities.append(1) |
| 310 | audio_pad_count += num_tokens |
| 311 | |
| 312 | if not pixel_values: |
| 313 | raise ValueError("NOTE video preprocessing produced no frames") |
| 314 | fragment_ids = self._token_ids("".join(fragments)) |
| 315 | vocab = self.tokenizer.get_vocab() |
| 316 | embed_mask = torch.isin( |
| 317 | torch.tensor(fragment_ids, dtype=torch.long), |
| 318 | torch.tensor([vocab[IMAGE_PAD], vocab[AUDIO_PAD]], dtype=torch.long), |
| 319 | ) |
| 320 | expected_embeds = image_pad_count + audio_pad_count |
| 321 | if int(embed_mask.sum()) != expected_embeds: |
| 322 | raise ValueError( |
| 323 | "NOTE video placeholder expansion produced an invalid embedding mask: " |
| 324 | f"expected={expected_embeds}, actual={int(embed_mask.sum())}" |
| 325 | ) |
| 326 | |
| 327 | empty_audio = torch.empty(0, dtype=torch.int32) |
| 328 | return { |
| 329 | "video_pixel_values": torch.cat(pixel_values), |
| 330 | "video_image_grid_thw": torch.stack(image_grids), |
| 331 | "video_audio_values": ( |
| 332 | torch.cat(audio_values) if audio_values else empty_audio |
| 333 | ), |
| 334 | "video_audio_lengths": torch.tensor( |
| 335 | [value.numel() for value in audio_values], |
| 336 | dtype=torch.long, |
| 337 | ), |
| 338 | "video_modalities": torch.tensor(modalities, dtype=torch.uint8), |
| 339 | "video_input_ids": torch.tensor(fragment_ids, dtype=torch.long), |
| 340 | "video_embed_mask": embed_mask, |
| 341 | "video_frame_counts": torch.tensor([len(image_grids)], dtype=torch.long), |
| 342 | "video_audio_counts": torch.tensor([len(audio_values)], dtype=torch.long), |
| 343 | "video_emission_counts": torch.tensor([len(modalities)], dtype=torch.long), |
| 344 | "video_prompt_lengths": torch.tensor([len(fragment_ids)], dtype=torch.long), |
| 345 | "video_patch_counts": torch.tensor( |
| 346 | [sum(int(grid.prod()) for grid in image_grids)], dtype=torch.long |
| 347 | ), |
| 348 | "video_audio_sample_counts": torch.tensor( |
| 349 | [sum(value.numel() for value in audio_values)], dtype=torch.long |
| 350 | ), |
| 351 | } |
| 352 | |
| 353 | def __call__(self, text: str, **kwargs: object) -> BatchFeature: |
| 354 | modality_order = [ |
| 355 | key for key in kwargs if key in ("images", "audios", "videos") |
| 356 | ] |
| 357 | raw_images = kwargs.pop("images", None) |
| 358 | raw_audios = kwargs.pop("audios", None) |
| 359 | raw_videos = kwargs.pop("videos", None) |
| 360 | images = ( |
| 361 | list(cast(Sequence[Image.Image], raw_images)) |
| 362 | if raw_images is not None |
| 363 | else [] |
| 364 | ) |
| 365 | audios = ( |
| 366 | list(cast(Sequence[np.ndarray], raw_audios)) |
| 367 | if raw_audios is not None |
| 368 | else [] |
| 369 | ) |
| 370 | videos = ( |
| 371 | list(cast(Sequence[object], raw_videos)) if raw_videos is not None else [] |
| 372 | ) |
| 373 | if videos and (images or audios): |
| 374 | raise ValueError( |
| 375 | "Dots3Note does not support mixing a native video with " |
| 376 | "separate image/audio inputs" |
| 377 | ) |
| 378 | if len(videos) > 1: |
| 379 | raise ValueError("Dots3Note supports one video per request") |
| 380 | detail = kwargs.pop("image_detail", "auto") |
| 381 | details = list(detail) if isinstance(detail, (list, tuple)) else None |
| 382 | size_overrides: dict[str, int | None] = {} |
| 383 | for key in ("min_pixels", "max_pixels", "target_height", "target_width"): |
| 384 | if key in kwargs: |
| 385 | value = kwargs.pop(key) |
| 386 | size_overrides[key] = None if value is None else int(cast(Any, value)) |
| 387 | video_seq = int(cast(Any, kwargs.pop("seq", self.max_model_len))) |
| 388 | reserve_value = kwargs.pop("output_reserve", None) |
| 389 | output_reserve = ( |
| 390 | None if reserve_value is None else int(cast(Any, reserve_value)) |
| 391 | ) |
| 392 | audio_cap = float(cast(Any, kwargs.pop("audio_cap", 1.0))) |
| 393 | audio_sample_rate = int(cast(Any, kwargs.pop("audio_sr", _DEFAULT_SAMPLE_RATE))) |
| 394 | k_mode = str(kwargs.pop("k_mode", "eval_ek")) |
| 395 | max_new_tokens = int(cast(Any, kwargs.pop("max_new_tokens", 0))) |
| 396 | question = str(kwargs.pop("video_question", text)) |
| 397 | |
| 398 | tokenized = self.tokenizer(text, **kwargs) |
| 399 | input_ids = tokenized["input_ids"] |
| 400 | if isinstance(input_ids, torch.Tensor): |
| 401 | if input_ids.ndim == 1: |
| 402 | input_ids = input_ids.unsqueeze(0) |
| 403 | elif input_ids and isinstance(input_ids[0], int): |
| 404 | input_ids = [input_ids] |
| 405 | data: dict[str, object] = {"input_ids": input_ids} |
| 406 | |
| 407 | for modality in modality_order: |
| 408 | if modality == "images" and images: |
| 409 | if self.image_processor is None: |
| 410 | raise ValueError("This NOTE checkpoint has no vision encoder") |
| 411 | pixel_values = [] |
| 412 | grids = [] |
| 413 | for idx, image in enumerate(images): |
| 414 | image_detail = details[idx] if details is not None else str(detail) |
| 415 | pixels, grid = self.image_processor.preprocess( |
| 416 | image, |
| 417 | detail=image_detail, |
| 418 | min_pixels=size_overrides.get("min_pixels"), |
| 419 | max_pixels=size_overrides.get("max_pixels"), |
| 420 | target_height=size_overrides.get("target_height"), |
| 421 | target_width=size_overrides.get("target_width"), |
| 422 | ) |
| 423 | pixel_values.append(pixels) |
| 424 | grids.append(grid) |
| 425 | data["pixel_values"] = torch.cat(pixel_values) |
| 426 | data["image_grid_thw"] = torch.stack(grids) |
| 427 | elif modality == "audios" and audios: |
| 428 | waveforms = [ |
| 429 | torch.from_numpy( |
| 430 | np.ascontiguousarray(audio, dtype=np.float32) |
| 431 | ).view(torch.int32) |
| 432 | for audio in audios |
| 433 | ] |
| 434 | data["audio_values"] = torch.cat(waveforms) |
| 435 | data["audio_lengths"] = torch.tensor( |
| 436 | [waveform.numel() for waveform in waveforms], |
| 437 | dtype=torch.long, |
| 438 | ) |
| 439 | elif modality == "videos" and videos: |
| 440 | video_detail = details[0] if details is not None else str(detail) |
| 441 | data.update( |
| 442 | self._process_video( |
| 443 | videos[0], |
| 444 | detail=video_detail, |
| 445 | size_overrides=size_overrides, |
| 446 | seq=video_seq, |
| 447 | output_reserve=output_reserve, |
| 448 | audio_cap=audio_cap, |
| 449 | audio_sample_rate=audio_sample_rate, |
| 450 | k_mode=k_mode, |
| 451 | max_new_tokens=max_new_tokens, |
| 452 | question=question, |
| 453 | ) |
| 454 | ) |
| 455 | return BatchFeature(data=data) |
| 456 | |
| 457 | |
| 458 | class Dots3NoteProcessingInfo(BaseProcessingInfo): |
| 459 | @cached_property |
| 460 | def vision_config(self) -> dict[str, Any] | None: |
| 461 | model_config = self.ctx.model_config |
| 462 | return load_note_config_section( |
| 463 | model_config.model, |
| 464 | model_config.revision, |
| 465 | "vision_config", |
| 466 | ) |
| 467 | |
| 468 | @cached_property |
| 469 | def audio_config(self) -> dict[str, Any] | None: |
| 470 | model_config = self.ctx.model_config |
| 471 | return load_note_config_section( |
| 472 | model_config.model, |
| 473 | model_config.revision, |
| 474 | "audio_config", |
| 475 | ) |
| 476 | |
| 477 | @cached_property |
| 478 | def image_processor(self) -> Dots3NoteImageProcessor | None: |
| 479 | if self.vision_config is None: |
| 480 | return None |
| 481 | model_config = self.ctx.model_config |
| 482 | processor_config = get_hf_file_to_dict( |
| 483 | "preprocessor_config.json", |
| 484 | model_config.model, |
| 485 | model_config.revision, |
| 486 | ) |
| 487 | preprocessor_config = (processor_config or {}).get("vision_config") |
| 488 | if not isinstance(preprocessor_config, dict): |
| 489 | raise ValueError( |
| 490 | "NOTE vision checkpoint is missing preprocessor_config.json" |
| 491 | ) |
| 492 | image_details = preprocessor_config.get("image_details", {}) |
| 493 | return Dots3NoteImageProcessor(preprocessor_config, image_details) |
| 494 | |
| 495 | @cached_property |
| 496 | def processor(self) -> Dots3NoteProcessor: |
| 497 | mm_config = self.ctx.get_mm_config() |
| 498 | audio_config = self.audio_config or {} |
| 499 | audio_token_stride = ( |
| 500 | _HOP_LENGTH |
| 501 | * (8 if audio_config.get("use_conv2d_stem", True) else 2) |
| 502 | * int(audio_config.get("merge_factor", 1)) |
| 503 | ) |
| 504 | return Dots3NoteProcessor( |
| 505 | self.get_tokenizer(), |
| 506 | self.image_processor, |
| 507 | max_model_len=self.ctx.model_config.max_model_len, |
| 508 | video_audio_enabled=( |
| 509 | self.audio_config is not None |
| 510 | and ( |
| 511 | mm_config.get_limit_per_prompt("audio") > 0 |
| 512 | or mm_config.get_limit_per_prompt("video") > 0 |
| 513 | ) |
| 514 | ), |
| 515 | audio_token_stride=audio_token_stride, |
| 516 | ) |
| 517 | |
| 518 | def get_hf_processor(self, **kwargs: object) -> Dots3NoteProcessor: |
| 519 | return self.processor |
| 520 | |
| 521 | def get_data_parser(self) -> MultiModalDataParser: |
| 522 | sample_rate = int((self.audio_config or {}).get("sampling_rate", 16000)) |
| 523 | return MultiModalDataParser(target_sr=float(sample_rate), target_channels=1) |
| 524 | |
| 525 | def get_supported_mm_limits(self) -> Mapping[str, int | None]: |
| 526 | limits: dict[str, int | None] = {} |
| 527 | if self.vision_config is not None: |
| 528 | limits["image"] = 512 |
| 529 | limits["video"] = 1 |
| 530 | if self.audio_config is not None: |
| 531 | limits["audio"] = 128 |
| 532 | return limits |
| 533 | |
| 534 | def get_max_image_size(self) -> tuple[int, int]: |
| 535 | image_processor = self.image_processor |
| 536 | if image_processor is None: |
| 537 | return 0, 0 |
| 538 | factor = image_processor.factor |
| 539 | max_tokens = image_processor.max_pixels // (factor * factor) |
| 540 | for height_factor in range(math.isqrt(max_tokens), 0, -1): |
| 541 | if max_tokens % height_factor == 0: |
| 542 | width_factor = max_tokens // height_factor |
| 543 | if width_factor / height_factor <= 200: |
| 544 | return factor * width_factor, factor * height_factor |
| 545 | return factor * max_tokens, factor |
| 546 | |
| 547 | def get_mm_max_tokens_per_item( |
| 548 | self, |
| 549 | seq_len: int, |
| 550 | mm_counts: Mapping[str, int], |
| 551 | ) -> Mapping[str, int]: |
| 552 | result: dict[str, int] = {} |
| 553 | if self.image_processor is not None: |
| 554 | factor = self.image_processor.factor |
| 555 | result["image"] = self.image_processor.max_pixels // (factor * factor) |
| 556 | if self.audio_config is not None: |
| 557 | stride = ( |
| 558 | _HOP_LENGTH |
| 559 | * (8 if self.audio_config.get("use_conv2d_stem", True) else 2) |
| 560 | * int(self.audio_config.get("merge_factor", 1)) |
| 561 | ) |
| 562 | chunk_samples = int( |
| 563 | self.audio_config.get("chunk_seconds", 60) |
| 564 | * self.audio_config.get("sampling_rate", _DEFAULT_SAMPLE_RATE) |
| 565 | ) |
| 566 | result["audio"] = min(seq_len, math.ceil(chunk_samples / stride)) |
| 567 | if self.image_processor is not None: |
| 568 | result["video"] = max(1, seq_len - seq_len // 4) |
| 569 | return result |
| 570 | |
| 571 | |
| 572 | class Dots3NoteDummyInputsBuilder(BaseDummyInputsBuilder[Dots3NoteProcessingInfo]): |
| 573 | def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: |
| 574 | return ( |
| 575 | f"{IMAGE_START}{IMAGE_PAD}{IMAGE_END}" * mm_counts.get("image", 0) |
| 576 | + f"{AUDIO_START}{AUDIO_PAD}{AUDIO_END}" * mm_counts.get("audio", 0) |
| 577 | + VIDEO_PLACEHOLDER * mm_counts.get("video", 0) |
| 578 | ) |
| 579 | |
| 580 | def get_dummy_mm_data( |
| 581 | self, |
| 582 | seq_len: int, |
| 583 | mm_counts: Mapping[str, int], |
| 584 | mm_options: MultiModalDummyOptions, |
| 585 | ) -> MultiModalDataDict: |
| 586 | data: dict[str, Any] = {} |
| 587 | num_images = mm_counts.get("image", 0) |
| 588 | if num_images: |
| 589 | width, height = self.info.get_max_image_size() |
| 590 | data["image"] = self._get_dummy_images( |
| 591 | width=width, |
| 592 | height=height, |
| 593 | num_images=num_images, |
| 594 | overrides=mm_options.get("image"), |
| 595 | ) |
| 596 | num_audios = mm_counts.get("audio", 0) |
| 597 | if num_audios: |
| 598 | audio_config = self.info.audio_config or {} |
| 599 | stride = ( |
| 600 | _HOP_LENGTH |
| 601 | * (8 if audio_config.get("use_conv2d_stem", True) else 2) |
| 602 | * int(audio_config.get("merge_factor", 1)) |
| 603 | ) |
| 604 | chunk_samples = int( |
| 605 | audio_config.get("chunk_seconds", 60) |
| 606 | * audio_config.get("sampling_rate", _DEFAULT_SAMPLE_RATE) |
| 607 | ) |
| 608 | data["audio"] = self._get_dummy_audios( |
| 609 | length=min(chunk_samples, max(1, seq_len) * stride), |
| 610 | num_audios=num_audios, |
| 611 | overrides=mm_options.get("audio"), |
| 612 | ) |
| 613 | num_videos = mm_counts.get("video", 0) |
| 614 | if num_videos: |
| 615 | width, height = self.info.get_max_image_size() |
| 616 | factor = self.info.image_processor.factor # type: ignore[union-attr] |
| 617 | frame_tokens = max(1, width * height // (factor * factor)) |
| 618 | video_tokens = max(1, seq_len - seq_len // 4) |
| 619 | num_frames = max(4, math.ceil(video_tokens / frame_tokens)) |
| 620 | data["video"] = self._get_dummy_videos( |
| 621 | width=width, |
| 622 | height=height, |
| 623 | num_frames=num_frames, |
| 624 | num_videos=num_videos, |
| 625 | overrides=mm_options.get("video"), |
| 626 | ) |
| 627 | return data |
| 628 | |
| 629 | |
| 630 | class Dots3NoteMultiModalProcessor(BaseMultiModalProcessor[Dots3NoteProcessingInfo]): |
| 631 | def _get_hf_mm_inputs( |
| 632 | self, |
| 633 | mm_items: MultiModalDataItems, |
| 634 | hf_kwargs: Mapping[str, object], |
| 635 | ) -> HFMultiModalInputs: |
| 636 | hf_inputs = super()._get_hf_mm_inputs(mm_items, hf_kwargs) |
| 637 | |
| 638 | # The Dots3Note processor accepts "audios" instead of "audio" |
| 639 | hf_data = hf_inputs.hf_data |
| 640 | if "audio" in hf_data: |
| 641 | hf_data["audios"] = hf_data.pop("audio") |
| 642 | |
| 643 | if "video" in mm_items: |
| 644 | videos = mm_items.get_items("video", VideoProcessorItems) |
| 645 | |
| 646 | raw_videos: list[object] = [] |
| 647 | for index, item in enumerate(videos.data): |
| 648 | if isinstance(item, MediaWithBytes): |
| 649 | raw_videos.append(item.original_bytes) |
| 650 | else: |
| 651 | raw_videos.append(videos.get(index)) |
| 652 | |
| 653 | hf_data["videos"] = raw_videos |
| 654 | |
| 655 | return hf_inputs |
| 656 | |
| 657 | def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: |
| 658 | return self.dummy_inputs.get_dummy_text(mm_counts) |
| 659 | |
| 660 | def _get_mm_fields_config( |
| 661 | self, |
| 662 | hf_inputs: BatchFeature, |
| 663 | hf_processor_mm_kwargs: Mapping[str, object], |
| 664 | ) -> Mapping[str, MultiModalFieldConfig]: |
| 665 | fields: dict[str, MultiModalFieldConfig] = {} |
| 666 | if "pixel_values" in hf_inputs: |
| 667 | grid_thw = hf_inputs["image_grid_thw"] |
| 668 | fields["pixel_values"] = MultiModalFieldConfig.flat_from_sizes( |
| 669 | "image", grid_thw.prod(-1) |
| 670 | ) |
| 671 | fields["image_grid_thw"] = MultiModalFieldConfig.batched( |
| 672 | "image", keep_on_cpu=True |
| 673 | ) |
| 674 | if "audio_values" in hf_inputs: |
| 675 | fields["audio_values"] = MultiModalFieldConfig.flat_from_sizes( |
| 676 | "audio", hf_inputs["audio_lengths"] |
| 677 | ) |
| 678 | fields["audio_lengths"] = MultiModalFieldConfig.batched( |
| 679 | "audio", keep_on_cpu=True |
| 680 | ) |
| 681 | if "video_input_ids" in hf_inputs: |
| 682 | prompt_lengths = hf_inputs["video_prompt_lengths"] |
| 683 | frame_counts = hf_inputs["video_frame_counts"] |
| 684 | audio_counts = hf_inputs["video_audio_counts"] |
| 685 | emission_counts = hf_inputs["video_emission_counts"] |
| 686 | fields["video_pixel_values"] = MultiModalFieldConfig.flat_from_sizes( |
| 687 | "video", hf_inputs["video_patch_counts"] |
| 688 | ) |
| 689 | fields["video_image_grid_thw"] = MultiModalFieldConfig.flat_from_sizes( |
| 690 | "video", frame_counts, keep_on_cpu=True |
| 691 | ) |
| 692 | fields["video_audio_values"] = MultiModalFieldConfig.flat_from_sizes( |
| 693 | "video", hf_inputs["video_audio_sample_counts"] |
| 694 | ) |
| 695 | fields["video_audio_lengths"] = MultiModalFieldConfig.flat_from_sizes( |
| 696 | "video", audio_counts, keep_on_cpu=True |
| 697 | ) |
| 698 | fields["video_modalities"] = MultiModalFieldConfig.flat_from_sizes( |
| 699 | "video", emission_counts, keep_on_cpu=True |
| 700 | ) |
| 701 | fields["video_input_ids"] = MultiModalFieldConfig.flat_from_sizes( |
| 702 | "video", prompt_lengths, keep_on_cpu=True |
| 703 | ) |
| 704 | fields["video_embed_mask"] = MultiModalFieldConfig.flat_from_sizes( |
| 705 | "video", prompt_lengths, keep_on_cpu=True |
| 706 | ) |
| 707 | for key in ( |
| 708 | "video_frame_counts", |
| 709 | "video_audio_counts", |
| 710 | "video_emission_counts", |
| 711 | ): |
| 712 | fields[key] = MultiModalFieldConfig.batched("video", keep_on_cpu=True) |
| 713 | return fields |
| 714 | |
| 715 | def _get_prompt_updates( |
| 716 | self, |
| 717 | mm_items: MultiModalDataItems, |
| 718 | hf_processor_mm_kwargs: Mapping[str, object], |
| 719 | out_mm_kwargs: MultiModalKwargsItems, |
| 720 | ) -> Sequence[PromptUpdate]: |
| 721 | vocab = self.info.get_tokenizer().get_vocab() |
| 722 | updates: list[PromptUpdate] = [] |
| 723 | |
| 724 | if "image" in out_mm_kwargs: |
| 725 | image_start_id = vocab[IMAGE_START] |
| 726 | image_pad_id = vocab[IMAGE_PAD] |
| 727 | image_end_id = vocab[IMAGE_END] |
| 728 | merge_size = self.info.image_processor.merge_size # type: ignore[union-attr] |
| 729 | |
| 730 | def image_replacement(item_idx: int) -> PromptUpdateDetails: |
| 731 | grid = out_mm_kwargs["image"][item_idx]["image_grid_thw"].data |
| 732 | assert isinstance(grid, torch.Tensor) |
| 733 | num_tokens = int(grid.prod()) // merge_size**2 |
| 734 | return PromptUpdateDetails.select_token_id( |
| 735 | [image_start_id] + [image_pad_id] * num_tokens + [image_end_id], |
| 736 | image_pad_id, |
| 737 | ) |
| 738 | |
| 739 | updates.append( |
| 740 | PromptReplacement( |
| 741 | modality="image", |
| 742 | target=[image_start_id, image_pad_id, image_end_id], |
| 743 | replacement=image_replacement, |
| 744 | ) |
| 745 | ) |
| 746 | |
| 747 | if "audio" in out_mm_kwargs: |
| 748 | audio_start_id = vocab[AUDIO_START] |
| 749 | audio_pad_id = vocab[AUDIO_PAD] |
| 750 | audio_end_id = vocab[AUDIO_END] |
| 751 | config = self.info.audio_config or {} |
| 752 | stride = ( |
| 753 | _HOP_LENGTH |
| 754 | * (8 if config.get("use_conv2d_stem", True) else 2) |
| 755 | * int(config.get("merge_factor", 1)) |
| 756 | ) |
| 757 | |
| 758 | def audio_replacement(item_idx: int) -> PromptUpdateDetails: |
| 759 | length = out_mm_kwargs["audio"][item_idx]["audio_lengths"].data |
| 760 | assert isinstance(length, torch.Tensor) |
| 761 | num_tokens = math.ceil(int(length.item()) / stride) |
| 762 | return PromptUpdateDetails.select_token_id( |
| 763 | [audio_start_id] + [audio_pad_id] * num_tokens + [audio_end_id], |
| 764 | audio_pad_id, |
| 765 | ) |
| 766 | |
| 767 | updates.append( |
| 768 | PromptReplacement( |
| 769 | modality="audio", |
| 770 | target=[audio_start_id, audio_pad_id, audio_end_id], |
| 771 | replacement=audio_replacement, |
| 772 | ) |
| 773 | ) |
| 774 | if "video" in out_mm_kwargs: |
| 775 | |
| 776 | def video_replacement(item_idx: int) -> PromptUpdateDetails: |
| 777 | item = out_mm_kwargs["video"][item_idx] |
| 778 | input_ids_data = item["video_input_ids"].data |
| 779 | embed_mask_data = item["video_embed_mask"].data |
| 780 | assert isinstance(input_ids_data, torch.Tensor) |
| 781 | assert isinstance(embed_mask_data, torch.Tensor) |
| 782 | input_ids = input_ids_data.tolist() |
| 783 | embed_mask = embed_mask_data.bool() |
| 784 | |
| 785 | def select_video_embeds(full) -> torch.Tensor: |
| 786 | del full |
| 787 | return embed_mask |
| 788 | |
| 789 | return PromptUpdateDetails( |
| 790 | full=input_ids, |
| 791 | is_embed=select_video_embeds, |
| 792 | ) |
| 793 | |
| 794 | updates.append( |
| 795 | PromptReplacement( |
| 796 | modality="video", |
| 797 | target=[vocab[VIDEO_PLACEHOLDER]], |
| 798 | replacement=video_replacement, |
| 799 | ) |
| 800 | ) |
| 801 | return updates |