# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. from typing import Optional, Union, List import numpy as np import torch from transformers.feature_extraction_utils import BatchFeature 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.video_utils import VideoInput class ZDTaichu5_0_ImagesKwargs(ImagesKwargs): min_pixels: Optional[int] max_pixels: Optional[int] patch_size: Optional[int] temporal_patch_size: Optional[int] merge_size: Optional[int] class ZDTaichu5_0_ProcessorKwargs(ProcessingKwargs, total=False): images_kwargs: ZDTaichu5_0_ImagesKwargs videos_kwargs: VideosKwargs _defaults = { "text_kwargs": { "padding": False, }, } class ZDTaichu5_0_Processor(ProcessorMixin): r""" Constructs a ZDTaichu-5.0 processor which wraps an image processor and a tokenizer into a single processor. [`ZDTaichu5_0_Processor`] offers all the functionalities of the image processor and tokenizer. See the [`~ZDTaichu5_0_Processor.__call__`] and [`~ZDTaichu5_0_Processor.decode`] for more information. Args: image_processor ([`AutoImageProcessor`], *optional*): The image processor is a required input. tokenizer ([`AutoTokenizer`], *optional*): The tokenizer is a required input. chat_template (`str`, *optional*): A Jinja template which will be used to convert lists of messages in a chat into a tokenizable string. """ attributes = ["image_processor", "tokenizer"] image_processor_class = "AutoImageProcessor" video_processor_class = "AutoVideoProcessor" tokenizer_class = ("AutoTokenizer") def __init__(self, image_processor=None, tokenizer=None, chat_template=None, **kwargs): # Defaults to Qwen3's built-in vision tokens; overridden by tokenizer_config.json attributes. self.image_token = getattr(tokenizer, "image_token", "<|image_pad|>") self.video_token = getattr(tokenizer, "video_token", "<|video_pad|>") self.image_start_token = getattr(tokenizer, "image_start_token", "<|vision_start|>") self.image_end_token = getattr(tokenizer, "image_end_token", "<|vision_end|>") 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) ) super().__init__(image_processor, tokenizer, chat_template=chat_template) def __call__( self, images: ImageInput = None, text: Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]] = None, videos: VideoInput = None, **kwargs: Unpack[ZDTaichu5_0_ProcessorKwargs], ) -> BatchFeature: """ Main method to prepare multimodal inputs (text, images, videos) for the model. This method processes text by replacing image/video tokens with appropriate placeholder sequences, processes images and videos through the image processor, and tokenizes the final text. Video-as-multi-image convention ─────────────────────────────── Videos are NOT processed as a separate temporal stream. Each frame is passed through the *image* pipeline with `max_num_tiles=1`, producing one 512x512 tile per frame, and the frame tile tensors are then APPENDED to the image stream (`pixel_values` / `num_patches` / `image_grid_thw`). The downstream model therefore sees a single uniform image batch with no separate video path. In the rendered prompt every frame is wrapped in `<|vision_start|> ... <|image_pad|> ... <|vision_end|>` — *image* tokens, not video tokens — and prefaced by a per-frame "Frame N sampled at T.TT seconds:" header. After tokenisation `mm_token_type_ids` therefore has no type=2 entries. The method performs the following key operations: 1. Processes images using the image processor to get pixel values and patch counts 2. Processes videos as multi-image (max_num_tiles=1) and appends frame data into the same pixel_values / num_patches / image_grid_thw containers 3. Replaces `<|image_pad|>` tokens in text with `<|vision_start|>` + image tokens + `<|vision_end|>` sequences 4. Replaces `<|video_pad|>` tokens in text with frame-by-frame descriptions including timestamps (if metadata provided) 5. Tokenizes the processed text and combines all outputs Args: images (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `List[PIL.Image.Image]`, `List[np.ndarray]`, `List[torch.Tensor]`, *optional*): The image or batch of images to be prepared. Each image can be a PIL image, NumPy array or PyTorch tensor. Both channels-first and channels-last formats are supported. text (`str`, `List[str]`, *optional*): The sequence or batch of sequences to be encoded. Each sequence should be a string. The text can contain special tokens `<|image_pad|>` and `<|video_pad|>` that will be replaced with appropriate token sequences. videos (`np.ndarray`, `torch.Tensor`, `List[np.ndarray]`, `List[torch.Tensor]`, *optional*): The video or batch of videos to be prepared. Each video should be a 4D NumPy array or PyTorch tensor with shape (num_frames, channels, height, width). Both channels-first and channels-last formats are supported. Note: Currently only supports batch size of 1 for videos. images_kwargs (`Dict`, *optional*): Additional keyword arguments for image processing, including: - `min_pixels` (`int`, *optional*): Minimum number of pixels for image processing - `max_pixels` (`int`, *optional*): Maximum number of pixels for image processing - `patch_size` (`int`, *optional*): Size of patches for image processing - `temporal_patch_size` (`int`, *optional*): Size of temporal patches - `merge_size` (`int`, *optional*): Size for merging patches videos_kwargs (`Dict`, *optional*): Additional keyword arguments for video processing, including: - `video_metadata` (`VideoMetadata`, *optional*): Metadata containing fps information for timestamp calculation text_kwargs (`Dict`, *optional*): Additional keyword arguments for text tokenization, including: - `return_tensors` (`str` or [`~utils.TensorType`], *optional*): Framework for returned tensors ('tf', 'pt', 'np', 'jax') - `padding` (`bool`, *optional*): Whether to pad sequences (defaults to False) Returns: [`BatchFeature`]: A [`BatchFeature`] with the following fields: - **input_ids** -- List of token ids to be fed to a model. Returned when `text` is not `None`. - **attention_mask** -- List of indices specifying which tokens should be attended to by the model. - **pixel_values** -- Concatenated tile pixel values from BOTH real images and video frames, in the order they appear in `input_ids` (real images first, then video frames). Returned when `images` is not `None` or `videos` is not `None`. - **num_patches** -- List of tile counts, one entry per real image followed by one entry per video frame (frame entries are always 1 because max_num_tiles=1). - **image_grid_thw** -- LongTensor[N_images + N_frames, 3] with [1, tile_rows, tile_cols] per real image and [1, 1, 1] per video frame. - **mm_token_type_ids** -- Per-token modality classification (0=text, 1=image incl. frames). Raises: AssertionError: If videos are provided with batch size > 1 (not currently supported). Note: - Image tokens `<|image_pad|>` in text are replaced with `<|vision_start|>` + repeated image tokens + `<|vision_end|>` - Video tokens `<|video_pad|>` in text are replaced with frame-by-frame descriptions, each frame using `<|image_pad|>` slots - When video metadata with fps is provided, frame descriptions include timestamps - Videos are processed with max_num_tiles=1 regardless of the images setting """ output_kwargs = self._merge_kwargs( ZDTaichu5_0_ProcessorKwargs, tokenizer_init_kwargs=self.tokenizer.init_kwargs, **kwargs, ) # Initialise as independent dicts so later `**image_inputs` merging # is well-defined whether or not images / videos are provided. image_inputs: dict = {} image_grid_thw = None # Frame counts default to empty so the video-text-expansion loop is a # no-op when `videos` is None. video_num_patches: list = [] if images is not None: image_inputs = self.image_processor(images=images, **output_kwargs["images_kwargs"]) image_num_patches = image_inputs["num_patches"] # image_grid_thw: list of [T=1, tile_rows, tile_cols] per image image_grid_thw = image_inputs.pop("image_grid_thw") image_pixel_values = image_inputs["pixel_values"] else: image_num_patches = [] if videos is not None: # ── Multi-image treatment of video ───────────────────────────────── # Every video frame is processed by the *image* pipeline with # max_num_tiles=1 so that one frame = one 512x512 tile = num_image_token # (e.g. 256) tokens. Frame tile tensors are then APPENDED to the # real-image stream: # # pixel_values : torch.cat([images, frames]) (total_tiles, C, H, W) # num_patches : image_num_patches + [1] * N_frames List[int] # image_grid_thw : torch.cat([image_grids, frame_grids], dim=0) # # Order matters: text expansion below replaces image tokens # before video tokens, so frame slots come *after* real-image # slots in input_ids — these tensors must follow the same order. # # In the rendered prompt every frame is wrapped in # <|vision_start|> ... <|image_pad|> x num_image_token ... <|vision_end|> # (image tokens, NOT video tokens) and prefaced by a per-frame # "Frame N sampled at T.TT seconds:" header. After tokenisation # mm_token_type_ids therefore has *no* type=2 entries — every # visual slot is type=1. The downstream model sees a single # uniform image stream and does not need a separate video path. orig_tiles = self.image_processor.max_num_tiles self.image_processor.max_num_tiles = 1 try: frame_inputs = self.image_processor( images=videos, **output_kwargs["images_kwargs"] ) finally: self.image_processor.max_num_tiles = orig_tiles frame_pixel_values = frame_inputs["pixel_values"] # (N_frames, C, H, W) frame_num_patches = list(frame_inputs["num_patches"]) frame_grid_thw = frame_inputs["image_grid_thw"] # (N_frames, 3) list/tensor video_num_patches = frame_num_patches # for text expansion below # Normalise grid containers to LongTensor so torch.cat works # whether the image processor returned lists or tensors. def _to_long_tensor(x): return x if isinstance(x, torch.Tensor) else torch.tensor(x, dtype=torch.long) if image_inputs: # Real images + video frames — concat along batch dim. image_inputs["pixel_values"] = torch.cat( [image_inputs["pixel_values"], frame_pixel_values], dim=0 ) image_inputs["num_patches"] = ( list(image_inputs["num_patches"]) + frame_num_patches ) image_grid_thw = torch.cat( [_to_long_tensor(image_grid_thw), _to_long_tensor(frame_grid_thw)], dim=0, ) image_num_patches = image_inputs["num_patches"] else: # Video-only — frames become the entire image stream. image_inputs = { "pixel_values": frame_pixel_values, "num_patches": frame_num_patches, } image_grid_thw = _to_long_tensor(frame_grid_thw) image_num_patches = frame_num_patches if not isinstance(text, list): text = [text] final_image_pixel_values = [] final_image_num_patches = [] final_image_grid_thw = [] text = text.copy() # below lines change text in-place if images is not None: index = 0 wrapped_token = self.image_start_token + self.image_token + self.image_end_token for i in range(len(text)): while self.image_token in text[i]: expansion = ( self.image_start_token + "<|placeholder|>" * image_num_patches[index] * self.image_processor.num_image_token + self.image_end_token ) # If the chat template already wrapped it, replace the whole # span — avoids double wrapping. # Otherwise fall back to replacing the bare token. search = wrapped_token if wrapped_token in text[i] else self.image_token text[i] = text[i].replace(search, expansion, 1) index += 1 #final_image_pixel_values.append(image_pixel_values[index]) #final_image_num_patches.append(i) text[i] = text[i].replace("<|placeholder|>", self.image_token) if videos is not None: assert len(text) == 1, "Video is not supported for batch size > 1" video_metadata = output_kwargs.get("videos_kwargs", {}).get("video_metadata", None) i = 0 wrapped_token = self.image_start_token + self.video_token + self.image_end_token if self.video_token in text[i]: each_frame = ( self.image_start_token + "<|placeholder|>" * self.image_processor.num_image_token + self.image_end_token ) video_prompt = "This is a video:\n" # One iteration per frame. video_num_patches has length N_frames # (always 1 per frame because max_num_tiles=1 was forced above), # so its length is the authoritative frame count even when # `images` is None and `image_num_patches` is unset. n_frames = len(video_num_patches) for j in range(n_frames): if video_metadata is not None and video_metadata.fps is not None: timestamp = j / video_metadata.fps video_prompt += f"Frame {j+1} sampled at {timestamp:.2f} seconds: {each_frame}\n" else: # Fallback to original format without timestamps video_prompt += f"Frame {j+1}: {each_frame}\n" # Strip the chat-template-applied <|vision_start|>...<|vision_end|> # wrapping if present; otherwise replace the bare <|video_pad|> # token. The fallback is video_token (NOT image_token), since by # this point image expansion has already consumed every # <|image_pad|> in the prompt. search = wrapped_token if wrapped_token in text[i] else self.video_token text[i] = text[i].replace(search, video_prompt, 1) text[i] = text[i].replace("<|placeholder|>", self.image_token) return_tensors = output_kwargs["text_kwargs"].pop("return_tensors", None) text_inputs = self.tokenizer(text, **output_kwargs["text_kwargs"]) # ── Build mm_token_type_ids from tokenized input_ids ───────────── # 0 = text, 1 = image. type=2 (video) is unreachable under the # multi-image-as-video convention because every frame gets expanded # to <|image_pad|> tokens — but we keep the video_token branch as a # defensive fallback for any tokenizer-injected video_pad token. input_ids = text_inputs["input_ids"] if isinstance(input_ids, list): mm_token_type_ids = [] for ids in input_ids: tt = [0] * len(ids) for j, tok_id in enumerate(ids): if tok_id == self.image_token_id: tt[j] = 1 elif tok_id == self.video_token_id: tt[j] = 2 mm_token_type_ids.append(tt) else: # Already a tensor (when return_tensors is set before tokenizer call) mm_token_type_ids = torch.zeros_like(input_ids) mm_token_type_ids[input_ids == self.image_token_id] = 1 mm_token_type_ids[input_ids == self.video_token_id] = 2 # ── Assemble output ────────────────────────────────────────────── # Note: video frames have already been merged into image_inputs above, # so there are no separate `pixel_values_videos` / `video_grid_thw` # outputs. Downstream code consumes a single image stream. data = {**text_inputs, **image_inputs} data["mm_token_type_ids"] = mm_token_type_ids if image_grid_thw is not None: data["image_grid_thw"] = image_grid_thw return BatchFeature(data=data, tensor_type=return_tensors) def _get_num_multimodal_tokens(self, image_sizes=None, video_sizes=None, **kwargs): """ Computes the number of placeholder tokens needed for multimodal inputs with the given sizes. Args: image_sizes (`list[list[int]]`, *optional*): The input sizes formatted as (height, width) per each image. video_sizes (`list[list[int]]`, *optional*): The input sizes formatted as (num_frames, height, width) per each video. Returns: `MultiModalData`: A `MultiModalData` object holding number of tokens per each of the provided input modalities, along with other useful data. """ vision_data = {} if image_sizes is not None: images_kwargs = ZDTaichu5_0_ProcessorKwargs._defaults.get("images_kwargs", {}) images_kwargs.update(kwargs) merge_size = images_kwargs.get("merge_size", None) or self.image_processor.merge_size num_image_patches = [ self.image_processor.get_number_of_image_patches(*image_size, images_kwargs) for image_size in image_sizes ] num_image_tokens = [(num_patches // merge_size**2) for num_patches in num_image_patches] vision_data.update({"num_image_tokens": num_image_tokens, "num_image_patches": num_image_patches}) return MultiModalData(**vision_data) def batch_decode(self, *args, **kwargs): """ This method forwards all its arguments to the tokenizer's [`~PreTrainedTokenizer.batch_decode`]. Please refer to the docstring of this method for more information. """ return self.tokenizer.batch_decode(*args, **kwargs) def decode(self, *args, **kwargs): """ This method forwards all its arguments to the tokenizer's [`~PreTrainedTokenizer.decode`]. Please refer to the docstring of this method for more information. """ return self.tokenizer.decode(*args, **kwargs) def post_process_image_text_to_text( self, generated_outputs, skip_special_tokens=True, clean_up_tokenization_spaces=False, **kwargs ): """ Post-process the output of the model to decode the text. Args: generated_outputs (`torch.Tensor` or `np.ndarray`): The output of the model `generate` function. The output is expected to be a tensor of shape `(batch_size, sequence_length)` or `(sequence_length,)`. skip_special_tokens (`bool`, *optional*, defaults to `True`): Whether or not to remove special tokens in the output. Argument passed to the tokenizer's `batch_decode` method. clean_up_tokenization_spaces (`bool`, *optional*, defaults to `False`): Whether or not to clean up the tokenization spaces. Argument passed to the tokenizer's `batch_decode` method. **kwargs: Additional arguments to be passed to the tokenizer's `batch_decode method`. Returns: `list[str]`: The decoded text. """ return self.tokenizer.batch_decode( generated_outputs, skip_special_tokens=skip_special_tokens, clean_up_tokenization_spaces=clean_up_tokenization_spaces, **kwargs, ) @property def model_input_names(self): tokenizer_input_names = self.tokenizer.model_input_names image_processor_input_names = self.image_processor.model_input_names names_from_processor = list(dict.fromkeys(tokenizer_input_names + image_processor_input_names)) # Note: video_grid_thw is NOT emitted under multi-image-as-video — # frame data lives in image_grid_thw alongside real images. return names_from_processor + ["mm_token_type_ids"] def from_messages( self, messages: list, return_tensors: str = "pt", add_vision_id: bool = True, **kwargs, ) -> BatchFeature: """ Prepare model inputs directly from Qwen-style structured messages. This is the high-level entry point that handles the full pipeline: structured messages → vision loading → chat template → tokenization. Supports messages with typed content lists:: messages = [ {"role": "user", "content": [ {"type": "image", "image": "photo.jpg"}, {"type": "text", "text": "What's in this image?"}, ]}, ] Video inputs (file path, URL, or list of frame paths):: messages = [ {"role": "user", "content": [ {"type": "video", "video": "clip.mp4", "fps": 2.0}, {"type": "text", "text": "Describe this video."}, ]}, ] Multi-image with automatic labelling:: messages = [ {"role": "user", "content": [ {"type": "image", "image": "a.jpg"}, {"type": "image", "image": "b.jpg"}, {"type": "text", "text": "Compare them."}, ]}, ] # With add_vision_id=True (default), the prompt includes: # Picture 1: <|vision_start|><|image_pad|><|vision_end|> # Picture 2: <|vision_start|><|image_pad|><|vision_end|> # Compare them. Args: messages: List of message dicts with structured ``content``. return_tensors: Framework for returned tensors (default ``"pt"``). add_vision_id: If ``True`` (default), the chat template prepends ``Picture N:`` / ``Video N:`` labels before each vision token. Set to ``False`` to omit labels. **kwargs: Forwarded to ``self.__call__``. Returns: ``BatchFeature`` ready for ``model.generate(**inputs)``. """ from .vision_utils import process_vision_info # 1. Load images and videos from the structured messages image_inputs, video_inputs, video_kwargs = process_vision_info(messages) # 2. Apply chat template — pass structured messages directly so the # template can iterate typed content dicts, count vision elements, # and emit "Picture N:" / "Video N:" labels when add_vision_id=True. prompt = self.tokenizer.apply_chat_template( messages, tokenize=False, add_generation_prompt=True, add_vision_id=add_vision_id, ) # 3. Prepare video inputs and metadata for timestamps videos_kwargs = {} flat_videos = None if video_inputs is not None: if len(video_inputs) > 1: raise ValueError( "Multiple videos in a single message batch are not yet " "supported. Please use one video per call." ) flat_videos = video_inputs[0] # List[Image.Image] if video_kwargs.get("metadata_list"): meta = video_kwargs["metadata_list"][0] fps = meta.get("sample_fps") or meta.get("fps") if fps: from transformers.video_utils import VideoMetadata videos_kwargs["video_metadata"] = VideoMetadata( fps=fps, total_num_frames=len(flat_videos), ) return self( images=image_inputs, text=prompt, videos=flat_videos, return_tensors=return_tensors, videos_kwargs=videos_kwargs, **kwargs, ) @property def ctx_image_token_id(self) -> int: return self.image_token_id __all__ = ["ZDTaichu5_0_Processor"]