| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """HuggingFace processor for the MuseGlimmer multimodal model (text + image + video). |
| 4 | |
| 5 | This packages MuseGlimmer's vision preprocessing and token-span expansion the standard |
| 6 | HF way, so you can do: |
| 7 | |
| 8 | processor = AutoProcessor.from_pretrained(hf_dir, trust_remote_code=True) |
| 9 | messages = [{"role": "user", "content": [ |
| 10 | {"type": "image"}, |
| 11 | {"type": "text", "text": "Describe this image."}, |
| 12 | ]}] |
| 13 | text = processor.apply_chat_template(messages, add_generation_prompt=True) |
| 14 | inputs = processor(text=text, images=[pil_image], return_tensors="pt") |
| 15 | out = model.generate(**inputs, max_new_tokens=256) |
| 16 | |
| 17 | It reproduces the exact token layout the model expects: |
| 18 | |
| 19 | image -> <|image_start|> + <|patch|> * N + <|image_end|> |
| 20 | video -> <|vid_start|> |
| 21 | ( "Time: X.Xs" + <|video|> * P [+ <|vid_frame_separator|>] )* |
| 22 | + <|vid_end|> |
| 23 | |
| 24 | where N = floor(H/(patch*ds)) * floor(W/(patch*ds)) (<= 4096) for images and |
| 25 | P = the same grid count per temporal frame-group (<= 144) for video, with the |
| 26 | target (H, W) chosen by MuseGlimmerImageProcessor.compute_image_size / |
| 27 | MuseGlimmerVideoProcessor.compute_video_frame_size (the same grid logic as |
| 28 | modeling_muse_glimmer.MuseGlimmerVisionEncoder). |
| 29 | |
| 30 | The chat template emits one sentinel per media item -- ``<|image|>`` for images, |
| 31 | ``<|video|>`` for videos -- which ``MuseGlimmerProcessor.__call__`` expands into the |
| 32 | spans above using each media item's computed token count. Vision features are |
| 33 | returned in ``pixel_values`` (a list of per-image / per-frame-group tensors) in |
| 34 | the order the sentinels appear, matching how ``MuseGlimmerModel`` merges them at the |
| 35 | ``<|patch|>`` / ``<|video|>`` positions. |
| 36 | """ |
| 37 | |
| 38 | from __future__ import annotations |
| 39 | |
| 40 | import itertools |
| 41 | import math |
| 42 | from pathlib import Path |
| 43 | |
| 44 | import torch |
| 45 | from PIL import Image |
| 46 | from torchvision import transforms as T |
| 47 | from transformers.feature_extraction_utils import BatchFeature |
| 48 | from transformers.image_processing_utils import BaseImageProcessor |
| 49 | from transformers.processing_utils import ProcessorMixin |
| 50 | from transformers.video_processing_utils import BaseVideoProcessor |
| 51 | |
| 52 | # Single-token sentinels emitted by the chat template, one per media item. |
| 53 | IMAGE_SENTINEL = "<|image|>" # id 200090; absent from the image/video spans |
| 54 | VIDEO_SENTINEL = "<|video|>" # id 200091; also the per-frame token (see __call__) |
| 55 | |
| 56 | # Jinja chat template (multimodal superset of convert_muse_glimmer_to_hf.CHAT_TEMPLATE). |
| 57 | # Message content may be a plain string OR a list of parts |
| 58 | # ({"type": "text"|"image"|"video", ...}); image/video parts render as a single |
| 59 | # sentinel that __call__ expands. |
| 60 | MUSE_GLIMMER_MM_CHAT_TEMPLATE = ( |
| 61 | "{{- bos_token -}}" |
| 62 | "{%- macro render_parts(content) -%}" |
| 63 | "{%- if content is string -%}{{- content -}}" |
| 64 | "{%- else -%}" |
| 65 | "{%- for part in content -%}" |
| 66 | "{%- if part['type'] == 'image' -%}{{- '<|image|>' -}}" |
| 67 | "{%- elif part['type'] == 'video' -%}{{- '<|video|>' -}}" |
| 68 | "{%- elif part['type'] == 'text' -%}{{- part['text'] -}}" |
| 69 | "{%- endif -%}" |
| 70 | "{%- endfor -%}" |
| 71 | "{%- endif -%}" |
| 72 | "{%- endmacro -%}" |
| 73 | # At inference (add_generation_prompt), inject a default system message when |
| 74 | # the caller supplied none. Full-transcript rendering |
| 75 | # (add_generation_prompt=False) is left unchanged. |
| 76 | "{%- set ns = namespace(has_system=false) -%}" |
| 77 | "{%- for m in messages -%}" |
| 78 | "{%- if m['role'] == 'system' -%}{%- set ns.has_system = true -%}{%- endif -%}" |
| 79 | "{%- endfor -%}" |
| 80 | "{%- if add_generation_prompt and not ns.has_system -%}" |
| 81 | "{{- '<|start|>system<|message|>You are a helpful assistant.<|eot|>' -}}" |
| 82 | "{%- endif -%}" |
| 83 | "{%- for message in messages -%}" |
| 84 | "{%- set role = message['role'] -%}" |
| 85 | "{%- if role == 'assistant' -%}" |
| 86 | "{%- set recipient = message.get('recipient') -%}" |
| 87 | "{%- set end_turn = message.get('end_turn') -%}" |
| 88 | "{%- if end_turn is none -%}" |
| 89 | "{%- set end_turn = not (recipient and recipient != 'user') -%}" |
| 90 | "{%- endif -%}" |
| 91 | "{{- '<|start|>assistant' -}}" |
| 92 | "{%- if recipient -%}{{- ' to=' + recipient -}}{%- endif -%}" |
| 93 | "{{- '<|message|>' -}}{{- render_parts(message['content']) -}}" |
| 94 | "{{- ('<|eot|>' if end_turn else '<|eom|>') -}}" |
| 95 | "{%- elif role == 'tool' -%}" |
| 96 | "{%- set name = message.get('name', '') -%}" |
| 97 | # Tool content is emitted as-is (string body or interleaved image/text parts). |
| 98 | # Any <tool_output ...> wrapper is baked into the SFT data content; the |
| 99 | # tokenizer does not add it, so the template must not either. |
| 100 | "{{- '<|start|>tool ' + name + '<|message|>' -}}" |
| 101 | "{{- render_parts(message['content']) -}}" |
| 102 | "{{- '<|eot|>' -}}" |
| 103 | "{%- else -%}" |
| 104 | "{%- set header = role -%}" |
| 105 | "{%- if message.get('name') -%}" |
| 106 | "{%- set header = role + ' ' + message['name'] -%}" |
| 107 | "{%- endif -%}" |
| 108 | "{{- '<|start|>' + header + '<|message|>' -}}" |
| 109 | "{{- render_parts(message['content']) -}}" |
| 110 | "{{- '<|eot|>' -}}" |
| 111 | "{%- endif -%}" |
| 112 | "{%- endfor -%}" |
| 113 | "{%- if add_generation_prompt -%}{{- '<|start|>assistant' -}}{%- endif -%}" |
| 114 | ) |
| 115 | |
| 116 | |
| 117 | def _grid_size( |
| 118 | img_w: int, img_h: int, patch_hw: int, max_tokens: int |
| 119 | ) -> tuple[int, int, int]: |
| 120 | """Pick the integer (H, W) grid closest to the aspect ratio under the token cap. |
| 121 | |
| 122 | Replicates MuseGlimmerVisionEncoder._compute_grid_size |
| 123 | (modeling_muse_glimmer.py) so the processor needs no torch model import. |
| 124 | Returns (target_h, target_w, n_tokens). |
| 125 | """ |
| 126 | i_nph = img_h / patch_hw |
| 127 | i_npw = img_w / patch_hw |
| 128 | ratio = i_npw / i_nph if i_nph > 0 else 1.0 |
| 129 | if i_nph * i_npw > max_tokens: |
| 130 | i_nph = (max_tokens / ratio) ** 0.5 |
| 131 | i_npw = i_nph * ratio |
| 132 | candidates = list( |
| 133 | set( |
| 134 | itertools.product( |
| 135 | [math.floor(i_nph), math.ceil(i_nph)], |
| 136 | [math.floor(i_npw), math.ceil(i_npw)], |
| 137 | ) |
| 138 | ) |
| 139 | ) |
| 140 | candidates = [ |
| 141 | (nph, npw) |
| 142 | for nph, npw in candidates |
| 143 | if nph >= 1 and npw >= 1 and nph * npw <= max_tokens |
| 144 | ] |
| 145 | if not candidates: |
| 146 | candidates = [(max(1, round(i_nph)), max(1, round(i_npw)))] |
| 147 | nph, npw = min(candidates, key=lambda c: abs(c[0] / c[1] - img_h / img_w)) |
| 148 | return nph * patch_hw, npw * patch_hw, nph * npw |
| 149 | |
| 150 | |
| 151 | class MuseGlimmerImageProcessor(BaseImageProcessor): |
| 152 | """Resize + normalize MuseGlimmer images and compute patch-token counts. |
| 153 | |
| 154 | Variable-resolution: each image is resized to the grid that best matches its |
| 155 | aspect ratio under the per-image token cap, then normalized with mean/std 0.5. |
| 156 | Returns per-image tensors (not stacked) because MuseGlimmer consumes a list of |
| 157 | variable-size images. Video frames are handled by MuseGlimmerVideoProcessor. |
| 158 | """ |
| 159 | |
| 160 | model_input_names = ["pixel_values"] |
| 161 | |
| 162 | def __init__( |
| 163 | self, |
| 164 | patch_size: int = 14, |
| 165 | downsample_factor: int = 2, |
| 166 | max_image_tokens: int = 4096, |
| 167 | image_mean: float = 0.5, |
| 168 | image_std: float = 0.5, |
| 169 | **kwargs, |
| 170 | ): |
| 171 | super().__init__(**kwargs) |
| 172 | self.patch_size = patch_size |
| 173 | self.downsample_factor = downsample_factor |
| 174 | self.max_image_tokens = max_image_tokens |
| 175 | self.image_mean = image_mean |
| 176 | self.image_std = image_std |
| 177 | |
| 178 | def _to_norm_tensor(self, image: Image.Image) -> torch.Tensor: |
| 179 | # Functional normalize -- no stored Normalize object, so the processor |
| 180 | # stays JSON-serializable for to_dict() / save_pretrained. |
| 181 | return T.functional.normalize( |
| 182 | T.functional.to_tensor(image), |
| 183 | [self.image_mean] * 3, |
| 184 | [self.image_std] * 3, |
| 185 | ) |
| 186 | |
| 187 | # -- size computation (mirrors modeling_muse_glimmer.MuseGlimmerVisionEncoder) -- |
| 188 | def compute_image_size(self, img_w: int, img_h: int) -> tuple[int, int, int]: |
| 189 | ph = self.patch_size * self.downsample_factor |
| 190 | return _grid_size(img_w, img_h, ph, self.max_image_tokens) |
| 191 | |
| 192 | # -- preprocessing --------------------------------------------------------- |
| 193 | def preprocess_image(self, image: Image.Image) -> tuple[torch.Tensor, int]: |
| 194 | """Return (pixel tensor [3, H, W], n_patch_tokens) for one image.""" |
| 195 | image = image.convert("RGB") |
| 196 | target_h, target_w, n_tokens = self.compute_image_size( |
| 197 | image.width, image.height |
| 198 | ) |
| 199 | image = image.resize((target_w, target_h), Image.LANCZOS) |
| 200 | return self._to_norm_tensor(image), n_tokens |
| 201 | |
| 202 | |
| 203 | # torchcodec is the training-faithful video decoder; other decoders (torchvision / |
| 204 | # PyAV) seek + color-convert differently and diverge from training. Optional at |
| 205 | # import time so text/image paths load without it; decode_video raises if used. |
| 206 | try: |
| 207 | import torchcodec |
| 208 | except Exception: |
| 209 | torchcodec = None |
| 210 | |
| 211 | |
| 212 | class MuseGlimmerVideoProcessor(BaseVideoProcessor): |
| 213 | """MuseGlimmer video preprocessing behind the standard HF |
| 214 | ``AutoVideoProcessor`` API. |
| 215 | |
| 216 | Exposes MuseGlimmer's training-faithful video handling as a first-class video |
| 217 | processor so you can call ``processor(videos="clip.mp4")`` and downstream |
| 218 | users can discover it via the standard HF ``AutoVideoProcessor`` API plus a |
| 219 | ``video_preprocessor_config.json``. It wraps: |
| 220 | |
| 221 | * torchcodec decode (matches training; other decoders diverge), |
| 222 | * uniform frame sampling to a whole multiple of ``patch_temporal``, |
| 223 | * real per-group PTS timestamps (rendered as ``Time: X.Xs`` by |
| 224 | MuseGlimmerProcessor), |
| 225 | * ``patch_temporal`` frame-grouping (frames cat on the channel axis -> |
| 226 | ``[patch_temporal * 3, H, W]``; the encoder detects video by channel count). |
| 227 | |
| 228 | It deliberately overrides ``preprocess`` instead of the BaseVideoProcessor fast |
| 229 | pipeline (group_videos_by_shape / smart_resize / single stacked tensor): MuseGlimmer |
| 230 | consumes a LIST of variable-size group tensors and needs the real per-group PTS, |
| 231 | neither of which the stacked fast path models. |
| 232 | """ |
| 233 | |
| 234 | model_input_names = ["pixel_values"] |
| 235 | |
| 236 | def __init__( |
| 237 | self, |
| 238 | patch_size: int = 14, |
| 239 | downsample_factor: int = 2, |
| 240 | patch_temporal: int = 2, |
| 241 | max_video_frame_tokens: int = 144, |
| 242 | image_mean: float = 0.5, |
| 243 | image_std: float = 0.5, |
| 244 | video_num_frames: int = 96, |
| 245 | video_sampling_fps: float = 2.0, |
| 246 | **kwargs, |
| 247 | ): |
| 248 | super().__init__(**kwargs) |
| 249 | self.patch_size = patch_size |
| 250 | self.downsample_factor = downsample_factor |
| 251 | self.patch_temporal = patch_temporal |
| 252 | self.max_video_frame_tokens = max_video_frame_tokens |
| 253 | self.image_mean = image_mean |
| 254 | self.image_std = image_std |
| 255 | self.video_num_frames = video_num_frames |
| 256 | self.video_sampling_fps = video_sampling_fps |
| 257 | |
| 258 | def _to_norm_tensor(self, image: Image.Image) -> torch.Tensor: |
| 259 | return T.functional.normalize( |
| 260 | T.functional.to_tensor(image), |
| 261 | [self.image_mean] * 3, |
| 262 | [self.image_std] * 3, |
| 263 | ) |
| 264 | |
| 265 | def decode_video(self, video_path: str) -> tuple[list[Image.Image], list[float]]: |
| 266 | """Sample frames + per-group PTS with torchcodec (the training decode path). |
| 267 | |
| 268 | ``timestamps[g]`` is the ACTUAL decoded PTS of the first frame in temporal |
| 269 | group ``g``; ``len(frames)`` is a whole multiple of ``patch_temporal``. |
| 270 | """ |
| 271 | if torchcodec is None: |
| 272 | raise RuntimeError( |
| 273 | "torchcodec is required for video decoding (it matches the training " |
| 274 | "decode path). See the Environment Setup steps in hf/README.md." |
| 275 | ) |
| 276 | pt = self.patch_temporal |
| 277 | reader = torchcodec.decoders.VideoDecoder(video_path) |
| 278 | total = len(reader) |
| 279 | assert reader.metadata.average_fps is not None, ( |
| 280 | f"Video has no FPS metadata: {video_path}" |
| 281 | ) |
| 282 | fps = reader.metadata.average_fps |
| 283 | assert self.video_sampling_fps and self.video_sampling_fps > 0, ( |
| 284 | f"video_sampling_fps must be positive, got {self.video_sampling_fps}" |
| 285 | ) |
| 286 | n = min( |
| 287 | int(total * self.video_sampling_fps / fps), self.video_num_frames, total |
| 288 | ) |
| 289 | n = max(pt, (n // pt) * pt) |
| 290 | n = min(n, total) |
| 291 | if n < pt: |
| 292 | raise ValueError( |
| 293 | f"Video has only {total} decodable frame(s) but needs at least " |
| 294 | f"{pt} (one temporal patch): {video_path}" |
| 295 | ) |
| 296 | indices = torch.linspace(0, total - 1, n).long().tolist() |
| 297 | frames: list[Image.Image] = [] |
| 298 | timestamps: list[float] = [] |
| 299 | for j, i in enumerate(indices): |
| 300 | fr = reader[i] |
| 301 | frames.append( |
| 302 | Image.fromarray(fr.data.permute(1, 2, 0).numpy()).convert("RGB") |
| 303 | ) |
| 304 | if j % pt == 0: |
| 305 | pts = getattr(fr, "pts_seconds", None) |
| 306 | timestamps.append(float(pts) if pts is not None else i / fps) |
| 307 | return frames, timestamps |
| 308 | |
| 309 | def compute_video_frame_size(self, img_w: int, img_h: int) -> tuple[int, int, int]: |
| 310 | """Pick (H, W, tokens_per_frame) for a video frame under the per-frame cap. |
| 311 | |
| 312 | Mirrors MuseGlimmerImageProcessor.compute_image_size but uses the (smaller) |
| 313 | per-frame video token budget; shares the same ``_grid_size`` grid logic as |
| 314 | modeling_muse_glimmer.MuseGlimmerVisionEncoder. |
| 315 | """ |
| 316 | ph = self.patch_size * self.downsample_factor |
| 317 | return _grid_size(img_w, img_h, ph, self.max_video_frame_tokens) |
| 318 | |
| 319 | def _group_frames( |
| 320 | self, frames: list[Image.Image] |
| 321 | ) -> tuple[list[torch.Tensor], int, int]: |
| 322 | """Resize/normalize frames and cat ``patch_temporal`` frames per group. |
| 323 | |
| 324 | Returns (group_tensors, n_groups, tokens_per_group). |
| 325 | """ |
| 326 | pt = self.patch_temporal |
| 327 | if pt <= 0: |
| 328 | raise ValueError(f"patch_temporal must be positive, got {pt}") |
| 329 | if not frames: |
| 330 | raise ValueError("video must contain at least one frame") |
| 331 | if padding := (-len(frames)) % pt: |
| 332 | frames = [*frames, *([frames[-1]] * padding)] |
| 333 | first = frames[0].convert("RGB") |
| 334 | target_h, target_w, n_tokens = self.compute_video_frame_size( |
| 335 | first.width, first.height |
| 336 | ) |
| 337 | groups: list[torch.Tensor] = [] |
| 338 | for i in range(0, len(frames), pt): |
| 339 | grp = [ |
| 340 | self._to_norm_tensor( |
| 341 | frames[i + j] |
| 342 | .convert("RGB") |
| 343 | .resize((target_w, target_h), Image.LANCZOS) |
| 344 | ) |
| 345 | for j in range(pt) |
| 346 | ] |
| 347 | groups.append(torch.cat(grp, dim=0)) |
| 348 | return groups, len(groups), n_tokens |
| 349 | |
| 350 | @staticmethod |
| 351 | def _normalize_videos(videos) -> list: |
| 352 | """Normalize the ``videos`` arg to a list of per-video items. |
| 353 | |
| 354 | A video item is a path (str/Path) or a list of PIL frames. Accepts a |
| 355 | single video (one path, or one list of frames) or a list of those. |
| 356 | """ |
| 357 | if isinstance(videos, (str, Path)): |
| 358 | return [videos] |
| 359 | if videos and not isinstance(videos[0], (list, tuple, str, Path)): |
| 360 | return [videos] # a single list of PIL frames |
| 361 | return list(videos or []) |
| 362 | |
| 363 | def preprocess_one( |
| 364 | self, |
| 365 | video: str | Path | list[Image.Image], |
| 366 | timestamps: list[float] | None = None, |
| 367 | ) -> tuple[list[torch.Tensor], int, int, list[float]]: |
| 368 | """Decode (if a path) + group ONE video. |
| 369 | |
| 370 | Returns (group_tensors, n_groups, tokens_per_group, group_timestamps). This |
| 371 | is the single per-video implementation shared by ``preprocess`` (the HF |
| 372 | AutoVideoProcessor batch API) and ``MuseGlimmerProcessor.__call__`` (prompt |
| 373 | building), so the two paths cannot drift. For a path, frames + real |
| 374 | per-group PTS come from ``decode_video``; for pre-decoded frames the given |
| 375 | ``timestamps`` are used verbatim (empty if not supplied). |
| 376 | """ |
| 377 | if isinstance(video, (str, Path)): |
| 378 | frames, ts = self.decode_video(str(video)) |
| 379 | else: |
| 380 | frames = [f.convert("RGB") for f in video] |
| 381 | ts = list(timestamps) if timestamps is not None else [] |
| 382 | groups, n_groups, tokens_per_group = self._group_frames(frames) |
| 383 | return groups, n_groups, tokens_per_group, ts |
| 384 | |
| 385 | def preprocess( |
| 386 | self, |
| 387 | videos, |
| 388 | video_timestamps: list[list[float]] | None = None, |
| 389 | return_tensors: str | None = "pt", |
| 390 | **kwargs, |
| 391 | ) -> BatchFeature: |
| 392 | """Preprocess video file path(s) or pre-decoded frame list(s). |
| 393 | |
| 394 | Accepts a single video (path str, or list of PIL frames) or a list of |
| 395 | those. Returns a BatchFeature carrying, in video order: |
| 396 | * ``pixel_values`` -- flat list of [pt*3, H, W] group tensors, |
| 397 | * ``video_num_groups`` -- groups per video, |
| 398 | * ``video_tokens_per_group`` -- patch-token count per group per video, |
| 399 | * ``video_timestamps`` -- per-group PTS per video, |
| 400 | i.e. exactly what MuseGlimmerProcessor needs to build the |
| 401 | ``<|vid_start|> ( Time: X.Xs <|video|>*P [sep] )* <|vid_end|>`` block. |
| 402 | """ |
| 403 | videos = self._normalize_videos(videos) |
| 404 | pixel_values: list[torch.Tensor] = [] |
| 405 | num_groups: list[int] = [] |
| 406 | tokens_per_group: list[int] = [] |
| 407 | out_ts: list[list[float]] = [] |
| 408 | for idx, v in enumerate(videos): |
| 409 | ts_in = video_timestamps[idx] if video_timestamps else None |
| 410 | groups, ng, tpg, ts = self.preprocess_one(v, ts_in) |
| 411 | pixel_values += groups |
| 412 | num_groups.append(ng) |
| 413 | tokens_per_group.append(tpg) |
| 414 | out_ts.append(ts) |
| 415 | batch = BatchFeature( |
| 416 | data={ |
| 417 | "video_num_groups": num_groups, |
| 418 | "video_tokens_per_group": tokens_per_group, |
| 419 | "video_timestamps": out_ts, |
| 420 | }, |
| 421 | tensor_type=None, |
| 422 | ) |
| 423 | # Variable-size list -- BatchFeature would fail to stack into one tensor. |
| 424 | batch["pixel_values"] = pixel_values |
| 425 | return batch |
| 426 | |
| 427 | |
| 428 | class MuseGlimmerProcessor(ProcessorMixin): |
| 429 | """Bundle MuseGlimmerImageProcessor + MuseGlimmerVideoProcessor + tokenizer; |
| 430 | expand media sentinels into spans. |
| 431 | |
| 432 | Images go through ``image_processor`` (MuseGlimmerImageProcessor) and videos through |
| 433 | ``video_processor`` (MuseGlimmerVideoProcessor) -- one preprocessing implementation |
| 434 | each, no overlap. |
| 435 | """ |
| 436 | |
| 437 | attributes = ["image_processor", "video_processor", "tokenizer"] |
| 438 | image_processor_class = "AutoImageProcessor" |
| 439 | video_processor_class = "AutoVideoProcessor" |
| 440 | tokenizer_class = "TokenizersBackend" |
| 441 | |
| 442 | def __init__( |
| 443 | self, |
| 444 | image_processor=None, |
| 445 | video_processor=None, |
| 446 | tokenizer=None, |
| 447 | chat_template=None, |
| 448 | **kwargs, |
| 449 | ): |
| 450 | if image_processor is None: |
| 451 | image_processor = MuseGlimmerImageProcessor() |
| 452 | if video_processor is None: |
| 453 | video_processor = MuseGlimmerVideoProcessor() |
| 454 | super().__init__( |
| 455 | image_processor, |
| 456 | video_processor, |
| 457 | tokenizer, |
| 458 | chat_template=chat_template or MUSE_GLIMMER_MM_CHAT_TEMPLATE, |
| 459 | **kwargs, |
| 460 | ) |
| 461 | |
| 462 | def _sid(self, token: str) -> int: |
| 463 | return self.tokenizer.convert_tokens_to_ids(token) |
| 464 | |
| 465 | def _image_block(self, n_tokens: int) -> list[int]: |
| 466 | return ( |
| 467 | [self._sid("<|image_start|>")] |
| 468 | + [self._sid("<|patch|>")] * n_tokens |
| 469 | + [self._sid("<|image_end|>")] |
| 470 | ) |
| 471 | |
| 472 | def _video_block( |
| 473 | self, |
| 474 | n_groups: int, |
| 475 | tokens_per_group: int, |
| 476 | timestamps: list[float] | None = None, |
| 477 | ) -> list[int]: |
| 478 | """Per-group ``Time: X.Xs`` + <|video|>*P, separated/terminated. |
| 479 | |
| 480 | ``timestamps`` (one per temporal group, in seconds) should be the ACTUAL |
| 481 | decoded frame times -- training renders the real factored frame PTS, not a |
| 482 | uniform grid. Falls back to ``g*patch_temporal/fps`` only when timestamps |
| 483 | are not supplied (e.g. callers that don't track frame times); that |
| 484 | approximation can differ from training by the per-frame jitter. |
| 485 | """ |
| 486 | vid = self._sid("<|video|>") |
| 487 | sep = self._sid("<|vid_frame_separator|>") |
| 488 | pt = self.video_processor.patch_temporal |
| 489 | fps = self.video_processor.video_sampling_fps |
| 490 | block = [self._sid("<|vid_start|>")] |
| 491 | for g in range(n_groups): |
| 492 | ts = timestamps[g] if timestamps is not None else g * pt / fps |
| 493 | block += self.tokenizer.encode(f"Time: {ts:.1f}s", add_special_tokens=False) |
| 494 | block += [vid] * tokens_per_group |
| 495 | block.append(sep if g < n_groups - 1 else self._sid("<|vid_end|>")) |
| 496 | return block |
| 497 | |
| 498 | def __call__( |
| 499 | self, |
| 500 | text: str | list[str] | None = None, |
| 501 | images: list[Image.Image] | None = None, |
| 502 | videos: list[list[Image.Image] | str] | None = None, |
| 503 | video_timestamps: list[list[float]] | None = None, |
| 504 | return_tensors: str | None = "pt", |
| 505 | **kwargs, |
| 506 | ) -> BatchFeature: |
| 507 | if text is None: |
| 508 | raise ValueError( |
| 509 | "`text` is required (use apply_chat_template to build it)." |
| 510 | ) |
| 511 | if isinstance(text, (list, tuple)): |
| 512 | if len(text) != 1: |
| 513 | raise ValueError( |
| 514 | "MuseGlimmerProcessor supports a single text sample per call." |
| 515 | ) |
| 516 | text = text[0] |
| 517 | |
| 518 | images = list(images or []) |
| 519 | videos = list(videos or []) |
| 520 | image_sentinel = self._sid(IMAGE_SENTINEL) |
| 521 | video_sentinel = self._sid(VIDEO_SENTINEL) |
| 522 | |
| 523 | # Preprocess media up front so we can expand sentinels in document order. |
| 524 | # Images via image_processor; videos via video_processor (each video item |
| 525 | # is a path -- decoded here -- or a pre-decoded list of PIL frames). |
| 526 | prepped_images = [self.image_processor.preprocess_image(im) for im in images] |
| 527 | prepped_videos = [ |
| 528 | self.video_processor.preprocess_one( |
| 529 | v, video_timestamps[i] if video_timestamps else None |
| 530 | ) |
| 531 | for i, v in enumerate(videos) |
| 532 | ] |
| 533 | |
| 534 | ids = self.tokenizer.encode(text, add_special_tokens=False) |
| 535 | n_img = sum(1 for t in ids if t == image_sentinel) |
| 536 | n_vid = sum(1 for t in ids if t == video_sentinel) |
| 537 | if n_img != len(prepped_images): |
| 538 | raise ValueError( |
| 539 | f"{n_img} image sentinel(s) in text but " |
| 540 | f"{len(prepped_images)} image(s) given." |
| 541 | ) |
| 542 | if n_vid != len(prepped_videos): |
| 543 | raise ValueError( |
| 544 | f"{n_vid} video sentinel(s) in text but " |
| 545 | f"{len(prepped_videos)} video(s) given." |
| 546 | ) |
| 547 | |
| 548 | out_ids: list[int] = [] |
| 549 | pixel_values: list[torch.Tensor] = [] |
| 550 | img_i = vid_i = 0 |
| 551 | for tid in ids: |
| 552 | if tid == image_sentinel: |
| 553 | tensor, n_tokens = prepped_images[img_i] |
| 554 | img_i += 1 |
| 555 | out_ids += self._image_block(n_tokens) |
| 556 | pixel_values.append(tensor) |
| 557 | elif tid == video_sentinel: |
| 558 | groups, n_groups, tokens_per_group, ts = prepped_videos[vid_i] |
| 559 | vid_i += 1 |
| 560 | out_ids += self._video_block(n_groups, tokens_per_group, ts or None) |
| 561 | pixel_values += groups |
| 562 | else: |
| 563 | out_ids.append(tid) |
| 564 | |
| 565 | data: dict = { |
| 566 | "input_ids": [out_ids], |
| 567 | "attention_mask": [[1] * len(out_ids)], |
| 568 | } |
| 569 | batch = BatchFeature(data=data, tensor_type=return_tensors) |
| 570 | # Keep pixel_values as a list of variable-size tensors (MuseGlimmerModel |
| 571 | # consumes a list); BatchFeature would fail to stack them into one tensor. |
| 572 | if pixel_values: |
| 573 | batch["pixel_values"] = pixel_values |
| 574 | return batch |
| 575 | |
| 576 | |
| 577 | __all__ = [ |
| 578 | "MuseGlimmerImageProcessor", |
| 579 | "MuseGlimmerVideoProcessor", |
| 580 | "MuseGlimmerProcessor", |
| 581 | "MUSE_GLIMMER_MM_CHAT_TEMPLATE", |
| 582 | ] |