modeling_ovis_u1.py

import logging
import math
from datetime import datetime
from importlib import import_module
from typing import List, Union, Optional, Dict

import numpy as np
import PIL.Image
import torch
from torch import Tensor
from torch.nn import init
from torch.nn.functional import softmax, gumbel_softmax, pad
from torchvision import transforms
import transformers
from transformers import AutoImageProcessor
from transformers import PreTrainedModel, AutoConfig, AutoModel, AutoTokenizer, AutoModelForCausalLM
from transformers.generation.utils import GenerateOutput
from transformers import CLIPImageProcessor

from .modeling_aimv2 import AIMv2Model
from .configuration_ovis_u1 import BaseVisualTokenizerConfig, Aimv2VisualTokenizerConfig
from .configuration_ovis_u1 import OvisU1Config, ConversationFormatter
from .configuration_ovis_u1 import IGNORE_ID, IMAGE_ATOM_ID, IMAGE_INDICATOR_IDS, IMAGE_TOKEN_ID, VIDEO_TOKEN_ID

# ----------------------------------------------------------------------
#                            Visual Tokenizer
# ----------------------------------------------------------------------
class BaseVisualTokenizer(PreTrainedModel):
    base_model_prefix = "backbone"
    main_input_name = None
    _image_processor_class = None
    _image_processor_kwargs = {}
    _backbone_class = None

    def __init__(self, config: BaseVisualTokenizerConfig, *inputs, **kwargs):
        super().__init__(config, *inputs, **kwargs)
        if kwargs.get('train_from_scratch'):
            # for key in self._image_processor_kwargs.keys():
            #     self._image_processor_kwargs[key] = getattr(self.config, key, self._image_processor_kwargs[key])
            image_processor = self._image_processor_class.from_pretrained(kwargs['backbone_name_or_path'],
                                                                               **self._image_processor_kwargs)

            self.backbone = self._backbone_class.from_pretrained(kwargs['backbone_name_or_path'], **self.config.backbone_kwargs)
            self.config.backbone_config = self.backbone.config

            config = image_processor.to_dict()
            if getattr(self.config, 'image_processor_new_kwargs', None) is not None:
                for key in self.config.image_processor_new_kwargs.keys():
                    config[key] = self.config.image_processor_new_kwargs[key]
            if 'patch_size' not in config:
                assert getattr(self.backbone.config, 'patch_size'), "Patch size must be set."
                config['patch_size'] = self.backbone.config.patch_size
            self.image_processor = self._image_processor_class.from_dict(config)

        else:
            self.image_processor = AutoImageProcessor.from_pretrained(kwargs['image_processor_name_or_path'])
            self.backbone = AutoModel.from_config(self.config.backbone_config)
        head_dim = self.config.vocab_size - len(IMAGE_INDICATOR_IDS)  # reserved tokens for IMAGE_INDICATORS
        self.head = torch.nn.Sequential(
            torch.nn.Linear(
                self.backbone.config.hidden_size * self.config.hidden_stride * self.config.hidden_stride, head_dim,
                bias=False
            ),
            torch.nn.LayerNorm(head_dim)
        )
        assert all((self.image_processor.do_resize,
                    not getattr(self.image_processor, 'do_center_crop', False),
                    self.image_processor.do_rescale,
                    self.image_processor.do_normalize
                    )), f"image_processor `{self.image_processor}` is not supported currently"

    def get_backbone(self):
        return self.backbone

    def get_monitor_tensors(self):
        raise NotImplementedError

    def get_image_processor(self):
        return self.image_processor

    def mock_input(self):
        height, width = self.get_image_size()
        return torch.zeros(1, 3, height, width), self.construct_image_placeholders((1, 1))

    def get_head(self):
        return self.head

    def get_image_size(self):
        raise NotImplementedError

    @staticmethod
    def construct_image_placeholders(grid, data_type='image'):
        if data_type == 'image':
            image_placeholders = [IMAGE_INDICATOR_IDS[0], IMAGE_ATOM_ID, IMAGE_INDICATOR_IDS[1]]
        elif data_type == 'video':
            image_placeholders = [IMAGE_INDICATOR_IDS[2], IMAGE_ATOM_ID, IMAGE_INDICATOR_IDS[2]]
        else:
            raise TypeError
        
        return image_placeholders

    @staticmethod
    def _partition(img_size, grid):
        w, h = img_size
        row_height = h // grid[0]
        col_width = w // grid[1]

        partition = []
        for row in range(grid[0]):
            for col in range(grid[1]):
                left = col * col_width
                upper = row * row_height
                right = w if col == grid[1] - 1 else (col + 1) * col_width
                lower = h if row == grid[0] - 1 else (row + 1) * row_height
                partition.append((left, upper, right, lower))

        return partition

    @staticmethod
    def get_best_grid(img_size, side, max_partition, covering_threshold):

        def _covering_area(left, upper, right, lower, side):
            w = right - left
            h = lower - upper
            w, h = max(w, h), min(w, h)
            if w > side:
                h = h / w * side
                w = side
            return w * h

        img_area = img_size[0] * img_size[1]

        candidate_grids = []
        for i in range(1, max_partition + 1):
            for j in range(1, max_partition + 1):
                if i * j <= max_partition:
                    candidate_grids.append((i, j))

        all_grids = []
        good_grids = []
        for grid in candidate_grids:
            partition = BaseVisualTokenizer._partition(img_size, grid)
            covering_ratio = sum([_covering_area(*p, side) for p in partition]) / img_area
            assert covering_ratio <= 1.0
            all_grids.append((grid, covering_ratio))
            if covering_ratio > covering_threshold:
                good_grids.append((grid, covering_ratio))

        if len(good_grids) > 0:
            # pick the good partition with minimum #sub_images and break the tie using covering_ratio
            return sorted(good_grids, key=lambda x: (x[0][0] * x[0][1], -x[1]))[0][0]
        else:
            # pick the partition with maximum covering_ratio and break the tie using #sub_images
            return sorted(all_grids, key=lambda x: (-x[1], x[0][0] * x[0][1]))[0][0]

    def preprocess_image(self, image: PIL.Image.Image, max_partition=4, covering_threshold=0.9, convert_to_rgb=True):
        def _preprocess(img: PIL.Image.Image, side):
            # first resize and preprocess
            w, h = img.size
            if w == h:
                new_width = new_height = side
            elif w > h:
                new_width = side
                new_height = int(h / w * new_width)
            else:
                new_height = side
                new_width = int(w / h * new_height)
            new_size = dict(height=new_height, width=new_width)
            pixel_values = self.image_processor.preprocess(img, size=new_size, return_tensors='pt')['pixel_values']

            # then pad to square
            square_values = torch.zeros([1, 3, side, side], dtype=pixel_values.dtype, device=pixel_values.device)
            new_height, new_width = pixel_values.shape[2:]
            if new_height == new_width:
                square_values[:, :, :, :] = pixel_values
            elif new_height > new_width:
                from_index = (side - new_width) // 2
                square_values[:, :, :, from_index:from_index + new_width] = pixel_values
            else:
                from_index = (side - new_height) // 2
                square_values[:, :, from_index:from_index + new_height, :] = pixel_values

            return square_values

        if convert_to_rgb and image.mode != 'RGB':
            image = image.convert('RGB')

        sides = self.get_image_size()
        if sides[0] != sides[1]:
            raise ValueError('get_image_size() returns non-square size')
        side = sides[0]
        grid = self.get_best_grid(image.size, side, max_partition, covering_threshold)
        partition = self._partition(image.size, grid)
        crops = [image.crop(p) for p in partition]
        if len(crops) > 1:
            crops.insert(0, image)
        pixel_values = torch.cat([_preprocess(crop, side) for crop in crops], dim=0)
        image_placeholders = self.construct_image_placeholders(grid)
        return pixel_values, image_placeholders

    def get_backbone_layer(self, index):
        if 'aimv2' in self.config.model_type:
            return self.backbone.trunk.blocks[index]
        else:
            return self.backbone.vision_model.encoder.layers[index]

    def tokenize(self, logits):
        def st_argmax(y_soft, dim):  # straight-through softmax
            index = y_soft.max(dim, keepdim=True)[1]
            y_hard = torch.zeros_like(y_soft, memory_format=torch.legacy_contiguous_format).scatter_(dim, index, 1.0)
            ret = y_hard - y_soft.detach() + y_soft
            return ret

        if self.config.tokenize_function == 'softmax':
            tokens = softmax(logits, dim=-1, dtype=torch.float32).to(logits.dtype)
        elif self.config.tokenize_function == 'gumbel_argmax':
            tokens = gumbel_softmax(logits, tau=self.config.tau, hard=True)
        elif self.config.tokenize_function == 'st_argmax':
            tokens = st_argmax(logits, dim=-1)
        else:
            raise ValueError(
                f'Invalid `max_type`, expected softmax or gumbel_argmax or st_argmax, but got {self.config.tokenize_function}')
        return tokens

    def encode(self, pixel_values):
        output = self.backbone(pixel_values, output_hidden_states=True, return_dict=True)
        features = output.hidden_states[-1]
        if self.config.drop_cls_token:
            features = features[:, 1:, :]

        # merge number of `hidden_stride * hidden_stride` hidden states together to reduce token sequence length
        # e.g., for hidden_stride=3, this leads to a token length reduction: 729 -> 81 for siglip
        if self.config.hidden_stride > 1:
            n, l, d = features.shape  # this `d` maybe different from the above `d
            sqrt_l = int(l ** 0.5)
            assert sqrt_l ** 2 == l, "The token sequence length should be a perfect square."
            features = features.reshape(n, sqrt_l, sqrt_l, d)
            pl = (self.config.hidden_stride - (sqrt_l % self.config.hidden_stride)) % self.config.hidden_stride
            features = pad(features, (0, 0, 0, pl, 0, pl), "constant", 0)
            sqrt_l += pl
            features = features.reshape(n, sqrt_l // self.config.hidden_stride, self.config.hidden_stride,
                                        sqrt_l // self.config.hidden_stride, self.config.hidden_stride, d)
            features = features.permute(0, 1, 3, 2, 4, 5)  # [n, sqrt_l/hs, sqrt_l/hs, hs, hs, d]
            features = features.flatten(3)  # [n, sqrt_l/hs, sqrt_l/hs, hs*hs*d]
            features = features.reshape(
                n, -1, self.config.hidden_stride * self.config.hidden_stride * d)

        return features

    def forward(self, pixel_values) -> torch.Tensor:  # [BatchSize, ImageShape] -> [BatchSize, #Token, VocabSize]
        features = self.encode(pixel_values)
        logits = self.head(features)
        tokens = self.tokenize(logits)
        # tokens' shape is [BatchSize, #Token, VocabSize-5], so padding with [BatchSize, #Token, 5], after
        # which, tokens' shape should become [BatchSize, #Token, VocabSize]
        batch_size, token_len, _ = tokens.shape
        padding_tensor = torch.zeros(size=(batch_size, token_len, len(IMAGE_INDICATOR_IDS)),
                                     dtype=tokens.dtype,
                                     device=tokens.device,
                                     layout=tokens.layout,
                                     requires_grad=False)
        tokens = torch.cat((tokens, padding_tensor), dim=2)
        return tokens

class Aimv2VisualTokenizer(BaseVisualTokenizer):
    config_class = Aimv2VisualTokenizerConfig
    supports_gradient_checkpointing = True
    _no_split_modules = ["AIMv2ViTPreprocessor", "AIMv2Block"]
    _image_processor_class = CLIPImageProcessor
    _image_processor_kwargs = dict(do_center_crop=False, crop_size={'height': -1, 'width': -1}, size={'shortest_edge':-1})
    _backbone_class = AIMv2Model
    
    # Copied from qwen2_vl
    def smart_resize(self, 
        height: int, width: int, factor: int = 28, min_pixels: int = 56 * 56, max_pixels: int = 14 * 14 * 4 * 1280
    ):
        """Rescales the image so that the following conditions are met:

        1. Both dimensions (height and width) are divisible by 'factor'.

        2. The total number of pixels is within the range ['min_pixels', 'max_pixels'].

        3. The aspect ratio of the image is maintained as closely as possible.

        """
        
        if height < factor or width < factor:
            print(f"height:{height} or width:{width} must be larger than factor:{factor}")
            if height < width:
                width = round(factor/height*width)
                height = factor
            else:
                height = round(factor/width*height)
                width = factor

        elif max(height, width) / min(height, width) > 200:
            print(
                f"absolute aspect ratio must be smaller than 200, got {max(height, width) / min(height, width)}"
            )
            if height > width:
                height = 200 * width
            else:
                width = 200 * height

        h_bar = round(height / factor) * factor
        w_bar = round(width / factor) * factor
        if h_bar * w_bar > max_pixels:
            beta = math.sqrt((height * width) / max_pixels)
            h_bar = math.floor(height / beta / factor) * factor
            w_bar = math.floor(width / beta / factor) * factor
        elif h_bar * w_bar < min_pixels:
            beta = math.sqrt(min_pixels / (height * width))
            h_bar = math.ceil(height * beta / factor) * factor
            w_bar = math.ceil(width * beta / factor) * factor
        return h_bar, w_bar

    def get_monitor_tensors(self):
        return dict(
            backbone_bottom=self.backbone.trunk.blocks[0].attn.qkv.weight,
            backbone_top=self.backbone.trunk.blocks[-1].attn.qkv.weight,
            head=self.head[0].weight
        )

    def get_min_image_size(self):
        min_pixels = self.image_processor.min_pixels
        max_pixels = self.image_processor.max_pixels
        height = int(min_pixels**0.5)
        width = int(min_pixels**0.5)
        patch_size = self.image_processor.patch_size
        hidden_stride = self.image_processor.hidden_stride
        height, width = self.smart_resize(height, width, patch_size * hidden_stride, min_pixels, max_pixels)
        return height, width
    
    def get_image_size(self):
        min_pixels = self.image_processor.min_pixels
        max_pixels = self.image_processor.max_pixels
        num_pixels = (min_pixels+max_pixels) / 2
        height = int(num_pixels**0.5)
        width = int(num_pixels**0.5)
        patch_size = self.image_processor.patch_size
        hidden_stride = self.image_processor.hidden_stride
        height, width = self.smart_resize(height, width, patch_size * hidden_stride, min_pixels, max_pixels)
        return height, width

    def get_token_length(self, width: int,
                            height: int, 
                            n_frames: int = 1,
                            num_images: int = 1):
        patch_size = self.image_processor.patch_size
        temporal_patch_size = self.image_processor.temporal_patch_size
        hidden_stride = self.image_processor.hidden_stride
        min_pixels = self.image_processor.min_pixels
        max_pixels = self.image_processor.max_pixels
        
        max_pixels = max_pixels // num_images
        min_pixels = min(max_pixels, min_pixels)
        
        resized_height, resized_width = height, width
        resized_height, resized_width = self.smart_resize(
                    height,
                    width,
                    factor=patch_size * hidden_stride,
                    min_pixels=min_pixels,
                    max_pixels=max_pixels,
                )
       
        if n_frames % temporal_patch_size != 0:
            n_frames = n_frames + temporal_patch_size - 1
        grid_t = n_frames // temporal_patch_size
        grid_h, grid_w = resized_height // patch_size // hidden_stride, resized_width // patch_size // hidden_stride

        return grid_t * grid_w * grid_h

    def mock_input(self):
        height, width = self.get_min_image_size()
        return torch.zeros(1, 3, height, width), self.construct_image_placeholders((1, 1))

    def preprocess_image(self, images: Union[PIL.Image.Image, List[PIL.Image.Image]], 
                            convert_to_rgb: Optional[bool] = True, 
                            num_images: Optional[int] = 1,
                            min_pixels: Optional[int] = None, 
                            max_pixels: Optional[int] = None,
                            multimodal_type: Optional[str] = 'single_image'):


        patch_size = self.image_processor.patch_size # 14
        temporal_patch_size = self.image_processor.temporal_patch_size # 1
        hidden_stride = self.image_processor.hidden_stride # 2
        min_pixels = min_pixels or self.image_processor.min_pixels # 200704
        max_pixels = max_pixels or self.image_processor.max_pixels # 3211264
        
        max_pixels = max_pixels // num_images
        min_pixels = min(max_pixels, min_pixels)

        if not isinstance(images, list):
            images = [images]
        if multimodal_type == 'video':
            assert len(images) >= 1
        else:
            assert len(images) == 1
        images = [image.convert("RGB") if convert_to_rgb and image.mode != 'RGB' else image for image in images ]
        # images = [np.array(image) for image in images]
        
        width, height = images[0].size
        resized_height, resized_width = height, width
        processed_images = []
        for image in images:
            resized_height, resized_width = self.smart_resize(
                height,
                width,
                factor=patch_size * hidden_stride,
                min_pixels=min_pixels,
                max_pixels=max_pixels,
            )
            new_size = dict(height=resized_height, width=resized_width)
            image_pt = self.image_processor.preprocess(image, size=new_size, return_tensors="np")['pixel_values'][0]
            
            processed_images.append(image_pt)

        patches = np.array(processed_images)
        # if data_format == ChannelDimension.LAST:
        #     patches = patches.transpose(0, 3, 1, 2)
        if patches.shape[0] % temporal_patch_size != 0:
            repeats = np.repeat(patches[-1][np.newaxis], temporal_patch_size - 1, axis=0)
            patches = np.concatenate([patches, repeats], axis=0)
        channel = patches.shape[1]
        grid_t = patches.shape[0] // temporal_patch_size # 1
        grid_h, grid_w = resized_height // patch_size, resized_width // patch_size # 32, 32
        
        patches = patches.reshape(
            grid_t,
            temporal_patch_size,
            channel,
            grid_h // hidden_stride,
            hidden_stride,
            patch_size,
            grid_w // hidden_stride,
            hidden_stride,
            patch_size,
        )
        patches = patches.transpose(0, 3, 6, 4, 7, 2, 1, 5, 8)
        flatten_patches = patches.reshape(
            grid_t * grid_h * grid_w, channel * temporal_patch_size * patch_size * patch_size
        )
        # 1024, 588

        image_placeholders = self.construct_image_placeholders((1, 1), data_type='video' if multimodal_type=='video' else 'image') # [-301, -300, -302, -305]
        
        # print(flatten_patches.shape, len(images))
        return torch.tensor(flatten_patches), torch.tensor([[grid_t, grid_h, grid_w]]), image_placeholders
    
    def encode(self, pixel_values, grid_thws):
        output = self.backbone(pixel_values, grid_thws, output_hidden_states=True, return_dict=True)
        features = output.hidden_states[-1]
        # default: false
        # if self.config.drop_cls_token:
        #     features = features[:, 1:, :]
        
        # refer to qwen2.5-vl patchmerger
        seq_len, _ = features.shape
        features = features.reshape(seq_len//(self.config.hidden_stride ** 2), -1)
        
        return features

    def forward(self, pixel_values, grid_thws) -> torch.Tensor:  # [BatchSize, ImageShape] -> [BatchSize, #Token, VocabSize]
        features = self.encode(pixel_values, grid_thws)
        logits = self.head(features)
        tokens = self.tokenize(logits)
        # tokens' shape is [#Token, VocabSize-5], so padding with [#Token, 5], after
        # which, tokens' shape should become [#Token, VocabSize];
        # this is different from original aimv2 which has [BatchSize, #Token, VocabSize-5]
        token_len, _ = tokens.shape
        padding_tensor = torch.zeros(size=(token_len, len(IMAGE_INDICATOR_IDS)),
                                     dtype=tokens.dtype,
                                     device=tokens.device,
                                     layout=tokens.layout,
                                     requires_grad=False)
        tokens = torch.cat((tokens, padding_tensor), dim=1)
        return tokens

AutoModel.register(Aimv2VisualTokenizerConfig, Aimv2VisualTokenizer)




# ----------------------------------------------------------------------
#                           Visual Generator
# ----------------------------------------------------------------------
from .configuration_yak import YakConfig
from .modeling_yak import YakModel
AutoConfig.register("yak", YakConfig)
AutoModel.register(YakConfig, YakModel)



# ----------------------------------------------------------------------
#                               OvisU1
# ----------------------------------------------------------------------
class VisualEmbedding(torch.nn.Embedding):
    def forward(self, visual_tokens: Tensor) -> Tensor:
        if visual_tokens.dtype in [torch.int8, torch.int16, torch.int32, torch.int64, torch.long]:
            return super().forward(visual_tokens)
        return torch.matmul(visual_tokens, self.weight)

    def reset_parameters(self, mean=0., std=1.) -> None:
        init.normal_(self.weight, mean=mean, std=std)
        self._fill_padding_idx_with_zero()


class OvisU1PreTrainedModel(PreTrainedModel):
    config_class = OvisU1Config
    base_model_prefix = "ovis_u1"


class OvisU1(OvisU1PreTrainedModel):
    
    def __init__(self, config: OvisU1Config, *inputs, **kwargs):
        super().__init__(config, *inputs, **kwargs)
        attn_kwargs = dict()
        if self.config.llm_attn_implementation:
            attn_kwargs['attn_implementation'] = self.config.llm_attn_implementation
        self.llm = AutoModelForCausalLM.from_config(self.config.llm_config, **attn_kwargs)
        assert self.config.hidden_size == self.llm.config.hidden_size, "hidden size mismatch"
        self.text_tokenizer = AutoTokenizer.from_pretrained(self.config.name_or_path)
        self.visual_tokenizer = AutoModel.from_config(self.config.visual_tokenizer_config,
                                                    image_processor_name_or_path=self.config.name_or_path)
        self.visual_generator = AutoModel.from_config(self.config.visual_generator_config)
        self.vte = VisualEmbedding(self.config.visual_tokenizer_config.vocab_size, self.config.hidden_size,
                                    device=self.visual_tokenizer.device, dtype=self.visual_tokenizer.dtype)

        def _merge_modules(modules_list: tuple):
            merged_modules = []
            for modules in modules_list:
                merged_modules.extend(modules if modules else [])
            return merged_modules

        self._no_split_modules = _merge_modules((self.llm._no_split_modules, self.visual_tokenizer._no_split_modules))
        self._skip_keys_device_placement = self.llm._skip_keys_device_placement
        self._keep_in_fp32_modules = _merge_modules(
            (self.llm._keep_in_fp32_modules, self.visual_tokenizer._keep_in_fp32_modules))
        self._supports_flash_attn_2 = True
        self.is_parallelizable = all((self.llm.is_parallelizable, self.visual_tokenizer.is_parallelizable, self.visual_generator.is_parallelizable))
        self.supports_gradient_checkpointing = all(
            (self.llm.supports_gradient_checkpointing, self.visual_tokenizer.supports_gradient_checkpointing, self.visual_generator.supports_gradient_checkpointing))
        self._supports_sdpa = all((self.llm._supports_sdpa, self.visual_tokenizer._supports_sdpa, self.visual_generator._supports_sdpa))

    def get_text_tokenizer(self):
        return self.text_tokenizer

    def get_visual_tokenizer(self):
        return self.visual_tokenizer
    
    def get_visual_generator(self):
        return self.visual_generator

    def tie_weights(self):
        if not self.config.disable_tie_weight:
            self.get_llm().tie_weights()

    def get_lm_head(self):
        return self.get_llm().get_output_embeddings()

    def get_llm(self):
        return self.llm

    def get_vte(self):
        return self.vte

    def get_wte(self):
        return self.llm.get_input_embeddings()

    def get_conversation_formatter(self) -> ConversationFormatter:
        if getattr(self, 'conversation_formatter', None) is None:
            self.conversation_formatter = getattr(import_module(".configuration_ovis_u1", __package__),
                                                  self.config.conversation_formatter_class)(self.text_tokenizer)
        return self.conversation_formatter

    def merge_multimodal(
            self,
            text_input_ids: torch.Tensor,
            text_attention_masks: torch.Tensor,
            text_labels: Optional[torch.Tensor],
            pixel_values: Optional[torch.Tensor],
            grid_thws: Optional[torch.Tensor],
            left_padding: bool = False
    ):
        input_device = text_input_ids.device
        visual_vocab_szie = self.get_visual_tokenizer().config.vocab_size
        visual_indicator_embeds = self.get_vte()(
            torch.tensor(
                list(range(visual_vocab_szie - 5, visual_vocab_szie)),
                dtype=torch.long,
                device=self.get_visual_tokenizer().device
            )
        ).to(device=input_device)

        if self.training:
            # When training, to be compatible with deepspeed zero, each sample has to include pixel_value tensor.
            # For text-only sample, one can simply use a full zero tensor as pixel_value, which will be ignored
            # (see below in this function); so, the gradient will not be affected.
            num_images = [x.prod() // (self.visual_tokenizer.config.hidden_stride**2) for x in grid_thws]
            
            visual_tokens = self.visual_tokenizer(pixel_values, grid_thws)

            visual_embeds_ = torch.split(self.get_vte()(visual_tokens).to(dtype=self.dtype, device=input_device),
                                        split_size_or_sections=num_images, dim=0)
            


            visual_input_ids_ = torch.split(torch.argmax(visual_tokens, dim=-1).to(device=input_device),
                                           split_size_or_sections=num_images, dim=0)


            visual_labels_ = [torch.full(x.shape, IGNORE_ID, dtype=torch.long, device=input_device) for x in
                             visual_input_ids_]

            
            visual_embeds = []
            visual_input_ids = []
            visual_labels = []
            ind = 0
            for text_input_id in text_input_ids:
                image_atom_positions = torch.where(torch.eq(text_input_id, IMAGE_ATOM_ID))[0].tolist()
                n = len(image_atom_positions)
                if n > 0:
                    visual_embeds.append(visual_embeds_[ind:ind+n])
                    visual_input_ids.append(visual_input_ids_[ind:ind+n])
                    visual_labels.append(visual_labels_[ind:ind+n])
                    ind += n
                else:
                    visual_embeds.append(visual_embeds_[ind:ind+1])
                    visual_input_ids.append(visual_input_ids_[ind:ind+1])
                    visual_labels.append(visual_labels_[ind:ind+1])
                    ind += 1
                

        else:
            # TODO: Not modified yet
            # When inference, sample can include only text with `None` pixel_value
            # num_images = [x.shape[0] if x is not None else 0 for x in pixel_values]
            num_images = [x.prod() // (self.visual_tokenizer.config.hidden_stride**2) if x is not None else 0 for x in grid_thws]
            if sum(num_images) > 0:
                visual_tokens = self.visual_tokenizer(pixel_values, grid_thws)
                try:
                    visual_embeds_ = torch.split(self.get_vte()(visual_tokens).to(dtype=self.dtype, device=input_device),
                                        split_size_or_sections=num_images, dim=0)
                except Exception as e:
                    print(e)
                    print(pixel_values.shape, grid_thws.shape, visual_tokens.shape, num_images)
            

                visual_input_ids_ = torch.split(torch.argmax(visual_tokens, dim=-1).to(device=input_device),
                                            split_size_or_sections=num_images, dim=0)


                visual_labels_ = [torch.full(x.shape, IGNORE_ID, dtype=torch.long, device=input_device) for x in
                                visual_input_ids_]
                
                visual_embeds = []
                visual_input_ids = []
                visual_labels = []
                ind = 0
                for text_input_id in text_input_ids:
                    image_atom_positions = torch.where(torch.eq(text_input_id, IMAGE_ATOM_ID))[0].tolist()
                    n = len(image_atom_positions)
                    if n > 0:
                        visual_embeds.append(visual_embeds_[ind:ind+n])
                        visual_input_ids.append(visual_input_ids_[ind:ind+n])
                        visual_labels.append(visual_labels_[ind:ind+n])
                        ind += n
                    else:
                        visual_embeds.append(visual_embeds_[ind:ind+1])
                        visual_input_ids.append(visual_input_ids_[ind:ind+1])
                        visual_labels.append(visual_labels_[ind:ind+1])
                        ind += 1
                        
            else:
                # just placeholders
                visual_embeds = [None] * len(num_images)
                visual_input_ids = [None] * len(num_images)
                visual_labels = [None] * len(num_images)
            
        # just placeholders
        if text_labels is None:
            text_labels = torch.full(text_input_ids.shape, IGNORE_ID, dtype=torch.long, device=input_device)

        input_embeds = []
        attention_masks = []
        labels = []
        input_img_poss = []
        for text_input_id, text_label, text_attention_mask, visual_embed, visual_input_id, visual_label in zip(
            text_input_ids, text_labels, text_attention_masks, visual_embeds, visual_input_ids, visual_labels
        ):
            placeholder_token_mask = torch.lt(text_input_id, 0)
            text_embed = self.get_wte()(torch.masked_fill(text_input_id, placeholder_token_mask, 0))
            for i, indicator_id in enumerate(IMAGE_INDICATOR_IDS):
                text_embed[text_input_id == indicator_id] = visual_indicator_embeds[i]
            image_atom_positions = torch.where(torch.eq(text_input_id, IMAGE_ATOM_ID))[0].tolist()
            if len(image_atom_positions) > 0:
                input_embed_parts = []
                attention_mask_parts = []
                label_parts = []
                input_img_pos_parts = []
                prev_image_atom_position = -1
                for index, image_atom_position in enumerate(image_atom_positions):
                    input_embed_parts.append(
                        text_embed[prev_image_atom_position + 1:image_atom_position, :])
                    label_parts.append(
                        text_label[prev_image_atom_position + 1:image_atom_position])
                    input_img_pos_parts.append(
                        torch.zeros_like(text_label[prev_image_atom_position + 1:image_atom_position])
                    )
                    attention_mask_parts.append(
                        text_attention_mask[prev_image_atom_position + 1:image_atom_position])
                    input_embed_parts.append(visual_embed[index])
                    attention_mask_parts.append(
                        torch.ones_like(visual_label[index], dtype=torch.bool))
                    label_parts.append(visual_label[index])
                    input_img_pos_parts.append(
                        torch.ones_like(visual_label[index])
                    )
                    prev_image_atom_position = image_atom_position
                if prev_image_atom_position + 1 < text_input_id.shape[0]:
                    input_embed_parts.append(
                        text_embed[prev_image_atom_position + 1:, :])
                    attention_mask_parts.append(
                        text_attention_mask[prev_image_atom_position + 1:])
                    label_parts.append(
                        text_label[prev_image_atom_position + 1:])
                    input_img_pos_parts.append(
                        torch.zeros_like(text_label[prev_image_atom_position + 1:])
                    )
                input_embed = torch.cat(input_embed_parts, dim=0)
                attention_mask = torch.cat(attention_mask_parts, dim=0)
                label = torch.cat(label_parts, dim=0)
                input_img_pos = torch.cat(input_img_pos_parts, dim=0)
            else:
                input_embed = text_embed
                attention_mask = text_attention_mask
                label = text_label
                input_img_pos = torch.zeros_like(text_label)
                if self.training:
                    # Make visual_embed & visual_indicator_embeds involved in the backward graph,
                    # to be compatible with deepspeed zero and ddp.
                    input_embed += torch.sum(visual_embed[0] * 0.0) + torch.sum(visual_indicator_embeds * 0.0)
            input_embeds.append(input_embed)
            attention_masks.append(attention_mask)
            labels.append(label)
            input_img_poss.append(input_img_pos)

        batch_input_embeds = self.pad_truncate_sequence(input_embeds, batch_first=True, padding_value=0.0, left_padding=left_padding)
        batch_attention_mask = self.pad_truncate_sequence(attention_masks, batch_first=True, padding_value=False, left_padding=left_padding)
        batch_labels = self.pad_truncate_sequence(labels, batch_first=True, padding_value=IGNORE_ID, left_padding=left_padding)
        batch_input_img_labels = self.pad_truncate_sequence(input_img_poss, batch_first=True, padding_value=0.0, left_padding=left_padding)

        return visual_input_ids, batch_input_embeds, batch_labels, batch_attention_mask, batch_input_img_labels

    def pad_truncate_sequence(self, sequences: List[torch.Tensor], batch_first: bool = True, padding_value: float = 0.0, left_padding: bool = False) -> torch.Tensor:
        if left_padding == False:
            pad_sequence = torch.nn.utils.rnn.pad_sequence(sequences, batch_first=batch_first, padding_value=padding_value)
            return pad_sequence[:,:self.config.multimodal_max_length]
        else:
            pad_sequence = torch.nn.utils.rnn.pad_sequence([i.flip(dims=[0]) for i in sequences],batch_first=True, padding_value=padding_value).flip(dims=[1])
            return pad_sequence[:,-self.config.multimodal_max_length:]

    def preprocess_inputs(
        self,
        text_or_conversations: Union[List[Dict], str],
        images: Optional[Union[List[PIL.Image.Image], List[List[PIL.Image.Image]]]],
        generation_preface='',
        return_labels=False,
        propagate_exception=True,
        frame_selector=None,
        multimodal_type="single_image",
        fix_sample_overall_length_navit=False,
        min_pixels=None,
        max_pixels=None,
        enable_thinking=False
    ):
        # convert text to conversations
        if isinstance(text_or_conversations, str):
            conversations = [{
                "from": "human",
                "value": text_or_conversations
            }]
        elif isinstance(text_or_conversations, list):
            conversations = text_or_conversations
        else:
            raise ValueError(f'[{datetime.now()}] Invalid type of `text_or_conversations`, expected `List[Dict]` or `str`,'
                             f' but got {type(text_or_conversations)}')

        if frame_selector is not None:
            conversations, images = frame_selector(conversations=conversations,frames=images,clear_prompt=True)

        # format conversations
        prompt, raw_input_ids, raw_labels = self.get_conversation_formatter().format(
            conversations, generation_preface=generation_preface, enable_thinking=enable_thinking)

        # place image placeholders
        input_ids = []
        labels = []
        pixel_values = []
        grid_thws = []
        invalidate_label = False
        image_token_indices = [i for i, v in enumerate(raw_input_ids) if v == IMAGE_TOKEN_ID or v == VIDEO_TOKEN_ID]
        last_image_token_index = -1
        for i in range(len(image_token_indices)):
            head = 0 if i == 0 else image_token_indices[i - 1] + 1
            tail = image_token_indices[i]
            last_image_token_index = tail
            input_ids.extend(raw_input_ids[head:tail])
            labels.extend(raw_labels[head:tail])
            try:
                # currently, do not support multiple videos
                if multimodal_type == "video":
                    image = images
                else:
                    image = images[i]
                raw_pixel_values, image_grid_thws, image_placeholders = self.visual_tokenizer.preprocess_image(
                    image, num_images=len(images) if fix_sample_overall_length_navit else 1, min_pixels=min_pixels, max_pixels=max_pixels,
                    multimodal_type=multimodal_type)
            except Exception as e:
                if propagate_exception:
                    raise e
                logging.exception(e)
                invalidate_label = True
                # raw_pixel_values, image_placeholders = self.visual_tokenizer.mock_input() # TODO
                raw_pixel_values, _ = self.visual_tokenizer.mock_input()
                mock_image = transforms.ToPILImage()(raw_pixel_values[0])
                raw_pixel_values, image_grid_thws, image_placeholders = self.visual_tokenizer.preprocess_image(
                            mock_image, min_pixels=min_pixels, max_pixels=max_pixels)
                
            input_ids.extend(image_placeholders)
            labels.extend([IGNORE_ID] * len(image_placeholders))
            pixel_values.append(raw_pixel_values)
            grid_thws.append(image_grid_thws)
        input_ids.extend(raw_input_ids[last_image_token_index + 1:])
        labels.extend(raw_labels[last_image_token_index + 1:])

        # return tensors
        input_ids = torch.tensor(input_ids, dtype=torch.long)
        labels = torch.tensor([IGNORE_ID] * len(labels) if invalidate_label else labels, dtype=torch.long)
        pixel_values = torch.cat(pixel_values, dim=0) if len(pixel_values) > 0 else None
        grid_thws = torch.cat(grid_thws, dim=0) if len(grid_thws) > 0 else None

        if return_labels:
            return prompt, input_ids, pixel_values, grid_thws, labels
        else:
            return prompt, input_ids, pixel_values, grid_thws

    def generate(
        self,
        inputs: Optional[torch.Tensor] = None,
        **kwargs,
    ) -> Union[GenerateOutput, torch.LongTensor]:
        # assert inputs.shape[0] == 1, 'Currently, only support `batch_size=1`'
        _, inputs_embeds, labels, attention_mask, input_img_labels = self.merge_multimodal(
            text_input_ids=inputs,
            text_attention_masks=kwargs.pop('attention_mask'),
            text_labels=None,
            pixel_values=kwargs.pop('pixel_values'),
            grid_thws=kwargs.pop('grid_thws'),
            left_padding=True
        )
        inputs_embeds = inputs_embeds.detach()
        torch.cuda.empty_cache()
        return self.llm.generate(inputs=None, inputs_embeds=inputs_embeds, attention_mask=attention_mask, **kwargs)

    def generate_condition(
            self,
            inputs: Optional[torch.Tensor] = None,
            **kwargs,
    ):
        # assert inputs.shape[0] == 1, 'Currently, only support `batch_size=1`'
        _, inputs_embeds, labels, attention_mask, input_img_labels = self.merge_multimodal(
            text_input_ids=inputs,
            text_attention_masks=kwargs.pop('attention_mask'),
            text_labels=None,
            pixel_values=kwargs.pop('pixel_values'),
            grid_thws=kwargs.pop('grid_thws'),
            left_padding=True
        )
        inputs_embeds = inputs_embeds.detach()
        torch.cuda.empty_cache()
        device = self.llm.device
        outputs = self.llm(inputs_embeds=inputs_embeds.to(device), 
                            labels=labels.to(device), 
                            attention_mask=attention_mask.to(device), 
                            output_hidden_states=True, 
                            **kwargs)
        semantic_cond_0 = outputs.hidden_states[-1]
        semantic_cond_1 = outputs.hidden_states[-2]
        semantic_cond = torch.cat([semantic_cond_0, semantic_cond_1], dim=-1)
        return dict(
            txt=semantic_cond
        )
    
    def generate_img(
        self,
        inputs: Optional[torch.Tensor] = None,
        cond = None,
        no_both_cond = None,
        no_txt_cond = None,
        **kwargs,
    ) -> Union[GenerateOutput, torch.LongTensor]:
        if cond is None:
            cond = self.generate_condition(inputs, **kwargs)
        
        height = kwargs.get('height', 1024)
        width = kwargs.get('width', 1024)
        num_steps = kwargs.get('num_steps', 50)
        seed = kwargs.get('seed', 42)
        img_cfg = kwargs.pop('img_cfg', 1.5)
        txt_cfg = kwargs.pop('txt_cfg', 5)
        yak_output = self.visual_generator.generate_image(
            cond=cond, no_txt_cond=no_txt_cond, no_both_cond=no_both_cond,
            height=height, width=width, 
            num_steps=num_steps, seed=seed, 
            img_cfg=img_cfg, txt_cfg=txt_cfg,
            output_type="pil")
        return yak_output