Transformers
File size: 9,971 Bytes
838c455
 
23b5fcc
 
 
 
 
 
 
 
 
838c455
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
23b5fcc
 
 
838c455
 
 
 
 
 
 
 
23b5fcc
838c455
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
23b5fcc
 
 
 
838c455
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
"""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"]