processing_molmoact.py

"""
Processor class for MolmoAct.
"""
from typing import List, Optional, Union, Dict, Tuple

import PIL
from PIL import ImageFile, ImageOps

try:
    from typing import Unpack
except ImportError:
    from typing_extensions import Unpack

import numpy as np
import torch

from transformers.image_utils import ImageInput
from transformers.processing_utils import (
    ProcessingKwargs,
    ProcessorMixin,
)
from transformers.feature_extraction_utils import BatchFeature
from transformers.tokenization_utils_base import TextInput, PreTokenizedInput
from transformers.utils import logging

from transformers import AutoTokenizer
from .image_processing_molmoact import MolmoActImagesKwargs, MolmoActImageProcessor


logger = logging.get_logger(__name__)


# Special tokens, these should be present in any tokenizer we use since the preprocessor uses them
IMAGE_PATCH_TOKEN = f"<im_patch>"  # Where to insert high-res tokens
IMAGE_LOW_RES_TOKEN = f"<im_low>"  # Where to insert low-res tokens
IM_START_TOKEN = f"<im_start>"
IM_END_TOKEN = f"<im_end>"
IM_COL_TOKEN = f"<im_col>"
IMAGE_PROMPT = "<|image|>"

EXTRA_TOKENS = (IM_START_TOKEN, IM_END_TOKEN, IMAGE_PATCH_TOKEN,
                IM_COL_TOKEN, IMAGE_PROMPT, IMAGE_LOW_RES_TOKEN)


DEMO_STYLES = [
    "point_count",
    "pointing",
    "cosyn_point",
    "user_qa",
    "long_caption",
    "short_caption",
    "correction_qa",
    "demo",
    "android_control",
]


def setup_pil():
    PIL.Image.MAX_IMAGE_PIXELS = None
    ImageFile.LOAD_TRUNCATED_IMAGES = True


def get_special_token_ids(tokenizer: AutoTokenizer) -> Dict[str, int]:
    ids = tokenizer.encode("".join(EXTRA_TOKENS), add_special_tokens=False)
    assert len(ids) == len(EXTRA_TOKENS)
    return {k: i for k, i in zip(EXTRA_TOKENS, ids)}


def load_image(image: Union[PIL.Image.Image, np.ndarray]) -> np.ndarray:
    """Load image"""
    setup_pil()
    if isinstance(image, PIL.Image.Image):
        image = image.convert("RGB")
        image = ImageOps.exif_transpose(image)
        return np.array(image)
    elif isinstance(image, np.ndarray):
        assert len(image.shape) == 3, "Image should have 3 dimensions"
        assert image.shape[2] == 3, "Image should have 3 channels"
        assert image.dtype == np.uint8, "Image should have uint8 type"
        return image
    else:
        raise ValueError("Image should be PIL.Image or np.ndarray")


class MolmoActProcessorKwargs(ProcessingKwargs, total=False):
    """MolmoAct processor kwargs"""
    images_kwargs: MolmoActImagesKwargs
    _defaults = {
        "text_kwargs": {
            "padding": False,
        },
    }


class MolmoActProcessor(ProcessorMixin):
    attributes = ["image_processor", "tokenizer"]
    optional_attributes = [
        "chat_template",
        "prompt_templates",
        "message_format",
        "system_prompt",
        "style",
        "always_start_with_space",
        "default_inference_len",
        "use_col_tokens",
        "image_padding_mask",
    ]
    image_processor_class = "AutoImageProcessor"
    tokenizer_class = "AutoTokenizer"

    def __init__(
        self,
        image_processor: MolmoActImageProcessor = None,
        tokenizer: AutoTokenizer = None,
        chat_template: Optional[str] = None,
        prompt_templates: Optional[str] = "uber_model",
        message_format: Optional[str] = "role",
        system_prompt: Optional[str] = "demo_or_style",
        style: Optional[str] = "demo",
        always_start_with_space: Optional[bool] = False,
        default_inference_len: Optional[int] = 65,
        use_col_tokens: Optional[bool] = True,
        image_padding_mask: bool = False,
        **kwargs
    ) -> None:
        if tokenizer.padding_side != "left":
            logger.warning(f"Tokenizer {tokenizer.name_or_path} is not left-padded, padding side will be set to left")
            tokenizer.padding_side = "left"  # type: ignore
        super().__init__(
            image_processor,
            tokenizer,
            chat_template=chat_template,
            prompt_templates=prompt_templates,
            message_format=message_format,
            system_prompt=system_prompt,
            style=style,
            always_start_with_space=always_start_with_space,
            default_inference_len=default_inference_len,
            use_col_tokens=use_col_tokens,
            image_padding_mask=image_padding_mask,
        )
        self._special_tokens = None

    @property
    def special_token_ids(self):
        if self._special_tokens is None:
            self._special_tokens = get_special_token_ids(self.tokenizer)
        return self._special_tokens
    
    def get_user_prompt(self, text: TextInput) -> str:
        """Get user prompt"""
        if self.prompt_templates == "none":
            return ""
        elif self.prompt_templates == "uber_model":
            return text
        else:
            raise NotImplementedError(self.prompt_templates)
    
    def get_prefix(self) -> str:
        """Get prefix"""
        if self.system_prompt == "style_and_length":  # captioner
            assert self.style in ["long_caption"]
            style = self.style
            n = None if self.default_inference_len is None else str(self.default_inference_len)
            if n is not None and len(n) > 0:  # allow empty string to signal unconditioned
                prefix = style + " " + n + ":"
            else:
                prefix = style + " :"
        elif self.system_prompt == "demo_or_style":  # demo model
            if self.style in DEMO_STYLES:
                prefix = ""
            else:
                prefix = self.style + ":"
        else:
            raise NotImplementedError(self.system_prompt)
        return prefix
    
    def format_prompt(self, prompt: str) -> str:
        """Format prompt"""
        if self.message_format == "none":
            pass
        elif self.message_format == "role":
            prompt = "User: " + prompt + " Assistant:"
        else:
            raise NotImplementedError(self.message_format)
        
        if self.always_start_with_space:
            prompt = " " + prompt
        
        return prompt
    
    def get_prompt(self, text: TextInput) -> str:
        prompt = self.get_user_prompt(text)
        if self.system_prompt and self.system_prompt != "none":
            prefix = self.get_prefix()
            if len(prefix) > 0 and len(prompt) > 0:
                prompt = prefix + " " + prompt
            elif len(prefix) > 0:
                prompt = prefix
        prompt = self.format_prompt(prompt)
        return prompt

    def get_image_tokens(self, image_grid: np.ndarray):
        joint = []
        for h, w in image_grid:
            per_row = np.full(w, IMAGE_PATCH_TOKEN)
            if self.use_col_tokens:
                per_row = np.concatenate([per_row, [IM_COL_TOKEN]], 0)
            extra_tokens = np.tile(per_row, [h])
            joint += [
                [IM_START_TOKEN],
                extra_tokens,
                [IM_END_TOKEN],
            ]
        return np.concatenate(joint)

    def insert_bos_numpy(
        self,
        input_ids: np.ndarray,
        attention_mask: np.ndarray,
        bos_token_id: int,
        pad_token_id: int,
    ):
        """
        Args:
            input_ids: [B, S] array with left padding
            attention_mask: [B, S] array (0 for pad, 1 for valid)
            bos_token_id: int
            pad_token_id: int
        Returns:
            input_ids_out: [B, S] or [B, S+1] array with bos inserted if needed
            attention_mask_out: same shape as input_ids_out
        """

        need_to_expand = len(input_ids.shape) == 1
        if need_to_expand:
            input_ids = input_ids[None, :]
            attention_mask = attention_mask[None, :]

        B, S = input_ids.shape

        # Handle zero-length sequence
        if S == 0:
            new_input_ids = np.full((B, 1), bos_token_id, dtype=input_ids.dtype)
            new_attention_mask = np.ones((B, 1), dtype=attention_mask.dtype)
            if need_to_expand:
                new_input_ids = new_input_ids[0]
                new_attention_mask = new_attention_mask[0]
            return new_input_ids, new_attention_mask

        first_valid_index = (attention_mask == 1).argmax(axis=-1)  # [B]
        bos_already_present = np.all(input_ids[np.arange(B), first_valid_index] == bos_token_id)

        if bos_already_present:
            if need_to_expand:
                input_ids = input_ids[0]
                attention_mask = attention_mask[0]
            return input_ids, attention_mask
        else:
            new_input_ids = np.full((B, S+1), pad_token_id, dtype=input_ids.dtype)
            new_attention_mask = np.zeros((B, S+1), dtype=attention_mask.dtype)

            src_idx = np.tile(np.arange(S), (B, 1))  # [B, S]
            valid_mask = src_idx >= first_valid_index[:, None]  # [B, S]
            tgt_idx = src_idx + 1  # shit right
            batch_idx = np.tile(np.arange(B)[:, None], (1, S))  # [B, S]

            # flatten valid_positions
            flat_vals = input_ids[valid_mask]
            flat_batch = batch_idx[valid_mask]
            flat_tgt = tgt_idx[valid_mask]

            new_input_ids[flat_batch, flat_tgt] = flat_vals
            new_attention_mask[flat_batch, flat_tgt] = 1
            
            insert_pos = first_valid_index
            new_input_ids[np.arange(B), insert_pos] = bos_token_id
            new_attention_mask[np.arange(B), insert_pos] = 1

            if need_to_expand:
                new_input_ids = new_input_ids[0]
                new_attention_mask = new_attention_mask[0]

            return new_input_ids, new_attention_mask

    def insert_bos_torch(
        self,
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor,
        bos_token_id: int,
        pad_token_id: int,
    ):
        """
        Args:
            input_ids: [B, S] tensor with left padding
            attention_mask: [B, S] tensor (0 for pad, 1 for valid)
            bos_token_id: int
            pad_token_id: int
        Returns:
            input_ids_out: [B, S] or [B, S+1] tensor with bos inserted if needed
            attention_mask_out: same shape as input_ids_out
        """
        
        B, S = input_ids.shape
        device = input_ids.device

        # Handle zero-length sequence
        if S == 0:
            new_input_ids = torch.full((B, 1), bos_token_id, dtype=input_ids.dtype, device=device)
            new_attention_mask = torch.ones((B, 1), dtype=attention_mask.dtype, device=device)
            return new_input_ids, new_attention_mask

        first_valid_index = (attention_mask == 1).long().argmax(dim=-1)  # [B]
        bos_already_present = (input_ids[torch.arange(B), first_valid_index] == bos_token_id).all()

        if bos_already_present:
            return input_ids, attention_mask
        else:
            new_input_ids = torch.full((B, S+1), pad_token_id, dtype=input_ids.dtype, device=device)
            new_attention_mask = torch.zeros((B, S+1), dtype=attention_mask.dtype, device=device)

            src_idx = torch.arange(S, device=device).expand(B, S)  # [B, S]
            valid_mask = src_idx >= first_valid_index.unsqueeze(1)  # [B, S]
            tgt_idx = src_idx + 1  # shift right
            batch_idx = torch.arange(B, device=device).unsqueeze(1).expand_as(src_idx)

            flat_vals = input_ids[valid_mask]
            flat_batch = batch_idx[valid_mask]
            flat_tgt = tgt_idx[valid_mask]

            new_input_ids[flat_batch, flat_tgt] = flat_vals
            new_attention_mask[flat_batch, flat_tgt] = 1

            insert_pos = first_valid_index
            batch_indices = torch.arange(B, device=device)
            new_input_ids[batch_indices, insert_pos] = bos_token_id
            new_attention_mask[batch_indices, insert_pos] = 1

            return new_input_ids, new_attention_mask

    def __call__(
        self,
        text: Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]] = None,
        images: Union[ImageInput, List[ImageInput]] = None,
        apply_chat_template: bool = False,
        **kwargs: Unpack[MolmoActProcessorKwargs],
    ) -> BatchFeature:
        if images is None and text is None:
            raise ValueError("You have to specify at least one of `images` or `text`.")

        output_kwargs = self._merge_kwargs(
            MolmoActProcessorKwargs,
            tokenizer_init_kwargs=self.tokenizer.init_kwargs,
            **kwargs,
        )

        if isinstance(text, (list, tuple)) and isinstance(images, (list, tuple)):
            if len(text) != len(images):
                raise ValueError("You have to provide the same number of text and images")
            if len(text) > 1 and not output_kwargs["text_kwargs"].get("padding", False):
                raise ValueError("You have to specify padding when you have multiple text inputs")

        if isinstance(text, str):
            text = [text]
        elif not isinstance(text, list) and not isinstance(text[0], str):
            raise ValueError("Invalid input text. Please provide a string, or a list of strings")

        if images is not None:
            image_inputs = self.image_processor(images, **output_kwargs["images_kwargs"])
        else:
            image_inputs = {}
        
        if apply_chat_template:
            text = [self.get_prompt(t) for t in text]
        
        prompt_strings = text
        if image_inputs.get("images", None) is not None:

            prompt_strings = []
            for idx, image_grids in enumerate(image_inputs.pop("image_grids")):
                if isinstance(image_grids, torch.Tensor):
                    image_grids = image_grids.cpu().numpy()
                if isinstance(images, (list, tuple)) and isinstance(images[idx], (list, tuple)):
                    image_grids = image_grids[~np.all(image_grids == -1, axis=-1)]
                    offset = 2 if len(images[idx]) < len(image_grids) else 1 # whether to use both low and high res images
                    all_image_strings = []
                    for i in range(0, len(image_grids), offset):
                        image_grids_i = image_grids[i:i+offset]
                        image_tokens = self.get_image_tokens(image_grids_i)
                        img_ix = i // offset
                        all_image_strings.append(f"Image {img_ix + 1}" + "".join(image_tokens))
                    image_string = "".join(all_image_strings)
                    prompt_strings.append(image_string + text[idx])
                else:
                    image_grids = image_grids[~np.all(image_grids == -1, axis=-1)]
                    assert len(image_grids) in [1, 2], "Only one or two crops are supported for single image inputs"
                    image_tokens = self.get_image_tokens(image_grids)
                    image_string = "".join(image_tokens)
                    prompt_strings.append(image_string + text[idx])
        
        text_inputs = self.tokenizer(prompt_strings, **output_kwargs["text_kwargs"])
        
        input_ids = text_inputs["input_ids"]
        attention_mask = text_inputs["attention_mask"]

        is_list = isinstance(input_ids, (list, tuple))
        if is_list:
            input_ids = np.array(input_ids)
            attention_mask = np.array(attention_mask)
        
        use_numpy = isinstance(attention_mask, np.ndarray)

        if use_numpy and np.issubdtype(input_ids.dtype, np.floating):
            input_ids = input_ids.astype(np.int64)
            attention_mask = attention_mask.astype(np.int64)
        elif not use_numpy and torch.is_floating_point(input_ids):
            input_ids = input_ids.to(torch.int64)
            attention_mask = attention_mask.to(torch.int64)
        
        bos = self.tokenizer.bos_token_id or self.tokenizer.eos_token_id
        if use_numpy:
            input_ids, attention_mask = self.insert_bos_numpy(
                input_ids, attention_mask, bos, self.tokenizer.pad_token_id
            )
        else:
            input_ids, attention_mask = self.insert_bos_torch(
                input_ids, attention_mask, bos, self.tokenizer.pad_token_id
            )
        if is_list:
            input_ids = input_ids.tolist()  # type: ignore
            attention_mask = attention_mask.tolist()  # type: ignore
        text_inputs["input_ids"] = input_ids
        text_inputs["attention_mask"] = attention_mask

        if kwargs.get("device", None) is not None:
            text_inputs = text_inputs.to(device=kwargs.get("device"), non_blocking=True)
        # there is no bos token in Qwen tokenizer
        return BatchFeature(
            data={**text_inputs, **image_inputs}, tensor_type=output_kwargs["common_kwargs"]["return_tensors"]
        )

    def batch_decode(self, *args, **kwargs):
        """
        This method forwards all its arguments to LlamaTokenizerFast'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 LlamaTokenizerFast's [`~PreTrainedTokenizer.decode`]. Please refer to
        the docstring of this method for more information.
        """
        return self.tokenizer.decode(*args, **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))


MolmoActProcessor.register_for_auto_class()