"""Standalone decord video sampling with QSV's absolute temporal-id contract.""" import math import os from functools import partial from itertools import chain import numpy as np import torch from torch import nn from torch.nn.init import trunc_normal_ from PIL import Image from transformers.feature_extraction_utils import BatchFeature from transformers.video_processing_utils import BaseVideoProcessor from .image_processing_quicksviewer import QuicksviewerImageProcessor, SIGLIP_SIZE, read_source TIME_SCALE = 0.1 def frame_temporal_ids(frame_ts, time_scale=TIME_SCALE, duration=None): if not len(frame_ts): return [] if duration is None: duration = float(frame_ts[-1]) + time_scale n_ticks = max(int(math.ceil(duration / time_scale)) - 1, 0) return [max(0, min(int(round(float(t) / time_scale)), n_ticks)) for t in frame_ts] def num_video_cubes(num_frames, fpq=5): return max(round(num_frames / fpq) - 1, 1) + 1 def is_frame(value): return isinstance(value, Image.Image) or ( isinstance(value, (np.ndarray, torch.Tensor)) and value.ndim == 3 ) def video_list(videos): if videos is None: return [] if isinstance(videos, (np.ndarray, torch.Tensor)) and videos.ndim == 5: return list(videos) if isinstance(videos, (list, tuple)): return [videos] if videos and all(is_frame(frame) for frame in videos) else list(videos) return [videos] MAX_TEMPORAL_ID = 72000 TIMESTAMP_TOKEN_COUNT = 6 # Qwen3 ordinary vocabulary tokens, not added or special tokens. _DIGIT_TOKEN_IDS = (15, 16, 17, 18, 19, 20, 21, 22, 23, 24) _DOT_TOKEN_ID = 13 TIMESTAMP_PLACEHOLDER_IDS = (15, 15, 15, 15, 13, 15) # "0000.0" def timestamp_token_count(temporal_id: int) -> int: tick = max(0, int(temporal_id)) return max(TIMESTAMP_TOKEN_COUNT, len(str(tick // 10)) + 2) def timestamp_placeholder_ids(token_count: int = TIMESTAMP_TOKEN_COUNT) -> tuple[int, ...]: if token_count < TIMESTAMP_TOKEN_COUNT: raise ValueError(f"timestamp slot needs at least {TIMESTAMP_TOKEN_COUNT} tokens, got {token_count}") return (_DIGIT_TOKEN_IDS[0],) * (token_count - 2) + (_DOT_TOKEN_ID, _DIGIT_TOKEN_IDS[0]) def canonical_timestamp_text(temporal_id: int, *, token_count: int | None = None) -> str: tick = max(0, int(temporal_id)) required = timestamp_token_count(tick) width = required if token_count is None else int(token_count) if width < required: raise ValueError(f"timestamp id {tick} needs {required} tokens, but the slot has {width}") return f"{tick // 10:0{width - 2}d}.{tick % 10}" def timestamp_token_ids(temporal_id: int, *, token_count: int | None = None) -> tuple[int, ...]: text = canonical_timestamp_text(temporal_id, token_count=token_count) return tuple(_DOT_TOKEN_ID if c == "." else _DIGIT_TOKEN_IDS[int(c)] for c in text) def timestamp_token_tensor(temporal_ids, *, device=None, token_count: int | None = None) -> torch.Tensor: temporal_ids = list(temporal_ids) width = (max((timestamp_token_count(t) for t in temporal_ids), default=TIMESTAMP_TOKEN_COUNT) if token_count is None else int(token_count)) rows = [timestamp_token_ids(t, token_count=width) for t in temporal_ids] return torch.tensor(rows, dtype=torch.long, device=device).reshape(-1, width) def patch_timestamp_slots(input_ids: torch.Tensor, timestamp_spans, temporal_ids) -> tuple[torch.Tensor, list[tuple[int, ...]]]: """Fill preallocated slots; replacement rows may have different widths per video.""" if input_ids.ndim != 2 or input_ids.shape[0] != 1: raise ValueError( f"timestamp patching expects input_ids [1, L], got {tuple(input_ids.shape)}" ) spans = torch.as_tensor(timestamp_spans, dtype=torch.long, device=input_ids.device) spans = spans.reshape(-1, 2) temporal_ids = [int(t) for t in temporal_ids] if spans.shape[0] != len(temporal_ids): raise ValueError( f"timestamp slots {spans.shape[0]} != actual video cubes {len(temporal_ids)}" ) invalid_ids = [t for t in temporal_ids if t < 0] if invalid_ids: raise ValueError(f"invalid cube temporal ids: {invalid_ids[:8]}") patched = input_ids.clone() replacement = [] seq_len = input_ids.shape[1] previous_end = -1 for tick, (start, end) in zip(temporal_ids, spans.tolist()): if not (0 <= start < end <= seq_len): raise ValueError(f"timestamp span {(start, end)} is outside sequence length {seq_len}") width = end - start placeholder_ids = timestamp_placeholder_ids(width) if start < previous_end: raise ValueError(f"timestamp spans overlap or are unsorted near {(start, end)}") placeholder = torch.tensor(placeholder_ids, dtype=input_ids.dtype, device=input_ids.device) actual = input_ids[0, start:end] if not torch.equal(actual, placeholder): raise ValueError( f"timestamp span {(start, end)} is not the {placeholder_ids} placeholder: " f"{actual.tolist()}" ) row = timestamp_token_ids(tick, token_count=width) patched[0, start:end] = torch.tensor(row, dtype=input_ids.dtype, device=input_ids.device) replacement.append(row) previous_end = end return patched, replacement def validate_timestamp_tokenizer(tokenizer) -> None: for tick in (0, 99, 100, 9999, 10000, MAX_TEMPORAL_ID, 99999, 100000, 1234567): text = canonical_timestamp_text(tick) actual = tuple(tokenizer.encode(text, add_special_tokens=False)) expected = timestamp_token_ids(tick) if actual != expected: raise ValueError( f"tokenizer timestamp contract mismatch for {text!r}: {actual} != {expected}" ) class QuicksviewerVideoProcessor(BaseVideoProcessor): model_input_names = ["pixel_values", "frames_per_media", "frame_temporal_ids"] def __init__(self, max_frames=360, num_threads=4, size=None, **kwargs): super().__init__(**kwargs) self.max_frames = int(max_frames) self.num_threads = int(num_threads) self.size = size or {"height": SIGLIP_SIZE, "width": SIGLIP_SIZE} if self.size != {"height": SIGLIP_SIZE, "width": SIGLIP_SIZE}: raise ValueError("Quicksviewer requires 980x980 frames") self.do_normalize = False self.do_rescale = False self._image_processor = QuicksviewerImageProcessor() def load_video(self, video, max_frames=None, fps=None, timestamps=None, duration=None, video_metadata=None): """Return uint8 [F,3,980,980], rounded seconds, and duration. File inputs use exactly serve.load_media's decord resize/get_batch path. Decoded frame inputs may supply fps or timestamps; the default fps is 25. """ limit = self.max_frames if max_frames is None else int(max_frames) if limit < 1: raise ValueError("max_frames must be positive") if video_metadata is not None: get = video_metadata.get if isinstance(video_metadata, dict) else lambda k, d=None: getattr(video_metadata, k, d) fps = get("fps", fps) timestamps = get("timestamps", timestamps) duration = get("duration", duration) if isinstance(video, (str, bytes, os.PathLike)) or hasattr(video, "read"): try: import decord except ImportError as error: raise ImportError("Install decord to decode video paths; decoded frame arrays are also accepted") from error reader = decord.VideoReader(read_source(video), width=SIGLIP_SIZE, height=SIGLIP_SIZE, num_threads=self.num_threads) count = len(reader) if count <= 0: raise ValueError(f"Video has no decodable frames: {video!r}") rate = float(reader.get_avg_fps()) or 25.0 indices = np.linspace(0, count - 1, max(min(limit, count), 2)).astype(int).tolist() batch = reader.get_batch(indices) array = batch.asnumpy() if hasattr(batch, "asnumpy") else np.asarray(batch) frames = torch.from_numpy(np.ascontiguousarray(array)).permute(0, 3, 1, 2).contiguous() return frames, [round(i / rate, 2) for i in indices], count / rate if isinstance(video, Image.Image): video = [video] if not isinstance(video, (list, tuple, np.ndarray, torch.Tensor)): raise TypeError(f"Unsupported video type: {type(video).__name__}") if isinstance(video, (np.ndarray, torch.Tensor)) and video.ndim != 4: raise ValueError("Decoded video must have shape FHWC or FCHW") count = len(video) if count == 0: raise ValueError("A video must contain at least one frame") rate = float(fps or 25.0) if not math.isfinite(rate) or rate <= 0: raise ValueError("fps must be positive and finite") indices = np.linspace(0, count - 1, max(min(limit, count), 2)).astype(int).tolist() if timestamps is not None and len(timestamps) != count: raise ValueError("timestamps must contain one value per supplied frame") ts = ([round(float(timestamps[i]), 2) for i in indices] if timestamps is not None else [round(i / rate, 2) for i in indices]) if duration is None: duration = float(timestamps[-1]) + TIME_SCALE if timestamps is not None else count / rate frames = torch.stack([self._image_processor.preprocess_one(video[i]) for i in indices]) return frames, ts, float(duration) def preprocess(self, videos, return_tensors="pt", **kwargs): frames, counts, tids = [], [], [] for video in video_list(videos): pixels, ts, duration = self.load_video(video, **kwargs) frames.append(pixels) counts.append(len(pixels)) tids.extend(frame_temporal_ids(ts, duration=duration or None)) data = { "pixel_values": torch.cat(frames) if frames else torch.empty((0, 3, SIGLIP_SIZE, SIGLIP_SIZE), dtype=torch.uint8), "frames_per_media": torch.tensor(counts, dtype=torch.long), "frame_temporal_ids": torch.tensor(tids, dtype=torch.long), } if return_tensors not in (None, "pt", "np"): raise ValueError("return_tensors must be 'pt', 'np', or None") if return_tensors == "np": data = {key: value.numpy() for key, value in data.items()} return BatchFeature(data) def to_dict(self): data = super().to_dict() data.pop("_image_processor", None) return data QuicksviewerVideoProcessor.register_for_auto_class("AutoVideoProcessor") def get_2d_sincos_pos_embed(embed_dim, image_size): """image_size: int or (h, w). Returns [h, w, embed_dim]. Port of the official function, including the `meshgrid(grid_w, grid_h)` ordering ("here w goes first") and the resulting grid[0]->h / grid[1]->w channel split. """ if isinstance(image_size, int): grid_h_size, grid_w_size = image_size, image_size else: grid_h_size, grid_w_size = image_size[0], image_size[1] grid_h = np.arange(grid_h_size, dtype=np.float32) grid_w = np.arange(grid_w_size, dtype=np.float32) grid = np.meshgrid(grid_w, grid_h) # here w goes first grid = np.stack(grid, axis=0) return get_2d_sincos_pos_embed_from_grid(embed_dim, grid) def get_2d_sincos_pos_embed_from_grid(embed_dim, grid): assert embed_dim % 2 == 0 # use half of dimensions to encode grid_h emb_h = get_1d_sincos_pos_embed_from_grid_new(embed_dim // 2, grid[0]) # (H, W, D/2) emb_w = get_1d_sincos_pos_embed_from_grid_new(embed_dim // 2, grid[1]) # (H, W, D/2) emb = np.concatenate([emb_h, emb_w], axis=-1) # (H, W, D) return emb def get_1d_sincos_pos_embed_from_grid_new(embed_dim, pos): """pos: (H, W) -> out: (H, W, D)""" assert embed_dim % 2 == 0 omega = np.arange(embed_dim // 2, dtype=np.float32) omega /= embed_dim / 2. omega = 1. / 10000 ** omega # (D/2,) out = np.einsum('hw,d->hwd', pos, omega) # (H, W, D/2), outer product emb_sin = np.sin(out) # (H, W, D/2) emb_cos = np.cos(out) # (H, W, D/2) emb = np.concatenate([emb_sin, emb_cos], axis=-1) # (H, W, D) return emb def get_1d_sincos_pos_embed_from_temporal_size(embed_dim, pos): """pos: (M,) -> out: (M, D). Full-width temporal embedding.""" assert embed_dim % 2 == 0 omega = np.arange(embed_dim // 2, dtype=np.float32) omega /= embed_dim / 2. omega = 1. / 10000 ** omega # (D/2,) pos = pos.reshape(-1) # (M,) out = np.einsum('m,d->md', pos, omega) # (M, D/2), outer product emb_sin = np.sin(out) # (M, D/2) emb_cos = np.cos(out) # (M, D/2) emb = np.concatenate([emb_sin, emb_cos], axis=1) # (M, D) return emb class Resampler(nn.Module): """Perceiver resampler with one cross-attention layer, learnable queries, 2D spatial sincos positions and an additive absolute temporal embedding. Output shape: (num_groups, num_queries, embed_dim). """ def __init__(self, num_queries=64, embed_dim=None, kv_dim=None, num_heads=None, norm_layer=partial(nn.LayerNorm, eps=1e-6), max_size=(70, 70), max_temporal_size=72000, adaptive=True, batch_forward=False, batch_infer=None, **kwargs): super().__init__() assert embed_dim is not None, "resampler embed_dim (LLM hidden size) is required" self.num_queries = num_queries self.kv_dim = kv_dim self.embed_dim = embed_dim self.adaptive = adaptive # max_size is a 2D (h, w) spatial extent now; accept a legacy 3-tuple # (t, h, w) from older configs and drop the temporal component, which is # covered by max_temporal_size instead. if len(max_size) == 3: max_size = tuple(max_size[1:]) self.max_size = tuple(max_size) self.max_temporal_size = max_temporal_size # `batch_infer` was the name used by the upstream inference-only path. The # operation is fully differentiable, so expose the behavior as # `batch_forward`; keep the old keyword as a compatibility alias. if batch_infer is not None: batch_forward = batch_infer self.batch_forward = bool(batch_forward) self.num_heads = (embed_dim // 128) if not num_heads else num_heads self.query = nn.Parameter(torch.zeros(self.num_queries, self.embed_dim)) trunc_normal_(self.query, std=.02) if self.kv_dim is not None and self.kv_dim != self.embed_dim: self.kv_proj = nn.Linear(self.kv_dim, self.embed_dim, bias=False) else: self.kv_proj = nn.Identity() self.attn = nn.MultiheadAttention(self.embed_dim, self.num_heads) self.ln_q = norm_layer(self.embed_dim) self.ln_kv = norm_layer(self.embed_dim) self.ln_post = norm_layer(self.embed_dim) self.proj = nn.Parameter((self.embed_dim ** -0.5) * torch.randn(self.embed_dim, self.embed_dim)) self._set_2d_pos_cache(self.max_size) self._set_temporal_pos_cache(self.max_temporal_size) self.apply(self._init_weights) def _set_2d_pos_cache(self, max_size, device='cpu'): pos_embed = torch.from_numpy( get_2d_sincos_pos_embed(self.embed_dim, max_size)).float().to(device) self.register_buffer("pos_embed", pos_embed, persistent=False) def _adjust_pos_cache(self, tgt_sizes, device): max_h = int(torch.max(tgt_sizes[:, 0])) max_w = int(torch.max(tgt_sizes[:, 1])) if max_h > self.max_size[0] or max_w > self.max_size[1]: self.max_size = (max(max_h, self.max_size[0]), max(max_w, self.max_size[1])) self._set_2d_pos_cache(self.max_size, device) def _set_temporal_pos_cache(self, max_temporal_size, device='cpu'): temporal_size = np.arange(max_temporal_size, dtype=np.float32) pos_embed = torch.from_numpy( get_1d_sincos_pos_embed_from_temporal_size(self.embed_dim, temporal_size) ).float().to(device) self.register_buffer("temporal_pos_embed", pos_embed, persistent=False) def _adjust_temporal_pos_cache(self, max_temporal_size, device): if max_temporal_size > self.max_temporal_size: self.max_temporal_size = max_temporal_size self._set_temporal_pos_cache(self.max_temporal_size, device) def _temporal_positions(self, temporal_ids, device, dtype): cache_size = self.temporal_pos_embed.shape[0] extra_ids = sorted({t for t in temporal_ids if t >= cache_size}) extra_positions = {} if extra_ids: # Absolute time is unbounded; cache only requested positions rather # than allocating a dense table proportional to video duration. values = torch.from_numpy(get_1d_sincos_pos_embed_from_temporal_size( self.embed_dim, np.asarray(extra_ids, dtype=np.float32))) values = values.to(device=device, dtype=self.temporal_pos_embed.dtype).to(dtype) extra_positions = dict(zip(extra_ids, values.unbind(0))) positions = [] for tick in temporal_ids: if tick == -1: positions.append(torch.zeros(self.embed_dim, dtype=dtype, device=device)) elif tick < -1: raise ValueError(f"invalid temporal id {tick}") elif tick < cache_size: positions.append(self.temporal_pos_embed[tick].to(dtype)) else: positions.append(extra_positions[tick]) return positions def _init_weights(self, m): if isinstance(m, nn.Linear): trunc_normal_(m.weight, std=.02) if isinstance(m, nn.Linear) and m.bias is not None: nn.init.constant_(m.bias, 0) elif isinstance(m, nn.LayerNorm): nn.init.constant_(m.bias, 0) nn.init.constant_(m.weight, 1.0) def forward(self, x, tgt_sizes=None, temporal_ids=None): """ Args: x: [n_frames, patches, kv_dim] per-frame ViT features. tgt_sizes: [n_frames, 2] (h, w) patch grid per frame. temporal_ids: list of per-group absolute temporal id lists, e.g. [[-1], [-1], [120, 123, 126]] -- two images then one 3-frame cube. sum(len(g)) must equal n_frames. Each group becomes ONE output row of num_queries tokens. Returns: [n_groups, num_queries, embed_dim] (n_groups == n_frames when temporal_ids is None). """ assert tgt_sizes is not None, "tgt_sizes is required" assert x.shape[0] == tgt_sizes.shape[0], f"{x.shape[0]} != {tgt_sizes.shape[0]}" bs = x.shape[0] device = x.device dtype = x.dtype patch_len = tgt_sizes[:, 0] * tgt_sizes[:, 1] self._adjust_pos_cache(tgt_sizes, device=device) temporal_pos_emb = False temporal_ids_flatten = None if temporal_ids is not None: # example: [[-1], [-1], [2, 6, 9]] temporal_ids_flatten = list(chain.from_iterable(temporal_ids)) assert len(temporal_ids_flatten) == bs, ( f"temporal_ids covers {len(temporal_ids_flatten)} frames but got {bs}") max_temporal_size = max(temporal_ids_flatten) + 1 if max_temporal_size > -1: temporal_pos_emb = True max_patch_len = int(torch.max(patch_len)) key_padding_mask = torch.zeros((bs, max_patch_len), dtype=torch.bool, device=device) x = self.kv_proj(x) # B * L * D x = self.ln_kv(x).permute(1, 0, 2) # L * B * D q = self.ln_q(self.query) # Q * D # Uniform-grid fast path: qsv interpolates every frame to the same # tgt_size, so the per-frame position list is B identical slabs. Skip building # it (and the padded [L,B,D] stack) and broadcast one [L,1,D] slab instead -- # the padded tensor would be a full-size copy (138k x 4096 = 1.1GB bf16 at # 240 frames) holding identical values. uniform_grid = bool((tgt_sizes == tgt_sizes[0]).all()) pos_embed_2d = [] pos_embed_temporal = (self._temporal_positions(temporal_ids_flatten, device, dtype) if temporal_pos_emb else []) for i in range(bs): if not uniform_grid: tgt_h, tgt_w = tgt_sizes[i] pos_embed_2d.append( self.pos_embed[:tgt_h, :tgt_w, :].reshape((tgt_h * tgt_w, -1)).to(dtype)) # patches * D key_padding_mask[i, patch_len[i]:] = True if uniform_grid: tgt_h, tgt_w = tgt_sizes[0] pos_slab = (self.pos_embed[:tgt_h, :tgt_w, :] .reshape(tgt_h * tgt_w, 1, -1).to(dtype)) # L * 1 * D k = x + pos_slab else: pos_embed_2d = torch.nn.utils.rnn.pad_sequence( pos_embed_2d, batch_first=True, padding_value=0.0).permute(1, 0, 2) # L * B * D k = x + pos_embed_2d v = x if self.batch_forward: out = self.batch_attn_forward(q, k, v, pos_embed_temporal, temporal_ids, key_padding_mask) else: # save gpu memory out = self.foreach_attn_forward(q, k, v, pos_embed_temporal, temporal_ids, key_padding_mask) # out: Q * B * D x = out.permute(1, 0, 2) # B * Q * D x = self.ln_post(x) x = x @ self.proj return x def _repeat(self, query, N: int): return query.unsqueeze(1).repeat(1, N, 1) def batch_attn_forward(self, q, k, v, pos_embed_temporal, temporal_ids, key_padding_mask): bs = k.shape[1] if pos_embed_temporal: # temporal embedding is ADDITIVE over the full channel width, then the # frames of one group are folded into a single long KV sequence. k = k + torch.stack(pos_embed_temporal, dim=0) bs = len(temporal_ids) merge_k = [] merge_v = [] merge_key_padding_mask = [] start = 0 for tp in temporal_ids: end = start + len(tp) # L * (end-start) * D -> (end-start) * L * D -> L*(end-start) * D merge_k.append(k[:, start: end, :].permute(1, 0, 2).reshape(-1, self.embed_dim)) merge_v.append(v[:, start: end, :].permute(1, 0, 2).reshape(-1, self.embed_dim)) merge_key_padding_mask.append(key_padding_mask[start: end, :].reshape(-1, 1)) start = end k = torch.nn.utils.rnn.pad_sequence(merge_k, batch_first=True, padding_value=0.0).permute(1, 0, 2) v = torch.nn.utils.rnn.pad_sequence(merge_v, batch_first=True, padding_value=0.0).permute(1, 0, 2) key_padding_mask = torch.nn.utils.rnn.pad_sequence( merge_key_padding_mask, batch_first=True, padding_value=True).squeeze(-1) out = self.attn( self._repeat(q, bs), # Q * B * D k, # L * B * D v, key_padding_mask=key_padding_mask, need_weights=False)[0] return out def foreach_attn_forward(self, q, k, v, pos_embed_temporal, temporal_ids, key_padding_mask): bs = k.shape[1] if pos_embed_temporal: # NOTE: the temporal embedding is added PER GROUP inside the loop rather # than to the whole `k` at once. Adding it globally would materialize a # second full-size copy of K (a 240-frame video is 138k tokens x 4096 = # 1.1GB in bf16), which is pure overhead since each group is consumed # independently. The arithmetic is identical. tpe = torch.stack(pos_embed_temporal, dim=0) # [B, D] out_list = [] start = 0 for tp in temporal_ids: end = start + len(tp) # one group at a time, no padding curr_k = ( k[:, start:end, :] + tpe[start:end] ).permute(1, 0, 2).reshape(-1, self.embed_dim) curr_v = v[:, start:end, :].permute(1, 0, 2).reshape(-1, self.embed_dim) curr_key_padding_mask = key_padding_mask[start:end, :].reshape(-1) curr_out = self.attn( q, curr_k, curr_v, key_padding_mask=curr_key_padding_mask, need_weights=False, )[0] out_list.append(curr_out) start = end out = torch.stack(out_list, dim=1) else: out = self.attn( self._repeat(q, bs), # Q * B * D k, # L * B * D v, key_padding_mask=key_padding_mask, need_weights=False)[0] return out @property def hidden_size(self): return self.embed_dim @property def batch_infer(self): """Compatibility alias for old configs and external callers.""" return self.batch_forward