processing_prts_qwen3_vl.py

# Copyright 2025 TeleAI Rhodes Team. 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.

"""Processor for PRTS built on Qwen3-VL (hub / trust_remote_code; no prts package required)."""

from __future__ import annotations

import logging
from typing import Optional, Union

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.utils.logging import get_logger
from transformers.video_utils import VideoInput

ACTION_START_TOKEN = "<|action_start|>"
ACTION_PLACEHOLDER_TOKEN = "<|action_pad|>"
ACTION_END_TOKEN = "<|action_end|>"
CRL_GOAL_REPR_TOKEN = "<|goal_repr|>"
CRL_OBS_REPR_TOKEN = "<|obs_repr|>"
VISION_START_TOKEN = "<|vision_start|>"         # beginning of vision input
IMAGE_PLACEHOLDER_TOKEN = "<|image_pad|>"       # image placeholder
VIDEO_PLACEHOLDER_TOKEN = "<|video_pad|>"       # video placeholder

logger = get_logger(__name__)
if not logger.handlers:
    handler = logging.StreamHandler()
    handler.setLevel(logging.INFO)
    handler.setFormatter(logging.Formatter("%(levelname)s:%(name)s:%(message)s"))
    logger.addHandler(handler)


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
    videos_kwargs: Qwen3VLVideosProcessorKwargs
    _defaults = {
        "text_kwargs": {
            "padding": False,
            "return_token_type_ids": False,
            "return_mm_token_type_ids": False,
        },
        "videos_kwargs": {"return_metadata": True},
    }


class PRTS_Qwen3VLProcessor(ProcessorMixin):
    r"""
    Constructs a PRTS processor which wraps a Qwen3-VL image processor and a Qwen2 tokenizer into a single processor.

    This processor is built independently (not inheriting from Qwen3VLProcessor) to avoid tight coupling,
    while maintaining compatibility with Qwen3-VL's timestamp-based video processing approach.

    [`PRTS_Qwen3VLProcessor`] offers all the functionalities needed for PRTS model with:
    - Action token handling (discrete and continuous)
    - State token handling for proprioceptive inputs
    - Expert trigger tokens for flow matching action prediction
    - Qwen3-VL compatible image/video processing with timestamp-based video handling

    Args:
        image_processor ([`Qwen2VLImageProcessor`], *optional*):
            The image processor is a required input.
        tokenizer ([`Qwen2TokenizerFast`], *optional*):
            The tokenizer is a required input.
        video_processor ([`Qwen3VLVideoProcessor`], *optional*):
            The video processor 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", "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):
        # Initialize base ProcessorMixin
        super().__init__(image_processor, tokenizer, video_processor, chat_template=chat_template)
        
        # Get image/video tokens from tokenizer
        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)
        )

        # Qwen3-VL vision tokens
        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.vision_start_token_id
            if getattr(tokenizer, "vision_start_token_id", None)
            else tokenizer.convert_tokens_to_ids(self.vision_start_token)
        )
        self.vision_end_token_id = (
            tokenizer.vision_end_token_id
            if getattr(tokenizer, "vision_end_token_id", None)
            else tokenizer.convert_tokens_to_ids(self.vision_end_token)
        )

        prts_special_tokens = [
            ACTION_START_TOKEN,
            ACTION_PLACEHOLDER_TOKEN,
            ACTION_END_TOKEN,
            CRL_GOAL_REPR_TOKEN,
            CRL_OBS_REPR_TOKEN,
        ]
        num_new_tokens = tokenizer.add_tokens(prts_special_tokens, special_tokens=True)
        logger.info(f"Added {num_new_tokens} new special tokens to the tokenizer.")

        self.action_token = getattr(tokenizer, "action_token", ACTION_PLACEHOLDER_TOKEN)
        self.action_token_id = tokenizer.convert_tokens_to_ids(self.action_token)
        token_dict = {
            "action_start_token_id": ACTION_START_TOKEN,
            "action_token_id": ACTION_PLACEHOLDER_TOKEN,
            "vision_start_token_id": VISION_START_TOKEN,
            "image_token_id": IMAGE_PLACEHOLDER_TOKEN,
            "video_token_id": VIDEO_PLACEHOLDER_TOKEN,
            "crl_goal_repr_token_id": CRL_GOAL_REPR_TOKEN,
            "crl_obs_repr_token_id": CRL_OBS_REPR_TOKEN,
        }
        self.token_ids = {key: tokenizer.convert_tokens_to_ids(value) for key, value in token_dict.items()}

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

        image_inputs = {}
        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_grid_thw = None

        videos_inputs = {}
        if videos is not None:
            videos_inputs = self.video_processor(videos=videos, **output_kwargs["videos_kwargs"])
            video_grid_thw = videos_inputs["video_grid_thw"]
            if "return_metadata" not in kwargs:
                video_metadata = videos_inputs.pop("video_metadata", None)
            else:
                video_metadata = videos_inputs.get("video_metadata", None)
        else:
            video_grid_thw = None
            video_metadata = 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]:
                    if video_metadata is not None and index < len(video_metadata):
                        metadata = video_metadata[index]
                        if metadata.fps is None:
                            logger.warning_once(
                                "Qwen3VL requires frame timestamps to construct prompts, but the `fps` of the input video could not be inferred. "
                                "Probably `video_metadata` was missing from inputs and you passed pre-sampled frames. "
                                "Defaulting to `fps=24`. Please provide `video_metadata` for more accurate results."
                            )
                            metadata.fps = 24 if metadata.fps is None else metadata.fps

                        curr_timestamp = self._calculate_timestamps(
                            metadata.frames_indices,
                            metadata.fps,
                            self.video_processor.merge_size,
                        )

                        video_placeholder = ""
                        frame_seqlen = video_grid_thw[index][1:].prod() // merge_length
                        for frame_idx in range(video_grid_thw[index][0]):
                            curr_time = curr_timestamp[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)
                    else:
                        num_video_tokens = video_grid_thw[index].prod() // merge_length
                        text[i] = text[i].replace(self.video_token, "<|placeholder|>" * num_video_tokens, 1)

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

        return_tensors = output_kwargs["text_kwargs"].pop("return_tensors", None)
        return_mm_token_type_ids = 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"])

        if return_mm_token_type_ids:
            array_ids = np.array(text_inputs["input_ids"])
            mm_token_type_ids = np.zeros_like(text_inputs["input_ids"])
            mm_token_type_ids[array_ids == self.image_token_id] = 1
            text_inputs["mm_token_type_ids"] = mm_token_type_ids.tolist()

        output_data = {**text_inputs, **image_inputs, **videos_inputs}
        if actions is not None:
            output_data["actions"] = actions

        return BatchFeature(data=output_data, tensor_type=return_tensors)

    def _calculate_timestamps(self, indices: Union[list[int], np.ndarray], video_fps: float, merge_size: int = 2):
        if not isinstance(indices, list):
            indices = indices.tolist()
        if len(indices) % merge_size != 0:
            indices.extend(indices[-1] for _ in range(merge_size - len(indices) % merge_size))
        timestamps = [idx / video_fps for idx in indices]
        timestamps = [
            (timestamps[i] + timestamps[i + merge_size - 1]) / 2 for i in range(0, len(timestamps), merge_size)
        ]
        return timestamps

    def _get_num_multimodal_tokens(self, image_sizes=None, video_sizes=None, **kwargs):
        vision_data = {}
        if image_sizes is not None:
            images_kwargs = Qwen3VLProcessorKwargs._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})

        if video_sizes is not None:
            videos_kwargs = Qwen3VLProcessorKwargs._defaults.get("videos_kwargs", {})
            videos_kwargs.update(kwargs)
            merge_size = videos_kwargs.get("merge_size", None) or self.video_processor.merge_size
            num_video_patches = [
                self.video_processor.get_number_of_video_patches(*video_size, videos_kwargs)
                for video_size in video_sizes
            ]
            num_video_tokens = [(num_patches // merge_size**2) for num_patches in num_video_patches]
            vision_data["num_video_tokens"] = num_video_tokens

        return MultiModalData(**vision_data)

    def set_action_tokenizer(self, action_tokenizer):
        self.action_tokenizer = action_tokenizer

        prts_fast_action_tokens = [f"<|action_token_{i}|>" for i in range(action_tokenizer.vocab_size)]
        num_new_tokens = self.tokenizer.add_tokens(prts_fast_action_tokens, special_tokens=True)
        logger.info(f"Added {num_new_tokens} FAST action tokens to the tokenizer.")

        self.action_token_start_index = self.tokenizer.convert_tokens_to_ids("<|action_token_0|>")
        self.action_vocab_size = action_tokenizer.vocab_size

        token_ids = self.tokenizer.convert_tokens_to_ids(prts_fast_action_tokens)
        self.action_mapper = {k: v for k, v in zip(prts_fast_action_tokens, token_ids, strict=True)}

    def preprocess_action(self, actions, **kwargs):
        raise NotImplementedError

    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,
        )

    @property
    def model_input_names(self):
        tokenizer_input_names = self.tokenizer.model_input_names
        image_processor_input_names = self.image_processor.model_input_names
        return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names))


PRTS_Qwen3VLProcessor.register_for_auto_class()

__all__ = ["PRTS_Qwen3VLProcessor"]