Quicksviewer-8B / video_processing_quicksviewer.py
Add files using upload-large-folder tool
0efe9d5 verified 15 days ago
25.4 kB
"""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