processing_qwen3_vl.py

"""Vision-Asym (BQA prefill + BKA decode) Processor for Qwen3-VL.

The processor only reports MODALITY: it emits ``input_ids`` + pixel tensors +
``visual_token_mask`` (True on image/video tokens). The always-keep/selectable
partition (sink + recent window) is a MODEL decision made at cache-finalize time
from the model config — the processor makes no assumptions about where system
text or instructions live and adds no padding tokens.

The video/image processors still align the visual grid so h/merge and w/merge are
multiples of ``focus_size`` (= R): each frame is a whole number of R x R spatial
regions, which the model's post-prefill region fold requires.
"""

from typing import Any, Optional, Union

import torch
import numpy as np
from transformers.image_utils import ImageInput
from transformers.processing_utils import ImagesKwargs, MultiModalData, ProcessingKwargs, ProcessorMixin, Unpack, VideosKwargs
from transformers.tokenization_utils_base import PreTokenizedInput, TextInput
from transformers.utils import logging
from transformers.video_utils import VideoInput


logger = logging.get_logger(__name__)


class MMFeature(dict):
    def __init__(self, data, tensor_type: str | None = None):
        super().__init__(data)
        self.tensor_type = tensor_type
        self.convert_to_tensor()

    def convert_to_tensor(self) -> "MMFeature":
        if self.tensor_type is None:
            return self
        match self.tensor_type:
            case "pt":
                for k, v in self.items():
                    if not isinstance(v, torch.Tensor):
                        try:
                            self[k] = torch.tensor(v)
                        except Exception:
                            pass
            case "np":
                for k, v in self.items():
                    if not isinstance(v, np.ndarray):
                        try:
                            self[k] = np.array(v)
                        except Exception:
                            pass
        return self

    def to(self, target: Any) -> "MMFeature":
        for k, v in self.items():
            if isinstance(v, torch.Tensor):
                self[k] = v.to(target)
        return self


class Qwen3VLVideosProcessorKwargs(VideosKwargs, total=False):
    pass


class Qwen3VLImagesKwargs(ImagesKwargs):
    min_pixels: Optional[int]
    max_pixels: Optional[int]
    patch_size: Optional[int]
    temporal_patch_size: Optional[int]
    merge_size: Optional[int]


class Qwen3VLProcessorKwargs(ProcessingKwargs, total=False):
    images_kwargs: Qwen3VLImagesKwargs  # type: ignore
    videos_kwargs: Qwen3VLVideosProcessorKwargs  # type: ignore
    _defaults = {  # type: ignore
        "text_kwargs": {
            "padding": False,
            "return_token_type_ids": False,
            "return_mm_token_type_ids": False,
        },
        "videos_kwargs": {"return_metadata": True},
    }


class VisionAsymQwen3VLProcessor(ProcessorMixin):
    """Processor for Vision-Asym (BQA prefill + BKA decode) Qwen3-VL.

    Emits ``input_ids`` + pixel tensors + ``visual_token_mask`` (True on
    image/video tokens). Modality only — the always/selectable partition is the
    model's decision (config sink_size / recent_window), not the processor's.
    """

    attributes = ["image_processor", "tokenizer", "video_processor"]
    image_processor_class = "AutoImageProcessor"
    video_processor_class = "AutoVideoProcessor"
    tokenizer_class = ("Qwen2Tokenizer", "Qwen2TokenizerFast")

    def __init__(self, image_processor=None, tokenizer=None, video_processor=None, chat_template=None,
                 **kwargs):
        super().__init__(image_processor, tokenizer, video_processor, chat_template=chat_template)
        self.image_token = "<|image_pad|>" if not hasattr(tokenizer, "image_token") else tokenizer.image_token
        self.video_token = "<|video_pad|>" if not hasattr(tokenizer, "video_token") else tokenizer.video_token
        self.image_token_id = (
            tokenizer.image_token_id
            if getattr(tokenizer, "image_token_id", None)
            else tokenizer.convert_tokens_to_ids(self.image_token)
        )
        self.video_token_id = (
            tokenizer.video_token_id
            if getattr(tokenizer, "video_token_id", None)
            else tokenizer.convert_tokens_to_ids(self.video_token)
        )
        self.vision_start_token = (
            "<|vision_start|>" if not hasattr(tokenizer, "vision_start_token") else tokenizer.vision_start_token
        )
        self.vision_end_token = (
            "<|vision_end|>" if not hasattr(tokenizer, "vision_end_token") else tokenizer.vision_end_token
        )
        self.vision_start_token_id = tokenizer.convert_tokens_to_ids(self.vision_start_token)


    def __call__(
        self,
        images: ImageInput = None,
        text: Union[TextInput, PreTokenizedInput, list[TextInput], list[PreTokenizedInput]] = None,
        videos: VideoInput = None,
        **kwargs: Unpack[Qwen3VLProcessorKwargs],
    ) -> MMFeature:
        output_kwargs = self._merge_kwargs(
            Qwen3VLProcessorKwargs,
            tokenizer_init_kwargs=self.tokenizer.init_kwargs,
            **kwargs,
        )

        if images is not None:
            image_inputs = self.image_processor(images=images, **output_kwargs["images_kwargs"])
            image_grid_thw = image_inputs["image_grid_thw"]
        else:
            image_inputs = {}
            image_grid_thw = None

        if videos is not None:
            videos_inputs = self.video_processor(videos=videos, **output_kwargs["videos_kwargs"])
            video_grid_thw = videos_inputs["video_grid_thw"]
            video_metadata = videos_inputs.pop("video_metadata", None)
        else:
            videos_inputs = {}
            video_grid_thw = None

        if not isinstance(text, list):
            text = [text]

        text = text.copy()
        if image_grid_thw is not None:
            merge_length = self.image_processor.merge_size**2
            index = 0
            for i in range(len(text)):
                while self.image_token in text[i]:
                    num_image_tokens = image_grid_thw[index].prod() // merge_length
                    text[i] = text[i].replace(self.image_token, "<|placeholder|>" * num_image_tokens, 1)
                    index += 1
                text[i] = text[i].replace("<|placeholder|>", self.image_token)

        if video_grid_thw is not None:
            merge_length = self.video_processor.merge_size**2
            index = 0
            for i in range(len(text)):
                while self.video_token in text[i]:
                    metadata = video_metadata[index]
                    if metadata.fps is None:
                        metadata.fps = 24

                    # Calculate timestamps
                    indices = metadata.frames_indices
                    if not isinstance(indices, list):
                        indices = indices.tolist()
                    ms = self.video_processor.merge_size
                    if len(indices) % ms != 0:
                        indices.extend(indices[-1] for _ in range(ms - len(indices) % ms))
                    timestamps = [idx / metadata.fps for idx in indices]
                    timestamps = [
                        (timestamps[j] + timestamps[j + ms - 1]) / 2
                        for j in range(0, len(timestamps), ms)
                    ]

                    # Pad timestamps to match grid_t (video processor may pad frames)
                    grid_t = int(video_grid_thw[index][0])
                    while len(timestamps) < grid_t:
                        timestamps.append(timestamps[-1])

                    video_placeholder = ""
                    frame_seqlen = video_grid_thw[index][1:].prod() // merge_length
                    for frame_idx in range(grid_t):
                        curr_time = timestamps[frame_idx]
                        video_placeholder += f"<{curr_time:.1f} seconds>"
                        video_placeholder += (
                            self.vision_start_token + "<|placeholder|>" * frame_seqlen + self.vision_end_token
                        )
                    if f"{self.vision_start_token}{self.video_token}{self.vision_end_token}" in text[i]:
                        text[i] = text[i].replace(
                            f"{self.vision_start_token}{self.video_token}{self.vision_end_token}", video_placeholder, 1
                        )
                    else:
                        text[i] = text[i].replace(self.video_token, video_placeholder, 1)
                    index += 1

                text[i] = text[i].replace("<|placeholder|>", self.video_token)

        return_tensors = output_kwargs["text_kwargs"].pop("return_tensors", None)
        output_kwargs["text_kwargs"].pop("return_mm_token_type_ids", None)
        text_inputs = self.tokenizer(text, **output_kwargs["text_kwargs"])

        self._check_special_mm_tokens(text, text_inputs, modalities=["image", "video"])

        ids = np.asarray(text_inputs["input_ids"])
        visual_token_mask = ((ids == self.image_token_id) | (ids == self.video_token_id)).astype(bool)
        if visual_token_mask.shape[0] == 1:
            visual_token_mask = visual_token_mask[0]

        return MMFeature(data={
            **text_inputs,
            **image_inputs,
            **videos_inputs,
            "visual_token_mask": visual_token_mask,
        }, tensor_type=return_tensors)

    def post_process_image_text_to_text(self, generated_outputs, skip_special_tokens=True,
                                         clean_up_tokenization_spaces=False, **kwargs):
        return self.tokenizer.batch_decode(
            generated_outputs,
            skip_special_tokens=skip_special_tokens,
            clean_up_tokenization_spaces=clean_up_tokenization_spaces,
            **kwargs,
        )


__all__ = ["VisionAsymQwen3VLProcessor"]