MSALab/PerceptionDLM

1 2import math3import torch4import warnings5import PIL.Image6 7from torch.nn import functional as F8from collections import UserDict, OrderedDict9from typing import Union, Optional, Tuple, List, Dict, Any10 11from transformers.image_utils import load_image12from transformers.feature_extraction_utils import BatchFeature13from .chat_template_utils import render_jinja_template14from transformers.processing_utils import ProcessorMixin, AllKwargsForChatTemplate15 16 17class PDMLLMProcessor(ProcessorMixin):18    attributes = ["tokenizer", "image_processor"]19    optional_attributes = ['chat_template']20    model_input_names = ['input_ids', 'attention_mask', 'pixel_values']21    image_processor_class = "AutoImageProcessor"22    tokenizer_class = "AutoTokenizer"23 24    def __init__(25            self, tokenizer, image_processor, chat_template=None,26            image_size=512,27            patch_size=16,28            downsample_ratio=0.5,29            max_sub_img=6,30            min_sub_img=1,31            image_token='<IMG_CONTEXT>',32            image_start_token='<img>',33            image_end_token='</img>',34            special_tokens=['<IMG_CONTEXT>', '<img>', '</img>'],35            **kwargs):36        if chat_template is None:37            chat_template = "{% for message in messages %}{% if loop.first and message['role'] != 'system' %}<|start_header_id|>system<|end_header_id|>\nYou are a helpful assistant.<|eot_id|>\n{% endif %}<|start_header_id|>{{ message['role'] }}<|end_header_id|>\n{% if message['role'] == 'assistant' %}{% generation %}{{ message['content'][0]['text'] }}<|eot_id|>{% endgeneration %}{% else %}{% for content in message['content'] %}{% if content['type'] == 'image' or 'image' in content or 'image_url' in content %}<img><IMG_CONTEXT></img>{% elif content['type'] == 'video' or 'video' in content %}<video><VIDEO_CONTEXT></video>{% elif 'text' in content %}{{ content['text'] }}{% endif %}{% endfor %}<|eot_id|>\n{% endif %}{% endfor %}{% if add_generation_prompt %}<|start_header_id|>assistant<|end_header_id|>\n{% endif %}"38        super().__init__(tokenizer=tokenizer, image_processor=image_processor, chat_template=chat_template)39        if isinstance(image_size, List) or isinstance(image_size, Tuple):40            image_size = image_size[0]41        self.num_image_token = int((image_size // patch_size) ** 2 * (downsample_ratio ** 2))42 43        self.vision_token_share_pe = kwargs.get('vision_token_share_pe', True)44        self.image_token_len = kwargs.pop('image_token_len', 256)45        self.max_sub_img = max_sub_img46        self.min_sub_img = min_sub_img47 48        self.image_token = image_token49        self.image_start_token = image_start_token50        self.image_end_token = image_end_token51        special_tokens = special_tokens + [f'<|Mask_Cap_{i}|>' for i in range(16)]52        self.tokenizer.add_special_tokens({'additional_special_tokens': special_tokens}, replace_additional_special_tokens=False)53        self.image_token_id = self.tokenizer.convert_tokens_to_ids(self.image_token)54        self.image_start_token_id = self.tokenizer.convert_tokens_to_ids(self.image_start_token)55        self.image_end_token_id = self.tokenizer.convert_tokens_to_ids(self.image_end_token)56        if 'llada' in tokenizer.name_or_path.lower():57            self._pad_token_id = self.tokenizer.convert_tokens_to_ids("<|eot_id|>")58 59        if isinstance(image_size, int):60            image_size = (image_size, image_size)61        else:62            image_size = image_size63        self.image_size = image_size64        assert image_size[0] == image_size[1]65 66    def apply_chat_template(self, conversation, chat_template = None, **kwargs) -> str:67        if chat_template is None:68            chat_template = self.chat_template69 70        # Split template kwargs from processor/tokenization kwargs so that71        # `tokenize=True` can reuse the processor pipeline without polluting72        # the template rendering inputs.73        tokenize = kwargs.pop("tokenize", False)74        return_dict = kwargs.pop("return_dict", False)75        return_tensors = kwargs.pop("return_tensors", None)76        images = kwargs.pop("images", [])77        videos = kwargs.pop("videos", None)78 79        if not images:80            for message in conversation:81                content = message.get("content", [])82                if isinstance(content, list):83                    for item in content:84                        if isinstance(item, dict) and (item.get("type") == "image" or "image" in item):85                            image = item.get("image") or item.get("image_url")86                            if image is not None:87                                images.append(image)88 89        processor_kwargs = {}90        for key in ("padding", "truncation", "max_length"):91            if key in kwargs:92                processor_kwargs[key] = kwargs.pop(key)93        if return_tensors is not None:94            processor_kwargs["return_tensors"] = return_tensors95 96        processed_kwargs = {97            "mm_load_kwargs": {},98            "template_kwargs": {},99        }100        # for kwarg_type in processed_kwargs:101        #     for key in AllKwargsForChatTemplate.__annotations__[kwarg_type].__annotations__.keys():102        #         kwarg_type_defaults = AllKwargsForChatTemplate.__annotations__[kwarg_type]103        #         default_value = getattr(kwarg_type_defaults, key, None)104        #         value = kwargs.pop(key, default_value)105        #         if value is not None and not isinstance(value, dict):106        #             processed_kwargs[kwarg_type][key] = value107 108        # Pass unprocessed custom kwargs109        processed_kwargs["template_kwargs"].update(kwargs)110        conversations = [conversation]111 112        prompt, generation_indices = render_jinja_template(113            conversations=conversations,114            chat_template=chat_template,115            return_assistant_tokens_mask=True,116            **processed_kwargs["template_kwargs"],  # different flags such as `return_assistant_mask`117            **self.tokenizer.special_tokens_map,  # tokenizer special tokens are used by some templates118        )119 120        if not tokenize:121            return prompt, generation_indices122 123        # Reuse the processor pipeline to produce tokenized inputs.124        model_inputs = self(125            text=prompt,126            images=images,127            videos=videos,128            generation_indices=generation_indices,129            **processor_kwargs,130        )131        # if return_dict:132        #     return model_inputs133        return model_inputs134 135    def __call__(self, text=None, images=[], videos=None, generation_indices=None, **kwargs) ->BatchFeature:136        inputs = self.tokenizer(text, padding=False, truncation=False, return_attention_mask=False)137        assistant_masks = []138        input_ids = inputs["input_ids"]139        for i in range(len(input_ids)):140            current_mask = [0] * len(input_ids[i])141            if 'llada' in self.tokenizer.name_or_path.lower():142                for assistant_start_char, assistant_end_char in generation_indices[i]:143                    start_token = inputs.char_to_token(i, assistant_start_char)144                    end_token = inputs.char_to_token(i, assistant_end_char - 1)145                    if start_token is None:146                        # start_token is out of bounds maybe due to truncation.147                        break148                    for token_id in range(start_token, end_token + 1 if end_token else len(input_ids[i])):149                        current_mask[token_id] = 1150            151            assistant_masks.append(current_mask)152 153        inputs["assistant_masks"] = assistant_masks[0]154        inputs['input_ids'] = input_ids[0]155 156        truncation = kwargs.pop('truncation', False)157        max_length = kwargs.pop('max_length', 1024)158        padding = kwargs.pop('padding', False)159 160        inputs = self.process_images(images, inputs=inputs)161        if isinstance(inputs, UserDict):162            inputs = inputs.data163        164        if 'attention_mask' not in inputs:165            inputs['attention_mask'] = [1] * len(inputs['input_ids'])166        if 'assistant_masks' in inputs:167            inputs['prompt_mask'] = [1-x for x in inputs.pop('assistant_masks')]168 169        inputs = self.process_inputs(inputs)170        if truncation and len(inputs['input_ids']) > max_length:171            inputs = self.truncate(inputs, max_length)172        if padding and len(inputs['input_ids']) < max_length:173            inputs = self.padding(inputs, max_length)174 175        inputs = self.to_tensor(inputs)176        self.check(inputs)177        if self.vision_token_share_pe:178            position_ids = self.get_position_ids(inputs)179            position_ids = torch.tensor([position_ids], dtype=torch.long)180            inputs['position_ids'] = position_ids181 182        inputs.pop('sub_image_nums', None)183 184        return BatchFeature(inputs)185 186    def get_position_ids(self, inputs: Dict[str, Any]):187        input_ids = inputs['input_ids'][0]188        image_token_lens = self.get_image_token_length(inputs)189        position_ids = []190        i, j = 0, 0191        while len(position_ids) < len(input_ids):192            if input_ids[len(position_ids)] == self.image_token_id:193                image_token_len = image_token_lens[j]194                assert image_token_len % self.image_token_len == 0195                num_views = image_token_len // self.image_token_len196                for _ in range(num_views):197                    position_ids += [i] * self.image_token_len # 同一个图像的所有 token 共享相同的位置编码198                    i += 1199                j += 1200            else:201                position_ids.append(i)202                i += 1203 204        assert j == len(image_token_lens) and len(position_ids) == len(input_ids), \205            f"Wrong position_ids, {j} != {len(image_token_lens)} or {len(position_ids)} != {len(input_ids)}"206 207        return position_ids208    209    def process_images(self, images, inputs):210        images = [load_image(img) for img in images]211        if len(images) > 0:212            processed_images = []213            sub_image_nums = []214            for image in images:215                if len(images) > 1:216                    # for multi images, remove the split strategy217                    sub_images = dynamic_preprocess(218                        image, min_num=1,219                        max_num=1,220                        image_size=self.image_size[0], use_thumbnail=True)221                else:222                    sub_images = dynamic_preprocess(223                        image, min_num=self.min_sub_img,224                        max_num=self.max_sub_img,225                        image_size=self.image_size[0], use_thumbnail=True)226 227                sub_image_nums.append(len(sub_images))228                processed_images += sub_images229            # print([_img.size for _img in processed_images])230            pixel_values = self.image_processor.preprocess(231                images=processed_images, return_tensors="pt"232            )["pixel_values"] # (N, c, h, w)233        else:234            pixel_values = torch.zeros((235                1, 3, self.image_size[0], self.image_size[1]), dtype=torch.float32236            )237            sub_image_nums = []238 239        inputs['pixel_values'] = pixel_values240        inputs['sub_image_nums'] = sub_image_nums241        return inputs242    243    def truncate(self, inputs: Dict[str, Any], max_length: int):244        assert self.image_token_id not in inputs['input_ids'][max_length:], f"Truncate image token is not allowed."245        inputs['input_ids'] = inputs['input_ids'][:max_length]246        inputs['attention_mask'] = inputs['attention_mask'][:max_length]247        if 'prompt_mask' in inputs:248            inputs['prompt_mask'] = inputs['prompt_mask'][:max_length]249        return inputs250 251    def get_image_token_length(self, inputs: Dict[str, Any]) -> List[int]:252        sub_image_nums = inputs.get('sub_image_nums', None)253        if sub_image_nums is None or len(sub_image_nums) == 0:254            return []255        image_token_lens = [_num * self.num_image_token for _num in sub_image_nums]256        return image_token_lens257 258    def process_inputs(self, inputs: Dict[str, Any]):259        graft_token_lens = self._get_graft_token_length(inputs)260        inputs['input_ids'] = self._graft_token(inputs['input_ids'], graft_token_lens, self.image_token_id)261        inputs['attention_mask'] = self._graft_token(inputs['attention_mask'], graft_token_lens, 'replicate')262        if 'prompt_mask' in inputs:263            inputs['prompt_mask'] = self._graft_token(inputs['prompt_mask'], graft_token_lens, 'replicate')264        return inputs265    266    def _graft_token(self, seq, graft_token_lens, value):267        if value == 'replicate':268            for i in reversed(graft_token_lens.keys()):269                seq[i:] = [seq[i]] * graft_token_lens[i] + seq[i+1:]270        else:271            for i in reversed(graft_token_lens.keys()):272                seq[i:] = [value] * graft_token_lens[i] + seq[i+1:]273        return seq274    275    def _get_graft_token_length(self, inputs: Dict[str, Any]) -> Dict[int, int]:276        image_token_pos = [i for i, x in enumerate(inputs['input_ids']) if x == self.image_token_id]277        image_token_lens = self.get_image_token_length(inputs)278        assert len(image_token_pos) == len(image_token_lens), \279            "Wrong image token count, " \280            f"image_token_count({len(image_token_pos)}) != image_count({len(image_token_lens)})"281 282        graft_token_lens = OrderedDict(item for item in zip(image_token_pos, image_token_lens))283        return graft_token_lens284    285    def check(self, inputs: Dict[str, Any]):286        image_embed_token_count = torch.count_nonzero(inputs['input_ids'] == self.image_token_id).item()287        image_embed_count = sum(self.get_image_token_length(inputs))288        assert image_embed_token_count == image_embed_count, \289            "Wrong image embed token count, " \290            f"image_embed_token_count({image_embed_token_count}) != image_embed_count({image_embed_count})"291 292    def padding(self, inputs: Dict[str, Any], max_length: int):293        padding_len = max_length - len(inputs['input_ids'])294        inputs['input_ids'] += [self.pad_token_id] * padding_len295        inputs['attention_mask'] += [0] * padding_len296        if 'prompt_mask' in inputs:297            inputs['prompt_mask'] += [0] * padding_len298        return inputs299    300    def decode(self, token_ids: Union[List[int], torch.Tensor], **kwargs):301        if isinstance(token_ids, torch.Tensor):302            token_ids = token_ids.tolist()303        text = self.tokenizer.decode(token_ids, **kwargs)304        return text305 306    def batch_decode(self, sequences: Union[List[List[int]], torch.Tensor], **kwargs):307        if isinstance(sequences, torch.Tensor):308            sequences = sequences.tolist()309        texts = self.tokenizer.batch_decode(sequences, **kwargs)310        return texts311    312    def to_tensor(self, inputs):313        inputs['input_ids'] = torch.tensor([inputs['input_ids']], dtype=torch.long)314        inputs['attention_mask'] = torch.tensor([inputs['attention_mask']], dtype=torch.bool)315        if 'prompt_mask' in inputs:316            inputs['prompt_mask'] = torch.tensor([inputs['prompt_mask']], dtype=torch.bool)317        return inputs318    319    @property320    def pad_token_id(self):321        return self._pad_token_id322    323    def __repr__(self):324        pass325 326    def __str__(self):327        return 'PDMLLMProcessor'328 329def find_closest_aspect_ratio(aspect_ratio, target_ratios, width, height, image_size):330    best_ratio_diff = float('inf')331    best_ratio = (1, 1)332    area = width * height333    for ratio in target_ratios:334        target_aspect_ratio = ratio[0] / ratio[1]335        ratio_diff = abs(aspect_ratio - target_aspect_ratio)336        if ratio_diff < best_ratio_diff:337            best_ratio_diff = ratio_diff338            best_ratio = ratio339        elif ratio_diff == best_ratio_diff:340            if area > 0.5 * image_size * image_size * ratio[0] * ratio[1]:341                best_ratio = ratio342    # print(f'width: {width}, height: {height}, best_ratio: {best_ratio}')343    return best_ratio344 345 346def dynamic_preprocess(image, min_num=1, max_num=6, image_size=512, use_thumbnail=True):347    orig_width, orig_height = image.size348    aspect_ratio = orig_width / orig_height349 350    # calculate the existing image aspect ratio351    target_ratios = set(352        (i, j) for n in range(min_num, max_num + 1) for i in range(1, n + 1) for j in range(1, n + 1) if353        i * j <= max_num and i * j >= min_num)354    target_ratios = sorted(target_ratios, key=lambda x: x[0] * x[1])355 356    # find the closest aspect ratio to the target357    target_aspect_ratio = find_closest_aspect_ratio(358        aspect_ratio, target_ratios, orig_width, orig_height, image_size)359 360    # calculate the target width and height361    target_width = image_size * target_aspect_ratio[0]362    target_height = image_size * target_aspect_ratio[1]363    blocks = target_aspect_ratio[0] * target_aspect_ratio[1]364 365    # resize the image366    resized_img = image.resize((target_width, target_height))367    processed_images = []368    for i in range(blocks):369        box = (370            (i % (target_width // image_size)) * image_size,371            (i // (target_width // image_size)) * image_size,372            ((i % (target_width // image_size)) + 1) * image_size,373            ((i // (target_width // image_size)) + 1) * image_size374        )375        # split the image376        split_img = resized_img.crop(box)377        processed_images.append(split_img)378    assert len(processed_images) == blocks379    if use_thumbnail and len(processed_images) != 1:380        thumbnail_img = image.resize((image_size, image_size))381        processed_images.append(thumbnail_img)382    return processed_images