SPRVLA
/

libero_spatial

@@ -1,463 +0,0 @@
-"""
-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()